mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Fix root cause: DB spend log session_id didn't match Langfuse trace_id
The proxy has two separate failure paths: 1. async_failure_handler → Langfuse callback (uses model_call_details with standard_logging_object containing the correct trace_id) 2. post_call_failure_hook → _ProxyDBLogger → spend log (uses request_data which did NOT have standard_logging_object, so session_id fell to random uuid4()) These two paths used different data dicts, so the DB session_id was a random UUID unrelated to the Langfuse trace_id. Users could not search by the Session ID from LiteLLM logs in Langfuse for failed requests. Fix: In _ProxyDBLogger.async_post_call_failure_hook, propagate standard_logging_object and litellm_trace_id from the litellm_logging_obj (already present in request_data) before writing the spend log. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Harshit28j
parent
15f9ad1f68
commit
7593ebc606
2 changed files with 650 additions and 279 deletions
|
|
@ -1,279 +1,297 @@
|
|||
import asyncio
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import Any, List, Optional, Union, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
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,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import log_db_metrics
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.utils import ProxyUpdateSpend
|
||||
from litellm.types.utils import (
|
||||
StandardLoggingPayload,
|
||||
StandardLoggingUserAPIKeyMetadata,
|
||||
)
|
||||
from litellm.utils import get_end_user_id_for_cost_tracking
|
||||
|
||||
|
||||
class _ProxyDBLogger(CustomLogger):
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
await self._PROXY_track_cost_callback(
|
||||
kwargs, response_obj, start_time, end_time
|
||||
)
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict,
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: Optional[str] = None,
|
||||
):
|
||||
request_route = user_api_key_dict.request_route
|
||||
if _ProxyDBLogger._should_track_errors_in_db() is False:
|
||||
return
|
||||
elif request_route is not None and not RouteChecks.is_llm_api_route(
|
||||
route=request_route
|
||||
):
|
||||
return
|
||||
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
_metadata = dict(
|
||||
StandardLoggingUserAPIKeyMetadata(
|
||||
user_api_key_hash=user_api_key_dict.api_key,
|
||||
user_api_key_alias=user_api_key_dict.key_alias,
|
||||
user_api_key_spend=user_api_key_dict.spend,
|
||||
user_api_key_max_budget=user_api_key_dict.max_budget,
|
||||
user_api_key_budget_reset_at=(
|
||||
user_api_key_dict.budget_reset_at.isoformat()
|
||||
if user_api_key_dict.budget_reset_at
|
||||
else None
|
||||
),
|
||||
user_api_key_user_email=user_api_key_dict.user_email,
|
||||
user_api_key_user_id=user_api_key_dict.user_id,
|
||||
user_api_key_team_id=user_api_key_dict.team_id,
|
||||
user_api_key_org_id=user_api_key_dict.org_id,
|
||||
user_api_key_team_alias=user_api_key_dict.team_alias,
|
||||
user_api_key_end_user_id=user_api_key_dict.end_user_id,
|
||||
user_api_key_request_route=user_api_key_dict.request_route,
|
||||
user_api_key_auth_metadata=user_api_key_dict.metadata,
|
||||
)
|
||||
)
|
||||
_metadata["user_api_key"] = user_api_key_dict.api_key
|
||||
_metadata["status"] = "failure"
|
||||
_metadata["error_information"] = (
|
||||
StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=original_exception,
|
||||
traceback_str=traceback_str,
|
||||
)
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
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
|
||||
|
||||
# Preserve model name and custom_llm_provider
|
||||
if "model" not in request_data:
|
||||
request_data["model"] = existing_litellm_params.get("model") or request_data.get("model", "")
|
||||
if "custom_llm_provider" not in request_data:
|
||||
request_data["custom_llm_provider"] = existing_litellm_params.get("custom_llm_provider") or request_data.get("custom_llm_provider", "")
|
||||
|
||||
await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
token=user_api_key_dict.api_key,
|
||||
response_cost=0.0,
|
||||
user_id=user_api_key_dict.user_id,
|
||||
end_user_id=user_api_key_dict.end_user_id,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
kwargs=request_data,
|
||||
completion_response=original_exception,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
org_id=user_api_key_dict.org_id,
|
||||
)
|
||||
|
||||
@log_db_metrics
|
||||
async def _PROXY_track_cost_callback(
|
||||
self,
|
||||
kwargs, # kwargs to completion
|
||||
completion_response: Optional[
|
||||
Union[litellm.ModelResponse, Any]
|
||||
], # response from completion
|
||||
start_time=None,
|
||||
end_time=None, # start/end time for completion
|
||||
):
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj, update_cache
|
||||
|
||||
verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback")
|
||||
try:
|
||||
verbose_proxy_logger.debug(
|
||||
f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}"
|
||||
)
|
||||
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs=kwargs)
|
||||
litellm_params = kwargs.get("litellm_params", {}) or {}
|
||||
end_user_id = get_end_user_id_for_cost_tracking(litellm_params)
|
||||
metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs)
|
||||
user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None))
|
||||
team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None))
|
||||
org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None))
|
||||
key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None))
|
||||
end_user_max_budget = metadata.get("user_api_end_user_max_budget", None)
|
||||
sl_object: Optional[StandardLoggingPayload] = kwargs.get(
|
||||
"standard_logging_object", None
|
||||
)
|
||||
response_cost = (
|
||||
sl_object.get("response_cost", None)
|
||||
if sl_object is not None
|
||||
else kwargs.get("response_cost", None)
|
||||
)
|
||||
tags: Optional[List[str]] = (
|
||||
sl_object.get("request_tags", None) if sl_object is not None else None
|
||||
)
|
||||
|
||||
if response_cost is not None:
|
||||
user_api_key = metadata.get("user_api_key", None)
|
||||
if kwargs.get("cache_hit", False) is True:
|
||||
response_cost = 0.0
|
||||
verbose_proxy_logger.debug(
|
||||
f"Cache Hit: response_cost {response_cost}, for user_id {user_id}"
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}"
|
||||
)
|
||||
if _should_track_cost_callback(
|
||||
user_api_key=user_api_key,
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
end_user_id=end_user_id,
|
||||
):
|
||||
## UPDATE DATABASE
|
||||
await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
token=user_api_key,
|
||||
response_cost=response_cost,
|
||||
user_id=user_id,
|
||||
end_user_id=end_user_id,
|
||||
team_id=team_id,
|
||||
kwargs=kwargs,
|
||||
completion_response=completion_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
org_id=org_id,
|
||||
)
|
||||
|
||||
# update cache
|
||||
asyncio.create_task(
|
||||
update_cache(
|
||||
token=user_api_key,
|
||||
user_id=user_id,
|
||||
end_user_id=end_user_id,
|
||||
response_cost=response_cost,
|
||||
team_id=team_id,
|
||||
parent_otel_span=parent_otel_span,
|
||||
tags=tags,
|
||||
)
|
||||
)
|
||||
|
||||
await proxy_logging_obj.slack_alerting_instance.customer_spend_alert(
|
||||
token=user_api_key,
|
||||
key_alias=key_alias,
|
||||
end_user_id=end_user_id,
|
||||
response_cost=response_cost,
|
||||
max_budget=end_user_max_budget,
|
||||
)
|
||||
else:
|
||||
if kwargs.get("stream") is not True or (
|
||||
kwargs.get("stream") is True and "complete_streaming_response" in kwargs
|
||||
):
|
||||
if sl_object is not None:
|
||||
cost_tracking_failure_debug_info: Union[dict, str] = (
|
||||
sl_object["response_cost_failure_debug_info"] # type: ignore
|
||||
or "response_cost_failure_debug_info is None in standard_logging_object"
|
||||
)
|
||||
else:
|
||||
cost_tracking_failure_debug_info = (
|
||||
"standard_logging_object not found"
|
||||
)
|
||||
model = kwargs.get("model")
|
||||
raise Exception(
|
||||
f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
|
||||
)
|
||||
except Exception as e:
|
||||
error_msg = f"Error in tracking cost callback - {str(e)}\n Traceback:{traceback.format_exc()}"
|
||||
model = kwargs.get("model", "")
|
||||
metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs)
|
||||
litellm_metadata = kwargs.get("litellm_params", {}).get(
|
||||
"litellm_metadata", {}
|
||||
)
|
||||
old_metadata = kwargs.get("litellm_params", {}).get("metadata", {})
|
||||
call_type = kwargs.get("call_type", "")
|
||||
error_msg += f"\n Args to _PROXY_track_cost_callback\n model: {model}\n chosen_metadata: {metadata}\n litellm_metadata: {litellm_metadata}\n old_metadata: {old_metadata}\n call_type: {call_type}\n"
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.failed_tracking_alert(
|
||||
error_message=error_msg,
|
||||
failing_model=model,
|
||||
)
|
||||
)
|
||||
|
||||
verbose_proxy_logger.exception(
|
||||
"Error in tracking cost callback - %s", str(e)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _should_track_errors_in_db():
|
||||
"""
|
||||
Returns True if errors should be tracked in the database
|
||||
|
||||
By default, errors are tracked in the database
|
||||
|
||||
If users want to disable error tracking, they can set the disable_error_logs flag in the general_settings
|
||||
"""
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
if general_settings.get("disable_error_logs") is True:
|
||||
return False
|
||||
return
|
||||
|
||||
|
||||
def _should_track_cost_callback(
|
||||
user_api_key: Optional[str],
|
||||
user_id: Optional[str],
|
||||
team_id: Optional[str],
|
||||
end_user_id: Optional[str],
|
||||
) -> bool:
|
||||
"""
|
||||
Determine if the cost callback should be tracked based on the kwargs
|
||||
"""
|
||||
|
||||
# don't run track cost callback if user opted into disabling spend
|
||||
if ProxyUpdateSpend.disable_spend_updates() is True:
|
||||
return False
|
||||
|
||||
if (
|
||||
user_api_key is not None
|
||||
or user_id is not None
|
||||
or team_id is not None
|
||||
or end_user_id is not None
|
||||
):
|
||||
return True
|
||||
return False
|
||||
import asyncio
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import Any, List, Optional, Union, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
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,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import log_db_metrics
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.utils import ProxyUpdateSpend
|
||||
from litellm.types.utils import (
|
||||
StandardLoggingPayload,
|
||||
StandardLoggingUserAPIKeyMetadata,
|
||||
)
|
||||
from litellm.utils import get_end_user_id_for_cost_tracking
|
||||
|
||||
|
||||
class _ProxyDBLogger(CustomLogger):
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
await self._PROXY_track_cost_callback(
|
||||
kwargs, response_obj, start_time, end_time
|
||||
)
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict,
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: Optional[str] = None,
|
||||
):
|
||||
request_route = user_api_key_dict.request_route
|
||||
if _ProxyDBLogger._should_track_errors_in_db() is False:
|
||||
return
|
||||
elif request_route is not None and not RouteChecks.is_llm_api_route(
|
||||
route=request_route
|
||||
):
|
||||
return
|
||||
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
_metadata = dict(
|
||||
StandardLoggingUserAPIKeyMetadata(
|
||||
user_api_key_hash=user_api_key_dict.api_key,
|
||||
user_api_key_alias=user_api_key_dict.key_alias,
|
||||
user_api_key_spend=user_api_key_dict.spend,
|
||||
user_api_key_max_budget=user_api_key_dict.max_budget,
|
||||
user_api_key_budget_reset_at=(
|
||||
user_api_key_dict.budget_reset_at.isoformat()
|
||||
if user_api_key_dict.budget_reset_at
|
||||
else None
|
||||
),
|
||||
user_api_key_user_email=user_api_key_dict.user_email,
|
||||
user_api_key_user_id=user_api_key_dict.user_id,
|
||||
user_api_key_team_id=user_api_key_dict.team_id,
|
||||
user_api_key_org_id=user_api_key_dict.org_id,
|
||||
user_api_key_team_alias=user_api_key_dict.team_alias,
|
||||
user_api_key_end_user_id=user_api_key_dict.end_user_id,
|
||||
user_api_key_request_route=user_api_key_dict.request_route,
|
||||
user_api_key_auth_metadata=user_api_key_dict.metadata,
|
||||
)
|
||||
)
|
||||
_metadata["user_api_key"] = user_api_key_dict.api_key
|
||||
_metadata["status"] = "failure"
|
||||
_metadata["error_information"] = (
|
||||
StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=original_exception,
|
||||
traceback_str=traceback_str,
|
||||
)
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
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
|
||||
|
||||
# Preserve model name and custom_llm_provider
|
||||
if "model" not in request_data:
|
||||
request_data["model"] = existing_litellm_params.get("model") or request_data.get("model", "")
|
||||
if "custom_llm_provider" not in request_data:
|
||||
request_data["custom_llm_provider"] = existing_litellm_params.get(
|
||||
"custom_llm_provider"
|
||||
) or request_data.get("custom_llm_provider", "")
|
||||
|
||||
# Propagate standard_logging_object and litellm_trace_id from the
|
||||
# Logging instance so that _get_session_id_for_spend_log uses the same
|
||||
# trace_id that Langfuse received (via async_failure_handler).
|
||||
# Without this, the DB session_id would be a random UUID that doesn't
|
||||
# match the Langfuse trace_id, making failed requests unsearchable.
|
||||
_litellm_logging_obj = request_data.get("litellm_logging_obj")
|
||||
if _litellm_logging_obj is not None:
|
||||
if "standard_logging_object" not in request_data:
|
||||
request_data["standard_logging_object"] = getattr(
|
||||
_litellm_logging_obj, "model_call_details", {}
|
||||
).get("standard_logging_object")
|
||||
if request_data.get("litellm_trace_id") is None:
|
||||
request_data["litellm_trace_id"] = getattr(
|
||||
_litellm_logging_obj, "litellm_trace_id", None
|
||||
)
|
||||
|
||||
await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
token=user_api_key_dict.api_key,
|
||||
response_cost=0.0,
|
||||
user_id=user_api_key_dict.user_id,
|
||||
end_user_id=user_api_key_dict.end_user_id,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
kwargs=request_data,
|
||||
completion_response=original_exception,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
org_id=user_api_key_dict.org_id,
|
||||
)
|
||||
|
||||
@log_db_metrics
|
||||
async def _PROXY_track_cost_callback(
|
||||
self,
|
||||
kwargs, # kwargs to completion
|
||||
completion_response: Optional[
|
||||
Union[litellm.ModelResponse, Any]
|
||||
], # response from completion
|
||||
start_time=None,
|
||||
end_time=None, # start/end time for completion
|
||||
):
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj, update_cache
|
||||
|
||||
verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback")
|
||||
try:
|
||||
verbose_proxy_logger.debug(
|
||||
f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}"
|
||||
)
|
||||
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs=kwargs)
|
||||
litellm_params = kwargs.get("litellm_params", {}) or {}
|
||||
end_user_id = get_end_user_id_for_cost_tracking(litellm_params)
|
||||
metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs)
|
||||
user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None))
|
||||
team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None))
|
||||
org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None))
|
||||
key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None))
|
||||
end_user_max_budget = metadata.get("user_api_end_user_max_budget", None)
|
||||
sl_object: Optional[StandardLoggingPayload] = kwargs.get(
|
||||
"standard_logging_object", None
|
||||
)
|
||||
response_cost = (
|
||||
sl_object.get("response_cost", None)
|
||||
if sl_object is not None
|
||||
else kwargs.get("response_cost", None)
|
||||
)
|
||||
tags: Optional[List[str]] = (
|
||||
sl_object.get("request_tags", None) if sl_object is not None else None
|
||||
)
|
||||
|
||||
if response_cost is not None:
|
||||
user_api_key = metadata.get("user_api_key", None)
|
||||
if kwargs.get("cache_hit", False) is True:
|
||||
response_cost = 0.0
|
||||
verbose_proxy_logger.debug(
|
||||
f"Cache Hit: response_cost {response_cost}, for user_id {user_id}"
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}"
|
||||
)
|
||||
if _should_track_cost_callback(
|
||||
user_api_key=user_api_key,
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
end_user_id=end_user_id,
|
||||
):
|
||||
## UPDATE DATABASE
|
||||
await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
token=user_api_key,
|
||||
response_cost=response_cost,
|
||||
user_id=user_id,
|
||||
end_user_id=end_user_id,
|
||||
team_id=team_id,
|
||||
kwargs=kwargs,
|
||||
completion_response=completion_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
org_id=org_id,
|
||||
)
|
||||
|
||||
# update cache
|
||||
asyncio.create_task(
|
||||
update_cache(
|
||||
token=user_api_key,
|
||||
user_id=user_id,
|
||||
end_user_id=end_user_id,
|
||||
response_cost=response_cost,
|
||||
team_id=team_id,
|
||||
parent_otel_span=parent_otel_span,
|
||||
tags=tags,
|
||||
)
|
||||
)
|
||||
|
||||
await proxy_logging_obj.slack_alerting_instance.customer_spend_alert(
|
||||
token=user_api_key,
|
||||
key_alias=key_alias,
|
||||
end_user_id=end_user_id,
|
||||
response_cost=response_cost,
|
||||
max_budget=end_user_max_budget,
|
||||
)
|
||||
else:
|
||||
if kwargs.get("stream") is not True or (
|
||||
kwargs.get("stream") is True and "complete_streaming_response" in kwargs
|
||||
):
|
||||
if sl_object is not None:
|
||||
cost_tracking_failure_debug_info: Union[dict, str] = (
|
||||
sl_object["response_cost_failure_debug_info"] # type: ignore
|
||||
or "response_cost_failure_debug_info is None in standard_logging_object"
|
||||
)
|
||||
else:
|
||||
cost_tracking_failure_debug_info = (
|
||||
"standard_logging_object not found"
|
||||
)
|
||||
model = kwargs.get("model")
|
||||
raise Exception(
|
||||
f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
|
||||
)
|
||||
except Exception as e:
|
||||
error_msg = f"Error in tracking cost callback - {str(e)}\n Traceback:{traceback.format_exc()}"
|
||||
model = kwargs.get("model", "")
|
||||
metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs)
|
||||
litellm_metadata = kwargs.get("litellm_params", {}).get(
|
||||
"litellm_metadata", {}
|
||||
)
|
||||
old_metadata = kwargs.get("litellm_params", {}).get("metadata", {})
|
||||
call_type = kwargs.get("call_type", "")
|
||||
error_msg += f"\n Args to _PROXY_track_cost_callback\n model: {model}\n chosen_metadata: {metadata}\n litellm_metadata: {litellm_metadata}\n old_metadata: {old_metadata}\n call_type: {call_type}\n"
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.failed_tracking_alert(
|
||||
error_message=error_msg,
|
||||
failing_model=model,
|
||||
)
|
||||
)
|
||||
|
||||
verbose_proxy_logger.exception(
|
||||
"Error in tracking cost callback - %s", str(e)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _should_track_errors_in_db():
|
||||
"""
|
||||
Returns True if errors should be tracked in the database
|
||||
|
||||
By default, errors are tracked in the database
|
||||
|
||||
If users want to disable error tracking, they can set the disable_error_logs flag in the general_settings
|
||||
"""
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
if general_settings.get("disable_error_logs") is True:
|
||||
return False
|
||||
return
|
||||
|
||||
|
||||
def _should_track_cost_callback(
|
||||
user_api_key: Optional[str],
|
||||
user_id: Optional[str],
|
||||
team_id: Optional[str],
|
||||
end_user_id: Optional[str],
|
||||
) -> bool:
|
||||
"""
|
||||
Determine if the cost callback should be tracked based on the kwargs
|
||||
"""
|
||||
|
||||
# don't run track cost callback if user opted into disabling spend
|
||||
if ProxyUpdateSpend.disable_spend_updates() is True:
|
||||
return False
|
||||
|
||||
if (
|
||||
user_api_key is not None
|
||||
or user_id is not None
|
||||
or team_id is not None
|
||||
or end_user_id is not None
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -126,3 +126,356 @@ async def test_async_post_call_failure_hook_non_llm_route():
|
|||
|
||||
# Assert that update_database was NOT called for non-LLM routes
|
||||
mock_update_database.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_track_cost_callback_skips_when_no_standard_logging_object():
|
||||
"""
|
||||
Reproduces the bug where _PROXY_track_cost_callback raises
|
||||
'Cost tracking failed for model=None' when kwargs has no
|
||||
standard_logging_object (e.g. call_type=afile_delete).
|
||||
|
||||
File operations have no model and no standard_logging_object.
|
||||
The callback should skip gracefully instead of raising.
|
||||
"""
|
||||
logger = _ProxyDBLogger()
|
||||
|
||||
kwargs = {
|
||||
"call_type": "afile_delete",
|
||||
"model": None,
|
||||
"litellm_call_id": "test-call-id",
|
||||
"litellm_params": {},
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj",
|
||||
) as mock_proxy_logging:
|
||||
mock_proxy_logging.failed_tracking_alert = AsyncMock()
|
||||
mock_proxy_logging.db_spend_update_writer = MagicMock()
|
||||
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
|
||||
|
||||
await logger._PROXY_track_cost_callback(
|
||||
kwargs=kwargs,
|
||||
completion_response=None,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
# update_database should NOT be called — nothing to track
|
||||
mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called()
|
||||
|
||||
# failed_tracking_alert should NOT be called — this is not an error
|
||||
mock_proxy_logging.failed_tracking_alert.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj():
|
||||
"""
|
||||
When an LLM call fails, the proxy calls post_call_failure_hook with
|
||||
request_data that doesn't contain standard_logging_object. But the
|
||||
litellm_logging_obj (set by function_setup) is in request_data and
|
||||
holds the standard_logging_object with the correct trace_id.
|
||||
|
||||
The failure hook should propagate this so the DB spend log's session_id
|
||||
matches the Langfuse trace_id.
|
||||
"""
|
||||
logger = _ProxyDBLogger()
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test_api_key",
|
||||
user_id="test_user_id",
|
||||
team_id="test_team_id",
|
||||
)
|
||||
|
||||
# Simulate a litellm_logging_obj with model_call_details containing
|
||||
# the standard_logging_object (as set by _failure_handler_helper_fn)
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.litellm_trace_id = "trace-id-from-logging-obj"
|
||||
mock_logging_obj.model_call_details = {
|
||||
"standard_logging_object": {
|
||||
"trace_id": "trace-id-from-logging-obj",
|
||||
"error_str": "InternalServerError",
|
||||
}
|
||||
}
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"metadata": {},
|
||||
"litellm_params": {},
|
||||
"litellm_logging_obj": mock_logging_obj,
|
||||
# Note: no "standard_logging_object" and no "litellm_trace_id"
|
||||
}
|
||||
|
||||
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("Provider error"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
mock_update_database.assert_called_once()
|
||||
call_kwargs = mock_update_database.call_args[1]["kwargs"]
|
||||
|
||||
# standard_logging_object should have been propagated from logging obj
|
||||
assert call_kwargs.get("standard_logging_object") is not None
|
||||
assert (
|
||||
call_kwargs["standard_logging_object"]["trace_id"]
|
||||
== "trace-id-from-logging-obj"
|
||||
)
|
||||
# litellm_trace_id should also be propagated as a fallback
|
||||
assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enrich_failure_metadata_with_team_alias():
|
||||
"""
|
||||
When team_id is set but team_alias is missing (and key_alias is present),
|
||||
_enrich_failure_metadata_with_key_info should look up the team from cache
|
||||
and populate user_api_key_team_alias.
|
||||
"""
|
||||
mock_team_obj = MagicMock()
|
||||
mock_team_obj.team_alias = "my-team-alias"
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_team_obj,
|
||||
):
|
||||
metadata = {
|
||||
"user_api_key": "hashed_key",
|
||||
"user_api_key_alias": "my-key-alias", # already set
|
||||
"user_api_key_team_id": "test_team_id",
|
||||
"user_api_key_team_alias": None,
|
||||
}
|
||||
result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
|
||||
assert result["user_api_key_team_alias"] == "my-team-alias"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enrich_failure_metadata_with_full_key_lookup():
|
||||
"""
|
||||
When all key fields are null (auth error 401 scenario), _enrich_failure_metadata_with_key_info
|
||||
should look up the key object from cache/DB and populate alias, user_id, team_id,
|
||||
then look up the team to get team_alias.
|
||||
"""
|
||||
mock_key_obj = MagicMock()
|
||||
mock_key_obj.key_alias = "fetched-key-alias"
|
||||
mock_key_obj.user_id = "fetched-user-id"
|
||||
mock_key_obj.team_id = "fetched-team-id"
|
||||
mock_key_obj.org_id = "fetched-org-id"
|
||||
|
||||
mock_team_obj = MagicMock()
|
||||
mock_team_obj.team_alias = "fetched-team-alias"
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_key_obj,
|
||||
), patch(
|
||||
"litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_team_obj,
|
||||
):
|
||||
metadata = {
|
||||
"user_api_key": "hashed_key",
|
||||
"user_api_key_alias": None, # all null - simulates auth error path
|
||||
"user_api_key_user_id": None,
|
||||
"user_api_key_team_id": None,
|
||||
"user_api_key_team_alias": None,
|
||||
"user_api_key_org_id": None,
|
||||
}
|
||||
result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
|
||||
assert result["user_api_key_alias"] == "fetched-key-alias"
|
||||
assert result["user_api_key_user_id"] == "fetched-user-id"
|
||||
assert result["user_api_key_team_id"] == "fetched-team-id"
|
||||
assert result["user_api_key_org_id"] == "fetched-org-id"
|
||||
assert result["user_api_key_team_alias"] == "fetched-team-alias"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enrich_failure_metadata_skips_when_team_alias_present():
|
||||
"""
|
||||
When team_alias is already populated, _enrich_failure_metadata_with_key_info
|
||||
should not perform a team cache lookup.
|
||||
"""
|
||||
with patch(
|
||||
"litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get_key, patch(
|
||||
"litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get_team:
|
||||
metadata = {
|
||||
"user_api_key": "hashed_key",
|
||||
"user_api_key_alias": "existing-alias",
|
||||
"user_api_key_team_id": "test_team_id",
|
||||
"user_api_key_team_alias": "already-set",
|
||||
}
|
||||
result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
|
||||
assert result["user_api_key_team_alias"] == "already-set"
|
||||
mock_get_key.assert_not_called()
|
||||
mock_get_team.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enrich_failure_metadata_skips_when_no_api_key():
|
||||
"""
|
||||
When api_key hash is absent, _enrich_failure_metadata_with_key_info should
|
||||
not perform any lookups.
|
||||
"""
|
||||
with patch(
|
||||
"litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get_key:
|
||||
metadata = {
|
||||
"user_api_key": None,
|
||||
"user_api_key_alias": None,
|
||||
"user_api_key_team_id": None,
|
||||
"user_api_key_team_alias": None,
|
||||
}
|
||||
result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
|
||||
mock_get_key.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_failure_hook_enriches_auth_error_metadata():
|
||||
"""
|
||||
Simulates a 401 ProxyException (e.g. can_key_call_model). In this case
|
||||
UserAPIKeyAuth is created with only api_key set. The failure hook should
|
||||
look up the key and team from cache/DB to populate all missing fields.
|
||||
"""
|
||||
logger = _ProxyDBLogger()
|
||||
|
||||
# This is what auth_exception_handler creates for 401 errors
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="hashed_key",
|
||||
# key_alias, user_id, team_id, team_alias are all None
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"model": "claude-haiku-4-5",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"metadata": {},
|
||||
"litellm_params": {},
|
||||
}
|
||||
|
||||
mock_key_obj = MagicMock()
|
||||
mock_key_obj.key_alias = "my-key-alias"
|
||||
mock_key_obj.user_id = "my-user-id"
|
||||
mock_key_obj.team_id = "my-team-id"
|
||||
mock_key_obj.org_id = None
|
||||
|
||||
mock_team_obj = MagicMock()
|
||||
mock_team_obj.team_alias = "my-team-alias"
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_update_database, patch(
|
||||
"litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_key_obj,
|
||||
), patch(
|
||||
"litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_team_obj,
|
||||
):
|
||||
await logger.async_post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=Exception("401 - model not allowed"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
mock_update_database.assert_called_once()
|
||||
call_args = mock_update_database.call_args[1]
|
||||
metadata = call_args["kwargs"]["litellm_params"]["metadata"]
|
||||
assert metadata["user_api_key_alias"] == "my-key-alias"
|
||||
assert metadata["user_api_key_user_id"] == "my-user-id"
|
||||
assert metadata["user_api_key_team_id"] == "my-team-id"
|
||||
assert metadata["user_api_key_team_alias"] == "my-team-alias"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_failure_hook_enriches_missing_team_alias():
|
||||
"""
|
||||
When user_api_key_dict has a team_id but no team_alias, async_post_call_failure_hook
|
||||
should look up the team from cache and populate user_api_key_team_alias in the
|
||||
spend log metadata written to the DB.
|
||||
"""
|
||||
logger = _ProxyDBLogger()
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test_api_key",
|
||||
key_alias="test_alias",
|
||||
user_id="test_user_id",
|
||||
team_id="test_team_id",
|
||||
team_alias=None, # Missing - simulates regular key auth where SQL view omits team_alias
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"metadata": {},
|
||||
"litellm_params": {},
|
||||
}
|
||||
|
||||
mock_team_obj = MagicMock()
|
||||
mock_team_obj.team_alias = "enriched-team-alias"
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_update_database, patch(
|
||||
"litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_team_obj,
|
||||
):
|
||||
await logger.async_post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=Exception("Provider rate limit"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
mock_update_database.assert_called_once()
|
||||
call_args = mock_update_database.call_args[1]
|
||||
metadata = call_args["kwargs"]["litellm_params"]["metadata"]
|
||||
assert metadata["user_api_key_team_alias"] == "enriched-team-alias"
|
||||
assert metadata["user_api_key_team_id"] == "test_team_id"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_value", [None, ""])
|
||||
async def test_track_cost_callback_skips_for_falsy_model_and_no_slo(model_value):
|
||||
"""
|
||||
Same bug as above but model can also be empty string (e.g. health check callbacks).
|
||||
The guard should catch all falsy model values when sl_object is missing.
|
||||
"""
|
||||
logger = _ProxyDBLogger()
|
||||
|
||||
kwargs = {
|
||||
"call_type": "acompletion",
|
||||
"model": model_value,
|
||||
"litellm_params": {},
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj",
|
||||
) as mock_proxy_logging:
|
||||
mock_proxy_logging.failed_tracking_alert = AsyncMock()
|
||||
mock_proxy_logging.db_spend_update_writer = MagicMock()
|
||||
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
|
||||
|
||||
await logger._PROXY_track_cost_callback(
|
||||
kwargs=kwargs,
|
||||
completion_response=None,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
mock_proxy_logging.failed_tracking_alert.assert_not_called()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue