From d80ff29b8fb87a05ada1f27a69931e4510e9e1c7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:51:35 +0000 Subject: [PATCH] fix(spend_tracking): preserve call_type and router metadata on failure spend logs --- .../proxy/hooks/proxy_track_cost_callback.py | 54 ++++++++++-- litellm/proxy/utils.py | 24 +++++- .../hooks/test_proxy_track_cost_callback.py | 57 +++++++++++++ tests/test_litellm/proxy/test_proxy_utils.py | 84 +++++++++++++++++++ 4 files changed, 207 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b839426fcda..3048809215e 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,6 +1,7 @@ import asyncio import traceback from datetime import datetime +from collections.abc import Mapping from typing import Any, List, Optional, Union, cast import litellm @@ -9,6 +10,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_litellm_metadata_from_kwargs, + get_metadata_variable_name_from_kwargs, ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.proxy._types import UserAPIKeyAuth @@ -106,23 +108,23 @@ class _ProxyDBLogger(CustomLogger): metadata=_metadata, ) - existing_metadata: dict = request_data.get("metadata", None) or {} - existing_metadata.update(_metadata) - if "litellm_params" not in request_data: request_data["litellm_params"] = {} existing_litellm_params = request_data.get("litellm_params", {}) - existing_litellm_metadata = existing_litellm_params.get("metadata", {}) or {} - - # Preserve tags from existing metadata - if existing_litellm_metadata.get("tags"): - existing_metadata["tags"] = existing_litellm_metadata.get("tags") + metadata_key = get_metadata_variable_name_from_kwargs(request_data) + merged_metadata = _merge_failure_metadata_buckets( + request_metadata=request_data.get(metadata_key), + litellm_params=existing_litellm_params, + trusted_metadata=_metadata, + ) request_data["litellm_params"]["proxy_server_request"] = ( request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} ) - request_data["litellm_params"]["metadata"] = existing_metadata + request_data["litellm_params"]["metadata"] = merged_metadata + if metadata_key != "metadata": + request_data["litellm_params"][metadata_key] = dict(merged_metadata) # Preserve model name and custom_llm_provider if "model" not in request_data: @@ -418,6 +420,40 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value +def _as_metadata_mapping(value: object) -> Mapping[str, object]: + return value if isinstance(value, Mapping) else {} + + +def _merge_failure_metadata_buckets( + request_metadata: Mapping[str, object] | None, + litellm_params: Mapping[str, object], + trusted_metadata: Mapping[str, object], +) -> Mapping[str, object]: + """ + Build the metadata for a failure spend log out of every bucket the request may + have used. + + Routes such as ``/v1/responses`` keep proxy-internal state (``model_group``, + ``model_info``, retry counts, tags) in ``litellm_metadata`` rather than + ``metadata``, so reading a single bucket drops router attribution. Key identity + fields always come from the authenticated key, never from the request body. + """ + caller_metadata = { + key: value + for key, value in (request_metadata or {}).items() + if not key.startswith("user_api_key") and key != "status" + } + base = { + **_as_metadata_mapping(litellm_params.get("metadata")), + **_as_metadata_mapping(litellm_params.get("litellm_metadata")), + **caller_metadata, + } + return { + **base, + **{key: value for key, value in trusted_metadata.items() if value is not None or base.get(key) is None}, + } + + def _should_track_cost_callback( user_api_key: Optional[str], user_id: Optional[str], diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 924189fed4b..e9756728299 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -401,6 +401,9 @@ class _CallbackCapabilities: resolved_callbacks: Tuple[Any, ...] = field(default_factory=tuple) +_KNOWN_CALL_TYPES: frozenset[str] = frozenset(call_type.value for call_type in CallTypes) + + class ProxyLogging: """ Logging/Custom Handlers for proxy. @@ -2117,6 +2120,15 @@ class ProxyLogging: if _first_handoff is not None: request_data["first_api_call_start_time"] = _first_handoff + # The spend log derives its call type from request_data, which carries the + # raw request body and never the route's call type. Lift it over before the + # logging object is popped so a failed /v1/responses call isn't recorded as + # a chat completion. + if not request_data.get("call_type"): + _call_type = _model_call_details.get("call_type") or getattr(_logging_obj, "call_type", None) + if _call_type in _KNOWN_CALL_TYPES: + request_data["call_type"] = _call_type + # A stream that broke mid-flight still billed the provider for the # chunks already delivered; the streaming handler stashes that # recovered usage and cost here. Lift them onto request_data so the @@ -2252,20 +2264,26 @@ class ProxyLogging: input: Union[list, str, dict] = "" normalized_call_type: Optional[str] = None + # A logging object built from a route string (rather than a litellm + # function) has no real call type, so infer one from the payload shape. + # A call type that litellm already resolved is authoritative and must + # survive; /v1/responses passes its prompt in ``input``, which would + # otherwise be mistaken for an embedding request. + can_infer_call_type = litellm_logging_obj.call_type not in _KNOWN_CALL_TYPES if "messages" in request_data and isinstance(request_data["messages"], list): input = request_data["messages"] litellm_logging_obj.model_call_details["messages"] = input - if litellm_logging_obj.call_type != CallTypes.pass_through.value: + if can_infer_call_type: normalized_call_type = CallTypes.acompletion.value elif "prompt" in request_data and isinstance(request_data["prompt"], str): input = request_data["prompt"] litellm_logging_obj.model_call_details["prompt"] = input - if litellm_logging_obj.call_type != CallTypes.pass_through.value: + if can_infer_call_type: normalized_call_type = CallTypes.atext_completion.value elif "input" in request_data and isinstance(request_data["input"], list): input = request_data["input"] litellm_logging_obj.model_call_details["input"] = input - if litellm_logging_obj.call_type != CallTypes.pass_through.value: + if can_infer_call_type: normalized_call_type = CallTypes.aembedding.value if normalized_call_type is not None: litellm_logging_obj.call_type = normalized_call_type diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index f289148101a..4ede6f7a12e 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -84,6 +84,63 @@ async def test_async_post_call_failure_hook(): assert metadata["original_key"] == "original_value" +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_preserves_router_metadata_from_litellm_metadata(): + """Routes that keep proxy state in ``litellm_metadata`` (e.g. /v1/responses) must + still get router attribution (model_group, deployment id, retries, tags) on the + failure spend log, and caller-supplied identity fields must be ignored.""" + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + logger = _ProxyDBLogger() + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", + user_id="test_user_id", + team_id="test_team_id", + ) + request_data = { + "model": "test-model", + "litellm_metadata": { + "model_group": "test-model-group", + "model_info": {"id": "test-deployment-id"}, + "attempted_retries": 2, + "max_retries": 2, + "user_api_key_user_id": "spoofed_user_id", + }, + "litellm_params": {"metadata": {"tags": ["failure-test"]}}, + "call_type": "aresponses", + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Test exception"), + user_api_key_dict=user_api_key_dict, + ) + + kwargs = mock_update_database.call_args[1]["kwargs"] + metadata = kwargs["litellm_params"]["metadata"] + assert metadata["model_group"] == "test-model-group" + assert metadata["model_info"] == {"id": "test-deployment-id"} + assert metadata["attempted_retries"] == 2 + assert metadata["max_retries"] == 2 + assert metadata["tags"] == ["failure-test"] + assert metadata["user_api_key_user_id"] == "test_user_id" + + payload = get_logging_payload( + kwargs=kwargs, + response_obj=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert payload["call_type"] == "aresponses" + assert payload["model_group"] == "test-model-group" + assert payload["model_id"] == "test-deployment-id" + assert payload["request_tags"] == '["failure-test"]' + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_non_llm_route(): # Setup diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 4673807a135..7f54ec339d3 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -175,6 +175,90 @@ async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params(): assert "litellm_metadata" not in captured["optional_params"] +@pytest.mark.asyncio +async def test_proxy_only_error_log_preserves_resolved_call_type(): + """A /v1/responses request already has ``aresponses`` resolved on its logging + object; the payload-shape inference must not rewrite it to ``aembedding``.""" + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import CallTypes + + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + logging_obj = Logging( + model="gpt-4o", + messages=[], + stream=False, + call_type=CallTypes.aresponses.value, + start_time=real_datetime.datetime.now(), + litellm_call_id="1234", + function_id="1234", + ) + + orig_pre_call = Logging.pre_call + orig_async_failure = Logging.async_failure_handler + + async def _noop_async_failure(self, *args, **kwargs): + return None + + Logging.pre_call = lambda self, *args, **kwargs: None + Logging.async_failure_handler = _noop_async_failure + try: + await proxy_logging_obj._handle_logging_proxy_only_error( + request_data={ + "model": "gpt-4o", + "input": [{"role": "user", "content": "hi"}], + "litellm_logging_obj": logging_obj, + }, + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-1234", request_route="/v1/responses" + ), + route="/v1/responses", + original_exception=HTTPException(status_code=429, detail="rate limited"), + ) + finally: + Logging.pre_call = orig_pre_call + Logging.async_failure_handler = orig_async_failure + + assert logging_obj.call_type == CallTypes.aresponses.value + assert logging_obj.model_call_details["call_type"] == CallTypes.aresponses.value + + +@pytest.mark.asyncio +async def test_post_call_failure_hook_lifts_call_type_onto_request_data(): + """The logging object is popped before the failure callbacks run, so the route's + call type must be lifted onto request_data or the spend log falls back to + ``acompletion`` for every failed request.""" + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import CallTypes + + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + logging_obj = Logging( + model="gpt-4o", + messages=[], + stream=False, + call_type=CallTypes.aresponses.value, + start_time=real_datetime.datetime.now(), + litellm_call_id="1234", + function_id="1234", + ) + request_data = { + "model": "gpt-4o", + "input": [{"role": "user", "content": "hi"}], + "litellm_logging_obj": logging_obj, + } + + await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=Exception("upstream boom"), + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-1234", request_route="/v1/responses" + ), + ) + + assert request_data["call_type"] == CallTypes.aresponses.value + + def test_get_model_group_info_order(): from litellm import Router from litellm.proxy.proxy_server import _get_model_group_info