add validation to get map

litellm/litellm_core_utils/get_model_cost_map.py

@@ -8,25 +8,45 @@ export LITELLM_LOCAL_MODEL_COST_MAP=True
 ```
 """
 
+import json
 import os
+from importlib.resources import files
 
 import httpx
 
+from litellm.constants import MIN_MODEL_COST_MAP_ENTRIES
+
+
+def _load_local_model_cost_map() -> dict:
+    """Load the model cost map from the bundled backup file."""
+    return json.loads(
+        files("litellm")
+        .joinpath("model_prices_and_context_window_backup.json")
+        .read_text(encoding="utf-8")
+    )
+
+
+def validate_model_cost_map(data: dict) -> bool:
+    """
+    Returns True if the model cost map looks structurally sound.
+
+    Checks:
+    - data is a dict
+    - has more than MIN_MODEL_COST_MAP_ENTRIES entries
+    """
+    if not isinstance(data, dict):
+        return False
+    if len(data) < MIN_MODEL_COST_MAP_ENTRIES:
+        return False
+    return True
+
 
 def get_model_cost_map(url: str) -> dict:
     if (
         os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", False)
         or os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", False) == "True"
     ):
-        from importlib.resources import files
-        import json
-
-        content = json.loads(
-            files("litellm")
-            .joinpath("model_prices_and_context_window_backup.json")
-            .read_text(encoding="utf-8")
-        )
-        return content
+        return _load_local_model_cost_map()
 
     try:
         response = httpx.get(
@@ -34,14 +54,13 @@ def get_model_cost_map(url: str) -> dict:
         )  # set a 5 second timeout for the get request
         response.raise_for_status()  # Raise an exception if the request is unsuccessful
         content = response.json()
+
+        if not validate_model_cost_map(content):
+            raise ValueError(
+                "Remote model cost map failed validation: "
+                f"got {type(content).__name__} with {len(content) if isinstance(content, dict) else 'N/A'} entries"
+            )
+
         return content
     except Exception:
-        from importlib.resources import files
-        import json
-
-        content = json.loads(
-            files("litellm")
-            .joinpath("model_prices_and_context_window_backup.json")
-            .read_text(encoding="utf-8")
-        )
-        return content
+        return _load_local_model_cost_map()