Coverage for agentos/agent/production.py: 35%
171 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-05 20:52 +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 dataclasses import dataclass, field
36from typing import Any, Callable, Generator, Optional
38from agentos.agent.tool_agent import (
39 ToolAgent,
40 ToolExecutor,
41 AgentConfig,
42 AgentStep,
43 AgentResult,
44)
45from agentos.agent.model_router import (
46 ModelRouter,
47 ModelSpec,
48 RequestSpec,
49 TaskComplexity,
50 TaskPriority,
51)
52from agentos.security.audit_logger import (
53 AuditLogger,
54 AuditEvent,
55 AuditSeverity,
56 AuditActionCategory,
57)
58from agentos.llm.base import LLMProvider, Message, MessageRole
59from agentos.llm.smart_cache import SmartCache, CacheConfig
62__all__ = [
63 "ProductionAgent",
64 "ProductionConfig",
65 "ComplexityEstimator",
66 "ComplexityEstimate",
67]
70# ── 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 "analyze", "compare", "research", "evaluate", "assess",
91 "代码审查", "架构", "重构", "安全审计",
92 "code review", "architecture", "refactor", "audit",
93 ]
94 EXPERT_KEYWORDS = [
95 "深度", "全面", "完整", "生产级", "企业级", "从零",
96 "comprehensive", "production", "enterprise", "from scratch",
97 "论文", "学术", "法律",
98 "paper", "academic", "legal",
99 ]
100 TRIVIAL_KEYWORDS = [
101 "天气", "时间", "翻译", "计算", "换算", "几点", "日期",
102 "weather", "time", "translate", "calculate", "convert",
103 ]
104 URGENT_KEYWORDS = [
105 "快", "紧急", "马上", "立即",
106 "urgent", "asap", "immediately", "now",
107 ]
109 def estimate(self, task: str) -> ComplexityEstimate:
110 task_lower = task.lower()
112 # check for urgency
113 priority = TaskPriority.NORMAL
114 for kw in self.URGENT_KEYWORDS:
115 if kw in task_lower:
116 priority = TaskPriority.HIGH
117 break
119 # complexity
120 complexity = TaskComplexity.MODERATE
122 expert_hits = sum(1 for kw in self.EXPERT_KEYWORDS if kw in task_lower)
123 complex_hits = sum(1 for kw in self.COMPLEX_KEYWORDS if kw in task_lower)
124 trivial_hits = sum(1 for kw in self.TRIVIAL_KEYWORDS if kw in task_lower)
126 if expert_hits >= 2 or len(task) > 500:
127 complexity = TaskComplexity.EXPERT
128 elif expert_hits >= 1 or complex_hits >= 3:
129 complexity = TaskComplexity.COMPLEX
130 elif complex_hits >= 1 or len(task) > 200:
131 complexity = TaskComplexity.MODERATE
132 elif trivial_hits >= 1:
133 complexity = TaskComplexity.TRIVIAL
134 else:
135 complexity = TaskComplexity.SIMPLE
137 # estimate token count
138 # rough heuristic: ~1.5 chars per token for Chinese, ~4 for English
139 char_count = len(task)
140 if any('\u4e00' <= c <= '\u9fff' for c in task):
141 estimated_tokens = max(50, char_count // 1.5)
142 else:
143 estimated_tokens = max(50, char_count // 4)
145 estimated_tokens = int(min(estimated_tokens, 100_000))
147 return ComplexityEstimate(
148 complexity=complexity,
149 priority=priority,
150 estimated_tokens=estimated_tokens,
151 reason=f"task_len={char_count} chars, expert_hits={expert_hits}, "
152 f"complex_hits={complex_hits}, trivial_hits={trivial_hits}",
153 )
156# ── ProductionConfig ────────────────────────────────────────────────
158@dataclass
159class ProductionConfig:
160 agent_config: AgentConfig = field(default_factory=AgentConfig)
161 enable_audit: bool = True
162 enable_routing: bool = True
163 enable_cache: bool = True
164 audit_log_dir: str = ""
165 session_id: str = ""
166 budget_usd: float = 50.0
167 fallback_on_error: bool = True
170# ── ProductionAgent ─────────────────────────────────────────────────
172class ProductionAgent:
173 """Production-grade ToolAgent wrapper with routing and auditing.
175 Architecture:
176 User Task
177 │
178 ▼
179 ComplexityEstimator → classifies task
180 │
181 ▼
182 ModelRouter → selects best model for this complexity
183 │
184 ▼
185 ToolAgent.run() ───→ AuditLogger (every step & tool call)
186 │
187 ▼
188 AgentResult (with routing metadata + audit trail)
190 All tool calls and agent steps are automatically audited with
191 SHA256-chained immutable entries. Model selection is automatic
192 based on task analysis.
193 """
195 def __init__(
196 self,
197 provider: LLMProvider,
198 tool_executor: ToolExecutor,
199 *,
200 config: ProductionConfig | None = None,
201 router: ModelRouter | None = None,
202 cache: SmartCache | None = None,
203 system_prompt: str = "",
204 ):
205 self._config = config or ProductionConfig()
207 # cache — wrap provider transparently
208 if cache and self._config.enable_cache:
209 self._cache = cache
210 self._provider = cache.wrap(provider)
211 else:
212 self._cache = None
213 self._provider = provider
215 self._executor = tool_executor
216 self._system_prompt = system_prompt
218 # routing
219 self._router = router or ModelRouter.with_defaults(
220 daily_budget_usd=self._config.budget_usd,
221 )
222 self._estimator = ComplexityEstimator()
224 # auditing
225 self._session_id = self._config.session_id or f"sess-{uuid.uuid4().hex[:8]}"
226 self._audit = AuditLogger(
227 log_dir=self._config.audit_log_dir,
228 max_events=100_000,
229 max_age_days=90,
230 ) if self._config.enable_audit else None
232 # track routing decisions
233 self._last_route: Optional[RequestSpec] = None
234 self._last_model: Optional[ModelSpec] = None
236 # ── public API ──────────────────────────────────────────────
238 def run(self, task: str) -> AgentResult:
239 t_start = time.time()
241 # 1. classify
242 estimate = self._estimator.estimate(task)
244 # 2. route
245 route_spec = RequestSpec(
246 estimated_input_tokens=estimate.estimated_tokens,
247 estimated_output_tokens=estimate.estimated_tokens // 2,
248 complexity=estimate.complexity,
249 priority=estimate.priority,
250 task_id=f"task-{uuid.uuid4().hex[:8]}",
251 session_id=self._session_id,
252 )
253 self._last_route = route_spec
254 route_result = self._router.route(route_spec)
256 if not route_result.success:
257 return AgentResult(
258 success=False,
259 error=f"Model routing failed: {route_result.reason}",
260 )
262 self._last_model = route_result.model
264 # audit: route decision
265 if self._audit:
266 self._audit.log(
267 agent="production",
268 action="model_route",
269 target=route_result.model.name,
270 result="success",
271 severity=AuditSeverity.INFO,
272 category=AuditActionCategory.CONFIG_CHANGE,
273 session_id=self._session_id,
274 details={
275 "complexity": estimate.complexity.name,
276 "priority": estimate.priority.name,
277 "estimated_tokens": estimate.estimated_tokens,
278 "estimated_cost": round(route_result.estimated_cost, 6),
279 "fallback_chain": route_result.fallback_chain,
280 "budget_remaining": round(self._router.daily_budget_remaining, 4),
281 },
282 )
284 # 3. build agent config with model info
285 agent_config = AgentConfig(
286 max_steps=self._config.agent_config.max_steps,
287 temperature=self._config.agent_config.temperature,
288 max_tokens=self._config.agent_config.max_tokens,
289 verbose=self._config.agent_config.verbose,
290 stop_on_error=self._config.agent_config.stop_on_error,
291 max_retries=self._config.agent_config.max_retries,
292 retry_delay=self._config.agent_config.retry_delay,
293 )
295 # 4. create tool agent and wrap tool executor with audit
296 agent = ToolAgent(
297 provider=self._provider,
298 tool_executor=self._make_audited_executor(),
299 config=agent_config,
300 system_prompt=self._system_prompt,
301 )
303 # 5. run
304 if self._audit:
305 self._audit.log(
306 agent="production",
307 action="agent_start",
308 target=task[:100],
309 result="success",
310 severity=AuditSeverity.INFO,
311 category=AuditActionCategory.AGENT_INVOKE,
312 session_id=self._session_id,
313 details={
314 "complexity": estimate.complexity.name,
315 "model": route_result.model.name,
316 "cost_estimate": round(route_result.estimated_cost, 6),
317 },
318 )
320 result = agent.run(task)
321 elapsed_ms = (time.time() - t_start) * 1000
323 # 6. record model stats
324 self._router.record_request(
325 model_name=route_result.model.name,
326 success=result.success,
327 tokens_used=result.total_tokens,
328 cost_usd=result.total_cost_usd,
329 latency_ms=elapsed_ms,
330 )
332 # 7. audit: agent completion
333 if self._audit:
334 self._audit.log(
335 agent="production",
336 action="agent_end",
337 result="success" if result.success else "failure",
338 severity=AuditSeverity.ERROR if not result.success else AuditSeverity.INFO,
339 category=AuditActionCategory.AGENT_INVOKE,
340 session_id=self._session_id,
341 duration_ms=elapsed_ms,
342 error_message=result.error or "",
343 details={
344 "steps": result.total_steps,
345 "tokens": result.total_tokens,
346 "cost": round(result.total_cost_usd, 6),
347 "model": route_result.model.name,
348 },
349 )
351 # attach routing metadata to result
352 result.total_cost_usd = result.total_cost_usd
353 result.total_duration_ms = elapsed_ms
355 # 8. capture cache stats
356 if self._cache and result.success:
357 # track estimated cost savings from cache
358 self._cache._stats.total_cost_saved_usd += route_result.estimated_cost
360 return result
362 # ── properties ───────────────────────────────────────────────
364 @property
365 def last_route(self) -> Optional[RequestSpec]:
366 """The last routing request spec."""
367 return self._last_route
369 @property
370 def last_model(self) -> Optional[ModelSpec]:
371 """The model used in the last run."""
372 return self._last_model
374 @property
375 def cache_stats(self) -> Optional[Any]:
376 """Cache statistics if cache is enabled, else None.
378 Returns a CacheStats dataclass with fields:
379 hits, misses, fuzzy_hits, exact_hits, evictions,
380 total_cost_saved_usd, total_entries, hit_rate, fuzzy_hit_rate.
381 """
382 return self._cache.stats if self._cache else None
384 @property
385 def cache_hit_rate(self) -> float:
386 """Cache hit rate (0.0-1.0). Returns 0.0 if cache disabled."""
387 if not self._cache:
388 return 0.0
389 s = self._cache.stats
390 total = s.hits + s.misses
391 return s.hits / total if total > 0 else 0.0
393 @property
394 def cache_savings(self) -> float:
395 """Estimated USD saved by cache hits."""
396 if not self._cache:
397 return 0.0
398 return self._cache.stats.total_cost_saved_usd
400 # ── streaming ────────────────────────────────────────────────
402 def run_stream(self, task: str) -> Generator[AgentStep, None, AgentResult]:
403 """Streaming version — yields steps as they complete."""
404 estimate = self._estimator.estimate(task)
405 route_spec = RequestSpec(
406 estimated_input_tokens=estimate.estimated_tokens,
407 estimated_output_tokens=estimate.estimated_tokens // 2,
408 complexity=estimate.complexity,
409 priority=estimate.priority,
410 task_id=f"task-{uuid.uuid4().hex[:8]}",
411 session_id=self._session_id,
412 )
413 self._last_route = route_spec
414 route_result = self._router.route(route_spec)
416 agent = ToolAgent(
417 provider=self._provider,
418 tool_executor=self._make_audited_executor(),
419 config=self._config.agent_config,
420 system_prompt=self._system_prompt,
421 )
423 yield from agent.run_stream(task)
425 # ── accessors ───────────────────────────────────────────────
427 @property
428 def router(self) -> ModelRouter:
429 return self._router
431 @property
432 def audit(self) -> Optional[AuditLogger]:
433 return self._audit
435 @property
436 def session_id(self) -> str:
437 return self._session_id
439 @property
440 def last_route(self) -> Optional[RequestSpec]:
441 return self._last_route
443 @property
444 def last_model(self) -> Optional[ModelSpec]:
445 return self._last_model
447 def route_summary(self) -> dict:
448 """Return routing + audit summary for the current session."""
449 summary = {
450 "session_id": self._session_id,
451 "router": self._router.summary() if self._router else {},
452 }
453 if self._audit:
454 summary["audit"] = self._audit.stats_summary()
455 if self._last_model:
456 summary["last_model"] = self._last_model.name
457 summary["last_model_tier"] = self._last_model.tier.name
458 return summary
460 # ── internal ────────────────────────────────────────────────
462 def _make_audited_executor(self) -> ToolExecutor:
463 """Wrap tool executor to auto-audit every call."""
464 if not self._audit:
465 return self._executor
467 audited = ToolExecutor()
469 # copy original tools with audit wrapping
470 for schema in self._executor.get_schemas():
471 original_name = schema.function.name
473 # Capture the original execute method for this tool
474 def make_wrapper(name: str) -> Callable[..., str]:
475 def wrapper(**kwargs: Any) -> str:
476 t_start = time.time()
477 try:
478 result = self._executor.execute(
479 type("FakeCall", (), {
480 "name": name,
481 "parsed_arguments": kwargs,
482 })()
483 )
484 elapsed = (time.time() - t_start) * 1000
485 self._audit.log(
486 agent="production",
487 action=f"tool:{name}",
488 target=str(list(kwargs.keys())),
489 result="success",
490 severity=AuditSeverity.DEBUG,
491 category=AuditActionCategory.TOOL_CALL,
492 session_id=self._session_id,
493 duration_ms=elapsed,
494 details={"arguments": kwargs},
495 )
496 return result
497 except Exception as exc:
498 elapsed = (time.time() - t_start) * 1000
499 self._audit.log(
500 agent="production",
501 action=f"tool:{name}",
502 target=str(list(kwargs.keys())),
503 result="failure",
504 severity=AuditSeverity.ERROR,
505 category=AuditActionCategory.TOOL_CALL,
506 session_id=self._session_id,
507 duration_ms=elapsed,
508 error_message=str(exc),
509 )
510 raise
511 return wrapper
513 audited.register(schema, make_wrapper(original_name))
515 return audited