Upload aco/trace_schema.py

aco/trace_schema.py

@@ -0,0 +1,234 @@
+"""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,
+        }