"""U3 测试: Tool 组合 (SequentialChain, ParallelFanOut, DynamicSelector) + MCPTool""" import asyncio import json import pytest from agentkit.tools.base import Tool from agentkit.tools.composition import DynamicSelector, ParallelFanOut, SequentialChain from agentkit.tools.function_tool import FunctionTool from agentkit.tools.mcp_tool import MCPTool # ── Helper: 创建简单 FunctionTool ──────────────────────── def _make_tool(name: str, func=None, **kw): if func is None: async def default_func(**kwargs): return {"tool": name, "input": kwargs} func = default_func return FunctionTool(name=name, description=f"Test tool {name}", func=func, **kw) # ── SequentialChain 测试 ───────────────────────────────── class TestSequentialChain: async def test_chain_two_tools(self): """顺序执行两个工具，前一个输出作为后一个输入""" async def step1(text: str) -> dict: return {"processed": text.upper(), "length": len(text)} async def step2(processed: str, length: int) -> dict: return {"result": f"{processed}({length})"} chain = SequentialChain( name="upper_and_format", description="Uppercase then format", tools=[ FunctionTool(name="step1", description="Step 1", func=step1), FunctionTool(name="step2", description="Step 2", func=step2), ], ) result = await chain.execute(text="hello") assert result["result"] == "HELLO(5)" async def test_chain_single_tool(self): """单工具链直接返回结果""" tool = _make_tool("solo") chain = SequentialChain(name="solo_chain", description="Solo", tools=[tool]) result = await chain.execute(x=1) assert result["tool"] == "solo" async def test_chain_preserves_dict(self): """链中每步都返回 dict""" async def add_field(x: int) -> dict: return {"x": x, "doubled": x * 2} async def add_label(x: int, doubled: int) -> dict: return {"x": x, "doubled": doubled, "label": f"{x}->{doubled}"} chain = SequentialChain( name="enrich", description="Enrich data", tools=[ FunctionTool(name="add_field", description="Add", func=add_field), FunctionTool(name="add_label", description="Label", func=add_label), ], ) result = await chain.execute(x=3) assert result["label"] == "3->6" # ── ParallelFanOut 测试 ────────────────────────────────── class TestParallelFanOut: async def test_parallel_execution(self): """并行执行三个工具""" async def search_web(query: str) -> dict: return {"web_results": [f"web: {query}"]} async def search_db(query: str) -> dict: return {"db_results": [f"db: {query}"]} async def search_cache(query: str) -> dict: return {"cache_results": [f"cache: {query}"]} fan_out = ParallelFanOut( name="multi_search", description="Search multiple sources", tools=[ FunctionTool(name="search_web", description="Web", func=search_web), FunctionTool(name="search_db", description="DB", func=search_db), FunctionTool(name="search_cache", description="Cache", func=search_cache), ], ) result = await fan_out.execute(query="AI") assert "web_results" in result assert "db_results" in result assert "cache_results" in result async def test_parallel_with_error(self): """并行执行中有工具失败，不影响其他""" async def good_tool(x: int) -> dict: return {"good": x * 2} async def bad_tool(x: int) -> dict: raise ValueError("boom") fan_out = ParallelFanOut( name="mixed", description="Mixed results", tools=[ FunctionTool(name="good", description="Good", func=good_tool), FunctionTool(name="bad", description="Bad", func=bad_tool), ], ) result = await fan_out.execute(x=5) assert result["good"] == 10 assert "_errors" in result assert any(e["tool"] == "bad" for e in result["_errors"]) async def test_parallel_namespace_merge(self): """namespace 合并策略：每个工具结果独立命名空间""" async def tool_a(x: int) -> dict: return {"result": "a"} async def tool_b(x: int) -> dict: return {"result": "b"} fan_out = ParallelFanOut( name="namespaced", description="Namespaced", tools=[ FunctionTool(name="tool_a", description="A", func=tool_a), FunctionTool(name="tool_b", description="B", func=tool_b), ], merge_strategy="namespace", ) result = await fan_out.execute(x=1) assert "tool_a" in result assert "tool_b" in result assert result["tool_a"]["result"] == "a" assert result["tool_b"]["result"] == "b" # ── DynamicSelector 测试 ───────────────────────────────── class TestDynamicSelector: async def test_keyword_selection_by_intent(self): """根据 _intent 参数选择工具""" search_tool = _make_tool("search_engine") calc_tool = _make_tool("calculator") selector = DynamicSelector( name="smart", description="Smart selector", tools=[search_tool, calc_tool], mode="keyword", ) result = await selector.execute(query="AI", _intent="search") assert result["tool"] == "search_engine" async def test_keyword_selection_by_input(self): """根据输入内容推断工具""" search_tool = _make_tool("search_data") calc_tool = _make_tool("calculate_result") selector = DynamicSelector( name="smart", description="Smart selector", tools=[search_tool, calc_tool], mode="keyword", ) result = await selector.execute(query="search for trends") assert result["tool"] == "search_data" async def test_no_matching_tool(self): """无匹配工具时返回错误""" tool_a = _make_tool("tool_a") selector = DynamicSelector( name="smart", description="Smart selector", tools=[tool_a], mode="keyword", ) result = await selector.execute(x=999) # 应该仍然选择 tool_a（best_score=0 的默认） assert "tool" in result or "error" in result async def test_llm_selection(self): """LLM 模式选择工具""" search_tool = _make_tool("search_engine") calc_tool = _make_tool("calculator") class MockLLM: async def chat(self, messages, **kwargs): return "0" # 选择第一个工具 selector = DynamicSelector( name="smart", description="Smart selector", tools=[search_tool, calc_tool], mode="llm", llm_client=MockLLM(), ) result = await selector.execute(query="find information") assert result["tool"] == "search_engine" async def test_llm_fallback_to_keyword(self): """LLM 失败时回退到关键字模式""" search_tool = _make_tool("search_engine") class BrokenLLM: async def chat(self, messages, **kwargs): raise ConnectionError("LLM unavailable") selector = DynamicSelector( name="smart", description="Smart selector", tools=[search_tool], mode="llm", llm_client=BrokenLLM(), ) result = await selector.execute(query="search something") assert result["tool"] == "search_engine" # ── MCPTool 测试 ───────────────────────────────────────── class TestMCPTool: async def test_mcp_tool_execute(self): """MCPTool 通过 client 调用远程工具""" class MockMCPClient: async def call_tool(self, name, arguments): return { "content": [ {"type": "text", "text": json.dumps({"status": "ok", "data": arguments})} ] } tool = MCPTool( name="remote_search", description="Remote search via MCP", client=MockMCPClient(), ) result = await tool.execute(query="test") assert result["status"] == "ok" assert result["data"]["query"] == "test" async def test_mcp_tool_text_response(self): """MCP 返回非 JSON 文本""" class MockMCPClient: async def call_tool(self, name, arguments): return { "content": [ {"type": "text", "text": "plain text result"} ] } tool = MCPTool( name="remote_echo", description="Remote echo", client=MockMCPClient(), ) result = await tool.execute(text="hello") assert result["result"] == "plain text result" async def test_mcp_tool_raw_response(self): """MCP 返回非标准格式""" class MockMCPClient: async def call_tool(self, name, arguments): return {"raw_key": "raw_value"} tool = MCPTool( name="remote_raw", description="Raw response", client=MockMCPClient(), ) result = await tool.execute(x=1) assert result["raw_key"] == "raw_value"