Coverage for agentos/agent/production.py: 35%
171 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 10:59 +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 AuditSeverity,
55 AuditActionCategory,
56)
57from agentos.llm.base import LLMProvider
58from agentos.llm.smart_cache import SmartCache
61__all__ = [
62 "ProductionAgent",
63 "ProductionConfig",
64 "ComplexityEstimator",
65 "ComplexityEstimate",
66]
69# ── Complexity Estimator ────────────────────────────────────────────
71@dataclass
72class ComplexityEstimate:
73 complexity: TaskComplexity
74 priority: TaskPriority
75 estimated_tokens: int
76 reason: str
79class ComplexityEstimator:
80 """Estimates task complexity from the user's request text.
82 Uses keyword heuristics to classify tasks from TRIVIAL to EXPERT.
83 Production use should integrate with a lightweight classifier model.
84 """
86 # keywords that indicate higher complexity
87 COMPLEX_KEYWORDS = [
88 "分析", "对比", "比较", "总结", "调研", "研究", "评估",
89 "analyze", "compare", "research", "evaluate", "assess",
90 "代码审查", "架构", "重构", "安全审计",
91 "code review", "architecture", "refactor", "audit",
92 ]
93 EXPERT_KEYWORDS = [
94 "深度", "全面", "完整", "生产级", "企业级", "从零",
95 "comprehensive", "production", "enterprise", "from scratch",
96 "论文", "学术", "法律",
97 "paper", "academic", "legal",
98 ]
99 TRIVIAL_KEYWORDS = [
100 "天气", "时间", "翻译", "计算", "换算", "几点", "日期",
101 "weather", "time", "translate", "calculate", "convert",
102 ]
103 URGENT_KEYWORDS = [
104 "快", "紧急", "马上", "立即",
105 "urgent", "asap", "immediately", "now",
106 ]
108 def estimate(self, task: str) -> ComplexityEstimate:
109 task_lower = task.lower()
111 # check for urgency
112 priority = TaskPriority.NORMAL
113 for kw in self.URGENT_KEYWORDS:
114 if kw in task_lower:
115 priority = TaskPriority.HIGH
116 break
118 # complexity
119 complexity = TaskComplexity.MODERATE
121 expert_hits = sum(1 for kw in self.EXPERT_KEYWORDS if kw in task_lower)
122 complex_hits = sum(1 for kw in self.COMPLEX_KEYWORDS if kw in task_lower)
123 trivial_hits = sum(1 for kw in self.TRIVIAL_KEYWORDS if kw in task_lower)
125 if expert_hits >= 2 or len(task) > 500:
126 complexity = TaskComplexity.EXPERT
127 elif expert_hits >= 1 or complex_hits >= 3:
128 complexity = TaskComplexity.COMPLEX
129 elif complex_hits >= 1 or len(task) > 200:
130 complexity = TaskComplexity.MODERATE
131 elif trivial_hits >= 1:
132 complexity = TaskComplexity.TRIVIAL
133 else:
134 complexity = TaskComplexity.SIMPLE
136 # estimate token count
137 # rough heuristic: ~1.5 chars per token for Chinese, ~4 for English
138 char_count = len(task)
139 if any('\u4e00' <= c <= '\u9fff' for c in task):
140 estimated_tokens = max(50, char_count // 1.5)
141 else:
142 estimated_tokens = max(50, char_count // 4)
144 estimated_tokens = int(min(estimated_tokens, 100_000))
146 return ComplexityEstimate(
147 complexity=complexity,
148 priority=priority,
149 estimated_tokens=estimated_tokens,
150 reason=f"task_len={char_count} chars, expert_hits={expert_hits}, "
151 f"complex_hits={complex_hits}, trivial_hits={trivial_hits}",
152 )
155# ── ProductionConfig ────────────────────────────────────────────────
157@dataclass
158class ProductionConfig:
159 agent_config: AgentConfig = field(default_factory=AgentConfig)
160 enable_audit: bool = True
161 enable_routing: bool = True
162 enable_cache: bool = True
163 audit_log_dir: str = ""
164 session_id: str = ""
165 budget_usd: float = 50.0
166 fallback_on_error: bool = True
169# ── ProductionAgent ─────────────────────────────────────────────────
171class ProductionAgent:
172 """Production-grade ToolAgent wrapper with routing and auditing.
174 Architecture:
175 User Task
176 │
177 ▼
178 ComplexityEstimator → classifies task
179 │
180 ▼
181 ModelRouter → selects best model for this complexity
182 │
183 ▼
184 ToolAgent.run() ───→ AuditLogger (every step & tool call)
185 │
186 ▼
187 AgentResult (with routing metadata + audit trail)
189 All tool calls and agent steps are automatically audited with
190 SHA256-chained immutable entries. Model selection is automatic
191 based on task analysis.
192 """
194 def __init__(
195 self,
196 provider: LLMProvider,
197 tool_executor: ToolExecutor,
198 *,
199 config: ProductionConfig | None = None,
200 router: ModelRouter | None = None,
201 cache: SmartCache | None = None,
202 system_prompt: str = "",
203 ):
204 self._config = config or ProductionConfig()
206 # cache — wrap provider transparently
207 if cache and self._config.enable_cache:
208 self._cache = cache
209 self._provider = cache.wrap(provider)
210 else:
211 self._cache = None
212 self._provider = provider
214 self._executor = tool_executor
215 self._system_prompt = system_prompt
217 # routing
218 self._router = router or ModelRouter.with_defaults(
219 daily_budget_usd=self._config.budget_usd,
220 )
221 self._estimator = ComplexityEstimator()
223 # auditing
224 self._session_id = self._config.session_id or f"sess-{uuid.uuid4().hex[:8]}"
225 self._audit = AuditLogger(
226 log_dir=self._config.audit_log_dir,
227 max_events=100_000,
228 max_age_days=90,
229 ) if self._config.enable_audit else None
231 # track routing decisions
232 self._last_route: Optional[RequestSpec] = None
233 self._last_model: Optional[ModelSpec] = None
235 # ── public API ──────────────────────────────────────────────
237 def run(self, task: str) -> AgentResult:
238 t_start = time.time()
240 # 1. classify
241 estimate = self._estimator.estimate(task)
243 # 2. route
244 route_spec = RequestSpec(
245 estimated_input_tokens=estimate.estimated_tokens,
246 estimated_output_tokens=estimate.estimated_tokens // 2,
247 complexity=estimate.complexity,
248 priority=estimate.priority,
249 task_id=f"task-{uuid.uuid4().hex[:8]}",
250 session_id=self._session_id,
251 )
252 self._last_route = route_spec
253 route_result = self._router.route(route_spec)
255 if not route_result.success:
256 return AgentResult(
257 success=False,
258 error=f"Model routing failed: {route_result.reason}",
259 )
261 self._last_model = route_result.model
263 # audit: route decision
264 if self._audit:
265 self._audit.log(
266 agent="production",
267 action="model_route",
268 target=route_result.model.name,
269 result="success",
270 severity=AuditSeverity.INFO,
271 category=AuditActionCategory.CONFIG_CHANGE,
272 session_id=self._session_id,
273 details={
274 "complexity": estimate.complexity.name,
275 "priority": estimate.priority.name,
276 "estimated_tokens": estimate.estimated_tokens,
277 "estimated_cost": round(route_result.estimated_cost, 6),
278 "fallback_chain": route_result.fallback_chain,
279 "budget_remaining": round(self._router.daily_budget_remaining, 4),
280 },
281 )
283 # 3. build agent config with model info
284 agent_config = AgentConfig(
285 max_steps=self._config.agent_config.max_steps,
286 temperature=self._config.agent_config.temperature,
287 max_tokens=self._config.agent_config.max_tokens,
288 verbose=self._config.agent_config.verbose,
289 stop_on_error=self._config.agent_config.stop_on_error,
290 max_retries=self._config.agent_config.max_retries,
291 retry_delay=self._config.agent_config.retry_delay,
292 )
294 # 4. create tool agent and wrap tool executor with audit
295 agent = ToolAgent(
296 provider=self._provider,
297 tool_executor=self._make_audited_executor(),
298 config=agent_config,
299 system_prompt=self._system_prompt,
300 )
302 # 5. run
303 if self._audit:
304 self._audit.log(
305 agent="production",
306 action="agent_start",
307 target=task[:100],
308 result="success",
309 severity=AuditSeverity.INFO,
310 category=AuditActionCategory.AGENT_INVOKE,
311 session_id=self._session_id,
312 details={
313 "complexity": estimate.complexity.name,
314 "model": route_result.model.name,
315 "cost_estimate": round(route_result.estimated_cost, 6),
316 },
317 )
319 result = agent.run(task)
320 elapsed_ms = (time.time() - t_start) * 1000
322 # 6. record model stats
323 self._router.record_request(
324 model_name=route_result.model.name,
325 success=result.success,
326 tokens_used=result.total_tokens,
327 cost_usd=result.total_cost_usd,
328 latency_ms=elapsed_ms,
329 )
331 # 7. audit: agent completion
332 if self._audit:
333 self._audit.log(
334 agent="production",
335 action="agent_end",
336 result="success" if result.success else "failure",
337 severity=AuditSeverity.ERROR if not result.success else AuditSeverity.INFO,
338 category=AuditActionCategory.AGENT_INVOKE,
339 session_id=self._session_id,
340 duration_ms=elapsed_ms,
341 error_message=result.error or "",
342 details={
343 "steps": result.total_steps,
344 "tokens": result.total_tokens,
345 "cost": round(result.total_cost_usd, 6),
346 "model": route_result.model.name,
347 },
348 )
350 # attach routing metadata to result
351 result.total_cost_usd = result.total_cost_usd
352 result.total_duration_ms = elapsed_ms
354 # 8. capture cache stats
355 if self._cache and result.success:
356 # track estimated cost savings from cache
357 self._cache._stats.total_cost_saved_usd += route_result.estimated_cost
359 return result
361 # ── properties ───────────────────────────────────────────────
363 @property
364 def last_route(self) -> Optional[RequestSpec]:
365 """The last routing request spec."""
366 return self._last_route
368 @property
369 def last_model(self) -> Optional[ModelSpec]:
370 """The model used in the last run."""
371 return self._last_model
373 @property
374 def cache_stats(self) -> Optional[Any]:
375 """Cache statistics if cache is enabled, else None.
377 Returns a CacheStats dataclass with fields:
378 hits, misses, fuzzy_hits, exact_hits, evictions,
379 total_cost_saved_usd, total_entries, hit_rate, fuzzy_hit_rate.
380 """
381 return self._cache.stats if self._cache else None
383 @property
384 def cache_hit_rate(self) -> float:
385 """Cache hit rate (0.0-1.0). Returns 0.0 if cache disabled."""
386 if not self._cache:
387 return 0.0
388 s = self._cache.stats
389 total = s.hits + s.misses
390 return s.hits / total if total > 0 else 0.0
392 @property
393 def cache_savings(self) -> float:
394 """Estimated USD saved by cache hits."""
395 if not self._cache:
396 return 0.0
397 return self._cache.stats.total_cost_saved_usd
399 # ── streaming ────────────────────────────────────────────────
401 def run_stream(self, task: str) -> Generator[AgentStep, None, AgentResult]:
402 """Streaming version — yields steps as they complete."""
403 estimate = self._estimator.estimate(task)
404 route_spec = RequestSpec(
405 estimated_input_tokens=estimate.estimated_tokens,
406 estimated_output_tokens=estimate.estimated_tokens // 2,
407 complexity=estimate.complexity,
408 priority=estimate.priority,
409 task_id=f"task-{uuid.uuid4().hex[:8]}",
410 session_id=self._session_id,
411 )
412 self._last_route = route_spec
413 route_result = self._router.route(route_spec)
415 agent = ToolAgent(
416 provider=self._provider,
417 tool_executor=self._make_audited_executor(),
418 config=self._config.agent_config,
419 system_prompt=self._system_prompt,
420 )
422 yield from agent.run_stream(task)
424 # ── accessors ───────────────────────────────────────────────
426 @property
427 def router(self) -> ModelRouter:
428 return self._router
430 @property
431 def audit(self) -> Optional[AuditLogger]:
432 return self._audit
434 @property
435 def session_id(self) -> str:
436 return self._session_id
438 @property
439 def last_route(self) -> Optional[RequestSpec]:
440 return self._last_route
442 @property
443 def last_model(self) -> Optional[ModelSpec]:
444 return self._last_model
446 def route_summary(self) -> dict:
447 """Return routing + audit summary for the current session."""
448 summary = {
449 "session_id": self._session_id,
450 "router": self._router.summary() if self._router else {},
451 }
452 if self._audit:
453 summary["audit"] = self._audit.stats_summary()
454 if self._last_model:
455 summary["last_model"] = self._last_model.name
456 summary["last_model_tier"] = self._last_model.tier.name
457 return summary
459 # ── internal ────────────────────────────────────────────────
461 def _make_audited_executor(self) -> ToolExecutor:
462 """Wrap tool executor to auto-audit every call."""
463 if not self._audit:
464 return self._executor
466 audited = ToolExecutor()
468 # copy original tools with audit wrapping
469 for schema in self._executor.get_schemas():
470 original_name = schema.function.name
472 # Capture the original execute method for this tool
473 def make_wrapper(name: str) -> Callable[..., str]:
474 def wrapper(**kwargs: Any) -> str:
475 t_start = time.time()
476 try:
477 result = self._executor.execute(
478 type("FakeCall", (), {
479 "name": name,
480 "parsed_arguments": kwargs,
481 })()
482 )
483 elapsed = (time.time() - t_start) * 1000
484 self._audit.log(
485 agent="production",
486 action=f"tool:{name}",
487 target=str(list(kwargs.keys())),
488 result="success",
489 severity=AuditSeverity.DEBUG,
490 category=AuditActionCategory.TOOL_CALL,
491 session_id=self._session_id,
492 duration_ms=elapsed,
493 details={"arguments": kwargs},
494 )
495 return result
496 except Exception as exc:
497 elapsed = (time.time() - t_start) * 1000
498 self._audit.log(
499 agent="production",
500 action=f"tool:{name}",
501 target=str(list(kwargs.keys())),
502 result="failure",
503 severity=AuditSeverity.ERROR,
504 category=AuditActionCategory.TOOL_CALL,
505 session_id=self._session_id,
506 duration_ms=elapsed,
507 error_message=str(exc),
508 )
509 raise
510 return wrapper
512 audited.register(schema, make_wrapper(original_name))
514 return audited