Coverage for agentos/agent/tool_agent.py: 25%
281 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 17:01 +0800
1"""
2Tool-Using Agent — 基于 LLM Function Calling 的自主 Agent 循环。
4核心模式:
5 用户任务 → LLM 推理(tool_calls) → 工具执行 → 结果回传 → 循环直到完成
7v1.16.1: +CircuitBreaker +ToolOutputValidator +Metrics integrated into ToolExecutor/Agent.
8v1.3.38: +streaming, retry, checkpoint/resume, tool error handling, mock provider.
9"""
11from __future__ import annotations
13import json
14import os
15import time
16from collections.abc import Callable, Generator
17from dataclasses import dataclass, field
19from agentos.llm.base import (
20 CompletionChoice,
21 CompletionResult,
22 CompletionUsage,
23 LLMProvider,
24 Message,
25 MessageRole,
26 Tool,
27 ToolCall,
28)
29from agentos.tools.circuit_breaker import CircuitBreaker
30from agentos.tools.metrics import MetricsCollector
31from agentos.tools.validation import ToolOutputValidator, ToolResult, ValidationResult
33__all__ = [
34 "ToolAgent",
35 "AgentConfig",
36 "AgentStep",
37 "AgentResult",
38 "ToolExecutor",
39 "MockLLMProvider",
40]
43# ── 数据类型 ─────────────────────────────────────────────────────
46@dataclass
47class AgentConfig:
48 max_steps: int = 10
49 temperature: float = 0.0
50 max_tokens: int = 4096
51 verbose: bool = False
52 stop_on_error: bool = True
53 max_retries: int = 2
54 retry_delay: float = 0.5
55 checkpoint_dir: str = ""
58@dataclass
59class AgentStep:
60 step: int
61 thought: str = ""
62 tool_calls: list[ToolCall] = field(default_factory=list)
63 tool_results: dict[str, str] = field(default_factory=dict)
64 finish_reason: str = ""
65 tokens_used: int = 0
66 cost_usd: float = 0.0
67 duration_ms: float = 0.0
70@dataclass
71class AgentResult:
72 success: bool = True
73 final_answer: str = ""
74 steps: list[AgentStep] = field(default_factory=list)
75 total_steps: int = 0
76 total_tokens: int = 0
77 total_cost_usd: float = 0.0
78 total_duration_ms: float = 0.0
79 error: str | None = None
82# ── 工具执行器 ───────────────────────────────────────────────────
85class ToolExecutor:
86 """工具注册与执行器。
88 v1.16.1: 集成 CircuitBreaker（熔断保护）、ToolOutputValidator（输出校验）、
89 MetricsCollector（指标收集）。所有参数均为可选，不传则退化为原始行为。
90 """
92 def __init__(
93 self,
94 circuit_breaker: CircuitBreaker | None = None,
95 validator: ToolOutputValidator | None = None,
96 metrics: MetricsCollector | None = None,
97 ):
98 self._tools: dict[str, Callable[..., str]] = {}
99 self._schemas: dict[str, Tool] = {}
100 self._cb = circuit_breaker
101 self._validator = validator
102 self._metrics = metrics
104 def register(self, tool: Tool, handler: Callable[..., str]) -> None:
105 self._tools[tool.function.name] = handler
106 self._schemas[tool.function.name] = tool
108 def get_schemas(self) -> list[Tool]:
109 return list(self._schemas.values())
111 def execute(self, tool_call: ToolCall) -> str:
112 """执行工具调用，经 CircuitBreaker → 执行 → Validator → Metrics 全链路。"""
113 handler = self._tools.get(tool_call.name)
114 if handler is None:
115 return json.dumps({"error": f"Unknown tool: {tool_call.name}"})
117 # 1. 熔断器检查
118 if self._cb is not None:
119 if not self._cb.allow_request():
120 self._cb.record_failure()
121 return json.dumps(
122 {
123 "error": f"Circuit breaker OPEN for '{self._cb.name}'",
124 "circuit_state": self._cb.state.name,
125 }
126 )
128 t0 = time.monotonic()
129 try:
130 raw_output = str(handler(**tool_call.parsed_arguments))
132 # 2. 输出校验
133 validation_msg = ""
134 if self._validator is not None:
135 tr = ToolResult(output=raw_output, tool_name=tool_call.name)
136 val_result: ValidationResult = self._validator.validate(tr)
137 if not val_result.is_valid:
138 issues = "; ".join(i.message for i in val_result.issues)
139 validation_msg = f" [validation: {issues}]"
141 # 3. 熔断器记录成功
142 if self._cb is not None:
143 self._cb.record_success()
145 # 4. 指标记录
146 elapsed = (time.monotonic() - t0) * 1000
147 if self._metrics is not None:
148 self._metrics.get_counter("tool_calls_total").inc(tool_call.name)
149 self._metrics.get_counter("tool_calls_success").inc(tool_call.name)
150 self._metrics.get_timer("tool_latency_ms").record(elapsed)
152 return raw_output if not validation_msg else raw_output + validation_msg
154 except Exception as e:
155 elapsed = (time.monotonic() - t0) * 1000
156 if self._cb is not None:
157 self._cb.record_failure()
158 if self._metrics is not None:
159 self._metrics.get_counter("tool_calls_total").inc(tool_call.name)
160 self._metrics.get_counter("tool_calls_errors").inc(tool_call.name)
161 return json.dumps({"error": str(e)})
164# ── MockLLMProvider ──────────────────────────────────────────────
167class MockLLMProvider(LLMProvider):
168 """可编程响应的 Mock Provider，供集成测试使用。"""
170 def __init__(self, responses: list[dict]):
171 super().__init__(model="mock", api_key="mock")
172 self._responses = responses
173 self._cursor = 0
174 self.calls: list[dict] = []
176 def chat(self, messages=None, *, temperature=0, max_tokens=4096, tools=None, **kwargs):
177 if self._cursor >= len(self._responses):
178 return self._build_result({"content": "done", "finish_reason": "stop"})
179 resp = self._responses[self._cursor]
180 self._cursor += 1
181 self.calls.append(
182 {
183 "tools": [t.function.name for t in (tools or [])],
184 "cursor": self._cursor - 1,
185 }
186 )
187 return self._build_result(resp)
189 async def achat(self, *args, **kwargs):
190 return self.chat(*args, **kwargs)
192 @property
193 def provider_name(self) -> str:
194 return "mock"
196 @staticmethod
197 def text_response(content: str, finish_reason: str = "stop") -> dict:
198 return {"content": content, "finish_reason": finish_reason}
200 @staticmethod
201 def tool_response(name: str, arguments: dict, tool_call_id: str = "") -> dict:
202 tid = tool_call_id or f"tc_{name}"
203 return {
204 "content": "",
205 "tool_calls": [ToolCall(id=tid, name=name, arguments=json.dumps(arguments))],
206 "finish_reason": "tool_calls",
207 }
209 def _build_result(self, resp: dict) -> CompletionResult:
210 msg = Message(
211 role=MessageRole.ASSISTANT,
212 content=resp.get("content", ""),
213 tool_calls=resp.get("tool_calls"),
214 )
215 choice = CompletionChoice(
216 index=0,
217 message=msg,
218 finish_reason=resp.get("finish_reason", "stop"),
219 )
220 return CompletionResult(
221 id=f"mock_{self._cursor}",
222 model="mock-model",
223 choices=[choice],
224 usage=CompletionUsage(
225 prompt_tokens=5,
226 completion_tokens=len(resp.get("content", "")) + 3,
227 total_tokens=len(resp.get("content", "")) + 8,
228 ),
229 )
232# ── Tool-Using Agent ─────────────────────────────────────────────
235class ToolAgent:
236 """基于 LLM Function Calling 的自主 Agent。
238 用法:
239 from agentos.agent import ToolAgent, ToolExecutor
240 from agentos.llm import create_provider, Tool
242 provider = create_provider("openai")
243 executor = ToolExecutor()
244 executor.register(
245 Tool.from_function("get_weather", "获取天气", {"city": ...}),
246 lambda city: f"{city}: 22°C sunny"
247 )
248 agent = ToolAgent(provider, executor)
249 result = agent.run("北京天气怎么样？")
250 print(result.final_answer)
251 """
253 def __init__(
254 self,
255 provider: LLMProvider,
256 tool_executor: ToolExecutor,
257 *,
258 config: AgentConfig | None = None,
259 system_prompt: str = "",
260 metrics: MetricsCollector | None = None,
261 ):
262 self._provider = provider
263 self._executor = tool_executor
264 self._config = config or AgentConfig()
265 self._system_prompt = system_prompt or (
266 "你是一个智能助手。你可以使用工具来获取信息。"
267 "当你可以给出最终答案时，直接回答，不要再调用工具。"
268 "用中文回答。"
269 )
270 self._metrics = metrics
272 # ── 同步 ──────────────────────────────────────────────────
274 def run(self, task: str) -> AgentResult:
275 t0 = time.monotonic()
276 steps: list[AgentStep] = []
277 tools = self._executor.get_schemas()
278 messages: list[Message] = [
279 Message(role=MessageRole.SYSTEM, content=self._system_prompt),
280 Message(role=MessageRole.USER, content=task),
281 ]
282 return self._run_loop(messages, task, tools, steps, 1, t0)
284 def _run_loop(
285 self,
286 messages,
287 task,
288 tools,
289 steps,
290 start_step,
291 t0,
292 ) -> AgentResult:
293 final_answer = ""
294 total_tokens = 0
295 total_cost = 0.0
296 step_num = start_step
298 try:
299 for step_num in range(start_step, self._config.max_steps + 1):
300 result = self._call_with_retry(messages, tools)
301 step, done, final = self._process_step(result, step_num)
302 total_tokens += step.tokens_used
303 total_cost += step.cost_usd
304 steps.append(step)
305 if done:
306 final_answer = final
307 break
308 messages.append(result.choices[0].message)
309 for tc in step.tool_calls:
310 messages.append(
311 Message(
312 role=MessageRole.TOOL,
313 content=step.tool_results.get(tc.id, ""),
314 tool_call_id=tc.id,
315 )
316 )
317 self._checkpoint(messages, task, step_num)
318 else:
319 return self._make_result(
320 False,
321 "",
322 steps,
323 total_tokens,
324 total_cost,
325 t0,
326 f"Reached max steps ({self._config.max_steps}) without final answer",
327 )
328 except Exception as e:
329 return self._make_result(False, "", steps, total_tokens, total_cost, t0, str(e))
331 return self._make_result(True, final_answer, steps, total_tokens, total_cost, t0)
333 # ── 流式 ──────────────────────────────────────────────────
335 def run_stream(self, task: str) -> Generator[AgentStep, None, AgentResult]:
336 t0 = time.monotonic()
337 steps: list[AgentStep] = []
338 tools = self._executor.get_schemas()
339 messages: list[Message] = [
340 Message(role=MessageRole.SYSTEM, content=self._system_prompt),
341 Message(role=MessageRole.USER, content=task),
342 ]
343 total_tokens = 0
344 total_cost = 0.0
345 final_answer = ""
347 try:
348 for step_num in range(1, self._config.max_steps + 1):
349 result = self._call_with_retry(messages, tools)
350 step, done, final = self._process_step(result, step_num)
351 total_tokens += step.tokens_used
352 total_cost += step.cost_usd
353 yield step
354 steps.append(step)
355 if done:
356 final_answer = final
357 break
358 messages.append(result.choices[0].message)
359 for tc in step.tool_calls:
360 messages.append(
361 Message(
362 role=MessageRole.TOOL,
363 content=step.tool_results.get(tc.id, ""),
364 tool_call_id=tc.id,
365 )
366 )
367 self._checkpoint(messages, task, step_num)
368 else:
369 return self._make_result(
370 False,
371 "",
372 steps,
373 total_tokens,
374 total_cost,
375 t0,
376 f"Reached max steps ({self._config.max_steps}) without final answer",
377 )
378 except Exception as e:
379 return self._make_result(False, "", steps, total_tokens, total_cost, t0, str(e))
381 return self._make_result(True, final_answer, steps, total_tokens, total_cost, t0)
383 # ── 异步 ──────────────────────────────────────────────────
385 async def arun(self, task: str) -> AgentResult:
386 t0 = time.monotonic()
387 steps: list[AgentStep] = []
388 tools = self._executor.get_schemas()
389 messages: list[Message] = [
390 Message(role=MessageRole.SYSTEM, content=self._system_prompt),
391 Message(role=MessageRole.USER, content=task),
392 ]
393 final_answer = ""
394 total_tokens = 0
395 total_cost = 0.0
397 try:
398 for step_num in range(1, self._config.max_steps + 1):
399 result = await self._acall_with_retry(messages, tools)
400 step, done, final = self._process_step(result, step_num)
401 total_tokens += step.tokens_used
402 total_cost += step.cost_usd
403 steps.append(step)
404 if done:
405 final_answer = final
406 break
407 messages.append(result.choices[0].message)
408 for tc in step.tool_calls:
409 messages.append(
410 Message(
411 role=MessageRole.TOOL,
412 content=step.tool_results.get(tc.id, ""),
413 tool_call_id=tc.id,
414 )
415 )
416 self._checkpoint(messages, task, step_num)
417 else:
418 return self._make_result(
419 False,
420 "",
421 steps,
422 total_tokens,
423 total_cost,
424 t0,
425 f"Reached max steps ({self._config.max_steps}) without final answer",
426 )
427 except Exception as e:
428 return self._make_result(False, "", steps, total_tokens, total_cost, t0, str(e))
430 return self._make_result(True, final_answer, steps, total_tokens, total_cost, t0)
432 # ── 共享步骤逻辑 ───────────────────────────────────────────────
434 def _process_step(
435 self,
436 result: CompletionResult,
437 step_num: int,
438 ) -> tuple[AgentStep, bool, str]:
439 """处理单步 LLM 结果：构建 AgentStep、执行工具、判断终止。
441 Returns:
442 (step, done, final_answer)
443 """
444 step_t0 = time.monotonic()
445 choice = result.choices[0]
446 assistant_msg = choice.message
448 step = AgentStep(
449 step=step_num,
450 thought=assistant_msg.content,
451 tool_calls=assistant_msg.tool_calls or [],
452 finish_reason=choice.finish_reason,
453 tokens_used=result.usage.total_tokens,
454 cost_usd=result.usage.cost_usd,
455 duration_ms=(time.monotonic() - step_t0) * 1000,
456 )
458 if self._config.verbose:
459 self._log_step(step)
461 # Metrics: track LLM calls and tokens
462 if self._metrics is not None:
463 self._metrics.get_counter("llm_calls_total").inc()
464 self._metrics.get_counter("llm_tokens_total").inc(result.usage.total_tokens)
465 self._metrics.get_counter("agent_steps_total").inc()
467 # 无工具调用 → 终止，内容即为答案
468 if not assistant_msg.tool_calls:
469 return step, True, assistant_msg.content
471 # 执行工具调用
472 for tc in assistant_msg.tool_calls:
473 tool_result = self._executor.execute(tc)
474 step.tool_results[tc.id] = tool_result
475 if "error" in tool_result and self._config.stop_on_error:
476 raise RuntimeError(f"Tool '{tc.name}' error: {tool_result}")
478 # finish_reason == "stop" → 提前终止
479 if choice.finish_reason == "stop":
480 return step, True, assistant_msg.content
482 return step, False, ""
484 def _make_result(
485 self,
486 success: bool,
487 answer: str,
488 steps: list[AgentStep],
489 total_tokens: int,
490 total_cost: float,
491 t0: float,
492 error: str = None,
493 ) -> AgentResult:
494 """统一构造 AgentResult。"""
495 return AgentResult(
496 success=success,
497 final_answer=answer,
498 steps=steps,
499 total_steps=len(steps),
500 total_tokens=total_tokens,
501 total_cost_usd=total_cost,
502 total_duration_ms=(time.monotonic() - t0) * 1000,
503 error=error,
504 )
506 # ── Checkpoint / Resume ───────────────────────────────────
508 def resume(self) -> AgentResult:
509 if not self._config.checkpoint_dir:
510 raise ValueError("checkpoint_dir not configured")
511 ckpt_path = os.path.join(self._config.checkpoint_dir, "agent_checkpoint.json")
512 if not os.path.exists(ckpt_path):
513 raise FileNotFoundError(f"No checkpoint found at {ckpt_path}")
514 with open(ckpt_path) as f:
515 data = json.load(f)
516 task = data["task"]
517 start_step = data["step"] + 1
518 messages_raw = data["messages"]
519 messages = [
520 Message(
521 role=MessageRole(m["role"]),
522 content=m["content"],
523 tool_call_id=m.get("tool_call_id"),
524 tool_calls=(
525 [ToolCall(**tc) for tc in m["tool_calls"]] if m.get("tool_calls") else None
526 ),
527 )
528 for m in messages_raw
529 ]
530 t0 = time.monotonic()
531 tools = self._executor.get_schemas()
532 steps: list[AgentStep] = []
533 return self._run_loop(messages, task, tools, steps, start_step, t0)
535 def _checkpoint(self, messages: list[Message], task: str, step: int) -> None:
536 if not self._config.checkpoint_dir:
537 return
538 ckpt_path = os.path.join(self._config.checkpoint_dir, "agent_checkpoint.json")
539 data = {
540 "task": task,
541 "step": step,
542 "messages": [
543 {
544 "role": m.role.value,
545 "content": m.content,
546 "tool_call_id": m.tool_call_id,
547 "tool_calls": (
548 [
549 {"id": tc.id, "name": tc.name, "arguments": tc.arguments}
550 for tc in m.tool_calls
551 ]
552 if m.tool_calls
553 else None
554 ),
555 }
556 for m in messages
557 ],
558 }
559 with open(ckpt_path, "w") as f:
560 json.dump(data, f, ensure_ascii=False)
562 # ── 内部方法 ──────────────────────────────────────────────
564 def _call_with_retry(self, messages: list[Message], tools: list[Tool]) -> CompletionResult:
565 last_error = None
566 for attempt in range(self._config.max_retries + 1):
567 try:
568 return self._provider.chat(
569 messages,
570 temperature=self._config.temperature,
571 max_tokens=self._config.max_tokens,
572 tools=tools if tools else None,
573 )
574 except Exception as e:
575 last_error = e
576 if attempt < self._config.max_retries:
577 time.sleep(self._config.retry_delay)
578 raise last_error # type: ignore
580 async def _acall_with_retry(
581 self, messages: list[Message], tools: list[Tool]
582 ) -> CompletionResult:
583 last_error = None
584 for attempt in range(self._config.max_retries + 1):
585 try:
586 return await self._provider.achat(
587 messages,
588 temperature=self._config.temperature,
589 max_tokens=self._config.max_tokens,
590 tools=tools if tools else None,
591 )
592 except Exception as e:
593 last_error = e
594 if attempt < self._config.max_retries:
595 import asyncio
597 await asyncio.sleep(self._config.retry_delay)
598 raise last_error # type: ignore
600 def _log_step(self, step: AgentStep) -> None:
601 print(
602 f"\n── Step {step.step} ({step.duration_ms:.0f}ms, {step.tokens_used}t, ${step.cost_usd:.6f}) ──"
603 )
604 if step.thought:
605 print(f" Thought: {step.thought}")
606 if step.tool_calls:
607 for tc in step.tool_calls:
608 result_preview = step.tool_results.get(tc.id, "")[:100]
609 print(f" Tool: {tc.name}({tc.arguments}) → {result_preview}")
610 print(f" Finish: {step.finish_reason}")