54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
|
from types import SimpleNamespace
|
||
|
|
|
||
|
|
from auto_reverse.agent import Agent
|
||
|
|
|
||
|
|
|
||
|
|
def _text_block(text):
|
||
|
|
return SimpleNamespace(type="text", text=text)
|
||
|
|
|
||
|
|
|
||
|
|
def _tool_use(tool_id, name, inp):
|
||
|
|
return SimpleNamespace(type="tool_use", id=tool_id, name=name, input=inp)
|
||
|
|
|
||
|
|
|
||
|
|
class FakeMessages:
|
||
|
|
def __init__(self, scripted):
|
||
|
|
self._scripted = list(scripted)
|
||
|
|
self.calls = []
|
||
|
|
|
||
|
|
def create(self, **kwargs):
|
||
|
|
self.calls.append(kwargs)
|
||
|
|
content = self._scripted.pop(0)
|
||
|
|
stop = "tool_use" if any(b.type == "tool_use" for b in content) else "end_turn"
|
||
|
|
return SimpleNamespace(content=content, stop_reason=stop, role="assistant")
|
||
|
|
|
||
|
|
|
||
|
|
class FakeClient:
|
||
|
|
def __init__(self, scripted):
|
||
|
|
self.messages = FakeMessages(scripted)
|
||
|
|
|
||
|
|
|
||
|
|
def test_agent_executes_tool_then_returns_text():
|
||
|
|
scripted = [
|
||
|
|
[_tool_use("t1", "flows_search", {"query": "users"})],
|
||
|
|
[_text_block("Found the users endpoint.")],
|
||
|
|
]
|
||
|
|
client = FakeClient(scripted)
|
||
|
|
registry = {
|
||
|
|
"flows_search": (
|
||
|
|
{"name": "flows_search", "input_schema": {"type": "object"}},
|
||
|
|
lambda inp: {"endpoints": [{"path": "/api/users"}]},
|
||
|
|
)
|
||
|
|
}
|
||
|
|
agent = Agent(client, registry, model="m", system="s")
|
||
|
|
reply = agent.run_turn("map users")
|
||
|
|
assert "users endpoint" in reply
|
||
|
|
# the tool result was fed back: second create call has >= 3 messages
|
||
|
|
assert len(client.messages.calls[1]["messages"]) >= 3
|
||
|
|
|
||
|
|
|
||
|
|
def test_agent_plain_text_no_tools():
|
||
|
|
client = FakeClient([[_text_block("Hello!")]])
|
||
|
|
agent = Agent(client, {}, model="m", system="s")
|
||
|
|
assert agent.run_turn("hi") == "Hello!"
|