feat(llm): cost tracker with per-tier pricing and breakdown

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-09 19:57:51 +02:00
parent a9261452e0
commit a6f32dd4d8
2 changed files with 105 additions and 0 deletions
@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from typing import Any
+
+from ..genome.hypothesis import ModelTier
+
+PRICE_PER_M_TOKENS: dict[ModelTier, dict[str, float]] = {
+    ModelTier.C: {"input": 0.40, "output": 0.40},
+    ModelTier.B: {"input": 3.00, "output": 15.00},
+}
+
+
+def estimate_cost(input_tokens: int, output_tokens: int, tier: ModelTier) -> float:
+    p = PRICE_PER_M_TOKENS[tier]
+    return (input_tokens / 1_000_000) * p["input"] + (output_tokens / 1_000_000) * p["output"]
+
+
+@dataclass
+class CostRecord:
+    ts: datetime
+    run_id: str
+    agent_id: str
+    tier: ModelTier
+    input_tokens: int
+    output_tokens: int
+    cost_usd: float
+
+
+@dataclass
+class CostTracker:
+    records: list[CostRecord] = field(default_factory=list)
+
+    def record(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        tier: ModelTier,
+        run_id: str,
+        agent_id: str,
+    ) -> CostRecord:
+        cost = estimate_cost(input_tokens, output_tokens, tier)
+        rec = CostRecord(
+            ts=datetime.now(UTC),
+            run_id=run_id,
+            agent_id=agent_id,
+            tier=tier,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cost_usd=cost,
+        )
+        self.records.append(rec)
+        return rec
+
+    def summary(self) -> dict[str, Any]:
+        by_tier: dict[str, dict[str, float]] = defaultdict(
+            lambda: {"calls": 0, "input_tokens": 0, "output_tokens": 0, "cost_usd": 0.0}
+        )
+        for r in self.records:
+            t = r.tier.value
+            by_tier[t]["calls"] += 1
+            by_tier[t]["input_tokens"] += r.input_tokens
+            by_tier[t]["output_tokens"] += r.output_tokens
+            by_tier[t]["cost_usd"] += r.cost_usd
+        return {
+            "calls": len(self.records),
+            "input_tokens": sum(r.input_tokens for r in self.records),
+            "output_tokens": sum(r.output_tokens for r in self.records),
+            "cost_usd": sum(r.cost_usd for r in self.records),
+            "by_tier": dict(by_tier),
+        }
@@ -0,0 +1,32 @@
+from multi_swarm.genome.hypothesis import ModelTier
+from multi_swarm.llm.cost_tracker import CostTracker, estimate_cost
+
+
+def test_estimate_cost_tier_c():
+    cost = estimate_cost(input_tokens=1_000_000, output_tokens=1_000_000, tier=ModelTier.C)
+    assert cost == 0.40 + 0.40
+
+
+def test_estimate_cost_tier_b():
+    cost = estimate_cost(input_tokens=1_000_000, output_tokens=1_000_000, tier=ModelTier.B)
+    assert cost == 3.00 + 15.00
+
+
+def test_tracker_accumulates():
+    t = CostTracker()
+    t.record(input_tokens=10_000, output_tokens=20_000, tier=ModelTier.C, run_id="r", agent_id="a")
+    t.record(input_tokens=5_000, output_tokens=15_000, tier=ModelTier.C, run_id="r", agent_id="b")
+    summary = t.summary()
+    assert summary["calls"] == 2
+    assert summary["input_tokens"] == 15_000
+    assert summary["output_tokens"] == 35_000
+    assert summary["cost_usd"] > 0
+
+
+def test_tracker_per_tier_breakdown():
+    t = CostTracker()
+    t.record(input_tokens=10_000, output_tokens=10_000, tier=ModelTier.C, run_id="r", agent_id="a")
+    t.record(input_tokens=10_000, output_tokens=10_000, tier=ModelTier.B, run_id="r", agent_id="b")
+    summary = t.summary()
+    assert "C" in summary["by_tier"]
+    assert "B" in summary["by_tier"]