Coverage for agentos/agent/production.py: 34%
166 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 07:12 +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 — wrap provider transparently
254 if cache and self._config.enable_cache:
255 self._cache = cache
256 self._provider = cache.wrap(provider)
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(
274 log_dir=self._config.audit_log_dir,
275 max_events=100_000,
276 max_age_days=90,
277 )
278 if self._config.enable_audit
279 else None
280 )
282 # track routing decisions
283 self._last_route: RequestSpec | None = None
284 self._last_model: ModelSpec | None = None
286 # ── public API ──────────────────────────────────────────────
288 def run(self, task: str) -> AgentResult:
289 t_start = time.time()
291 # 1. classify
292 estimate = self._estimator.estimate(task)
294 # 2. route
295 route_spec = RequestSpec(
296 estimated_input_tokens=estimate.estimated_tokens,
297 estimated_output_tokens=estimate.estimated_tokens // 2,
298 complexity=estimate.complexity,
299 priority=estimate.priority,
300 task_id=f"task-{uuid.uuid4().hex[:8]}",
301 session_id=self._session_id,
302 )
303 self._last_route = route_spec
304 route_result = self._router.route(route_spec)
306 if not route_result.success:
307 return AgentResult(
308 success=False,
309 error=f"Model routing failed: {route_result.reason}",
310 )
312 self._last_model = route_result.model
314 # audit: route decision
315 if self._audit:
316 self._audit.log(
317 agent="production",
318 action="model_route",
319 target=route_result.model.name,
320 result="success",
321 severity=AuditSeverity.INFO,
322 category=AuditActionCategory.CONFIG_CHANGE,
323 session_id=self._session_id,
324 details={
325 "complexity": estimate.complexity.name,
326 "priority": estimate.priority.name,
327 "estimated_tokens": estimate.estimated_tokens,
328 "estimated_cost": round(route_result.estimated_cost, 6),
329 "fallback_chain": route_result.fallback_chain,
330 "budget_remaining": round(self._router.daily_budget_remaining, 4),
331 },
332 )
334 # 3. build agent config with model info
335 agent_config = AgentConfig(
336 max_steps=self._config.agent_config.max_steps,
337 temperature=self._config.agent_config.temperature,
338 max_tokens=self._config.agent_config.max_tokens,
339 verbose=self._config.agent_config.verbose,
340 stop_on_error=self._config.agent_config.stop_on_error,
341 max_retries=self._config.agent_config.max_retries,
342 retry_delay=self._config.agent_config.retry_delay,
343 )
345 # 4. create tool agent and wrap tool executor with audit
346 agent = ToolAgent(
347 provider=self._provider,
348 tool_executor=self._make_audited_executor(),
349 config=agent_config,
350 system_prompt=self._system_prompt,
351 )
353 # 5. run
354 if self._audit:
355 self._audit.log(
356 agent="production",
357 action="agent_start",
358 target=task[:100],
359 result="success",
360 severity=AuditSeverity.INFO,
361 category=AuditActionCategory.AGENT_INVOKE,
362 session_id=self._session_id,
363 details={
364 "complexity": estimate.complexity.name,
365 "model": route_result.model.name,
366 "cost_estimate": round(route_result.estimated_cost, 6),
367 },
368 )
370 result = agent.run(task)
371 elapsed_ms = (time.time() - t_start) * 1000
373 # 6. record model stats
374 self._router.record_request(
375 model_name=route_result.model.name,
376 success=result.success,
377 tokens_used=result.total_tokens,
378 cost_usd=result.total_cost_usd,
379 latency_ms=elapsed_ms,
380 )
382 # 7. audit: agent completion
383 if self._audit:
384 self._audit.log(
385 agent="production",
386 action="agent_end",
387 result="success" if result.success else "failure",
388 severity=AuditSeverity.ERROR if not result.success else AuditSeverity.INFO,
389 category=AuditActionCategory.AGENT_INVOKE,
390 session_id=self._session_id,
391 duration_ms=elapsed_ms,
392 error_message=result.error or "",
393 details={
394 "steps": result.total_steps,
395 "tokens": result.total_tokens,
396 "cost": round(result.total_cost_usd, 6),
397 "model": route_result.model.name,
398 },
399 )
401 # attach routing metadata to result
402 result.total_cost_usd = result.total_cost_usd
403 result.total_duration_ms = elapsed_ms
405 # 8. capture cache stats
406 if self._cache and result.success:
407 # track estimated cost savings from cache
408 self._cache._stats.total_cost_saved_usd += route_result.estimated_cost
410 return result
412 # ── properties ───────────────────────────────────────────────
414 @property
415 def last_route(self) -> RequestSpec | None:
416 """The last routing request spec."""
417 return self._last_route
419 @property
420 def last_model(self) -> ModelSpec | None:
421 """The model used in the last run."""
422 return self._last_model
424 @property
425 def cache_stats(self) -> Any | None:
426 """Cache statistics if cache is enabled, else None.
428 Returns a CacheStats dataclass with fields:
429 hits, misses, fuzzy_hits, exact_hits, evictions,
430 total_cost_saved_usd, total_entries, hit_rate, fuzzy_hit_rate.
431 """
432 return self._cache.stats if self._cache else None
434 @property
435 def cache_hit_rate(self) -> float:
436 """Cache hit rate (0.0-1.0). Returns 0.0 if cache disabled."""
437 if not self._cache:
438 return 0.0
439 s = self._cache.stats
440 total = s.hits + s.misses
441 return s.hits / total if total > 0 else 0.0
443 @property
444 def cache_savings(self) -> float:
445 """Estimated USD saved by cache hits."""
446 if not self._cache:
447 return 0.0
448 return self._cache.stats.total_cost_saved_usd
450 # ── streaming ────────────────────────────────────────────────
452 def run_stream(self, task: str) -> Generator[AgentStep, None, AgentResult]:
453 """Streaming version — yields steps as they complete."""
454 estimate = self._estimator.estimate(task)
455 route_spec = RequestSpec(
456 estimated_input_tokens=estimate.estimated_tokens,
457 estimated_output_tokens=estimate.estimated_tokens // 2,
458 complexity=estimate.complexity,
459 priority=estimate.priority,
460 task_id=f"task-{uuid.uuid4().hex[:8]}",
461 session_id=self._session_id,
462 )
463 self._last_route = route_spec
464 self._router.route(route_spec)
466 agent = ToolAgent(
467 provider=self._provider,
468 tool_executor=self._make_audited_executor(),
469 config=self._config.agent_config,
470 system_prompt=self._system_prompt,
471 )
473 yield from agent.run_stream(task)
475 # ── accessors ───────────────────────────────────────────────
477 @property
478 def router(self) -> ModelRouter:
479 return self._router
481 @property
482 def audit(self) -> AuditLogger | None:
483 return self._audit
485 @property
486 def session_id(self) -> str:
487 return self._session_id
489 def route_summary(self) -> dict:
490 """Return routing + audit summary for the current session."""
491 summary = {
492 "session_id": self._session_id,
493 "router": self._router.summary() if self._router else {},
494 }
495 if self._audit:
496 summary["audit"] = self._audit.stats_summary()
497 if self._last_model:
498 summary["last_model"] = self._last_model.name
499 summary["last_model_tier"] = self._last_model.tier.name
500 return summary
502 # ── internal ────────────────────────────────────────────────
504 def _make_audited_executor(self) -> ToolExecutor:
505 """Wrap tool executor to auto-audit every call."""
506 if not self._audit:
507 return self._executor
509 audited = ToolExecutor()
511 # copy original tools with audit wrapping
512 for schema in self._executor.get_schemas():
513 original_name = schema.function.name
515 # Capture the original execute method for this tool
516 def make_wrapper(name: str) -> Callable[..., str]:
517 def wrapper(**kwargs: Any) -> str:
518 t_start = time.time()
519 try:
520 result = self._executor.execute(
521 type(
522 "FakeCall",
523 (),
524 {
525 "name": name,
526 "parsed_arguments": kwargs,
527 },
528 )()
529 )
530 elapsed = (time.time() - t_start) * 1000
531 self._audit.log(
532 agent="production",
533 action=f"tool:{name}",
534 target=str(list(kwargs.keys())),
535 result="success",
536 severity=AuditSeverity.DEBUG,
537 category=AuditActionCategory.TOOL_CALL,
538 session_id=self._session_id,
539 duration_ms=elapsed,
540 details={"arguments": kwargs},
541 )
542 return result
543 except Exception as exc:
544 elapsed = (time.time() - t_start) * 1000
545 self._audit.log(
546 agent="production",
547 action=f"tool:{name}",
548 target=str(list(kwargs.keys())),
549 result="failure",
550 severity=AuditSeverity.ERROR,
551 category=AuditActionCategory.TOOL_CALL,
552 session_id=self._session_id,
553 duration_ms=elapsed,
554 error_message=str(exc),
555 )
556 raise
558 return wrapper
560 audited.register(schema, make_wrapper(original_name))
562 return audited