"""Normalized trace schema for Agent Cost Optimizer.""" from dataclasses import dataclass, field from typing import Optional, List, Dict, Any from enum import Enum from datetime import datetime class TaskType(Enum): QUICK_ANSWER = "quick_answer" RESEARCH = "research" CODING = "coding" DOCUMENT_DRAFTING = "document_drafting" LEGAL_REGULATED = "legal_regulated" TOOL_HEAVY = "tool_heavy" RETRIEVAL_HEAVY = "retrieval_heavy" LONG_HORIZON = "long_horizon" UNKNOWN_AMBIGUOUS = "unknown_ambiguous" class Outcome(Enum): SUCCESS = "success" PARTIAL_SUCCESS = "partial_success" FAILURE = "failure" FALSE_DONE = "false_done" BLOCKED = "blocked" ESCALATED_HUMAN = "escalated_human" STOPPED_DOOM = "stopped_doom" class FailureTag(Enum): MODEL_TOO_WEAK = "model_too_weak" CONTEXT_TOO_SMALL = "context_too_small" TOOL_FAILED = "tool_failed" TOOL_UNNECESSARY = "tool_unnecessary" TOOL_MISSED = "tool_missed" VERIFIER_FALSE_PASS = "verifier_false_pass" VERIFIER_FALSE_REJECT = "verifier_false_reject" RETRY_LOOP = "retry_loop" CACHE_BREAK = "cache_break" HALLUCINATION = "hallucination" TIMEOUT = "timeout" COST_EXCEEDED = "cost_exceeded" UNSAFE_CHEAP_MODEL = "unsafe_cheap_model" MISSED_ESCALATION = "missed_escalation" @dataclass class ToolCall: tool_name: str tool_input: Dict[str, Any] tool_output: Optional[str] = None tool_cost: float = 0.0 tool_latency_ms: float = 0.0 cache_hit: bool = False repeated: bool = False ignored_result: bool = False failed: bool = False @dataclass class ModelCall: model_id: str provider: str input_tokens: int = 0 output_tokens: int = 0 reasoning_tokens: int = 0 cost_per_1k_input: float = 0.0 cost_per_1k_output: float = 0.0 cache_hit_input_tokens: int = 0 latency_ms: float = 0.0 @property def total_cost(self) -> float: input_cost = (self.input_tokens / 1000) * self.cost_per_1k_input output_cost = (self.output_tokens / 1000) * self.cost_per_1k_output cache_discount = (self.cache_hit_input_tokens / 1000) * self.cost_per_1k_input * 0.5 # 50% discount return input_cost + output_cost - cache_discount @dataclass class VerifierCall: verifier_model_id: str target_step_id: str passed: bool = False confidence: float = 0.0 cost: float = 0.0 latency_ms: float = 0.0 @dataclass class TraceStep: step_id: str timestamp: datetime task_type: TaskType model_call: ModelCall tool_calls: List[ToolCall] = field(default_factory=list) verifier_calls: List[VerifierCall] = field(default_factory=list) context_size_tokens: int = 0 context_sources: List[str] = field(default_factory=list) cache_boundary_reached: bool = False retry_count: int = 0 recovery_action: Optional[str] = None planned_next: Optional[str] = None user_correction: Optional[str] = None artifacts_created: List[str] = field(default_factory=list) step_outcome: Optional[Outcome] = None @property def step_cost(self) -> float: mc = self.model_call.total_cost if self.model_call else 0.0 tc = sum(t.tool_cost for t in self.tool_calls) vc = sum(v.cost for v in self.verifier_calls) return mc + tc + vc @property def step_latency_ms(self) -> float: ml = self.model_call.latency_ms if self.model_call else 0.0 tl = sum(t.tool_latency_ms for t in self.tool_calls) vl = sum(v.latency_ms for v in self.verifier_calls) return ml + tl + vl @dataclass class AgentTrace: trace_id: str user_request: str task_type: TaskType steps: List[TraceStep] = field(default_factory=list) final_outcome: Optional[Outcome] = None final_artifacts: List[str] = field(default_factory=list) failure_tags: List[FailureTag] = field(default_factory=list) user_satisfaction: Optional[float] = None # 0-1 total_cost_saved_vs_frontier: Optional[float] = None total_cost: Optional[float] = None optimal_cost: Optional[float] = None # oracle routing cost metadata: Dict[str, Any] = field(default_factory=dict) @property def total_cost_computed(self) -> float: return sum(s.step_cost for s in self.steps) @property def total_latency_ms(self) -> float: return sum(s.step_latency_ms for s in self.steps) @property def total_retries(self) -> int: return sum(s.retry_count for s in self.steps) @property def total_tool_calls(self) -> int: return sum(len(s.tool_calls) for s in self.steps) @property def total_verifier_calls(self) -> int: return sum(len(s.verifier_calls) for s in self.steps) @property def total_context_tokens(self) -> int: return sum(s.context_size_tokens for s in self.steps) @property def cache_hit_rate(self) -> float: model_calls = [s.model_call for s in self.steps if s.model_call] if not model_calls: return 0.0 total_input = sum(m.input_tokens for m in model_calls) total_cache_hit = sum(m.cache_hit_input_tokens for m in model_calls) return total_cache_hit / total_input if total_input > 0 else 0.0 def to_dict(self) -> Dict[str, Any]: return { "trace_id": self.trace_id, "user_request": self.user_request, "task_type": self.task_type.value, "steps": [ { "step_id": s.step_id, "timestamp": s.timestamp.isoformat(), "task_type": s.task_type.value, "model_call": { "model_id": s.model_call.model_id, "provider": s.model_call.provider, "input_tokens": s.model_call.input_tokens, "output_tokens": s.model_call.output_tokens, "reasoning_tokens": s.model_call.reasoning_tokens, "cost": s.model_call.total_cost, "latency_ms": s.model_call.latency_ms, "cache_hit_input_tokens": s.model_call.cache_hit_input_tokens, }, "tool_calls": [ { "tool_name": t.tool_name, "tool_cost": t.tool_cost, "tool_latency_ms": t.tool_latency_ms, "cache_hit": t.cache_hit, "repeated": t.repeated, "ignored_result": t.ignored_result, "failed": t.failed, } for t in s.tool_calls ], "verifier_calls": [ { "verifier_model_id": v.verifier_model_id, "passed": v.passed, "confidence": v.confidence, "cost": v.cost, } for v in s.verifier_calls ], "context_size_tokens": s.context_size_tokens, "retry_count": s.retry_count, "recovery_action": s.recovery_action, "step_outcome": s.step_outcome.value if s.step_outcome else None, "step_cost": s.step_cost, "step_latency_ms": s.step_latency_ms, } for s in self.steps ], "final_outcome": self.final_outcome.value if self.final_outcome else None, "failure_tags": [f.value for f in self.failure_tags], "total_cost": self.total_cost_computed, "total_latency_ms": self.total_latency_ms, "total_retries": self.total_retries, "total_tool_calls": self.total_tool_calls, "total_verifier_calls": self.total_verifier_calls, "total_context_tokens": self.total_context_tokens, "cache_hit_rate": self.cache_hit_rate, "user_satisfaction": self.user_satisfaction, "total_cost_saved_vs_frontier": self.total_cost_saved_vs_frontier, "optimal_cost": self.optimal_cost, }