Coverage for agentos/agent/production.py: 94%
167 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 10:28 +0800
1"""ProductionAgent — ToolAgent with ModelRouter + AuditLogger + SmartCache.
3ProductionAgent wraps the standard ToolAgent with production-grade features:
4 - ModelRouter: auto-selects the best model per task (complexity-aware)
5 - AuditLogger: immutable audit trail of every tool call and step
6 - SmartCache: LLM response caching (exact + fuzzy match) reduces API costs
7 - Automatic cost/latency tracking and budget enforcement
8 - Structured task classification (trivial → expert)
10Usage:
11 from agentos.agent import ProductionAgent, ToolExecutor
12 from agentos.llm import create_provider, SmartCache
13 from agentos.agent.model_router import ModelRouter
15 router = ModelRouter.with_defaults(daily_budget_usd=50.0)
16 cache = SmartCache()
17 provider = create_provider("openai", api_key="...")
18 executor = ToolExecutor()
19 executor.register(...)
21 agent = ProductionAgent(provider, executor, router=router, cache=cache)
22 result = agent.run("总结这篇论文的核心观点")
23 # → automatically routes to gpt-4o for complex analysis
24 # → caches responses to save cost on repeated queries
25 # → all tool calls audited in audit.jsonl
27v1.9.12: +SmartCache integration, cache-aware cost tracking.
28v1.9.10: Initial — ModelRouter + AuditLogger bidirectional integration.
29"""
31from __future__ import annotations
33import time
34import uuid
35from collections.abc import Callable, Generator
36from dataclasses import dataclass, field
37from typing import Any
39from agentos.agent.model_router import (
40 ModelRouter,
41 ModelSpec,
42 RequestSpec,
43 TaskComplexity,
44 TaskPriority,
45)
46from agentos.agent.tool_agent import (
47 AgentConfig,
48 AgentResult,
49 AgentStep,
50 ToolAgent,
51 ToolExecutor,
52)
53from agentos.llm.base import LLMProvider
54from agentos.llm.smart_cache import SmartCache
55from agentos.security.audit_logger import (
56 AuditActionCategory,
57 AuditLogger,
58 AuditSeverity,
59)
61__all__ = [
62 "ProductionAgent",
63 "ProductionConfig",
64 "ComplexityEstimator",
65 "ComplexityEstimate",
66]
69# ── Complexity Estimator ────────────────────────────────────────────
72@dataclass
73class ComplexityEstimate:
74 complexity: TaskComplexity
75 priority: TaskPriority
76 estimated_tokens: int
77 reason: str
80class ComplexityEstimator:
81 """Estimates task complexity from the user's request text.
83 Uses keyword heuristics to classify tasks from TRIVIAL to EXPERT.
84 Production use should integrate with a lightweight classifier model.
85 """
87 # keywords that indicate higher complexity
88 COMPLEX_KEYWORDS = [
89 "分析",
90 "对比",
91 "比较",
92 "总结",
93 "调研",
94 "研究",
95 "评估",
96 "analyze",
97 "compare",
98 "research",
99 "evaluate",
100 "assess",
101 "代码审查",
102 "架构",
103 "重构",
104 "安全审计",
105 "code review",
106 "architecture",
107 "refactor",
108 "audit",
109 ]
110 EXPERT_KEYWORDS = [
111 "深度",
112 "全面",
113 "完整",
114 "生产级",
115 "企业级",
116 "从零",
117 "comprehensive",
118 "production",
119 "enterprise",
120 "from scratch",
121 "论文",
122 "学术",
123 "法律",
124 "paper",
125 "academic",
126 "legal",
127 ]
128 TRIVIAL_KEYWORDS = [
129 "天气",
130 "时间",
131 "翻译",
132 "计算",
133 "换算",
134 "几点",
135 "日期",
136 "weather",
137 "time",
138 "translate",
139 "calculate",
140 "convert",
141 ]
142 URGENT_KEYWORDS = [
143 "快",
144 "紧急",
145 "马上",
146 "立即",
147 "urgent",
148 "asap",
149 "immediately",
150 "now",
151 ]
153 def estimate(self, task: str) -> ComplexityEstimate:
154 task_lower = task.lower()
156 # check for urgency
157 priority = TaskPriority.NORMAL
158 for kw in self.URGENT_KEYWORDS:
159 if kw in task_lower:
160 priority = TaskPriority.HIGH
161 break
163 # complexity
164 complexity = TaskComplexity.MODERATE
166 expert_hits = sum(1 for kw in self.EXPERT_KEYWORDS if kw in task_lower)
167 complex_hits = sum(1 for kw in self.COMPLEX_KEYWORDS if kw in task_lower)
168 trivial_hits = sum(1 for kw in self.TRIVIAL_KEYWORDS if kw in task_lower)
170 if expert_hits >= 2 or len(task) > 500:
171 complexity = TaskComplexity.EXPERT
172 elif expert_hits >= 1 or complex_hits >= 3:
173 complexity = TaskComplexity.COMPLEX
174 elif complex_hits >= 1 or len(task) > 200:
175 complexity = TaskComplexity.MODERATE
176 elif trivial_hits >= 1:
177 complexity = TaskComplexity.TRIVIAL
178 else:
179 complexity = TaskComplexity.SIMPLE
181 # estimate token count
182 # rough heuristic: ~1.5 chars per token for Chinese, ~4 for English
183 char_count = len(task)
184 if any("\u4e00" <= c <= "\u9fff" for c in task):
185 estimated_tokens = max(50, char_count // 1.5)
186 else:
187 estimated_tokens = max(50, char_count // 4)
189 estimated_tokens = int(min(estimated_tokens, 100_000))
191 return ComplexityEstimate(
192 complexity=complexity,
193 priority=priority,
194 estimated_tokens=estimated_tokens,
195 reason=f"task_len={char_count} chars, expert_hits={expert_hits}, "
196 f"complex_hits={complex_hits}, trivial_hits={trivial_hits}",
197 )
200# ── ProductionConfig ────────────────────────────────────────────────
203@dataclass
204class ProductionConfig:
205 agent_config: AgentConfig = field(default_factory=AgentConfig)
206 enable_audit: bool = True
207 enable_routing: bool = True
208 enable_cache: bool = True
209 audit_log_dir: str = ""
210 session_id: str = ""
211 budget_usd: float = 50.0
212 fallback_on_error: bool = True
215# ── ProductionAgent ─────────────────────────────────────────────────
218class ProductionAgent:
219 """Production-grade ToolAgent wrapper with routing and auditing.
221 Architecture:
222 User Task
223 │
224 ▼
225 ComplexityEstimator → classifies task
226 │
227 ▼
228 ModelRouter → selects best model for this complexity
229 │
230 ▼
231 ToolAgent.run() ───→ AuditLogger (every step & tool call)
232 │
233 ▼
234 AgentResult (with routing metadata + audit trail)
236 All tool calls and agent steps are automatically audited with
237 SHA256-chained immutable entries. Model selection is automatic
238 based on task analysis.
239 """
241 def __init__(
242 self,
243 provider: LLMProvider,
244 tool_executor: ToolExecutor,
245 *,
246 config: ProductionConfig | None = None,
247 router: ModelRouter | None = None,
248 cache: SmartCache | None = None,
249 system_prompt: str = "",
250 ):
251 self._config = config or ProductionConfig()
253 # cache (SmartCache API: get/set/contains/clear/size)
254 if cache and self._config.enable_cache:
255 self._cache = cache
256 self._provider = provider # SmartCache does not wrap providers
257 else:
258 self._cache = None
259 self._provider = provider
261 self._executor = tool_executor
262 self._system_prompt = system_prompt
264 # routing
265 self._router = router or ModelRouter.with_defaults(
266 daily_budget_usd=self._config.budget_usd,
267 )
268 self._estimator = ComplexityEstimator()
270 # auditing
271 self._session_id = self._config.session_id or f"sess-{uuid.uuid4().hex[:8]}"
272 self._audit = (
273 AuditLogger(log_dir=self._config.audit_log_dir)
274 if self._config.enable_audit
275 else None
276 )
278 # track routing decisions
279 self._last_route: RequestSpec | None = None
280 self._last_model: ModelSpec | None = None
282 # ── public API ──────────────────────────────────────────────
284 def run(self, task: str) -> AgentResult:
285 t_start = time.time()
287 # 1. classify
288 estimate = self._estimator.estimate(task)
290 # 2. route
291 route_spec = RequestSpec(
292 estimated_input_tokens=estimate.estimated_tokens,
293 estimated_output_tokens=estimate.estimated_tokens // 2,
294 complexity=estimate.complexity,
295 priority=estimate.priority,
296 task_id=f"task-{uuid.uuid4().hex[:8]}",
297 session_id=self._session_id,
298 )
299 self._last_route = route_spec
300 route_result = self._router.route(route_spec)
302 if not route_result.success:
303 return AgentResult(
304 success=False,
305 error=f"Model routing failed: {route_result.reason}",
306 )
308 self._last_model = route_result.model
310 # audit: route decision
311 if self._audit:
312 self._audit.log(
313 agent="production",
314 action="model_route",
315 target=route_result.model.name,
316 result="success",
317 severity=AuditSeverity.INFO,
318 category=AuditActionCategory.CONFIG_CHANGE,
319 session_id=self._session_id,
320 details={
321 "complexity": estimate.complexity.name,
322 "priority": estimate.priority.name,
323 "estimated_tokens": estimate.estimated_tokens,
324 "estimated_cost": round(route_result.estimated_cost, 6),
325 "fallback_chain": route_result.fallback_chain,
326 "budget_remaining": round(self._router.daily_budget_remaining, 4),
327 },
328 )
330 # 3. build agent config with model info
331 agent_config = AgentConfig(
332 max_steps=self._config.agent_config.max_steps,
333 temperature=self._config.agent_config.temperature,
334 max_tokens=self._config.agent_config.max_tokens,
335 verbose=self._config.agent_config.verbose,
336 stop_on_error=self._config.agent_config.stop_on_error,
337 max_retries=self._config.agent_config.max_retries,
338 retry_delay=self._config.agent_config.retry_delay,
339 )
341 # 4. create tool agent and wrap tool executor with audit
342 agent = ToolAgent(
343 provider=self._provider,
344 tool_executor=self._make_audited_executor(),
345 config=agent_config,
346 system_prompt=self._system_prompt,
347 )
349 # 5. run
350 if self._audit:
351 self._audit.log(
352 agent="production",
353 action="agent_start",
354 target=task[:100],
355 result="success",
356 severity=AuditSeverity.INFO,
357 category=AuditActionCategory.AGENT_INVOKE,
358 session_id=self._session_id,
359 details={
360 "complexity": estimate.complexity.name,
361 "model": route_result.model.name,
362 "cost_estimate": round(route_result.estimated_cost, 6),
363 },
364 )
366 result = agent.run(task)
367 elapsed_ms = (time.time() - t_start) * 1000
369 # 6. record model stats
370 self._router.record_request(
371 model_name=route_result.model.name,
372 success=result.success,
373 tokens_used=result.total_tokens,
374 cost_usd=result.total_cost_usd,
375 latency_ms=elapsed_ms,
376 )
378 # 7. audit: agent completion
379 if self._audit:
380 self._audit.log(
381 agent="production",
382 action="agent_end",
383 result="success" if result.success else "failure",
384 severity=AuditSeverity.ERROR if not result.success else AuditSeverity.INFO,
385 category=AuditActionCategory.AGENT_INVOKE,
386 session_id=self._session_id,
387 duration_ms=elapsed_ms,
388 error_message=result.error or "",
389 details={
390 "steps": result.total_steps,
391 "tokens": result.total_tokens,
392 "cost": round(result.total_cost_usd, 6),
393 "model": route_result.model.name,
394 },
395 )
397 # attach routing metadata to result
398 result.total_cost_usd = result.total_cost_usd
399 result.total_duration_ms = elapsed_ms
401 # 8. capture cache stats
402 if self._cache and result.success:
403 # SmartCache does not track stats natively; no-op
404 pass
406 return result
408 # ── properties ───────────────────────────────────────────────
410 @property
411 def last_route(self) -> RequestSpec | None:
412 """The last routing request spec."""
413 return self._last_route
415 @property
416 def last_model(self) -> ModelSpec | None:
417 """The model used in the last run."""
418 return self._last_model
420 @property
421 def cache_stats(self) -> Any | None:
422 """Cache statistics if cache is enabled, else None.
424 SmartCache does not natively expose stats; returns basic info.
425 """
426 if not self._cache:
427 return None
428 return {
429 "size": self._cache.size,
430 }
432 @property
433 def cache_hit_rate(self) -> float:
434 """Cache hit rate (0.0-1.0). Returns 0.0 if cache disabled."""
435 if not self._cache:
436 return 0.0
437 # SmartCache does not expose hit/miss counters; return 0.0
438 return 0.0
440 @property
441 def cache_savings(self) -> float:
442 """Estimated USD saved by cache hits."""
443 if not self._cache:
444 return 0.0
445 return 0.0 # SmartCache does not track cost savings
447 # ── streaming ────────────────────────────────────────────────
449 def run_stream(self, task: str) -> Generator[AgentStep, None, AgentResult]:
450 """Streaming version — yields steps as they complete."""
451 estimate = self._estimator.estimate(task)
452 route_spec = RequestSpec(
453 estimated_input_tokens=estimate.estimated_tokens,
454 estimated_output_tokens=estimate.estimated_tokens // 2,
455 complexity=estimate.complexity,
456 priority=estimate.priority,
457 task_id=f"task-{uuid.uuid4().hex[:8]}",
458 session_id=self._session_id,
459 )
460 self._last_route = route_spec
461 self._router.route(route_spec)
463 agent = ToolAgent(
464 provider=self._provider,
465 tool_executor=self._make_audited_executor(),
466 config=self._config.agent_config,
467 system_prompt=self._system_prompt,
468 )
470 yield from agent.run_stream(task)
472 # ── accessors ───────────────────────────────────────────────
474 @property
475 def router(self) -> ModelRouter:
476 return self._router
478 @property
479 def audit(self) -> AuditLogger | None:
480 return self._audit
482 @property
483 def session_id(self) -> str:
484 return self._session_id
486 def route_summary(self) -> dict:
487 """Return routing + audit summary for the current session."""
488 summary = {
489 "session_id": self._session_id,
490 "router": self._router.summary() if self._router else {},
491 }
492 if self._audit:
493 summary["audit"] = self._audit.stats_summary()
494 if self._last_model:
495 summary["last_model"] = self._last_model.name
496 if hasattr(self._last_model, "tier"):
497 summary["last_model_tier"] = self._last_model.tier.name
498 return summary
500 # ── internal ────────────────────────────────────────────────
502 def _make_audited_executor(self) -> ToolExecutor:
503 """Wrap tool executor to auto-audit every call."""
504 if not self._audit:
505 return self._executor
507 audited = ToolExecutor()
509 # copy original tools with audit wrapping
510 for schema in self._executor.get_schemas():
511 original_name = schema.function.name
513 # Capture the original execute method for this tool
514 def make_wrapper(name: str) -> Callable[..., str]:
515 def wrapper(**kwargs: Any) -> str:
516 t_start = time.time()
517 try:
518 result = self._executor.execute(
519 type(
520 "FakeCall",
521 (),
522 {
523 "name": name,
524 "parsed_arguments": kwargs,
525 },
526 )()
527 )
528 elapsed = (time.time() - t_start) * 1000
529 self._audit.log(
530 agent="production",
531 action=f"tool:{name}",
532 target=str(list(kwargs.keys())),
533 result="success",
534 severity=AuditSeverity.DEBUG,
535 category=AuditActionCategory.TOOL_CALL,
536 session_id=self._session_id,
537 duration_ms=elapsed,
538 details={"arguments": kwargs},
539 )
540 return result
541 except Exception as exc:
542 elapsed = (time.time() - t_start) * 1000
543 self._audit.log(
544 agent="production",
545 action=f"tool:{name}",
546 target=str(list(kwargs.keys())),
547 result="failure",
548 severity=AuditSeverity.ERROR,
549 category=AuditActionCategory.TOOL_CALL,
550 session_id=self._session_id,
551 duration_ms=elapsed,
552 error_message=str(exc),
553 )
554 raise
556 return wrapper
558 audited.register(schema, make_wrapper(original_name))
560 return audited