Coverage for agentos/agent/pipeline.py: 33%
122 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 07:12 +0800
1"""
2多Agent编排管道 — Conditional / Parallel / Router。
4v1.5.1: 支持条件路由(ConditionalPipeline)、并行扇出(ParallelPipeline)、
5 动态路由(RouterAgent) 三种生产级编排拓扑。
6"""
8from __future__ import annotations
10import concurrent.futures
11from collections.abc import Callable
12from dataclasses import dataclass, field
14from agentos.agent.tool_agent import AgentConfig, AgentResult, ToolAgent
17@dataclass
18class PipelineAgent:
19 """管道中的单个 Agent 节点。"""
21 name: str
22 agent: ToolAgent
23 config: AgentConfig | None = None
26@dataclass
27class PipelineResult:
28 """管道执行结果。"""
30 success: bool = True
31 steps: list[dict] = field(default_factory=list)
32 final_output: str = ""
33 total_tokens: int = 0
34 total_cost_usd: float = 0.0
35 total_duration_ms: float = 0.0
36 error: str = ""
38 @property
39 def output(self) -> str:
40 return self.final_output
43@dataclass
44class StepResult:
45 """单个步骤的结果包装。"""
47 agent_name: str
48 result: AgentResult
49 output_key: str | None = None
52# ── ConditionalPipeline — 条件路由 ──────────────────────────────────
54ConditionFn = Callable[[str], str] # 输入 → 下一个 agent 名称
57class ConditionalPipeline:
58 """基于条件路由的多 Agent 管道。
60 每个 Agent 执行完成后，通过条件函数决定下一个调用的 Agent。
61 支持 if-else / switch-case 风格的决策路由。
63 Usage::
65 cp = ConditionalPipeline()
66 cp.add("classifier", classifier_agent)
67 cp.add("legal", legal_agent)
68 cp.add("tech", tech_agent)
70 # 根据分类器输出决定下一跳
71 def route(output: str) -> str:
72 if "法律" in output: return "legal"
73 if "技术" in output: return "tech"
74 return "__END__"
76 result = cp.run("这份合同有问题吗？", router=route)
77 """
79 def __init__(self, max_hops: int = 5):
80 self._agents: dict[str, PipelineAgent] = {}
81 self._max_hops = max_hops
83 def add(self, name: str, agent: ToolAgent, config: AgentConfig | None = None):
84 self._agents[name] = PipelineAgent(name=name, agent=agent, config=config)
86 def run(
87 self, task: str, start_agent: str | None = None, router: ConditionFn | None = None
88 ) -> PipelineResult:
89 if start_agent is None and self._agents:
90 start_agent = next(iter(self._agents))
91 if start_agent not in self._agents:
92 return PipelineResult(success=False, error=f"Unknown start agent: {start_agent}")
94 current = start_agent
95 pipeline_output = ""
96 steps: list[dict] = []
97 total_tokens = 0
98 total_cost = 0.0
99 total_ms = 0.0
101 for hop in range(self._max_hops):
102 pa = self._agents[current]
103 result = pa.agent.run(task)
105 steps.append(
106 {
107 "hop": hop,
108 "agent": current,
109 "output": result.final_answer,
110 "tokens": result.total_tokens,
111 "cost": result.total_cost_usd,
112 "duration_ms": result.total_duration_ms,
113 }
114 )
115 total_tokens += result.total_tokens
116 total_cost += result.total_cost_usd
117 total_ms += result.total_duration_ms
118 pipeline_output = result.final_answer
120 if not result.success:
121 return PipelineResult(
122 success=False,
123 steps=steps,
124 final_output=pipeline_output,
125 total_tokens=total_tokens,
126 total_cost_usd=total_cost,
127 total_duration_ms=total_ms,
128 error=result.error,
129 )
131 if router is None:
132 break
134 next_agent = router(result.final_answer)
135 if next_agent == "__END__" or next_agent not in self._agents:
136 break
138 task = result.final_answer # 下一跳的输入是当前输出
139 current = next_agent
141 return PipelineResult(
142 success=True,
143 steps=steps,
144 final_output=pipeline_output,
145 total_tokens=total_tokens,
146 total_cost_usd=total_cost,
147 total_duration_ms=total_ms,
148 )
151# ── ParallelPipeline — 并行扇出 ──────────────────────────────────
154class ParallelPipeline:
155 """并行执行多个 Agent，聚合结果。
157 所有 Agent 同时接收同一个 task，各自独立运行，最后合并输出。
159 Usage::
161 pp = ParallelPipeline()
162 pp.add("analyst_1", agent_1)
163 pp.add("analyst_2", agent_2)
164 pp.add("analyst_3", agent_3)
166 result = pp.run("分析 Q3 财报",
167 aggregator=lambda results: "\\n---\\n".join(results.values()))
168 """
170 def __init__(self, max_workers: int = 5):
171 self._agents: dict[str, PipelineAgent] = {}
172 self._max_workers = max_workers
174 def add(self, name: str, agent: ToolAgent, config: AgentConfig | None = None):
175 self._agents[name] = PipelineAgent(name=name, agent=agent, config=config)
177 def run(
178 self, task: str, aggregator: Callable[[dict[str, str]], str] | None = None
179 ) -> PipelineResult:
180 if not self._agents:
181 return PipelineResult(success=False, error="No agents registered")
183 def _run_one(pa: PipelineAgent) -> tuple[str, AgentResult]:
184 return pa.name, pa.agent.run(task)
186 results: dict[str, AgentResult] = {}
187 total_tokens = 0
188 total_cost = 0.0
189 total_ms = 0.0
190 errors: list[str] = []
192 with concurrent.futures.ThreadPoolExecutor(max_workers=self._max_workers) as pool:
193 futures = {pool.submit(_run_one, pa): pa.name for pa in self._agents.values()}
194 for fut in concurrent.futures.as_completed(futures):
195 try:
196 name, result = fut.result()
197 results[name] = result
198 total_tokens += result.total_tokens
199 total_cost += result.total_cost_usd
200 total_ms = max(total_ms, result.total_duration_ms)
201 except Exception as e:
202 errors.append(f"{futures[fut]}: {e}")
204 if errors and not results:
205 return PipelineResult(success=False, error="; ".join(errors))
207 raw_outputs = {name: r.final_answer for name, r in results.items()}
208 if aggregator:
209 final = aggregator(raw_outputs)
210 else:
211 parts = [f"## {name}\n{out}" for name, out in raw_outputs.items()]
212 final = "\n\n".join(parts)
214 return PipelineResult(
215 success=len(errors) == 0,
216 steps=[
217 {"agent": name, "output": r.final_answer, "tokens": r.total_tokens}
218 for name, r in results.items()
219 ],
220 final_output=final,
221 total_tokens=total_tokens,
222 total_cost_usd=total_cost,
223 total_duration_ms=total_ms,
224 error="; ".join(errors) if errors else "",
225 )
228# ── RouterAgent — 动态路由 ──────────────────────────────────────
230RouterFn = Callable[[str], tuple[str, str]] # (task, str) → (next_agent_name, rewritten_task)
233class RouterAgent:
234 """动态路由编排器 — 根据初始化内容选择合适的 Agent。
236 使用一个分类器 Agent 先分析任务，再动态路由到目标 Agent。
238 Usage::
240 ra = RouterAgent(classifier_agent)
241 ra.register("code", code_agent, description="代码生成/调试任务")
242 ra.register("writing", writer_agent, description="写作/翻译/总结任务")
243 ra.register("research", research_agent, description="调研/搜索/分析任务")
245 result = ra.run("帮我写一个 Python 快速排序")
246 """
248 def __init__(self, classifier: ToolAgent):
249 self._classifier = classifier
250 self._routes: dict[str, tuple[ToolAgent, str]] = {}
252 def register(self, name: str, agent: ToolAgent, description: str = ""):
253 self._routes[name] = (agent, description)
255 def run(self, task: str) -> PipelineResult:
256 if not self._routes:
257 return PipelineResult(success=False, error="No routes registered")
259 # 构建分类提示
260 route_desc = "\n".join(f"- {name}: {desc}" for name, (_, desc) in self._routes.items())
261 classify_task = (
262 f"Analyze the following task and output ONLY the best matching route name "
263 f"from the list below. Output just the name, nothing else.\n\n"
264 f"Available routes:\n{route_desc}\n\n"
265 f"Task: {task}\n\n"
266 f"Route:"
267 )
269 class_result = self._classifier.run(classify_task)
270 target = class_result.final_answer.strip().lower()
272 # 模糊匹配
273 matched = None
274 for name in self._routes:
275 if name.lower() in target:
276 matched = name
277 break
279 if matched is None:
280 # 回退到第一个
281 matched = next(iter(self._routes))
283 agent, _ = self._routes[matched]
284 result = agent.run(task)
286 return PipelineResult(
287 success=result.success,
288 steps=[
289 {"agent": "router", "output": f"Classified as: {matched}"},
290 {"agent": matched, "output": result.final_answer, "tokens": result.total_tokens},
291 ],
292 final_output=result.final_answer,
293 total_tokens=result.total_tokens,
294 total_cost_usd=result.total_cost_usd,
295 total_duration_ms=result.total_duration_ms,
296 )