diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index cecc35ee1c1..6ac56a666fb 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -164,18 +164,39 @@ def remove_items_at_indices(items: Optional[List[Any]], indices: Iterable[int]) items.pop(index) +ROUTER_DEPLOYMENT_METADATA_KEYS = frozenset( + { + "model_group", + "model_group_alias", + "model_info", + "deployment", + "deployment_model_name", + "api_base", + } +) + + def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metadata: dict) -> dict: """ - Helper to get litellm metadata for spend tracking + Merge spend-tracking fields that live in `metadata` into `litellm_metadata`. - PATCH for issue where both `litellm_metadata` and `metadata` are present in the kwargs - and user_api_key values are in 'metadata'. + Both dicts can be present on the same request: the proxy and the Router write auth + (`user_api_key*`) and deployment attribution (`model_group`, `model_info`, ...) into + `metadata`, while a caller or a pre-call guardrail can create `litellm_metadata`. + Reading only `litellm_metadata` then loses those fields. + + `user_api_key*` values in `metadata` are proxy-authoritative and win. Deployment + attribution is only filled in when absent from `litellm_metadata`, so the Router keeps + precedence on the endpoints where it writes to `litellm_metadata` directly. """ - 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 + return { + **litellm_metadata, + **{ + key: value + for key, value in metadata.items() + if "user_api_key" in key or (key in ROUTER_DEPLOYMENT_METADATA_KEYS and key not in litellm_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..d0a25e72efa 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, + get_litellm_metadata_from_kwargs, get_or_create_metadata_bucket, map_finish_reason, reconstruct_model_name, @@ -51,6 +52,65 @@ class TestGetOrCreateMetadataBucket: assert bucket is request_data["litellm_metadata"] +class TestGetLitellmMetadataFromKwargs: + """A pre-call guardrail can create `litellm_metadata` while the Router keeps writing + deployment attribution to `metadata`; spend tracking must see both.""" + + def test_router_deployment_fields_survive_guardrail_litellm_metadata(self): + kwargs = { + "litellm_params": { + "metadata": { + "model_group": "router-model-group", + "model_info": {"id": "deployment-id"}, + "deployment": "bedrock/claude-sonnet-4-6", + "user_api_key": "hashed-key", + }, + "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["deployment"] == "bedrock/claude-sonnet-4-6" + assert metadata["user_api_key"] == "hashed-key" + assert metadata["user_api_key_user_id"] == "test-user" + + def test_litellm_metadata_keeps_precedence_on_deployment_fields(self): + kwargs = { + "litellm_params": { + "metadata": {"model_group": "caller-supplied"}, + "litellm_metadata": {"model_group": "router-model-group"}, + } + } + + assert get_litellm_metadata_from_kwargs(kwargs)["model_group"] == "router-model-group" + + def test_unrelated_metadata_keys_are_not_merged(self): + kwargs = { + "litellm_params": { + "metadata": {"model_group": "router-model-group", "caller_key": "caller-value"}, + "litellm_metadata": {"user_api_key_user_id": "test-user"}, + } + } + + assert "caller_key" not in get_litellm_metadata_from_kwargs(kwargs) + + 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": "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.""" 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..af29b98c70c 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 @@ -1065,6 +1065,37 @@ def test_get_logging_payload_includes_agent_id_from_kwargs(): ), f"Expected agent_id '{test_agent_id}', got '{payload.get('agent_id')}'" +def test_get_logging_payload_keeps_router_model_group_when_guardrail_adds_litellm_metadata(): + """ + A pre-call guardrail creates `litellm_metadata` holding only the auth fields, while the + Router records deployment attribution in `metadata`. Spend logs must still be attributable + to the model group and the deployment. + """ + kwargs = { + "model": "claude-sonnet-4-6", + "custom_llm_provider": "anthropic", + "litellm_params": { + "metadata": { + "user_api_key": "sk-test-key", + "model_group": "router-model-group", + "model_info": {"id": "deployment-id-123"}, + }, + "litellm_metadata": {"user_api_key_user_id": "test-user"}, + }, + } + + payload = get_logging_payload( + kwargs=kwargs, + 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"] == "router-model-group" + assert payload["model_id"] == "deployment-id-123" + assert payload["user"] == "test-user" + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_includes_overhead_in_spend_logs_metadata():