Mplan/simple-chat-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化目录结构

Browse Source
This commit is contained in:
Mplan
2026-07-03 20:49:36 +08:00
parent 04fefcf488
commit 8073dd63a7
10 changed files with 220 additions and 238 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_agent.py
+31 -13
View File
@@ -58,20 +58,26 @@ def test_agent_executes_tool_and_persists_full_history(tmp_path: Path) -> None:
[
completion(
content=None,
tool_calls=[tool_call("call-1", "calculate", '{"expression":"2+2"}')],
tool_calls=[
tool_call(
"call-1",
"get_current_time",
'{"timezone":"UTC"}',
)
],
),
completion(content="结果是 4。", prompt_tokens=20, completion_tokens=3),
completion(content="当前时间已查询。", prompt_tokens=20, completion_tokens=3),
]
)
storage = JsonSessionStorage(tmp_path)
session = await storage.create("Use tools when needed.")
service = ChatService(storage, provider, "test-model") # type: ignore[arg-type]
response = await service.generate_response(session.session_id, "2+2 等于多少？")
response = await service.generate_response(session.session_id, "现在几点？")
assert response.message.content == "结果是 4。"
assert response.message.content == "当前时间已查询。"
assert "tool_calls" not in response.message.model_dump()
assert response.tools_use == ["calculate"]
assert response.tools_use == ["get_current_time"]
assert response.usage is not None
assert response.usage.model_dump() == {
"prompt_tokens": 30,
@@ -86,7 +92,7 @@ def test_agent_executes_tool_and_persists_full_history(tmp_path: Path) -> None:
"assistant",
"tool",
]
assert json.loads(second_messages[-1]["content"])["result"] == 4
assert json.loads(second_messages[-1]["content"])["timezone"] == "UTC"
saved = await storage.read(session.session_id)
assert [message.role for message in saved.messages] == [
@@ -108,11 +114,15 @@ def test_agent_executes_multiple_tools_in_order(tmp_path: Path) -> None:
completion(
content=None,
tool_calls=[
tool_call("call-1", "calculate", '{"expression":"6*7"}'),
tool_call(
"call-1",
"get_current_time",
'{"timezone":"UTC"}',
),
tool_call(
"call-2",
"get_current_time",
'{"timezone":"UTC"}',
'{"timezone":"Asia/Shanghai"}',
),
],
),
@@ -130,7 +140,7 @@ def test_agent_executes_multiple_tools_in_order(tmp_path: Path) -> None:
"call-1",
"call-2",
]
assert response.tools_use == ["calculate", "get_current_time"]
assert response.tools_use == ["get_current_time"]
asyncio.run(scenario())
@@ -141,7 +151,7 @@ def test_tool_error_is_returned_to_model(tmp_path: Path) -> None:
[
completion(
content=None,
tool_calls=[tool_call("bad-1", "calculate", "not-json")],
tool_calls=[tool_call("bad-1", "get_current_time", "not-json")],
),
completion(content="参数无效。"),
]
@@ -164,13 +174,21 @@ def test_tool_error_is_returned_to_model(tmp_path: Path) -> None:
@pytest.mark.parametrize(
("max_rounds", "max_calls", "calls"),
[
(1, 10, [tool_call("call-1", "calculate", '{"expression":"1+1"}')]),
(
1,
10,
[tool_call("call-1", "get_current_time", '{"timezone":"UTC"}')],
),
(
5,
1,
[
tool_call("call-1", "calculate", '{"expression":"1+1"}'),
tool_call("call-2", "calculate", '{"expression":"2+2"}'),
tool_call("call-1", "get_current_time", '{"timezone":"UTC"}'),
tool_call(
"call-2",
"get_current_time",
'{"timezone":"Asia/Shanghai"}',
),
],
),
],
tests/test_tools.py
+12 -21
View File
@@ -2,7 +2,7 @@ import asyncio
import json
import time
from tools import CalculatorArguments, ToolRegistry, ToolSpec
from tools import CurrentTimeArguments, ToolRegistry, ToolSpec
def test_builtin_tool_definitions() -> None:
@@ -11,37 +11,28 @@ def test_builtin_tool_definitions() -> None:
definitions = registry.definitions()
assert [item["function"]["name"] for item in definitions] == [ # type: ignore[index]
"get_current_time",
"calculate",
"get_current_time"
]
def test_calculator_and_time_tools() -> None:
def test_current_time_tool() -> None:
async def scenario() -> None:
registry = ToolRegistry()
calculation = await registry.execute(
"calculate", '{"expression":"(2 + 3) * 4"}'
)
current_time = await registry.execute(
"get_current_time", '{"timezone":"Asia/Shanghai"}'
)
assert calculation.is_error is False
assert json.loads(calculation.content)["result"] == 20
assert current_time.is_error is False
assert json.loads(current_time.content)["timezone"] == "Asia/Shanghai"
asyncio.run(scenario())
def test_tool_argument_and_expression_errors_are_safe() -> None:
def test_tool_argument_and_execution_errors_are_safe() -> None:
async def scenario() -> None:
registry = ToolRegistry()
invalid_json = await registry.execute("calculate", "not-json")
invalid_json = await registry.execute("get_current_time", "not-json")
unknown = await registry.execute("missing", "{}")
unsafe = await registry.execute(
"calculate", '{"expression":"__import__(\\"os\\").system(\\"id\\")"}'
)
invalid_timezone = await registry.execute(
"get_current_time", '{"timezone":"Not/A_Real_Zone"}'
)
@@ -50,8 +41,6 @@ def test_tool_argument_and_expression_errors_are_safe() -> None:
assert "invalid_arguments" in invalid_json.content
assert unknown.is_error is True
assert "unknown_tool" in unknown.content
assert unsafe.is_error is True
assert "execution_failed" in unsafe.content
assert invalid_timezone.is_error is True
assert "execution_failed" in invalid_timezone.content
@@ -61,18 +50,20 @@ def test_tool_argument_and_expression_errors_are_safe() -> None:
def test_tool_timeout_returns_error() -> None:
def slow_handler(arguments):
time.sleep(0.05)
return {"result": 1}
return {"timezone": "UTC"}
async def scenario() -> None:
registry = ToolRegistry(timeout_seconds=0.01)
registry._specs["calculate"] = ToolSpec(
name="calculate",
registry._specs["get_current_time"] = ToolSpec(
name="get_current_time",
description="slow",
arguments_model=CalculatorArguments,
arguments_model=CurrentTimeArguments,
handler=slow_handler,
)
result = await registry.execute("calculate", '{"expression":"1+1"}')
result = await registry.execute(
"get_current_time", '{"timezone":"UTC"}'
)
assert result.is_error is True
assert "timeout" in result.content