初始化提交
This commit is contained in:
chenxin committed
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,405 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai_provider.types.chat import TokenUsage
|
||||
from app.ai_provider.types.messages import HumanMessage
|
||||
from app.application.agent import (
|
||||
AfterToolCallDecision,
|
||||
Agent,
|
||||
AgentConfig,
|
||||
AgentContext,
|
||||
AgentEvent,
|
||||
AgentRunContext,
|
||||
AgentTurnResult,
|
||||
BeforeToolCallDecision,
|
||||
ExtensionRegistry,
|
||||
PendingMessageQueue,
|
||||
QueueMode,
|
||||
ToolCallContext,
|
||||
TurnSnapshot,
|
||||
)
|
||||
from app.modules.tool_runtime.definition import ToolResult
|
||||
|
||||
|
||||
def test_pending_message_queue_supports_both_modes() -> None:
|
||||
queue = PendingMessageQueue(QueueMode.one_at_a_time)
|
||||
queue.put(HumanMessage(content="一"))
|
||||
queue.put(HumanMessage(content="二"))
|
||||
|
||||
assert [message.content for message in queue.drain()] == ["一"]
|
||||
assert [message.content for message in queue.drain()] == ["二"]
|
||||
|
||||
queue.mode = QueueMode.all
|
||||
queue.put(HumanMessage(content="三"))
|
||||
queue.put(HumanMessage(content="四"))
|
||||
assert [message.content for message in queue.drain()] == ["三", "四"]
|
||||
|
||||
|
||||
async def test_extension_registry_orders_context_and_tool_hooks() -> None:
|
||||
registry = ExtensionRegistry()
|
||||
run = AgentRunContext(run_id="run_1")
|
||||
|
||||
async def first(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
return AgentContext(
|
||||
messages=(*context.messages, HumanMessage(content="first")),
|
||||
tools=context.tools,
|
||||
model_hint=context.model_hint,
|
||||
)
|
||||
|
||||
async def second(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
return AgentContext(
|
||||
messages=(*context.messages, HumanMessage(content="second")),
|
||||
tools=context.tools,
|
||||
model_hint=context.model_hint,
|
||||
)
|
||||
|
||||
async def rewrite(
|
||||
_call: ToolCallContext, _run: AgentRunContext
|
||||
) -> BeforeToolCallDecision:
|
||||
return BeforeToolCallDecision(input={"value": 2})
|
||||
|
||||
async def terminate(
|
||||
_call: ToolCallContext,
|
||||
result: ToolResult,
|
||||
_run: AgentRunContext,
|
||||
) -> AfterToolCallDecision:
|
||||
return AfterToolCallDecision(result=result, terminate=True)
|
||||
|
||||
registry.add_context_transformer(first)
|
||||
registry.add_context_transformer(second)
|
||||
registry.add_before_tool_hook(rewrite)
|
||||
registry.add_after_tool_hook(terminate)
|
||||
|
||||
transformed = await registry.transform_context(AgentContext((), ()), run)
|
||||
assert [message.content for message in transformed.messages] == ["first", "second"]
|
||||
|
||||
call = ToolCallContext("call_1", "tool", {"value": 1}, ())
|
||||
before = await registry.before_tool_call(call, run)
|
||||
assert before.input == {"value": 2}
|
||||
|
||||
result = ToolResult(True, 0, "ok", "", 1, False, False, "trace_1")
|
||||
after = await registry.after_tool_call(call, result, run)
|
||||
assert after.result is result
|
||||
assert after.terminate is True
|
||||
|
||||
|
||||
async def test_extension_registry_extend_preserves_registration_order() -> None:
|
||||
registry = ExtensionRegistry()
|
||||
skill_registry = ExtensionRegistry()
|
||||
|
||||
async def base(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
return AgentContext((*context.messages, HumanMessage(content="base")), context.tools)
|
||||
|
||||
async def skill(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||||
return AgentContext((*context.messages, HumanMessage(content="skill")), context.tools)
|
||||
|
||||
registry.add_context_transformer(base)
|
||||
skill_registry.add_context_transformer(skill)
|
||||
registry.extend(skill_registry)
|
||||
|
||||
transformed = await registry.transform_context(AgentContext((), ()), AgentRunContext())
|
||||
|
||||
assert [message.content for message in transformed.messages] == ["base", "skill"]
|
||||
|
||||
|
||||
async def test_stateful_agent_consumes_follow_up_and_publishes_events() -> None:
|
||||
calls: list[list[str]] = []
|
||||
events: list[str] = []
|
||||
|
||||
async def runner(*, messages: list[Any], **_kwargs: object) -> AgentTurnResult:
|
||||
calls.append([message.content for message in messages])
|
||||
messages.append(HumanMessage(content=f"assistant-{len(calls)}"))
|
||||
return AgentTurnResult(
|
||||
content=f"done-{len(calls)}",
|
||||
reasoning="",
|
||||
iterations=1,
|
||||
tool_history=[{"name": f"tool-{len(calls)}"}],
|
||||
usage=TokenUsage(total_tokens=10),
|
||||
)
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
|
||||
async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
events.append(event.type)
|
||||
|
||||
extensions.add_event_listener(listen)
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=MagicMock(), max_iterations=2),
|
||||
extensions=extensions,
|
||||
runner=runner,
|
||||
)
|
||||
|
||||
async def enqueue_follow_up(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
if event.type == "agent.started":
|
||||
assert agent.follow_up("再做一步") is True
|
||||
|
||||
extensions.add_event_listener(enqueue_follow_up)
|
||||
result = await agent.prompt("开始")
|
||||
|
||||
assert result.content == "done-2"
|
||||
assert calls == [["开始"], ["开始", "assistant-1", "再做一步"]]
|
||||
assert result.iterations == 2
|
||||
assert result.usage.total_tokens == 20
|
||||
assert [item["name"] for item in result.tool_history] == ["tool-1", "tool-2"]
|
||||
assert events == ["agent.started", "agent.ended", "agent.settled"]
|
||||
assert agent.state.status == "completed"
|
||||
assert agent.state.streaming_message == "done-2"
|
||||
|
||||
|
||||
async def test_stateful_agent_does_not_continue_after_runner_truncates() -> None:
|
||||
calls = 0
|
||||
|
||||
async def runner(*, messages: list[Any], **_kwargs: object) -> AgentTurnResult:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
messages.append(HumanMessage(content="assistant"))
|
||||
return AgentTurnResult(content="partial", reasoning="", iterations=1, truncated=True)
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=None, max_iterations=3),
|
||||
extensions=extensions,
|
||||
runner=runner,
|
||||
)
|
||||
|
||||
async def enqueue(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
if event.type == "agent.started":
|
||||
assert agent.follow_up("稍后处理") is True
|
||||
|
||||
extensions.add_event_listener(enqueue)
|
||||
result = await agent.prompt("开始")
|
||||
|
||||
assert calls == 1
|
||||
assert result.truncated is True
|
||||
assert [message.content for message in agent.follow_ups.drain()] == ["稍后处理"]
|
||||
|
||||
|
||||
async def test_follow_up_is_preserved_when_iteration_budget_is_exhausted() -> None:
|
||||
async def runner(*, messages: list[Any], **_kwargs: object) -> AgentTurnResult:
|
||||
messages.append(HumanMessage(content="assistant"))
|
||||
return AgentTurnResult(content="done", reasoning="", iterations=1)
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=None, max_iterations=1),
|
||||
extensions=extensions,
|
||||
runner=runner,
|
||||
)
|
||||
|
||||
async def enqueue(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
if event.type == "agent.started":
|
||||
assert agent.follow_up("稍后处理") is True
|
||||
|
||||
extensions.add_event_listener(enqueue)
|
||||
result = await agent.prompt("开始")
|
||||
|
||||
assert result.truncated is True
|
||||
assert bool(agent.follow_ups) is True
|
||||
assert [message.content for message in agent.follow_ups.drain()] == ["稍后处理"]
|
||||
|
||||
|
||||
async def test_started_listener_failure_restores_idle_state() -> None:
|
||||
extensions = ExtensionRegistry()
|
||||
|
||||
async def fail(_event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
raise RuntimeError("event store unavailable")
|
||||
|
||||
extensions.add_event_listener(fail)
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=MagicMock()),
|
||||
extensions=extensions,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="event store unavailable"):
|
||||
await agent.prompt("开始")
|
||||
await agent.wait_for_idle()
|
||||
assert agent.state.status == "failed"
|
||||
|
||||
|
||||
async def test_stateful_agent_abort_cancels_active_run() -> None:
|
||||
started = asyncio.Event()
|
||||
events: list[str] = []
|
||||
|
||||
async def runner(**_kwargs: object) -> AgentTurnResult:
|
||||
started.set()
|
||||
await asyncio.Event().wait()
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
|
||||
async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
events.append(event.type)
|
||||
|
||||
extensions.add_event_listener(listen)
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=MagicMock()),
|
||||
extensions=extensions,
|
||||
runner=runner,
|
||||
)
|
||||
task = asyncio.create_task(agent.prompt("开始"))
|
||||
await started.wait()
|
||||
|
||||
assert await agent.abort() is True
|
||||
assert agent.run_context.cancellation.cancelled is True
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
await agent.wait_for_idle()
|
||||
assert agent.state.status == "cancelled"
|
||||
assert events == ["agent.started", "agent.cancelled", "agent.settled"]
|
||||
|
||||
|
||||
async def test_agent_replaces_cancelled_token_on_next_run() -> None:
|
||||
async def runner(**_kwargs: object) -> AgentTurnResult:
|
||||
return AgentTurnResult(content="done", reasoning="", iterations=1)
|
||||
|
||||
agent = Agent(AgentConfig(tools=(), tool_runtime=None), runner=runner)
|
||||
agent.run_context.cancellation.cancel()
|
||||
|
||||
await agent.prompt("重新开始")
|
||||
|
||||
assert agent.run_context.cancellation.cancelled is False
|
||||
|
||||
|
||||
async def test_extension_registry_should_stop_short_circuits_on_first_policy() -> None:
|
||||
registry = ExtensionRegistry()
|
||||
run = AgentRunContext(run_id="run_1")
|
||||
turn = TurnSnapshot(
|
||||
iteration=1,
|
||||
content="done",
|
||||
reasoning="",
|
||||
tool_calls=(),
|
||||
tool_results=(),
|
||||
)
|
||||
calls: list[str] = []
|
||||
|
||||
async def keep_going(snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
|
||||
calls.append("keep_going")
|
||||
assert snapshot is turn
|
||||
return False
|
||||
|
||||
async def stop_now(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
|
||||
calls.append("stop_now")
|
||||
return True
|
||||
|
||||
async def never_called(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
|
||||
calls.append("never_called")
|
||||
return False
|
||||
|
||||
registry.add_turn_policy(keep_going)
|
||||
registry.add_turn_policy(stop_now)
|
||||
registry.add_turn_policy(never_called)
|
||||
|
||||
assert await registry.should_stop(turn, run) is True
|
||||
assert calls == ["keep_going", "stop_now"]
|
||||
|
||||
|
||||
async def test_extension_registry_should_stop_false_when_all_policies_decline() -> None:
|
||||
registry = ExtensionRegistry()
|
||||
turn = TurnSnapshot(
|
||||
iteration=2,
|
||||
content="",
|
||||
reasoning="",
|
||||
tool_calls=({"name": "tool", "args": "{}"},),
|
||||
tool_results=(),
|
||||
)
|
||||
|
||||
async def decline(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
|
||||
return False
|
||||
|
||||
registry.add_turn_policy(decline)
|
||||
registry.add_turn_policy(decline)
|
||||
|
||||
assert await registry.should_stop(turn, AgentRunContext()) is False
|
||||
|
||||
|
||||
async def test_extension_registry_should_stop_respects_cancellation() -> None:
|
||||
registry = ExtensionRegistry()
|
||||
run = AgentRunContext()
|
||||
run.cancellation.cancel()
|
||||
turn = TurnSnapshot(
|
||||
iteration=1,
|
||||
content="",
|
||||
reasoning="",
|
||||
tool_calls=(),
|
||||
tool_results=(),
|
||||
)
|
||||
|
||||
async def policy(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
|
||||
return False
|
||||
|
||||
# raise_if_cancelled 在 policy 遍历循环内，故至少注册一个 policy 才会触发；
|
||||
# 这与 transform_context / before_tool_call / after_tool_call 一致：空注册时
|
||||
# 方法无 await 点，不额外检查取消信号。
|
||||
registry.add_turn_policy(policy)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await registry.should_stop(turn, run)
|
||||
|
||||
|
||||
async def test_on_tool_event_publishes_tool_lifecycle_to_extensions() -> None:
|
||||
"""_on_tool_event 既转发 SSE，又把同一事件发布进扩展事件流。"""
|
||||
events: list[tuple[str, dict[str, Any]]] = []
|
||||
sse_events: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||||
sse_events.append((event_type, payload))
|
||||
|
||||
async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
events.append((event.type, dict(event.payload)))
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
extensions.add_event_listener(listen)
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=None, on_tool_event=on_tool_event),
|
||||
extensions=extensions,
|
||||
)
|
||||
|
||||
await agent._on_tool_event(
|
||||
"tool.started", {"tool_call_id": "c1", "tool_name": "search"}
|
||||
)
|
||||
await agent._on_tool_event(
|
||||
"tool.completed", {"tool_call_id": "c1", "status": "success"}
|
||||
)
|
||||
await agent._on_tool_event(
|
||||
"tool.completed", {"tool_call_id": "c2", "status": "denied"}
|
||||
)
|
||||
|
||||
assert [item[0] for item in sse_events] == [
|
||||
"tool.started",
|
||||
"tool.completed",
|
||||
"tool.completed",
|
||||
]
|
||||
assert [item[0] for item in events] == [
|
||||
"tool.started",
|
||||
"tool.completed",
|
||||
"tool.completed",
|
||||
]
|
||||
assert events[1][1]["status"] == "success"
|
||||
# pending 状态维护未被桥接改动破坏
|
||||
assert "c1" not in agent.state.pending_tool_calls
|
||||
|
||||
|
||||
async def test_on_tool_event_swallows_listener_failure() -> None:
|
||||
"""扩展 listener 抛异常不应阻断 SSE 转发或工具事件主流程。"""
|
||||
sse: list[str] = []
|
||||
|
||||
async def on_tool_event(event_type: str, _payload: dict[str, Any]) -> None:
|
||||
sse.append(event_type)
|
||||
|
||||
async def boom(_event: AgentEvent, _run: AgentRunContext) -> None:
|
||||
raise RuntimeError("listener boom")
|
||||
|
||||
extensions = ExtensionRegistry()
|
||||
extensions.add_event_listener(boom)
|
||||
agent = Agent(
|
||||
AgentConfig(tools=(), tool_runtime=None, on_tool_event=on_tool_event),
|
||||
extensions=extensions,
|
||||
)
|
||||
|
||||
# 不抛：异常被吞，SSE 仍照常转发
|
||||
await agent._on_tool_event("tool.started", {"tool_call_id": "c1"})
|
||||
assert sse == ["tool.started"]
|
||||
Reference in new issue
Block a user