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:
2026-07-06 22:43:22 -04:00
parent c2ff2855b7
commit 9ca9035190
2 changed files with 588 additions and 0 deletions
+347
View File
@@ -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: LLMProviderProtocol 类型化)
- tool_fn: Callable → ToolDispatcher.dispatch()Protocol + context
- Step 新增 call_id(从 LLMResponse.call_id 透传)
- thinking 从 getattr(msg, "reasoning_content") → response.thinkingadapters 已统一剥离)
- token 用量从 response.usage.prompt_tokens → response.prompt_tokensLLMResponse 扁平化)
"""
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)}
+241
View File
@@ -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