Общая стоимость = Σ (input_tokens × input_price + output_tokens × output_price)
                  для каждого вызова API в задаче

import anthropic
import time
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class CostRecord:
    step: int
    model: str
    input_tokens: int
    output_tokens: int
    cache_read_tokens: int = 0
    cache_creation_tokens: int = 0
    input_cost: float = 0.0
    output_cost: float = 0.0
    total_cost: float = 0.0
    duration_ms: float = 0.0

# Pricing per million tokens (as of early 2026)
MODEL_PRICING = {
    "claude-haiku": {"input": 0.25, "output": 1.25, "cache_read": 0.025, "cache_write": 0.30},
    "claude-sonnet": {"input": 3.00, "output": 15.00, "cache_read": 0.30, "cache_write": 3.75},
    "claude-opus": {"input": 15.00, "output": 75.00, "cache_read": 1.50, "cache_write": 18.75},
}

@dataclass
class TaskCostTracker:
    task_id: str
    budget_limit: Optional[float] = None
    records: list = field(default_factory=list)
    total_cost: float = 0.0

    def record_call(self, step: int, model: str, usage) -> CostRecord:
        pricing = MODEL_PRICING.get(model, MODEL_PRICING["claude-sonnet"])

        input_cost = (usage.input_tokens / 1_000_000) * pricing["input"]
        output_cost = (usage.output_tokens / 1_000_000) * pricing["output"]
        cache_read_cost = (getattr(usage, 'cache_read_input_tokens', 0) / 1_000_000) * pricing["cache_read"]
        cache_write_cost = (getattr(usage, 'cache_creation_input_tokens', 0) / 1_000_000) * pricing["cache_write"]

        total = input_cost + output_cost + cache_read_cost + cache_write_cost

        record = CostRecord(
            step=step,
            model=model,
            input_tokens=usage.input_tokens,
            output_tokens=usage.output_tokens,
            cache_read_tokens=getattr(usage, 'cache_read_input_tokens', 0),
            cache_creation_tokens=getattr(usage, 'cache_creation_input_tokens', 0),
            input_cost=input_cost + cache_read_cost + cache_write_cost,
            output_cost=output_cost,
            total_cost=total,
        )

        self.records.append(record)
        self.total_cost += total

        print(f"  Step {step} [{model}]: {usage.input_tokens} in / {usage.output_tokens} out = ${total:.4f} (cumulative: ${self.total_cost:.4f})")

        if self.budget_limit and self.total_cost > self.budget_limit:
            raise BudgetExceededError(
                f"Task {self.task_id} exceeded budget: ${self.total_cost:.4f} > ${self.budget_limit:.4f}"
            )

        return record

    def summary(self) -> dict:
        return {
            "task_id": self.task_id,
            "total_steps": len(self.records),
            "total_input_tokens": sum(r.input_tokens for r in self.records),
            "total_output_tokens": sum(r.output_tokens for r in self.records),
            "total_cost": self.total_cost,
            "cost_by_model": self._cost_by_model(),
        }

    def _cost_by_model(self) -> dict:
        by_model = {}
        for r in self.records:
            if r.model not in by_model:
                by_model[r.model] = {"calls": 0, "cost": 0.0}
            by_model[r.model]["calls"] += 1
            by_model[r.model]["cost"] += r.total_cost
        return by_model


class BudgetExceededError(Exception):
    pass