92 lines
4.1 KiB
Python
92 lines
4.1 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
from sqlmodel import SQLModel, Session, create_engine, select
|
||
|
|
|
||
|
|
from app.agents.runtime import RuntimeContext, RuntimeTool, run_agent, set_tool_approval
|
||
|
|
from app.llm import tools_unsupported
|
||
|
|
from app.models import Agent, AgentSession, RunEvent
|
||
|
|
|
||
|
|
|
||
|
|
class AgentRuntimeTests(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.engine = create_engine("sqlite://")
|
||
|
|
SQLModel.metadata.create_all(self.engine)
|
||
|
|
self.db = Session(self.engine)
|
||
|
|
self.agent = Agent(id=101, name="测试 Agent", role="production")
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
self.db.close()
|
||
|
|
|
||
|
|
def test_tool_call_continues_to_final_output(self):
|
||
|
|
calls = []
|
||
|
|
|
||
|
|
def tool(_: RuntimeContext, args: dict) -> dict:
|
||
|
|
return {"ok": True, "value": args["query"]}
|
||
|
|
|
||
|
|
def model_call(_db, messages, *_args, **_kwargs):
|
||
|
|
calls.append(messages)
|
||
|
|
if any(item.get("role") == "tool" for item in messages):
|
||
|
|
return {"status": "ok", "text": "最终正文", "tool_calls": []}
|
||
|
|
return {"status": "ok", "text": "", "tool_calls": [{
|
||
|
|
"id": "call-1", "type": "function",
|
||
|
|
"function": {"name": "knowledge_search", "arguments": '{"query":"钩子"}'},
|
||
|
|
}]}
|
||
|
|
|
||
|
|
result = run_agent(
|
||
|
|
self.db, run_id=11, campaign_id=1, segment_id=2, node_id="n1", agent=self.agent,
|
||
|
|
provider_id=1, model="test-model", phase="production",
|
||
|
|
messages=[{"role": "user", "content": "生成一段内容"}],
|
||
|
|
tools=[RuntimeTool("knowledge_search", "检索", {
|
||
|
|
"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"],
|
||
|
|
}, tool)],
|
||
|
|
model_call=model_call,
|
||
|
|
prompt_meta={"template_id": 7, "template_name": "测试模板", "template_version": 3},
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertEqual(result["status"], "completed")
|
||
|
|
self.assertEqual(result["text"], "最终正文")
|
||
|
|
self.assertEqual(len(calls), 2)
|
||
|
|
row = self.db.exec(select(AgentSession)).one()
|
||
|
|
self.assertEqual(row.status, "completed")
|
||
|
|
self.assertEqual(row.turn, 2)
|
||
|
|
self.assertEqual(row.result_meta["prompt"]["template_version"], 3)
|
||
|
|
started = self.db.exec(select(RunEvent).where(RunEvent.event_type == "model_started")).all()
|
||
|
|
self.assertEqual(len(started), 2)
|
||
|
|
self.assertEqual(started[0].payload["prompt"]["template_id"], 7)
|
||
|
|
|
||
|
|
def test_approval_waits_then_resumes_same_session(self):
|
||
|
|
def guarded(_: RuntimeContext, __: dict) -> dict:
|
||
|
|
return {"ok": True, "approved_result": "done"}
|
||
|
|
|
||
|
|
def model_call(_db, messages, *_args, **_kwargs):
|
||
|
|
if any(item.get("role") == "tool" for item in messages):
|
||
|
|
return {"status": "ok", "text": "审批后完成", "tool_calls": []}
|
||
|
|
return {"status": "ok", "text": "", "tool_calls": [{
|
||
|
|
"id": "call-approve", "type": "function",
|
||
|
|
"function": {"name": "guarded_tool", "arguments": "{}"},
|
||
|
|
}]}
|
||
|
|
|
||
|
|
args = dict(
|
||
|
|
run_id=12, campaign_id=1, segment_id=3, node_id="n2", agent=self.agent,
|
||
|
|
provider_id=1, model="test-model", phase="production",
|
||
|
|
messages=[{"role": "user", "content": "需要审批的任务"}],
|
||
|
|
tools=[RuntimeTool("guarded_tool", "需要审批", {"type": "object", "properties": {}},
|
||
|
|
guarded, requires_approval=True)],
|
||
|
|
model_call=model_call,
|
||
|
|
)
|
||
|
|
first = run_agent(self.db, **args)
|
||
|
|
self.assertEqual(first["status"], "waiting_approval")
|
||
|
|
set_tool_approval(self.db, first["session_id"], True)
|
||
|
|
resumed = run_agent(self.db, **args)
|
||
|
|
self.assertEqual(resumed["status"], "completed")
|
||
|
|
self.assertEqual(resumed["text"], "审批后完成")
|
||
|
|
self.assertEqual(resumed["session_id"], first["session_id"])
|
||
|
|
|
||
|
|
def test_provider_without_tools_is_detected_for_fallback(self):
|
||
|
|
self.assertTrue(tools_unsupported(Exception("400 tools are not supported by this provider")))
|
||
|
|
self.assertFalse(tools_unsupported(Exception("401 invalid api key")))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|