scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(spend_tracking): preserve call_type and router metadata on failure spend logs

Browse source
This commit is contained in:
Devin AI 2026-07-29 15:51:35 +00:00
parent c274cf321c
commit d80ff29b8f
4 changed files with 207 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -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],

24
litellm/proxy/utils.py
View file

@ -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

57
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -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

84
tests/test_litellm/proxy/test_proxy_utils.py
View file

@ -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