diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index cecc35ee1c1..49457135915 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -168,14 +168,17 @@ def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metad """ 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'. + Both `litellm_metadata` and `metadata` can be present in the kwargs (e.g. a pre-call + guardrail creates `litellm_metadata` while the router populated `metadata`). Spend + tracking reads a single bucket, so fill in every key `litellm_metadata` is missing + (`model_group`, `model_info`, ...) from `metadata`, and let `metadata` win for the + `user_api_key*` keys, which are always the authoritative auth values. """ 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 + auth_metadata = { + key: value for key, value in metadata.items() if potential_spend_tracking_metadata_substring in key + } + return {**metadata, **litellm_metadata, **auth_metadata} def get_metadata_variable_name_from_kwargs( diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index b4f539da286..6ebdd13cee7 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/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, + add_missing_spend_metadata_to_litellm_metadata, get_or_create_metadata_bucket, map_finish_reason, reconstruct_model_name, @@ -242,3 +243,28 @@ class TestRedactNestedMatchAndRegexKeys: def test_passes_through_none_and_str(self): assert redact_nested_match_and_regex_keys(None) is None assert redact_nested_match_and_regex_keys("plain") == "plain" + + +class TestAddMissingSpendMetadataToLitellmMetadata: + def test_fills_missing_router_fields_without_clobbering_litellm_metadata(self): + merged = add_missing_spend_metadata_to_litellm_metadata( + litellm_metadata={"guardrail_added": True, "model_group": "from-litellm-metadata"}, + metadata={ + "model_group": "from-metadata", + "model_info": {"id": "deployment-123"}, + "user_api_key_team_id": "team-1", + }, + ) + + assert merged["model_info"] == {"id": "deployment-123"} + assert merged["user_api_key_team_id"] == "team-1" + assert merged["guardrail_added"] is True + assert merged["model_group"] == "from-litellm-metadata" + + def test_auth_keys_always_come_from_metadata(self): + merged = add_missing_spend_metadata_to_litellm_metadata( + litellm_metadata={"user_api_key": "stale"}, + metadata={"user_api_key": "authoritative"}, + ) + + assert merged["user_api_key"] == "authoritative" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index cc1e2943c8f..6e8938f624e 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2902,3 +2902,27 @@ async def test_compression_savings_survive_to_spend_log_payload_metadata(monkeyp "tokens_saved": 7000, "source": "compression_interception", } + + +def test_get_logging_payload_keeps_router_fields_when_guardrail_adds_litellm_metadata(): + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_params": { + "metadata": { + "user_api_key": "test-key", + "user_api_key_team_id": "team-1", + "model_group": "my-model-group", + "model_info": {"id": "deployment-123"}, + }, + "litellm_metadata": {"guardrail_added": True}, + }, + }, + response_obj=litellm.ModelResponse(id="chatcmpl-test", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["model_group"] == "my-model-group" + assert payload["model_id"] == "deployment-123" + assert json.loads(payload["metadata"])["user_api_key_team_id"] == "team-1"