import asyncio import traceback from collections.abc import Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED 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.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( get_key_object, get_team_object, log_db_metrics, ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.db_spend_update_writer import ( debitable_model_access_groups, get_llm_router, ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.spend_log_error_logger import ( should_suppress_spend_log_tracebacks, spend_log_error, ) from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, get_request_model_access_groups, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( CallTypes, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: from litellm.proxy.utils import ProxyLogging _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( { CallTypes.pass_through.value, CallTypes.llm_passthrough_route.value, CallTypes.allm_passthrough_route.value, # CheckBatchCost's synthetic logging_obj for a completed managed batch carries # whatever LiteLLM_ManagedObjectTable stored at create time, and all of it is # None for a batch created before those columns were persisted, or by the master # key. The batch already incurred real provider cost, so track it regardless. CallTypes.aretrieve_batch.value, } ) # Both spellings, because call_type reaches the callback as str(...) of either the # enum member or its value. _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( ( CallTypes.aretrieve_batch.value, str(CallTypes.aretrieve_batch), ) ) 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: str | None = None, ): try: await _release_budget_reservation(budget_reservation=user_api_key_dict.budget_reservation) except Exception: verbose_proxy_logger.exception("Failed to release budget reservation during failure handling") try: await _invalidate_budget_reservation_counters(budget_reservation=user_api_key_dict.budget_reservation) if user_api_key_dict.budget_reservation is not None: user_api_key_dict.budget_reservation["finalized"] = True except Exception: verbose_proxy_logger.exception( "Failed to invalidate budget reservation counters after failure release failed" ) request_route: Final = user_api_key_dict.request_route if ( _ProxyDBLogger._should_track_errors_in_db() is False or request_route is not None and not ( RouteChecks.is_llm_api_route(route=request_route) or RouteChecks.is_info_route(route=request_route) ) ): return from litellm.proxy.proxy_server import proxy_logging_obj _metadata = dict( LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) ) _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["status"] = "failure" _error_information = StandardLoggingPayloadSetup.get_error_information( original_exception=original_exception, traceback_str=traceback_str, ) if should_suppress_spend_log_tracebacks(): # Drop the traceback key entirely so the per-row Metadata pane in # the UI (which renders the JSON blob verbatim) doesn't show a # noisy ``"traceback": ""`` line. Downstream consumers all use # ``.get("traceback")`` / truthy checks, and the TypedDict marks # the field as optional, so omitting is type-safe. _error_information.pop("traceback", None) # Strip echoed request input + apply DB-size cap before storing in # the spend-log metadata column (LIT-2992). Result is never None # here because the input above is constructed non-None. _error_information = cast( StandardLoggingPayloadErrorInformation, _sanitize_error_information_for_spend_logs(_error_information), ) _metadata["error_information"] = _error_information _metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( metadata=_metadata, ) existing_metadata: Final[dict] = request_data.get("metadata", None) or {} existing_metadata.update(_metadata) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") if ( isinstance(litellm_metadata_bucket, dict) and "standard_logging_guardrail_information" not in existing_metadata ): guardrail_info: Final = litellm_metadata_bucket.get("standard_logging_guardrail_information") if guardrail_info is not None: existing_metadata["standard_logging_guardrail_information"] = guardrail_info if "litellm_params" not in request_data: request_data["litellm_params"] = {} existing_litellm_params: Final = request_data.get("litellm_params", {}) existing_litellm_metadata: Final = 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: Final = request_data.get("litellm_logging_obj") if _litellm_logging_obj is not None: if not request_data.get("standard_logging_object"): 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) # Use the actual request start time from the logging object so that # failed requests record the real duration instead of 0. actual_start_time = datetime.now() if _litellm_logging_obj is not None: obj_start: Final = getattr(_litellm_logging_obj, "start_time", None) if obj_start is not None: actual_start_time = obj_start # A stream that broke mid-flight still billed the provider for the # chunks already delivered. ``post_call_failure_hook`` lifts that # recovered cost onto request_data (the usage rides along in # ``combined_usage_object`` for the token columns), so attribute the # real partial spend to this failure row instead of zero. recovered_stream_cost: Final = ( max(float(request_data.get("response_cost") or 0.0), 0.0) if isinstance(request_data.get("combined_usage_object"), litellm.Usage) else 0.0 ) recovered_response_cost: Final = recovered_stream_cost + guardrail_information_cost( existing_metadata.get("standard_logging_guardrail_information") ) await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key_dict.api_key, response_cost=recovered_response_cost, 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=actual_start_time, 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: litellm.ModelResponse | Any | None, # response from completion start_time=None, end_time=None, # start/end time for completion ): from litellm.proxy.proxy_server import ( increment_spend_counters, proxy_logging_obj, update_cache, ) verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") try: verbose_proxy_logger.debug( "kwargs stream: %s + complete streaming response: %s", kwargs.get("stream", None), kwargs.get("complete_streaming_response", None), ) parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs=kwargs) litellm_params: Final = kwargs.get("litellm_params", {}) or {} end_user_id: Final = get_end_user_id_for_cost_tracking(litellm_params) metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) # Only fetch key details when user_id wasn't already populated (e.g. direct MCP REST calls). # Avoids a cache/DB lookup on every normal LLM request. if metadata.get("user_api_key") and not metadata.get("user_api_key_user_id"): metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( # rebind-ok: enriched metadata replaces the original metadata=metadata, resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES, ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None)) end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None) sl_object: Final[StandardLoggingPayload | None] = 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: Final = _get_request_tags_for_cost_tracking( sl_object=sl_object, metadata=metadata, ) model_access_groups: Final = debitable_model_access_groups( attributed=get_request_model_access_groups(kwargs), served_model_id=sl_object.get("model_id") if sl_object is not None else None, router=get_llm_router(), ) if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) if kwargs.get("cache_hit", False) is True: response_cost = 0.0 verbose_proxy_logger.debug("Cache Hit: response_cost %s, for user_id %s", response_cost, user_id) verbose_proxy_logger.debug( "user_api_key %s, user_id %s, team_id %s, end_user_id %s", user_api_key, user_id, team_id, end_user_id, ) call_type: str | None = kwargs.get("call_type") 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, call_type=call_type, ): ## UPDATE DATABASE await _update_database_and_spend_counters( proxy_logging_obj=proxy_logging_obj, increment_spend_counters=increment_spend_counters, user_api_key=user_api_key, user_id=user_id, end_user_id=end_user_id, team_id=team_id, org_id=org_id, kwargs=kwargs, completion_response=completion_response, start_time=start_time, end_time=end_time, response_cost=response_cost, budget_reservation=budget_reservation, request_tags=tags, model_access_groups=model_access_groups, ) # update cache (fire-and-forget for backward compat: # cached object fields, soft budget alerts, etc.) 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, ) elif budget_reservation is not None: await _release_budget_reservation(budget_reservation=budget_reservation) else: if _is_unbilled_interaction_response(completion_response): if BACKGROUND_INTERACTION_COST_POLLING_ENABLED and _is_unbilled_in_progress_interaction( completion_response ): verbose_proxy_logger.debug( "Cost tracking deferred for in-progress background interaction; " "the budget reservation stays open until the poll task logs the final usage" ) return await _release_budget_reservation(budget_reservation=budget_reservation) verbose_proxy_logger.debug( "Released the budget reservation for an interaction create with no usage " "that no poll task will settle" ) return await _release_budget_reservation(budget_reservation=budget_reservation) # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. # Use .get() for "stream" to avoid KeyError on health checks. # WS session wrappers (_aresponses_websocket, _arealtime) also reach here with # result=None; their per-turn costs are tracked on the inner aresponses/realtime calls. if sl_object is None and ( not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime") ): verbose_proxy_logger.warning( "Cost tracking - skipping, no standard_logging_object for call_type=%s", kwargs.get("call_type", "unknown"), ) return 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: dict | str = ( sl_object["response_cost_failure_debug_info"] 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 - {e}\n Traceback:{traceback.format_exc()}" model = kwargs.get("model", "") metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) litellm_metadata: Final = kwargs.get("litellm_params", {}).get("litellm_metadata", {}) old_metadata: Final = 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, ) ) spend_log_error("Error in tracking cost callback - %s", str(e), exc=e) @staticmethod async def _enrich_failure_metadata_with_key_info(metadata: dict, resolve_missing_key_identity: bool = True) -> dict: """ Enriches failure spend log metadata by looking up the key object (and team object) from cache/DB when key fields are missing. This handles two scenarios: 1. Auth errors (401): UserAPIKeyAuth is created with only api_key set, all other fields are null. We look up the full key object to fill in alias, user_id, team_id, etc. 2. Post-auth failures (provider errors, rate limits): key fields are populated but team_alias is missing because LiteLLM_VerificationTokenView SQL view doesn't include it. We look up the team object to fill in team_alias. Scenario 1 reads the key's identity as it stands right now, so it is only correct for a log emitted within the request it describes. Callers that log after a delay, against an identity captured earlier, pass resolve_missing_key_identity=False and keep their own user_id, team_id and org_id. """ api_key_hash: Final = metadata.get("user_api_key") if not api_key_hash: return metadata from litellm.proxy.proxy_server import ( prisma_client, proxy_logging_obj, user_api_key_cache, ) # Step 1: If key fields are missing, look up the full key object if resolve_missing_key_identity and metadata.get("user_api_key_alias") is None: try: key_obj: Final = await get_key_object( hashed_token=api_key_hash, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) if metadata.get("user_api_key_alias") is None: metadata["user_api_key_alias"] = key_obj.key_alias if metadata.get("user_api_key_user_id") is None: metadata["user_api_key_user_id"] = key_obj.user_id if metadata.get("user_api_key_team_id") is None: metadata["user_api_key_team_id"] = key_obj.team_id if metadata.get("user_api_key_org_id") is None: metadata["user_api_key_org_id"] = key_obj.org_id except Exception: verbose_proxy_logger.debug( "Failed to enrich failure metadata with key info for api_key=%s", api_key_hash, ) # Step 2: If team_id is known but team_alias is missing, look up the team object team_id: Final = metadata.get("user_api_key_team_id") if team_id and metadata.get("user_api_key_team_alias") is None: try: team_obj: Final = await get_team_object( team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) if team_obj.team_alias is not None: metadata["user_api_key_team_alias"] = team_obj.team_alias except Exception: verbose_proxy_logger.debug( "Failed to enrich failure metadata with team_alias for team_id=%s", team_id, ) return metadata @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 _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: patch = {k: v for k, v in metadata.items() if (k.startswith("user_api_key") or k == "tags") and v is not None} if not patch: return litellm_params: Final = kwargs.setdefault("litellm_params", {}) for bucket_name in ("litellm_metadata", "metadata"): bucket = litellm_params.get(bucket_name) if isinstance(bucket, dict): for key, value in patch.items(): if bucket.get(key) is None: bucket[key] = value def _is_unbilled_interaction_response(completion_response: object) -> bool: from litellm.interactions.background_cost_polling import missing_usage_is_expected from litellm.types.interactions import InteractionsAPIResponse if not isinstance(completion_response, InteractionsAPIResponse): return False return completion_response.usage is None and missing_usage_is_expected(completion_response) def _is_unbilled_in_progress_interaction(completion_response: object) -> bool: from litellm.interactions.background_cost_polling import is_pollable_background_interaction from litellm.types.interactions import InteractionsAPIResponse if not isinstance(completion_response, InteractionsAPIResponse): return False return completion_response.usage is None and is_pollable_background_interaction(completion_response) def _should_track_cost_callback( user_api_key: str | None, user_id: str | None, team_id: str | None, end_user_id: str | None, call_type: str | None = None, ) -> bool: """ Determine if the cost callback should be tracked based on the kwargs Pass-through endpoints can be configured with ``auth=false``, which leaves the request with no key/user/team/end-user to attribute spend to. Those requests still forward real provider traffic that operators expect to see in request/usage logs, so they are tracked even when unauthenticated. The same reasoning applies to a completed managed batch's cost event (see _UNATTRIBUTED_TRACKABLE_CALL_TYPES). """ # 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 call_type in _UNATTRIBUTED_TRACKABLE_CALL_TYPES def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: metadata_budget_reservation: Final = metadata.get("user_api_key_budget_reservation") if isinstance(metadata_budget_reservation, dict): return metadata_budget_reservation user_api_key_auth_obj: Final = metadata.get("user_api_key_auth") if user_api_key_auth_obj is None: return None if isinstance(user_api_key_auth_obj, dict): budget_reservation: Final = user_api_key_auth_obj.get("budget_reservation") return budget_reservation if isinstance(budget_reservation, dict) else None return getattr(user_api_key_auth_obj, "budget_reservation", None) def _get_request_tags_for_cost_tracking( sl_object: StandardLoggingPayload | None, metadata: dict, ) -> list[str] | None: if sl_object is not None: request_tags: Final = sl_object.get("request_tags", None) if isinstance(request_tags, list): return request_tags metadata_tags: Final = metadata.get("tags", None) if isinstance(metadata_tags, list): return metadata_tags return None async def _update_database_and_spend_counters( proxy_logging_obj: "ProxyLogging", increment_spend_counters: Any, user_api_key: str | None, user_id: str | None, end_user_id: str | None, team_id: str | None, org_id: str | None, kwargs: dict, completion_response: litellm.ModelResponse | Any | None, start_time: Any, end_time: Any, response_cost: float, budget_reservation: dict | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, ) -> None: try: spend_log_request_id = 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, ) except Exception: if budget_reservation is not None: try: await _release_budget_reservation(budget_reservation=budget_reservation) except Exception: verbose_proxy_logger.exception("Failed to release budget reservation after database update failed") try: await _invalidate_budget_reservation_counters(budget_reservation=budget_reservation) except Exception: verbose_proxy_logger.exception( "Failed to invalidate budget reservation counters after release failed" ) raise try: await increment_spend_counters( token=user_api_key, team_id=team_id, user_id=user_id, response_cost=response_cost, org_id=org_id, budget_reservation=budget_reservation, end_user_id=end_user_id, tags=request_tags, request_id=spend_log_request_id, request_started_at=start_time, model_access_groups=model_access_groups, ) except Exception: if budget_reservation is not None: try: await _invalidate_budget_reservation_counters(budget_reservation=budget_reservation) except Exception: verbose_proxy_logger.exception( "Failed to invalidate budget reservation counters after spend counter update failed" ) finally: budget_reservation["finalized"] = True raise async def _release_budget_reservation(budget_reservation: dict | None) -> None: if budget_reservation is None: return from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation, ) await release_budget_reservation( budget_reservation=budget_reservation, ) async def _invalidate_budget_reservation_counters( budget_reservation: dict | None, ) -> None: if budget_reservation is None: return from litellm.proxy.spend_tracking.budget_reservation import ( invalidate_budget_reservation_counters, ) await invalidate_budget_reservation_counters( budget_reservation=budget_reservation, )