Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/hooks/proxy_track_cost_callback.py: 56%
313 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import asyncio
2import traceback
3from collections.abc import Callable, Mapping, Sequence
4from datetime import datetime
5from typing import TYPE_CHECKING, Any, Final, Protocol, cast
7import litellm
8from litellm._logging import verbose_proxy_logger
9from litellm.batches.batch_utils import batch_cost_is_final
10from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED
11from litellm.integrations.custom_logger import CustomLogger
12from litellm.litellm_core_utils.core_helpers import (
13 _get_parent_otel_span_from_kwargs,
14 budget_reservation_from_metadata,
15 get_litellm_metadata_from_kwargs,
16)
17from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
18from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
19from litellm.proxy._types import UserAPIKeyAuth
20from litellm.proxy.auth.auth_checks import (
21 get_key_object,
22 get_team_object,
23 log_db_metrics,
24)
25from litellm.proxy.auth.route_checks import RouteChecks
26from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded
27from litellm.proxy.db.db_spend_update_writer import (
28 DBSpendUpdateWriter,
29 debitable_model_access_groups,
30 get_llm_router,
31)
32from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
33from litellm.proxy.spend_tracking.spend_event import (
34 ObjectMapping,
35 SpendEventBuildError,
36 SpendEventDecodeError,
37 build_spend_event,
38 decode_spend_event,
39 is_offloadable_success,
40 spend_event_callback_args,
41)
42from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer
43from litellm.proxy.spend_tracking.spend_log_error_logger import (
44 should_suppress_spend_log_tracebacks,
45 spend_log_error,
46)
47from litellm.proxy.spend_tracking.spend_tracking_utils import (
48 _sanitize_error_information_for_spend_logs,
49 get_request_model_access_groups,
50 should_store_prompts_and_responses_in_spend_logs,
51)
52from litellm.proxy.utils import ProxyUpdateSpend
53from litellm.types.utils import (
54 CallTypes,
55 LiteLLMBatch,
56 StandardLoggingPayload,
57 StandardLoggingPayloadErrorInformation,
58)
59from litellm.utils import get_end_user_id_for_cost_tracking
61if TYPE_CHECKING: 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true
62 from litellm.proxy.utils import ProxyLogging
64_UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset(
65 {
66 CallTypes.pass_through.value,
67 CallTypes.llm_passthrough_route.value,
68 CallTypes.allm_passthrough_route.value,
69 CallTypes.call_mcp_tool.value,
70 # CheckBatchCost's synthetic logging_obj for a completed managed batch carries
71 # whatever LiteLLM_ManagedObjectTable stored at create time, and all of it is
72 # None for a batch created before those columns were persisted, or by the master
73 # key. The batch already incurred real provider cost, so track it regardless.
74 CallTypes.aretrieve_batch.value,
75 }
76)
78# Both spellings, because call_type reaches the callback as str(...) of either the
79# enum member or its value.
80_CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset(
81 (
82 CallTypes.aretrieve_batch.value,
83 str(CallTypes.aretrieve_batch),
84 )
85)
88def _proxy_spend_writer() -> DBSpendUpdateWriter:
89 from litellm.proxy.proxy_server import proxy_logging_obj
91 return proxy_logging_obj.db_spend_update_writer
94class _ProxyDBLogger(CustomLogger):
95 def __init__(
96 self,
97 spend_event_producer: SpendEventProducer | None = None,
98 *,
99 turn_off_message_logging: bool = False,
100 message_logging: bool = True,
101 spend_writer: Callable[[], DBSpendUpdateWriter] = _proxy_spend_writer,
102 ) -> None:
103 super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging)
104 self.spend_event_producer = spend_event_producer
105 self._spend_writer: Final = spend_writer
107 async def async_log_success_event(
108 self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime
109 ) -> None:
110 if self.spend_event_producer is None or not is_offloadable_success(response_obj): 110 ↛ 113line 110 didn't jump to line 113 because the condition on line 110 was always true
111 await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time)
112 return
113 event: Final = build_spend_event(
114 kwargs,
115 response_obj,
116 start_time,
117 end_time,
118 store_bodies=should_store_prompts_and_responses_in_spend_logs(),
119 )
120 if isinstance(event, SpendEventBuildError):
121 verbose_proxy_logger.warning("collector: tracking cost in-process, event not buildable: %s", event.reason)
122 await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time)
123 return
124 await self.spend_event_producer.publish(event)
126 async def run_spend_event(self, line: bytes) -> None:
127 """Run the unchanged cost pipeline on a serialized spend event (sidecar consumer and in-process fallback)."""
128 event: Final = decode_spend_event(line)
129 if isinstance(event, SpendEventDecodeError):
130 verbose_proxy_logger.error("collector: discarding undecodable spend event: %s", event.reason)
131 return
132 args: Final = spend_event_callback_args(event)
133 await self._PROXY_track_cost_callback(args.kwargs, args.response_obj, args.start_time, args.end_time)
135 async def async_post_call_failure_hook(
136 self,
137 request_data: dict,
138 original_exception: Exception,
139 user_api_key_dict: UserAPIKeyAuth,
140 traceback_str: str | None = None,
141 ):
142 try:
143 await _release_budget_reservation(budget_reservation=user_api_key_dict.budget_reservation)
144 except Exception:
145 verbose_proxy_logger.exception("Failed to release budget reservation during failure handling")
146 try:
147 await _invalidate_budget_reservation_counters(budget_reservation=user_api_key_dict.budget_reservation)
148 if user_api_key_dict.budget_reservation is not None:
149 user_api_key_dict.budget_reservation["finalized"] = True
150 except Exception:
151 verbose_proxy_logger.exception(
152 "Failed to invalidate budget reservation counters after failure release failed"
153 )
155 request_route: Final = user_api_key_dict.request_route
156 if (
157 _ProxyDBLogger._should_track_errors_in_db() is False
158 or request_route is not None
159 and not (
160 RouteChecks.is_llm_api_route(route=request_route) or RouteChecks.is_info_route(route=request_route)
161 )
162 ):
163 return
165 _metadata = dict(
166 LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict)
167 )
168 _metadata["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict)
169 _metadata["status"] = "failure"
170 _error_information = StandardLoggingPayloadSetup.get_error_information(
171 original_exception=original_exception,
172 traceback_str=traceback_str,
173 )
174 if should_suppress_spend_log_tracebacks(): 174 ↛ 180line 174 didn't jump to line 180 because the condition on line 174 was never true
175 # Drop the traceback key entirely so the per-row Metadata pane in
176 # the UI (which renders the JSON blob verbatim) doesn't show a
177 # noisy ``"traceback": ""`` line. Downstream consumers all use
178 # ``.get("traceback")`` / truthy checks, and the TypedDict marks
179 # the field as optional, so omitting is type-safe.
180 _error_information.pop("traceback", None)
181 # Strip echoed request input + apply DB-size cap before storing in
182 # the spend-log metadata column (LIT-2992). Result is never None
183 # here because the input above is constructed non-None.
184 _error_information = cast(
185 StandardLoggingPayloadErrorInformation,
186 _sanitize_error_information_for_spend_logs(_error_information, original_exception=original_exception),
187 )
188 _metadata["error_information"] = _error_information
190 _metadata = await _ProxyDBLogger._enrich_failure_metadata_unless_db_stalled(
191 metadata=_metadata, original_exception=original_exception
192 )
194 existing_metadata: Final[dict] = request_data.get("metadata", None) or {}
195 existing_metadata.update(_metadata)
197 litellm_metadata_bucket: Final = request_data.get("litellm_metadata")
198 if (
199 isinstance(litellm_metadata_bucket, dict)
200 and "standard_logging_guardrail_information" not in existing_metadata
201 ):
202 guardrail_info: Final = litellm_metadata_bucket.get("standard_logging_guardrail_information")
203 if guardrail_info is not None: 203 ↛ 204line 203 didn't jump to line 204 because the condition on line 203 was never true
204 existing_metadata["standard_logging_guardrail_information"] = guardrail_info
206 if "litellm_params" not in request_data:
207 request_data["litellm_params"] = {}
209 existing_litellm_params: Final = request_data.get("litellm_params", {})
210 existing_litellm_metadata: Final = existing_litellm_params.get("metadata", {}) or {}
212 # Preserve tags from existing metadata
213 if existing_litellm_metadata.get("tags"): 213 ↛ 214line 213 didn't jump to line 214 because the condition on line 213 was never true
214 existing_metadata["tags"] = existing_litellm_metadata.get("tags")
216 request_data["litellm_params"]["proxy_server_request"] = (
217 request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {}
218 )
219 request_data["litellm_params"]["metadata"] = existing_metadata
221 # Preserve model name and custom_llm_provider
222 if "model" not in request_data:
223 request_data["model"] = existing_litellm_params.get("model") or request_data.get("model", "")
224 if "custom_llm_provider" not in request_data:
225 request_data["custom_llm_provider"] = existing_litellm_params.get(
226 "custom_llm_provider"
227 ) or request_data.get("custom_llm_provider", "")
229 # Propagate standard_logging_object and litellm_trace_id from the Logging
230 # instance so the failure row carries the same trace_id Langfuse received.
231 _litellm_logging_obj: Final = request_data.get("litellm_logging_obj")
232 if _litellm_logging_obj is not None: 232 ↛ 233line 232 didn't jump to line 233 because the condition on line 232 was never true
233 if not request_data.get("standard_logging_object"):
234 request_data["standard_logging_object"] = getattr(_litellm_logging_obj, "model_call_details", {}).get(
235 "standard_logging_object"
236 )
237 if request_data.get("litellm_trace_id") is None:
238 request_data["litellm_trace_id"] = getattr(_litellm_logging_obj, "litellm_trace_id", None)
240 lifted_start_time: Final = request_data.get("start_time")
241 actual_start_time: Final = (
242 lifted_start_time
243 if isinstance(lifted_start_time, datetime)
244 else getattr(_litellm_logging_obj, "start_time", None) or datetime.now()
245 )
247 # A stream that broke mid-flight still billed the provider for the
248 # chunks already delivered. ``post_call_failure_hook`` lifts that
249 # recovered cost onto request_data (the usage rides along in
250 # ``combined_usage_object`` for the token columns), so attribute the
251 # real partial spend to this failure row instead of zero.
252 recovered_stream_cost: Final = (
253 max(float(request_data.get("response_cost") or 0.0), 0.0)
254 if isinstance(request_data.get("combined_usage_object"), litellm.Usage)
255 else 0.0
256 )
257 recovered_response_cost: Final = recovered_stream_cost + guardrail_information_cost(
258 existing_metadata.get("standard_logging_guardrail_information")
259 )
261 await self._spend_writer().update_database(
262 token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict),
263 response_cost=recovered_response_cost,
264 user_id=user_api_key_dict.user_id,
265 end_user_id=user_api_key_dict.end_user_id,
266 team_id=user_api_key_dict.team_id,
267 kwargs=request_data,
268 completion_response=original_exception,
269 start_time=actual_start_time,
270 end_time=datetime.now(),
271 org_id=user_api_key_dict.org_id,
272 project_id=user_api_key_dict.project_id,
273 )
275 @log_db_metrics
276 async def _PROXY_track_cost_callback(
277 self,
278 kwargs, # kwargs to completion
279 completion_response: litellm.ModelResponse | Any | None, # response from completion
280 start_time=None,
281 end_time=None, # start/end time for completion
282 ):
283 from litellm.proxy.proxy_server import (
284 increment_spend_counters,
285 proxy_logging_obj,
286 update_cache,
287 )
289 verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback")
290 try:
291 verbose_proxy_logger.debug(
292 "kwargs stream: %s + complete streaming response: %s",
293 kwargs.get("stream", None),
294 kwargs.get("complete_streaming_response", None),
295 )
296 parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs=kwargs)
297 litellm_params: Final = kwargs.get("litellm_params", {}) or {}
298 end_user_id: Final = get_end_user_id_for_cost_tracking(litellm_params)
299 metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs)
300 # Only fetch key details when user_id wasn't already populated (e.g. direct MCP REST calls).
301 # Avoids a cache/DB lookup on every normal LLM request.
302 if metadata.get("user_api_key") and not metadata.get("user_api_key_user_id"): 302 ↛ 303line 302 didn't jump to line 303 because the condition on line 302 was never true
303 metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( # rebind-ok: enriched metadata replaces the original
304 metadata=metadata,
305 resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES,
306 )
307 _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
308 budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata)
309 if ( 309 ↛ 314line 309 didn't jump to line 314 because the condition on line 309 was never true
310 isinstance(completion_response, LiteLLMBatch)
311 and kwargs.get("call_type") == CallTypes.aretrieve_batch.value
312 and not batch_cost_is_final(completion_response)
313 ):
314 verbose_proxy_logger.debug(
315 "Cost tracking deferred for batch %s still in status %s",
316 completion_response.id,
317 completion_response.status,
318 )
319 await _release_budget_reservation(budget_reservation=budget_reservation)
320 return
321 user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None))
322 team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None))
323 org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None))
324 project_id: Final = (
325 project_id_value
326 if isinstance(project_id_value := metadata.get("user_api_key_project_id"), str)
327 else None
328 )
329 key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None))
330 end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None)
331 sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
332 response_cost: Final = (
333 sl_object.get("response_cost", None) if sl_object is not None else kwargs.get("response_cost", None)
334 )
335 tags: Final = _get_request_tags_for_cost_tracking(
336 sl_object=sl_object,
337 metadata=metadata,
338 )
339 model_access_groups: Final = debitable_model_access_groups(
340 attributed=get_request_model_access_groups(kwargs),
341 served_model_id=sl_object.get("model_id") if sl_object is not None else None,
342 router=get_llm_router(),
343 )
345 if response_cost is not None:
346 user_api_key: Final = metadata.get("user_api_key", None)
347 verbose_proxy_logger.debug(
348 "user_api_key %s, user_id %s, team_id %s, end_user_id %s",
349 user_api_key,
350 user_id,
351 team_id,
352 end_user_id,
353 )
354 call_type: str | None = kwargs.get("call_type")
355 if _should_track_cost_callback(
356 user_api_key=user_api_key,
357 user_id=user_id,
358 team_id=team_id,
359 end_user_id=end_user_id,
360 call_type=call_type,
361 ):
362 ## UPDATE DATABASE
363 charged: Final = await _update_database_and_spend_counters(
364 proxy_logging_obj=proxy_logging_obj,
365 increment_spend_counters=increment_spend_counters,
366 user_api_key=user_api_key,
367 user_id=user_id,
368 end_user_id=end_user_id,
369 team_id=team_id,
370 org_id=org_id,
371 kwargs=kwargs,
372 completion_response=completion_response,
373 start_time=start_time,
374 end_time=end_time,
375 response_cost=response_cost,
376 budget_reservation=budget_reservation,
377 request_tags=tags,
378 model_access_groups=model_access_groups,
379 project_id=project_id,
380 )
381 if not charged: 381 ↛ 382line 381 didn't jump to line 382 because the condition on line 381 was never true
382 return
384 # update cache (fire-and-forget for backward compat:
385 # cached object fields, soft budget alerts, etc.)
386 asyncio.create_task(
387 update_cache(
388 token=user_api_key,
389 user_id=user_id,
390 end_user_id=end_user_id,
391 response_cost=response_cost,
392 team_id=team_id,
393 parent_otel_span=parent_otel_span,
394 tags=tags,
395 )
396 )
398 await proxy_logging_obj.slack_alerting_instance.customer_spend_alert(
399 token=user_api_key,
400 key_alias=key_alias,
401 end_user_id=end_user_id,
402 response_cost=response_cost,
403 max_budget=end_user_max_budget,
404 )
405 elif budget_reservation is not None: 405 ↛ 406line 405 didn't jump to line 406 because the condition on line 405 was never true
406 await _release_budget_reservation(budget_reservation=budget_reservation)
407 else:
408 if _is_unbilled_interaction_response(completion_response): 408 ↛ 409line 408 didn't jump to line 409 because the condition on line 408 was never true
409 if BACKGROUND_INTERACTION_COST_POLLING_ENABLED and _is_unbilled_in_progress_interaction(
410 completion_response
411 ):
412 verbose_proxy_logger.debug(
413 "Cost tracking deferred for in-progress background interaction; "
414 "the budget reservation stays open until the poll task logs the final usage"
415 )
416 return
417 await _release_budget_reservation(budget_reservation=budget_reservation)
418 verbose_proxy_logger.debug(
419 "Released the budget reservation for an interaction create with no usage "
420 "that no poll task will settle"
421 )
422 return
423 await _release_budget_reservation(budget_reservation=budget_reservation)
424 # Non-model call types (health checks, afile_delete) have no model or standard_logging_object.
425 # Use .get() for "stream" to avoid KeyError on health checks.
426 # WS session wrappers (_aresponses_websocket, _arealtime) also reach here with
427 # result=None; their per-turn costs are tracked on the inner aresponses/realtime calls.
428 if sl_object is None and ( 428 ↛ 436line 428 didn't jump to line 436 because the condition on line 428 was always true
429 not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime")
430 ):
431 verbose_proxy_logger.warning(
432 "Cost tracking - skipping, no standard_logging_object for call_type=%s",
433 kwargs.get("call_type", "unknown"),
434 )
435 return
436 if kwargs.get("stream") is not True or (
437 kwargs.get("stream") is True and "complete_streaming_response" in kwargs
438 ):
439 if sl_object is not None:
440 cost_tracking_failure_debug_info: dict | str = (
441 sl_object["response_cost_failure_debug_info"]
442 or "response_cost_failure_debug_info is None in standard_logging_object"
443 )
444 else:
445 cost_tracking_failure_debug_info = "standard_logging_object not found"
446 model = kwargs.get("model")
447 raise Exception(
448 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"
449 )
450 except Exception as e:
451 failing_model: Final = kwargs.get("model", "")
452 failing_call_type: Final = kwargs.get("call_type", "")
453 error_msg: Final = (
454 f"Error in tracking cost callback - {e}\n Traceback:{traceback.format_exc()}\n"
455 f" Args to _PROXY_track_cost_callback\n model: {failing_model}\n call_type: {failing_call_type}\n"
456 )
457 failing_litellm_params: Final = kwargs.get("litellm_params") or {}
458 verbose_proxy_logger.debug(
459 "Cost tracking callback failed for model=%s call_type=%s;"
460 " chosen_metadata keys=%s litellm_metadata keys=%s old_metadata keys=%s",
461 failing_model,
462 failing_call_type,
463 _metadata_keys(get_litellm_metadata_from_kwargs(kwargs=kwargs)),
464 _metadata_keys(failing_litellm_params.get("litellm_metadata")),
465 _metadata_keys(failing_litellm_params.get("metadata")),
466 )
467 asyncio.create_task(
468 proxy_logging_obj.failed_tracking_alert(
469 error_message=error_msg,
470 failing_model=failing_model,
471 )
472 )
474 spend_log_error("Error in tracking cost callback - %s", str(e), exc=e)
476 @staticmethod
477 async def _enrich_failure_metadata_unless_db_stalled(metadata: dict, original_exception: Exception) -> dict:
478 if isinstance(original_exception, DBLookupDeadlineExceeded): 478 ↛ 479line 478 didn't jump to line 479 because the condition on line 478 was never true
479 return metadata
480 return await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata)
482 @staticmethod
483 async def _enrich_failure_metadata_with_key_info(metadata: dict, resolve_missing_key_identity: bool = True) -> dict:
484 """
485 Enriches failure spend log metadata by looking up the key object (and team object)
486 from cache/DB when key fields are missing.
488 This handles two scenarios:
489 1. Auth errors (401): UserAPIKeyAuth is created with only api_key set, all other
490 fields are null. We look up the full key object to fill in alias, user_id,
491 team_id, etc.
492 2. Post-auth failures (provider errors, rate limits): key fields are populated
493 but team_alias is missing because LiteLLM_VerificationTokenView SQL view
494 doesn't include it. We look up the team object to fill in team_alias.
496 Scenario 1 reads the key's identity as it stands right now, so it is only correct
497 for a log emitted within the request it describes. Callers that log after a delay,
498 against an identity captured earlier, pass resolve_missing_key_identity=False and
499 keep their own user_id, team_id and org_id.
500 """
501 api_key_hash: Final = metadata.get("user_api_key")
502 if not api_key_hash:
503 return metadata
505 from litellm.proxy.proxy_server import (
506 prisma_client,
507 proxy_logging_obj,
508 user_api_key_cache,
509 )
511 # Step 1: If key fields are missing, look up the full key object
512 if resolve_missing_key_identity and metadata.get("user_api_key_alias") is None: 512 ↛ 537line 512 didn't jump to line 537 because the condition on line 512 was always true
513 try:
514 key_obj: Final = await get_key_object(
515 hashed_token=api_key_hash,
516 prisma_client=prisma_client,
517 user_api_key_cache=user_api_key_cache,
518 proxy_logging_obj=proxy_logging_obj,
519 )
520 if metadata.get("user_api_key_alias") is None: 520 ↛ 522line 520 didn't jump to line 522 because the condition on line 520 was always true
521 metadata["user_api_key_alias"] = key_obj.key_alias
522 if metadata.get("user_api_key_user_id") is None:
523 metadata["user_api_key_user_id"] = key_obj.user_id
524 if metadata.get("user_api_key_team_id") is None: 524 ↛ 526line 524 didn't jump to line 526 because the condition on line 524 was always true
525 metadata["user_api_key_team_id"] = key_obj.team_id
526 if metadata.get("user_api_key_org_id") is None: 526 ↛ 528line 526 didn't jump to line 528 because the condition on line 526 was always true
527 metadata["user_api_key_org_id"] = key_obj.org_id
528 if metadata.get("user_api_key_project_id") is None: 528 ↛ 537line 528 didn't jump to line 537 because the condition on line 528 was always true
529 metadata["user_api_key_project_id"] = key_obj.project_id
530 except Exception:
531 verbose_proxy_logger.debug(
532 "Failed to enrich failure metadata with key info for api_key=%s",
533 api_key_hash,
534 )
536 # Step 2: If team_id is known but team_alias is missing, look up the team object
537 team_id: Final = metadata.get("user_api_key_team_id")
538 if team_id and metadata.get("user_api_key_team_alias") is None: 538 ↛ 539line 538 didn't jump to line 539 because the condition on line 538 was never true
539 try:
540 team_obj: Final = await get_team_object(
541 team_id=team_id,
542 prisma_client=prisma_client,
543 user_api_key_cache=user_api_key_cache,
544 proxy_logging_obj=proxy_logging_obj,
545 )
546 if team_obj.team_alias is not None:
547 metadata["user_api_key_team_alias"] = team_obj.team_alias
548 except Exception:
549 verbose_proxy_logger.debug(
550 "Failed to enrich failure metadata with team_alias for team_id=%s",
551 team_id,
552 )
553 return metadata
555 @staticmethod
556 def _should_track_errors_in_db():
557 """
558 Returns True if errors should be tracked in the database
560 By default, errors are tracked in the database
562 If users want to disable error tracking, they can set the disable_error_logs flag in the general_settings
563 """
564 from litellm.proxy.proxy_server import general_settings
566 if general_settings.get("disable_error_logs") is True: 566 ↛ 567line 566 didn't jump to line 567 because the condition on line 566 was never true
567 return False
568 return
571def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None:
572 patch = {k: v for k, v in metadata.items() if (k.startswith("user_api_key") or k == "tags") and v is not None}
573 if not patch:
574 return
576 litellm_params: Final = kwargs.setdefault("litellm_params", {})
577 for bucket_name in ("litellm_metadata", "metadata"):
578 bucket = litellm_params.get(bucket_name)
579 if isinstance(bucket, dict):
580 for key, value in patch.items():
581 if bucket.get(key) is None:
582 bucket[key] = value
585async def run_spend_event(line: bytes) -> None:
586 await _ProxyDBLogger().run_spend_event(line)
589def _is_unbilled_interaction_response(completion_response: object) -> bool:
590 from litellm.interactions.background_cost_polling import missing_usage_is_expected
591 from litellm.types.interactions import InteractionsAPIResponse
593 if not isinstance(completion_response, InteractionsAPIResponse): 593 ↛ 595line 593 didn't jump to line 595 because the condition on line 593 was always true
594 return False
595 return completion_response.usage is None and missing_usage_is_expected(completion_response)
598def _is_unbilled_in_progress_interaction(completion_response: object) -> bool:
599 from litellm.interactions.background_cost_polling import is_pollable_background_interaction
600 from litellm.types.interactions import InteractionsAPIResponse
602 if not isinstance(completion_response, InteractionsAPIResponse):
603 return False
604 return completion_response.usage is None and is_pollable_background_interaction(completion_response)
607def _should_track_cost_callback(
608 user_api_key: str | None,
609 user_id: str | None,
610 team_id: str | None,
611 end_user_id: str | None,
612 call_type: str | None = None,
613) -> bool:
614 """
615 Determine if the cost callback should be tracked based on the kwargs
617 Pass-through endpoints can be configured with ``auth=false``, which leaves
618 the request with no key/user/team/end-user to attribute spend to. Those
619 requests still forward real provider traffic that operators expect to see
620 in request/usage logs, so they are tracked even when unauthenticated.
621 The same reasoning applies to a completed managed batch's cost event
622 (see _UNATTRIBUTED_TRACKABLE_CALL_TYPES).
623 """
625 # don't run track cost callback if user opted into disabling spend
626 if ProxyUpdateSpend.disable_spend_updates() is True: 626 ↛ 627line 626 didn't jump to line 627 because the condition on line 626 was never true
627 return False
629 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:
630 return True
631 return call_type in _UNATTRIBUTED_TRACKABLE_CALL_TYPES
634def _metadata_keys(metadata: object) -> tuple[str, ...]:
635 if not isinstance(metadata, Mapping):
636 return ()
637 return tuple(sorted(str(key) for key in metadata))
640def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None:
641 return budget_reservation_from_metadata(metadata)
644def _get_request_tags_for_cost_tracking(
645 sl_object: StandardLoggingPayload | None,
646 metadata: dict,
647) -> list[str] | None:
648 if sl_object is not None:
649 request_tags: Final = sl_object.get("request_tags", None)
650 if isinstance(request_tags, list): 650 ↛ 653line 650 didn't jump to line 653 because the condition on line 650 was always true
651 return request_tags
653 metadata_tags: Final = metadata.get("tags", None)
654 if isinstance(metadata_tags, list): 654 ↛ 655line 654 didn't jump to line 655 because the condition on line 654 was never true
655 return metadata_tags
657 return None
660class _IncrementSpendCounters(Protocol):
661 """The ``increment_spend_counters`` coroutine :func:`_update_database_and_spend_counters` awaits."""
663 async def __call__( 663 ↛ exitline 663 didn't return from function '__call__' because
664 self,
665 token: str | None,
666 team_id: str | None,
667 user_id: str | None,
668 response_cost: float | None,
669 org_id: str | None = None,
670 budget_reservation: dict[str, object] | None = None,
671 end_user_id: str | None = None,
672 tags: list[str] | None = None,
673 request_started_at: datetime | None = None,
674 model_access_groups: Sequence[str] | None = None,
675 ) -> None: ...
678async def _update_database_and_spend_counters(
679 proxy_logging_obj: "ProxyLogging",
680 increment_spend_counters: _IncrementSpendCounters,
681 user_api_key: str | None,
682 user_id: str | None,
683 end_user_id: str | None,
684 team_id: str | None,
685 org_id: str | None,
686 kwargs: dict,
687 completion_response: object,
688 start_time: datetime | None,
689 end_time: datetime | None,
690 response_cost: float,
691 budget_reservation: dict | None,
692 request_tags: list[str] | None = None,
693 model_access_groups: Sequence[str] | None = None,
694 project_id: str | None = None,
695) -> bool:
696 if budget_reservation is not None: 696 ↛ 697line 696 didn't jump to line 697 because the condition on line 696 was never true
697 await _reconcile_budget_reservation_before_db_update(
698 budget_reservation=budget_reservation, response_cost=response_cost
699 )
700 try:
701 charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database(
702 token=user_api_key,
703 response_cost=response_cost,
704 user_id=user_id,
705 end_user_id=end_user_id,
706 team_id=team_id,
707 kwargs=kwargs,
708 completion_response=completion_response,
709 start_time=start_time,
710 end_time=end_time,
711 org_id=org_id,
712 project_id=project_id,
713 )
714 except Exception:
715 if budget_reservation is not None:
716 try:
717 await _release_budget_reservation(budget_reservation=budget_reservation)
718 except Exception:
719 verbose_proxy_logger.exception("Failed to release budget reservation after database update failed")
720 try:
721 await _invalidate_budget_reservation_counters(budget_reservation=budget_reservation)
722 except Exception:
723 verbose_proxy_logger.exception(
724 "Failed to invalidate budget reservation counters after release failed"
725 )
726 raise
727 if not charged: 727 ↛ 728line 727 didn't jump to line 728 because the condition on line 727 was never true
728 await _release_budget_reservation(budget_reservation=budget_reservation)
729 return False
731 try:
732 await increment_spend_counters(
733 token=user_api_key,
734 team_id=team_id,
735 user_id=user_id,
736 response_cost=response_cost,
737 org_id=org_id,
738 budget_reservation=budget_reservation,
739 end_user_id=end_user_id,
740 tags=request_tags,
741 request_started_at=start_time,
742 model_access_groups=model_access_groups,
743 project_id=project_id,
744 )
745 except Exception:
746 if budget_reservation is not None:
747 try:
748 await _invalidate_budget_reservation_counters(budget_reservation=budget_reservation)
749 except Exception:
750 verbose_proxy_logger.exception(
751 "Failed to invalidate budget reservation counters after spend counter update failed"
752 )
753 finally:
754 budget_reservation["finalized"] = True
755 raise
756 return True
759async def _reconcile_budget_reservation_before_db_update(
760 budget_reservation: dict, # mutable-ok: reconcile_budget_reservation stamps applied_adjustment on the caller's shared reservation dict
761 response_cost: float,
762) -> None:
763 from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation
765 try:
766 await reconcile_budget_reservation(
767 budget_reservation=budget_reservation, actual_cost=response_cost, finalize=False
768 )
769 except Exception: # noqa: BLE001 # a failed reconcile must not block the spend write; the counters are dropped instead
770 verbose_proxy_logger.warning(
771 "Failed to reconcile budget reservation before persisting spend; invalidating reserved counters"
772 )
773 try:
774 await _invalidate_budget_reservation_counters(budget_reservation=budget_reservation)
775 except Exception: # noqa: BLE001 # nothing left to try; the finalized stamp below keeps it from being reprocessed
776 verbose_proxy_logger.exception(
777 "Failed to invalidate budget reservation counters after pre-persist reconcile failed"
778 )
779 finally:
780 budget_reservation["finalized"] = True # rebind-ok: stamps the caller's shared dict for the counter update
783async def _release_budget_reservation(budget_reservation: dict | None) -> None:
784 if budget_reservation is None: 784 ↛ 787line 784 didn't jump to line 787 because the condition on line 784 was always true
785 return
787 from litellm.proxy.spend_tracking.budget_reservation import (
788 release_budget_reservation,
789 )
791 await release_budget_reservation(
792 budget_reservation=budget_reservation,
793 )
796async def _invalidate_budget_reservation_counters(
797 budget_reservation: dict | None,
798) -> None:
799 if budget_reservation is None:
800 return
802 from litellm.proxy.spend_tracking.budget_reservation import (
803 invalidate_budget_reservation_counters,
804 )
806 await invalidate_budget_reservation_counters(
807 budget_reservation=budget_reservation,
808 )