feat(core/agent): 实现 AgentLoop 推理循环引擎

保真 TRM4 算法 #11: json_repair 兜底、submit_answer 终止、 pluggy hook 生命周期、无效工具不计步。 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-07-06 22:43:22 -04:00
parent c2ff2855b7
commit 9ca9035190
2 changed files with 588 additions and 0 deletions
@@ -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