Cost Tracking

# Cost tracking for LLM API usage
from dataclasses import dataclass, field
from typing import Dict

# Pricing per 1000 tokens (approximate, check current rates)
MODEL_PRICING = {
    "gpt-4o":           {"input": 0.005, "output": 0.015},
    "gpt-4o-mini":      {"input": 0.00015, "output": 0.0006},
    "claude-3-5-sonnet":{"input": 0.003, "output": 0.015},
    "claude-3-haiku":   {"input": 0.00025, "output": 0.00125},
}

@dataclass
class CostTracker:
    agent_id: str
    task_id: str
    model: str = "gpt-4o-mini"
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    step_costs: list = field(default_factory=list)

    def record_llm_call(self, step: int, input_tokens: int, output_tokens: int):
        pricing = MODEL_PRICING.get(self.model, {"input": 0, "output": 0})
        step_cost = (input_tokens / 1000 * pricing["input"] +
                     output_tokens / 1000 * pricing["output"])

        self.total_input_tokens  += input_tokens
        self.total_output_tokens += output_tokens
        self.step_costs.append({"step": step, "cost_usd": round(step_cost, 6)})

    @property
    def total_cost_usd(self) -> float:
        return sum(s["cost_usd"] for s in self.step_costs)

    def report(self) -> dict:
        return {
            "agent_id":      self.agent_id,
            "task_id":       self.task_id,
            "model":         self.model,
            "input_tokens":  self.total_input_tokens,
            "output_tokens": self.total_output_tokens,
            "total_tokens":  self.total_input_tokens + self.total_output_tokens,
            "total_cost_usd": round(self.total_cost_usd, 6),
            "cost_per_step": round(self.total_cost_usd / max(len(self.step_costs), 1), 6),
            "steps":         len(self.step_costs),
        }

# Simulate a 5-step agent task
tracker = CostTracker("agent-support-1", "TASK-001", model="gpt-4o-mini")

step_usage = [
    (1, 1500, 200),   # Reasoning step
    (2, 800,  100),   # After get_order
    (3, 1200, 150),   # After eligibility check
    (4, 900,  80),    # After issue_refund
    (5, 600,  200),   # Final FINISH step
]

for step, inp, out in step_usage:
    tracker.record_llm_call(step, inp, out)

report = tracker.report()
for k, v in report.items():
    print(f"  {k:20}: {v}")