294 lines
11 KiB
Python
294 lines
11 KiB
Python
"""集成测试 - ReWOO 渐进式回退链
|
||
|
||
测试 ReWOOEngine 的 planning → simplified_rewoo → react → direct 回退策略。
|
||
仅 mock LLMGateway（外部 API），使用真实 ReWOOEngine 实例。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
import pytest
|
||
|
||
from agentkit.core.rewoo import ReWOOEngine, ReWOOStep
|
||
from agentkit.core.react import ReActResult, ReActStep
|
||
from agentkit.llm.gateway import LLMGateway
|
||
from agentkit.llm.protocol import LLMResponse, TokenUsage, ToolCall
|
||
from agentkit.tools.base import Tool
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class FakeTool(Tool):
|
||
"""用于测试的 Fake Tool"""
|
||
|
||
def __init__(
|
||
self,
|
||
name: str = "fake_tool",
|
||
description: str = "A fake tool for testing",
|
||
result: dict | None = None,
|
||
should_fail: bool = False,
|
||
):
|
||
super().__init__(
|
||
name=name,
|
||
description=description,
|
||
)
|
||
self._result = result or {"status": "ok"}
|
||
self._should_fail = should_fail
|
||
self.call_count = 0
|
||
|
||
async def execute(self, **kwargs) -> dict:
|
||
self.call_count += 1
|
||
if self._should_fail:
|
||
raise RuntimeError(f"Tool '{self.name}' execution failed")
|
||
return self._result
|
||
|
||
|
||
def make_response(
|
||
content: str = "",
|
||
tool_calls: list[ToolCall] | None = None,
|
||
prompt_tokens: int = 10,
|
||
completion_tokens: int = 20,
|
||
) -> LLMResponse:
|
||
return LLMResponse(
|
||
content=content,
|
||
model="test-model",
|
||
usage=TokenUsage(
|
||
prompt_tokens=prompt_tokens,
|
||
completion_tokens=completion_tokens,
|
||
),
|
||
tool_calls=tool_calls or [],
|
||
)
|
||
|
||
|
||
def make_plan_response(
|
||
steps: list[dict],
|
||
reasoning: str = "Plan reasoning",
|
||
) -> LLMResponse:
|
||
plan_json = json.dumps({
|
||
"reasoning": reasoning,
|
||
"steps": steps,
|
||
})
|
||
return make_response(content=plan_json)
|
||
|
||
|
||
def make_mock_gateway(responses: list[LLMResponse]) -> MagicMock:
|
||
gateway = MagicMock(spec=LLMGateway)
|
||
gateway.chat = AsyncMock(side_effect=responses)
|
||
return gateway
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 1: Planning succeeds → no fallback
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestReWOOPlanningSucceeds:
|
||
"""规划成功，不使用回退"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_planning_succeeds_no_fallback(self):
|
||
tool = FakeTool(name="calculator", result={"value": 42})
|
||
|
||
plan_response = make_plan_response([
|
||
{"step_id": 1, "tool_name": "calculator", "arguments": {"expr": "6*7"}, "reasoning": "Calculate"},
|
||
])
|
||
synthesis_response = make_response(content="The result is 42")
|
||
|
||
gateway = make_mock_gateway([plan_response, synthesis_response])
|
||
engine = ReWOOEngine(llm_gateway=gateway)
|
||
|
||
result = await engine.execute(
|
||
messages=[{"role": "user", "content": "Calculate 6*7"}],
|
||
tools=[tool],
|
||
)
|
||
|
||
assert isinstance(result, ReActResult)
|
||
assert result.output == "The result is 42"
|
||
assert result.fallback_strategy is None
|
||
assert result.status == "success"
|
||
# 1 tool_call + 1 final_answer = 2 steps
|
||
assert result.total_steps == 2
|
||
assert tool.call_count == 1
|
||
|
||
# Verify ReWOOStep has plan_step_id
|
||
tool_steps = [s for s in result.trajectory if s.action == "tool_call"]
|
||
assert len(tool_steps) == 1
|
||
assert isinstance(tool_steps[0], ReWOOStep)
|
||
assert tool_steps[0].plan_step_id == 1
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 2: Planning fails → simplified planning succeeds → fallback_strategy="simplified_rewoo"
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestReWOOSimplifiedFallback:
|
||
"""规划失败，简化规划成功"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_planning_fails_simplified_succeeds(self):
|
||
tool = FakeTool(name="search", result={"results": ["found"]})
|
||
|
||
# First plan fails (invalid JSON), second (simplified) succeeds
|
||
invalid_plan_response = make_response(content="I cannot create a plan for this task.")
|
||
simplified_plan_response = make_plan_response([
|
||
{"step_id": 1, "tool_name": "search", "arguments": {"query": "test"}, "reasoning": "Simplified search"},
|
||
])
|
||
synthesis_response = make_response(content="Simplified result")
|
||
|
||
gateway = make_mock_gateway([invalid_plan_response, simplified_plan_response, synthesis_response])
|
||
engine = ReWOOEngine(llm_gateway=gateway)
|
||
|
||
result = await engine.execute(
|
||
messages=[{"role": "user", "content": "Complex task"}],
|
||
tools=[tool],
|
||
)
|
||
|
||
assert result.output == "Simplified result"
|
||
assert result.fallback_strategy == "simplified_rewoo"
|
||
assert tool.call_count == 1
|
||
|
||
# Verify trajectory still has proper structure
|
||
tool_steps = [s for s in result.trajectory if s.action == "tool_call"]
|
||
assert len(tool_steps) == 1
|
||
assert tool_steps[0].tool_name == "search"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 3: All planning fails → ReAct fallback → fallback_strategy="react"
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestReWOOReActFallback:
|
||
"""所有规划失败，回退到 ReAct"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_planning_and_simplified_fail_react_succeeds(self):
|
||
# Both plan attempts fail (invalid JSON), ReAct succeeds
|
||
invalid_plan1 = make_response(content="Not a plan")
|
||
invalid_plan2 = make_response(content="Still not a plan")
|
||
react_response = make_response(content="ReAct fallback answer")
|
||
|
||
gateway = make_mock_gateway([invalid_plan1, invalid_plan2, react_response])
|
||
engine = ReWOOEngine(llm_gateway=gateway)
|
||
|
||
result = await engine.execute(
|
||
messages=[{"role": "user", "content": "Complex task"}],
|
||
)
|
||
|
||
assert result.output == "ReAct fallback answer"
|
||
assert result.fallback_strategy == "react"
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_react_fallback_with_tool_calls(self):
|
||
"""ReAct 回退时带工具调用"""
|
||
tool = FakeTool(name="search", result={"results": ["found"]})
|
||
|
||
invalid_plan1 = make_response(content="Cannot plan")
|
||
invalid_plan2 = make_response(content="Still cannot plan")
|
||
react_tool_response = make_response(
|
||
content="",
|
||
tool_calls=[ToolCall(id="tc_1", name="search", arguments={"query": "test"})],
|
||
)
|
||
react_final_response = make_response(content="ReAct answer with tool")
|
||
|
||
gateway = make_mock_gateway([
|
||
invalid_plan1,
|
||
invalid_plan2,
|
||
react_tool_response,
|
||
react_final_response,
|
||
])
|
||
engine = ReWOOEngine(llm_gateway=gateway)
|
||
|
||
result = await engine.execute(
|
||
messages=[{"role": "user", "content": "Search task"}],
|
||
tools=[tool],
|
||
)
|
||
|
||
assert result.output == "ReAct answer with tool"
|
||
assert result.fallback_strategy == "react"
|
||
assert tool.call_count == 1
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_malformed_json_triggers_react_fallback(self):
|
||
"""格式错误的 JSON 触发 ReAct 回退"""
|
||
malformed_response = make_response(content='{"reasoning": "plan", "steps": [invalid json')
|
||
simplified_fail_response = make_response(content="Also not a plan")
|
||
react_response = make_response(content="ReAct answer")
|
||
|
||
gateway = make_mock_gateway([malformed_response, simplified_fail_response, react_response])
|
||
engine = ReWOOEngine(llm_gateway=gateway)
|
||
|
||
result = await engine.execute(
|
||
messages=[{"role": "user", "content": "Task"}],
|
||
)
|
||
|
||
assert result.output == "ReAct answer"
|
||
assert result.fallback_strategy == "react"
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_missing_steps_key_triggers_react_fallback(self):
|
||
"""缺少 steps 键触发 ReAct 回退"""
|
||
no_steps_response = make_response(content='{"reasoning": "no steps here"}')
|
||
simplified_fail_response = make_response(content="Also no steps")
|
||
react_response = make_response(content="ReAct fallback")
|
||
|
||
gateway = make_mock_gateway([no_steps_response, simplified_fail_response, react_response])
|
||
engine = ReWOOEngine(llm_gateway=gateway)
|
||
|
||
result = await engine.execute(
|
||
messages=[{"role": "user", "content": "Task"}],
|
||
)
|
||
|
||
assert result.output == "ReAct fallback"
|
||
assert result.fallback_strategy == "react"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test: Multi-step plan with fallback chain integration
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestReWOOMultiStepWithFallbackChain:
|
||
"""多步计划与回退链的集成测试"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_three_step_plan_no_fallback(self):
|
||
"""三步计划成功，无回退"""
|
||
search_tool = FakeTool(name="search", result={"results": ["Python is great"]})
|
||
calc_tool = FakeTool(name="calculator", result={"value": 100})
|
||
weather_tool = FakeTool(name="weather", result={"temp": 25, "city": "Shanghai"})
|
||
|
||
plan_response = make_plan_response([
|
||
{"step_id": 1, "tool_name": "search", "arguments": {"query": "Python"}, "reasoning": "Search first"},
|
||
{"step_id": 2, "tool_name": "calculator", "arguments": {"expr": "10*10"}, "reasoning": "Calculate"},
|
||
{"step_id": 3, "tool_name": "weather", "arguments": {"city": "Shanghai"}, "reasoning": "Check weather"},
|
||
])
|
||
synthesis_response = make_response(content="Based on search, calculation (100), and weather (25°C)")
|
||
|
||
gateway = make_mock_gateway([plan_response, synthesis_response])
|
||
engine = ReWOOEngine(llm_gateway=gateway)
|
||
|
||
result = await engine.execute(
|
||
messages=[{"role": "user", "content": "Search, calculate and check weather"}],
|
||
tools=[search_tool, calc_tool, weather_tool],
|
||
)
|
||
|
||
assert result.fallback_strategy is None
|
||
assert result.total_steps == 4 # 3 tool_calls + 1 final_answer
|
||
assert search_tool.call_count == 1
|
||
assert calc_tool.call_count == 1
|
||
assert weather_tool.call_count == 1
|
||
assert "100" in result.output
|
||
assert "25" in result.output
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fallback_strategies_constant(self):
|
||
"""验证 FALLBACK_STRATEGIES 常量"""
|
||
assert ReWOOEngine.FALLBACK_STRATEGIES == ["simplified_rewoo", "react", "direct"]
|