Spaces:
Sleeping
Sleeping
| from types import SimpleNamespace | |
| from agent import graph | |
| def _part(text=None, thought=False, function_call=None): | |
| return SimpleNamespace(text=text, thought=thought, function_call=function_call) | |
| def _fc(name, args): | |
| return SimpleNamespace(name=name, args=args) | |
| def _response(parts): | |
| content = SimpleNamespace(role="model", parts=parts) | |
| return SimpleNamespace(candidates=[SimpleNamespace(content=content)]) | |
| def test_split_response_separates_thought_answer_calls(): | |
| resp = _response([ | |
| _part(text="thinking hard", thought=True), | |
| _part(function_call=_fc("calculator", {"expression": "2+2"})), | |
| ]) | |
| thought, answer, calls, content = graph.split_response(resp) | |
| assert thought == "thinking hard" | |
| assert answer == "" | |
| assert len(calls) == 1 and calls[0].name == "calculator" | |
| assert content.role == "model" | |
| def test_split_response_final_answer(): | |
| resp = _response([ | |
| _part(text="done thinking", thought=True), | |
| _part(text="The answer is 42.", thought=False), | |
| ]) | |
| thought, answer, calls, _ = graph.split_response(resp) | |
| assert thought == "done thinking" | |
| assert answer == "The answer is 42." | |
| assert calls == [] | |
| def test_classify_weak(): | |
| assert graph.classify_weak("wikipedia_search", "") | |
| assert graph.classify_weak("wikipedia_search", "tiny") | |
| assert graph.classify_weak("wikipedia_search", "France may refer to: the country") | |
| assert graph.classify_weak("wikipedia_search", "ERROR: boom") | |
| assert not graph.classify_weak("wikipedia_search", "Paris is the capital of France " * 5) | |
| assert not graph.classify_weak("calculator", "42") | |
| assert graph.classify_weak("calculator", "ERROR: bad expression") | |
| def test_signature_is_stable_and_distinguishes_args(): | |
| a = graph.signature(_fc("wikipedia_search", {"query": "France"})) | |
| b = graph.signature(_fc("wikipedia_search", {"query": "France"})) | |
| c = graph.signature(_fc("wikipedia_search", {"query": "Paris"})) | |
| assert a == b | |
| assert a != c | |
| class ScriptedClient: | |
| """Returns queued fake responses; mimics client.models.generate_content.""" | |
| def __init__(self, responses): | |
| self._responses = list(responses) | |
| self.models = self | |
| def generate_content(self, *, model, contents, config): | |
| return self._responses.pop(0) | |
| class LoopingClient: | |
| """Always returns the same response (to exercise the step cap).""" | |
| def __init__(self, response): | |
| self._response = response | |
| self.models = self | |
| def generate_content(self, *, model, contents, config): | |
| return self._response | |
| def _events(task, client, tool_fns, **kw): | |
| return list(graph.stream_run( | |
| task, client=client, tool_fns=tool_fns, | |
| declarations=[], system_prompt="sys", sleep=lambda s: None, **kw, | |
| )) | |
| def test_full_loop_shows_revision_and_final(): | |
| wiki_calls = {"n": 0} | |
| def fake_wiki(query): | |
| wiki_calls["n"] += 1 | |
| if wiki_calls["n"] == 1: | |
| return "France may refer to: the country, a film, ..." | |
| return "Paris is the capital of France. " * 6 | |
| tool_fns = {"wikipedia_search": fake_wiki, "calculator": lambda expression: "2140"} | |
| responses = [ | |
| _response([_part("Find the capital.", thought=True), | |
| _part(function_call=_fc("wikipedia_search", {"query": "France"}))]), | |
| _response([_part("That was ambiguous; searching more specifically.", thought=True), | |
| _part(function_call=_fc("wikipedia_search", {"query": "Paris capital population"}))]), | |
| _response([_part("Now compute.", thought=True), | |
| _part(function_call=_fc("calculator", {"expression": "2140000/1000"}))]), | |
| _response([_part("Done.", thought=True), _part("The answer is 2140.", thought=False)]), | |
| ] | |
| events = _events("q", ScriptedClient(responses), tool_fns) | |
| kinds = [e["kind"] for e in events] | |
| assert "thought" in kinds and "tool_call" in kinds and "observation" in kinds | |
| thoughts = [e for e in events if e["kind"] == "thought"] | |
| assert thoughts[0]["revision"] is False | |
| assert any(t["revision"] for t in thoughts) # step 2 thought is a revision | |
| obs = [e for e in events if e["kind"] == "observation"] | |
| assert obs[0]["weak"] is True | |
| assert obs[1]["weak"] is False | |
| finals = [e for e in events if e["kind"] == "final"] | |
| assert finals and finals[-1]["text"] == "The answer is 2140." | |
| assert len([e for e in events if e["kind"] == "tool_call"]) == 3 | |
| def test_step_cap_emits_limit(): | |
| looping = LoopingClient(_response([ | |
| _part("loop", thought=True), | |
| _part(function_call=_fc("calculator", {"expression": "1+1"})), | |
| ])) | |
| tool_fns = {"calculator": lambda expression: "2"} | |
| events = _events("q", looping, tool_fns, max_steps=3) | |
| assert any(e["kind"] == "limit" for e in events) | |
| assert len([e for e in events if e["kind"] == "tool_call"]) == 3 | |