core/ledger.py

from dataclasses import dataclass, field
from typing import Dict, Any, Optional, List

@dataclass
class Ledger:
    spend_usd: float = 0.0
    tokens_prompt: int = 0
    tokens_completion: int = 0
    tokens_cached_prompt: int = 0
    tokens_reasoning: int = 0

    latency_ms_total: int = 0
    llm_calls: int = 0
    tool_calls: int = 0

    spend_by_agent: Dict[str, float] = field(default_factory=dict)
    spend_by_model: Dict[str, float] = field(default_factory=dict)

    alerts: List[str] = field(default_factory=list)

    def add_cost(self, usd: float, agent_id: Optional[str], model_key: Optional[str]):
        self.spend_usd += float(usd)
        if agent_id:
            self.spend_by_agent[agent_id] = self.spend_by_agent.get(agent_id, 0.0) + float(usd)
        if model_key:
            self.spend_by_model[model_key] = self.spend_by_model.get(model_key, 0.0) + float(usd)

    def add_tokens(self, tokens: Dict[str, Any]):
        self.tokens_prompt += int(tokens.get("prompt_tokens", 0))
        self.tokens_completion += int(tokens.get("completion_tokens", 0))
        self.tokens_cached_prompt += int(tokens.get("cached_prompt_tokens", 0))
        self.tokens_reasoning += int(tokens.get("reasoning_tokens", 0))

    def add_latency(self, latency_ms: int):
        self.latency_ms_total += int(latency_ms)

    def mark_llm_call(self):
        self.llm_calls += 1

    def mark_tool_call(self):
        self.tool_calls += 1

    def check_budget(self, soft: float, hard: float):
        if self.spend_usd >= hard:
            msg = f"BUDGET_HARD_EXCEEDED: spend=${self.spend_usd:.4f} >= ${hard:.2f}"
            if not self.alerts or self.alerts[-1] != msg:
                self.alerts.append(msg)
        elif self.spend_usd >= soft:
            msg = f"BUDGET_SOFT_EXCEEDED: spend=${self.spend_usd:.4f} >= ${soft:.2f}"
            if not self.alerts or self.alerts[-1] != msg:
                self.alerts.append(msg)