fix(spend_tracking): keep router model_group when guardrails add litellm_metadata

litellm/litellm_core_utils/core_helpers.py

@@ -164,18 +164,25 @@ def remove_items_at_indices(items: Optional[List[Any]], indices: Iterable[int])
             items.pop(index)
 
 
-def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metadata: dict) -> dict:
+ROUTER_SPEND_TRACKING_METADATA_KEYS = frozenset({"model_group", "model_info"})
+
+
+def add_missing_spend_metadata_to_litellm_metadata(
+    litellm_metadata: dict[str, object], metadata: dict[str, object]
+) -> dict[str, object]:
     """
     Helper to get litellm metadata for spend tracking
 
-    PATCH for issue where both `litellm_metadata` and `metadata` are present in the kwargs
-    and user_api_key values are in 'metadata'.
+    Handles the case where both `litellm_metadata` and `metadata` are present in the kwargs:
+    the auth values (`user_api_key*`) and the router deployment values (`model_group`,
+    `model_info`) live in `metadata`, so they'd otherwise be lost when `litellm_metadata` wins.
     """
-    potential_spend_tracking_metadata_substring = "user_api_key"
-    for key, value in metadata.items():
-        if potential_spend_tracking_metadata_substring in key:
-            litellm_metadata[key] = value
-    return litellm_metadata
+    carried_over = {
+        key: value
+        for key, value in metadata.items()
+        if "user_api_key" in key or (key in ROUTER_SPEND_TRACKING_METADATA_KEYS and not litellm_metadata.get(key))
+    }
+    return {**litellm_metadata, **carried_over}
 
 
 def get_metadata_variable_name_from_kwargs(

tests/test_litellm/litellm_core_utils/test_core_helpers.py

@@ -4,6 +4,7 @@ import pytest
 
 from litellm.litellm_core_utils.core_helpers import (
     _FINISH_REASON_MAP,
+    get_litellm_metadata_from_kwargs,
     get_or_create_metadata_bucket,
     map_finish_reason,
     reconstruct_model_name,
@@ -51,6 +52,68 @@ class TestGetOrCreateMetadataBucket:
         assert bucket is request_data["litellm_metadata"]
 
 
+class TestGetLitellmMetadataFromKwargs:
+    """Spend tracking reads the router deployment values from the merged metadata, so a
+    guardrail-created `litellm_metadata` must not shadow them away."""
+
+    def test_carries_router_deployment_values_from_metadata(self):
+        kwargs = {
+            "litellm_params": {
+                "metadata": {
+                    "model_group": "router-model-group",
+                    "model_info": {"id": "deployment-id"},
+                },
+                "litellm_metadata": {"user_api_key_user_id": "test-user"},
+            }
+        }
+
+        metadata = get_litellm_metadata_from_kwargs(kwargs)
+
+        assert metadata["model_group"] == "router-model-group"
+        assert metadata["model_info"] == {"id": "deployment-id"}
+        assert metadata["user_api_key_user_id"] == "test-user"
+
+    def test_carries_auth_values_from_metadata(self):
+        kwargs = {
+            "litellm_params": {
+                "metadata": {"user_api_key_hash": "hashed-key", "unrelated": "dropped"},
+                "litellm_metadata": {"model_group": "litellm-metadata-group"},
+            }
+        }
+
+        metadata = get_litellm_metadata_from_kwargs(kwargs)
+
+        assert metadata["user_api_key_hash"] == "hashed-key"
+        assert metadata["model_group"] == "litellm-metadata-group"
+        assert "unrelated" not in metadata
+
+    def test_litellm_metadata_deployment_values_win(self):
+        kwargs = {
+            "litellm_params": {
+                "metadata": {"model_group": "stale-group"},
+                "litellm_metadata": {"model_group": "current-group"},
+            }
+        }
+
+        assert get_litellm_metadata_from_kwargs(kwargs)["model_group"] == "current-group"
+
+    def test_does_not_mutate_the_request_metadata(self):
+        litellm_metadata = {"user_api_key_user_id": "test-user"}
+        kwargs = {
+            "litellm_params": {
+                "metadata": {
+                    "model_group": "router-model-group",
+                    "user_api_key_hash": "hashed-key",
+                },
+                "litellm_metadata": litellm_metadata,
+            }
+        }
+
+        get_litellm_metadata_from_kwargs(kwargs)
+
+        assert litellm_metadata == {"user_api_key_user_id": "test-user"}
+
+
 def test_reconstruct_model_name_prefers_deployment_value():
     """Ensure deployment metadata wins when reconstructing the model name."""