增加工具使用功能
This commit is contained in:
@@ -0,0 +1,202 @@
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from service import ChatProviderError, ChatService
|
||||
from storage import JsonSessionStorage
|
||||
|
||||
|
||||
def completion(
|
||||
*,
|
||||
content: str | None,
|
||||
tool_calls: list[SimpleNamespace] | None = None,
|
||||
prompt_tokens: int = 10,
|
||||
completion_tokens: int = 2,
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
choices=[
|
||||
SimpleNamespace(
|
||||
message=SimpleNamespace(content=content, tool_calls=tool_calls)
|
||||
)
|
||||
],
|
||||
usage=SimpleNamespace(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def tool_call(call_id: str, name: str, arguments: str) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
id=call_id,
|
||||
type="function",
|
||||
function=SimpleNamespace(name=name, arguments=arguments),
|
||||
)
|
||||
|
||||
|
||||
class ScriptedClient:
|
||||
def __init__(self, responses: list[SimpleNamespace]) -> None:
|
||||
self.responses = responses
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
self.chat = SimpleNamespace(
|
||||
completions=SimpleNamespace(create=self.create)
|
||||
)
|
||||
|
||||
async def create(self, **kwargs: Any) -> SimpleNamespace:
|
||||
self.calls.append(kwargs)
|
||||
return self.responses.pop(0)
|
||||
|
||||
|
||||
def test_agent_executes_tool_and_persists_full_history(tmp_path: Path) -> None:
|
||||
async def scenario() -> None:
|
||||
provider = ScriptedClient(
|
||||
[
|
||||
completion(
|
||||
content=None,
|
||||
tool_calls=[tool_call("call-1", "calculate", '{"expression":"2+2"}')],
|
||||
),
|
||||
completion(content="结果是 4。", 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 等于多少？")
|
||||
|
||||
assert response.message.content == "结果是 4。"
|
||||
assert "tool_calls" not in response.message.model_dump()
|
||||
assert response.tools_use == ["calculate"]
|
||||
assert response.usage is not None
|
||||
assert response.usage.model_dump() == {
|
||||
"prompt_tokens": 30,
|
||||
"completion_tokens": 5,
|
||||
"total_tokens": 35,
|
||||
}
|
||||
assert "tools" in provider.calls[0]
|
||||
second_messages = provider.calls[1]["messages"]
|
||||
assert [message["role"] for message in second_messages] == [
|
||||
"system",
|
||||
"user",
|
||||
"assistant",
|
||||
"tool",
|
||||
]
|
||||
assert json.loads(second_messages[-1]["content"])["result"] == 4
|
||||
|
||||
saved = await storage.read(session.session_id)
|
||||
assert [message.role for message in saved.messages] == [
|
||||
"user",
|
||||
"assistant",
|
||||
"tool",
|
||||
"assistant",
|
||||
]
|
||||
assert saved.messages[1].tool_calls[0].id == "call-1" # type: ignore[union-attr]
|
||||
assert saved.messages[2].tool_call_id == "call-1" # type: ignore[union-attr]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_agent_executes_multiple_tools_in_order(tmp_path: Path) -> None:
|
||||
async def scenario() -> None:
|
||||
provider = ScriptedClient(
|
||||
[
|
||||
completion(
|
||||
content=None,
|
||||
tool_calls=[
|
||||
tool_call("call-1", "calculate", '{"expression":"6*7"}'),
|
||||
tool_call(
|
||||
"call-2",
|
||||
"get_current_time",
|
||||
'{"timezone":"UTC"}',
|
||||
),
|
||||
],
|
||||
),
|
||||
completion(content="完成"),
|
||||
]
|
||||
)
|
||||
storage = JsonSessionStorage(tmp_path)
|
||||
session = await storage.create("system")
|
||||
service = ChatService(storage, provider, "test-model") # type: ignore[arg-type]
|
||||
|
||||
response = await service.generate_response(session.session_id, "计算并查询时间")
|
||||
|
||||
messages = provider.calls[1]["messages"]
|
||||
assert [message.get("tool_call_id") for message in messages[-2:]] == [
|
||||
"call-1",
|
||||
"call-2",
|
||||
]
|
||||
assert response.tools_use == ["calculate", "get_current_time"]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_tool_error_is_returned_to_model(tmp_path: Path) -> None:
|
||||
async def scenario() -> None:
|
||||
provider = ScriptedClient(
|
||||
[
|
||||
completion(
|
||||
content=None,
|
||||
tool_calls=[tool_call("bad-1", "calculate", "not-json")],
|
||||
),
|
||||
completion(content="参数无效。"),
|
||||
]
|
||||
)
|
||||
storage = JsonSessionStorage(tmp_path)
|
||||
session = await storage.create("system")
|
||||
service = ChatService(storage, provider, "test-model") # type: ignore[arg-type]
|
||||
|
||||
await service.generate_response(session.session_id, "test")
|
||||
|
||||
saved = await storage.read(session.session_id)
|
||||
tool_message = saved.messages[2]
|
||||
assert tool_message.role == "tool"
|
||||
assert tool_message.is_error is True # type: ignore[union-attr]
|
||||
assert "invalid_arguments" in tool_message.content
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("max_rounds", "max_calls", "calls"),
|
||||
[
|
||||
(1, 10, [tool_call("call-1", "calculate", '{"expression":"1+1"}')]),
|
||||
(
|
||||
5,
|
||||
1,
|
||||
[
|
||||
tool_call("call-1", "calculate", '{"expression":"1+1"}'),
|
||||
tool_call("call-2", "calculate", '{"expression":"2+2"}'),
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_agent_limits_do_not_persist_partial_turn(
|
||||
tmp_path: Path,
|
||||
max_rounds: int,
|
||||
max_calls: int,
|
||||
calls: list[SimpleNamespace],
|
||||
) -> None:
|
||||
async def scenario() -> None:
|
||||
provider = ScriptedClient([completion(content=None, tool_calls=calls)])
|
||||
storage = JsonSessionStorage(tmp_path)
|
||||
session = await storage.create("system")
|
||||
service = ChatService(
|
||||
storage,
|
||||
provider, # type: ignore[arg-type]
|
||||
"test-model",
|
||||
max_tool_rounds=max_rounds,
|
||||
max_tool_calls_per_turn=max_calls,
|
||||
)
|
||||
|
||||
with pytest.raises(ChatProviderError):
|
||||
await service.generate_response(session.session_id, "test")
|
||||
|
||||
saved = await storage.read(session.session_id)
|
||||
assert saved.messages == []
|
||||
|
||||
asyncio.run(scenario())
|
||||
@@ -136,7 +136,9 @@ def test_multi_turn_request_contains_full_history(
|
||||
"completion_tokens": 2,
|
||||
"total_tokens": 12,
|
||||
}
|
||||
assert second.json()["tools_use"] == []
|
||||
datetime.fromisoformat(second.json()["message"]["created_at"])
|
||||
assert "tool_calls" not in second.json()["message"]
|
||||
data = json.loads((tmp_path / f"{session_id}.json").read_text())
|
||||
assert [message["content"] for message in data["messages"]] == [
|
||||
"first",
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
|
||||
from tools import CalculatorArguments, ToolRegistry, ToolSpec
|
||||
|
||||
|
||||
def test_builtin_tool_definitions() -> None:
|
||||
registry = ToolRegistry()
|
||||
|
||||
definitions = registry.definitions()
|
||||
|
||||
assert [item["function"]["name"] for item in definitions] == [ # type: ignore[index]
|
||||
"get_current_time",
|
||||
"calculate",
|
||||
]
|
||||
|
||||
|
||||
def test_calculator_and_time_tools() -> 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:
|
||||
async def scenario() -> None:
|
||||
registry = ToolRegistry()
|
||||
invalid_json = await registry.execute("calculate", "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"}'
|
||||
)
|
||||
|
||||
assert invalid_json.is_error is True
|
||||
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
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_tool_timeout_returns_error() -> None:
|
||||
def slow_handler(arguments):
|
||||
time.sleep(0.05)
|
||||
return {"result": 1}
|
||||
|
||||
async def scenario() -> None:
|
||||
registry = ToolRegistry(timeout_seconds=0.01)
|
||||
registry._specs["calculate"] = ToolSpec(
|
||||
name="calculate",
|
||||
description="slow",
|
||||
arguments_model=CalculatorArguments,
|
||||
handler=slow_handler,
|
||||
)
|
||||
|
||||
result = await registry.execute("calculate", '{"expression":"1+1"}')
|
||||
|
||||
assert result.is_error is True
|
||||
assert "timeout" in result.content
|
||||
|
||||
asyncio.run(scenario())
|
||||
Reference in New Issue
Block a user