164 lines
6.0 KiB
Python
164 lines
6.0 KiB
Python
"""ReAct Engine 集成测试 - 完整 ReAct 循环"""
|
||
|
||
import pytest
|
||
|
||
from agentkit.llm.gateway import LLMGateway
|
||
from agentkit.llm.protocol import LLMResponse, TokenUsage, ToolCall
|
||
from agentkit.tools.base import Tool
|
||
|
||
|
||
class KnowledgeTool(Tool):
|
||
"""知识检索工具"""
|
||
|
||
def __init__(self):
|
||
super().__init__(
|
||
name="retrieve_knowledge",
|
||
description="Retrieve knowledge from the knowledge base",
|
||
)
|
||
|
||
async def execute(self, **kwargs) -> dict:
|
||
query = kwargs.get("query", "")
|
||
return {"knowledge": f"Knowledge about {query}", "relevance": 0.95}
|
||
|
||
|
||
class GenerateTool(Tool):
|
||
"""内容生成工具"""
|
||
|
||
def __init__(self):
|
||
super().__init__(
|
||
name="generate_content",
|
||
description="Generate content based on input",
|
||
)
|
||
|
||
async def execute(self, **kwargs) -> dict:
|
||
topic = kwargs.get("topic", "")
|
||
return {"content": f"Generated content about {topic}"}
|
||
|
||
|
||
class TestReActFullLoop:
|
||
"""完整 ReAct 循环:检索知识 → 生成内容 → 返回结果"""
|
||
|
||
async def test_knowledge_then_generate_loop(self):
|
||
from agentkit.core.react import ReActEngine, ReActResult
|
||
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
knowledge_tool = KnowledgeTool()
|
||
generate_tool = GenerateTool()
|
||
|
||
gateway = MagicMock(spec=LLMGateway)
|
||
gateway.chat = AsyncMock(side_effect=[
|
||
# Step 1: LLM 决定检索知识
|
||
LLMResponse(
|
||
content="",
|
||
model="test-model",
|
||
usage=TokenUsage(prompt_tokens=50, completion_tokens=10),
|
||
tool_calls=[ToolCall(id="tc_1", name="retrieve_knowledge", arguments={"query": "AI agents"})],
|
||
),
|
||
# Step 2: LLM 决定生成内容
|
||
LLMResponse(
|
||
content="",
|
||
model="test-model",
|
||
usage=TokenUsage(prompt_tokens=80, completion_tokens=10),
|
||
tool_calls=[ToolCall(id="tc_2", name="generate_content", arguments={"topic": "AI agents"})],
|
||
),
|
||
# Step 3: LLM 返回最终答案
|
||
LLMResponse(
|
||
content="Based on the knowledge retrieved and content generated, here is the answer about AI agents.",
|
||
model="test-model",
|
||
usage=TokenUsage(prompt_tokens=100, completion_tokens=30),
|
||
),
|
||
])
|
||
|
||
engine = ReActEngine(llm_gateway=gateway)
|
||
result = await engine.execute(
|
||
messages=[{"role": "user", "content": "Tell me about AI agents"}],
|
||
tools=[knowledge_tool, generate_tool],
|
||
system_prompt="You are a knowledgeable AI assistant.",
|
||
)
|
||
|
||
assert isinstance(result, ReActResult)
|
||
assert result.total_steps == 3
|
||
assert "AI agents" in result.output
|
||
assert result.total_tokens == 50 + 10 + 80 + 10 + 100 + 30
|
||
|
||
# 验证轨迹
|
||
assert result.trajectory[0].tool_name == "retrieve_knowledge"
|
||
assert result.trajectory[1].tool_name == "generate_content"
|
||
assert result.trajectory[2].action == "final_answer"
|
||
|
||
async def test_react_with_error_recovery(self):
|
||
"""带错误恢复的 ReAct 循环"""
|
||
from agentkit.core.react import ReActEngine
|
||
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
class FlakyTool(Tool):
|
||
def __init__(self):
|
||
super().__init__(name="flaky_api", description="A flaky API tool")
|
||
self._call_count = 0
|
||
|
||
async def execute(self, **kwargs) -> dict:
|
||
self._call_count += 1
|
||
if self._call_count == 1:
|
||
raise ConnectionError("API timeout")
|
||
return {"data": "success on retry"}
|
||
|
||
flaky_tool = FlakyTool()
|
||
|
||
gateway = MagicMock(spec=LLMGateway)
|
||
gateway.chat = AsyncMock(side_effect=[
|
||
# Step 1: LLM 调用 flaky API(第一次失败)
|
||
LLMResponse(
|
||
content="",
|
||
model="test-model",
|
||
usage=TokenUsage(prompt_tokens=50, completion_tokens=10),
|
||
tool_calls=[ToolCall(id="tc_1", name="flaky_api", arguments={})],
|
||
),
|
||
# Step 2: LLM 收到错误后重试
|
||
LLMResponse(
|
||
content="",
|
||
model="test-model",
|
||
usage=TokenUsage(prompt_tokens=80, completion_tokens=10),
|
||
tool_calls=[ToolCall(id="tc_2", name="flaky_api", arguments={})],
|
||
),
|
||
# Step 3: LLM 返回最终答案
|
||
LLMResponse(
|
||
content="After retrying, I got the data successfully.",
|
||
model="test-model",
|
||
usage=TokenUsage(prompt_tokens=100, completion_tokens=20),
|
||
),
|
||
])
|
||
|
||
engine = ReActEngine(llm_gateway=gateway)
|
||
result = await engine.execute(
|
||
messages=[{"role": "user", "content": "Call the flaky API"}],
|
||
tools=[flaky_tool],
|
||
)
|
||
|
||
assert result.total_steps == 3
|
||
# 第一次调用失败,但错误信息被包含在观察中
|
||
assert "error" in str(result.trajectory[0].result).lower() or "failed" in str(result.trajectory[0].result).lower()
|
||
# 第二次调用成功
|
||
assert result.trajectory[1].result == {"data": "success on retry"}
|
||
assert result.output == "After retrying, I got the data successfully."
|
||
|
||
|
||
class TestQualityGatePlaceholder:
|
||
"""Quality Gate 集成占位(将在 U5 实现)"""
|
||
|
||
async def test_react_result_has_quality_metrics_placeholder(self):
|
||
"""验证 ReActResult 可扩展以支持 Quality Gate"""
|
||
from agentkit.core.react import ReActResult, ReActStep
|
||
|
||
result = ReActResult(
|
||
output="test",
|
||
trajectory=[ReActStep(step=1, action="final_answer", content="test")],
|
||
total_steps=1,
|
||
total_tokens=10,
|
||
)
|
||
# ReActResult 应是一个 dataclass,可以正常访问属性
|
||
assert result.output == "test"
|
||
assert result.total_steps == 1
|
||
# 未来可以扩展添加 quality_score 等字段
|