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