From 9ca9035190f8166dc7020f0435bf2ac6aaf419d7 Mon Sep 17 00:00:00 2001 From: iomgaa Date: Mon, 6 Jul 2026 22:43:22 -0400 Subject: [PATCH] =?UTF-8?q?feat(core/agent):=20=E5=AE=9E=E7=8E=B0=20AgentL?= =?UTF-8?q?oop=20=E6=8E=A8=E7=90=86=E5=BE=AA=E7=8E=AF=E5=BC=95=E6=93=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 保真 TRM4 算法 #11: json_repair 兜底、submit_answer 终止、 pluggy hook 生命周期、无效工具不计步。 Co-Authored-By: Claude Opus 4.6 (1M context) --- core/agent/loop.py | 347 ++++++++++++++++++++++++++++++++++ tests/unit/test_agent_loop.py | 241 +++++++++++++++++++++++ 2 files changed, 588 insertions(+) create mode 100644 core/agent/loop.py create mode 100644 tests/unit/test_agent_loop.py diff --git a/core/agent/loop.py b/core/agent/loop.py new file mode 100644 index 0000000..24d646f --- /dev/null +++ b/core/agent/loop.py @@ -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)} diff --git a/tests/unit/test_agent_loop.py b/tests/unit/test_agent_loop.py new file mode 100644 index 0000000..85cdcfa --- /dev/null +++ b/tests/unit/test_agent_loop.py @@ -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