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,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