From 537e8ac068d45a866160dcd245247a36d8fc5b6c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 19:49:26 -0700 Subject: [PATCH 1/3] feat(cost): warn and count $0 cost on billable requests (#42345) * feat(cost): warn and count $0 cost on billable requests A request that carries usage but prices to $0 on a model whose pricing entry has a non-zero rate now logs one warning naming the model, the pricing entry, and the missing rate, and increments litellm_zero_cost_requests_total{requested_model, model, model_id, api_provider, reason}. Free models (every used rate is 0), requests without usage, and unmapped models stay silent. The diagnostic rides on the standard logging payload as zero_cost_diagnostic * fix(cost): keep the zero-cost diagnostic importable on 3.10 and recursion-free * fix(cost): warn once per request when a $0 result is priced again * fix(cost): judge a free deployment by its own pricing and keep it silent on calculator errors * fix(cost): warn once per request when a usage-less evaluation sits between two zero-cost findings * fix(cost): judge zero-cost findings by the priced entry, skip cache hits, count failure rows * test(cost): type the zero-cost diagnostic test helpers * test(logging): flag a $0 terminal Responses stream event by its inner response * chore: restore the lazy OpenAPI snapshot as CI's Python 3.12 generates it --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../grafana_dashboard.json | 57 ++ litellm/cost_calculator.py | 43 +- litellm/integrations/prometheus.py | 47 ++ litellm/litellm_core_utils/litellm_logging.py | 149 ++++- .../llm_cost_calc/zero_cost_diagnostic.py | 146 +++++ litellm/types/integrations/prometheus.py | 10 + litellm/types/utils.py | 10 + .../test_prometheus_zero_cost_metric.py | 179 ++++++ .../test_zero_cost_diagnostic.py | 157 ++++++ .../test_litellm_logging.py | 531 ++++++++++++++++-- 10 files changed, 1268 insertions(+), 61 deletions(-) create mode 100644 litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py create mode 100644 tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py create mode 100644 tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py diff --git a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json index d8cb122417a..af88708166f 100644 --- a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json +++ b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json @@ -6267,6 +6267,63 @@ ], "title": "Spend update queue sizes (litellm__size)", "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "Requests that carried usage but were logged at $0 on a model whose pricing entry has a non-zero rate, by requested model and reason", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "drawStyle": "line", + "fillOpacity": 10, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": false + }, + "unit": "short" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 430 + }, + "id": 110, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "sum(rate(litellm_zero_cost_requests_total[$__rate_interval])) by (requested_model, reason)", + "legendFormat": "{{requested_model}} / {{reason}}", + "range": true, + "refId": "A" + } + ], + "title": "litellm_zero_cost_requests rate", + "type": "timeseries" } ], "preload": false, diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index b317e356e1d..37743a9ce33 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -948,7 +948,7 @@ def _extract_service_tier(source: object) -> str | None: return None -def _get_usage_object( +def get_usage_object( completion_response: object, ) -> Usage | None: usage_obj: Final = cast( @@ -1336,7 +1336,7 @@ def completion_cost( cache_creation_input_tokens: int | None = None cache_read_input_tokens: int | None = None audio_transcription_file_duration: float = 0.0 - provider_usage_object: Final = _get_usage_object(completion_response=completion_response) + provider_usage_object: Final = get_usage_object(completion_response=completion_response) cost_per_token_usage_object: Final[Usage | None] = ( _without_provider_stated_cost(provider_usage_object) if custom_pricing else provider_usage_object ) @@ -2033,6 +2033,45 @@ def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelIn return None +def _raw_cost_map_entry(key: str) -> Mapping[str, object] | None: + raw_entry: Final = litellm.model_cost.get(key) + return raw_entry if isinstance(raw_entry, Mapping) else None + + +def pricing_entry_for_cost_calc( + model: str | None, + completion_response: object | None, + custom_llm_provider: str | None, + custom_pricing: bool | None, + base_model: str | None, + router_model_id: str | None, + region_name: str | None, + litellm_logging_obj: LitellmLoggingObject | None, +) -> tuple[str, Mapping[str, object]] | None: + deployment_entry: Final = _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id) + deployment_key: Final = router_model_id or model + if deployment_entry is not None and deployment_key is not None: + registered_entry: Final = _raw_cost_map_entry(router_model_id) if router_model_id is not None else None + return deployment_key, registered_entry or deployment_entry + selected_model: Final = _select_model_name_for_cost_calc( + model=model, + completion_response=completion_response, + base_model=base_model, + custom_pricing=custom_pricing, + custom_llm_provider=custom_llm_provider, + router_model_id=router_model_id, + region_name=region_name, + ) + candidates: Final = (selected_model, _get_response_model(completion_response), model) + resolved: Final = next( + (info for info in (_cost_map_model_info(name, custom_llm_provider) for name in candidates if name) if info), + None, + ) + if resolved is None: + return None + return resolved["key"], _raw_cost_map_entry(resolved["key"]) or resolved + + def ocr_cost( model: str, custom_llm_provider: str | None, diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 37b7344917e..28ac9f5cdae 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -10,6 +10,7 @@ import sys from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import replace from datetime import datetime, timedelta +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast from pydantic import BaseModel @@ -66,6 +67,7 @@ from litellm.types.proxy.carried_budget_state import ( from litellm.types.utils import ( StandardLoggingGuardrailInformation, StandardLoggingPayload, + StandardLoggingZeroCostDiagnostic, ) if TYPE_CHECKING: @@ -713,6 +715,15 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_requests_metric"), ) + self.litellm_zero_cost_requests_total = self._counter_factory( + name="litellm_zero_cost_requests_total", + documentation=( + "Requests that carried usage but were logged at $0 on a model whose pricing entry " + "has a non-zero rate, by reason (missing_pricing_key, pricing_not_applied, cost_calculation_error)" + ), + labelnames=self.get_labels_for_metric("litellm_zero_cost_requests_total"), + ) + # Cache metrics self.litellm_cache_hits_metric = self._counter_factory( name="litellm_cache_hits_metric", @@ -1410,6 +1421,11 @@ class PrometheusLogger(CustomLogger): enum_values=enum_values, label_context=label_context, ) + self._increment_zero_cost_requests_metric( + zero_cost_diagnostic=standard_logging_payload.get("zero_cost_diagnostic"), + enum_values=enum_values, + label_context=label_context, + ) # input, output, total token metrics self._increment_token_metrics( @@ -1983,6 +1999,30 @@ class PrometheusLogger(CustomLogger): amount=float(response_cost), ) + def _increment_zero_cost_requests_metric( + self, + zero_cost_diagnostic: StandardLoggingZeroCostDiagnostic | None, + enum_values: UserAPIKeyLabelValues, + label_context: PrometheusLabelFactoryContext, + ) -> None: + if zero_cost_diagnostic is None: + return + supported_labels: Final = self.get_labels_for_metric("litellm_zero_cost_requests_total") + reason_label: Final = ( + MappingProxyType({ZERO_COST_REASON_LABEL: zero_cost_diagnostic["reason"]}) + if ZERO_COST_REASON_LABEL in supported_labels + else MappingProxyType({}) + ) + labels: Final = MappingProxyType( + { + **prometheus_label_factory( + supported_enum_labels=supported_labels, enum_values=enum_values, label_context=label_context + ), + **reason_label, + } + ) + self.litellm_zero_cost_requests_total.labels(**labels).inc() + @staticmethod def _get_remaining_from_v3_rate_limit_headers( standard_logging_payload: StandardLoggingPayload | None, @@ -2333,6 +2373,8 @@ class PrometheusLogger(CustomLogger): team_alias=user_api_team_alias, user=user_id, model_id=standard_logging_payload.get("model_id", ""), + requested_model=standard_logging_payload.get("model_group"), + api_provider=standard_logging_payload.get("custom_llm_provider"), custom_metadata_labels=get_custom_labels_from_metadata( metadata=_get_combined_custom_metadata_from_standard_logging_payload( standard_logging_payload=standard_logging_payload @@ -2345,6 +2387,11 @@ class PrometheusLogger(CustomLogger): "litellm_llm_api_failed_requests_metric", enum_values, ) + self._increment_zero_cost_requests_metric( + zero_cost_diagnostic=standard_logging_payload.get("zero_cost_diagnostic"), + enum_values=enum_values, + label_context=PrometheusLabelFactoryContext(enum_values), + ) self.set_llm_deployment_failure_metrics(kwargs) await self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7bad711940e..a486cdeff19 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -50,6 +50,8 @@ from litellm.cost_calculator import ( RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, _select_model_name_for_cost_calc, + get_usage_object, + pricing_entry_for_cost_calc, ) from litellm.exceptions import ( BudgetExceededError, @@ -89,6 +91,10 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( InteractionsUsageObjectTransformation, ) +from litellm.litellm_core_utils.llm_cost_calc.zero_cost_diagnostic import ( + diagnose_zero_cost, + zero_cost_warning, +) from litellm.litellm_core_utils.logging_utils import ( truncate_base64_in_messages, truncate_base64_in_messages_async, @@ -157,6 +163,7 @@ from litellm.types.utils import ( StandardLoggingPayloadStatusFields, StandardLoggingPromptManagementMetadata, StandardLoggingVectorStoreRequest, + StandardLoggingZeroCostDiagnostic, TextCompletionResponse, TranscriptionResponse, Usage, @@ -614,6 +621,7 @@ class Logging(LiteLLMLoggingBaseClass): self.truncated_messages_for_logging: str | list | dict | None = None # mutable-ok: logged messages shape ## TIME TO FIRST TOKEN LOGGING ## self.completion_start_time: datetime.datetime | None = None + self.zero_cost_warned: bool = False self._llm_caching_handler: LLMCachingHandler | None = None # INITIAL LITELLM_PARAMS @@ -1764,11 +1772,6 @@ class Logging(LiteLLMLoggingBaseClass): ) result_hidden_params: Final = getattr(priced_result, "_hidden_params", None) or MappingProxyType({}) - result_additional_headers: Final = ( - result_hidden_params.get("additional_headers") - if isinstance(result_hidden_params, dict) - else getattr(result_hidden_params, "additional_headers", None) - ) if isinstance(priced_result, (BaseModel, HttpxBinaryResponseContent)) and hasattr( priced_result, "_hidden_params" ): @@ -1776,6 +1779,12 @@ class Logging(LiteLLMLoggingBaseClass): if ( "response_cost" in hidden_params and hidden_params["response_cost"] is not None ): # use cost if already calculated + self._record_zero_cost_diagnostic( + priced_result, + hidden_params["response_cost"], + litellm_model_name=litellm_model_name, + router_model_id=router_model_id or hidden_params.get("model_id"), + ) return hidden_params["response_cost"] elif router_model_id is None and "model_id" in hidden_params: # use model_id if not already set router_model_id = hidden_params["model_id"] @@ -1787,18 +1796,7 @@ class Logging(LiteLLMLoggingBaseClass): router_model_id = self.get_router_model_id() ## RESPONSE COST ## - spilled_over: Final = is_spilled_over_ptu_request( - model_info=_deployment_model_info(self.litellm_params if hasattr(self, "litellm_params") else None), - response_headers=self.model_call_details.get("response_headers"), - additional_headers=result_additional_headers, - ) - custom_pricing: Final = ( - False - if spilled_over - else use_custom_pricing_for_model( - litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None) - ) - ) + custom_pricing: Final = self._custom_pricing_for(priced_result) prompt = self._prompt_for_cost_calculation() @@ -1850,9 +1848,18 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug("response_cost: %s", response_cost) additional_response_cost: Final[object] = self.model_call_details.get("additional_response_cost") - if isinstance(additional_response_cost, (int, float)) and additional_response_cost > 0: - return (response_cost or 0.0) + additional_response_cost - return response_cost + total_response_cost: Final = ( + (response_cost or 0.0) + additional_response_cost + if isinstance(additional_response_cost, (int, float)) and additional_response_cost > 0 + else response_cost + ) + self._record_zero_cost_diagnostic( + priced_result, + total_response_cost, + litellm_model_name=litellm_model_name, + router_model_id=router_model_id, + ) + return total_response_cost except Exception as e: # error calculating cost debug_info = StandardLoggingModelCostFailureDebugInformation( error_str=str(e), @@ -1866,9 +1873,108 @@ class Logging(LiteLLMLoggingBaseClass): ) verbose_logger.debug("response_cost_failure_debug_information: %s", debug_info) self.model_call_details["response_cost_failure_debug_information"] = debug_info + self._record_zero_cost_diagnostic( + priced_result, + None, + calculation_failed=True, + litellm_model_name=litellm_model_name, + router_model_id=router_model_id, + ) return None + def _record_zero_cost_diagnostic( + self, + result: object, + response_cost: float | None, + *, + calculation_failed: bool = False, + litellm_model_name: str | None = None, + router_model_id: str | None = None, + ) -> None: + if response_cost is None and not calculation_failed: + return + if self.model_call_details.get("cache_hit") is True: + self.model_call_details["zero_cost_diagnostic"] = None + return + try: + finding: Final = self._zero_cost_finding( + result, + response_cost, + calculation_failed=calculation_failed, + litellm_model_name=litellm_model_name, + router_model_id=router_model_id, + ) + except Exception as e: # noqa: BLE001 # the pricing helpers raise plain Exception and a diagnostic must never break cost tracking + verbose_logger.debug("zero_cost_diagnostic skipped: %s", e) + return + self.model_call_details["zero_cost_diagnostic"] = finding[0] if finding is not None else None + if finding is None or self.zero_cost_warned: + return + self.zero_cost_warned = True + verbose_logger.warning(finding[1]) + + def _zero_cost_finding( + self, + result: object, + response_cost: float | None, + *, + calculation_failed: bool, + litellm_model_name: str | None, + router_model_id: str | None, + ) -> tuple[StandardLoggingZeroCostDiagnostic, str] | None: + metadata: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(self.litellm_params) + if response_cost or is_unbilled_non_inference_call(self.call_type, metadata, result): + return None + usage: Final = get_usage_object(completion_response=result) + if usage is None: + return None + model: Final = litellm_model_name or self.model + custom_llm_provider: Final = self.model_call_details.get("custom_llm_provider") + pricing: Final = pricing_entry_for_cost_calc( + model=model, + completion_response=result, + custom_llm_provider=custom_llm_provider, + custom_pricing=self._custom_pricing_for(result), + base_model=_get_base_model_from_metadata(model_call_details=self.model_call_details), + router_model_id=router_model_id or self.get_router_model_id(), + region_name=_resolve_mantle_region_for_cost( + custom_llm_provider=custom_llm_provider, + litellm_params=self.model_call_details.get("litellm_params"), + ), + litellm_logging_obj=self, + ) + if pricing is None: + return None + diagnostic: Final = diagnose_zero_cost( + usage=usage, pricing_model=pricing[0], pricing_entry=pricing[1], calculation_failed=calculation_failed + ) + if diagnostic is None: + return None + model_group: Final = metadata.get("model_group") + return diagnostic, zero_cost_warning( + diagnostic, + model_group=model_group if isinstance(model_group, str) else None, + model=model, + custom_llm_provider=custom_llm_provider, + usage=usage, + ) + + def _custom_pricing_for(self, result: object) -> bool: + litellm_params: Final = getattr(self, "litellm_params", None) + result_hidden_params: Final = getattr(result, "_hidden_params", None) or MappingProxyType({}) + additional_headers: Final = ( + result_hidden_params.get("additional_headers") + if isinstance(result_hidden_params, dict) + else getattr(result_hidden_params, "additional_headers", None) + ) + spilled_over: Final = is_spilled_over_ptu_request( + model_info=_deployment_model_info(litellm_params), + response_headers=self.model_call_details.get("response_headers"), + additional_headers=additional_headers, + ) + return False if spilled_over else use_custom_pricing_for_model(litellm_params=litellm_params) + def _prompt_for_cost_calculation(self) -> str: """ The raw input string is only priced directly for text-to-speech, which bills per character. @@ -2213,6 +2319,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = 0.0 elif "response_cost" in hidden_params: self.model_call_details["response_cost"] = hidden_params["response_cost"] + self._record_zero_cost_diagnostic(logging_result, hidden_params["response_cost"]) elif (existing_cost := self.model_call_details.get("response_cost")) is not None and existing_cost != 0: # Preserve response_cost if already calculated (e.g., by pass-through # handlers like Gemini/Vertex which call completion_cost directly). @@ -6507,6 +6614,7 @@ def get_standard_logging_object_payload( error_str=error_str, error_information=error_information, response_cost_failure_debug_info=kwargs.get("response_cost_failure_debug_information"), + zero_cost_diagnostic=kwargs.get("zero_cost_diagnostic"), guardrail_information=metadata.get("standard_logging_guardrail_information", None), standard_built_in_tools_params=standard_built_in_tools_params, ) @@ -6685,6 +6793,7 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: response_cost=response_cost, autorouter_savings=None, response_cost_failure_debug_info=None, + zero_cost_diagnostic=None, status="success", total_tokens=int(DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT), prompt_tokens=int(DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT), diff --git a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py new file mode 100644 index 00000000000..6331d815bdc --- /dev/null +++ b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py @@ -0,0 +1,146 @@ +from collections.abc import Mapping +from functools import reduce +from typing import Final + +from pydantic import TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.types.utils import StandardLoggingZeroCostDiagnostic, Usage + +ZERO_COST_COUNTER_NAME: Final = "litellm_zero_cost_requests_total" + +_TEXT_INPUT_RATE: Final = "input_cost_per_token" +_AUDIO_INPUT_RATE: Final = "input_cost_per_audio_token" +_TEXT_OUTPUT_RATE: Final = "output_cost_per_token" +_AUDIO_OUTPUT_RATE: Final = "output_cost_per_audio_token" +_RATE_KEY_MARKERS: Final = ("cost", "pricing") +_NESTED_PRICING: Final = TypeAdapter(Mapping[str, object] | tuple[object, ...]) +_MAX_PRICING_DEPTH: Final = 4 + + +def _audio_tokens(details: object) -> int: + audio_tokens: Final = getattr(details, "audio_tokens", None) + return audio_tokens if isinstance(audio_tokens, int) and audio_tokens > 0 else 0 + + +def _tokens(value: object) -> int: + return value if isinstance(value, int) and value > 0 else 0 + + +def used_pricing_keys(usage: Usage) -> tuple[str, ...]: + prompt_audio: Final = _audio_tokens(usage.prompt_tokens_details) + completion_audio: Final = _audio_tokens(usage.completion_tokens_details) + prompt_text: Final = _tokens(usage.prompt_tokens) - prompt_audio + completion_text: Final = _tokens(usage.completion_tokens) - completion_audio + components: Final = ( + (_TEXT_INPUT_RATE, prompt_text), + (_AUDIO_INPUT_RATE, prompt_audio), + (_TEXT_OUTPUT_RATE, completion_text), + (_AUDIO_OUTPUT_RATE, completion_audio), + ) + return tuple(key for key, count in components if count > 0) + + +def _nested_pricing(value: object) -> Mapping[str, object] | tuple[object, ...] | None: + try: + return _NESTED_PRICING.validate_python(value) + except ValidationError: + return None + + +def _is_rate_key(key: str) -> bool: + return any(marker in key for marker in _RATE_KEY_MARKERS) + + +def _rate_values(value: object) -> tuple[object, ...]: + nested: Final = _nested_pricing(value) + if isinstance(nested, Mapping): + return tuple(child for key, child in nested.items() if _is_rate_key(key)) + if nested is None: + return (value,) + return nested + + +def _expand_rate_values(values: tuple[object, ...], _depth: int) -> tuple[object, ...]: + return tuple(nested for value in values for nested in _rate_values(value)) + + +def _is_positive_number(value: object) -> bool: + return not isinstance(value, bool) and isinstance(value, (int, float)) and value > 0 + + +def _declares_a_rate(pricing_entry: Mapping[str, object]) -> bool: + leaves: Final = reduce(_expand_rate_values, range(_MAX_PRICING_DEPTH), (pricing_entry,)) + return any(_is_positive_number(leaf) for leaf in leaves) + + +def _is_explicit_zero(value: object) -> bool: + return not isinstance(value, bool) and isinstance(value, (int, float)) and value == 0 + + +def diagnose_zero_cost( + usage: Usage, + pricing_model: str, + pricing_entry: Mapping[str, object], + calculation_failed: bool, +) -> StandardLoggingZeroCostDiagnostic | None: + used_keys: Final = used_pricing_keys(usage) + if not used_keys: + return None + missing_keys: Final = tuple(key for key in used_keys if pricing_entry.get(key) is None) + if not missing_keys and all(_is_explicit_zero(pricing_entry[key]) for key in used_keys): + return None + if not _declares_a_rate(pricing_entry): + return None + if calculation_failed: + return StandardLoggingZeroCostDiagnostic( + reason="cost_calculation_error", pricing_model=pricing_model, missing_pricing_keys=() + ) + if missing_keys: + return StandardLoggingZeroCostDiagnostic( + reason="missing_pricing_key", pricing_model=pricing_model, missing_pricing_keys=missing_keys + ) + return StandardLoggingZeroCostDiagnostic( + reason="pricing_not_applied", pricing_model=pricing_model, missing_pricing_keys=() + ) + + +def _cause(diagnostic: StandardLoggingZeroCostDiagnostic) -> str: + reason: Final = diagnostic["reason"] + match reason: + case "missing_pricing_key": + return ( + f"pricing entry '{diagnostic['pricing_model']}' has no {', '.join(diagnostic['missing_pricing_keys'])}. " + "Set the missing rate in the deployment's model_info or in the model cost map, " + "or set every rate to 0 to mark the model free" + ) + case "pricing_not_applied": + return ( + f"pricing entry '{diagnostic['pricing_model']}' declares non-zero rates for this usage, " + "but the cost calculator returned $0" + ) + case "cost_calculation_error": + return ( + f"cost calculation raised for pricing entry '{diagnostic['pricing_model']}', " + "see response_cost_failure_debug_information" + ) + case _: + return assert_never(reason) + + +def zero_cost_warning( + diagnostic: StandardLoggingZeroCostDiagnostic, + *, + model_group: str | None, + model: str, + custom_llm_provider: str | None, + usage: Usage, +) -> str: + request: Final = ( + f"model_group={model_group or model} model={model} provider={custom_llm_provider or 'unknown'} " + f"prompt_tokens={_tokens(usage.prompt_tokens)} completion_tokens={_tokens(usage.completion_tokens)}" + ) + return ( + f"Billable request priced at $0 and logged as such ({request}): {_cause(diagnostic)}. " + f'Counted in {ZERO_COST_COUNTER_NAME}{{reason="{diagnostic["reason"]}"}}' + ) diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index f4893e857d1..c929ee2ee79 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -131,6 +131,7 @@ EXCEPTION_STATUS: Final = "exception_status" EXCEPTION_CLASS: Final = "exception_class" RATE_LIMIT_CATEGORY: Final = "rate_limit_category" RATE_LIMIT_TYPE: Final = "rate_limit_type" +ZERO_COST_REASON_LABEL: Final = "reason" STATUS_CODE: Final = "status_code" EXCEPTION_LABELS: Final = [EXCEPTION_STATUS, EXCEPTION_CLASS] LATENCY_BUCKETS: Final = ( @@ -279,6 +280,7 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_guardrail_latency_seconds", "litellm_guardrail_errors_total", "litellm_guardrail_requests_total", + "litellm_zero_cost_requests_total", # Cache metrics "litellm_cache_hits_metric", "litellm_cache_misses_metric", @@ -590,6 +592,14 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.SERVICE_TIER.value, ] + litellm_zero_cost_requests_total = ( + UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.MODEL_ID.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + ZERO_COST_REASON_LABEL, + ) + litellm_input_tokens_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6cc2637357d..35d967fc1d5 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3206,6 +3206,15 @@ class StandardLoggingModelCostFailureDebugInformation(TypedDict, total=False): custom_pricing: bool | None +ZeroCostReason = Literal["missing_pricing_key", "pricing_not_applied", "cost_calculation_error"] + + +class StandardLoggingZeroCostDiagnostic(TypedDict): + reason: ReadOnly[ZeroCostReason] + pricing_model: ReadOnly[str] + missing_pricing_keys: ReadOnly[tuple[str, ...]] + + class StandardLoggingPayloadErrorInformation(TypedDict, total=False): error_code: str | None error_class: str | None @@ -3524,6 +3533,7 @@ class StandardLoggingPayload(ClassifierAudit): autorouter_savings_estimate: ReadOnly[Mapping[str, JsonValue] | None] autorouter_baseline_observation: ReadOnly[str | None] response_cost_failure_debug_info: StandardLoggingModelCostFailureDebugInformation | None + zero_cost_diagnostic: NotRequired[ReadOnly[StandardLoggingZeroCostDiagnostic | None]] status: StandardLoggingPayloadStatus status_fields: StandardLoggingPayloadStatusFields custom_llm_provider: str | None diff --git a/tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py b/tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py new file mode 100644 index 00000000000..989009b309a --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py @@ -0,0 +1,179 @@ +import datetime +from typing import Final + +import pytest +from prometheus_client import REGISTRY +from prometheus_client.samples import Sample + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.utils import StandardLoggingZeroCostDiagnostic + +METRIC: Final = "litellm_zero_cost_requests_total" +MISSING_KEY_DIAGNOSTIC: Final[StandardLoggingZeroCostDiagnostic] = { + "reason": "missing_pricing_key", + "pricing_model": "dep-1", + "missing_pricing_keys": ("input_cost_per_token", "output_cost_per_token"), +} + + +def _clear_prometheus_registry() -> None: + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +def _samples(metric_name: str) -> list[Sample]: + return [sample for metric in REGISTRY.collect() for sample in metric.samples if sample.name == metric_name] + + +def _payload(zero_cost_diagnostic: StandardLoggingZeroCostDiagnostic | None) -> dict[str, object]: + return { + "id": "t", + "call_type": "completion", + "response_cost": 0.0, + "status": "success", + "total_tokens": 30, + "prompt_tokens": 20, + "completion_tokens": 10, + "startTime": 1.0, + "endTime": 2.0, + "completionStartTime": 1.5, + "model": "openai/gpt-5.4-nano", + "model_id": "dep-1", + "model_group": "per-second-priced-chat", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "request_tags": [], + "end_user": None, + "cache_hit": False, + "stream": False, + "response": {"id": "chatcmpl-1"}, + "model_parameters": {}, + "zero_cost_diagnostic": zero_cost_diagnostic, + "metadata": { + "user_api_key_hash": "h", + "user_api_key_alias": "a", + "user_api_key_team_id": "t", + "user_api_key_team_alias": "ta", + "user_api_key_user_id": "u", + "user_api_key_user_email": "e@x.com", + "user_api_key_org_id": None, + "user_api_key_org_alias": None, + "requester_metadata": None, + "user_api_key_end_user_id": None, + "usage_object": None, + }, + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + } + + +async def _log_success( + logger: PrometheusLogger, zero_cost_diagnostic: StandardLoggingZeroCostDiagnostic | None +) -> None: + now: Final = datetime.datetime.now() + kwargs: Final = { + "model": "openai/gpt-5.4-nano", + "litellm_params": {"metadata": {}}, + "standard_logging_object": _payload(zero_cost_diagnostic), + "stream": False, + "start_time": now - datetime.timedelta(seconds=3), + "api_call_start_time": now - datetime.timedelta(seconds=2), + "completion_start_time": now - datetime.timedelta(seconds=1), + "end_time": now, + } + await logger.async_log_success_event(kwargs, None, now, now) + + +async def _log_failure( + logger: PrometheusLogger, zero_cost_diagnostic: StandardLoggingZeroCostDiagnostic | None +) -> None: + now: Final = datetime.datetime.now() + kwargs: Final = { + "model": "openai/gpt-5.4-nano", + "litellm_params": {"metadata": {}}, + "standard_logging_object": {**_payload(zero_cost_diagnostic), "status": "failure"}, + "exception": Exception("stream cut off after the usage chunk"), + "stream": True, + "start_time": now - datetime.timedelta(seconds=3), + "end_time": now, + } + await logger.async_log_failure_event(kwargs, None, now, now) + + +@pytest.mark.asyncio +async def test_failure_event_counts_a_zero_cost_request_by_model_and_reason() -> None: + _clear_prometheus_registry() + try: + logger: Final = PrometheusLogger() + await _log_failure(logger, None) + assert _samples(METRIC) == [] + + await _log_failure(logger, MISSING_KEY_DIAGNOSTIC) + + samples: Final = _samples(METRIC) + assert len(samples) == 1 + assert samples[0].labels == { + "requested_model": "per-second-priced-chat", + "model": "openai/gpt-5.4-nano", + "model_id": "dep-1", + "api_provider": "openai", + "reason": "missing_pricing_key", + } + assert samples[0].value == 1.0 + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_success_event_counts_a_zero_cost_request_by_model_and_reason() -> None: + _clear_prometheus_registry() + try: + logger: Final = PrometheusLogger() + await _log_success(logger, MISSING_KEY_DIAGNOSTIC) + await _log_success(logger, MISSING_KEY_DIAGNOSTIC) + + samples: Final = _samples(METRIC) + assert len(samples) == 1 + assert samples[0].labels == { + "requested_model": "per-second-priced-chat", + "model": "openai/gpt-5.4-nano", + "model_id": "dep-1", + "api_provider": "openai", + "reason": "missing_pricing_key", + } + assert samples[0].value == 2.0 + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_request_without_a_diagnostic_leaves_the_counter_untouched() -> None: + _clear_prometheus_registry() + try: + await _log_success(PrometheusLogger(), None) + + assert _samples(METRIC) == [] + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_label_filter_that_drops_reason_still_counts_the_request() -> None: + _clear_prometheus_registry() + previous_config: Final = litellm.prometheus_metrics_config + litellm.prometheus_metrics_config = [ + {"group": "zero_cost", "metrics": [METRIC], "include_labels": ["requested_model"]} + ] + try: + await _log_success(PrometheusLogger(), MISSING_KEY_DIAGNOSTIC) + + samples: Final = _samples(METRIC) + assert len(samples) == 1 + assert samples[0].labels == {"requested_model": "per-second-priced-chat"} + assert samples[0].value == 1.0 + finally: + litellm.prometheus_metrics_config = previous_config + _clear_prometheus_registry() diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py new file mode 100644 index 00000000000..0e453e3f5eb --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py @@ -0,0 +1,157 @@ +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.litellm_core_utils.llm_cost_calc.zero_cost_diagnostic import ( + ZERO_COST_COUNTER_NAME, + diagnose_zero_cost, + used_pricing_keys, + zero_cost_warning, +) +from litellm.types.utils import CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, Usage + +PER_SECOND_ENTRY: Final = {"input_cost_per_second": 0.00042, "output_cost_per_second": 0.00042} +FREE_ENTRY: Final = {"input_cost_per_token": 0, "output_cost_per_token": 0, "cache_read_input_token_cost": 2e-08} +PRICED_ENTRY: Final = {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06} +TEXT_USAGE: Final = Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + + +def test_missing_pricing_key_names_every_rate_the_usage_needs() -> None: + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PER_SECOND_ENTRY, calculation_failed=False + ) + + assert diagnostic == { + "reason": "missing_pricing_key", + "pricing_model": "dep-1", + "missing_pricing_keys": ("input_cost_per_token", "output_cost_per_token"), + } + + +def test_only_the_absent_rate_is_reported() -> None: + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry={"input_cost_per_token": 1e-06}, calculation_failed=False + ) + + assert diagnostic is not None + assert diagnostic["missing_pricing_keys"] == ("output_cost_per_token",) + + +def test_free_model_stays_silent() -> None: + assert ( + diagnose_zero_cost(usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=FREE_ENTRY, calculation_failed=False) + is None + ) + + +@pytest.mark.parametrize("calculation_failed", [False, True]) +def test_request_without_usage_stays_silent(calculation_failed: bool) -> None: + usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + + assert ( + diagnose_zero_cost( + usage=usage, pricing_model="dep-1", pricing_entry=PER_SECOND_ENTRY, calculation_failed=calculation_failed + ) + is None + ) + + +@pytest.mark.parametrize( + "entry", + [ + {"litellm_provider": "openai", "mode": "chat", "supports_prompt_caching": True}, + {"tiered_pricing": [{"range": [0, 128000], "input_cost_per_token": 0, "output_cost_per_token": 0}]}, + {"tiered_pricing": "not a tier table", "litellm_provider": "openai"}, + ], +) +def test_entry_that_declares_no_rate_stays_silent(entry: Mapping[str, object]) -> None: + assert ( + diagnose_zero_cost(usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=entry, calculation_failed=False) + is None + ) + + +def test_tiered_rate_counts_as_a_declared_rate() -> None: + entry = {"tiered_pricing": [{"range": [0, 128000], "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}]} + + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=entry, calculation_failed=False + ) + + assert diagnostic is not None + assert diagnostic["reason"] == "missing_pricing_key" + + +def test_priced_entry_that_still_prices_to_zero_is_pricing_not_applied() -> None: + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PRICED_ENTRY, calculation_failed=False + ) + + assert diagnostic == {"reason": "pricing_not_applied", "pricing_model": "dep-1", "missing_pricing_keys": ()} + + +def test_calculator_failure_on_a_priced_entry_is_cost_calculation_error() -> None: + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PRICED_ENTRY, calculation_failed=True + ) + + assert diagnostic == {"reason": "cost_calculation_error", "pricing_model": "dep-1", "missing_pricing_keys": ()} + + +def test_calculator_failure_on_a_free_entry_stays_silent() -> None: + assert ( + diagnose_zero_cost(usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=FREE_ENTRY, calculation_failed=True) + is None + ) + + +def test_calculator_failure_on_an_entry_that_declares_no_rate_stays_silent() -> None: + entry: Final = {"litellm_provider": "openai", "mode": "chat", "supports_prompt_caching": True} + assert ( + diagnose_zero_cost(usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=entry, calculation_failed=True) + is None + ) + + +def test_audio_tokens_need_the_audio_rates() -> None: + usage = Usage( + prompt_tokens=10, + completion_tokens=20, + total_tokens=30, + prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=10, text_tokens=0), + completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=5, text_tokens=15), + ) + + assert used_pricing_keys(usage) == ( + "input_cost_per_audio_token", + "output_cost_per_token", + "output_cost_per_audio_token", + ) + diagnostic = diagnose_zero_cost( + usage=usage, pricing_model="gemini-audio", pricing_entry=PRICED_ENTRY, calculation_failed=False + ) + assert diagnostic is not None + assert diagnostic["missing_pricing_keys"] == ("input_cost_per_audio_token", "output_cost_per_audio_token") + + +def test_warning_names_the_request_the_entry_the_missing_keys_and_the_counter() -> None: + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PER_SECOND_ENTRY, calculation_failed=False + ) + assert diagnostic is not None + + message = zero_cost_warning( + diagnostic, + model_group="per-second-priced-chat", + model="openai/gpt-5.4-nano", + custom_llm_provider="openai", + usage=TEXT_USAGE, + ) + + assert "model_group=per-second-priced-chat" in message + assert "model=openai/gpt-5.4-nano" in message + assert "provider=openai" in message + assert "prompt_tokens=10 completion_tokens=20" in message + assert "pricing entry 'dep-1' has no input_cost_per_token, output_cost_per_token" in message + assert f'{ZERO_COST_COUNTER_NAME}{{reason="missing_pricing_key"}}' in message diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 959b4f01986..a9bbb4992c6 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1,9 +1,10 @@ import asyncio import contextlib import datetime +import logging import os import sys -from collections.abc import Callable +from collections.abc import Callable, Iterator, Mapping from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -302,6 +303,398 @@ def test_response_cost_calculator_uses_router_model_id_from_litellm_metadata(): litellm.model_cost.pop(custom_model_id, None) +class TestZeroCostDiagnostic: + DEPLOYMENT_ID: Final = "lit7898-per-second-priced-deployment" + MODEL_GROUP: Final = "per-second-priced-chat" + PER_SECOND_PRICING: Final = {"input_cost_per_second": 0.00042, "output_cost_per_second": 0.00042} + FREE_PRICING: Final = {"input_cost_per_token": 0, "output_cost_per_token": 0} + + @pytest.fixture(params=["per_second", "free"]) + def deployment_pricing(self, request: pytest.FixtureRequest) -> Iterator[Mapping[str, float]]: + pricing: Final = self.PER_SECOND_PRICING if request.param == "per_second" else self.FREE_PRICING + litellm.register_model(model_cost={self.DEPLOYMENT_ID: pricing}, persist_across_reloads=False) + try: + yield pricing + finally: + litellm.model_cost.pop(self.DEPLOYMENT_ID, None) + + def _logging_obj( + self, + pricing: Mapping[str, object], + stream: bool = False, + model: str = "openai/gpt-5.4-nano", + call_type: str = "completion", + deployment_id: str | None = DEPLOYMENT_ID, + custom_llm_provider: str = "openai", + ) -> LitellmLogging: + logging_obj: Final = LitellmLogging( + model=model, + messages=[{"role": "user", "content": "Hi"}], + stream=stream, + call_type=call_type, + start_time=time.time(), + litellm_call_id="lit7898", + function_id="fn", + ) + self._route_to_deployment( + logging_obj, pricing, model=model, deployment_id=deployment_id, custom_llm_provider=custom_llm_provider + ) + return logging_obj + + def _route_to_deployment( + self, + logging_obj: LitellmLogging, + pricing: Mapping[str, object], + model: str = "openai/gpt-5.4-nano", + deployment_id: str | None = DEPLOYMENT_ID, + custom_llm_provider: str = "openai", + ) -> None: + model_info: Final = pricing if deployment_id is None else {"id": deployment_id, **pricing} + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params={"metadata": {"model_group": self.MODEL_GROUP, "model_info": model_info}}, + custom_llm_provider=custom_llm_provider, + ) + + @staticmethod + def _response( + usage: litellm.Usage | None = None, model: str = "gpt-5.4-nano", **hidden_params: object + ) -> ModelResponse: + response: Final = ModelResponse( + model=model, + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="hello"))], + usage=usage, + ) + response._hidden_params = {"custom_llm_provider": "openai", **hidden_params} + return response + + @staticmethod + def _zero_cost_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.name == "LiteLLM" and record.levelno == logging.WARNING and "priced at $0" in record.getMessage() + ] + + def _assert_flagged(self, logging_obj: LitellmLogging, caplog: pytest.LogCaptureFixture) -> None: + assert logging_obj.model_call_details["zero_cost_diagnostic"] == { + "reason": "missing_pricing_key", + "pricing_model": self.DEPLOYMENT_ID, + "missing_pricing_keys": ("input_cost_per_token", "output_cost_per_token"), + } + warnings: Final = self._zero_cost_warnings(caplog) + assert len(warnings) == 1 + assert f"model_group={self.MODEL_GROUP}" in warnings[0] + assert f"pricing entry '{self.DEPLOYMENT_ID}' has no input_cost_per_token, output_cost_per_token" in warnings[0] + + def test_zero_cost_with_a_missing_rate_warns_once_and_is_recorded( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + first_cost: Final = logging_obj._response_cost_calculator(result=self._response(usage)) + second_cost: Final = logging_obj._response_cost_calculator(result=self._response(usage)) + + assert first_cost == 0.0 + assert second_cost == 0.0 + if deployment_pricing is self.FREE_PRICING: + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + return + self._assert_flagged(logging_obj, caplog) + + def test_usage_less_stream_chunk_does_not_hide_the_final_response_diagnostic( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=8, completion_tokens=2, total_tokens=10) + logging_obj: Final = self._logging_obj(deployment_pricing, stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._response_cost_calculator(result=self._response(usage=None)) + logging_obj._response_cost_calculator(result=self._response(usage)) + + if deployment_pricing is self.FREE_PRICING: + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + return + self._assert_flagged(logging_obj, caplog) + + def test_terminal_responses_stream_event_is_judged_by_its_inner_response( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + logging_obj: Final = self._logging_obj(deployment_pricing, stream=True, call_type="aresponses") + event: Final = ResponseCompletedEvent( + type="response.completed", + response=ResponsesAPIResponse( + id="resp-lit7898", + created_at=1, + object="response", + status="completed", + model="gpt-5.4-nano", + output=[], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=20, total_tokens=30), + ), + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + cost: Final = logging_obj._response_cost_calculator(result=event) + + assert cost == 0.0 + if deployment_pricing is self.FREE_PRICING: + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + return + self._assert_flagged(logging_obj, caplog) + + def test_precomputed_zero_hidden_cost_is_flagged_and_lands_in_the_payload( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing) + response: Final = self._response(usage, response_cost=0.0, model_id=self.DEPLOYMENT_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._process_hidden_params_and_response_cost( + response, start_time=datetime.datetime.now(), end_time=datetime.datetime.now() + ) + + payload: Final = logging_obj.model_call_details["standard_logging_object"] + assert payload["response_cost"] == 0.0 + if deployment_pricing is self.FREE_PRICING: + assert payload["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + return + self._assert_flagged(logging_obj, caplog) + assert payload["zero_cost_diagnostic"] == logging_obj.model_call_details["zero_cost_diagnostic"] + + def test_uncomputed_hidden_cost_is_not_a_zero_cost( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing) + response: Final = self._response(usage, response_cost=None, model_id=self.DEPLOYMENT_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._process_hidden_params_and_response_cost( + response, start_time=datetime.datetime.now(), end_time=datetime.datetime.now() + ) + + assert logging_obj.model_call_details["standard_logging_object"]["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + + def test_unbilled_read_route_with_usage_stays_silent( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing, call_type="aget_responses") + response: Final = self._response(usage, response_cost=0.0, model_id=self.DEPLOYMENT_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._process_hidden_params_and_response_cost( + response, start_time=datetime.datetime.now(), end_time=datetime.datetime.now() + ) + + assert logging_obj.model_call_details["standard_logging_object"]["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + + def test_unmapped_model_that_fails_cost_calculation_stays_silent(self, caplog: pytest.LogCaptureFixture) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj( + {}, model="openai/lit7898-unmapped-model", deployment_id="lit7898-unmapped-deployment" + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + cost: Final = logging_obj._response_cost_calculator( + result=self._response(usage, model="lit7898-unmapped-model") + ) + + assert cost is None + assert logging_obj.model_call_details["response_cost_failure_debug_information"] is not None + assert logging_obj.model_call_details.get("zero_cost_diagnostic") is None + assert self._zero_cost_warnings(caplog) == [] + + def test_malformed_usage_never_raises_out_of_the_cost_calculator( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + logging_obj: Final = self._logging_obj(deployment_pricing) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + cost: Final = logging_obj._response_cost_calculator( + result={"model": "gpt-5.4-nano", "usage": {"prompt_tokens": "n/a", "completion_tokens": 3}} + ) + + assert cost is None + assert logging_obj.model_call_details.get("zero_cost_diagnostic") is None + assert self._zero_cost_warnings(caplog) == [] + + def test_usage_less_evaluation_between_two_zero_cost_findings_does_not_warn_twice( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=8, completion_tokens=2, total_tokens=10) + logging_obj: Final = self._logging_obj(deployment_pricing, stream=True, call_type="anthropic_messages") + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._response_cost_calculator(result=self._response(usage=None)) + logging_obj._response_cost_calculator(result=self._response(usage)) + logging_obj._response_cost_calculator(result=self._response(usage=None)) + logging_obj._response_cost_calculator(result=self._response(usage)) + + if deployment_pricing is self.FREE_PRICING: + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + return + self._assert_flagged(logging_obj, caplog) + + def test_retry_that_prices_clears_the_diagnostic_and_a_later_zero_cost_is_recorded_silently( + self, caplog: pytest.LogCaptureFixture + ) -> None: + priced_id: Final = "lit7898-priced-deployment" + priced_pricing: Final = {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06} + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + litellm.register_model( + model_cost={self.DEPLOYMENT_ID: self.PER_SECOND_PRICING, priced_id: priced_pricing}, + persist_across_reloads=False, + ) + try: + logging_obj: Final = self._logging_obj(self.PER_SECOND_PRICING) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert logging_obj._response_cost_calculator(result=self._response(usage)) == 0.0 + self._assert_flagged(logging_obj, caplog) + + self._route_to_deployment(logging_obj, priced_pricing, deployment_id=priced_id) + assert logging_obj._response_cost_calculator(result=self._response(usage)) == pytest.approx(5e-05) + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + + self._route_to_deployment(logging_obj, self.PER_SECOND_PRICING) + assert logging_obj._response_cost_calculator(result=self._response(usage)) == 0.0 + + assert logging_obj.model_call_details["zero_cost_diagnostic"]["reason"] == "missing_pricing_key" + assert len(self._zero_cost_warnings(caplog)) == 1 + finally: + litellm.model_cost.pop(self.DEPLOYMENT_ID, None) + litellm.model_cost.pop(priced_id, None) + + def test_one_request_evaluated_against_two_cost_map_entries_warns_once( + self, caplog: pytest.LogCaptureFixture + ) -> None: + dated_model: Final = "lit7898-nano-2026-03-17" + requested_model: Final = "lit7898-nano" + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + cost_map_entry: Final = {"litellm_provider": "openai", "mode": "chat", **self.PER_SECOND_PRICING} + litellm.register_model( + model_cost={dated_model: cost_map_entry, requested_model: cost_map_entry}, persist_across_reloads=False + ) + try: + logging_obj: Final = self._logging_obj( + {}, model=f"openai/{requested_model}", deployment_id="lit7898-cost-map-deployment" + ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert logging_obj._response_cost_calculator(result=self._response(usage, model=dated_model)) == 0.0 + assert logging_obj._response_cost_calculator(result=self._response(usage, model=requested_model)) == 0.0 + + assert logging_obj.model_call_details["zero_cost_diagnostic"]["pricing_model"] == requested_model + warnings: Final = self._zero_cost_warnings(caplog) + assert len(warnings) == 1 + assert f"pricing entry '{dated_model}' has no input_cost_per_token, output_cost_per_token" in warnings[0] + finally: + litellm.model_cost.pop(dated_model, None) + litellm.model_cost.pop(requested_model, None) + + def test_free_deployment_without_a_router_id_is_judged_by_its_own_pricing( + self, caplog: pytest.LogCaptureFixture + ) -> None: + global_model: Final = "lit7898-priced-global" + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + litellm.register_model( + model_cost={ + global_model: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + try: + logging_obj: Final = self._logging_obj(self.FREE_PRICING, model=global_model, deployment_id=None) + response: Final = self._response(usage, model=global_model, response_cost=0.0) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._process_hidden_params_and_response_cost( + response, start_time=datetime.datetime.now(), end_time=datetime.datetime.now() + ) + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + finally: + litellm.model_cost.pop(global_model, None) + + def test_cache_hit_priced_for_saved_cost_stays_silent( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing) + logging_obj.model_call_details["cache_hit"] = True + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert logging_obj._response_cost_calculator(result=self._response(usage), cache_hit=False) == 0.0 + + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + + @pytest.mark.parametrize("spilled_over", [True, False]) + def test_ptu_deployment_is_judged_by_the_entry_the_calculator_priced_with( + self, spilled_over: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + router_model_id: Final = "lit7898-ptu-router-model-id" + served_model: Final = "azure/lit7898-ptu-served-model" + ptu_model_info: Final = { + "team_id": "team-1", + "ptu_count": 100, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01", + **self.FREE_PRICING, + } + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + litellm.register_model( + model_cost={ + router_model_id: {**self.FREE_PRICING, "litellm_provider": "azure", "mode": "chat"}, + served_model: {**self.PER_SECOND_PRICING, "litellm_provider": "azure", "mode": "chat"}, + }, + persist_across_reloads=False, + ) + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + try: + logging_obj: Final = self._logging_obj( + ptu_model_info, model=served_model, deployment_id=router_model_id, custom_llm_provider="azure" + ) + spillover_headers: Final = {"llm_provider-x-ms-is-spilled-over": "true"} if spilled_over else {} + response: Final = self._response( + usage, model=served_model, custom_llm_provider="azure", additional_headers=spillover_headers + ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert logging_obj._response_cost_calculator(result=response) == 0.0 + + warnings: Final = self._zero_cost_warnings(caplog) + if not spilled_over: + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert warnings == [] + return + assert logging_obj.model_call_details["zero_cost_diagnostic"] == { + "reason": "missing_pricing_key", + "pricing_model": served_model, + "missing_pricing_keys": ("input_cost_per_token", "output_cost_per_token"), + } + assert len(warnings) == 1 + assert f"pricing entry '{served_model}' has no input_cost_per_token, output_cost_per_token" in warnings[0] + finally: + litellm.model_cost.pop(router_model_id, None) + litellm.model_cost.pop(served_model, None) + + class TestGetRouterModelId: """Tests for the get_router_model_id helper method.""" @@ -407,7 +800,6 @@ class TestGetRouterDeploymentModelInfo: logging_obj.litellm_params = {"api_base": ""} assert logging_obj.get_router_deployment_model_info() is None - def test_a_published_batch_rate_never_displaces_a_declared_standard_rate(self) -> None: """Ownership is per token direction, not per field. @@ -1111,7 +1503,9 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch): @pytest.mark.asyncio -async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(monkeypatch: pytest.MonkeyPatch): +async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log( + monkeypatch: pytest.MonkeyPatch, +): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.responses.main import base_llm_http_handler @@ -7066,22 +7460,41 @@ async def test_classifier_audit_matches_provider_transport(provider: str) -> Non return httpx.Response(200, json=mock_responses_api_response(content).model_dump()) if provider == "anthropic": - return httpx.Response(200, json={ - "id": "msg-audit", "type": "message", "role": "assistant", "model": "claude-haiku-4-5", - "content": [{"type": "text", "text": content}], "stop_reason": "end_turn", - "usage": {"input_tokens": 10, "output_tokens": 5}, - }) + return httpx.Response( + 200, + json={ + "id": "msg-audit", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": content}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) if provider == "bedrock": - return httpx.Response(200, json={ - "output": {"message": {"role": "assistant", "content": [{"text": content}]}}, - "stopReason": "end_turn", "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, - "metrics": {"latencyMs": 1}, - }) - return httpx.Response(200, json={ - "id": "chatcmpl-audit", "object": "chat.completion", "created": 0, "model": "gpt-5.6", - "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - }) + return httpx.Response( + 200, + json={ + "output": {"message": {"role": "assistant", "content": [{"text": content}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + "metrics": {"latencyMs": 1}, + }, + ) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-audit", + "object": "chat.completion", + "created": 0, + "model": "gpt-5.6", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) async def capture(kwargs, response_obj, start_time, end_time): logs.put_nowait(kwargs["standard_logging_object"]) @@ -7092,11 +7505,15 @@ async def test_classifier_audit_matches_provider_transport(provider: str) -> Non handler.client = http_client client: Final = ( AsyncAzureOpenAI( - api_key="transport-only", azure_endpoint="https://azure.invalid", - api_version="2025-04-01-preview", http_client=http_client, + api_key="transport-only", + azure_endpoint="https://azure.invalid", + api_version="2025-04-01-preview", + http_client=http_client, ) - if provider == "azure" else AsyncOpenAI(api_key="transport-only", http_client=http_client) - if provider == "openai" else handler + if provider == "azure" + else AsyncOpenAI(api_key="transport-only", http_client=http_client) + if provider == "openai" + else handler ) model: Final = { "openai": "openai/gpt-5.6", @@ -7109,23 +7526,44 @@ async def test_classifier_audit_matches_provider_transport(provider: str) -> Non async def run(marker: str) -> None: if provider == "responses": await litellm.aresponses( - model=model, api_key="transport-only", client=client, max_output_tokens=128, - instructions="classifier-rubric", input=marker, + model=model, + api_key="transport-only", + client=client, + max_output_tokens=128, + instructions="classifier-rubric", + input=marker, metadata={"internal_call_origin": "autorouter_classifier"}, proxy_server_request={"body": {}, "originating_request_masked": {"input": f"source-only-{marker}"}}, - success_callback=[capture], num_retries=0, + success_callback=[capture], + num_retries=0, ) return await litellm.acompletion( - model=model, api_key="transport-only", client=client, max_tokens=128, - aws_access_key_id="transport-only", aws_secret_access_key="transport-only", aws_region_name="us-east-1", + model=model, + api_key="transport-only", + client=client, + max_tokens=128, + aws_access_key_id="transport-only", + aws_secret_access_key="transport-only", + aws_region_name="us-east-1", messages=[{"role": "system", "content": "classifier-rubric"}, {"role": "user", "content": marker}], metadata={"internal_call_origin": "autorouter_classifier"}, proxy_server_request={"body": {}, "originating_request_masked": {"input": f"source-only-{marker}"}}, - success_callback=[capture], num_retries=0, - **({"api_base": "https://azure.invalid", "api_version": "2025-04-01-preview"} if provider == "azure" else {}), - **({"extra_body": {"audit_context": "provider-extra"}, "extra_headers": {"X-Audit": "header-only-secret"}} - if provider in ("openai", "azure") else {}), + success_callback=[capture], + num_retries=0, + **( + {"api_base": "https://azure.invalid", "api_version": "2025-04-01-preview"} + if provider == "azure" + else {} + ), + **( + { + "extra_body": {"audit_context": "provider-extra"}, + "extra_headers": {"X-Audit": "header-only-secret"}, + } + if provider in ("openai", "azure") + else {} + ), ) await asyncio.gather(run("request-one"), run("request-two")) @@ -7149,14 +7587,17 @@ async def test_classifier_audit_matches_provider_transport(provider: str) -> Non @pytest.mark.parametrize("redaction", ["none", "global", "request", "header"]) @pytest.mark.parametrize("status", ["success", "failure"]) @pytest.mark.parametrize("call_type", ["completion", "acompletion", "responses", "aresponses"]) -def test_classifier_audit_obeys_message_logging_before_payload_emission(logging_obj, monkeypatch, redaction, status, call_type): +def test_classifier_audit_obeys_message_logging_before_payload_emission( + logging_obj, monkeypatch, redaction, status, call_type +): from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload monkeypatch.setattr(litellm, "turn_off_message_logging", redaction == "global") params: Final = { - "metadata": {"internal_call_origin": "autorouter_classifier", **( - {"headers": {"x-litellm-enable-message-redaction": "true"}} if redaction == "header" else {} - )}, + "metadata": { + "internal_call_origin": "autorouter_classifier", + **({"headers": {"x-litellm-enable-message-redaction": "true"}} if redaction == "header" else {}), + }, "proxy_server_request": {"body": {}, "originating_request_masked": {"input": "source-only"}}, } logging_obj.call_type = call_type @@ -7169,8 +7610,12 @@ def test_classifier_audit_obeys_message_logging_before_payload_emission(logging_ ) now: Final = datetime.datetime.now() payload: Final = get_standard_logging_object_payload( - kwargs={**logging_obj.model_call_details, "call_type": call_type}, init_response_obj={}, - start_time=now, end_time=now, logging_obj=logging_obj, status=status, + kwargs={**logging_obj.model_call_details, "call_type": call_type}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status=status, ) assert payload is not None if redaction == "none": @@ -7421,7 +7866,13 @@ def _completed_responses_event(usage: ResponseAPIUsage) -> ResponseCompletedEven return ResponseCompletedEvent( type="response.completed", response=ResponsesAPIResponse( - id="resp-1", created_at=1, object="response", status="completed", model="codex-mini-latest", output=[], usage=usage + id="resp-1", + created_at=1, + object="response", + status="completed", + model="codex-mini-latest", + output=[], + usage=usage, ), ) @@ -7441,7 +7892,9 @@ def test_get_assembled_streaming_response_bills_a_provider_reported_usage_cost() now = datetime.datetime.now() assembled = logging_obj._get_assembled_streaming_response( - result=_completed_responses_event(ResponseAPIUsage(input_tokens=12, output_tokens=2, total_tokens=14, cost=0.0042)), + result=_completed_responses_event( + ResponseAPIUsage(input_tokens=12, output_tokens=2, total_tokens=14, cost=0.0042) + ), start_time=now, end_time=now, is_async=True, From e7cd97c6b6ade3afa4121de9f424ca12924b8b5b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 19:51:12 -0700 Subject: [PATCH 2/3] fix(proxy): release unclaimed budget reservations at request end (#42304) * fix(proxy): release unclaimed budget reservations at request end * fix(proxy): release unclaimed budget reservations of websocket sessions too * test(proxy): drop the structural middleware inheritance check * fix(proxy): claim the budget reservation on streaming pass-through before its cost callback The SSE chunk processor hands its success handler to the logging worker after the response, so the request-end release freed the reservation first and left the key unguarded until the worker drained. Claim it at both end-of-stream hand-offs, the immediate enqueue and the coroutine parked for deferred dispatch. Give the xai realtime test double the litellm_params attribute every real Logging object carries, since the wrapper now reads it. * test(pass-through): give the vertex streaming test doubles a litellm_params dict The spec'd Logging mocks in test_vertex_ai_anthropic_streaming_cost_injection.py lacked the instance attribute the chunk processor now reads to claim the budget reservation. Also restores main's _lazy_openapi_snapshot.json: the branch's copy had been regenerated under Python 3.14, which dedents one docstring description that the CI regeneration on Python 3.12 keeps indented, and the PR adds no lazily loaded route, so main's file is the correct one. * fix(pass-through): claim the budget reservation only after its cost callback is enqueued Every pass-through success hand-off stamped callback_bound before handing the coroutine to the logging worker. When that enqueue raised, the reservation stayed claimed with no callback left to reconcile it, so the request-end release skipped it and the reserved cost stayed pinned on the key's counter. Enqueue first, then claim, so a failed hand-off leaves the reservation for the request-end release. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 40 ++ litellm/litellm_core_utils/litellm_logging.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 48 +-- .../proxy/hooks/proxy_track_cost_callback.py | 13 +- .../budget_reservation_release_middleware.py | 33 ++ .../pass_through_endpoints.py | 4 + .../streaming_handler.py | 4 + litellm/proxy/proxy_server.py | 9 +- .../spend_tracking/budget_reservation.py | 14 + .../rust_bridge/callbacks_legacy_python.py | 19 +- litellm/utils.py | 9 +- ...x_ai_anthropic_streaming_cost_injection.py | 4 + .../litellm_core_utils/test_core_helpers.py | 55 +++ .../test_litellm_logging.py | 17 + .../llms/xai/test_xai_key_fallback.py | 3 + .../proxy/auth/test_user_api_key_auth.py | 88 +++++ ...t_budget_reservation_release_middleware.py | 349 ++++++++++++++++++ .../test_pass_through_endpoints.py | 106 ++++++ .../test_streaming_handler_interrupt.py | 69 ++++ .../spend_tracking/test_budget_reservation.py | 39 ++ .../test_callbacks_legacy_python.py | 79 +++- tests/test_litellm/test_utils.py | 94 +++++ 22 files changed, 1057 insertions(+), 41 deletions(-) create mode 100644 litellm/proxy/middleware/budget_reservation_release_middleware.py create mode 100644 tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index ce6f77f78a0..5bcde688521 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -359,6 +359,46 @@ def get_litellm_metadata_from_kwargs(kwargs: dict): return {} +def _budget_reservation_on_auth_object(user_api_key_auth: object) -> object: + if isinstance(user_api_key_auth, Mapping): + return user_api_key_auth.get("budget_reservation") + return getattr(user_api_key_auth, "budget_reservation", None) + + +def budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict | None: + stamped: Final = metadata.get("user_api_key_budget_reservation") + if isinstance(stamped, dict): + return stamped + on_auth_object: Final = _budget_reservation_on_auth_object(metadata.get("user_api_key_auth")) + return on_auth_object if isinstance(on_auth_object, dict) else None + + +def _stamp_budget_reservation_callback_bound(litellm_params: Mapping[str, object], callback_bound: bool) -> None: + for metadata_variable_name in ("metadata", "litellm_metadata"): + metadata = litellm_params.get(metadata_variable_name) + if not isinstance(metadata, Mapping): + continue + budget_reservation = budget_reservation_from_metadata(metadata) + if budget_reservation is not None: + budget_reservation["callback_bound"] = callback_bound + + +def bind_budget_reservation_to_callbacks(litellm_params: Mapping[str, object]) -> None: + """Mark the request's budget reservation as owned by the success callbacks of this call. + + The proxy releases any reservation still unbound when the request ends; one bound here + is left for the cost callback, which may finish after the response has been sent. Bind + only where a success handler is guaranteed to run: a logging object merely existing is + not that, since the proxy builds one for every route before calling anything. + """ + _stamp_budget_reservation_callback_bound(litellm_params, True) + + +def unbind_budget_reservation_from_callbacks(litellm_params: Mapping[str, object]) -> None: + """Hand a failed call's reservation back to the request-end release: failure handlers never settle it.""" + _stamp_budget_reservation_callback_bound(litellm_params, False) + + def reconstruct_model_name( model_name: str, custom_llm_provider: str | None, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a486cdeff19..f5967185af0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5493,7 +5493,7 @@ def request_model_access_groups_from_litellm_params(litellm_params: Mapping[str, """Access groups the auth layer stamped onto this request, from whichever metadata field carries them. Detached internal sub-calls only inherit the identity keys, so the auth object is the - fallback there, exactly as _get_budget_reservation_from_metadata does for reservations. + fallback there, exactly as budget_reservation_from_metadata does for reservations. """ for metadata_variable_name in ("metadata", "litellm_metadata"): metadata = litellm_params.get(metadata_variable_name) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 08f64e610bd..a6c0792a86f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -650,6 +650,7 @@ async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str "type": "http", "headers": scope_headers, "path": ws_scope.get("path", ""), + "state": ws_scope.setdefault("state", {}), # mutable-ok: Starlette's socket state, shared with the request } for key in ("root_path", "app_root_path"): if key in ws_scope: @@ -3086,31 +3087,30 @@ async def _reserve_budget_after_common_checks( request: Request | None = None, ) -> None: user_api_key_auth_obj.budget_reservation = None - if skip_budget_checks: - return - if general_settings.get("disable_budget_reservation") is True: - return + if not skip_budget_checks and general_settings.get("disable_budget_reservation") is not True: + from litellm.proxy.spend_tracking.budget_reservation import ( + reserve_budget_for_request, + ) - from litellm.proxy.spend_tracking.budget_reservation import ( - reserve_budget_for_request, - ) - - user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request( - request_body=request_data, - route=route, - llm_router=llm_router, - valid_token=user_api_key_auth_obj, - team_object=team_object, - user_object=user_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - end_user_id=end_user_id, - end_user_object=end_user_object, - apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, - fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, - raw_body=await read_raw_json_body(request=request), - ) + user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request( + request_body=request_data, + route=route, + llm_router=llm_router, + valid_token=user_api_key_auth_obj, + team_object=team_object, + user_object=user_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + end_user_id=end_user_id, + end_user_object=end_user_object, + apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, + fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, + raw_body=await read_raw_json_body(request=request), + ) + if request is not None: + reservation: Final = user_api_key_auth_obj.budget_reservation + request.state.budget_reservation = reservation # rebind-ok: read by the release middleware def _should_skip_budget_checks( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 08e8e4f8c10..81894a5ff12 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -11,6 +11,7 @@ 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, + budget_reservation_from_metadata, get_litellm_metadata_from_kwargs, ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -630,17 +631,7 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]: 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) + return budget_reservation_from_metadata(metadata) def _get_request_tags_for_cost_tracking( diff --git a/litellm/proxy/middleware/budget_reservation_release_middleware.py b/litellm/proxy/middleware/budget_reservation_release_middleware.py new file mode 100644 index 00000000000..f7ac885274e --- /dev/null +++ b/litellm/proxy/middleware/budget_reservation_release_middleware.py @@ -0,0 +1,33 @@ +from collections.abc import Awaitable, Callable, Mapping +from typing import Final + +from starlette.types import ASGIApp, Receive, Scope, Send + +_SCOPES_AUTH_STAMPS: Final = frozenset({"http", "websocket"}) + + +class BudgetReservationReleaseMiddleware: + """Releases the budget reservation auth made for a request once no callback owns it. + + Auth stamps the reservation on the request or socket state; a call that starts + claims it for the cost callbacks, which settle it on success or failure. When the + response has been sent or the socket has closed and the reservation is still + unclaimed, nothing else ever would, so it is released here instead of pinning the + spend counter until its TTL. + """ + + def __init__(self, app: ASGIApp, release: Callable[[Mapping[str, object]], Awaitable[None]]) -> None: + self.app = app + self.release = release + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] not in _SCOPES_AUTH_STAMPS: + await self.app(scope, receive, send) + return + try: + await self.app(scope, receive, send) + finally: + state: Final = scope.get("state") + budget_reservation: Final = state.get("budget_reservation") if isinstance(state, Mapping) else None + if isinstance(budget_reservation, Mapping): + await self.release(budget_reservation) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index bbbf5e4b94f..8902e599788 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -47,6 +47,7 @@ from litellm.constants import ( from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( + bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, ) @@ -1637,6 +1638,7 @@ async def pass_through_request( **kwargs, ) ) + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) ## CUSTOM HEADERS - `x-litellm-*` custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -2561,6 +2563,7 @@ async def websocket_passthrough_request( **success_kwargs, ) ) + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) # Call the proxy logging success hook if proxy_logging_obj: @@ -2732,6 +2735,7 @@ async def _relay_passthrough_response_bytes( **success_handler_kwargs, ) ) + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) def _extract_model_from_vertex_ai_setup(setup_response: Mapping[str, object]) -> str | None: diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index fe9e104789b..fae26b5a72d 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -9,6 +9,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy._types import PassThroughEndpointLoggingResultValues @@ -218,6 +219,7 @@ class PassThroughStreamingHandler: and response.status_code < 400 ): logging_scheduled = True + bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params) litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),) except Exception as e: verbose_proxy_logger.error("Error in chunk_processor: %s", e) @@ -250,6 +252,8 @@ class PassThroughStreamingHandler: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine()) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) + else: + bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params) @staticmethod async def _route_streaming_logging_to_handler( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 87da8a6e44e..6fa7a7d761f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -654,6 +654,9 @@ from litellm.proxy.middleware.billable_request_metrics_middleware import ( BillableRequestMetricsMiddleware, BillingRecorder, ) +from litellm.proxy.middleware.budget_reservation_release_middleware import ( + BudgetReservationReleaseMiddleware, +) from litellm.proxy.plugin_routes import ( register_plugins_from_config, ) @@ -729,7 +732,10 @@ from litellm.proxy.shutdown.scheduled_jobs import ( pause_scheduled_jobs, stop_in_flight_scheduler_jobs, ) -from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.budget_reservation import ( + get_budget_window_start, + release_unbound_budget_reservation, +) from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( run_scheduled_daily_global_spend_reconcile, ) @@ -2358,6 +2364,7 @@ app.add_middleware( # it sees prisma_client as of the first request rather than import time. sink_factory=lambda: gateway_request_accumulator if prisma_client is not None else None, ) +app.add_middleware(BudgetReservationReleaseMiddleware, release=release_unbound_budget_reservation) app.add_middleware(InFlightRequestsMiddleware) app.add_middleware(SecurityHeadersMiddleware) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index f9e5c4ff1e4..ac9b07a55f7 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -366,6 +366,7 @@ async def reserve_budget_for_request( "reserved_cost": reservation_cost, "entries": applied_entries, "finalized": False, + "callback_bound": False, "input_cost": min(float(input_cost or 0.0), reservation_cost), "input_tokens": max(input_token_counts.values(), default=None), } @@ -474,6 +475,19 @@ async def release_or_invalidate_budget_reservation( budget_reservation["finalized"] = True +async def release_unbound_budget_reservation(budget_reservation: Mapping[str, object]) -> None: + """Release a reservation no logging callback took ownership of, once the request ended. + + A handler whose litellm call never builds a logging object (batch cancel, file + content, anything without the client decorator) runs no cost callback, so nothing + else would ever reconcile its reservation. A bound reservation is left alone: its + success or failure handler settles it, possibly after the response has been sent. + """ + if not isinstance(budget_reservation, dict) or budget_reservation.get("callback_bound") is True: + return + await release_or_invalidate_budget_reservation(budget_reservation=budget_reservation) + + async def _get_budget_counters( request_body: dict, valid_token: UserAPIKeyAuth, diff --git a/litellm/rust_bridge/callbacks_legacy_python.py b/litellm/rust_bridge/callbacks_legacy_python.py index 30aa1d97bfc..6bbf2ffed6b 100644 --- a/litellm/rust_bridge/callbacks_legacy_python.py +++ b/litellm/rust_bridge/callbacks_legacy_python.py @@ -59,9 +59,17 @@ def setup( } supplied: Final = arguments.get("litellm_logging_obj") if isinstance(supplied, Logging): - return CallSetup(supplied, arguments) + return _claim_budget_reservation(CallSetup(supplied, arguments), asynchronous) logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments) - return CallSetup(logger, prepared) + return _claim_budget_reservation(CallSetup(logger, prepared), asynchronous) + + +def _claim_budget_reservation(call_setup: CallSetup, asynchronous: bool) -> CallSetup: + from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks + + if asynchronous and not is_internal_call(): + bind_budget_reservation_to_callbacks(call_setup.logger.litellm_params) + return call_setup def check_limits(kwargs: Mapping[str, object]) -> None: @@ -96,6 +104,9 @@ def finalize( class LoggingSurface(Protocol): + @property + def litellm_params(self) -> Mapping[str, object]: ... + def update_from_kwargs( self, kwargs: dict[str, object], @@ -236,8 +247,12 @@ def sync_success_for_async_call( def failure_handler( logger: LoggingSurface, error: Exception, start: datetime.datetime, end: datetime.datetime, asynchronous: bool ) -> Coroutine[object, object, None] | None: + from litellm.litellm_core_utils.core_helpers import unbind_budget_reservation_from_callbacks + trace: Final = "".join(traceback.format_exception(error)) if asynchronous: + if not is_internal_call(): + unbind_budget_reservation_from_callbacks(logger.litellm_params) return logger.async_failure_handler(error, trace, start, end) logger.failure_handler(error, trace, start, end) return None diff --git a/litellm/utils.py b/litellm/utils.py index d96a1c5ae3f..812299560c7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -81,7 +81,11 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import ( + bind_budget_reservation_to_callbacks, + normalize_drop_params, + unbind_budget_reservation_from_callbacks, +) from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, @@ -1880,6 +1884,8 @@ def client(original_function): # Type assertion: logging_obj is guaranteed to be non-None after function_setup assert logging_obj is not None, "logging_obj should not be None after function_setup" + if not _is_litellm_internal_call: + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) kwargs["litellm_logging_obj"] = logging_obj modified_kwargs: Final = await async_pre_call_deployment_hook(kwargs, call_type) @@ -2081,6 +2087,7 @@ def client(original_function): # the failure hook ran, so a slow callback doesn't inflate the reported duration. end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with if logging_obj and not _is_litellm_internal_call: + unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) try: logging_obj.failure_handler( e, traceback_exception, start_time, end_time diff --git a/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py b/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py index 498f0a734a3..ab5fd7f80ee 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py @@ -53,6 +53,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_enabled(): # Setup logging object with model info litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() @@ -132,6 +133,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_disabled(): response.aiter_bytes = mock_aiter_bytes litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() @@ -194,6 +196,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_no_usage_chunk(): response.aiter_bytes = mock_aiter_bytes litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() @@ -249,6 +252,7 @@ async def test_vertex_ai_anthropic_streaming_model_extraction(): response.aiter_bytes = mock_aiter_bytes litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index b2ad13c205e..bca61a0e76f 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -6,6 +6,8 @@ import pytest from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, + bind_budget_reservation_to_callbacks, + budget_reservation_from_metadata, drop_params_env_flag, drop_params_flag, get_or_create_metadata_bucket, @@ -13,7 +15,60 @@ from litellm.litellm_core_utils.core_helpers import ( normalize_drop_params, reconstruct_model_name, redact_nested_match_and_regex_keys, + unbind_budget_reservation_from_callbacks, ) +from litellm.proxy._types import UserAPIKeyAuth + + +class TestBudgetReservationBinding: + """The request-end release skips a reservation a cost callback has claimed, so the claim + must land on the one dict auth stamped, through whichever metadata field or auth object + carries it, and a failed call must be able to hand it back.""" + + @staticmethod + def _reservation() -> dict: + return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + @pytest.mark.parametrize("metadata_variable_name", ["metadata", "litellm_metadata"]) + def test_reservation_stamped_on_the_metadata_is_bound(self, metadata_variable_name: str): + reservation = self._reservation() + + bind_budget_reservation_to_callbacks({metadata_variable_name: {"user_api_key_budget_reservation": reservation}}) + + assert reservation["callback_bound"] is True + + def test_reservation_reachable_only_through_the_auth_object_is_bound(self): + reservation = self._reservation() + user_api_key_auth = UserAPIKeyAuth(token="hashed") + user_api_key_auth.budget_reservation = reservation + + bind_budget_reservation_to_callbacks({"metadata": {"user_api_key_auth": user_api_key_auth}}) + + assert reservation["callback_bound"] is True + + def test_reservation_reachable_only_through_a_dumped_auth_object_is_bound(self): + reservation = self._reservation() + + bind_budget_reservation_to_callbacks({"metadata": {"user_api_key_auth": {"budget_reservation": reservation}}}) + + assert reservation["callback_bound"] is True + + def test_unbind_hands_a_claimed_reservation_back(self): + reservation = self._reservation() + litellm_params = {"litellm_metadata": {"user_api_key_budget_reservation": reservation}} + bind_budget_reservation_to_callbacks(litellm_params) + + unbind_budget_reservation_from_callbacks(litellm_params) + + assert reservation["callback_bound"] is False + + def test_request_without_a_reservation_binds_nothing(self): + metadata = {"user_api_key_auth": UserAPIKeyAuth(token="hashed")} + + bind_budget_reservation_to_callbacks({"metadata": metadata, "litellm_metadata": None}) + + assert budget_reservation_from_metadata(metadata) is None + assert "user_api_key_budget_reservation" not in metadata class TestGetOrCreateMetadataBucket: diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index a9bbb4992c6..1bfc55d1486 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -26,6 +26,7 @@ from litellm.litellm_core_utils.litellm_logging import ( set_callbacks, ) from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, @@ -7941,3 +7942,19 @@ def test_response_cost_calculator_prices_terminal_responses_event_from_its_respo assert event_cost is not None and event_cost > 0 assert event_cost == inner_cost assert logging_obj.cost_breakdown["input_cost"] is not None and logging_obj.cost_breakdown["input_cost"] > 0 + + +class TestBudgetReservationBinding: + """The proxy builds a logging object for every route before calling anything, so a + logging object seeing the reservation is no promise that a cost callback will settle + it: the claim belongs to the call wrapper, and this object must leave it unbound.""" + + def test_update_environment_variables_leaves_the_reservation_unbound(self, logging_obj): + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + logging_obj.update_environment_variables( + litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}, optional_params={} + ) + + assert logging_obj.litellm_params["metadata"]["user_api_key_budget_reservation"] is reservation + assert reservation["callback_bound"] is False diff --git a/tests/test_litellm/llms/xai/test_xai_key_fallback.py b/tests/test_litellm/llms/xai/test_xai_key_fallback.py index 092e4951547..cbc507c5ee3 100644 --- a/tests/test_litellm/llms/xai/test_xai_key_fallback.py +++ b/tests/test_litellm/llms/xai/test_xai_key_fallback.py @@ -12,6 +12,9 @@ from litellm.types.router import GenericLiteLLMParams class FakeLogging: + def __init__(self) -> None: + self.litellm_params: dict = {} + def update_from_kwargs(self, **kwargs): pass diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index da36071a5b4..f03abe8f124 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -59,6 +59,7 @@ from litellm.proxy.auth.user_api_key_auth import ( _user_api_key_auth_builder, get_api_key, user_api_key_auth, + user_api_key_auth_websocket_for_model, ) from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata @@ -9043,3 +9044,90 @@ async def test_router_settings_model_group_alias_authorizes_target_for_team(monk await authorize() assert (await request.json())["model"] == target assert get_client_requested_model(request) == "AgentX-LLM" + + +@pytest.mark.asyncio +async def test_reserve_budget_after_common_checks_hands_the_reservation_to_the_request_state(): + from fastapi import Request + + request = Request(scope={"type": "http"}) + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=reservation), + ): + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/batches/batch_123/cancel", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={}, + request=request, + ) + + assert user_api_key_auth_obj.budget_reservation is reservation + assert request.state.budget_reservation is reservation + assert request.scope["state"]["budget_reservation"] is reservation + + +@pytest.mark.asyncio +async def test_reserve_budget_after_common_checks_clears_the_request_state_when_budget_checks_skip(): + from fastapi import Request + + request = Request(scope={"type": "http", "state": {"budget_reservation": {"reserved_cost": 0.5}}}) + + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=UserAPIKeyAuth(token="test_token"), + request_data={"model": "free-model"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=True, + general_settings={}, + request=request, + ) + + assert request.state.budget_reservation is None + + +@pytest.mark.asyncio +async def test_websocket_auth_hands_the_reservation_to_the_socket_state(): + from fastapi import WebSocket + + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + websocket = WebSocket( + scope={ + "type": "websocket", + "path": "/v1/realtime", + "headers": [(b"authorization", b"Bearer sk-1234")], + "query_string": b"model=gpt-realtime", + }, + receive=AsyncMock(), + send=AsyncMock(), + ) + + async def auth_that_reserves(request, api_key): + request.state.budget_reservation = reservation + return UserAPIKeyAuth(token="hashed", budget_reservation=reservation) + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + new=AsyncMock(side_effect=auth_that_reserves), + ): + result = await user_api_key_auth_websocket_for_model(websocket, model="gpt-realtime") + + assert result.budget_reservation == reservation + assert websocket.state.budget_reservation is reservation + assert websocket.scope["state"]["budget_reservation"] is reservation diff --git a/tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py b/tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py new file mode 100644 index 00000000000..f37a20dff8b --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py @@ -0,0 +1,349 @@ +""" +Tests for BudgetReservationReleaseMiddleware. + +Auth reserves budget before the handler runs and hands the reservation to the +request or socket state. A litellm call made through the async client wrapper +claims it for the cost callback that runs after the call; anything still unclaimed +when the response is done or the socket has closed would keep the spend counter +pinned until its TTL, so the middleware releases it. +""" + +import asyncio +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from datetime import datetime +from typing import Final + +import pytest +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse, Response, StreamingResponse +from starlette.routing import Route +from starlette.types import ASGIApp, Message, Receive, Scope, Send +from starlette.websockets import WebSocket + +import litellm +from litellm.caching import DualCache +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.middleware.budget_reservation_release_middleware import ( + BudgetReservationReleaseMiddleware, +) +from litellm.proxy.spend_tracking.budget_reservation import ( + reconcile_budget_reservation, + release_unbound_budget_reservation, + reserve_budget_for_request, +) +from litellm.proxy.utils import ProxyLogging +from litellm.utils import Rules, function_setup + +KEY_TOKEN: Final = "hashed-release-middleware-key" +COUNTER_KEY: Final = f"spend:key:{KEY_TOKEN}" +CHAT_BODY: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + + +@pytest.fixture +def spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> DualCache: + cache: Final = DualCache() + monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", None) + return cache + + +@pytest.fixture +def no_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: + for callback_list_name in ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, callback_list_name, []) + + +async def _reserve() -> dict: + reservation: Final = await reserve_budget_for_request( + request_body=CHAT_BODY, + route="/v1/chat/completions", + llm_router=None, + valid_token=UserAPIKeyAuth(token=KEY_TOKEN, max_budget=1.0, spend=0.0), + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + ) + assert reservation is not None + assert reservation["reserved_cost"] > 0 + return reservation + + +async def _chat(reservation: dict, **kwargs: object) -> object: + return await litellm.acompletion( + **CHAT_BODY, + metadata={"user_api_key_budget_reservation": reservation}, + **kwargs, + ) + + +def _proxy_pre_call_setup(route_type: str, reservation: dict) -> None: + function_setup( + original_function=route_type, + rules_obj=Rules(), + start_time=datetime.now(), + **CHAT_BODY, + litellm_call_id="proxy-pre-call-setup", + metadata={"user_api_key_budget_reservation": reservation}, + ) + + +def _app( + handler: Callable[[Request], Awaitable[Response]], + release: Callable[[Mapping[str, object]], Awaitable[None]] = release_unbound_budget_reservation, +) -> Starlette: + app: Final = Starlette(routes=[Route("/", handler, methods=["POST"])]) + app.add_middleware(BudgetReservationReleaseMiddleware, release=release) + return app + + +async def _post(app: ASGIApp) -> None: + scope: Final = { + "type": "http", + "method": "POST", + "path": "/", + "raw_path": b"/", + "headers": [], + "query_string": b"", + "scheme": "http", + "server": ("testserver", 80), + "client": ("testclient", 1), + } + + body_delivered: Final = asyncio.Event() + client_never_disconnects: Final = asyncio.Event() + + async def receive() -> Message: + if body_delivered.is_set(): + await client_never_disconnects.wait() + body_delivered.set() + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message: Message) -> None: + return None + + await app(scope, receive, send) + + +def _counter(spend_counter_cache: DualCache) -> float | None: + return spend_counter_cache.in_memory_cache.get_cache(key=COUNTER_KEY) + + +@pytest.mark.asyncio +async def test_unbound_reservation_is_released_after_the_response(spend_counter_cache: DualCache): + reservation: Final = await _reserve() + assert _counter(spend_counter_cache) == pytest.approx(reservation["reserved_cost"]) + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + return JSONResponse({"id": "batch_123", "status": "cancelling"}) + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_unbound_reservation_is_released_when_the_handler_raises(spend_counter_cache: DualCache): + reservation: Final = await _reserve() + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + raise RuntimeError("upstream refused the cancel") + + with pytest.raises(RuntimeError): + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reservation_seen_only_by_the_proxy_pre_call_logging_object_is_released( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + + async def cancel_batch_without_a_client_wrapper() -> dict: + return {"id": "batch_123", "status": "cancelling"} + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acancel_batch", reservation) + return JSONResponse(await cancel_batch_without_a_client_wrapper()) + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reservation_of_a_failed_call_is_released_after_the_error_response( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + refused: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o") + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acompletion", reservation) + try: + await _chat(reservation, mock_response=refused) + except litellm.AuthenticationError: + return JSONResponse({"error": {"message": "bad key"}}, status_code=401) + raise AssertionError("the mocked call must fail") + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reservation_claimed_by_a_completed_call_is_left_for_the_callback( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + reserved_cost: Final = reservation["reserved_cost"] + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acompletion", reservation) + response: Final = await _chat(reservation, mock_response="ok") + return JSONResponse(response.model_dump()) + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(reserved_cost) + assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_reservation_claimed_by_a_streaming_call_is_left_for_the_callback_that_finishes_after_the_response( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + reserved_cost: Final = reservation["reserved_cost"] + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acompletion", reservation) + stream: Final = await _chat(reservation, mock_response="ok", stream=True) + + async def sse() -> AsyncIterator[bytes]: + async for chunk in stream: + yield f"data: {chunk.model_dump_json()}\n\n".encode() + yield b"data: [DONE]\n\n" + + return StreamingResponse(sse(), media_type="text/event-stream") + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(reserved_cost) + assert reservation["finalized"] is False + + actual_cost: Final = reserved_cost / 4 + await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=actual_cost) + + assert _counter(spend_counter_cache) == pytest.approx(actual_cost) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_unbound_reservation_of_a_websocket_session_is_released_when_the_socket_closes( + spend_counter_cache: DualCache, +): + reservation: Final = await _reserve() + + async def listen_without_a_provider_key(scope: Scope, receive: Receive, send: Send) -> None: + websocket: Final = WebSocket(scope, receive, send) + websocket.state.budget_reservation = reservation + await websocket.close(code=1011, reason="Required 'DEEPGRAM_API_KEY' in environment") + + async def receive() -> Message: + return {"type": "websocket.connect"} + + async def send(message: Message) -> None: + return None + + middleware: Final = BudgetReservationReleaseMiddleware( + listen_without_a_provider_key, release=release_unbound_budget_reservation + ) + await middleware({"type": "websocket", "path": "/deepgram/v1/listen", "headers": []}, receive, send) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_release_runs_once_per_request_with_the_stamped_reservation(): + released: Final = [] + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + async def release(budget_reservation: Mapping[str, object]) -> None: + released.append(budget_reservation) + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + return JSONResponse({}) + + await _post(_app(handler, release=release)) + + assert released == [reservation] + assert released[0] is reservation + + +@pytest.mark.asyncio +async def test_request_without_a_reservation_releases_nothing(): + released: Final = [] + + async def release(budget_reservation: Mapping[str, object]) -> None: + released.append(budget_reservation) + + async def unauthenticated(request: Request) -> Response: + return JSONResponse({}) + + async def budget_checks_skipped(request: Request) -> Response: + request.state.budget_reservation = None + return JSONResponse({}) + + await _post(_app(unauthenticated, release=release)) + await _post(_app(budget_checks_skipped, release=release)) + + assert released == [] + + +@pytest.mark.asyncio +async def test_lifespan_scopes_pass_through(): + released: Final = [] + seen: Final = [] + + async def release(budget_reservation: Mapping[str, object]) -> None: + released.append(budget_reservation) + + async def inner(scope: Scope, receive: Receive, send: Send) -> None: + seen.append(scope["type"]) + + async def receive() -> Message: + return {"type": "lifespan.startup"} + + async def send(message: Message) -> None: + return None + + middleware: Final = BudgetReservationReleaseMiddleware(inner, release=release) + await middleware({"type": "lifespan", "state": {"budget_reservation": {"reserved_cost": 1.0}}}, receive, send) + + assert seen == ["lifespan"] + assert released == [] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 6cd537489ac..6ad850866b7 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4256,6 +4256,112 @@ async def test_pass_through_request_non_streaming_success_unchanged(): mock_success_handler.assert_called_once() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "upstream_status_code, claimed_by_the_success_handler", + [(200, True), (500, False)], + ids=["success-claims-the-reservation", "upstream-error-leaves-it-for-the-request-end-release"], +) +async def test_pass_through_request_claims_the_budget_reservation_only_when_its_success_handler_runs( + upstream_status_code: int, claimed_by_the_success_handler: bool +): + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + user_api_key_dict: Final = UserAPIKeyAuth(api_key="hashed") + user_api_key_dict.budget_reservation = reservation + upstream_response: Final = httpx.Response( + status_code=upstream_status_code, + headers={"content-type": "application/json"}, + content=b'{"status": "upstream"}', + request=httpx.Request("POST", "http://target-api.com/api/generate"), + ) + + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client") as mock_get_client, + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker, + ): + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/mock-upstream/api/generate" + mock_request.body = AsyncMock(return_value=b'{"prompt": "hi"}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + response = await pass_through_request( + request=mock_request, + target="http://target-api.com/api/generate", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + ) + + assert response.status_code == upstream_status_code + assert reservation["callback_bound"] is claimed_by_the_success_handler + assert mock_worker.ensure_initialized_and_enqueue.call_count == int(claimed_by_the_success_handler) + + +@pytest.mark.asyncio +async def test_pass_through_request_leaves_the_budget_reservation_for_the_request_end_release_when_its_success_handler_cannot_be_enqueued(): + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + user_api_key_dict: Final = UserAPIKeyAuth(api_key="hashed") + user_api_key_dict.budget_reservation = reservation + upstream_response: Final = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"status": "upstream"}', + request=httpx.Request("POST", "http://target-api.com/api/generate"), + ) + + def refuse_to_enqueue(async_coroutine): + async_coroutine.close() + raise RuntimeError("logging worker is shutting down") + + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client") as mock_get_client, + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker, + ): + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=refuse_to_enqueue) + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/mock-upstream/api/generate" + mock_request.body = AsyncMock(return_value=b'{"prompt": "hi"}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + with pytest.raises(ProxyException): + await pass_through_request( + request=mock_request, + target="http://target-api.com/api/generate", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + ) + + assert reservation["callback_bound"] is False + + @pytest.mark.asyncio async def test_pass_through_request_internal_failure_still_raises_proxy_exception(): """ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index ea6adc35b9a..88b82349c83 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -838,3 +838,72 @@ async def test_chunk_processor_bills_partial_google_usage_on_mid_stream_exceptio assert failure_payload["completion_tokens"] == 12 assert failure_payload["response_cost"] > 12 * 3.75e-06 assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deferred_dispatch_armed", + [False, True], + ids=["enqueued-at-end-of-stream", "parked-for-deferred-dispatch"], +) +async def test_chunk_processor_claims_the_budget_reservation_before_handing_it_to_the_cost_callback( + deferred_dispatch_armed: bool, +): + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + response = _make_streaming_response([b"event-1", b"event-2"]) + logging_obj = _unarmed_logging_obj() + logging_obj.litellm_params = {"metadata": {"user_api_key_budget_reservation": reservation}} + if deferred_dispatch_armed: + logging_obj._on_deferred_stream_complete = AsyncMock() + claimed_when_the_callback_ran = [] + + async def cost_callback(**kwargs): + claimed_when_the_callback_ran.append(reservation["callback_bound"]) + + async for _ in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/bedrock/model/claude/invoke-with-response-stream", + route_streaming_logging=cost_callback, + ): + pass + + if deferred_dispatch_armed: + (parked_cost_callback,) = logging_obj._deferred_stream_complete_args + await parked_cost_callback + else: + await GLOBAL_LOGGING_WORKER.flush() + + assert reservation["callback_bound"] is True + assert claimed_when_the_callback_ran == [True] + + +@pytest.mark.asyncio +async def test_chunk_processor_leaves_the_budget_reservation_for_the_request_end_release_when_the_cost_callback_cannot_be_enqueued(): + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + response = _make_streaming_response([b"event-1", b"event-2"]) + logging_obj = _unarmed_logging_obj() + logging_obj.litellm_params = {"metadata": {"user_api_key_budget_reservation": reservation}} + + def refuse_to_enqueue(async_coroutine): + async_coroutine.close() + raise RuntimeError("logging worker is shutting down") + + with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", side_effect=refuse_to_enqueue): + async for _ in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/bedrock/model/claude/invoke-with-response-stream", + route_streaming_logging=AsyncMock(), + ): + pass + + assert reservation["callback_bound"] is False diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 0e0025f5194..214b5cde7da 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -26,6 +26,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( _get_team_member_budget_counter, count_request_input_tokens, estimate_request_max_cost, + release_unbound_budget_reservation, reserve_budget_for_request, ) from litellm.proxy.utils import ProxyLogging @@ -546,3 +547,41 @@ async def test_team_member_reservation_counter_adds_temp_increase_to_live_team_d assert counter is not None assert counter.max_budget == expected_max_budget assert counter.fallback_spend == 0.5 + + +@pytest.mark.asyncio +async def test_reservation_starts_unbound_to_any_callback(): + reservation: Final = await _reserve("/v1/responses") + + assert reservation is not None + assert reservation["callback_bound"] is False + + +@pytest.mark.asyncio +async def test_release_unbound_budget_reservation_frees_the_counter(spend_counter_cache: DualCache): + counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}" + reservation: Final = await _reserve_for_tiny_budget_key( + "/v1/chat/completions", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + ) + assert reservation is not None + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"]) + + await release_unbound_budget_reservation(reservation) + + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_release_unbound_budget_reservation_leaves_a_bound_one_to_its_callback(spend_counter_cache: DualCache): + counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}" + reservation: Final = await _reserve_for_tiny_budget_key( + "/v1/chat/completions", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + ) + assert reservation is not None + reservation["callback_bound"] = True + + await release_unbound_budget_reservation(reservation) + + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"]) + assert reservation["finalized"] is False diff --git a/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py b/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py index 1f6a214398a..05f2d13a079 100644 --- a/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py +++ b/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py @@ -9,9 +9,10 @@ import pytest from pydantic import TypeAdapter import litellm +from litellm._internal_context import is_internal_call from litellm.litellm_core_utils.litellm_logging import Logging from litellm.rust_bridge import callbacks_legacy_python as legacy -from litellm.rust_bridge.callbacks_legacy_python import check_limits, setup +from litellm.rust_bridge.callbacks_legacy_python import check_limits, failure_handler, setup _OCR_KWARGS: Final = MappingProxyType( { @@ -81,6 +82,82 @@ def test_setup_builds_a_logger_when_none_is_supplied(call_type: str, kwargs: Map assert result.logger.litellm_call_id == result.kwargs["litellm_call_id"] +def _budget_reservation() -> dict: + return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + +def _kwargs_with_a_budget_reservation(reservation: dict) -> dict[str, object]: + return {**_OCR_KWARGS, "metadata": {"user_api_key_budget_reservation": reservation}} + + +def test_setup_claims_the_budget_reservation_for_an_async_call() -> None: + reservation: Final = _budget_reservation() + + setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=True) + + assert reservation["callback_bound"] is True + + +def test_setup_claims_the_budget_reservation_a_supplied_logger_already_saw() -> None: + reservation: Final = _budget_reservation() + supplied: Final = _supplied_logger() + supplied.update_environment_variables( + litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}, optional_params={} + ) + assert reservation["callback_bound"] is False + + setup("aocr", (), {**_OCR_KWARGS, "litellm_logging_obj": supplied}, datetime.datetime.now(), asynchronous=True) + + assert reservation["callback_bound"] is True + + +def test_setup_leaves_the_budget_reservation_alone_for_a_sync_call() -> None: + reservation: Final = _budget_reservation() + + setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=False) + + assert reservation["callback_bound"] is False + + +def test_setup_leaves_the_budget_reservation_alone_for_an_internal_call() -> None: + reservation: Final = _budget_reservation() + token: Final = is_internal_call.set(True) + try: + setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=True) + finally: + is_internal_call.reset(token) + + assert reservation["callback_bound"] is False + + +def test_failure_handler_hands_the_budget_reservation_back_for_an_async_call() -> None: + reservation: Final = _budget_reservation() + now: Final = datetime.datetime.now() + result: Final = setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), now, asynchronous=True) + assert reservation["callback_bound"] is True + + pending: Final = failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=True) + + assert reservation["callback_bound"] is False + assert pending is not None + pending.close() + + +def test_failure_handler_of_an_internal_call_leaves_the_outer_budget_reservation_claim_in_place() -> None: + reservation: Final = _budget_reservation() + now: Final = datetime.datetime.now() + result: Final = setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), now, asynchronous=True) + token: Final = is_internal_call.set(True) + try: + pending: Final = failure_handler(result.logger, RuntimeError("inner step failed"), now, now, asynchronous=True) + finally: + is_internal_call.reset(token) + + assert reservation["callback_bound"] is True + assert pending is not None + pending.close() + + CONTRACT_PATH: Final = ( Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy-python/python_contract.json" ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 1e59f4d878e..cf61a6d9f65 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4984,6 +4984,100 @@ async def test_wrapper_async_fires_post_call_failure_deployment_hook_on_internal assert isinstance(recorder.calls[0][1], litellm.AuthenticationError) +def _budget_reservation(callback_bound: bool = False) -> dict: + return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": callback_bound} + + +_BUDGET_RESERVATION_CALL_KWARGS: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]} +_BUDGET_RESERVATION_REFUSAL: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o") + + +@pytest.mark.asyncio +async def test_wrapper_async_claims_the_budget_reservation_for_the_cost_callback() -> None: + reservation = _budget_reservation() + + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response="ok", + metadata={"user_api_key_budget_reservation": reservation}, + ) + + assert reservation["callback_bound"] is True + + +@pytest.mark.asyncio +async def test_wrapper_async_claims_the_budget_reservation_before_the_stream_is_consumed() -> None: + reservation = _budget_reservation() + + stream = await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response="ok", + stream=True, + metadata={"user_api_key_budget_reservation": reservation}, + ) + + assert reservation["callback_bound"] is True + async for _ in stream: + pass + + +@pytest.mark.asyncio +async def test_wrapper_async_claims_the_budget_reservation_a_supplied_logging_object_already_saw() -> None: + reservation = _budget_reservation() + logging_obj, kwargs = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + **_BUDGET_RESERVATION_CALL_KWARGS, + litellm_call_id="proxy-pre-call-setup", + metadata={"user_api_key_budget_reservation": reservation}, + ) + assert reservation["callback_bound"] is False + + await litellm.acompletion(**kwargs, litellm_logging_obj=logging_obj, mock_response="ok") + + assert reservation["callback_bound"] is True + + +@pytest.mark.asyncio +async def test_wrapper_async_hands_the_budget_reservation_back_when_the_call_fails() -> None: + reservation = _budget_reservation() + + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response=_BUDGET_RESERVATION_REFUSAL, + metadata={"user_api_key_budget_reservation": reservation}, + ) + + assert reservation["callback_bound"] is False + + +@pytest.mark.asyncio +async def test_wrapper_async_leaves_the_budget_reservation_alone_on_internal_calls() -> None: + claimed_by_the_outer_call = _budget_reservation(callback_bound=True) + never_claimed = _budget_reservation() + + token = is_internal_call.set(True) + try: + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response="ok", + metadata={"user_api_key_budget_reservation": never_claimed}, + ) + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response=_BUDGET_RESERVATION_REFUSAL, + metadata={"user_api_key_budget_reservation": claimed_by_the_outer_call}, + ) + finally: + is_internal_call.reset(token) + + assert never_claimed["callback_bound"] is False + assert claimed_by_the_outer_call["callback_bound"] is True + + @pytest.mark.asyncio async def test_wrapper_async_does_not_fire_failure_hook_for_pre_call_budget_error( monkeypatch: pytest.MonkeyPatch, From 2f76aa1b1bcd64bd5e3101592fbd9be48f4b011b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 19:51:26 -0700 Subject: [PATCH 3/3] test(integration): move Xiaomi MiMo coverage from live e2e to the providers wire shard (#42395) --- .../coverage_registry/llm_conversational.yaml | 3 - .../llm_translation/test_xiaomi_mimo_e2e.py | 214 --------------- tests/integration/contracts.json | 12 + .../providers/test_xiaomi_mimo_wire.py | 258 ++++++++++++++++++ 4 files changed, 270 insertions(+), 217 deletions(-) delete mode 100644 tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py create mode 100644 tests/integration/providers/test_xiaomi_mimo_wire.py diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 64335aa560c..49d4d92ff0b 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -89,9 +89,6 @@ - {id: llm.chat_completions.together_ai.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool result round trip"} - {id: llm.chat_completions.together_ai.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together cost header and spend row match the registry price"} - {id: llm.chat_completions.together_ai.thinking.nonstream.effort_none_disables, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [effort_none_disables], source: "llm_translation/test_together_ai_e2e.py", rationale: "reasoning_effort=none maps to Together's reasoning disable toggle on hybrid models"} -- {id: llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "Native MiMo v2.6 rows price the cost header and spend row from the cost map"} -- {id: llm.chat_completions.xiaomi_mimo.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo reasoning deltas stream as reasoning_content"} -- {id: llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo tool calls are not dropped"} - {id: llm.chat_completions.together_ai.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: structured_output, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "response_format json_schema reaches Together and constrains the reply"} - {id: llm.chat_completions.together_ai.prompt_cache_5m.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: prompt_cache_5m, streaming: nonstream, assertions: [cache_hit, cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together prefix-cache reads bill at cache_read_input_token_cost, not full input price"} - {id: llm.messages.together_ai.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together over /v1/messages streaming"} diff --git a/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py b/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py deleted file mode 100644 index efca216634b..00000000000 --- a/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py +++ /dev/null @@ -1,214 +0,0 @@ -"""Live e2e: Xiaomi MiMo v2.6 through the gateway on /chat/completions. - -Both native ``xiaomi_mimo/`` v2.6 rows (pro and flash) are registered via -``/model/new`` and driven against Xiaomi's own endpoint. What the gateway owes -us is that the reasoning chain surfaces as ``reasoning_content``, tool calls -survive translation, and the cost header plus spend row follow the proxy's own -cost-map price for the row (read back from ``/model/info``, never pinned here). -Requires XIAOMI_MIMO_API_KEY on the proxy; no skip gate. -""" - -from __future__ import annotations - -from typing import Final - -import pytest -from e2e_config import unique_marker -from e2e_http import StreamingResponse, require_successful_call, unwrap -from lifecycle import ResourceManager -from models import ( - ChatBody, - ChatMessage, - ChatResponse, - ChatTool, - ChatToolFunction, - CostMapEntry, - LiteLLMParamsBody, - OutMessage, - SpendLogRow, -) -from passthrough_client import PassthroughClient -from pydantic import BaseModel - -pytestmark = pytest.mark.e2e - -BACKENDS: Final = ("xiaomi_mimo/mimo-v2.6-pro", "xiaomi_mimo/mimo-v2.6-flash") -ARITHMETIC_PROMPT = "What is 17 + 26? Answer with just the number." -WEATHER_PROMPT = "What is the weather in Paris? Use the tool." -COUNTING_PROMPT = "Count from 1 to 50, one number per line." - -WEATHER_TOOL = ChatTool( - function=ChatToolFunction( - name="get_weather", - description="Get the current weather for a location.", - parameters={ - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"], - }, - ) -) - - -class _WeatherArgs(BaseModel): - location: str - - -class _StreamDelta(BaseModel): - content: str | None = None - reasoning_content: str | None = None - - -class _StreamChoice(BaseModel): - delta: _StreamDelta | None = None - - -class _StreamChunk(BaseModel): - choices: list[_StreamChoice] = [] - - -def _approx_equal(actual: float, expected: float) -> bool: - return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) - - -@pytest.fixture(scope="module") -def registry(client: PassthroughClient) -> dict[str, CostMapEntry]: - return client.proxy.model_cost_map() - - -def _register(client: PassthroughClient, resources: ResourceManager, backend: str) -> tuple[str, str]: - model = f"e2e-xiaomi-{unique_marker()}" - model_id = client.proxy.create_model( - model, LiteLLMParamsBody(model=backend, api_key="os.environ/XIAOMI_MIMO_API_KEY") - ) - resources.defer(lambda: client.proxy.delete_model(model_id)) - return model, resources.key() - - -def _message(response: ChatResponse) -> OutMessage: - assert response.choices, f"Xiaomi returned no choices: {response}" - message = response.choices[0].message - assert message is not None, f"Xiaomi choice has no message: {response}" - return message - - -def _deltas(result: StreamingResponse) -> list[_StreamDelta]: - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_done, f"stream never reached [DONE]: {result.stream_events[-3:]}" - return [ - choice.delta - for event in result.stream_events - for choice in _StreamChunk.model_validate_json(event).choices - if choice.delta is not None - ] - - -@pytest.mark.parametrize("backend", BACKENDS) -class TestXiaomiMimoChatCompletions: - @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged") - def test_cost_header_and_spend_row_match_the_registry_price( - self, - client: PassthroughClient, - resources: ResourceManager, - registry: dict[str, CostMapEntry], - backend: str, - ) -> None: - price = registry.get(backend) - assert price is not None, f"{backend} has no row in the proxy's cost map, so native calls would bill $0" - assert price.litellm_provider == "xiaomi_mimo", f"{backend} is filed under the wrong provider: {price}" - assert price.input_cost_per_token and price.output_cost_per_token, f"{backend} carries no price: {price}" - model, key = _register(client, resources, backend) - - result = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(key), - json=ChatBody( - model=model, - messages=[ChatMessage(role="user", content=f"{ARITHMETIC_PROMPT} {unique_marker()}")], - max_tokens=1024, - ), - ) - require_successful_call(result) - response = ChatResponse.model_validate_json(result.body) - message = _message(response) - assert message.content and "43" in message.content, f"answer lost: {message}" - assert message.reasoning_content, f"{backend} reasons, but no reasoning_content came back: {message}" - - usage = response.usage - assert usage is not None and usage.prompt_tokens and usage.completion_tokens, ( - f"response carries no usage, so the cost cannot be real: {result.body[:300]}" - ) - header_cost = result.response_cost - assert header_cost is not None and header_cost > 0, ( - f"x-litellm-response-cost header missing or non-positive: {result.headers}" - ) - cached = (usage.prompt_tokens_details.cached_tokens or 0) if usage.prompt_tokens_details else 0 - expected = ( - (usage.prompt_tokens - cached) * price.input_cost_per_token - + cached * (price.cache_read_input_token_cost or 0.0) - + usage.completion_tokens * price.output_cost_per_token - ) - assert _approx_equal(header_cost, expected), ( - f"header cost {header_cost} disagrees with the registry price for {backend} at {usage}: expected {expected}" - ) - - def _priced(rows: list[SpendLogRow]) -> bool: - return any(row.spend is not None and row.spend > 0 for row in rows) - - rows = client.proxy.poll_logs_for_key(key, predicate=_priced) - priced = [row for row in rows if row.spend is not None and row.spend > 0] - assert priced, f"no priced spend row landed for key {key}; got {rows}" - row = priced[0] - assert row.custom_llm_provider == "xiaomi_mimo", f"spend row misattributed: {row}" - assert row.spend is not None and _approx_equal(row.spend, header_cost), ( - f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}" - ) - - @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.thinking.stream.works") - def test_reasoning_and_answer_stream_as_deltas( - self, client: PassthroughClient, resources: ResourceManager, backend: str - ) -> None: - model, key = _register(client, resources, backend) - - deltas = _deltas( - client.proxy.chat_stream( - key, - ChatBody( - model=model, - messages=[ChatMessage(role="user", content=COUNTING_PROMPT)], - max_tokens=2048, - stream=True, - ), - ) - ) - reasoning = "".join(delta.reasoning_content or "" for delta in deltas) - content = "".join(delta.content or "" for delta in deltas) - assert reasoning, f"stream carried no reasoning_content deltas: {deltas[:5]}" - assert "50" in content, f"streamed answer lost: {content[:300]!r}" - - @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works") - def test_tool_call_is_returned(self, client: PassthroughClient, resources: ResourceManager, backend: str) -> None: - model, key = _register(client, resources, backend) - - message = _message( - unwrap( - client.proxy.chat( - key, - ChatBody( - model=model, - messages=[ChatMessage(role="user", content=WEATHER_PROMPT)], - tools=[WEATHER_TOOL], - max_tokens=1024, - ), - ) - ) - ) - assert message.tool_calls, f"{backend} dropped the tool call: {message}" - call = message.tool_calls[0] - assert call.id, f"tool call carries no id, so a tool result cannot answer it: {call}" - assert call.function.name == "get_weather", f"wrong tool called: {call}" - assert call.function.arguments, f"tool call carries no arguments: {call}" - args = _WeatherArgs.model_validate_json(call.function.arguments) - assert "paris" in args.location.lower(), f"tool arguments lost the location: {args}" diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 2cdd4c17e39..6629b1fa1f4 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -190,6 +190,18 @@ "tests/integration/providers/test_fal_ai_chat_wire.py::test_fal_moondream3_chat_sends_prompt_image_and_reasoning": [ "other.provider_wire.fal_ai.moondream3_chat_query_wire_and_token_pricing" ], + "tests/integration/providers/test_xiaomi_mimo_wire.py::test_xiaomi_mimo_nonstream_surfaces_reasoning_and_charges_registry_price[mimo-v2.6-pro]": [ + "other.provider_wire.xiaomi_mimo.reasoning_content_and_registry_pricing" + ], + "tests/integration/providers/test_xiaomi_mimo_wire.py::test_xiaomi_mimo_nonstream_surfaces_reasoning_and_charges_registry_price[mimo-v2.6-flash]": [ + "other.provider_wire.xiaomi_mimo.reasoning_content_and_registry_pricing" + ], + "tests/integration/providers/test_xiaomi_mimo_wire.py::test_xiaomi_mimo_stream_delivers_reasoning_then_answer_deltas": [ + "other.provider_wire.xiaomi_mimo.reasoning_and_answer_stream_as_deltas" + ], + "tests/integration/providers/test_xiaomi_mimo_wire.py::test_xiaomi_mimo_tool_call_is_forwarded_and_returned": [ + "other.provider_wire.xiaomi_mimo.tool_call_survives_translation" + ], "tests/integration/providers/test_fal_ai_video_wire.py::test_fal_h3_video_create_uses_canonical_body_and_status_path": [ "other.provider_wire.fal_ai.video_queue_create_status_and_content_download" ], diff --git a/tests/integration/providers/test_xiaomi_mimo_wire.py b/tests/integration/providers/test_xiaomi_mimo_wire.py new file mode 100644 index 00000000000..96b9dc17bdf --- /dev/null +++ b/tests/integration/providers/test_xiaomi_mimo_wire.py @@ -0,0 +1,258 @@ +import json +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +_BACKENDS: Final = ("mimo-v2.6-pro", "mimo-v2.6-flash") +_API_KEY: Final = "synthetic-xiaomi-key" +_ARITHMETIC_PROMPT: Final = "What is 17 + 26? Answer with just the number." +_WEATHER_PROMPT: Final = "What is the weather in Paris? Use the tool." +_COUNTING_PROMPT: Final = "Count from 1 to 5, one number per line." +_WEATHER_TOOL: Final[JsonValue] = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} +_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]]) + + +class _Delta(BaseModel): + model_config = ConfigDict(extra="ignore") + content: str | None = None + reasoning_content: str | None = None + + +class _Choice(BaseModel): + model_config = ConfigDict(extra="ignore") + delta: _Delta + finish_reason: str | None = None + + +class _Chunk(BaseModel): + model_config = ConfigDict(extra="ignore") + id: str + choices: tuple[_Choice, ...] + + +def _catalog_cost(backend: str, field: str) -> float: + cost_map: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes()) + cost_value: Final = cost_map[f"xiaomi_mimo/{backend}"][field] + assert isinstance(cost_value, (int, float)) + return float(cost_value) + + +def _approx(value: float) -> object: + return pytest.approx(value, rel=1e-6) # pyright: ignore[reportUnknownMemberType] # pytest lacks typed approx stubs + + +def _completion(identity: str, backend: str, message: Mapping[str, object], finish: str) -> bytes: + return json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": backend, + "choices": [{"index": 0, "message": message, "finish_reason": finish}], + "usage": {"prompt_tokens": 23, "completion_tokens": 41, "total_tokens": 64}, + } + ).encode() + + +def _frame(identity: str, backend: str, delta: Mapping[str, object], finish: str | None = None) -> bytes: + value: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": backend, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + return b"data: " + json.dumps(value).encode() + b"\n\n" + + +def _assert_provider_request(request: Request, backend: str, prompt: str) -> dict[str, JsonValue]: + assert request.method == "POST" + assert request.target == "/chat/completions" + assert request.headers["authorization"] == f"Bearer {_API_KEY}" + assert request.headers["content-type"] == "application/json" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == backend + assert body["messages"] == [{"role": "user", "content": prompt}] + return body + + +@pytest.mark.covers("other.provider_wire.xiaomi_mimo.reasoning_content_and_registry_pricing") +@pytest.mark.parametrize("backend", _BACKENDS) +def test_xiaomi_mimo_nonstream_surfaces_reasoning_and_charges_registry_price(gateway: Gateway, backend: str) -> None: + identity: Final = f"xiaomi-cost-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _assert_provider_request(request, backend, _ARITHMETIC_PROMPT) + assert body["max_tokens"] == 256 + assert "max_completion_tokens" not in body + return Reply( + body=_completion( + identity, + backend, + {"role": "assistant", "content": "43", "reasoning_content": "17 plus 26 is 43."}, + "stop", + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"xiaomi_mimo/{backend}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _ARITHMETIC_PROMPT}], + "max_completion_tokens": 256, + }, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["id"] == identity + assert payload["choices"] == [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "role": "assistant", + "content": "43", + "reasoning_content": "17 plus 26 is 43.", + "provider_specific_fields": {"refusal": None}, + }, + "provider_specific_fields": {}, + } + ] + assert payload["usage"] == {"prompt_tokens": 23, "completion_tokens": 41, "total_tokens": 64} + expected_cost: Final = 23 * _catalog_cost(backend, "input_cost_per_token") + 41 * _catalog_cost( + backend, "output_cost_per_token" + ) + assert float(response.headers["x-litellm-response-cost"]) == _approx(expected_cost) + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == (23, 41) + spend: Final = rows[0]["spend"] + assert isinstance(spend, (int, float, str)) + assert float(spend) == _approx(expected_cost) + + +@pytest.mark.covers("other.provider_wire.xiaomi_mimo.reasoning_and_answer_stream_as_deltas") +def test_xiaomi_mimo_stream_delivers_reasoning_then_answer_deltas(gateway: Gateway) -> None: + backend: Final = _BACKENDS[0] + identity: Final = f"xiaomi-stream-{uuid.uuid4().hex}" + frames: Final = ( + _frame(identity, backend, {"role": "assistant", "reasoning_content": "Count "}), + _frame(identity, backend, {"reasoning_content": "up by one."}), + _frame(identity, backend, {"content": "1\n2\n"}), + _frame(identity, backend, {"content": "3\n4\n5"}), + _frame(identity, backend, {}, finish="stop"), + b"data: [DONE]\n\n", + ) + + def respond(request: Request) -> Reply: + body: Final = _assert_provider_request(request, backend, _COUNTING_PROMPT) + assert body["stream"] is True + return Reply(content_type="text/event-stream", chunks=frames) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"xiaomi_mimo/{backend}", api_base=wire.url, api_key=_API_KEY) + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json={"model": model, "messages": [{"role": "user", "content": _COUNTING_PROMPT}], "stream": True}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read() + lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: ")) + assert lines[-1] == "data: [DONE]" + chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]) + assert {chunk.id for chunk in chunks} == {identity} + choices: Final = tuple(choice for chunk in chunks for choice in chunk.choices) + assert "".join(choice.delta.reasoning_content or "" for choice in choices) == "Count up by one." + assert "".join(choice.delta.content or "" for choice in choices) == "1\n2\n3\n4\n5" + assert tuple(choice.finish_reason for choice in choices if choice.finish_reason) == ("stop",) + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + +@pytest.mark.covers("other.provider_wire.xiaomi_mimo.tool_call_survives_translation") +def test_xiaomi_mimo_tool_call_is_forwarded_and_returned(gateway: Gateway) -> None: + backend: Final = _BACKENDS[1] + identity: Final = f"xiaomi-tool-{uuid.uuid4().hex}" + tool_call: Final = { + "id": "call_paris", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps({"city": "Paris"})}, + } + + def respond(request: Request) -> Reply: + body: Final = _assert_provider_request(request, backend, _WEATHER_PROMPT) + assert body["tools"] == [_WEATHER_TOOL] + assert body["tool_choice"] == "auto" + return Reply( + body=_completion( + identity, + backend, + { + "role": "assistant", + "content": None, + "reasoning_content": "Need the tool.", + "tool_calls": [tool_call], + }, + "tool_calls", + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"xiaomi_mimo/{backend}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _WEATHER_PROMPT}], + "tools": [_WEATHER_TOOL], + "tool_choice": "auto", + }, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": None, + "reasoning_content": "Need the tool.", + "tool_calls": [tool_call], + "provider_specific_fields": {"refusal": None}, + }, + "provider_specific_fields": {}, + } + ] + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]