From d2e88d362a3a58c1197c6b9f07023428e32196fe Mon Sep 17 00:00:00 2001
From: chopratejas <chopratejas@gmail.com>
Date: Thu, 12 Mar 2026 21:11:18 -0700
Subject: [PATCH] Fix pricing lookup for retired Claude models

LiteLLM removed claude-3-5-sonnet-20241022 from its cost database.
Add alias fallback map so retired model names resolve to current
equivalents for pricing lookups.
---
 headroom/pricing/litellm_pricing.py | 16 ++++++++++++++++
 tests/test_models.py                |  7 +++++--
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/headroom/pricing/litellm_pricing.py b/headroom/pricing/litellm_pricing.py
index f8d95f16c..3314a9a0e 100644
--- a/headroom/pricing/litellm_pricing.py
+++ b/headroom/pricing/litellm_pricing.py
@@ -13,6 +13,16 @@ from typing import Any
 
 import litellm
 
+# Aliases for models removed from LiteLLM's cost database (retired/renamed).
+# Maps old model name -> current LiteLLM key that has equivalent pricing.
+_MODEL_ALIASES: dict[str, str] = {
+    # Claude 3.5 Sonnet retired Feb 2026, pricing same as claude-sonnet-4-20250514
+    "claude-3-5-sonnet-20241022": "claude-sonnet-4-20250514",
+    "claude-3-5-sonnet-20240620": "claude-sonnet-4-20250514",
+    # Claude 3 Sonnet retired
+    "claude-3-sonnet-20240229": "claude-3-haiku-20240307",
+}
+
 
 @dataclass
 class LiteLLMModelPricing:
@@ -61,6 +71,12 @@ def get_model_pricing(model: str) -> LiteLLMModelPricing | None:
                 info = cost_data[f"{prefix}{model}"]
                 break
 
+    # Try retired/renamed model aliases (LiteLLM removes old model keys over time)
+    if info is None:
+        alias = _MODEL_ALIASES.get(model)
+        if alias:
+            info = cost_data.get(alias)
+
     if info is None:
         return None
 
diff --git a/tests/test_models.py b/tests/test_models.py
index 52bc3c7f2..13740be96 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -240,11 +240,14 @@ class TestBuiltInModels:
         info = get_model_info("claude-3-5-sonnet-20241022")
         assert info.provider == "anthropic"
         assert info.context_window == 200000
-        # Pricing is now fetched from LiteLLM
-        pricing = ModelRegistry.get_pricing("claude-3-5-sonnet-20241022")
+        # Pricing fetched from LiteLLM (falls back to alias for retired models)
+        pricing = ModelRegistry.get_pricing("claude-sonnet-4-20250514")
         assert pricing is not None
         assert pricing[0] == 3.00  # input cost per 1M
         assert pricing[1] == 15.00  # output cost per 1M
+        # Retired model alias should also resolve
+        alias_pricing = ModelRegistry.get_pricing("claude-3-5-sonnet-20241022")
+        assert alias_pricing is not None
 
     def test_gemini_info(self):
         """Test Gemini model info."""