From d7ee22e5579af8e7fe331b2b4a038731c5fec557 Mon Sep 17 00:00:00 2001
From: Vadym Samoilenko <vadymsamoilenko@oliver.agency>
Date: Fri, 24 Apr 2026 18:57:25 +0100
Subject: [PATCH] Fix backfill pricing: read from model_pricing collection +
 --delete-existing-estimates flag

---
 backend/scripts/backfill_usage.py | 56 ++++++++++++++++++++++---------
 1 file changed, 41 insertions(+), 15 deletions(-)

diff --git a/backend/scripts/backfill_usage.py b/backend/scripts/backfill_usage.py
index 2c5b11c0..f1e13bc2 100644
--- a/backend/scripts/backfill_usage.py
+++ b/backend/scripts/backfill_usage.py
@@ -45,22 +45,39 @@ def _estimate_tokens(text: str, model: str) -> dict:
     return {"prompt": n, "completion": 0}
 
 
-def _estimate_cost(prompt_tokens: int, completion_tokens: int, model: str) -> float:
-    """Rough cost estimate in USD."""
-    rate_per_m = {
-        "gemini": (0.35, 1.05),
-        "gpt-4":  (30.00, 60.00),
-        "gpt-3":  (0.50, 1.50),
-    }
-    key = "gemini"
-    if model:
-        m = model.lower()
-        if "gpt-4" in m or "gpt-5" in m:
-            key = "gpt-4"
-        elif "gpt-3" in m:
-            key = "gpt-3"
+_pricing_cache: dict = {}
 
-    input_rate, output_rate = rate_per_m[key]
+def _load_pricing(db) -> None:
+    """Load current pricing from model_pricing collection into cache."""
+    for row in db.model_pricing.find({"effective_until": None}):
+        model = row.get("model", "")
+        tiers = row.get("tiers") or []
+        if tiers:
+            t = tiers[0]
+            _pricing_cache[model] = (
+                t.get("input_per_mtok", 2.0),
+                t.get("output_per_mtok", 12.0),
+            )
+
+
+def _estimate_cost(prompt_tokens: int, completion_tokens: int, model: str) -> float:
+    """Cost estimate in USD using model_pricing collection rates."""
+    # Try exact match, then prefix match
+    rates = _pricing_cache.get(model)
+    if not rates:
+        for key, val in _pricing_cache.items():
+            if model and key and (key in model or model in key):
+                rates = val
+                break
+    # Final fallback matching seed_model_pricing.py values
+    if not rates:
+        m = (model or "").lower()
+        if "gpt-5" in m or "gpt-4" in m:
+            rates = (2.50, 15.00)   # gpt-5.4 pricing from seed
+        else:
+            rates = (2.00, 12.00)   # gemini-3.1-pro-preview pricing from seed
+
+    input_rate, output_rate = rates
     cost = (prompt_tokens / 1_000_000) * input_rate + (completion_tokens / 1_000_000) * output_rate
     return round(cost, 8)
 
@@ -257,6 +274,8 @@ def backfill_personas(db, dry_run: bool) -> int:
 def main():
     parser = argparse.ArgumentParser(description="Backfill usage_events from existing data")
     parser.add_argument("--dry-run", action="store_true", help="Preview without writing")
+    parser.add_argument("--delete-existing-estimates", action="store_true",
+                        help="Delete previously created estimated events before backfilling")
     args = parser.parse_args()
 
     if args.dry_run:
@@ -264,6 +283,13 @@ def main():
 
     db = connect()
 
+    if args.delete_existing_estimates and not args.dry_run:
+        result = db.usage_events.delete_many({"is_estimated": True})
+        print(f"Deleted {result.deleted_count} existing estimated events\n")
+
+    _load_pricing(db)
+    print(f"Loaded {len(_pricing_cache)} pricing rows: {list(_pricing_cache.keys())}")
+
     total = 0
     total += backfill_messages(db, args.dry_run)
     total += backfill_personas(db, args.dry_run)