feat(core/agent): 实现 AgentLoop 推理循环引擎
保真 TRM4 算法 #11: json_repair 兜底、submit_answer 终止、 pluggy hook 生命周期、无效工具不计步。 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,347 @@
|
|||||||
|
"""Agent Loop 引擎 — Thinking+JSON 推理循环,pluggy 驱动 hook。
|
||||||
|
|
||||||
|
算法保真 #11: 完整保留 TRM4 core/loop.py 逻辑:
|
||||||
|
- json_repair 兜底解析
|
||||||
|
- submit_answer 终止
|
||||||
|
- 无效工具(ValueError)不计步
|
||||||
|
- pluggy hook 生命周期(before_step / after_tool / after_step / on_finish)
|
||||||
|
|
||||||
|
TRM4 → TRM5 有意变更(非简化):
|
||||||
|
- 同步 → 全异步(async/await)
|
||||||
|
- client: Any → llm: LLMProvider(Protocol 类型化)
|
||||||
|
- tool_fn: Callable → ToolDispatcher.dispatch()(Protocol + context)
|
||||||
|
- Step 新增 call_id(从 LLMResponse.call_id 透传)
|
||||||
|
- thinking 从 getattr(msg, "reasoning_content") → response.thinking(adapters 已统一剥离)
|
||||||
|
- token 用量从 response.usage.prompt_tokens → response.prompt_tokens(LLMResponse 扁平化)
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
import pluggy
|
||||||
|
from json_repair import repair_json
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
from core.agent.protocols import AgentLoopSpec, ToolDispatcher
|
||||||
|
from core.agent.types import LoopResult, Step
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from core.protocols import LLMProvider
|
||||||
|
from core.types import LLMResponse
|
||||||
|
|
||||||
|
|
||||||
|
async def _call_hook(hook: Any, **kwargs: Any) -> list[Any]:
|
||||||
|
"""调用 pluggy hook 并 await 异步返回值。
|
||||||
|
|
||||||
|
pluggy 本身是同步调度,但 hookimpl 可以是 async def,
|
||||||
|
此时 hook() 返回 coroutine 列表,需要逐个 await。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
hook: pluggy hook caller(如 pm.hook.before_step)。
|
||||||
|
**kwargs: 传递给 hook 的关键字参数。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
已 resolve 的返回值列表。
|
||||||
|
"""
|
||||||
|
results = hook(**kwargs)
|
||||||
|
if results is not None:
|
||||||
|
resolved = []
|
||||||
|
for r in results:
|
||||||
|
if hasattr(r, "__await__"):
|
||||||
|
resolved.append(await r)
|
||||||
|
else:
|
||||||
|
resolved.append(r)
|
||||||
|
return resolved
|
||||||
|
return []
|
||||||
|
|
||||||
|
|
||||||
|
class AgentLoop:
|
||||||
|
"""Thinking+JSON 推理循环引擎。
|
||||||
|
|
||||||
|
类比 nn.Module: 接收 prompt + 工具调度器,返回 LoopResult。
|
||||||
|
不感知视频树、QA、数据库等领域概念。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
llm: LLMProvider 实例(Protocol 类型化,提供 async chat 方法)。
|
||||||
|
max_steps: 最大有效步数(每次成功工具调用计一步)。
|
||||||
|
max_retries: JSON 解析连续失败的最大容忍次数。
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
llm: LLMProvider,
|
||||||
|
max_steps: int = 15,
|
||||||
|
max_retries: int = 3,
|
||||||
|
) -> None:
|
||||||
|
self._llm = llm
|
||||||
|
self._max_steps = max_steps
|
||||||
|
self._max_retries = max_retries
|
||||||
|
|
||||||
|
async def run(
|
||||||
|
self,
|
||||||
|
system_prompt: str,
|
||||||
|
user_prompt: str,
|
||||||
|
tool_dispatcher: ToolDispatcher,
|
||||||
|
plugins: list[object] | None = None,
|
||||||
|
*,
|
||||||
|
session_id: str | None = None,
|
||||||
|
) -> LoopResult:
|
||||||
|
"""执行 Thinking+JSON 推理循环。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
system_prompt: 系统提示词。
|
||||||
|
user_prompt: 用户提示词。
|
||||||
|
tool_dispatcher: 工具调度器,ToolDispatcher Protocol 实例。
|
||||||
|
plugins: pluggy 插件列表。
|
||||||
|
session_id: 会话 ID,透传给 LLMProvider。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
LoopResult 实例,包含推理步骤、token 用量、终止原因。
|
||||||
|
"""
|
||||||
|
pm = self._create_plugin_manager(plugins)
|
||||||
|
messages: list[dict[str, Any]] = [
|
||||||
|
{"role": "system", "content": system_prompt},
|
||||||
|
{"role": "user", "content": user_prompt},
|
||||||
|
]
|
||||||
|
steps: list[Step] = []
|
||||||
|
token_usage: dict[str, int] = {"prompt_tokens": 0, "completion_tokens": 0}
|
||||||
|
step_count = 0
|
||||||
|
retry_count = 0
|
||||||
|
iteration = 0
|
||||||
|
|
||||||
|
while step_count < self._max_steps:
|
||||||
|
await _call_hook(pm.hook.before_step, iteration=iteration, messages=messages)
|
||||||
|
|
||||||
|
# Phase 1: LLM 调用
|
||||||
|
try:
|
||||||
|
response = await self._call_llm(messages, token_usage, session_id=session_id)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("LLM API 调用失败: {}", e)
|
||||||
|
result = LoopResult(
|
||||||
|
steps=steps,
|
||||||
|
steps_used=step_count,
|
||||||
|
token_usage=token_usage,
|
||||||
|
stop_reason="error",
|
||||||
|
)
|
||||||
|
await _call_hook(pm.hook.on_finish, result=result)
|
||||||
|
return result
|
||||||
|
|
||||||
|
# Phase 2: 解析响应
|
||||||
|
parsed = self._parse_response(response)
|
||||||
|
if parsed is None:
|
||||||
|
retry_count += 1
|
||||||
|
logger.warning("响应解析失败 (retry {}/{})", retry_count, self._max_retries)
|
||||||
|
messages.append({"role": "assistant", "content": response.content})
|
||||||
|
messages.append(
|
||||||
|
{
|
||||||
|
"role": "user",
|
||||||
|
"content": (
|
||||||
|
"你的输出不是合法 JSON。请严格输出 JSON 格式:"
|
||||||
|
'{"reflect": {...}, "plan": {...}, '
|
||||||
|
'"action": {"tool": "...", "args": {...}}}'
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
if retry_count >= self._max_retries:
|
||||||
|
result = LoopResult(
|
||||||
|
steps=steps,
|
||||||
|
steps_used=step_count,
|
||||||
|
token_usage=token_usage,
|
||||||
|
stop_reason="parse_error",
|
||||||
|
)
|
||||||
|
await _call_hook(pm.hook.after_step, iteration=iteration, messages=messages)
|
||||||
|
await _call_hook(pm.hook.on_finish, result=result)
|
||||||
|
return result
|
||||||
|
await _call_hook(pm.hook.after_step, iteration=iteration, messages=messages)
|
||||||
|
iteration += 1
|
||||||
|
continue
|
||||||
|
|
||||||
|
thought, reflect, plan, raw_content, action, call_id = parsed
|
||||||
|
retry_count = 0
|
||||||
|
messages.append({"role": "assistant", "content": raw_content})
|
||||||
|
|
||||||
|
# Phase 3: 执行工具
|
||||||
|
tool_name: str = action["tool"]
|
||||||
|
tool_args: dict[str, Any] = action["args"]
|
||||||
|
context: dict[str, Any] = {
|
||||||
|
"session_id": session_id,
|
||||||
|
"iteration": iteration,
|
||||||
|
}
|
||||||
|
output, is_valid = await self._execute_tool(
|
||||||
|
tool_dispatcher, tool_name, tool_args, context=context
|
||||||
|
)
|
||||||
|
if not is_valid:
|
||||||
|
messages.append(
|
||||||
|
{
|
||||||
|
"role": "user",
|
||||||
|
"content": f"[工具调用无效: {tool_name}] {output}",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
await _call_hook(pm.hook.after_step, iteration=iteration, messages=messages)
|
||||||
|
iteration += 1
|
||||||
|
continue
|
||||||
|
|
||||||
|
step_count += 1
|
||||||
|
step = Step(
|
||||||
|
thought=thought,
|
||||||
|
reflect=reflect,
|
||||||
|
plan=plan,
|
||||||
|
tool_call={"tool": tool_name, "args": tool_args},
|
||||||
|
tool_output=output,
|
||||||
|
raw_content=raw_content,
|
||||||
|
call_id=call_id,
|
||||||
|
)
|
||||||
|
steps.append(step)
|
||||||
|
|
||||||
|
# Phase 4: Hook + 反馈
|
||||||
|
hints = await _call_hook(pm.hook.after_tool, iteration=iteration, step=step)
|
||||||
|
feedback = self._build_feedback(tool_name, output, hints)
|
||||||
|
messages.append(feedback)
|
||||||
|
await _call_hook(pm.hook.after_step, iteration=iteration, messages=messages)
|
||||||
|
|
||||||
|
# Phase 5: 终止检查
|
||||||
|
if tool_name == "submit_answer":
|
||||||
|
result = LoopResult(
|
||||||
|
result=tool_args,
|
||||||
|
steps=steps,
|
||||||
|
steps_used=step_count,
|
||||||
|
token_usage=token_usage,
|
||||||
|
stop_reason="finished",
|
||||||
|
)
|
||||||
|
await _call_hook(pm.hook.on_finish, result=result)
|
||||||
|
return result
|
||||||
|
|
||||||
|
iteration += 1
|
||||||
|
|
||||||
|
# 预算耗尽
|
||||||
|
result = LoopResult(
|
||||||
|
steps=steps,
|
||||||
|
steps_used=step_count,
|
||||||
|
token_usage=token_usage,
|
||||||
|
stop_reason="budget_exceeded",
|
||||||
|
)
|
||||||
|
await _call_hook(pm.hook.on_finish, result=result)
|
||||||
|
return result
|
||||||
|
|
||||||
|
def _create_plugin_manager(self, plugins: list[object] | None) -> pluggy.PluginManager:
|
||||||
|
"""创建并注册 plugins 的 PluginManager。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
plugins: pluggy 插件列表,可为 None。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
配置好的 PluginManager 实例。
|
||||||
|
"""
|
||||||
|
pm = pluggy.PluginManager("agent_loop")
|
||||||
|
pm.add_hookspecs(AgentLoopSpec)
|
||||||
|
for plugin in plugins or []:
|
||||||
|
pm.register(plugin)
|
||||||
|
return pm
|
||||||
|
|
||||||
|
async def _call_llm(
|
||||||
|
self,
|
||||||
|
messages: list[dict[str, Any]],
|
||||||
|
token_usage: dict[str, int],
|
||||||
|
*,
|
||||||
|
session_id: str | None = None,
|
||||||
|
) -> LLMResponse:
|
||||||
|
"""调用 LLMProvider 并累加 token 使用量。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
messages: 消息历史。
|
||||||
|
token_usage: 可变字典,就地累加。
|
||||||
|
session_id: 会话 ID,透传给 LLMProvider。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
LLMResponse 实例。
|
||||||
|
"""
|
||||||
|
response = await self._llm.chat(messages, session_id=session_id)
|
||||||
|
token_usage["prompt_tokens"] += response.prompt_tokens
|
||||||
|
token_usage["completion_tokens"] += response.completion_tokens
|
||||||
|
return response
|
||||||
|
|
||||||
|
def _parse_response(
|
||||||
|
self, response: LLMResponse
|
||||||
|
) -> tuple[str, dict, dict, str, dict, str] | None:
|
||||||
|
"""从 LLMResponse 中提取结构化决策数据。
|
||||||
|
|
||||||
|
解析流程: content → repair_json → json.loads → 校验 action/tool/args。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
response: LLMResponse 实例。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
解析成功返回 (thought, reflect, plan, raw_content, action, call_id);
|
||||||
|
解析失败返回 None。
|
||||||
|
"""
|
||||||
|
content = response.content
|
||||||
|
thought = response.thinking
|
||||||
|
|
||||||
|
if not content.strip():
|
||||||
|
return None
|
||||||
|
|
||||||
|
repaired = repair_json(content)
|
||||||
|
try:
|
||||||
|
data = json.loads(repaired)
|
||||||
|
except (json.JSONDecodeError, ValueError):
|
||||||
|
return None
|
||||||
|
|
||||||
|
if not isinstance(data, dict) or "action" not in data:
|
||||||
|
return None
|
||||||
|
|
||||||
|
action = data["action"]
|
||||||
|
if not isinstance(action, dict) or "tool" not in action or "args" not in action:
|
||||||
|
return None
|
||||||
|
|
||||||
|
reflect = data.get("reflect", {})
|
||||||
|
plan = data.get("plan", {})
|
||||||
|
return thought, reflect, plan, content, action, response.call_id
|
||||||
|
|
||||||
|
async def _execute_tool(
|
||||||
|
self,
|
||||||
|
dispatcher: ToolDispatcher,
|
||||||
|
name: str,
|
||||||
|
args: dict[str, Any],
|
||||||
|
*,
|
||||||
|
context: dict[str, Any],
|
||||||
|
) -> tuple[str, bool]:
|
||||||
|
"""执行工具调用。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
dispatcher: 工具调度器。
|
||||||
|
name: 工具名称。
|
||||||
|
args: 工具参数。
|
||||||
|
context: 调用上下文(session_id、iteration 等)。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
(output, is_valid) — ValueError 时 is_valid=False 且不计步数。
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
output = await dispatcher.dispatch(name, args, context=context)
|
||||||
|
return output, True
|
||||||
|
except ValueError as e:
|
||||||
|
return f"工具调用失败: {e}", False
|
||||||
|
|
||||||
|
def _build_feedback(
|
||||||
|
self,
|
||||||
|
tool_name: str,
|
||||||
|
tool_output: str,
|
||||||
|
hints: list[str | None],
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""组装工具结果反馈消息。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
tool_name: 工具名称。
|
||||||
|
tool_output: 工具原始输出。
|
||||||
|
hints: hook 返回的 hint 列表(含 None)。
|
||||||
|
|
||||||
|
返回:
|
||||||
|
user role 消息字典。
|
||||||
|
"""
|
||||||
|
parts = [f"[工具执行结果: {tool_name}]", tool_output]
|
||||||
|
for hint in hints:
|
||||||
|
if hint is not None:
|
||||||
|
parts.append(hint)
|
||||||
|
return {"role": "user", "content": "\n".join(parts)}
|
||||||
@@ -0,0 +1,241 @@
|
|||||||
|
"""core/agent/loop.py 单元测试。
|
||||||
|
|
||||||
|
算法保真 #11 — AgentLoop 推理循环引擎。
|
||||||
|
9 个测试覆盖: 终止、预算、无效工具、解析错误、JSON 修复、
|
||||||
|
thinking 捕获、token 累加、call_id 透传、pluggy hook。
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
from unittest.mock import AsyncMock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from core.agent.loop import AgentLoop
|
||||||
|
from core.agent.protocols import hookimpl
|
||||||
|
from core.types import LLMResponse
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from core.agent.types import LoopResult, Step
|
||||||
|
|
||||||
|
|
||||||
|
# ── 测试基础设施 ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class _StubDispatcher:
|
||||||
|
"""测试用工具调度器。"""
|
||||||
|
|
||||||
|
async def dispatch(
|
||||||
|
self, tool_name: str, args: dict[str, Any], *, context: dict[str, Any]
|
||||||
|
) -> str:
|
||||||
|
if tool_name == "submit_answer":
|
||||||
|
return "答案已提交"
|
||||||
|
if tool_name == "search_tree":
|
||||||
|
return "搜索结果: 找到节点 L2-3"
|
||||||
|
raise ValueError(f"未知工具: {tool_name}")
|
||||||
|
|
||||||
|
|
||||||
|
def _make_response(content: str, thinking: str = "") -> LLMResponse:
|
||||||
|
"""构造测试用 LLMResponse。"""
|
||||||
|
return LLMResponse(
|
||||||
|
content=content,
|
||||||
|
thinking=thinking,
|
||||||
|
model="test-model",
|
||||||
|
provider="test",
|
||||||
|
prompt_tokens=10,
|
||||||
|
completion_tokens=10,
|
||||||
|
latency_ms=100,
|
||||||
|
ttft_ms=50.0,
|
||||||
|
max_inter_token_ms=10.0,
|
||||||
|
cache_hit=False,
|
||||||
|
call_id="test-call-id",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _submit_json(answer: str = "42") -> str:
|
||||||
|
"""构造 submit_answer 的 JSON 响应。"""
|
||||||
|
return json.dumps(
|
||||||
|
{
|
||||||
|
"reflect": {"observation": "找到答案"},
|
||||||
|
"plan": {"next_step": "提交"},
|
||||||
|
"action": {"tool": "submit_answer", "args": {"answer": answer}},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _search_json() -> str:
|
||||||
|
"""构造 search_tree 的 JSON 响应。"""
|
||||||
|
return json.dumps(
|
||||||
|
{
|
||||||
|
"reflect": {"observation": "需要搜索"},
|
||||||
|
"plan": {"next_step": "搜索"},
|
||||||
|
"action": {"tool": "search_tree", "args": {"query": "test"}},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _invalid_tool_json() -> str:
|
||||||
|
"""构造无效工具的 JSON 响应。"""
|
||||||
|
return json.dumps(
|
||||||
|
{
|
||||||
|
"reflect": {},
|
||||||
|
"plan": {},
|
||||||
|
"action": {"tool": "unknown_tool", "args": {}},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── 测试用例 ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class TestAgentLoop:
|
||||||
|
"""AgentLoop 推理循环引擎测试。"""
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_submit_answer_terminates_loop(self) -> None:
|
||||||
|
"""submit_answer 终止循环 → finished, result=args, steps_used=1。"""
|
||||||
|
llm = AsyncMock()
|
||||||
|
llm.chat.return_value = _make_response(_submit_json())
|
||||||
|
|
||||||
|
loop = AgentLoop(llm=llm, max_steps=10)
|
||||||
|
result = await loop.run("system", "user", _StubDispatcher())
|
||||||
|
|
||||||
|
assert result.stop_reason == "finished"
|
||||||
|
assert result.result == {"answer": "42"}
|
||||||
|
assert result.steps_used == 1
|
||||||
|
assert len(result.steps) == 1
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_budget_exceeded(self) -> None:
|
||||||
|
"""max_steps=3 用完 → budget_exceeded, steps_used=3。"""
|
||||||
|
llm = AsyncMock()
|
||||||
|
llm.chat.return_value = _make_response(_search_json())
|
||||||
|
|
||||||
|
loop = AgentLoop(llm=llm, max_steps=3)
|
||||||
|
result = await loop.run("system", "user", _StubDispatcher())
|
||||||
|
|
||||||
|
assert result.stop_reason == "budget_exceeded"
|
||||||
|
assert result.steps_used == 3
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_invalid_tool_not_counted_as_step(self) -> None:
|
||||||
|
"""无效工具(ValueError)不计步 → steps_used=1。"""
|
||||||
|
llm = AsyncMock()
|
||||||
|
llm.chat.side_effect = [
|
||||||
|
_make_response(_invalid_tool_json()),
|
||||||
|
_make_response(_submit_json()),
|
||||||
|
]
|
||||||
|
|
||||||
|
loop = AgentLoop(llm=llm, max_steps=10)
|
||||||
|
result = await loop.run("system", "user", _StubDispatcher())
|
||||||
|
|
||||||
|
assert result.stop_reason == "finished"
|
||||||
|
assert result.steps_used == 1
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_parse_error_after_max_retries(self) -> None:
|
||||||
|
"""非 JSON 内容连续失败 → parse_error, steps_used=0。"""
|
||||||
|
llm = AsyncMock()
|
||||||
|
llm.chat.return_value = _make_response("这不是JSON内容")
|
||||||
|
|
||||||
|
loop = AgentLoop(llm=llm, max_steps=10, max_retries=3)
|
||||||
|
result = await loop.run("system", "user", _StubDispatcher())
|
||||||
|
|
||||||
|
assert result.stop_reason == "parse_error"
|
||||||
|
assert result.steps_used == 0
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_json_repair_handles_malformed(self) -> None:
|
||||||
|
"""轻微 JSON 缺陷(缺少闭合花括号)被 json_repair 修复。"""
|
||||||
|
malformed = (
|
||||||
|
'{"reflect": {}, "plan": {}, '
|
||||||
|
'"action": {"tool": "submit_answer", "args": {"answer": "42"}}'
|
||||||
|
)
|
||||||
|
llm = AsyncMock()
|
||||||
|
llm.chat.return_value = _make_response(malformed)
|
||||||
|
|
||||||
|
loop = AgentLoop(llm=llm, max_steps=10)
|
||||||
|
result = await loop.run("system", "user", _StubDispatcher())
|
||||||
|
|
||||||
|
assert result.stop_reason == "finished"
|
||||||
|
assert result.result == {"answer": "42"}
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_thinking_content_captured_in_step(self) -> None:
|
||||||
|
"""LLMResponse.thinking → Step.thought。"""
|
||||||
|
llm = AsyncMock()
|
||||||
|
llm.chat.return_value = _make_response(_submit_json(), thinking="深度思考过程")
|
||||||
|
|
||||||
|
loop = AgentLoop(llm=llm, max_steps=10)
|
||||||
|
result = await loop.run("system", "user", _StubDispatcher())
|
||||||
|
|
||||||
|
assert result.steps[0].thought == "深度思考过程"
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_token_usage_accumulated(self) -> None:
|
||||||
|
"""多步 token 累加: 3 次调用 × 10 tokens = 30。"""
|
||||||
|
llm = AsyncMock()
|
||||||
|
llm.chat.side_effect = [
|
||||||
|
_make_response(_search_json()),
|
||||||
|
_make_response(_search_json()),
|
||||||
|
_make_response(_submit_json()),
|
||||||
|
]
|
||||||
|
|
||||||
|
loop = AgentLoop(llm=llm, max_steps=10)
|
||||||
|
result = await loop.run("system", "user", _StubDispatcher())
|
||||||
|
|
||||||
|
assert result.token_usage["prompt_tokens"] == 30
|
||||||
|
assert result.token_usage["completion_tokens"] == 30
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_call_id_propagated_to_step(self) -> None:
|
||||||
|
"""LLMResponse.call_id → Step.call_id。"""
|
||||||
|
llm = AsyncMock()
|
||||||
|
llm.chat.return_value = _make_response(_submit_json())
|
||||||
|
|
||||||
|
loop = AgentLoop(llm=llm, max_steps=10)
|
||||||
|
result = await loop.run("system", "user", _StubDispatcher())
|
||||||
|
|
||||||
|
assert result.steps[0].call_id == "test-call-id"
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_pluggy_hooks_called(self) -> None:
|
||||||
|
"""TrackingPlugin 验证 before_step/after_tool/after_step/on_finish 全部触发。"""
|
||||||
|
|
||||||
|
class TrackingPlugin:
|
||||||
|
"""记录 hook 调用事件的测试插件。"""
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.events: list[str] = []
|
||||||
|
|
||||||
|
@hookimpl
|
||||||
|
async def before_step(self, iteration: int, messages: list[dict[str, Any]]) -> None:
|
||||||
|
self.events.append(f"before_step:{iteration}")
|
||||||
|
|
||||||
|
@hookimpl
|
||||||
|
async def after_tool(self, iteration: int, step: Step) -> str | None:
|
||||||
|
self.events.append(f"after_tool:{iteration}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
@hookimpl
|
||||||
|
async def after_step(self, iteration: int, messages: list[dict[str, Any]]) -> None:
|
||||||
|
self.events.append(f"after_step:{iteration}")
|
||||||
|
|
||||||
|
@hookimpl
|
||||||
|
async def on_finish(self, result: LoopResult) -> None:
|
||||||
|
self.events.append(f"on_finish:{result.stop_reason}")
|
||||||
|
|
||||||
|
tracker = TrackingPlugin()
|
||||||
|
llm = AsyncMock()
|
||||||
|
llm.chat.return_value = _make_response(_submit_json())
|
||||||
|
|
||||||
|
loop = AgentLoop(llm=llm, max_steps=10)
|
||||||
|
result = await loop.run("system", "user", _StubDispatcher(), plugins=[tracker])
|
||||||
|
|
||||||
|
assert result.stop_reason == "finished"
|
||||||
|
assert "before_step:0" in tracker.events
|
||||||
|
assert "after_tool:0" in tracker.events
|
||||||
|
assert "after_step:0" in tracker.events
|
||||||
|
assert "on_finish:finished" in tracker.events
|
||||||
Reference in New Issue
Block a user