From 369f3cb5f0b4284ab6b7d8faba2019137fb145d0 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 20:08:45 -0700 Subject: [PATCH 01/64] feat(cost-map): add baseten joey-1 decision model (#49123) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 9 +++++++++ model_prices_and_context_window.json | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1a1982fddd7..77ac4f49e8b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -81059,5 +81059,14 @@ "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": true + }, + "baseten/decision-models/baseten-joey-2b-latest": { + "input_cost_per_token": 1e-08, + "litellm_provider": "baseten", + "max_input_tokens": 16384, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://inference.baseten.co/v1/models", + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1a1982fddd7..77ac4f49e8b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -81059,5 +81059,14 @@ "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": true + }, + "baseten/decision-models/baseten-joey-2b-latest": { + "input_cost_per_token": 1e-08, + "litellm_provider": "baseten", + "max_input_tokens": 16384, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://inference.baseten.co/v1/models", + "supports_vision": true } } From 733cfeba935284056e77baf2dcc9517fdc96a2ba Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 29 Oct 2026 20:30:42 -0700 Subject: [PATCH 02/64] fix(cost): bill a provider reported usage.cost on every path unless the deployment is priced (#48953) * fix(cost): bill a provider reported usage.cost on every path unless the deployment is priced * test(integration): a reported cost is attributed to output in the gateway parity cases * fix(cost): skip a chat usage.cost that already holds litellm's stamped bill and Responses API objects * fix(cost): record litellm's bill on the Responses terminal event so a later pricing pass does not read it as the upstream's charge * fix(cost): mark usage.cost litellm wrote as its own bill and keep a breakdown for a reported cost * fix(cost): keep the token breakdown when a header reported cost repeats an already priced bill * fix(cost): read the stored breakdown total through the logging object's typed field * refactor(cost): name the providers billed from their own usage fields for what they are * fix(cost): keep the upstream's usage.cost beside litellm's bill and stamp a zero operator price on the stream frame --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 21 +- litellm/litellm_core_utils/core_helpers.py | 37 ++- .../streaming_chunk_builder_utils.py | 23 +- litellm/responses/streaming_iterator.py | 3 +- .../cost_calculation/cases/openai.py | 225 ++++++++++++++++++ .../integration/cost_calculation/cost_map.py | 6 + .../litellm_core_utils/test_core_helpers.py | 78 +++--- .../test_streaming_handler.py | 6 +- .../llms/openai_like/test_kyma_provider.py | 8 +- tests/unit/test_main.py | 40 ++++ 10 files changed, 387 insertions(+), 60 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6677b7327dc..ee102e82ca4 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1544,6 +1544,12 @@ def _response_time_ms_for_cost( return litellm_logging_obj.get_response_ms() +def _breakdown_total(litellm_logging_obj: LitellmLoggingObject | None) -> float | None: + breakdown: Final = None if litellm_logging_obj is None else litellm_logging_obj.cost_breakdown + total: Final[object] = None if breakdown is None else breakdown.get("total_cost") + return float(total) if isinstance(total, (int, float)) and not isinstance(total, bool) else None + + def _adjusted_provider_billed_cost( reported_cost: float, custom_llm_provider: str | None, @@ -1561,9 +1567,9 @@ def _adjusted_provider_billed_cost( ) _store_cost_breakdown_in_logging_obj( litellm_logging_obj=litellm_logging_obj, - prompt_tokens_cost_usd_dollar=None, - completion_tokens_cost_usd_dollar=None, - cost_for_built_in_tools_cost_usd_dollar=None, + prompt_tokens_cost_usd_dollar=0.0, + completion_tokens_cost_usd_dollar=adjusted_cost, + cost_for_built_in_tools_cost_usd_dollar=0.0, total_cost_usd_dollar=adjusted_cost, original_cost=reported_cost, discount_percent=discount_percent, @@ -2376,6 +2382,15 @@ def response_cost_calculator( ) provider_response_cost: Final = get_response_cost_from_hidden_params(hidden_params) if provider_response_cost is not None and not custom_pricing: + if _breakdown_total(litellm_logging_obj) != provider_response_cost: + _store_cost_breakdown_in_logging_obj( + litellm_logging_obj=litellm_logging_obj, + prompt_tokens_cost_usd_dollar=0.0, + completion_tokens_cost_usd_dollar=provider_response_cost, + cost_for_built_in_tools_cost_usd_dollar=0.0, + total_cost_usd_dollar=provider_response_cost, + original_cost=provider_response_cost, + ) return provider_response_cost response_cost = completion_cost( diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index d940baf8b56..28f45d2bb2c 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -1017,13 +1017,44 @@ def chunk_usage(chunk: object) -> object: return getattr(chunk, "usage", None) -USAGE_COST_BILLED_PROVIDERS: Final = frozenset({"kyma"}) +PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS: Final = frozenset({"xai", "perplexity", "opencode_go"}) + + +UPSTREAM_USAGE_COST: Final = "upstream_usage_cost" + + +def _numeric(value: object) -> float | None: + return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None + + +def stamp_litellm_bill_on_usage(response: object, usage: object, bill: float) -> None: + """Writes litellm's bill into ``usage.cost`` and the hidden ``response_cost`` cache, keeping the upstream's own + ``usage.cost`` under ``upstream_usage_cost`` so a later pricing pass still bills the upstream's charge instead of + reading litellm's bill back as if the upstream had reported it.""" + hidden: Final[object] = getattr(response, "_hidden_params", None) + if isinstance(hidden, dict): + hidden_params: Final = cast(dict[str, object], hidden) # cast-ok: hidden params ride on the response as a dict + upstream: Final = ( + hidden_params[UPSTREAM_USAGE_COST] + if UPSTREAM_USAGE_COST in hidden_params + else _numeric(getattr(usage, "cost", None)) + ) + hidden_params.update({"response_cost": bill, UPSTREAM_USAGE_COST: upstream}) + setattr(usage, "cost", bill) def billed_usage_cost(response: object, custom_llm_provider: str | None) -> float | None: - if custom_llm_provider not in USAGE_COST_BILLED_PROVIDERS: + """The charge an OpenAI-compatible upstream reports in ``usage.cost``, skipped for providers billed from their own + usage fields. Once litellm wrote its own bill into ``usage.cost`` (the stream builder and the Responses terminal + event do), the upstream's figure is read from ``upstream_usage_cost`` instead.""" + if custom_llm_provider in PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS: return None - cost: Final = provider_reported_usage_cost(_usage_field(response, "usage"), custom_llm_provider) + hidden: Final = as_str_mapping(getattr(response, "_hidden_params", None)) + cost: Final = ( + _numeric(hidden[UPSTREAM_USAGE_COST]) + if hidden is not None and UPSTREAM_USAGE_COST in hidden + else provider_reported_usage_cost(_usage_field(response, "usage"), custom_llm_provider) + ) return cost if cost is not None and math.isfinite(cost) and cost >= 0 else None diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index c69d9277d65..94351e72ab1 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -13,9 +13,10 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import ( - USAGE_COST_BILLED_PROVIDERS, + PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS, chunk_usage, provider_reported_usage_cost, + stamp_litellm_bill_on_usage, ) from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping @@ -1380,9 +1381,6 @@ def stream_chunk_builder_text_completion( return TextCompletionResponse(**response) -_CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value}) | USAGE_COST_BILLED_PROVIDERS - - def _logged_provider(logging_obj: Optional["Logging"]) -> object: if logging_obj is None: return None @@ -1390,10 +1388,6 @@ def _logged_provider(logging_obj: Optional["Logging"]) -> object: return provider -def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool: - return _logged_provider(logging_obj) in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS - - def _stream_builder_response_cost( response: ModelResponse, logging_obj: Optional["Logging"], chunks: Sequence[object] ) -> float | None: @@ -1404,9 +1398,6 @@ def _stream_builder_response_cost( for chunk in reversed(chunks) ) return next((cost for cost in reported if cost is not None), None) - usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None) - if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj): - return float(usage_cost) if logging_obj is not None: return None provider_hint: Final = response._hidden_params.get( # pyright: ignore[reportPrivateUsage] # no public accessor @@ -1453,13 +1444,15 @@ def _set_stream_builder_response_cost( def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None: if logging_obj is None: return - if isinstance(getattr(usage, "cost", None), (int, float)): + if _logged_provider(logging_obj) in PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS and isinstance( + getattr(usage, "cost", None), (int, float) + ): return computed_cost: Final = logging_obj.response_cost_calculator(result=response) - if not isinstance(computed_cost, (int, float)) or computed_cost <= 0: + if isinstance(computed_cost, bool) or not isinstance(computed_cost, (int, float)) or computed_cost < 0: return - setattr(usage, "cost", computed_cost) - _set_hidden_response_cost(response, float(computed_cost)) + if computed_cost > 0 or isinstance(getattr(usage, "cost", None), (int, float)): + stamp_litellm_bill_on_usage(response, usage, float(computed_cost)) def _field_of(value: object, name: str) -> object: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 6a82304b76a..f0bb2201162 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -43,6 +43,7 @@ from litellm.litellm_core_utils.core_helpers import ( process_response_headers, provider_reported_usage_cost, set_response_cost_in_hidden_params, + stamp_litellm_bill_on_usage, ) from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base from litellm.litellm_core_utils.internal_call_metadata import ( @@ -1869,7 +1870,7 @@ def stamp_responses_usage_cost( if isinstance(bill, bool) or not isinstance(bill, (int, float)): return if bill > 0 or isinstance(getattr(usage_obj, "cost", None), (int, float)): - setattr(usage_obj, "cost", float(bill)) + stamp_litellm_bill_on_usage(response_obj, usage_obj, float(bill)) def opening_item_payload(payload: Mapping[str, object]) -> Mapping[str, object]: diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index 53fd5c1cec2..1443cc5ab59 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -2,6 +2,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" +from collections.abc import Mapping from typing import Final, cast from integration.cost_calculation.cost_tracking_case import ( @@ -11878,7 +11879,215 @@ GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM: Final = GPT_5_6_TERRA_MESSAGES_BRIDGED_JS } ) + +# An OpenAI-compatible gateway that reports what it charged under usage.cost (#47533). The reported total is the bill, +# split across input and output by the ratio of the cost map rates (100 * 1e-6 : 20 * 3e-6), on plain and streamed +# replies alike. A deployment the operator priced keeps the operator's rates on both. +_GATEWAY_REPORTED_USAGE: Final = {"prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120, "cost": 0.037} +_GATEWAY_REPORTED_EXPECTED: Final = ExactExpected( + spend=0.037, input_cost=0.0, output_cost=0.037, prompt_tokens=100, completion_tokens=20 +) +_GATEWAY_OPERATOR_PRICED: Final = Deployment(input_cost_per_token=2e-06, output_cost_per_token=6e-06) +_GATEWAY_OPERATOR_PRICED_EXPECTED: Final = ExactExpected( + spend=0.00032, input_cost=0.0002, output_cost=0.00012, prompt_tokens=100, completion_tokens=20 +) +_GATEWAY_CHUNK: Final = { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1, + "model": "gateway-reported-cost-fixture", +} + +GATEWAY_REPORTED_COST: Final = CostTrackingTestCase( + name="gateway-reported-cost", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gateway-reported-cost-fixture", + endpoint="/v1/chat/completions", + request={"model": "$MODEL", "messages": [{"role": "user", "content": "reported cost"}], "stream": False}, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1, + "model": "gateway-reported-cost-fixture", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": _GATEWAY_REPORTED_USAGE, + }, + ), + expected=_GATEWAY_REPORTED_EXPECTED, +) + +GATEWAY_REPORTED_COST_STREAM: Final = CostTrackingTestCase( + name="gateway-reported-cost-stream", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gateway-reported-cost-fixture", + endpoint="/v1/chat/completions", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "reported cost"}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + {**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}]}, + {**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {**_GATEWAY_CHUNK, "choices": [], "usage": _GATEWAY_REPORTED_USAGE}, + done=True, + ), + ), + expected=_GATEWAY_REPORTED_EXPECTED, +) + +GATEWAY_REPORTED_COST_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST.model_copy( + update={ + "name": "gateway-reported-cost-operator_priced", + "deployment": _GATEWAY_OPERATOR_PRICED, + "expected": _GATEWAY_OPERATOR_PRICED_EXPECTED, + } +) + +GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM: Final = GATEWAY_REPORTED_COST_STREAM.model_copy( + update={ + "name": "gateway-reported-cost-operator_priced-stream", + "deployment": _GATEWAY_OPERATOR_PRICED, + "expected": _GATEWAY_OPERATOR_PRICED_EXPECTED, + } +) + +_GATEWAY_RESPONSES_USAGE: Final[Mapping[str, JsonValue]] = {"input_tokens": 100, "output_tokens": 20, "total_tokens": 120, "cost": 0.037} +_GATEWAY_RESPONSES_ITEM: Final[Mapping[str, JsonValue]] = { + "id": "msg_gateway", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], +} +_GATEWAY_RESPONSES_BODY: Final[Mapping[str, JsonValue]] = { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gateway-reported-cost-fixture", + "output": [_GATEWAY_RESPONSES_ITEM], + "usage": _GATEWAY_RESPONSES_USAGE, +} + +GATEWAY_REPORTED_COST_RESPONSES: Final = CostTrackingTestCase( + name="gateway-reported-cost-responses", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gateway-reported-cost-fixture", + endpoint="/v1/responses", + request={"model": "$MODEL", "input": "reported cost", "cache": {"no-cache": True}}, + response=JsonResponse(content_type="application/json", body=_GATEWAY_RESPONSES_BODY), + expected=_GATEWAY_REPORTED_EXPECTED, +) + +GATEWAY_REPORTED_COST_RESPONSES_STREAM: Final = GATEWAY_REPORTED_COST_RESPONSES.model_copy( + update={ + "name": "gateway-reported-cost-responses-stream", + "request": {**GATEWAY_REPORTED_COST_RESPONSES.request, "stream": True}, + "response": SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "response.created", + "sequence_number": 0, + "response": {**_GATEWAY_RESPONSES_BODY, "status": "in_progress", "output": [], "usage": None}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_gateway", + "output_index": 0, + "content_index": 0, + "delta": "ok", + }, + {"type": "response.completed", "sequence_number": 2, "response": _GATEWAY_RESPONSES_BODY}, + ), + ), + } +) + +GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST_RESPONSES.model_copy( + update={ + "name": "gateway-reported-cost-responses-operator_priced", + "deployment": _GATEWAY_OPERATOR_PRICED, + "expected": _GATEWAY_OPERATOR_PRICED_EXPECTED, + } +) + +GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED_STREAM: Final = GATEWAY_REPORTED_COST_RESPONSES_STREAM.model_copy( + update={ + "name": "gateway-reported-cost-responses-operator_priced-stream", + "deployment": _GATEWAY_OPERATOR_PRICED, + "expected": _GATEWAY_OPERATOR_PRICED_EXPECTED, + } +) + + +_GATEWAY_ZERO_USAGE: Final = {**_GATEWAY_REPORTED_USAGE, "cost": 0} + +GATEWAY_REPORTED_ZERO_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST_OPERATOR_PRICED.model_copy( + update={ + "name": "gateway-reported-zero-operator_priced", + "response": JsonResponse( + content_type="application/json", + body={**GATEWAY_REPORTED_COST.response.body, "usage": _GATEWAY_ZERO_USAGE}, + ), + } +) + +GATEWAY_REPORTED_ZERO_OPERATOR_PRICED_STREAM: Final = GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM.model_copy( + update={ + "name": "gateway-reported-zero-operator_priced-stream", + "response": SseResponse( + content_type="text/event-stream", + frames=sse_frames( + {**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}]}, + {**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {**_GATEWAY_CHUNK, "choices": [], "usage": _GATEWAY_ZERO_USAGE}, + done=True, + ), + ), + } +) + +_GATEWAY_OPENROUTER_OPERATOR_PRICED: Final = _GATEWAY_OPERATOR_PRICED.model_copy( + update={"model": "openrouter/vendor/gateway-reported-cost-fixture"} +) + +GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST_OPERATOR_PRICED.model_copy( + update={ + "name": "gateway-reported-cost-openrouter-operator_priced", + "deployment": _GATEWAY_OPENROUTER_OPERATOR_PRICED, + } +) + +GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED_STREAM: Final = ( + GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM.model_copy( + update={ + "name": "gateway-reported-cost-openrouter-operator_priced-stream", + "deployment": _GATEWAY_OPENROUTER_OPERATOR_PRICED, + } + ) +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( + GATEWAY_REPORTED_COST, + GATEWAY_REPORTED_COST_STREAM, + GATEWAY_REPORTED_COST_OPERATOR_PRICED, + GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM, + GATEWAY_REPORTED_COST_RESPONSES, + GATEWAY_REPORTED_COST_RESPONSES_STREAM, + GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED, + GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED_STREAM, + GATEWAY_REPORTED_ZERO_OPERATOR_PRICED, + GATEWAY_REPORTED_ZERO_OPERATOR_PRICED_STREAM, + GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED, + GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED_STREAM, GPT_5_6_SLASH_ALIAS, GPT_5_6_SLASH_ALIAS_STREAM, GPT_6_SOL_RESPONSES_DEPLOYMENT_JSON, @@ -12120,6 +12329,22 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( + StreamParityTestCase(plain=GATEWAY_REPORTED_COST, streamed=GATEWAY_REPORTED_COST_STREAM), + StreamParityTestCase( + plain=GATEWAY_REPORTED_COST_OPERATOR_PRICED, streamed=GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM + ), + StreamParityTestCase( + plain=GATEWAY_REPORTED_ZERO_OPERATOR_PRICED, streamed=GATEWAY_REPORTED_ZERO_OPERATOR_PRICED_STREAM + ), + StreamParityTestCase( + plain=GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED, + streamed=GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED_STREAM, + ), + StreamParityTestCase(plain=GATEWAY_REPORTED_COST_RESPONSES, streamed=GATEWAY_REPORTED_COST_RESPONSES_STREAM), + StreamParityTestCase( + plain=GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED, + streamed=GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED_STREAM, + ), StreamParityTestCase(plain=GPT_5_3_CODEX_INPUT_TEXT, streamed=GPT_5_3_CODEX_STREAM), StreamParityTestCase(plain=GPT_5_3_CODEX_INPUT_TEXT, streamed=GPT_5_3_CODEX_STREAM_INCOMPLETE), StreamParityTestCase(plain=GPT_5_3_CODEX_INPUT_TEXT, streamed=GPT_5_3_CODEX_STREAM_UNVALIDATED), diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 32722b37c9b..da4c2b6e859 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -26,6 +26,12 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "input_cost_per_token_above_128k_tokens": 2.2e-07, "output_cost_per_token_above_128k_tokens": 8.8e-07, }, + "gateway-reported-cost-fixture": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + }, "openrouter/anthropic/claude-sonnet-5": { "litellm_provider": "openrouter", "mode": "chat", diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index a4a3ae285b1..1c86c4e30d8 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -69,9 +69,7 @@ def test_preserve_upstream_non_openai_attributes_keeps_scalar_and_structured_val scalar_attribute="scalar", structured_attribute={"items": ["value"]}, ) - model_response: Final = ModelResponseStream( - choices=[StreamingChoices(index=0, delta=Delta(content="hello"))] - ) + model_response: Final = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="hello"))]) preserve_upstream_non_openai_attributes(model_response, original_chunk) @@ -262,9 +260,7 @@ class TestMapFinishReasonAnthropic: ("content_filtered", "content_filter"), ], ) - def test_anthropic_finish_reasons( - self, provider_reason: str, expected: str - ) -> None: + def test_anthropic_finish_reasons(self, provider_reason: str, expected: str) -> None: assert map_finish_reason(provider_reason) == expected def test_refusal(self): @@ -336,9 +332,7 @@ class TestMapFinishReasonZhipu: class TestMapFinishReasonOpenAIPassthrough: - @pytest.mark.parametrize( - "reason", ["stop", "length", "tool_calls", "function_call", "content_filter"] - ) + @pytest.mark.parametrize("reason", ["stop", "length", "tool_calls", "function_call", "content_filter"]) def test_openai_values_pass_through(self, reason): assert map_finish_reason(reason) == reason @@ -349,9 +343,7 @@ class TestMapFinishReasonGenericError: assert map_finish_reason("error") == "stop" def test_lowercase_error_does_not_warn(self, mocker): - warn = mocker.patch( - "litellm.litellm_core_utils.core_helpers.verbose_logger.warning" - ) + warn = mocker.patch("litellm.litellm_core_utils.core_helpers.verbose_logger.warning") assert map_finish_reason("error") == "stop" warn.assert_not_called() @@ -369,8 +361,7 @@ class TestFinishReasonMapOutputsAreValid: """Every value in _FINISH_REASON_MAP must be a valid OpenAI finish reason.""" for provider_reason, openai_reason in _FINISH_REASON_MAP.items(): assert openai_reason in VALID_OPENAI_FINISH_REASONS, ( - f"Mapped value '{openai_reason}' (from '{provider_reason}') " - f"is not a valid OpenAI finish reason" + f"Mapped value '{openai_reason}' (from '{provider_reason}') is not a valid OpenAI finish reason" ) @@ -380,28 +371,18 @@ class TestRedactNestedMatchAndRegexKeys: "assessments": [ { "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "secret-name", "action": "BLOCKED"} - ] - }, - "wordPolicy": { - "customWords": [{"match": "badword", "action": "BLOCKED"}] + "piiEntities": [{"type": "NAME", "match": "secret-name", "action": "BLOCKED"}] }, + "wordPolicy": {"customWords": [{"match": "badword", "action": "BLOCKED"}]}, } ], "regex": "should-redact-key-named-regex", } out = redact_nested_match_and_regex_keys(payload) - assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] == "[REDACTED]" - assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == ( - "[REDACTED]" - ) + assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" + assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == ("[REDACTED]") assert out["regex"] == "[REDACTED]" - assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][ - 0 - ]["match"] == "secret-name" + assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "secret-name" def test_passes_through_none_and_str(self): assert redact_nested_match_and_regex_keys(None) is None @@ -564,13 +545,17 @@ class TestIsExpectedClientError: assert is_expected_client_error(over_budget) is True litellm_limit = RateLimitError( - message="key over rpm", llm_provider="anthropic", model="claude-haiku-4-5", + message="key over rpm", + llm_provider="anthropic", + model="claude-haiku-4-5", category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, ) assert is_expected_client_error(litellm_limit) is True vendor_limit = RateLimitError( - message="rate limited upstream", llm_provider="anthropic", model="claude-haiku-4-5", + message="rate limited upstream", + llm_provider="anthropic", + model="claude-haiku-4-5", category=RateLimitErrorCategory.VENDOR_RATE_LIMIT, ) assert is_expected_client_error(vendor_limit) is False @@ -803,3 +788,34 @@ class TestReceiveStampSurvivesASecondMetadataBucket: } assert get_litellm_metadata_from_kwargs(kwargs)["litellm_received_at"] is own + + +def test_billed_usage_cost_bills_any_chat_upstream_without_its_own_reader() -> None: + from types import SimpleNamespace + + from litellm.litellm_core_utils.core_helpers import billed_usage_cost + from litellm.types.utils import ModelResponse, Usage + + usage = Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120, cost=0.037) + chat = ModelResponse(usage=usage) + assert billed_usage_cost(chat, "openai") == 0.037 + assert billed_usage_cost(chat, "openrouter") == 0.037 + assert billed_usage_cost(chat, "xai") is None + assert billed_usage_cost(SimpleNamespace(usage=usage), "openai") == 0.037 + + +def test_billed_usage_cost_keeps_the_upstream_charge_after_litellm_wrote_its_bill_into_usage_cost() -> None: + from litellm.litellm_core_utils.core_helpers import billed_usage_cost, stamp_litellm_bill_on_usage + from litellm.types.utils import ModelResponse, Usage + + reported = ModelResponse(usage=Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120, cost=0.037)) + stamp_litellm_bill_on_usage(reported, reported.usage, 0.0407) + stamp_litellm_bill_on_usage(reported, reported.usage, 0.0407) + assert reported.usage.cost == 0.0407 + assert reported._hidden_params["response_cost"] == 0.0407 + assert billed_usage_cost(reported, "openai") == 0.037 + + unreported = ModelResponse(usage=Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120)) + stamp_litellm_bill_on_usage(unreported, unreported.usage, 0.0015) + assert unreported.usage.cost == 0.0015 + assert billed_usage_cost(unreported, "openai") is None diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index f4502937fce..41e8c50fa85 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -2659,7 +2659,7 @@ def test_perplexity_streaming_dict_cost_bills_through_its_own_calculator(): ) == pytest.approx(0.00503) -def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map(): +def test_openai_compatible_streaming_reported_cost_is_billed_instead_of_the_cost_map(): import litellm from litellm.cost_calculator import ( get_response_cost_from_hidden_params, @@ -2681,7 +2681,7 @@ def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map(): id="chatcmpl-openai-compatible", model=model, choices=[], - usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=3_144_000), + usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.037), ) CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openai") @@ -2693,7 +2693,7 @@ def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map(): custom_llm_provider="openai", call_type="completion", optional_params={}, - ) == pytest.approx(2e-5) + ) == pytest.approx(0.037) def test_xai_streaming_reported_cost_still_takes_the_margin(monkeypatch): diff --git a/tests/unit/llms/openai_like/test_kyma_provider.py b/tests/unit/llms/openai_like/test_kyma_provider.py index 2b3740ca676..e834e0fc0e9 100644 --- a/tests/unit/llms/openai_like/test_kyma_provider.py +++ b/tests/unit/llms/openai_like/test_kyma_provider.py @@ -254,7 +254,7 @@ def test_interaction_on_kyma_is_billed_at_the_cost_kyma_reports(): @pytest.mark.parametrize("model", ("openrouter/openai/gpt-4o-mini", "deepseek/deepseek-chat")) -def test_interaction_on_another_provider_keeps_its_upstream_cost_out_of_the_body(model: str): +def test_interaction_on_another_openai_compatible_provider_is_billed_at_its_reported_cost(model: str): chat_usage: Final = {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7, "cost": 0.00721} responses_body: Final = { "id": "resp_upstream", @@ -278,7 +278,7 @@ def test_interaction_on_another_provider_keeps_its_upstream_cost_out_of_the_body upstream.route(path__regex=r".*/responses$").respond(200, json=responses_body) interaction: Final = litellm.interactions.create(model=model, input="ping", api_key="k") - assert "cost" not in interaction.model_dump(exclude_none=True)["usage"] + assert interaction.model_dump(exclude_none=True)["usage"]["cost"] == pytest.approx(0.00721) def test_sdk_completion_cost_prices_kyma_at_the_cost_kyma_reports(): @@ -410,8 +410,8 @@ def test_cost_headers_show_the_discount_and_margin_taken_from_the_cost_kyma_repo assert float(headers["x-litellm-response-cost-discount-amount"] or "nan") == pytest.approx(0.004) assert float(headers["x-litellm-response-cost-margin-amount"] or "nan") == pytest.approx(0.0004) assert float(headers["x-litellm-response-cost"] or "nan") == pytest.approx(0.0044) - assert headers["x-litellm-response-cost-input"] is None - assert headers["x-litellm-response-cost-output"] is None + assert float(headers["x-litellm-response-cost-input"] or "nan") == 0.0 + assert float(headers["x-litellm-response-cost-output"] or "nan") == pytest.approx(0.0044) @pytest.mark.parametrize("thinking", ({"type": "enabled", "budget_tokens": 1024}, {"type": "adaptive"})) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 571f890a4e6..b2ae6b43d26 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -5443,6 +5443,46 @@ def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): assert response._hidden_params["response_cost"] == pytest.approx(0.5) +def test_completion_cost_on_a_stamped_stream_still_bills_the_reported_charge(): + usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") + usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello ", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() + ) + + assert response is not None + assert litellm.completion_cost(completion_response=response) == pytest.approx(0.5) + + +def test_stream_chunk_builder_stamps_a_zero_operator_price_over_a_reported_cost(): + logging_obj: Final = _stream_builder_logging_obj() + logging_obj.litellm_params = { + **logging_obj.litellm_params, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + logging_obj.custom_cost_per_token = {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0} + usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") + usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello ", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + assert response is not None + assert getattr(response.usage, "cost", None) == 0.0 + assert response._hidden_params["response_cost"] == 0.0 + + def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk(): from openai.types.completion_usage import CompletionUsage From 342ce6ae4a4549cacaf96e47aa80a096ae1035bd Mon Sep 17 00:00:00 2001 From: DewellNaa Date: Thu, 29 Oct 2026 20:30:50 -0700 Subject: [PATCH 03/64] feat(zai): support response_format and max_completion_tokens, add glm-5v-turbo pricing (#48335) * feat(zai): fall back to BIGMODEL/ZHIPUAI api keys and add GLM-5.3 aliases and glm-5v-turbo Resolve the zai api key from ZAI_API_KEY, BIGMODEL_API_KEY or ZHIPUAI_API_KEY so deployments predating the zai rename keep working. Accept response_format and max_completion_tokens as supported openai params. Register zai/GLM-5.3 and zai/GLM-5.3-Flash aliases plus the glm-5v-turbo vision model in both the main and backup cost maps * fix(zai): read only ZAI_API_KEY, dropping the BigModel key fallbacks * fix(zai): map max_completion_tokens to max_tokens, correct glm-5v-turbo pricing, drop redundant aliases * fix(zai): reject json_schema response_format, add glm-5v-turbo cache read price * fix(zai): forward every response_format type so allowed_openai_params keeps working * fix(zai): follow json_schema natively on glm-5.3, prompt other models, drop over-limit max_tokens under drop_params * fix(zai): keep BigModel's json_schema as sent and never drop the caller's max_tokens under drop_params * fix(zai): write the prompted JSON schema as JSON so models answer with valid JSON * test(zai): type the zai provider tests and mark their locals Final --------- Co-authored-by: DewellNaa <210961864+DewellNaa@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../prompt_templates/factory.py | 34 ++- litellm/llms/bigmodel/chat/transformation.py | 4 + litellm/llms/zai/chat/transformation.py | 94 ++++++++- ...odel_prices_and_context_window_backup.json | 17 ++ model_prices_and_context_window.json | 17 ++ .../llms/bigmodel/test_bigmodel_provider.py | 23 ++ tests/unit/llms/zai/test_zai_provider.py | 197 ++++++++++++++++++ 7 files changed, 364 insertions(+), 22 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index fad4caa59a0..834d26b4fe8 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5878,38 +5878,36 @@ def function_call_prompt( ] -def response_schema_prompt(model: str, response_schema: dict, llm_provider: str) -> str: +def response_schema_prompt(model: str, response_schema: Mapping[str, object], llm_provider: str) -> str: """ Decides if a user-defined custom prompt or default needs to be used Returns the prompt str that's passed to the model as a user message """ - custom_prompt_details: dict | None = None - response_schema_as_message: Final = [{"role": "user", "content": f"{response_schema}"}] - if f"{model}/response_schema_prompt" in litellm.custom_prompt_dict: - custom_prompt_details = litellm.custom_prompt_dict[ - f"{model}/response_schema_prompt" - ] # allow user to define custom response schema prompt by model - elif "response_schema_prompt" in litellm.custom_prompt_dict: - custom_prompt_details = litellm.custom_prompt_dict["response_schema_prompt"] - - if custom_prompt_details is not None: - return render_custom_prompt_template( - template=custom_prompt_details, messages=response_schema_as_message, model=model, llm_provider=llm_provider - ) - else: + prompts: Final[Mapping[str, Mapping[str, object]]] = litellm.custom_prompt_dict + model_prompt_key: Final = f"{model}/response_schema_prompt" + custom_prompt_details: Final = prompts.get( + model_prompt_key if model_prompt_key in prompts else "response_schema_prompt" + ) + if custom_prompt_details is None: return default_response_schema_prompt(response_schema=response_schema) + return render_custom_prompt_template( + template=custom_prompt_details, + messages=[{"role": "user", "content": f"{response_schema}"}], + model=model, + llm_provider=llm_provider, + ) -def default_response_schema_prompt(response_schema: dict) -> str: +def default_response_schema_prompt(response_schema: Mapping[str, object]) -> str: """ Used if provider/model doesn't support 'response_schema' param. This is the default prompt. Allow user to override this with a custom_prompt. """ prompt_str: Final = f"""Use this JSON schema: - ```json - {response_schema} + ```json + {json.dumps(response_schema, default=str)} ```""" return prompt_str diff --git a/litellm/llms/bigmodel/chat/transformation.py b/litellm/llms/bigmodel/chat/transformation.py index 12d5e912c15..93e9e24fe7b 100644 --- a/litellm/llms/bigmodel/chat/transformation.py +++ b/litellm/llms/bigmodel/chat/transformation.py @@ -26,6 +26,10 @@ class BigModelChatConfig(ZAIChatConfig): def custom_llm_provider(self) -> str: return "bigmodel" + @property + def prompts_unsupported_json_schema(self) -> bool: + return False + def get_openai_compatible_provider_info(self, api_base: str | None, api_key: str | None) -> tuple[str, str | None]: return get_bigmodel_url(api_base, "chat"), get_bigmodel_api_key(api_key) diff --git a/litellm/llms/zai/chat/transformation.py b/litellm/llms/zai/chat/transformation.py index e0e42d772f8..9e2b97d2bed 100644 --- a/litellm/llms/zai/chat/transformation.py +++ b/litellm/llms/zai/chat/transformation.py @@ -1,11 +1,43 @@ +from collections.abc import Mapping from typing import Final +from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 + +import litellm +from litellm.litellm_core_utils.prompt_templates.factory import response_schema_prompt from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam, ChatCompletionUserMessage +from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration from ...openai.chat.gpt_transformation import OpenAIGPTConfig ZAI_API_BASE: Final = "https://api.z.ai/api/paas/v4" +TOKEN_LIMIT_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens"}) + + +def _is_json_object( + value: object, +) -> TypeIs[Mapping[str, object]]: # guard-ok: cost-map rows and request params are JSON + return isinstance(value, Mapping) + + +def _json_field(value: object, key: str) -> object: + return value.get(key) if _is_json_object(value) else None + + +def _exceeds_output_limit(model: str, max_tokens: object) -> bool: + info: Final = lookup_model_info_for_registration(model=model, custom_llm_provider="zai") + limit: Final = None if isinstance(info, ModelInfoLookupMiss) else info.get("max_output_tokens") + return isinstance(max_tokens, int) and isinstance(limit, int) and max_tokens > limit + + +def _schema_without_native_support(model: str, response_format: object) -> Mapping[str, object] | None: + if _json_field(response_format, "type") != "json_schema": + return None + if litellm.supports_response_schema(model=model, custom_llm_provider="zai"): + return None + schema: Final = _json_field(_json_field(response_format, "json_schema"), "schema") + return schema if _is_json_object(schema) else None class ZAIChatConfig(OpenAIGPTConfig): @@ -13,6 +45,10 @@ class ZAIChatConfig(OpenAIGPTConfig): def custom_llm_provider(self) -> str | None: return "zai" + @property + def prompts_unsupported_json_schema(self) -> bool: + return True + def _get_openai_compatible_provider_info( self, api_base: str | None, api_key: str | None ) -> tuple[str | None, str | None]: @@ -40,9 +76,11 @@ class ZAIChatConfig(OpenAIGPTConfig): # GLM/ZAI supports cache_control, so return messages and tools unchanged return messages, tools - def get_supported_openai_params(self, model: str) -> list: + def get_supported_openai_params(self, model: str) -> list[str]: base_params: Final = [ "max_tokens", + "max_completion_tokens", + "response_format", "stream", "stream_options", "temperature", @@ -52,8 +90,6 @@ class ZAIChatConfig(OpenAIGPTConfig): "tool_choice", ] - import litellm - try: if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): base_params.append("thinking") @@ -61,3 +97,53 @@ class ZAIChatConfig(OpenAIGPTConfig): pass return base_params + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: request pipeline adds parameters to the returned body + supported_params: Final = self.get_supported_openai_params(model) + sent: Final = {param: value for param, value in non_default_params.items() if param in supported_params} + max_completion_tokens: Final = sent.get("max_completion_tokens") + falls_back_to_max_tokens: Final = max_completion_tokens is None or ( + (litellm.drop_params or drop_params) and _exceeds_output_limit(model, max_completion_tokens) + ) + max_tokens: Final = sent.get("max_tokens") if falls_back_to_max_tokens else max_completion_tokens + return { + **optional_params, + **{param: value for param, value in sent.items() if param not in TOKEN_LIMIT_PARAMS}, + **({} if max_tokens is None else {"max_tokens": max_tokens}), + } + + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: OpenAIGPTConfig.transform_request signature + optional_params: dict[str, object], # mutable-ok: OpenAIGPTConfig.transform_request signature + litellm_params: dict[str, object], # mutable-ok: OpenAIGPTConfig.transform_request signature + headers: dict[str, str], # mutable-ok: OpenAIGPTConfig.transform_request signature + ) -> dict[str, object]: # mutable-ok: inherited contract returns the request body + schema: Final = ( + _schema_without_native_support(model, optional_params.get("response_format")) + if self.prompts_unsupported_json_schema + else None + ) + schema_message: Final = ( + () + if schema is None + else ( + ChatCompletionUserMessage( + role="user", content=response_schema_prompt(model=model, response_schema=schema, llm_provider="zai") + ), + ) + ) + return super().transform_request( + model, + [*messages, *schema_message], + optional_params if schema is None else {**optional_params, "response_format": {"type": "json_object"}}, + litellm_params, + headers, + ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 77ac4f49e8b..96a04c96f44 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -54512,6 +54512,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true }, "zai/glm-5.3-flash": { @@ -54526,10 +54527,26 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true, "source": "https://docs.z.ai/guides/overview/pricing", "supports_vision": true }, + "zai/glm-5v-turbo": { + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 1.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://vercel.com/ai-gateway/models/glm-5v-turbo", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "zai/glm-5.1": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 2.6e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 77ac4f49e8b..96a04c96f44 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -54512,6 +54512,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true }, "zai/glm-5.3-flash": { @@ -54526,10 +54527,26 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true, "source": "https://docs.z.ai/guides/overview/pricing", "supports_vision": true }, + "zai/glm-5v-turbo": { + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 1.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://vercel.com/ai-gateway/models/glm-5v-turbo", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "zai/glm-5.1": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 2.6e-07, diff --git a/tests/unit/llms/bigmodel/test_bigmodel_provider.py b/tests/unit/llms/bigmodel/test_bigmodel_provider.py index cdf98c8f9a4..6abfcb060d0 100644 --- a/tests/unit/llms/bigmodel/test_bigmodel_provider.py +++ b/tests/unit/llms/bigmodel/test_bigmodel_provider.py @@ -197,6 +197,29 @@ def test_chat_does_not_send_the_openai_organization_to_bigmodel(monkeypatch: pyt assert "openai-organization" not in route.calls.last.request.headers +@respx.mock +def test_chat_sends_json_schema_response_format_as_given() -> None: + route: Final = respx.post(f"{BASE}/api/paas/v4/chat/completions").respond( + json={ + "id": "chat_test", + "object": "chat.completion", + "created": 1, + "model": "glm-4.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "{}"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ) + messages: Final = [{"role": "user", "content": "Ann likes 7"}] + response_format: Final = { + "type": "json_schema", + "json_schema": {"name": "profile", "schema": {"type": "object", "properties": {"n": {"type": "integer"}}}}, + } + _ = litellm.completion(model="bigmodel/glm-4.6", messages=messages, response_format=response_format) + body: Final = json.loads(route.calls.last.request.content) + assert body["response_format"] == response_format + assert body["messages"] == messages + + @pytest.mark.parametrize("is_async", (False, True)) @pytest.mark.asyncio @respx.mock diff --git a/tests/unit/llms/zai/test_zai_provider.py b/tests/unit/llms/zai/test_zai_provider.py index 32849d5eef1..fb126bd8f41 100644 --- a/tests/unit/llms/zai/test_zai_provider.py +++ b/tests/unit/llms/zai/test_zai_provider.py @@ -2,13 +2,27 @@ Tests for Z.AI (Zhipu AI) provider - GLM models """ +import json import math +from collections.abc import Mapping +from typing import Final import pytest +import respx import litellm from litellm import completion from litellm.cost_calculator import cost_per_token +from litellm.llms.zai.chat.transformation import ZAIChatConfig +from litellm.types.utils import PromptTokensDetailsWrapper, Usage + +JSON_SCHEMA_RESPONSE_FORMAT: Final = { + "type": "json_schema", + "json_schema": { + "name": "profile", + "schema": {"type": "object", "properties": {"favorite_number": {"type": "integer"}}}, + }, +} @pytest.fixture @@ -99,3 +113,186 @@ def test_zai_sync_completion(respx_mock, zai_response, monkeypatch): assert response.choices[0].message.content == "Hello! How can I help you today?" assert response.usage.total_tokens == 25 + + +def test_zai_sends_response_format_and_maps_max_completion_tokens( + respx_mock: respx.MockRouter, zai_response: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("ZAI_API_KEY", "test-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response) + + completion( + model="zai/glm-5.3", + messages=[{"role": "user", "content": "Reply in JSON"}], + response_format={"type": "json_object"}, + max_completion_tokens=64, + ) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body["response_format"] == {"type": "json_object"} + assert body["max_tokens"] == 64 + assert "max_completion_tokens" not in body + + +@pytest.mark.parametrize( + "token_limits", + [{"max_tokens": 300, "max_completion_tokens": 15}, {"max_completion_tokens": 15, "max_tokens": 300}], +) +def test_zai_max_completion_tokens_wins_over_max_tokens(token_limits: dict[str, int]) -> None: + mapped: Final = ZAIChatConfig().map_openai_params(token_limits, {}, "glm-5.3", drop_params=False) + + assert mapped["max_tokens"] == 15 + assert "max_completion_tokens" not in mapped + + +@pytest.mark.parametrize( + ("over_limit_by", "drop_params", "forwarded"), + [(1, True, False), (1, False, True), (0, True, True)], +) +def test_zai_drop_params_drops_max_completion_tokens_above_model_limit( + local_model_cost_map: None, + respx_mock: respx.MockRouter, + zai_response: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, + over_limit_by: int, + drop_params: bool, + forwarded: bool, +) -> None: + monkeypatch.setenv("ZAI_API_KEY", "test-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response) + requested: Final = litellm.model_cost["zai/glm-4.5-flash"]["max_output_tokens"] + over_limit_by + + completion( + model="zai/glm-4.5-flash", + messages=[{"role": "user", "content": "Count"}], + max_completion_tokens=requested, + drop_params=drop_params, + ) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body.get("max_tokens") == (requested if forwarded else None) + + +@pytest.mark.parametrize("with_over_limit_max_completion_tokens", [False, True]) +def test_zai_drop_params_keeps_max_tokens_above_model_limit( + local_model_cost_map: None, + respx_mock: respx.MockRouter, + zai_response: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, + with_over_limit_max_completion_tokens: bool, +) -> None: + monkeypatch.setenv("ZAI_API_KEY", "test-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response) + over_limit: Final = litellm.model_cost["zai/glm-4.5-flash"]["max_output_tokens"] + 1 + over_limit_max_completion_tokens: Final = ( + {"max_completion_tokens": over_limit + 1} if with_over_limit_max_completion_tokens else {} + ) + + completion( + model="zai/glm-4.5-flash", + messages=[{"role": "user", "content": "Count"}], + max_tokens=over_limit, + drop_params=True, + **over_limit_max_completion_tokens, + ) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body["max_tokens"] == over_limit + + +@pytest.mark.parametrize("pass_through", [{}, {"allowed_openai_params": ["response_format"]}]) +def test_zai_forwards_json_schema_to_models_that_enforce_it( + local_model_cost_map: None, + respx_mock: respx.MockRouter, + zai_response: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, + pass_through: dict[str, list[str]], +) -> None: + monkeypatch.setenv("ZAI_API_KEY", "test-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response) + messages: Final = [{"role": "user", "content": "Ann likes 7"}] + + completion(model="zai/glm-5.3", messages=messages, response_format=JSON_SCHEMA_RESPONSE_FORMAT, **pass_through) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body["response_format"] == JSON_SCHEMA_RESPONSE_FORMAT + assert body["messages"] == messages + + +@pytest.mark.parametrize("pass_through", [{}, {"allowed_openai_params": ["response_format"]}]) +def test_zai_puts_json_schema_in_prompt_for_models_that_ignore_it( + local_model_cost_map: None, + respx_mock: respx.MockRouter, + zai_response: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, + pass_through: dict[str, list[str]], +) -> None: + monkeypatch.setenv("ZAI_API_KEY", "test-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response) + + completion( + model="zai/glm-4.6", + messages=[{"role": "user", "content": "Ann likes 7"}], + response_format=JSON_SCHEMA_RESPONSE_FORMAT, + **pass_through, + ) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body["response_format"] == {"type": "json_object"} + assert body["messages"][0] == {"role": "user", "content": "Ann likes 7"} + assert body["messages"][-1]["role"] == "user" + prompt: Final = body["messages"][-1]["content"] + rendered_schema: Final = prompt[prompt.index("{") : prompt.rindex("}") + 1] + assert json.loads(rendered_schema) == JSON_SCHEMA_RESPONSE_FORMAT["json_schema"]["schema"] + + +def test_zai_glm_5v_turbo_bills_cached_input_at_cache_read_price(local_model_cost_map: None) -> None: + usage: Final = Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800), + ) + + prompt_cost, _ = cost_per_token(model="glm-5v-turbo", custom_llm_provider="zai", usage_object=usage) + + entry: Final = litellm.model_cost["zai/glm-5v-turbo"] + assert math.isclose( + prompt_cost, 200 * entry["input_cost_per_token"] + 800 * entry["cache_read_input_token_cost"] + ) + + +@pytest.mark.parametrize("model", ["glm-5v-turbo", "GLM-5V-Turbo", "GLM-5.3", "GLM-5.3-Flash"]) +def test_zai_cost_resolves_case_insensitively(local_model_cost_map: None, model: str) -> None: + prompt_cost, completion_cost = cost_per_token( + model=model, custom_llm_provider="zai", prompt_tokens=1000, completion_tokens=1000 + ) + + entry: Final = litellm.model_cost[f"zai/{model.lower()}"] + assert math.isclose(prompt_cost, 1000 * entry["input_cost_per_token"]) + assert math.isclose(completion_cost, 1000 * entry["output_cost_per_token"]) + + +def test_zai_glm_5v_turbo_accepts_thinking( + local_model_cost_map: None, + respx_mock: respx.MockRouter, + zai_response: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("ZAI_API_KEY", "test-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response) + + completion( + model="zai/glm-5v-turbo", + messages=[{"role": "user", "content": "Hello"}], + thinking={"type": "disabled"}, + ) + + body: Final = json.loads(respx_mock.calls[0].request.content) + assert body["thinking"] == {"type": "disabled"} From b91fc92b6d0815df038e3496b669871fd917c056 Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Thu, 29 Oct 2026 20:32:04 -0700 Subject: [PATCH 04/64] fix(router): cap the default OTPM reservation and reconcile token counters after each call (#49054) * fix(router): cap the default OTPM reservation for requests without max_tokens * fix(router): reconcile ITPM and OTPM reservations after every successful call * fix(router): read the default OTPM reservation from the served model and make its fraction configurable The default reservation now looks up max output from litellm_params.model, as main does, so a pricing base_model never changes it. NO_MAX_TOKENS_TPM_RESERVATION_FRACTION sets the fraction for both the router check and the key and team TPM floor. A max_tokens that is not a number or string counts as no cap instead of raising inside the check. * fix(router): refund a failed attempt's io token reservation in make_call The proxy passes one logging object to every retry, and it runs failure callbacks only for the first attempt, so the last attempt's ITPM and OTPM reservation stayed held until the minute ended. Each attempt that raises now gives its reservation back before the retry or fallback * fix(router): count Anthropic cache writes once in ITPM reconcile Native /v1/messages usage reports input_tokens without cache reads and writes, so the reconcile now maps it through AnthropicConfig.calculate_usage before subtracting cached tokens. * fix(router): use a deployment max_tokens, charge late usage to its finish minute, and read the reservation fraction as a fraction * fix(router): reserve with the deployment cap only through the param the request's API sends * fix(router): never let a deployment's output cap raise the reservation above the default * test(router): move the reservation fraction env var test to the integration suite The unit test started a child interpreter, which the unit guide forbids. The integration test boots a real proxy with the env var set and checks admission against the held reservation. * test(router): retry the reservation fraction scenario when the minute rolls over The OTPM counter is keyed by minute, so a rollover between the held stream and the capped requests gave the capped request a fresh counter. The control request now leaves room for usage the first attempt left in the new minute. * test(router): give the rollover retry its own deployment --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 4 + litellm/litellm_core_utils/env_utils.py | 10 +- .../hooks/parallel_request_limiter_v3.py | 17 +- litellm/router.py | 30 +- .../io_token_rate_limit_check.py | 185 +++--- .../pre_call_checks/model_rate_limit_check.py | 37 +- .../test_otpm_default_reservation_fraction.py | 137 +++++ .../unit/litellm_core_utils/test_env_utils.py | 10 + tests/unit/test_constants.py | 9 +- .../test_router/test_io_token_rate_limits.py | 530 +++++++++++++++++- 10 files changed, 831 insertions(+), 138 deletions(-) create mode 100644 tests/integration/routing/test_otpm_default_reservation_fraction.py diff --git a/litellm/constants.py b/litellm/constants.py index e07d727973d..cfee777f36a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -7,6 +7,7 @@ from typing import Final, Literal from litellm.litellm_core_utils.env_utils import ( get_env_float, + get_env_float_in_range, get_env_float_or_none, get_env_int, get_env_int_in_range, @@ -149,6 +150,9 @@ DEFAULT_FAILURE_THRESHOLD_PERCENT: Final = get_env_float( "DEFAULT_FAILURE_THRESHOLD_PERCENT", 0.5 ) # default cooldown a deployment if 50% of requests fail in a given minute DEFAULT_MAX_TOKENS: Final = get_env_int("DEFAULT_MAX_TOKENS", 4096) +NO_MAX_TOKENS_TPM_RESERVATION_FRACTION: Final = get_env_float_in_range( + "NO_MAX_TOKENS_TPM_RESERVATION_FRACTION", 0.25, 0.01, 1.0 +) DEFAULT_ALLOWED_FAILS: Final = get_env_int("DEFAULT_ALLOWED_FAILS", 3) DEFAULT_REDIS_SYNC_INTERVAL: Final = get_env_int("DEFAULT_REDIS_SYNC_INTERVAL", 1) DEFAULT_COOLDOWN_TIME_SECONDS: Final = get_env_int("DEFAULT_COOLDOWN_TIME_SECONDS", 5) diff --git a/litellm/litellm_core_utils/env_utils.py b/litellm/litellm_core_utils/env_utils.py index 6883af98bbd..462b0b3df47 100644 --- a/litellm/litellm_core_utils/env_utils.py +++ b/litellm/litellm_core_utils/env_utils.py @@ -71,7 +71,15 @@ def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int) Values outside the range fall back to the default and warn, so a misconfigured knob can neither crash the caller nor silently change the meaning of what it computes. """ - value: Final = get_env_int(env_var, default) + return _within_range(env_var, get_env_int(env_var, default), default, minimum, maximum) + + +def get_env_float_in_range(env_var: str, default: float, minimum: float, maximum: float) -> float: + """Float counterpart of `get_env_int_in_range`.""" + return _within_range(env_var, get_env_float(env_var, default), default, minimum, maximum) + + +def _within_range(env_var: str, value: _NumberT, default: _NumberT, minimum: _NumberT, maximum: _NumberT) -> _NumberT: if minimum <= value <= maximum: return value logging.getLogger("LiteLLM").warning( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 404aff84994..14b08a09116 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -42,7 +42,11 @@ from litellm.caching.redis_batch import ( active_request_redis_batch, ) from litellm.caching.redis_cache import log_redis_failure -from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.constants import ( + DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, + INTERNAL_CALL_ORIGIN_METADATA_KEY, + NO_MAX_TOKENS_TPM_RESERVATION_FRACTION, +) from litellm.cost_calculator import get_usage_object from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens @@ -508,11 +512,6 @@ REDIS_NODE_HASHTAG_NAME: Final = "all_keys" # *some* output budget; these define that fallback estimate. DEFAULT_MAX_TOKENS_ESTIMATE: Final = 4096 DEFAULT_CHARS_PER_TOKEN: Final = 4 -# Fraction of the available output budget reserved as the upfront floor when -# the request omits max_tokens. Applied to both DEFAULT_MAX_TOKENS_ESTIMATE -# (baseline floor) and to the smallest configured TPM limit (capped floor for -# small per-tenant TPM caps). -_TPM_FLOOR_FRACTION: Final = 4 # Both embeddings and the Responses API put their prompt in data["input"], # but only embeddings have no output tokens. Every "is this an embedding" # check on data["input"] must exclude these call types, or a Responses call @@ -1003,10 +1002,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): per-tenant cap can't be tripped by the floor alone. Returns the baseline floor when no limit is provided. """ - baseline: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION + baseline: Final = int(DEFAULT_MAX_TOKENS_ESTIMATE * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) if min_configured_tpm_limit is None: return baseline - return min(baseline, max(1, min_configured_tpm_limit // _TPM_FLOOR_FRACTION)) + return min(baseline, max(1, int(min_configured_tpm_limit * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION))) @staticmethod def _is_embedding_request(data: object, call_type: str | None) -> bool: @@ -1138,7 +1137,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if call_type in RESPONSES_API_CALL_TYPES else base_capped_floor ) - baseline_floor: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION + baseline_floor: Final = int(DEFAULT_MAX_TOKENS_ESTIMATE * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) is_embedding: Final = PROXY_MaxParallelRequestsHandler_v3._is_embedding_request(data, call_type) if ( capped_floor >= baseline_floor diff --git a/litellm/router.py b/litellm/router.py index 3a423b79e49..20af2e5649e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -377,6 +377,7 @@ from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check impor EncryptedContentAffinityCheck, ) from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + async_io_token_refund_failure, build_io_token_rate_limit_headers, deployment_has_io_token_limits, refund_stale_reservation_before_retry, @@ -9500,7 +9501,7 @@ class Router(DeprecatedRouterMembers): raise original_exception - async def make_call(self, original_function: Any, *args, **kwargs): + async def make_call(self, original_function: Any, *args: object, **kwargs: object): """ Handler for making a call to the .completion()/.embeddings()/etc. functions. """ @@ -9509,17 +9510,26 @@ class Router(DeprecatedRouterMembers): endpoint_function: Final = ( generic_function if callable(generic_function) else _OPAQUE_VALUE_ADAPTER.validate_python(original_function) ) - with endpoint_selection( - getattr(endpoint_function, "__name__", None), - getattr(endpoint_function, "__module__", None), - model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group), - ): - response = original_function(*args, **kwargs) - if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): - response = await response + try: + with endpoint_selection( + getattr(endpoint_function, "__name__", None), + getattr(endpoint_function, "__module__", None), + model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group), + ): + response = original_function(*args, **kwargs) + if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): + response = await response + except Exception: + with contextlib.suppress(Exception): + await async_io_token_refund_failure( + self.cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs) + ) + raise await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) ## PROCESS RESPONSE HEADERS - response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs) + response = await self.set_response_headers( + response=response, model_group=model_group if isinstance(model_group, str) else None, request_kwargs=kwargs + ) return response diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 0b3d48c9e1d..7580c7dadd3 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -1,7 +1,8 @@ """ Separate ITPM/OTPM (input/output tokens per minute) deployment rate limits. -- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM +- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM; a request + without max_tokens reserves the model's max output, capped at a fraction of the OTPM limit - Post-call: reconcile ITPM to actual input tokens and OTPM to actual output tokens - Cached prompt-read tokens are excluded from ITPM post-call accounting @@ -22,6 +23,7 @@ from litellm import token_counter from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET @@ -127,10 +129,15 @@ def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] | return itpm_key, otpm_key -def _tokenizer_model(deployment: object) -> str: +def _served_model(deployment: object) -> str: params: Final = deployment.get("litellm_params") if is_str_keyed_mapping(deployment) else None served: Final = params.get("model") if is_str_keyed_mapping(params) else None - return deployment_openai_model(deployment) or (served if isinstance(served, str) else "") + return served if isinstance(served, str) else "" + + +def _tokenizer_and_served_models(deployment: object) -> tuple[str, str]: + served: Final = _served_model(deployment) + return deployment_openai_model(deployment) or served, served def _estimate_input_tokens(request_kwargs: dict[str, Any] | None, model: str = "") -> int: @@ -160,24 +167,61 @@ def _model_max_output_tokens(model_name: str) -> int | None: return None -def _resolve_max_tokens(request_kwargs: dict[str, Any] | None, deployment: dict) -> int: - if request_kwargs: - # An explicit max_tokens=0 must be honored, not treated as absent and - # replaced by the model default. - explicit = request_kwargs.get("max_tokens") - if explicit is None: - explicit = request_kwargs.get("max_completion_tokens") - if explicit is None: - explicit = request_kwargs.get("max_output_tokens") - if explicit is not None: - return max(0, int(explicit)) +_CHAT_OUTPUT_CAP_KEYS: Final = ("max_tokens", "max_completion_tokens") +_RESPONSES_OUTPUT_CAP_KEYS: Final = ("max_output_tokens",) +_OUTPUT_CAP_KEYS: Final = _CHAT_OUTPUT_CAP_KEYS + _RESPONSES_OUTPUT_CAP_KEYS - model_name: Final = (deployment.get("litellm_params") or {}).get("model") - if model_name: - model_max: Final = _model_max_output_tokens(model_name) - if model_max is not None: - return model_max - return 4096 + +def _explicit_max_tokens( + request_kwargs: Mapping[str, object] | None, cap_keys: tuple[str, ...] = _OUTPUT_CAP_KEYS +) -> int | None: + if not request_kwargs: + return None + explicit: Final = next( + (request_kwargs[key] for key in cap_keys if request_kwargs.get(key) is not None), + None, + ) + match explicit: + case int() | float() | str(): + return max(0, int(explicit)) + case _: + return None + + +def _sent_deployment_cap_keys(request_kwargs: Mapping[str, object] | None) -> tuple[str, ...]: + is_responses_request: Final = ( + request_kwargs is not None and "input" in request_kwargs and "messages" not in request_kwargs + ) + return _RESPONSES_OUTPUT_CAP_KEYS if is_responses_request else _CHAT_OUTPUT_CAP_KEYS + + +def _deployment_output_cap(request_kwargs: Mapping[str, object] | None, deployment: Mapping[str, object]) -> int | None: + params: Final = deployment.get("litellm_params") + return ( + _explicit_max_tokens(params, _sent_deployment_cap_keys(request_kwargs)) + if is_str_keyed_mapping(params) + else None + ) + + +def _default_otpm_reservation(model_name: str, otpm_limit: int) -> int: + model_max: Final = _model_max_output_tokens(model_name) if model_name else None + share_of_limit: Final = max(1, int(otpm_limit * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)) + return min(4096 if model_max is None else model_max, share_of_limit) + + +def _otpm_reservation( + request_kwargs: Mapping[str, object] | None, + deployment: Mapping[str, object], + model_name: str, + otpm_limit: int, +) -> int: + from_request: Final = _explicit_max_tokens(request_kwargs) + if from_request is not None: + return from_request + default: Final = _default_otpm_reservation(model_name, otpm_limit) + from_deployment: Final = _deployment_output_cap(request_kwargs, deployment) + return default if from_deployment is None else min(from_deployment, default) def _get_usage_tokens(usage: object) -> tuple[int, int, int]: @@ -210,6 +254,12 @@ def _extract_response_usage(response_obj: object) -> object: return getattr(response_obj, "usage", None) +def _openai_shaped_usage(usage: object) -> object: + if isinstance(usage, dict) and litellm.AnthropicConfig.is_anthropic_usage_object(usage): + return litellm.AnthropicConfig().calculate_usage(usage_object=usage, reasoning_content=None) + return usage + + def _usage_is_present(usage: object) -> bool: """ True only if usage carries an actual input/output breakdown. @@ -238,7 +288,7 @@ def _resolve_reconcile_usage_tokens( fields. When usage cannot be resolved, return ``usage_resolved=False`` so callers keep the pre-call reservation instead of refunding it as zero usage. """ - usage: Final = _extract_response_usage(response_obj) + usage: Final = _openai_shaped_usage(_extract_response_usage(response_obj)) if _usage_is_present(usage): prompt_tokens, completion_tokens, cached_tokens = _get_usage_tokens(usage) return max(0, prompt_tokens - cached_tokens), completion_tokens, True @@ -261,6 +311,21 @@ def _resolve_reconcile_usage_tokens( return 0, 0, False +def _key_for_this_minute(reserved_key: str) -> str: + counter, _, _ = reserved_key.rpartition(":") + return f"{counter}:{get_utc_datetime().strftime('%H-%M')}" + + +def _reconcile_increments( + reservation: tuple[int, int, str | None, str | None], billable_input: int, completion_tokens: int +) -> tuple[tuple[str, int], ...]: + itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation + deltas: Final = ((itpm_key, billable_input - itpm_reserved), (otpm_key, completion_tokens - otpm_reserved)) + return tuple( + (key if delta < 0 else _key_for_this_minute(key), delta) for key, delta in deltas if key is not None and delta + ) + + def _stash_reservation_in_metadata( request_kwargs: dict[str, object] | None, *, @@ -311,7 +376,9 @@ def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object, return () top_level: Final = (kwargs.get("metadata"), kwargs.get("litellm_metadata")) litellm_params: Final = _as_mutable_mapping(kwargs.get("litellm_params")) - from_params: Final = () if litellm_params is None else (litellm_params.get("metadata"),) + from_params: Final = ( + () if litellm_params is None else (litellm_params.get("metadata"), litellm_params.get("litellm_metadata")) + ) standard_logging_object: Final = _as_mutable_mapping(kwargs.get("standard_logging_object")) from_logging_object: Final = () if standard_logging_object is None else (standard_logging_object.get("metadata"),) return top_level + from_params + from_logging_object @@ -433,9 +500,11 @@ def io_token_pre_call_check( return deployment request_kwargs: Final = get_io_token_rate_limit_request_kwargs() - _model: Final = _tokenizer_model(deployment) + _model, served_model = _tokenizer_and_served_models(deployment) estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model) - max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) + otpm_to_reserve: Final = ( + _otpm_reservation(request_kwargs, deployment, served_model, otpm_limit) if otpm_limit is not None else 0 + ) dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") @@ -458,7 +527,7 @@ def io_token_pre_call_check( ) if otpm_limit is not None: - otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit) + otpm_reserved = otpm_to_reserve try: _sync_increment_with_rollback( dual_cache, @@ -497,9 +566,11 @@ async def async_io_token_pre_call_check( return deployment request_kwargs: Final = get_io_token_rate_limit_request_kwargs() - _model: Final = _tokenizer_model(deployment) + _model, served_model = _tokenizer_and_served_models(deployment) estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model) - max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) + otpm_to_reserve: Final = ( + _otpm_reservation(request_kwargs, deployment, served_model, otpm_limit) if otpm_limit is not None else 0 + ) dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") @@ -523,7 +594,7 @@ async def async_io_token_pre_call_check( ) if otpm_limit is not None: - otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit) + otpm_reserved = otpm_to_reserve try: await _increment_with_rollback( dual_cache, @@ -564,7 +635,8 @@ def io_token_reconcile_success( ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj - itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs) + reservation: Final = _read_reservation_from_kwargs(request_kwargs) + itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation if itpm_key is None and otpm_key is None: return @@ -572,23 +644,8 @@ def io_token_reconcile_success( try: if usage_resolved: - if itpm_key is not None: - itpm_delta: Final = billable_input - itpm_reserved - if itpm_delta != 0: - dual_cache.increment_cache( - key=itpm_key, - value=itpm_delta, - ttl=RoutingArgsTTL, - ) - - if otpm_key is not None: - otpm_delta: Final = completion_tokens - otpm_reserved - if otpm_delta != 0: - dual_cache.increment_cache( - key=otpm_key, - value=otpm_delta, - ttl=RoutingArgsTTL, - ) + for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens): + dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL) else: verbose_router_logger.debug( "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", @@ -618,39 +675,23 @@ async def async_io_token_reconcile_success( ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj - itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs) + reservation: Final = _read_reservation_from_kwargs(request_kwargs) + itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation if itpm_key is None and otpm_key is None: return billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response) - # Reconcile against the exact key that held the reservation (which encodes - # the reservation's minute), not a key recomputed at response time. This - # tracks actual usage even when the pre-call estimate was 0, and avoids a - # minute-boundary mismatch for calls that span into the next minute. Always - # clear the stash afterwards (even if an increment throws) so a retry or a - # duplicate success event can't re-process it. + # A refund goes back to the minute that held the reservation, while usage above the reservation is charged + # to the minute the response finished in, so a call that runs past its minute still counts against a live + # window. Always clear the stash afterwards (even if an increment throws) so a retry or a duplicate success + # event can't re-process it. try: if usage_resolved: - if itpm_key is not None: - itpm_delta: Final = billable_input - itpm_reserved - if itpm_delta != 0: - await dual_cache.async_increment_cache( - key=itpm_key, - value=itpm_delta, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) - - if otpm_key is not None: - otpm_delta: Final = completion_tokens - otpm_reserved - if otpm_delta != 0: - await dual_cache.async_increment_cache( - key=otpm_key, - value=otpm_delta, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) + for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens): + await dual_cache.async_increment_cache( + key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span + ) else: verbose_router_logger.debug( "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index f4b41211c6e..f41d19cd62c 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -26,7 +26,6 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( - ITPM_RESERVED_KEY, async_io_token_pre_call_check, async_io_token_reconcile_success, async_io_token_refund_failure, @@ -485,24 +484,12 @@ class ModelRateLimitingCheck(CustomLogger): try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") - - # IO token reconciliation works purely from the cache keys stashed in - # kwargs/metadata, so it must run before the model_id guard below - # (which only the TPM-tracking path needs). Otherwise a request whose - # standard_logging_object lacks model_id would never return its - # reservation, leaving the counter elevated until the TTL expires. - slo_metadata: Final = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {} - kwargs_metadata: Final = kwargs.get("metadata") or {} - if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata: - await async_io_token_reconcile_success( - self.dual_cache, - kwargs, - response_obj, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - ) - # Fall through: a deployment can also configure tpm/rpm alongside - # itpm/otpm, and that path's pre-call check reads the tpm_key - # counter tracked below, so it must still be incremented here. + await async_io_token_reconcile_success( + self.dual_cache, + kwargs, + response_obj, + parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), + ) if standard_logging_object is None: return @@ -557,17 +544,7 @@ class ModelRateLimitingCheck(CustomLogger): """ try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") - slo_metadata: Final = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {} - kwargs_metadata: Final = kwargs.get("metadata") or {} - if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata: - io_token_reconcile_success( - self.dual_cache, - kwargs, - response_obj, - ) - # Fall through: a deployment can also configure tpm/rpm alongside - # itpm/otpm, and that path's pre-call check reads the tpm_key - # counter tracked below, so it must still be incremented here. + io_token_reconcile_success(self.dual_cache, kwargs, response_obj) if standard_logging_object is None: return diff --git a/tests/integration/routing/test_otpm_default_reservation_fraction.py b/tests/integration/routing/test_otpm_default_reservation_fraction.py new file mode 100644 index 00000000000..589a2535a03 --- /dev/null +++ b/tests/integration/routing/test_otpm_default_reservation_fraction.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +import dataclasses +import threading +import uuid +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, gateway_from_environment +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import owned_config +from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +GROUPS: Final = ("otpm-fraction-first", "otpm-fraction-retry") +OTPM_LIMIT: Final = 1000 +FRACTION: Final = 0.5 +DEFAULT_RESERVATION: Final = int(OTPM_LIMIT * FRACTION) + + +@dataclass(frozen=True, slots=True) +class Rig: + owned: OwnedProxy + upstream: Wire + gates: SimpleQueue[threading.Event] + + +@dataclass(frozen=True, slots=True) +class Admissions: + held_forwarded_max_tokens: bool + statuses: tuple[int, int, int] + one_minute: bool + + +def _reply_for(gates: SimpleQueue[threading.Event]) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + if is_model_info_probe(request): + return model_list_reply() + if JSON_OBJECT.validate_json(request.body).get("stream") is True: + return dataclasses.replace( + chat_reply("chatcmpl-held", OPENAI_MODEL, "held", stream=True), gate_before_first=gates.get(timeout=60) + ) + return chat_reply("chatcmpl-capped", OPENAI_MODEL, "capped", stream=False) + + return reply + + +def _config(directory: Path, upstream: Wire) -> Path: + return owned_config( + directory, + [ + { + "model_name": group, + "litellm_params": { + "model": OPENAI_MODEL, + "api_base": upstream.url + "/v1", + "api_key": "synthetic-openai-key", + "otpm": OTPM_LIMIT, + }, + "model_info": {"id": group}, + } + for group in GROUPS + ], + litellm_settings={"cache": False}, + router_settings={"optional_pre_call_checks": ["enforce_model_rate_limits"], "num_retries": 0}, + ) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("otpm-fraction") + gates: Final[SimpleQueue[threading.Event]] = SimpleQueue() + with ( + gateway_from_environment() as gateway, + wire_server(_reply_for(gates)) as upstream, + owned_proxy_process( + gateway, + directory, + {"NO_MAX_TOKENS_TPM_RESERVATION_FRACTION": str(FRACTION)}, + config=_config(directory, upstream), + ) as owned, + ): + yield Rig(owned, upstream, gates) + + +def _chat(rig: Rig, group: str, **params: int | bool) -> httpx.Response: + return rig.owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": group, "messages": [{"role": "user", "content": f"fraction {uuid.uuid4().hex}"}], **params}, + ) + + +def _await_forwarded_chat(upstream: Wire) -> Request: + request: Final = upstream.received.get(timeout=60) + return _await_forwarded_chat(upstream) if is_model_info_probe(request) else request + + +def _utc_minute() -> str: + return datetime.now(UTC).strftime("%Y-%m-%d %H:%M") + + +def _admissions_beside_a_held_stream(rig: Rig, group: str) -> Admissions: + release: Final = threading.Event() + rig.gates.put(release) + started: Final = _utc_minute() + with ThreadPoolExecutor(max_workers=1) as pool: + held: Final = pool.submit(_chat, rig, group, stream=True) + try: + forwarded: Final = JSON_OBJECT.validate_json(_await_forwarded_chat(rig.upstream).body) + over: Final = _chat(rig, group, max_tokens=OTPM_LIMIT - DEFAULT_RESERVATION + 1) + fits: Final = _chat(rig, group, max_tokens=OTPM_LIMIT - DEFAULT_RESERVATION) + ended: Final = _utc_minute() + finally: + release.set() + return Admissions( + held_forwarded_max_tokens="max_tokens" in forwarded, + statuses=(over.status_code, fits.status_code, held.result(timeout=60).status_code), + one_minute=started == ended, + ) + + +def test_reservation_fraction_env_var_sets_what_a_request_without_max_tokens_holds(rig: Rig) -> None: + first: Final = _admissions_beside_a_held_stream(rig, GROUPS[0]) + admissions: Final = first if first.one_minute else _admissions_beside_a_held_stream(rig, GROUPS[1]) + + assert admissions == Admissions(held_forwarded_max_tokens=False, statuses=(429, 200, 200), one_minute=True) diff --git a/tests/unit/litellm_core_utils/test_env_utils.py b/tests/unit/litellm_core_utils/test_env_utils.py index 1ddd172e274..89b091e47c2 100644 --- a/tests/unit/litellm_core_utils/test_env_utils.py +++ b/tests/unit/litellm_core_utils/test_env_utils.py @@ -6,6 +6,7 @@ import pytest from litellm._logging import verbose_logger from litellm.litellm_core_utils.env_utils import ( get_env_float, + get_env_float_in_range, get_env_float_or_none, get_env_int, get_env_int_or_none, @@ -94,3 +95,12 @@ def test_unparsable_stdout_log_limit_does_not_break_logging( assert "first line" in messages assert "second line" in messages assert len([message for message in messages if raw in message]) == 1 + + +@pytest.mark.parametrize(("raw", "expected"), [("0.5", 0.5), ("1", 1.0), ("4", 0.25), ("0", 0.25), ("half", 0.25)]) +def test_float_in_range_keeps_values_inside_the_range_and_falls_back_otherwise( + monkeypatch: pytest.MonkeyPatch, raw: str, expected: float +) -> None: + monkeypatch.setenv("ENV_UTILS_TEST_FRACTION", raw) + + assert get_env_float_in_range("ENV_UTILS_TEST_FRACTION", 0.25, 0.01, 1.0) == expected diff --git a/tests/unit/test_constants.py b/tests/unit/test_constants.py index 99b049f1803..d9c34e95364 100644 --- a/tests/unit/test_constants.py +++ b/tests/unit/test_constants.py @@ -19,7 +19,14 @@ from unittest.mock import patch _ENV_HELPERS = frozenset( - {"get_env_int", "get_env_float", "get_env_int_in_range", "get_env_int_or_none", "get_env_float_or_none"} + { + "get_env_int", + "get_env_float", + "get_env_int_in_range", + "get_env_float_in_range", + "get_env_int_or_none", + "get_env_float_or_none", + } ) diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index 05921aa70b3..e5ff4531f84 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -3,6 +3,8 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). """ import asyncio +from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime from typing import Final import pytest @@ -10,13 +12,16 @@ import pytest import litellm from litellm import Router from litellm.caching.dual_cache import DualCache +from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( ITPM_CACHE_KEY, ITPM_RESERVED_KEY, OTPM_CACHE_KEY, OTPM_RESERVED_KEY, + _explicit_max_tokens, _reservation_value, - _resolve_max_tokens, async_io_token_pre_call_check, async_io_token_reconcile_success, build_io_token_rate_limit_headers, @@ -30,7 +35,16 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) +from litellm.types.router import RouterCacheEnum from litellm.types.utils import ModelResponse, Usage +from litellm.utils import get_utc_datetime + +MANTLE_OPUS = "bedrock_mantle/anthropic.claude-opus-4-7" +MANTLE_OPUS_MAX_OUTPUT = litellm.get_model_info(model=MANTLE_OPUS)["max_output_tokens"] +UNMAPPED_MODEL = "openai/litellm-io-unmapped-model" +UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT = 4096 +AZURE_BASE_MODEL = "azure/gpt-5" +AZURE_BASE_MODEL_MAX_OUTPUT = litellm.get_model_info(model=AZURE_BASE_MODEL)["max_output_tokens"] class TestIOTokenRateLimitHelpers: @@ -48,13 +62,23 @@ class TestIOTokenRateLimitHelpers: # A real non-zero estimate is reserved as-is. assert _reservation_value(42, 100) == 42 - def test_resolve_max_tokens_respects_explicit_zero(self): - deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}} - # An explicit max_tokens=0 is honored, not replaced by the model default. - assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0 - # max_completion_tokens is the fallback only when max_tokens is absent. - assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12 - assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9 + @pytest.mark.parametrize( + "request_kwargs, expected", + [ + ({"max_tokens": 0}, 0), + ({"max_completion_tokens": 12}, 12), + ({"max_output_tokens": 9}, 9), + ({"max_tokens": 5, "max_completion_tokens": 12}, 5), + ({"max_tokens": None, "max_output_tokens": 9}, 9), + ({"max_tokens": [100]}, None), + ({"messages": []}, None), + (None, None), + ], + ) + def test_explicit_max_tokens_reads_every_output_cap_key( + self, request_kwargs: dict[str, object] | None, expected: int | None + ): + assert _explicit_max_tokens(request_kwargs) == expected def test_build_io_token_rate_limit_headers(self): headers = build_io_token_rate_limit_headers( @@ -386,6 +410,244 @@ class TestModelRateLimitingCheckIOTokens: normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_check", [False, True]) + @pytest.mark.parametrize( + "model, otpm_limit", + [ + (MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT // 2), + (MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT - 1), + (UNMAPPED_MODEL, 1000), + (UNMAPPED_MODEL, 3), + ], + ) + async def test_request_without_max_tokens_is_admitted_below_model_max_output( + self, model: str, otpm_limit: int, use_sync_check: bool + ): + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment: Final = { + "litellm_params": {"model": model, "otpm": otpm_limit}, + "model_info": {"id": f"io-no-max-tokens-{otpm_limit}"}, + "model_name": "no-max-tokens", + } + request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + + if use_sync_check: + check.pre_call_check(deployment) + else: + await check.async_pre_call_check(deployment) + + reserved: Final = request_kwargs["metadata"][OTPM_RESERVED_KEY] + assert isinstance(reserved, int) + assert 0 < reserved < otpm_limit + assert await dual_cache.async_get_cache(key=request_kwargs["metadata"][OTPM_CACHE_KEY]) == reserved + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "model, otpm_limit", + [(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT), (UNMAPPED_MODEL, UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT)], + ) + async def test_concurrent_requests_without_max_tokens_share_otpm_without_overshoot( + self, model: str, otpm_limit: int + ): + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment: Final = { + "litellm_params": {"model": model, "otpm": otpm_limit}, + "model_info": {"id": "io-no-max-tokens-race-id"}, + "model_name": "no-max-tokens", + } + + async def _attempt() -> dict[str, object] | None: + request_kwargs: Final = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + try: + await check.async_pre_call_check(deployment) + except litellm.RateLimitError: + return None + return request_kwargs["metadata"] + + admitted: Final = [m for m in await asyncio.gather(*[_attempt() for _ in range(8)]) if m is not None] + otpm_key: Final = admitted[0][OTPM_CACHE_KEY] + + assert len(admitted) == round(1 / NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) + assert {m[OTPM_CACHE_KEY] for m in admitted} == {otpm_key} + assert await dual_cache.async_get_cache(key=otpm_key) == sum(m[OTPM_RESERVED_KEY] for m in admitted) + assert await dual_cache.async_get_cache(key=otpm_key) <= otpm_limit + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "model, model_max_output", + [(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT), (UNMAPPED_MODEL, UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT)], + ) + async def test_request_without_max_tokens_reserves_model_max_output_when_otpm_allows( + self, model: str, model_max_output: int + ): + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment: Final = { + "litellm_params": {"model": model, "otpm": model_max_output * 10}, + "model_info": {"id": "io-no-max-tokens-roomy-id"}, + "model_name": "roomy", + } + request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + + await check.async_pre_call_check(deployment) + + assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == model_max_output + + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_check", [False, True]) + async def test_base_model_does_not_change_the_default_otpm_reservation(self, use_sync_check: bool): + check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + otpm_limit: Final = AZURE_BASE_MODEL_MAX_OUTPUT * 10 + + async def _reserved(model_info: dict[str, str]) -> int: + request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + deployment: Final = { + "litellm_params": {"model": "azure/litellm-io-custom-deployment", "otpm": otpm_limit}, + "model_info": model_info, + "model_name": "custom-name", + } + if use_sync_check: + check.pre_call_check(deployment) + else: + await check.async_pre_call_check(deployment) + return request_kwargs["metadata"][OTPM_RESERVED_KEY] + + with_base_model: Final = await _reserved({"id": "io-with-base-model", "base_model": AZURE_BASE_MODEL}) + without_base_model: Final = await _reserved({"id": "io-without-base-model"}) + + assert with_base_model == without_base_model + + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_check", [False, True]) + async def test_unparseable_max_tokens_leaves_itpm_uncharged(self, use_sync_check: bool): + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment: Final = { + "litellm_params": {"model": MANTLE_OPUS, "itpm": 2000, "otpm": 2000}, + "model_info": {"id": "io-bad-max-tokens-id"}, + "model_name": "opus", + } + + async def _pre_call(request_kwargs: dict[str, object]) -> None: + set_io_token_rate_limit_request_kwargs(request_kwargs) + if use_sync_check: + check.pre_call_check(deployment) + else: + await check.async_pre_call_check(deployment) + + await _pre_call({"messages": [{"role": "user", "content": "hello there"}], "max_tokens": "abc", "metadata": {}}) + valid_kwargs: Final = { + "messages": [{"role": "user", "content": "hello there"}], + "max_tokens": 1, + "metadata": {}, + } + await _pre_call(valid_kwargs) + + itpm_counter: Final = await dual_cache.async_get_cache(key=valid_kwargs["metadata"][ITPM_CACHE_KEY]) + assert itpm_counter == valid_kwargs["metadata"][ITPM_RESERVED_KEY] + + @pytest.mark.asyncio + async def test_explicit_max_tokens_above_default_share_is_reserved_in_full(self): + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment: Final = { + "litellm_params": {"model": MANTLE_OPUS, "otpm": 100}, + "model_info": {"id": "io-explicit-max-tokens-id"}, + "model_name": "opus", + } + request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 100, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + + await check.async_pre_call_check(deployment) + + assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == 100 + set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": "hi"}], "metadata": {}}) + with pytest.raises(litellm.RateLimitError): + await check.async_pre_call_check(deployment) + + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_check", [False, True]) + async def test_deployment_max_tokens_sets_the_reservation_unless_the_request_sets_one(self, use_sync_check: bool): + check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + deployment: Final = { + "litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, "max_tokens": 100}, + "model_info": {"id": "io-deployment-max-tokens-id"}, + "model_name": "opus", + } + without_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + with_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 300, "metadata": {}} + + for request_kwargs in (without_cap, with_cap): + set_io_token_rate_limit_request_kwargs(request_kwargs) + if use_sync_check: + check.pre_call_check(deployment) + else: + await check.async_pre_call_check(deployment) + + assert without_cap["metadata"][OTPM_RESERVED_KEY] == 100 + assert with_cap["metadata"][OTPM_RESERVED_KEY] == 300 + + @pytest.mark.asyncio + async def test_deployment_cap_sets_the_reservation_only_through_the_param_its_api_sends(self): + check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + chat: Final = {"messages": [{"role": "user", "content": "hello"}]} + responses: Final = {"input": "hello"} + + async def reserved(deployment_cap: Mapping[str, int], request_shape: Mapping[str, object]) -> int: + request_kwargs: Final = {**request_shape, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check( + { + "litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, **deployment_cap}, + "model_info": {"id": f"io-sent-cap-{sorted(deployment_cap)}-{sorted(request_shape)}"}, + "model_name": "opus", + } + ) + return request_kwargs["metadata"][OTPM_RESERVED_KEY] + + chat_default: Final = await reserved({}, chat) + responses_default: Final = await reserved({}, responses) + + assert chat_default == responses_default > 100 + assert await reserved({"max_tokens": 100}, chat) == 100 + assert await reserved({"max_completion_tokens": 100}, chat) == 100 + assert await reserved({"max_output_tokens": 100}, chat) == chat_default + assert await reserved({"max_output_tokens": 100}, responses) == 100 + assert await reserved({"max_tokens": 100}, responses) == responses_default + assert await reserved({"max_completion_tokens": 100}, responses) == responses_default + + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_check", [False, True]) + async def test_deployment_cap_above_the_default_reserves_the_default(self, use_sync_check: bool): + check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + otpm_limit: Final = UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT * 5 + + async def reserved(deployment_cap: Mapping[str, int]) -> int: + request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + deployment: Final = { + "litellm_params": {"model": UNMAPPED_MODEL, "otpm": otpm_limit, **deployment_cap}, + "model_info": {"id": f"io-deployment-cap-above-default-{deployment_cap}"}, + "model_name": "unmapped", + } + if use_sync_check: + check.pre_call_check(deployment) + else: + await check.async_pre_call_check(deployment) + return request_kwargs["metadata"][OTPM_RESERVED_KEY] + + default: Final = await reserved({}) + + assert await reserved({"max_tokens": default + 1}) == default + assert await reserved({"max_tokens": otpm_limit * 2}) == default + def test_sync_io_pre_call_reserves_and_reconciles(self): from litellm.utils import get_utc_datetime @@ -415,9 +677,9 @@ class TestModelRateLimitingCheckIOTokens: "standard_logging_object": { "model_id": "io-sync-id", "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), + "metadata": {}, }, - "metadata": request_kwargs["metadata"], + "litellm_params": {"metadata": request_kwargs["metadata"]}, } response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], @@ -709,6 +971,32 @@ class TestModelRateLimitingCheckIOTokens: assert await dual_cache.async_get_cache(key=itpm_key) == 4 + @pytest.mark.asyncio + @pytest.mark.parametrize("use_sync_reconcile", [False, True]) + async def test_reconcile_charges_usage_above_the_reservation_to_the_minute_it_finished_in( + self, use_sync_reconcile: bool + ): + dual_cache: Final = DualCache() + counter: Final = "global_router:io-late-id:bedrock_mantle/test:otpm" + reserved_minute_key: Final = f"{counter}:99-99" + await dual_cache.async_increment_cache(key=reserved_minute_key, value=10, ttl=60) + kwargs: Final = {"metadata": {ITPM_RESERVED_KEY: 0, OTPM_RESERVED_KEY: 10, OTPM_CACHE_KEY: reserved_minute_key}} + response: Final = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=4, completion_tokens=25, total_tokens=29), + ) + minute_before: Final = get_utc_datetime().strftime("%H-%M") + + if use_sync_reconcile: + io_token_reconcile_success(dual_cache, kwargs, response) + else: + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + finished_minutes: Final = {minute_before, get_utc_datetime().strftime("%H-%M")} + charged: Final = [await dual_cache.async_get_cache(key=f"{counter}:{minute}") for minute in finished_minutes] + assert await dual_cache.async_get_cache(key=reserved_minute_key) == 10 + assert sum(value or 0 for value in charged) == 25 - 10 + @pytest.mark.asyncio async def test_reconcile_missing_usage_keeps_reservation(self): dual_cache = DualCache() @@ -799,16 +1087,16 @@ class TestModelRateLimitingCheckIOTokens: assert await dual_cache.async_get_cache(key=itpm_key) == 4 - def test_sync_reconcile_anthropic_dict_usage(self): + def test_sync_reconcile_anthropic_dict_usage_counts_uncached_input_once(self): dual_cache = DualCache() itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" - dual_cache.set_cache(key=itpm_key, value=6, ttl=60) + dual_cache.set_cache(key=itpm_key, value=10, ttl=60) dual_cache.set_cache(key=otpm_key, value=4, ttl=60) kwargs = { "metadata": { - ITPM_RESERVED_KEY: 6, + ITPM_RESERVED_KEY: 10, OTPM_RESERVED_KEY: 4, ITPM_CACHE_KEY: itpm_key, OTPM_CACHE_KEY: otpm_key, @@ -816,12 +1104,17 @@ class TestModelRateLimitingCheckIOTokens: } response = { "type": "message", - "usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1}, + "usage": { + "input_tokens": 3, + "output_tokens": 2, + "cache_read_input_tokens": 1, + "cache_creation_input_tokens": 4, + }, } io_token_reconcile_success(dual_cache, kwargs, response) - assert dual_cache.get_cache(key=itpm_key) == 2 + assert dual_cache.get_cache(key=itpm_key) == 3 + 4 assert dual_cache.get_cache(key=otpm_key) == 2 @pytest.mark.asyncio @@ -958,7 +1251,214 @@ class TestModelRateLimitingCheckIOTokens: assert current == 0 +ROUTER_IO_MODEL: Final = "anthropic/claude-haiku-5-5" +RouterCall = Callable[[Router, str | Exception], Awaitable[tuple[int, int]]] + + +def _router_with_io_limits() -> Router: + return Router( + model_list=[ + { + "model_name": "limited", + "litellm_params": {"model": ROUTER_IO_MODEL, "api_key": "sk-test", "itpm": 50000, "otpm": 50000}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + num_retries=1, + ) + + +def _current_minute() -> str: + return get_utc_datetime().strftime("%H-%M") + + +def _counter(router: Router, template: RouterCacheEnum, minutes: frozenset[str]) -> int: + deployment: Final = router.model_list[0] + return sum( + router.cache.get_cache( + key=template.value.format( + id=deployment["model_info"]["id"], + model=deployment["litellm_params"]["model"], + current_minute=minute, + ) + ) + or 0 + for minute in minutes + ) + + +async def _io_counters_since(router: Router, first_minute: str) -> tuple[int, int]: + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + minutes: Final = frozenset({first_minute, _current_minute()}) + return _counter(router, RouterCacheEnum.ITPM, minutes), _counter(router, RouterCacheEnum.OTPM, minutes) + + +async def _sdk_chat(router: Router, mock_response: str | Exception) -> tuple[int, int]: + response: Final = await router.acompletion( + model="limited", + messages=[{"role": "user", "content": "hello there"}], + max_tokens=100, + mock_response=mock_response, + ) + return response.usage.prompt_tokens, response.usage.completion_tokens + + +async def _proxy_chat(router: Router, mock_response: str | Exception) -> tuple[int, int]: + response: Final = await router.acompletion( + model="limited", + messages=[{"role": "user", "content": "hello there"}], + max_tokens=100, + metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_response, + ) + return response.usage.prompt_tokens, response.usage.completion_tokens + + +async def _proxy_responses(router: Router, mock_response: str | Exception) -> tuple[int, int]: + response: Final = await router.aresponses( + model="limited", + input="hello there", + max_output_tokens=100, + litellm_metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_response, + ) + return response.usage.input_tokens, response.usage.output_tokens + + +async def _proxy_messages(router: Router, mock_response: str | Exception) -> tuple[int, int]: + response: Final = await router.aanthropic_messages( + model="limited", + messages=[{"role": "user", "content": "hello there"}], + max_tokens=100, + litellm_metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_response, + ) + return response["usage"]["input_tokens"], response["usage"]["output_tokens"] + + +ProxyCallSharingLogging = Callable[[Router, str | Exception, LiteLLMLoggingObject], Awaitable[object]] +PROXY_MESSAGES: Final = [{"role": "user", "content": "hello there"}] + + +def _logging_obj_shared_across_attempts(original_function: str) -> LiteLLMLoggingObject: + logging_obj, _ = litellm.utils.function_setup( + original_function, + litellm.utils.Rules(), + datetime.now(), + model="limited", + messages=PROXY_MESSAGES, + litellm_call_id="proxy-call", + ) + return logging_obj + + +async def _proxy_chat_sharing_logging( + router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject +) -> object: + return await router.acompletion( + model="limited", + messages=PROXY_MESSAGES, + max_tokens=100, + metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_error, + litellm_call_id="proxy-call", + litellm_logging_obj=logging_obj, + ) + + +async def _proxy_responses_sharing_logging( + router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject +) -> object: + return await router.aresponses( + model="limited", + input="hello there", + max_output_tokens=100, + litellm_metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_error, + litellm_call_id="proxy-call", + litellm_logging_obj=logging_obj, + ) + + +async def _proxy_messages_sharing_logging( + router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject +) -> object: + return await router.aanthropic_messages( + model="limited", + messages=PROXY_MESSAGES, + max_tokens=100, + litellm_metadata={"user_api_key_hash": "hashed-key"}, + mock_response=mock_error, + litellm_call_id="proxy-call", + litellm_logging_obj=logging_obj, + ) + + class TestRouterIOTokenIntegration: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "call", + [_sdk_chat, _proxy_chat, _proxy_responses, _proxy_messages], + ids=["sdk_chat", "proxy_chat_metadata", "proxy_responses_litellm_metadata", "proxy_messages_litellm_metadata"], + ) + async def test_success_reconciles_io_counters_to_real_usage(self, call: RouterCall): + router: Final = _router_with_io_limits() + first_minute: Final = _current_minute() + + usage: Final = await call(router, "hi") + + assert await _io_counters_since(router, first_minute) == usage + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call", "mock_error"), + [ + ( + _proxy_responses, + litellm.InternalServerError(message="mock 500", llm_provider="anthropic", model=ROUTER_IO_MODEL), + ), + (_proxy_messages, "litellm.InternalServerError"), + ], + ids=["proxy_responses_litellm_metadata", "proxy_messages_litellm_metadata"], + ) + async def test_retried_failure_refunds_each_io_reservation_once( + self, call: RouterCall, mock_error: str | Exception + ): + router: Final = _router_with_io_limits() + first_minute: Final = _current_minute() + + with pytest.raises(litellm.InternalServerError): + await call(router, mock_error) + + assert await _io_counters_since(router, first_minute) == (0, 0) + + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call", "original_function", "mock_error"), + [ + (_proxy_chat_sharing_logging, "acompletion", "litellm.InternalServerError"), + ( + _proxy_responses_sharing_logging, + "aresponses", + litellm.InternalServerError(message="mock 500", llm_provider="anthropic", model=ROUTER_IO_MODEL), + ), + (_proxy_messages_sharing_logging, "anthropic_messages", "litellm.InternalServerError"), + ], + ids=["proxy_chat", "proxy_responses", "proxy_messages"], + ) + async def test_retried_failure_sharing_one_logging_object_refunds_every_attempt( + self, call: ProxyCallSharingLogging, original_function: str, mock_error: str | Exception + ): + router: Final = _router_with_io_limits() + first_minute: Final = _current_minute() + + with pytest.raises(litellm.InternalServerError): + await call(router, mock_error, _logging_obj_shared_across_attempts(original_function)) + + assert await _io_counters_since(router, first_minute) == (0, 0) + @pytest.mark.asyncio async def test_model_group_info_aggregates_io_limits(self): router = Router( From ccc56d27adeeece2a1069bde587c76486ba9ab39 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:04:22 -0700 Subject: [PATCH 05/64] fix(azure): add batch prices to the unversioned o3, o4-mini, o3-mini, gpt-4o and gpt-4o-mini rows (#49182) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 15 +++++++++++++++ model_prices_and_context_window.json | 15 +++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 96a04c96f44..0799a298942 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6014,11 +6014,14 @@ "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -6313,16 +6316,19 @@ "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "litellm_provider": "azure", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { "search_context_size_high": 0.014, "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -11303,16 +11309,19 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_batches": 4e-06, "search_context_cost_per_query": { "search_context_size_high": 0.014, "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -11411,11 +11420,14 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -11513,16 +11525,19 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, "search_context_cost_per_query": { "search_context_size_high": 0.014, "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 96a04c96f44..0799a298942 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6014,11 +6014,14 @@ "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -6313,16 +6316,19 @@ "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "litellm_provider": "azure", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { "search_context_size_high": 0.014, "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -11303,16 +11309,19 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_batches": 4e-06, "search_context_cost_per_query": { "search_context_size_high": 0.014, "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -11411,11 +11420,14 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -11513,16 +11525,19 @@ "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, + "input_cost_per_token_batches": 5.5e-07, "litellm_provider": "azure", "max_input_tokens": 200000, "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_batches": 2.2e-06, "search_context_cost_per_query": { "search_context_size_high": 0.014, "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", From f32ac94ad1fc3bce79da08752c7bb6a17ebdbfa1 Mon Sep 17 00:00:00 2001 From: mika-lt <349045450+mika-lt@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:30:02 -0700 Subject: [PATCH 06/64] fix(openai): keep the upstream error body on e.response for OpenAI and Azure errors (#49079) * fix(openai): keep the upstream error response on mapped exceptions so e.response.text has the wire body * fix: keep the upstream error body on 429 and 5xx errors RateLimitError, InternalServerError, BadGatewayError, and ServiceUnavailableError rebuilt their response and dropped the upstream body. They now keep the body and exactly the upstream headers, and an unread streaming response no longer raises. * fix: keep a gzip-encoded 429 or 5xx error readable The kept body is already decoded, so the rebuilt response drops content-encoding instead of decoding the body a second time and raising DecodingError. * fix: keep the upstream headers and drop credentials on every mapped provider error The copied error response now keeps every upstream header and carries a request with only the method and the URL, so no provider key, query secret, or prompt reaches e.request. Timeout and APIError get the response too, and azure/ chat, streaming, embedding, and image errors attach it * fix: keep the upstream error body on gateway timeouts and re-raised gateway errors * test(azure): cover the 404, 500, 502, and 504 error bodies and the missing-key placeholder * fix: keep no e.response on timeouts and errors litellm raises itself * fix: keep the upstream elapsed time on error responses and the upstream body on azure_text errors * fix: keep the upstream body without credentials on Azure content filter errors and message-matched OpenAI errors --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/exceptions.py | 93 +++--- .../exception_mapping_utils.py | 57 +++- litellm/llms/azure/azure.py | 12 +- litellm/llms/azure/completion/handler.py | 22 +- litellm/llms/azure/exception_mapping.py | 5 +- litellm/llms/openai/common_utils.py | 9 + litellm/llms/openai/completion/handler.py | 23 +- .../llms/openai/image_variations/handler.py | 16 +- litellm/llms/openai/openai.py | 24 +- .../test_exception_mapping_utils.py | 82 +++++- .../azure/test_azure_exception_mapping.py | 201 +++++++++++++ .../completion/test_completion_handler.py | 24 +- .../llms/openai/image_variations/__init__.py | 0 .../test_openai_image_variations_handler.py | 51 ++++ tests/unit/llms/openai/test_openai.py | 274 +++++++++++++++++- .../proxy/test_common_request_processing.py | 63 ++++ .../test_exception_header_preservation.py | 19 ++ 17 files changed, 895 insertions(+), 80 deletions(-) create mode 100644 tests/unit/llms/openai/image_variations/__init__.py create mode 100644 tests/unit/llms/openai/image_variations/test_openai_image_variations_handler.py diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 28211b3dab2..12425c45b01 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -9,6 +9,7 @@ ## LiteLLM versions of the OpenAI Exception Types +import datetime import enum from collections.abc import Mapping, Sequence from typing import Any, Final @@ -115,6 +116,54 @@ def validate_rate_limit_type(value: object) -> str | None: return None +def _upstream_body(response: object) -> bytes: + if not isinstance(response, httpx.Response): + return b"" + try: + return response.content + except httpx.ResponseNotRead: + return b"" + + +def _request_without_credentials(upstream: object, fallback_url: str) -> httpx.Request: + if not isinstance(upstream, httpx.Response): + return httpx.Request(method="POST", url=fallback_url) + try: + upstream_request: Final = upstream.request + except RuntimeError: + return httpx.Request(method="POST", url=fallback_url) + return httpx.Request( + method=upstream_request.method, + url=upstream_request.url.copy_with(userinfo=b"", query=None, fragment=None), + ) + + +def _upstream_elapsed(upstream: object) -> datetime.timedelta | None: + if not isinstance(upstream, httpx.Response): + return None + try: + return upstream.elapsed + except RuntimeError: + return None + + +def error_response_without_credentials(status_code: int, upstream: object, fallback_url: str) -> httpx.Response: + response: Final = httpx.Response( + status_code=status_code, + stream=httpx.ByteStream(_upstream_body(upstream)), + request=_request_without_credentials(upstream, fallback_url), + ) + _ = response.read() + response.headers = httpx.Headers(upstream.headers if isinstance(upstream, httpx.Response) else None) + elapsed: Final = _upstream_elapsed(upstream) + if elapsed is not None: + response.elapsed = elapsed + return response + + +_NO_UPSTREAM_REQUEST_URL: Final = "https://cloud.google.com/vertex-ai/" + + _MINIMAL_ERROR_RESPONSE: httpx.Response | None = None @@ -519,20 +568,11 @@ class RateLimitError(openai.RateLimitError): # headers stay reachable on `e.response.headers` for callers that # explicitly want them; only the proxy-supplied `headers=` kwarg # makes it onto `self.headers`. - _response_headers: Final = getattr(response, "headers", None) if response is not None else None self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None # Mirrors FastAPI HTTPException.detail so the same instance can be # serialized through both the ProxyException and HTTPException paths. self.detail = detail if detail is not None else self.message - self.response = httpx.Response( - status_code=429, - headers=_response_headers, - content=response.content if response is not None else None, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ) + self.response = error_response_without_credentials(429, response, _NO_UPSTREAM_REQUEST_URL) super().__init__( self.message, response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime @@ -757,15 +797,7 @@ class ServiceUnavailableError(openai.APIStatusError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers: Final = getattr(response, "headers", None) if response is not None else None - self.response = httpx.Response( - status_code=self.status_code, - headers=_response_headers, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ) + self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL) super().__init__( self.message, response=self.response, body=None ) # Call the base class constructor with the parameters it needs @@ -808,15 +840,7 @@ class BadGatewayError(openai.APIStatusError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers: Final = getattr(response, "headers", None) if response is not None else None - self.response = httpx.Response( - status_code=self.status_code, - headers=_response_headers, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ) + self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL) super().__init__( self.message, response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime @@ -862,15 +886,7 @@ class InternalServerError(openai.InternalServerError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers: Final = getattr(response, "headers", None) if response is not None else None - self.response = httpx.Response( - status_code=self.status_code, - headers=_response_headers, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ) + self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL) super().__init__( self.message, response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime @@ -908,6 +924,7 @@ class APIError(openai.APIError): max_retries: int | None = None, num_retries: int | None = None, body: object | None = None, + response: httpx.Response | None = None, ): self.status_code = status_code self.message = f"litellm.APIError: {message}" @@ -916,6 +933,8 @@ class APIError(openai.APIError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries + if response is not None: + self.response = response if request is None: request = httpx.Request(method="POST", url="https://api.openai.com/v1") super().__init__(self.message, request=request, body=body) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 64a95eb5f28..5a27f5446ad 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -37,6 +37,7 @@ from ..exceptions import ( ServiceUnavailableError, Timeout, UnprocessableEntityError, + error_response_without_credentials, ) @@ -305,6 +306,16 @@ def _litellm_proxy_response( ) +def _without_credentials(response: httpx.Response | None) -> httpx.Response | None: + if not isinstance(response, httpx.Response): + return response + return error_response_without_credentials(response.status_code, response, "https://docs.litellm.ai/docs") + + +def _upstream_reply(response: httpx.Response | None) -> httpx.Response | None: + return response if response is not None and (response.headers or response.content) else None + + def _map_openai_exception( *, model: str, @@ -315,7 +326,7 @@ def _map_openai_exception( exception_provider: str, extra_information: str, ) -> None: - response: Final = _litellm_proxy_response(original_exception, custom_llm_provider) + response: Final = _without_credentials(_litellm_proxy_response(original_exception, custom_llm_provider)) if (cyber_policy_body := _openai_cyber_policy_body(original_exception)) is not None: raise ContentPolicyViolationError( message=f"ContentPolicyViolationError: {exception_provider} - {original_exception}", @@ -378,6 +389,7 @@ def _map_openai_exception( model=model, llm_provider=custom_llm_provider, litellm_debug_info=extra_information, + response=_upstream_reply(response), ) elif ( ("invalid_request_error" in error_str and "content_policy_violation" in error_str) @@ -449,6 +461,7 @@ def _map_openai_exception( model=model, llm_provider=custom_llm_provider, body=getattr(original_exception, "body", None), + response=_upstream_reply(response), ) elif "Request too large" in error_str: raise RateLimitError( @@ -479,6 +492,7 @@ def _map_openai_exception( model=model, request=_request, litellm_debug_info=extra_information, + response=_upstream_reply(response), ) elif hasattr(original_exception, "status_code"): if original_exception.status_code == 400: @@ -512,6 +526,7 @@ def _map_openai_exception( model=model, llm_provider=custom_llm_provider, litellm_debug_info=extra_information, + response=_upstream_reply(response), ) elif original_exception.status_code == 422: raise BadRequestError( @@ -563,6 +578,7 @@ def _map_openai_exception( llm_provider=custom_llm_provider, litellm_debug_info=extra_information, exception_status_code=original_exception.status_code, + response=_upstream_reply(response), ) else: raise APIError( @@ -572,6 +588,7 @@ def _map_openai_exception( model=model, request=getattr(original_exception, "request", None), litellm_debug_info=extra_information, + response=_upstream_reply(response), ) else: # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors @@ -2019,6 +2036,7 @@ def _map_azure_exception( exception_provider: str, extra_information: str, ) -> None: + response: Final = _without_credentials(getattr(original_exception, "response", None)) message = get_error_message(error_obj=original_exception) if message is None: if hasattr(original_exception, "message"): @@ -2054,7 +2072,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif "DeploymentNotFound" in error_str: raise NotFoundError( @@ -2062,7 +2080,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif azure_error_code == "content_policy_violation" or ExceptionCheckers.is_azure_content_policy_violation_error( error_str @@ -2076,6 +2094,7 @@ def _map_azure_exception( model=model, extra_information=extra_information, original_exception=original_exception, + response=response, ) elif azure_error_code == "invalid_encrypted_content" or "could not be verified" in error_str: helpful_message: Final = ( @@ -2094,7 +2113,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, body=getattr(original_exception, "body", None), ) elif getattr(original_exception, "status_code", None) in (None, 400) and ( @@ -2106,7 +2125,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, body=getattr(original_exception, "body", None), ) elif "invalid_request_error" in error_str and getattr(original_exception, "status_code", None) in (None, 400): @@ -2115,7 +2134,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, body=getattr(original_exception, "body", None), ) elif "The api_key client option must be set either by passing api_key to the client or by setting" in error_str: @@ -2124,7 +2143,7 @@ def _map_azure_exception( llm_provider=custom_llm_provider, model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif "Connection error" in error_str: raise APIConnectionError( @@ -2140,7 +2159,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, body=getattr(original_exception, "body", None), ) elif original_exception.status_code == 401: @@ -2149,7 +2168,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 404: raise NotFoundError( @@ -2157,7 +2176,7 @@ def _map_azure_exception( llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 408: raise Timeout( @@ -2165,6 +2184,7 @@ def _map_azure_exception( model=model, litellm_debug_info=extra_information, llm_provider="azure", + response=_upstream_reply(response), ) elif original_exception.status_code == 422: raise BadRequestError( @@ -2172,7 +2192,7 @@ def _map_azure_exception( model=model, llm_provider="azure", litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 429: raise RateLimitError( @@ -2180,7 +2200,7 @@ def _map_azure_exception( model=model, llm_provider="azure", litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 502: raise BadGatewayError( @@ -2188,7 +2208,7 @@ def _map_azure_exception( model=model, llm_provider="azure", litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 503: raise ServiceUnavailableError( @@ -2196,7 +2216,7 @@ def _map_azure_exception( model=model, llm_provider="azure", litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, ) elif original_exception.status_code == 504: # gateway timeout error raise Timeout( @@ -2205,6 +2225,7 @@ def _map_azure_exception( litellm_debug_info=extra_information, llm_provider="azure", exception_status_code=original_exception.status_code, + response=_upstream_reply(response), ) else: raise APIError( @@ -2214,6 +2235,7 @@ def _map_azure_exception( litellm_debug_info=extra_information, model=model, request=httpx.Request(method="POST", url="https://openai.com/"), + response=_upstream_reply(response), ) else: # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors @@ -2589,17 +2611,20 @@ def exception_type( model=model, llm_provider=custom_llm_provider, litellm_debug_info=extra_information, + response=_upstream_reply( + _without_credentials(_litellm_proxy_response(mappable_exception, custom_llm_provider)) + ), ) if ( custom_llm_provider == "litellm_proxy" ): # handle special case where calling litellm proxy + exception str contains error message extract_and_raise_litellm_exception( - response=_litellm_proxy_response(mappable_exception, custom_llm_provider), + response=_without_credentials(_litellm_proxy_response(mappable_exception, custom_llm_provider)), error_str=error_str, model=model, custom_llm_provider=custom_llm_provider, - body=getattr(original_exception, "body", None), + body=getattr(mappable_exception, "body", None), ) if ( custom_llm_provider == "openai" diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index a1ba8de6d7f..69ffac5d6f9 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -42,6 +42,7 @@ from ..openai.common_utils import ( build_output_token_limit_response, is_output_token_limit_error, model_dump_provider_response, + upstream_error_response, ) from .common_utils import ( AzureOpenAIError, @@ -399,6 +400,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): raise AzureOpenAIError( status_code=status_code, message=str(e), + response=upstream_error_response(e), headers=error_headers, body=error_body, ) @@ -662,6 +664,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): raise AzureOpenAIError( status_code=status_code, message=message, + response=upstream_error_response(e), headers=error_headers, body=error_body, ) @@ -855,7 +858,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) error_text = error_response.text - raise AzureOpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise AzureOpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) async def make_async_azure_httpx_request( self, @@ -1190,7 +1198,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except Exception as e: error_code: Final[object] = getattr(e, "status_code", None) if error_code is not None: - raise AzureOpenAIError(status_code=error_code, message=str(e)) + raise AzureOpenAIError(status_code=error_code, message=str(e), response=upstream_error_response(e)) else: raise AzureOpenAIError(status_code=500, message=str(e)) diff --git a/litellm/llms/azure/completion/handler.py b/litellm/llms/azure/completion/handler.py index 34e51eaf1ec..c2744e7f750 100644 --- a/litellm/llms/azure/completion/handler.py +++ b/litellm/llms/azure/completion/handler.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.prompt_templates.factory import prompt_factory from litellm.utils import CustomStreamWrapper, ModelResponse, TextCompletionResponse +from ...openai.common_utils import upstream_error_response from ...openai.completion.transformation import OpenAITextCompletionConfig from ..common_utils import AzureOpenAIError, BaseAzureLLM @@ -183,7 +184,12 @@ class AzureTextCompletion(BaseAzureLLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers) + raise AzureOpenAIError( + status_code=status_code, + message=str(e), + headers=error_headers, + response=upstream_error_response(e), + ) async def acompletion( self, @@ -244,7 +250,12 @@ class AzureTextCompletion(BaseAzureLLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers) + raise AzureOpenAIError( + status_code=status_code, + message=str(e), + headers=error_headers, + response=upstream_error_response(e), + ) def streaming( self, @@ -355,4 +366,9 @@ class AzureTextCompletion(BaseAzureLLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers) + raise AzureOpenAIError( + status_code=status_code, + message=str(e), + headers=error_headers, + response=upstream_error_response(e), + ) diff --git a/litellm/llms/azure/exception_mapping.py b/litellm/llms/azure/exception_mapping.py index f7146a36384..6eb5161f7e4 100644 --- a/litellm/llms/azure/exception_mapping.py +++ b/litellm/llms/azure/exception_mapping.py @@ -1,5 +1,7 @@ from typing import Final +import httpx + from litellm.exceptions import ContentPolicyViolationError @@ -14,6 +16,7 @@ class AzureOpenAIExceptionMapping: model: str, extra_information: str, original_exception: Exception, + response: httpx.Response | None, ) -> ContentPolicyViolationError: """ Create a content policy violation error @@ -39,7 +42,7 @@ class AzureOpenAIExceptionMapping: llm_provider="azure", model=model, litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), + response=response, provider_specific_fields={ # Preserve legacy key for backward compatibility. "innererror": inner_error, diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 6c416124d1a..aa14a985afe 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -29,6 +29,7 @@ if TYPE_CHECKING: from aiohttp import ClientSession import litellm +from litellm.exceptions import error_response_without_credentials from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object from litellm.litellm_core_utils.token_counter import token_counter, usage_tokenizer_model @@ -148,6 +149,14 @@ class OpenAIError(BaseLLMException): ) +def upstream_error_response(error: Exception) -> httpx.Response | None: + if not isinstance(error, (openai.APIStatusError, httpx.HTTPStatusError)): + return None + return error_response_without_credentials( + error.response.status_code, error.response, "https://docs.litellm.ai/docs" + ) + + _DECODED_JSON: Final = TypeAdapter(object) diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 278427e8d77..8bfdd536f90 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -10,7 +10,7 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser from litellm.types.utils import LlmProviders, ModelResponse, TextCompletionResponse from litellm.utils import ProviderConfigManager -from ..common_utils import BaseOpenAILLM, OpenAIError +from ..common_utils import BaseOpenAILLM, OpenAIError, upstream_error_response from ..workload_identity import build_async_openai_client, build_openai_client from .transformation import OpenAITextCompletionConfig @@ -164,7 +164,12 @@ class OpenAITextCompletion(BaseLLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) async def acompletion( self, @@ -220,7 +225,12 @@ class OpenAITextCompletion(BaseLLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) def streaming( self, @@ -261,7 +271,12 @@ class OpenAITextCompletion(BaseLLM): error_response: object = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) streamwrapper: Final = CustomStreamWrapper( completion_stream=response, model=model, diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index 01f01d3de1d..504dfb50015 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -14,7 +14,7 @@ from litellm.utils import ProviderConfigManager from ...base_llm.image_variations.transformation import BaseImageVariationConfig from ...custom_httpx.llm_http_handler import LiteLLMLoggingObj -from ..common_utils import OpenAIError +from ..common_utils import OpenAIError, upstream_error_response from ..workload_identity import build_async_openai_client, build_openai_client @@ -129,7 +129,12 @@ class OpenAIImageVariationsHandler: error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) def image_variations( self, @@ -245,4 +250,9 @@ class OpenAIImageVariationsHandler: error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index af2e85c5f20..e0a23d6c99f 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -64,6 +64,7 @@ from .common_utils import ( is_openai_backed_api_base, is_output_token_limit_error, model_dump_provider_response, + upstream_error_response, ) from .workload_identity import ( build_async_openai_client, @@ -983,6 +984,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): raise OpenAIError( status_code=status_code, message=error_text, + response=upstream_error_response(e), headers=error_headers, body=error_body, ) @@ -1110,6 +1112,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): raise OpenAIError( status_code=status_code, message=message, + response=upstream_error_response(e), headers=error_headers, body=exception_body, ) @@ -1282,6 +1285,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): raise OpenAIError( status_code=getattr(e, "status_code", 500), message=str(e), + response=upstream_error_response(e), headers=error_headers, body=exception_body, ) @@ -1409,7 +1413,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) def embedding( self, @@ -1499,7 +1508,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) + raise OpenAIError( + status_code=status_code, + message=error_text, + response=upstream_error_response(e), + headers=error_headers, + ) async def aimage_generation( self, @@ -1666,7 +1680,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): original_response=str(e), ) if hasattr(e, "status_code"): - raise OpenAIError(status_code=getattr(e, "status_code", 500), message=str(e)) + raise OpenAIError( + status_code=getattr(e, "status_code", 500), + message=str(e), + response=upstream_error_response(e), + ) else: raise OpenAIError(status_code=500, message=str(e)) diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 19e5387da9a..88501a23235 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -24,6 +24,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.proxy.common_utils.openai_error_payload import openai_error_type from litellm.types.utils import LlmProviders @@ -54,8 +55,8 @@ def test_openai_cyber_policy_uses_the_structured_error_code(shape: str, code: st assert caught.value.code == code assert caught.value.body == body assert caught.value.status_code == 400 - assert caught.value.response is response - assert caught.value.request is request + assert caught.value.response.status_code == 400 + assert caught.value.request.url == request.url assert caught.value.response.headers["x-request-id"] == "cyber-policy-request" @@ -2940,3 +2941,80 @@ def test_transport_timeout_maps_to_timeout( assert exc_info.value.status_code == 408 assert f"Request timed out. Error_str: {reason}" in str(exc_info.value) + + +def test_bedrock_server_error_keeps_the_aws_body_without_the_signed_request() -> None: + body: Final = '{"message":"Bedrock is unable to process your request."}' + signed_request: Final = httpx.Request( + "POST", + "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-opus-5-5/converse", + headers={"authorization": "AWS4-HMAC-SHA256 Credential=AKIDSIGNED/20261029/us-east-1/bedrock/aws4_request"}, + content=b'{"messages":[{"role":"user","content":[{"text":"hi"}]}]}', + ) + upstream: Final = httpx.Response( + 500, content=body, headers={"x-amzn-requestid": "req-aws"}, request=signed_request + ) + + with pytest.raises(litellm.ServiceUnavailableError) as raised: + exception_type( + model="bedrock/anthropic.claude-opus-5-5", + original_exception=BedrockError(status_code=500, message=body, response=upstream), + custom_llm_provider="bedrock", + ) + + assert raised.value.response.text == body + assert raised.value.response.headers["x-amzn-requestid"] == "req-aws" + assert "AKIDSIGNED" not in str(raised.value.request.headers) + assert raised.value.request.content == b"" + + +@pytest.mark.parametrize("provider", ("gemini", "vertex_ai")) +def test_gemini_server_error_response_carries_the_upstream_error_body(provider: str) -> None: + body: Final = '{"error":{"code":500,"message":"Internal error encountered.","status":"INTERNAL"}}' + + with pytest.raises(litellm.InternalServerError) as raised: + exception_type( + model="gemini-3.1-pro", + original_exception=VertexAIError(status_code=500, message=body, headers=None), + custom_llm_provider=provider, + ) + + assert raised.value.response.status_code == 500 + assert raised.value.response.json() == json.loads(body) + + +@pytest.mark.parametrize( + ("status_code", "message", "mapped_error"), + ( + (500, "The server had an error processing your request. Sorry about that!", litellm.InternalServerError), + (520, "Web server is returning an unknown error", litellm.InternalServerError), + (504, "A timeout occurred", litellm.Timeout), + (500, "Mistral API raised a streaming error", litellm.APIError), + ), +) +def test_openai_error_matched_by_its_message_keeps_the_upstream_body( + status_code: int, message: str, mapped_error: type[Exception] +) -> None: + body: Final = json.dumps({"error": {"message": message, "type": "server_error", "param": None, "code": None}}) + sent_request: Final = httpx.Request( + "POST", + "https://api.openai.com/v1/chat/completions", + headers={"authorization": "Bearer sk-openai-secret"}, + content=b'{"messages":[{"role":"user","content":"hi"}]}', + ) + upstream: Final = httpx.Response( + status_code, content=body, headers={"x-request-id": "req-openai"}, request=sent_request + ) + + with pytest.raises(mapped_error) as raised: + exception_type( + model="gpt-5.6", + original_exception=openai.APIStatusError(message, response=upstream, body=json.loads(body)), + custom_llm_provider="openai", + ) + + error_response: Final = getattr(raised.value, "response") + assert error_response.text == body + assert error_response.headers["x-request-id"] == "req-openai" + assert "sk-openai-secret" not in str(error_response.request.headers) + assert error_response.request.content == b"" diff --git a/tests/unit/llms/azure/test_azure_exception_mapping.py b/tests/unit/llms/azure/test_azure_exception_mapping.py index cbbf70c691b..018b499f401 100644 --- a/tests/unit/llms/azure/test_azure_exception_mapping.py +++ b/tests/unit/llms/azure/test_azure_exception_mapping.py @@ -1,3 +1,4 @@ +import datetime import json from collections.abc import Awaitable, Callable, Mapping from typing import Final @@ -12,6 +13,7 @@ import respx import litellm from litellm.exceptions import BadRequestError, ContentPolicyViolationError, ContextWindowExceededError from litellm.litellm_core_utils.exception_mapping_utils import exception_type +from litellm.llms.azure.common_utils import AzureOpenAIError @respx.mock @@ -632,3 +634,202 @@ def test_azure_content_policy_wins_over_context_length_code(monkeypatch: pytest. error: Final = _azure_chat_error(monkeypatch, _azure_error_response(400, policy_error)) assert type(error) is ContentPolicyViolationError + + +_AZURE_REFUSAL_BODY: Final = '{"error":{"message":"azure refused","type":"upstream_refusal","code":"refused"}}' +_AZURE_DEPLOYMENT_NOT_FOUND_BODY: Final = '{"error":{"code":"DeploymentNotFound","message":"no such deployment"}}' +_AZURE_INTERNAL_ERROR_BODY: Final = '{"error":{"code":"InternalServerError","message":"Internal server error"}}' +_AZURE_CONTENT_FILTER_BODY: Final = json.dumps( + { + "error": { + "message": "The response was filtered due to the prompt triggering Azure OpenAI's content management policy.", + "type": None, + "param": "prompt", + "code": "content_filter", + "status": 400, + "innererror": { + "code": "ResponsibleAIPolicyViolation", + "content_filter_result": {"violence": {"filtered": True, "severity": "high"}}, + }, + } + } +) +_AZURE_KEY: Final = "azure-provider-secret" +_AZURE_BASE: Final = "https://azure-upstream.openai.azure.com" +_AZURE_ARGS: Final = {"api_base": _AZURE_BASE, "api_version": "2024-10-21", "api_key": _AZURE_KEY, "max_retries": 0} + + +async def _azure_completion() -> None: + litellm.completion(model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], **_AZURE_ARGS) + + +async def _azure_completion_stream() -> None: + list( + litellm.completion( + model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **_AZURE_ARGS + ) + ) + + +async def _azure_acompletion() -> None: + await litellm.acompletion(model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], **_AZURE_ARGS) + + +async def _azure_acompletion_stream() -> None: + stream: Final = await litellm.acompletion( + model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **_AZURE_ARGS + ) + async for _ in stream: + pass + + +async def _azure_embedding() -> None: + litellm.embedding(model="azure/text-embedding-3-large", input=["hi"], **_AZURE_ARGS) + + +async def _azure_aembedding() -> None: + await litellm.aembedding(model="azure/text-embedding-3-large", input=["hi"], **_AZURE_ARGS) + + +async def _azure_image_generation() -> None: + litellm.image_generation(model="azure/gpt-image-2", prompt="a cat", **_AZURE_ARGS) + + +async def _azure_text_completion() -> None: + litellm.text_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", **_AZURE_ARGS) + + +async def _azure_text_completion_stream() -> None: + list(litellm.text_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", stream=True, **_AZURE_ARGS)) + + +async def _azure_atext_completion() -> None: + await litellm.atext_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", **_AZURE_ARGS) + + +async def _azure_atext_completion_stream() -> None: + stream: Final = await litellm.atext_completion( + model="azure_text/gpt-35-turbo-instruct", prompt="hi", stream=True, **_AZURE_ARGS + ) + async for _ in stream: + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + ( + _azure_completion, + _azure_completion_stream, + _azure_acompletion, + _azure_acompletion_stream, + _azure_embedding, + _azure_aembedding, + _azure_image_generation, + _azure_text_completion, + _azure_text_completion_stream, + _azure_atext_completion, + _azure_atext_completion_stream, + ), + ids=lambda route: route.__name__.removeprefix("_azure_"), +) +@pytest.mark.parametrize( + ("status_code", "body", "mapped_error"), + ( + (400, _AZURE_REFUSAL_BODY, BadRequestError), + (400, _AZURE_CONTENT_FILTER_BODY, BadRequestError), + (403, _AZURE_REFUSAL_BODY, litellm.APIError), + (404, _AZURE_DEPLOYMENT_NOT_FOUND_BODY, litellm.NotFoundError), + (500, _AZURE_INTERNAL_ERROR_BODY, openai.APIError), + (502, _AZURE_REFUSAL_BODY, litellm.BadGatewayError), + (504, _AZURE_REFUSAL_BODY, litellm.Timeout), + ), +) +@respx.mock +async def test_azure_error_keeps_the_upstream_body_without_the_api_key( + route: Callable[[], Awaitable[None]], + status_code: int, + body: str, + mapped_error: type[Exception], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx.post(url__regex=r"https://azure-upstream\.openai\.azure\.com/.*").mock( + return_value=httpx.Response(status_code, content=body, headers={"content-type": "application/json"}) + ) + + with pytest.raises(mapped_error) as raised: + await route() + + error_response: Final = getattr(raised.value, "response") + assert error_response.status_code == status_code + assert error_response.text == body + assert error_response.json() == json.loads(body) + for exposed in (error_response.request, getattr(raised.value, "request")): + assert _AZURE_KEY not in str(exposed.headers) + assert exposed.content == b"" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + ( + _azure_completion, + _azure_completion_stream, + _azure_acompletion, + _azure_acompletion_stream, + _azure_embedding, + _azure_aembedding, + _azure_text_completion, + _azure_atext_completion, + _azure_atext_completion_stream, + ), + ids=lambda route: route.__name__.removeprefix("_azure_"), +) +@respx.mock +async def test_azure_error_response_keeps_the_upstream_elapsed_time( + route: Callable[[], Awaitable[None]], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx.post(url__regex=r"https://azure-upstream\.openai\.azure\.com/.*").mock( + return_value=httpx.Response(400, content=_AZURE_REFUSAL_BODY, headers={"content-type": "application/json"}) + ) + + with pytest.raises(BadRequestError) as raised: + await route() + + assert raised.value.response.elapsed > datetime.timedelta(0) + + +def test_azure_client_side_missing_key_error_keeps_the_empty_placeholder_response() -> None: + missing_key: Final = openai.OpenAIError( + "The api_key client option must be set either by passing api_key to the client or by setting the " + "OPENAI_API_KEY environment variable" + ) + + with pytest.raises(litellm.AuthenticationError) as raised: + exception_type(model="azure/gpt-4.1-mini", original_exception=missing_key, custom_llm_provider="azure") + + assert raised.value.response.status_code == 401 + assert raised.value.response.text == "" + + +@pytest.mark.parametrize( + ("status_code", "message", "mapped_error"), + ( + (500, "azure_client is not an instance of AzureOpenAI or OpenAI", litellm.APIError), + (408, "the deployment did not answer in time", litellm.Timeout), + (504, "the deployment did not answer in time", litellm.Timeout), + ), +) +def test_azure_error_litellm_raises_itself_has_no_upstream_response( + status_code: int, message: str, mapped_error: type[Exception] +) -> None: + with pytest.raises(mapped_error) as raised: + exception_type( + model="azure/gpt-4.1-mini", + original_exception=AzureOpenAIError(status_code=status_code, message=message), + custom_llm_provider="azure", + ) + + assert getattr(raised.value, "response", None) is None diff --git a/tests/unit/llms/openai/completion/test_completion_handler.py b/tests/unit/llms/openai/completion/test_completion_handler.py index 9e2754eef68..628c9bbd783 100644 --- a/tests/unit/llms/openai/completion/test_completion_handler.py +++ b/tests/unit/llms/openai/completion/test_completion_handler.py @@ -93,24 +93,34 @@ async def test_acompletion_forwards_client_headers_to_provider( assert request_headers["x-mycorp-llmcall-id"] == "abc-123" +_REJECTION_BODY: Final = '{"error":{"message":"bad prompt","type":"invalid_request_error"}}' + + @pytest.fixture def rejected_completions_endpoint(): return respx.post("https://api.openai.com/v1/completions").mock( return_value=Response( 400, - json={"error": {"message": "bad prompt", "type": "invalid_request_error"}}, - headers={"x-request-id": "req-1"}, + content=_REJECTION_BODY, + headers={"x-request-id": "req-1", "content-type": "application/json"}, ) ) +def _assert_carries_the_upstream_rejection(rejected: litellm.BadRequestError) -> None: + assert rejected.status_code == 400 + assert rejected.litellm_response_headers["x-request-id"] == "req-1" + assert rejected.response.text == _REJECTION_BODY + assert rejected.response.json() == json.loads(_REJECTION_BODY) + assert rejected.response.headers["x-request-id"] == "req-1" + + @respx.mock def test_completion_surfaces_the_status_and_headers_of_a_rejected_request(rejected_completions_endpoint): with pytest.raises(litellm.BadRequestError) as rejected: text_completion(model="gpt-3.5-turbo-instruct", prompt="hello") - assert rejected.value.status_code == 400 - assert rejected.value.litellm_response_headers["x-request-id"] == "req-1" + _assert_carries_the_upstream_rejection(rejected.value) @respx.mock @@ -118,8 +128,7 @@ def test_streaming_completion_surfaces_the_status_and_headers_of_a_rejected_requ with pytest.raises(litellm.BadRequestError) as rejected: list(text_completion(model="gpt-3.5-turbo-instruct", prompt="hello", stream=True)) - assert rejected.value.status_code == 400 - assert rejected.value.litellm_response_headers["x-request-id"] == "req-1" + _assert_carries_the_upstream_rejection(rejected.value) @respx.mock @@ -130,8 +139,7 @@ async def test_acompletion_surfaces_the_status_and_headers_of_a_rejected_request with pytest.raises(litellm.BadRequestError) as rejected: await atext_completion(model="gpt-3.5-turbo-instruct", prompt="hello") - assert rejected.value.status_code == 400 - assert rejected.value.litellm_response_headers["x-request-id"] == "req-1" + _assert_carries_the_upstream_rejection(rejected.value) @respx.mock diff --git a/tests/unit/llms/openai/image_variations/__init__.py b/tests/unit/llms/openai/image_variations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/image_variations/test_openai_image_variations_handler.py b/tests/unit/llms/openai/image_variations/test_openai_image_variations_handler.py new file mode 100644 index 00000000000..83ddbec6038 --- /dev/null +++ b/tests/unit/llms/openai/image_variations/test_openai_image_variations_handler.py @@ -0,0 +1,51 @@ +import json +from io import BytesIO +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.llms.openai.common_utils import OpenAIError + +_UPSTREAM_ERROR_BODY: Final = '{"error":{"message":"image is not a valid PNG","type":"invalid_request_error"}}' + + +def _upstream_rejection(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + content=_UPSTREAM_ERROR_BODY, + headers={"x-request-id": "req_variation_400", "content-type": "application/json"}, + ) + + +_VARIATION_KWARGS: Final = { + "custom_llm_provider": "openai", + "api_key": "transport-only", + "api_base": "https://upstream.example/v1", + "max_retries": 0, +} + + +async def _request_variation(is_async: bool) -> None: + if is_async: + await litellm.aimage_variation(image=BytesIO(b"image-bytes"), **_VARIATION_KWARGS) + return + litellm.image_variation(image=BytesIO(b"image-bytes"), **_VARIATION_KWARGS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +async def test_rejected_image_variation_carries_the_upstream_error_response( + is_async: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + transport: Final = httpx.MockTransport(_upstream_rejection) + monkeypatch.setattr(litellm, "client_session", httpx.Client(transport=transport)) + monkeypatch.setattr(litellm, "aclient_session", httpx.AsyncClient(transport=transport)) + + with pytest.raises((litellm.BadRequestError, OpenAIError)) as raised: + await _request_variation(is_async) + + assert raised.value.response.text == _UPSTREAM_ERROR_BODY + assert raised.value.response.json() == json.loads(_UPSTREAM_ERROR_BODY) + assert raised.value.response.headers["x-request-id"] == "req_variation_400" diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 5b27a446fb3..3e1cbad9308 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -1,8 +1,9 @@ import asyncio +import gzip import importlib import json import os -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from types import ModuleType from typing import Final, Literal, Protocol from unittest.mock import MagicMock, Mock, patch @@ -375,6 +376,277 @@ async def test_acompletion_preserves_public_errors_for_both_http_clients( assert not http_client.is_closed +_UPSTREAM_ERROR_BODY: Final = ( + '{"error":{"message":"upstream rejected the request","type":"invalid_request_error",' + '"param":"response_format.type","code":"invalid_value"}}' +) +_UPSTREAM_BASE: Final = "https://upstream.example/v1" +_UPSTREAM_TRANSPORT: Final = httpx.MockTransport( + lambda _request: httpx.Response( + 400, + content=_UPSTREAM_ERROR_BODY, + headers={"x-request-id": "req_upstream_400", "x-upstream-trace": "trace-7", "content-type": "application/json"}, + ) +) +_MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +def _rejecting_client() -> OpenAI: + return OpenAI( + api_key="transport-only", + base_url=_UPSTREAM_BASE, + max_retries=0, + http_client=httpx.Client(transport=_UPSTREAM_TRANSPORT), + ) + + +def _rejecting_async_client() -> AsyncOpenAI: + return AsyncOpenAI( + api_key="transport-only", + base_url=_UPSTREAM_BASE, + max_retries=0, + http_client=httpx.AsyncClient(transport=_UPSTREAM_TRANSPORT), + ) + + +def _upstream(provider: str, model: str) -> dict[str, object]: + return {"model": f"{provider}/{model}", "api_base": _UPSTREAM_BASE, "api_key": "transport-only", "num_retries": 0} + + +async def _completion(provider: str) -> None: + litellm.completion(messages=_MESSAGES, client=_rejecting_client(), **_upstream(provider, "gpt-5.6")) + + +async def _completion_stream(provider: str) -> None: + stream: Final = litellm.completion( + messages=_MESSAGES, client=_rejecting_client(), stream=True, **_upstream(provider, "gpt-5.6") + ) + list(stream) + + +async def _acompletion(provider: str) -> None: + await litellm.acompletion(messages=_MESSAGES, client=_rejecting_async_client(), **_upstream(provider, "gpt-5.6")) + + +async def _acompletion_stream(provider: str) -> None: + stream: Final = await litellm.acompletion( + messages=_MESSAGES, client=_rejecting_async_client(), stream=True, **_upstream(provider, "gpt-5.6") + ) + async for _ in stream: + pass + + +async def _embedding(provider: str) -> None: + litellm.embedding(input=["hi"], client=_rejecting_client(), **_upstream(provider, "text-embedding-3-large")) + + +async def _aembedding(provider: str) -> None: + await litellm.aembedding( + input=["hi"], client=_rejecting_async_client(), **_upstream(provider, "text-embedding-3-large") + ) + + +async def _image_generation(provider: str) -> None: + litellm.image_generation(prompt="a cat", client=_rejecting_client(), **_upstream(provider, "gpt-image-2")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "litellm_proxy")) +@pytest.mark.parametrize( + "route", + (_completion, _completion_stream, _acompletion, _acompletion_stream, _embedding, _aembedding, _image_generation), + ids=lambda route: route.__name__.lstrip("_"), +) +async def test_mapped_error_carries_the_upstream_error_response( + provider: str, route: Callable[[str], Awaitable[None]] +) -> None: + with pytest.raises(litellm.BadRequestError) as raised: + await route(provider) + + assert raised.value.response.text == _UPSTREAM_ERROR_BODY + assert raised.value.response.json() == json.loads(_UPSTREAM_ERROR_BODY) + assert raised.value.response.headers["x-request-id"] == "req_upstream_400" + assert raised.value.response.headers["x-upstream-trace"] == "trace-7" + + +_UPSTREAM_SERVER_ERROR_BODY: Final = ( + '{"error":{"message":"The server had an error while processing your request.","type":"server_error",' + '"param":null,"code":null}}' +) +_CHUNKED_UPSTREAM_HEADERS: Final = { + "content-type": "application/json", + "transfer-encoding": "chunked", + "x-request-id": "req_upstream_5xx", +} + + +def _chunked_upstream_client(status_code: int, content_encoding: Literal["gzip"] | None) -> httpx.AsyncClient: + body: Final = _UPSTREAM_SERVER_ERROR_BODY.encode() + encoding_headers: Final = {"content-encoding": content_encoding} if content_encoding else {} + return httpx.AsyncClient( + transport=httpx.MockTransport( + lambda _request: httpx.Response( + status_code, + headers={**_CHUNKED_UPSTREAM_HEADERS, **encoding_headers}, + stream=httpx.ByteStream(gzip.compress(body) if content_encoding else body), + ) + ) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "litellm_proxy")) +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize( + ("status_code", "mapped_error"), + ( + (429, litellm.RateLimitError), + (500, litellm.InternalServerError), + (502, litellm.BadGatewayError), + (503, litellm.ServiceUnavailableError), + ), +) +@pytest.mark.parametrize("content_encoding", (None, "gzip")) +async def test_retryable_error_keeps_the_decoded_upstream_body_and_its_headers( + provider: str, + stream: bool, + status_code: int, + mapped_error: type[Exception], + content_encoding: Literal["gzip"] | None, +) -> None: + client: Final = AsyncOpenAI( + api_key="transport-only", + base_url=_UPSTREAM_BASE, + max_retries=0, + http_client=_chunked_upstream_client(status_code, content_encoding), + ) + + with pytest.raises(mapped_error) as raised: + await litellm.acompletion(messages=_MESSAGES, client=client, stream=stream, **_upstream(provider, "gpt-5.6")) + + error_response: Final = getattr(raised.value, "response") + assert error_response.status_code == status_code + assert error_response.text == _UPSTREAM_SERVER_ERROR_BODY + assert error_response.json() == json.loads(_UPSTREAM_SERVER_ERROR_BODY) + assert dict(error_response.headers) == { + **_CHUNKED_UPSTREAM_HEADERS, + **({"content-encoding": content_encoding} if content_encoding else {}), + } + + +_REFUSAL_BODY: Final = '{"error":{"message":"upstream refused","type":"upstream_refusal","param":null,"code":null}}' +_GATEWAY_TIMEOUT_BODY: Final = ( + '{"error":{"message":"litellm.Timeout: APITimeoutError - Request timed out.","type":"invalid_request_error",' + '"param":null,"code":"408"}}' +) +_PROVIDER_KEY: Final = "sk-provider-secret" +_QUERY_SECRET: Final = "query-secret" + + +@pytest.mark.parametrize("provider", ("openai", "litellm_proxy")) +@pytest.mark.parametrize( + ("status_code", "mapped_error", "body"), + ( + (400, litellm.BadRequestError, _REFUSAL_BODY), + (403, litellm.APIError, _REFUSAL_BODY), + (408, litellm.Timeout, _REFUSAL_BODY), + (408, litellm.Timeout, _GATEWAY_TIMEOUT_BODY), + (409, litellm.APIError, _REFUSAL_BODY), + (413, litellm.APIError, _REFUSAL_BODY), + (429, litellm.RateLimitError, _REFUSAL_BODY), + (500, litellm.InternalServerError, _REFUSAL_BODY), + (504, litellm.Timeout, _REFUSAL_BODY), + ), +) +def test_mapped_error_keeps_the_upstream_body_without_the_provider_credentials( + provider: str, status_code: int, mapped_error: type[Exception], body: str +) -> None: + client: Final = OpenAI( + api_key=_PROVIDER_KEY, + base_url=_UPSTREAM_BASE, + default_query={"key": _QUERY_SECRET}, + max_retries=0, + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda _request: httpx.Response( + status_code, content=body, headers={"content-type": "application/json"} + ) + ) + ), + ) + + with pytest.raises(mapped_error) as raised: + litellm.completion( + model=f"{provider}/gpt-5.6", + messages=_MESSAGES, + client=client, + api_base=_UPSTREAM_BASE, + api_key=_PROVIDER_KEY, + num_retries=0, + ) + + error_response: Final = getattr(raised.value, "response") + assert error_response.status_code == status_code + assert error_response.text == body + assert error_response.json() == json.loads(body) + assert str(error_response.request.url) == f"{_UPSTREAM_BASE}/chat/completions" + for exposed in (error_response.request, getattr(raised.value, "request")): + assert _QUERY_SECRET not in str(exposed.url) + assert _PROVIDER_KEY not in str(exposed.headers) + assert exposed.content == b"" + + +def _time_out(request: httpx.Request) -> httpx.Response: + raise httpx.ReadTimeout("no reply", request=request) + + +async def _sync_completion_timing_out(provider: str) -> None: + litellm.completion( + model=f"{provider}/gpt-5.6", + messages=_MESSAGES, + client=OpenAI( + api_key=_PROVIDER_KEY, + base_url=_UPSTREAM_BASE, + max_retries=0, + http_client=httpx.Client(transport=httpx.MockTransport(_time_out)), + ), + api_base=_UPSTREAM_BASE, + api_key=_PROVIDER_KEY, + num_retries=0, + ) + + +async def _async_stream_timing_out(provider: str) -> None: + stream: Final = await litellm.acompletion( + model=f"{provider}/gpt-5.6", + messages=_MESSAGES, + stream=True, + client=AsyncOpenAI( + api_key=_PROVIDER_KEY, + base_url=_UPSTREAM_BASE, + max_retries=0, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(_time_out)), + ), + api_base=_UPSTREAM_BASE, + api_key=_PROVIDER_KEY, + num_retries=0, + ) + async for _ in stream: + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "litellm_proxy")) +@pytest.mark.parametrize("call", (_sync_completion_timing_out, _async_stream_timing_out)) +async def test_timeout_before_any_reply_has_no_upstream_response( + provider: str, call: Callable[[str], Awaitable[None]] +) -> None: + with pytest.raises(litellm.Timeout) as raised: + await call(provider) + + assert getattr(raised.value, "response", None) is None + + @pytest.mark.asyncio async def test_acompletion_streams_text_deltas_over_injected_transport( http_backend: ModuleType, sdk_provider: str, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index aeab7a3d348..c867db2fd02 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -2,6 +2,7 @@ import asyncio import contextvars import copy import datetime +import gzip import json from collections.abc import AsyncIterable, Mapping from types import MappingProxyType, SimpleNamespace @@ -13,6 +14,7 @@ from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stub import anthropic import httpx +import openai import pytest import respx from fastapi import HTTPException, Request, Response, status @@ -11675,6 +11677,67 @@ async def test_handle_llm_api_exception_forwards_litellm_response_headers_when_r assert exc_info.value.headers["llm_provider-x-request-id"] == "req_openai_400" +_UPSTREAM_REJECTION_BODY: Final = b'{"error":{"message":"upstream said no","type":"upstream_rejection"}}' + + +def _upstream_rejection_wire(content_encoding: Literal["gzip"] | None) -> tuple[bytes, dict[str, str]]: + wire_body: Final = gzip.compress(_UPSTREAM_REJECTION_BODY) if content_encoding else _UPSTREAM_REJECTION_BODY + return wire_body, { + "x-request-id": "req_upstream", + "x-upstream-trace": "trace-7", + "content-type": "application/json", + "content-length": str(len(wire_body)), + **({"content-encoding": content_encoding} if content_encoding else {}), + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status_code", "mapped_error"), + ( + (400, litellm.BadRequestError), + (429, litellm.RateLimitError), + (503, litellm.ServiceUnavailableError), + ), +) +@pytest.mark.parametrize("content_encoding", (None, "gzip")) +async def test_handle_llm_api_exception_forwards_exactly_the_upstream_headers_of_a_mapped_openai_error( + status_code: int, mapped_error: type[Exception], content_encoding: Literal["gzip"] | None +) -> None: + wire_body, upstream_headers = _upstream_rejection_wire(content_encoding) + transport: Final = httpx.MockTransport( + lambda _request: httpx.Response(status_code, content=wire_body, headers=upstream_headers) + ) + with pytest.raises(mapped_error) as mapped: + await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + num_retries=0, + client=openai.AsyncOpenAI( + api_key="transport-only", + base_url="https://upstream.example/v1", + max_retries=0, + http_client=httpx.AsyncClient(transport=transport), + ), + ) + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + with pytest.raises(ProxyException) as forwarded: + await ProxyBaseLLMRequestProcessing(data={}).handle_llm_api_exception( + e=mapped.value, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=proxy_logging_obj, + ) + + assert forwarded.value.code == str(status_code) + assert "upstream said no" in forwarded.value.message + assert {name: value for name, value in forwarded.value.headers.items() if name.startswith("llm_provider-")} == { + f"llm_provider-{name}": value for name, value in upstream_headers.items() + } + + @pytest.mark.asyncio async def test_handle_llm_api_exception_logs_bug_report_for_unmapped_error( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/test_exception_header_preservation.py b/tests/unit/test_exception_header_preservation.py index 7acacf05226..c5a0868f36d 100644 --- a/tests/unit/test_exception_header_preservation.py +++ b/tests/unit/test_exception_header_preservation.py @@ -13,6 +13,7 @@ from typing import Final import httpx import pytest +import litellm from litellm.exceptions import ( BadRequestError, ContentPolicyViolationError, @@ -418,3 +419,21 @@ class TestBedrockErrorHeaders: headers = get_response_headers(dict(_response_headers)) assert headers.get("llm_provider-x-amzn-requestid") == "req-proxy-500" + + +@pytest.mark.parametrize( + "error_class", + (RateLimitError, ServiceUnavailableError, litellm.BadGatewayError, litellm.InternalServerError), +) +def test_error_built_from_an_unread_streaming_response_keeps_its_headers(error_class: type[Exception]) -> None: + unread: Final = httpx.Response( + 500, + headers={"x-request-id": "req-unread"}, + stream=httpx.ByteStream(b'{"error":"never read"}'), + ) + + error: Final = error_class(message="upstream failed", llm_provider="openai", model="gpt-5.6", response=unread) + + error_response: Final = getattr(error, "response") + assert error_response.text == "" + assert dict(error_response.headers) == {"x-request-id": "req-unread"} From dbcacc7ab7880d5115376cd6c22d22898c50d08f Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:34:57 -0700 Subject: [PATCH 07/64] fix(health): probe only models a partial wildcard serves (#49060) * fix(health): probe only models a partial wildcard serves A partial wildcard health check now probes the cheapest cost-map chat models the wildcard matches, keeps provider route prefixes such as Bedrock converse/ on the probe, and falls back only within the wildcard. A wildcard with no matching cost-map chat model reports an error naming model_info.health_check_model * fix(health): keep the provider-wide probe for an unmatched wildcard, skip retired models, and ignore Bedrock ARNs A partial wildcard with no served cost-map chat model falls back to the provider's cheapest chat models instead of failing, and fails with a message naming model_info.health_check_model only when the provider has none. The cheapest-model picker skips models past their cost-map deprecation_date, and a Bedrock ARN segment is never read as a route prefix. * fix(health): fall back to the provider's cheapest probe when a wildcard's own models fail A partial wildcard now probes up to three of its own cheapest chat models first and then the provider-wide cheapest ones, all from one cost-map scan, so stale or unpriced rows the wildcard matches cannot turn a healthy deployment red. Retired models are read by their resolved cost-map key, both deprecation checks use the UTC date, Bedrock only treats converse, invoke, and region names as route segments, and equal-price ties break by model id * fix(health): cap wildcard health probes at three and rank retired cost-map rows last A partial wildcard now probes two of its own cheapest cost-map chat models and then one provider-wide model, so a failing deployment makes no more calls than before. Rows past their cost-map deprecation date go last instead of being dropped, since some still serve, and the shared cheapest-model picker returns what it returned before. * fix(health): probe a wildcard's own models first, then the same three models main probes A partial wildcard's health check tries its two cheapest own cost-map models, then the provider's three cheapest models main already probes, so a row main shows healthy stays healthy. The deprecation check uses the server's local date, like the Baseten listing * fix(health): keep main's provider-wide probe order so equal-price ties match main --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../health_check_helpers.py | 80 +++- .../litellm_core_utils/llm_request_utils.py | 58 ++- .../test_health_check_helpers.py | 343 +++++++++++++++++- .../test_llm_request_utils.py | 84 ++++- .../proxy/test_health_check_max_tokens.py | 7 +- 5 files changed, 545 insertions(+), 27 deletions(-) diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 940730eeb9f..e40299a218f 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -3,6 +3,8 @@ Helper functions for health check calls. """ import base64 +import itertools +import re from collections.abc import Awaitable, Callable, Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal @@ -26,6 +28,9 @@ IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = ( ) ANTHROPIC_MESSAGES_HEALTH_CHECK_MAX_TOKENS: Final = 16 +WILDCARD_HEALTH_CHECK_SERVED_MODEL_COUNT: Final = 2 +WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT: Final = 3 +BEDROCK_ANY_MODEL_ROUTES: Final = frozenset({"converse", "invoke"}) EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( {"reachable": MappingProxyType({"type": "text", "instructions": "Any word from the text"})} ) @@ -90,6 +95,62 @@ def _ocr_health_check_document(model: str, custom_llm_provider: str) -> Document return native(model, custom_llm_provider) +def _provider_base_model(custom_llm_provider: str, model: str) -> str | None: + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + try: + model_info: Final = ProviderConfigManager.get_provider_model_info( + model=None, provider=LlmProviders(custom_llm_provider) + ) + return None if model_info is None else model_info.get_base_model(model) + except (ValueError, NotImplementedError): + return None + + +def _is_route_segment(custom_llm_provider: str, segment: str, model_after_segment: str) -> bool: + if custom_llm_provider == "bedrock": + from litellm.llms.bedrock.common_utils import get_all_bedrock_regions + + return segment in BEDROCK_ANY_MODEL_ROUTES or segment in get_all_bedrock_regions() + base_model: Final = _provider_base_model(custom_llm_provider, f"{segment}/{model_after_segment}") + return base_model is not None and not base_model.startswith(f"{segment}/") + + +def wildcard_health_check_models(wildcard: str, custom_llm_provider: str) -> tuple[str, ...]: + from litellm.litellm_core_utils.llm_request_utils import rank_chat_models_from_llm_provider + from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + + provider_prefix: Final = f"{custom_llm_provider}/" + segments: Final = wildcard.removeprefix(provider_prefix).split("/") + route_segments: Final = tuple( + itertools.takewhile( + lambda segment_index: _is_route_segment( + custom_llm_provider, segments[segment_index], "/".join(segments[segment_index + 1 :]) + ), + range(len(segments) - 1), + ) + ) + route_prefix: Final = "".join(f"{segments[segment_index]}/" for segment_index in route_segments) + model_regex: Final = re.compile(PatternMatchRouter().pattern_to_regex("/".join(segments[len(route_segments) :]))) + ranked_model_ids: Final = tuple( + (ranked_model.model.removeprefix(provider_prefix), ranked_model) + for ranked_model in rank_chat_models_from_llm_provider(custom_llm_provider) + ) + served_model_ids: Final = sorted( + (ranked for ranked in ranked_model_ids if model_regex.fullmatch(ranked[0])), + key=lambda ranked: (ranked[1].is_past_deprecation_date, ranked[1].route_rank, ranked[1].cost, ranked[0]), + ) + served_probes: Final = itertools.islice( + dict.fromkeys(f"{provider_prefix}{route_prefix}{model_id}" for model_id, _ in served_model_ids), + WILDCARD_HEALTH_CHECK_SERVED_MODEL_COUNT, + ) + provider_probes: Final = ( + f"{provider_prefix}{model_id}" for model_id, _ in ranked_model_ids[:WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT] + ) + return tuple(dict.fromkeys((*served_probes, *provider_probes))) + + class HealthCheckHelpers: @staticmethod async def ahealth_check_wildcard_models( @@ -99,23 +160,16 @@ class HealthCheckHelpers: litellm_logging_obj: "Logging", ) -> dict: from litellm import acompletion - from litellm.litellm_core_utils.llm_request_utils import ( - pick_cheapest_chat_models_from_llm_provider, - ) - # this is a wildcard model, we need to pick a random model from the provider - cheapest_models = pick_cheapest_chat_models_from_llm_provider(custom_llm_provider=custom_llm_provider, n=3) - if len(cheapest_models) == 0: + probe_models: Final = wildcard_health_check_models(wildcard=model, custom_llm_provider=custom_llm_provider) + if len(probe_models) == 0: raise Exception( - f"Unable to health check wildcard model for provider {custom_llm_provider}. Add a model on your config.yaml or contribute here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json" + f"Unable to health check wildcard model {model} for provider {custom_llm_provider}: the cost map has " + "no chat model for this provider. Set model_info.health_check_model to a model this wildcard serves" ) - if len(cheapest_models) > 1: - fallback_models = cheapest_models[1:] # Pick the last 2 models from the shuffled list - else: - fallback_models = None - model_params["model"] = cheapest_models[0] + model_params["model"] = probe_models[0] model_params["litellm_logging_obj"] = litellm_logging_obj - model_params["fallbacks"] = fallback_models + model_params["fallbacks"] = list(probe_models[1:]) or None model_params["max_tokens"] = model_params.get("max_tokens", 16) # GPT-5 models require max_output_tokens >= 16 await acompletion(**model_params) return {} diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index db881aa16c6..9c2ec0b8984 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -1,6 +1,11 @@ +import datetime from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass from typing import Final, cast +from pydantic import TypeAdapter +from typing_extensions import NotRequired, ReadOnly, TypedDict + import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider @@ -151,19 +156,52 @@ def pick_cheapest_chat_models_from_llm_provider(custom_llm_provider: str, n=1): Returns: list[str]: A list of the n cheapest chat models. """ - if custom_llm_provider not in litellm.models_by_provider: - return [] + return [ranked_model.model for ranked_model in rank_chat_models_from_llm_provider(custom_llm_provider)[:n]] + +@dataclass(frozen=True, slots=True) +class RankedChatModel: + model: str + route_rank: int + cost: float + is_past_deprecation_date: bool + + +def rank_chat_models_from_llm_provider(custom_llm_provider: str) -> tuple[RankedChatModel, ...]: known_models: Final[Iterable[str]] = litellm.models_by_provider.get(custom_llm_provider, ()) - ranked_rows: Final = ( - (route_rank, cost, model) - for model, cost in _chat_row_costs(known_models, custom_llm_provider) + ranked_models: Final = ( + RankedChatModel( + model=model, route_rank=route_rank, cost=cost, is_past_deprecation_date=is_past_deprecation_date + ) + for model, cost, is_past_deprecation_date in _chat_row_costs(known_models, custom_llm_provider) if (route_rank := _route_rank(model, custom_llm_provider)) is not None ) - return [model for _, _, model in sorted(ranked_rows, key=lambda row: row[:2])[:n]] + return tuple(sorted(ranked_models, key=lambda ranked_model: (ranked_model.route_rank, ranked_model.cost))) -def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float]]: +class _CostMapDeprecation(TypedDict): + deprecation_date: NotRequired[ReadOnly[object]] + + +_COST_MAP_DEPRECATION: Final = TypeAdapter(_CostMapDeprecation) + + +def _is_past_deprecation_date(cost_map_key: str) -> bool: + deprecation_date: Final = _COST_MAP_DEPRECATION.validate_python( + litellm.model_cost.get(cost_map_key) or {} # pyright: ignore[reportUnknownMemberType] # untyped cost map JSON + ).get("deprecation_date") + if not isinstance(deprecation_date, str): + return False + try: + return ( + datetime.date.fromisoformat(deprecation_date) + <= datetime.datetime.now(datetime.timezone.utc).astimezone().date() + ) + except ValueError: + return False + + +def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float, bool]]: for model in models: try: model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) @@ -171,7 +209,11 @@ def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator continue if model_info.get("mode") != "chat": continue - yield model, (model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0) + yield ( + model, + (model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0), + _is_past_deprecation_date(model_info["key"]), + ) def _route_rank(model: str, custom_llm_provider: str) -> int | None: diff --git a/tests/unit/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py index 6979d72ff14..950a91beb81 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -8,6 +8,7 @@ from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import responses import respx @@ -19,6 +20,7 @@ from litellm.litellm_core_utils.health_check_helpers import ( HealthCheckHelpers, default_health_check_mode, native_health_check_mode, + wildcard_health_check_models, ) from litellm.main import ahealth_check from litellm.proxy._types import UserAPIKeyAuth @@ -1449,6 +1451,339 @@ async def test_groq_wildcard_health_check_uses_a_concrete_model( assert response == {} +_WILDCARD_CATALOG: Final = MappingProxyType( + { + "databricks": ( + "databricks/databricks-claude-fixture-d", + "databricks/databricks-claude-fixture-c", + "databricks/databricks-claude-fixture-b", + "databricks/databricks-claude-fixture-a", + "databricks/databricks-cheap-chat", + "databricks/databricks-claude-fixture-retired", + ), + "bedrock": ("meta.llama-fixture", "anthropic.claude-fixture"), + "cloudflare": ("cloudflare/@cf/meta/fixture-chat",), + "vertex_ai": ("claude-fixture@default",), + } +) +_WILDCARD_CATALOG_COST: Final = MappingProxyType( + { + "databricks/databricks-claude-fixture-retired": 1e-10, + "databricks/databricks-cheap-chat": 1e-9, + "databricks/databricks-claude-fixture-a": 2e-9, + "databricks/databricks-claude-fixture-b": 3e-9, + "databricks/databricks-claude-fixture-c": 4e-9, + "databricks/databricks-claude-fixture-d": 5e-9, + "meta.llama-fixture": 1e-9, + "anthropic.claude-fixture": 2e-9, + "cloudflare/@cf/meta/fixture-chat": 1e-9, + "claude-fixture@default": 1e-9, + } +) +_DATABRICKS_FIXTURE_API_BASE: Final = "https://fixture.cloud.databricks.com/serving-endpoints" + + +@pytest.fixture +def wildcard_catalog(monkeypatch: pytest.MonkeyPatch) -> None: + providers: Final = {model: provider for provider, models in _WILDCARD_CATALOG.items() for model in models} + litellm.register_model( + { + model: { + "litellm_provider": providers[model], + "mode": "chat", + "max_tokens": 1024, + "input_cost_per_token": cost, + "output_cost_per_token": cost, + **({"deprecation_date": "2000-01-01"} if model.endswith("-retired") else {}), + } + for model, cost in _WILDCARD_CATALOG_COST.items() + } + ) + for provider, models in _WILDCARD_CATALOG.items(): + monkeypatch.setitem(litellm.models_by_provider, provider, list(models)) + + +_DATABRICKS_PROVIDER_PROBES: Final = ( + "databricks/databricks-claude-fixture-retired", + "databricks/databricks-cheap-chat", + "databricks/databricks-claude-fixture-a", +) +_BEDROCK_PROVIDER_PROBES: Final = ("bedrock/meta.llama-fixture", "bedrock/anthropic.claude-fixture") +_DATABRICKS_SERVED_CLAUDE_PROBES: Final = ( + "databricks/databricks-claude-fixture-a", + "databricks/databricks-claude-fixture-b", +) + + +@pytest.mark.parametrize( + ("custom_llm_provider", "wildcard", "expected"), + [ + ("databricks", "databricks-claude-*", (*_DATABRICKS_SERVED_CLAUDE_PROBES, *_DATABRICKS_PROVIDER_PROBES[:2])), + ( + "databricks", + "databricks/databricks-claude-*", + (*_DATABRICKS_SERVED_CLAUDE_PROBES, *_DATABRICKS_PROVIDER_PROBES[:2]), + ), + ("databricks", "*", (*_DATABRICKS_PROVIDER_PROBES[1:], _DATABRICKS_PROVIDER_PROBES[0])), + ("databricks", "databricks-claude-fixture-r*", _DATABRICKS_PROVIDER_PROBES), + ("databricks", "databricks-llama-*", _DATABRICKS_PROVIDER_PROBES), + ("databricks", "*-fixture", _DATABRICKS_PROVIDER_PROBES), + ("databricks", "team-endpoints/*", _DATABRICKS_PROVIDER_PROBES), + ("bedrock", "converse/anthropic.*", ("bedrock/converse/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)), + ( + "bedrock", + "converse/*", + ( + "bedrock/converse/meta.llama-fixture", + "bedrock/converse/anthropic.claude-fixture", + *_BEDROCK_PROVIDER_PROBES, + ), + ), + ("bedrock", "invoke/anthropic.*", ("bedrock/invoke/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)), + ("bedrock", "us-east-1/anthropic.*", ("bedrock/us-east-1/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)), + ("bedrock", "agent/anthropic.*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "nova/*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "nova-2/*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "mantle/*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "openai/*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "chat_completions/*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/*", _BEDROCK_PROVIDER_PROBES), + ( + "bedrock", + "converse/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/*", + _BEDROCK_PROVIDER_PROBES, + ), + ("cloudflare", "@cf/*", ("cloudflare/@cf/meta/fixture-chat",)), + ("vertex_ai", "claude-*", ("vertex_ai/claude-fixture@default",)), + ("vertex_ai", "meta/*", ("vertex_ai/claude-fixture@default",)), + ("fixture-gateway", "team/*", ()), + ], +) +@pytest.mark.usefixtures("wildcard_catalog") +def test_wildcard_health_check_probes_the_cheapest_served_chat_models_before_the_provider_cheapest( + custom_llm_provider: str, wildcard: str, expected: tuple[str, ...] +) -> None: + assert wildcard_health_check_models(wildcard=wildcard, custom_llm_provider=custom_llm_provider) == expected + + +def _fail_first_databricks_claude_probe(request: httpx.Request) -> httpx.Response: + model: Final = json.loads(request.content)["model"] + if model == "databricks-claude-fixture-a": + return httpx.Response(500, json={"error_code": "INTERNAL_ERROR", "message": "fixture outage"}) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-health", + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "2"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("wildcard_catalog") +async def test_partial_wildcard_health_check_falls_back_within_the_wildcard_before_any_other_model( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock( + side_effect=_fail_first_databricks_claude_probe + ) + + response: Final = await ahealth_check( + { + "model": "databricks/databricks-claude-*", + "api_key": "fake-key", + "api_base": _DATABRICKS_FIXTURE_API_BASE, + "messages": [{"role": "user", "content": "What is 1 + 1?"}], + } + ) + + probed_models: Final = [json.loads(call.request.content)["model"] for call in upstream.calls] + assert probed_models[0] == "databricks-claude-fixture-a" + assert probed_models[-1] == "databricks-claude-fixture-b" + assert all(model.startswith("databricks-claude-") for model in probed_models) + assert "databricks-claude-fixture-a" not in probed_models[1:] + assert response == {} + + +def _answer_only_the_retired_databricks_claude_model(request: httpx.Request) -> httpx.Response: + model: Final = json.loads(request.content)["model"] + if model != "databricks-claude-fixture-retired": + return httpx.Response(404, json={"error_code": "ENDPOINT_NOT_FOUND", "message": f"{model} does not exist"}) + return _fail_first_databricks_claude_probe(request) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("wildcard_catalog") +async def test_partial_wildcard_health_check_stays_healthy_while_any_model_the_provider_probes_answers( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock( + side_effect=_answer_only_the_retired_databricks_claude_model + ) + + response: Final = await ahealth_check( + { + "model": "databricks/databricks-claude-*", + "api_key": "fake-key", + "api_base": _DATABRICKS_FIXTURE_API_BASE, + "messages": [{"role": "user", "content": "What is 1 + 1?"}], + } + ) + + assert [json.loads(call.request.content)["model"] for call in upstream.calls] == [ + "databricks-claude-fixture-a", + "databricks-claude-fixture-b", + "databricks-claude-fixture-retired", + ] + assert response == {} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("wildcard_catalog") +async def test_partial_wildcard_health_check_that_fails_on_every_probe_tries_its_own_models_then_the_provider_probes( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").respond( + 401, json={"error_code": "UNAUTHENTICATED", "message": "fixture bad token"} + ) + + response: Final = await ahealth_check( + { + "model": "databricks/databricks-claude-*", + "api_key": "fake-key", + "api_base": _DATABRICKS_FIXTURE_API_BASE, + "messages": [{"role": "user", "content": "What is 1 + 1?"}], + } + ) + + assert [json.loads(call.request.content)["model"] for call in upstream.calls] == [ + "databricks-claude-fixture-a", + "databricks-claude-fixture-b", + "databricks-claude-fixture-retired", + "databricks-cheap-chat", + ] + assert "error" in response + + +_DATABRICKS_EQUAL_PRICE_MODELS: Final = ( + "databricks/databricks-free-d", + "databricks/databricks-free-c", + "databricks/databricks-free-b", + "databricks/databricks-free-a", +) + + +def _register_equal_price_databricks_models(monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, ...]) -> None: + litellm.register_model( + { + model: { + "litellm_provider": "databricks", + "mode": "chat", + "max_tokens": 1024, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + for model in listed_order + } + ) + monkeypatch.setitem(litellm.models_by_provider, "databricks", list(listed_order)) + + +@pytest.mark.parametrize("listed_order", [_DATABRICKS_EQUAL_PRICE_MODELS, _DATABRICKS_EQUAL_PRICE_MODELS[::-1]]) +def test_partial_wildcard_health_check_orders_its_own_equal_price_models_by_id_whatever_the_listing_order( + monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, ...] +) -> None: + _register_equal_price_databricks_models(monkeypatch, listed_order) + + assert wildcard_health_check_models(wildcard="databricks/databricks-free-*", custom_llm_provider="databricks")[ + :2 + ] == ("databricks/databricks-free-a", "databricks/databricks-free-b") + + +def _answer_only_the_first_listed_databricks_free_model(request: httpx.Request) -> httpx.Response: + model: Final = json.loads(request.content)["model"] + if model != "databricks-free-d": + return httpx.Response(404, json={"error_code": "ENDPOINT_NOT_FOUND", "message": f"{model} does not exist"}) + return _fail_first_databricks_claude_probe(request) + + +@pytest.mark.asyncio +async def test_wildcard_health_check_falls_back_to_equal_price_provider_models_in_the_provider_listing_order( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + _register_equal_price_databricks_models(monkeypatch, _DATABRICKS_EQUAL_PRICE_MODELS) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock( + side_effect=_answer_only_the_first_listed_databricks_free_model + ) + + response: Final = await ahealth_check( + { + "model": "databricks/databricks-claude-*", + "api_key": "fake-key", + "api_base": _DATABRICKS_FIXTURE_API_BASE, + "messages": [{"role": "user", "content": "What is 1 + 1?"}], + } + ) + + assert [json.loads(call.request.content)["model"] for call in upstream.calls] == ["databricks-free-d"] + assert response == {} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("wildcard_catalog") +async def test_partial_wildcard_health_check_without_a_served_cost_map_model_probes_the_provider_cheapest_chat_model( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock( + side_effect=_fail_first_databricks_claude_probe + ) + + response: Final = await ahealth_check( + { + "model": "databricks/databricks-llama-*", + "api_key": "fake-key", + "api_base": _DATABRICKS_FIXTURE_API_BASE, + "messages": [{"role": "user", "content": "What is 1 + 1?"}], + } + ) + + assert json.loads(upstream.calls[0].request.content)["model"] == "databricks-claude-fixture-retired" + assert response == {} + + +@pytest.mark.asyncio +async def test_partial_wildcard_health_check_on_a_provider_without_cost_map_chat_models_names_health_check_model( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setitem(litellm.models_by_provider, "databricks", []) + + response: Final = await ahealth_check( + { + "model": "databricks/databricks-llama-*", + "api_key": "fake-key", + "api_base": _DATABRICKS_FIXTURE_API_BASE, + "messages": [{"role": "user", "content": "What is 1 + 1?"}], + } + ) + + assert "model_info.health_check_model" in response["error"], response["error"] + assert respx_mock.calls.call_count == 0 + + @pytest.mark.asyncio async def test_cohere_rerank_health_check_returns_a_successful_response( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter @@ -1632,7 +1967,12 @@ async def test_wildcard_health_check_with_pinned_provider_calls_a_model_without_ monkeypatch.setitem( litellm.model_cost, "fake-together-chat", - {"input_cost_per_token": 1e-9, "output_cost_per_token": 1e-9, "mode": "chat", "litellm_provider": "together_ai"}, + { + "input_cost_per_token": 1e-9, + "output_cost_per_token": 1e-9, + "mode": "chat", + "litellm_provider": "together_ai", + }, ) monkeypatch.setitem(litellm.models_by_provider, "together_ai", ("fake-together-chat",)) litellm.in_memory_llm_clients_cache.flush_cache() @@ -1647,4 +1987,3 @@ async def test_wildcard_health_check_with_pinned_provider_calls_a_model_without_ assert "error" not in result, result assert json.loads(together_route.calls.last.request.content)["model"] == "fake-together-chat" assert not other_hosts.called - diff --git a/tests/unit/litellm_core_utils/test_llm_request_utils.py b/tests/unit/litellm_core_utils/test_llm_request_utils.py index 0a2bd5cec1b..e7350fae395 100644 --- a/tests/unit/litellm_core_utils/test_llm_request_utils.py +++ b/tests/unit/litellm_core_utils/test_llm_request_utils.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_request_utils import ( ensure_extra_body_is_safe, flatten_form_field_values, pick_cheapest_chat_models_from_llm_provider, + rank_chat_models_from_llm_provider, serialize_multipart_form_fields, ) @@ -149,6 +150,24 @@ def deepseek_rows_whose_cheapest_ids_route_elsewhere(monkeypatch: pytest.MonkeyP "deepseek/fake-cheapest-embedding": {"input_cost_per_token": 1e-13, "mode": "embedding"}, "openai/fake-free-chat": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0, "mode": "chat"}, "deepseek/fake-cheap-chat": {"input_cost_per_token": 1e-12, "output_cost_per_token": 1e-12, "mode": "chat"}, + "deepseek/fake-retired-chat": { + "input_cost_per_token": 1e-14, + "output_cost_per_token": 1e-14, + "mode": "chat", + "deprecation_date": "2000-01-01", + }, + "deepseek/fake-retiring-chat": { + "input_cost_per_token": 2e-12, + "output_cost_per_token": 2e-12, + "mode": "chat", + "deprecation_date": "2999-01-01", + }, + "deepseek/fake-misdated-chat": { + "input_cost_per_token": 2.5e-12, + "output_cost_per_token": 2.5e-12, + "mode": "chat", + "deprecation_date": "next spring", + }, "fake-free-chat-without-provider": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0, "mode": "chat"}, } for name, row in rows.items(): @@ -156,12 +175,73 @@ def deepseek_rows_whose_cheapest_ids_route_elsewhere(monkeypatch: pytest.MonkeyP monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, tuple(rows)) -def test_pick_cheapest_chat_models_prefers_provider_routed_ids_and_skips_ids_routed_elsewhere( +def test_pick_cheapest_chat_models_prefers_provider_routed_ids_skips_ids_routed_elsewhere_and_keeps_retired_ones( deepseek_rows_whose_cheapest_ids_route_elsewhere: None, ) -> None: assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=10) == [ + "deepseek/fake-retired-chat", "deepseek/fake-cheap-chat", + "deepseek/fake-retiring-chat", + "deepseek/fake-misdated-chat", "deepseek/fake-pricey-chat", "fake-free-chat-without-provider", ] - assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=1) == ["deepseek/fake-cheap-chat"] + assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=1) == ["deepseek/fake-retired-chat"] + + +def test_rank_chat_models_flags_only_the_rows_whose_cost_map_deprecation_date_has_passed( + deepseek_rows_whose_cheapest_ids_route_elsewhere: None, +) -> None: + assert { + ranked_model.model: ranked_model.is_past_deprecation_date + for ranked_model in rank_chat_models_from_llm_provider(PICKER_PROVIDER) + } == { + "deepseek/fake-retired-chat": True, + "deepseek/fake-cheap-chat": False, + "deepseek/fake-retiring-chat": False, + "deepseek/fake-misdated-chat": False, + "deepseek/fake-pricey-chat": False, + "fake-free-chat-without-provider": False, + } + + +@pytest.mark.parametrize( + "listed_order", + [("deepseek/fake-tie-b", "deepseek/fake-tie-a"), ("deepseek/fake-tie-a", "deepseek/fake-tie-b")], +) +def test_pick_cheapest_chat_models_keeps_the_provider_listing_order_for_equal_price_models( + monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, str] +) -> None: + for name in listed_order: + monkeypatch.setitem( + litellm.model_cost, + name, + { + "input_cost_per_token": 1e-12, + "output_cost_per_token": 1e-12, + "mode": "chat", + "litellm_provider": PICKER_PROVIDER, + }, + ) + monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, listed_order) + + assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=2) == list(listed_order) + + +def test_rank_chat_models_flags_a_retired_model_the_provider_lists_without_its_cost_map_prefix( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rows: Final = { + "vertex_ai/fake-retired-vertex-chat": {"input_cost_per_token": 1e-14, "deprecation_date": "2000-01-01"}, + "vertex_ai/fake-live-vertex-chat": {"input_cost_per_token": 1e-12, "deprecation_date": "2999-01-01"}, + } + for name, row in rows.items(): + monkeypatch.setitem(litellm.model_cost, name, {**row, "mode": "chat", "litellm_provider": "vertex_ai"}) + monkeypatch.setitem( + litellm.models_by_provider, "vertex_ai", tuple(name.removeprefix("vertex_ai/") for name in rows) + ) + + assert [ + (ranked_model.model, ranked_model.is_past_deprecation_date) + for ranked_model in rank_chat_models_from_llm_provider("vertex_ai") + ] == [("fake-retired-vertex-chat", True), ("fake-live-vertex-chat", False)] diff --git a/tests/unit/proxy/test_health_check_max_tokens.py b/tests/unit/proxy/test_health_check_max_tokens.py index e48bf4edc09..df0429316ed 100644 --- a/tests/unit/proxy/test_health_check_max_tokens.py +++ b/tests/unit/proxy/test_health_check_max_tokens.py @@ -11,6 +11,7 @@ import respx import litellm from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers +from litellm.litellm_core_utils.llm_request_utils import RankedChatModel from litellm.proxy import health_check as hc_module from litellm.proxy.health_check import ( _is_strategy_router_deployment, @@ -85,8 +86,10 @@ async def test_ahealth_check_wildcard_models_respects_max_tokens(): """ with ( patch( - "litellm.litellm_core_utils.llm_request_utils.pick_cheapest_chat_models_from_llm_provider", - return_value=["gpt-4o-mini"], + "litellm.litellm_core_utils.llm_request_utils.rank_chat_models_from_llm_provider", + return_value=( + RankedChatModel(model="gpt-4o-mini", route_rank=0, cost=0.0, is_past_deprecation_date=False), + ), ), patch("litellm.acompletion", new_callable=AsyncMock), ): From 4a5c5665faec4efd3bd2dc31436c9f6f8b25cf2e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:49:55 -0700 Subject: [PATCH 08/64] fix(azure_ai)!: keep OpenAI prompt cache breakpoints on the Foundry chat completions path (#45299) * fix(azure_ai): keep OpenAI prompt cache breakpoints on the Foundry chat completions path * fix(azure_ai): keep only OpenAI content part types when preserving prompt cache breakpoints * fix(azure_ai): share the breakpoint check with the cache-mark step and keep only text parts A deployment's model_info opt-in is registered under its provider-prefixed cost-map key without a litellm_provider, so that prefix now names the provider the flag speaks for. * refactor(azure_ai): read text parts through a walrus and mark new test values Final * fix(azure_ai): read a router-registered breakpoint opt-in only for azure_ai and keep a message mark off dropped file parts * fix(azure_ai): read the breakpoint opt-in from each deployment's own model_info and base_model A router deployment without its own supports_prompt_cache_breakpoint flag now reads the cost map as loaded, so a sibling's opt-in on the same model string no longer turns its requests into 400s. The shared hook no longer reads router registered keys for azure_ai. * fix(azure_ai): send breakpoints in Foundry's shape under drop_params and fall back to the model name * test(integration): expect the injected breakpoint on the Foundry GPT-6 chat wire --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../anthropic_cache_control_hook.py | 33 +- litellm/llms/azure_ai/chat/transformation.py | 111 +++++- ...ect_prompt_cache_breakpoint_owned_proxy.py | 17 +- .../chat/test_azure_ai_transformation.py | 356 ++++++++++++++++++ 4 files changed, 499 insertions(+), 18 deletions(-) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index fc1f1830129..45c82589d2d 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -23,6 +23,7 @@ from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.integrations.prompt_management_base import PromptManagementClient +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.prompt_templates.common_utils import ( cache_control_as_prompt_cache_breakpoint, with_prompt_cache_breakpoint, @@ -175,7 +176,10 @@ def _hosted_entry_flag(entry: Mapping[str, object], resolve_provider: Callable[[ def _hosted_openai_dialect_flag( - model: str, custom_llm_provider: str | None, resolve_provider: Callable[[str], str | None] + model: str, + custom_llm_provider: str | None, + resolve_provider: Callable[[str], str | None], + model_cost: Mapping[str, Mapping[str, object]] | None = None, ) -> bool | None: """ Explicit opt-in for an OpenAI-shaped deployment served by another provider. @@ -194,14 +198,19 @@ def _hosted_openai_dialect_flag( """ import litellm - exact_entry: Final = litellm.model_cost.get(model) - exact_provider: Final = exact_entry.get("litellm_provider") if isinstance(exact_entry, dict) else None + cost_map: Final = ( + cast(Mapping[str, Mapping[str, object]], litellm.model_cost) # cast-ok: the live map holds catalog entries + if model_cost is None + else model_cost + ) + exact_entry: Final = cost_map.get(model) + exact_provider: Final = exact_entry.get("litellm_provider") if isinstance(exact_entry, Mapping) else None if custom_llm_provider is None and (exact_provider == "openai" or ("/" not in model and exact_provider is None)): return None provider: Final = custom_llm_provider or resolve_provider(model) if provider is None or provider == "openai": return None - if isinstance(exact_entry, dict) and exact_provider == provider: + if isinstance(exact_entry, Mapping) and exact_provider == provider: return _hosted_entry_flag(exact_entry, lambda: provider) from litellm.utils import get_potential_model_names @@ -214,8 +223,8 @@ def _hosted_openai_dialect_flag( names["stripped_model_name"], names["provider_prefixed_model_name"], ) - entries: Final = (litellm.model_cost.get(candidate) for candidate in candidates) - flags: Final = (_hosted_entry_flag(entry, lambda: provider) for entry in entries if isinstance(entry, dict)) + entries: Final = (cost_map.get(candidate) for candidate in candidates) + flags: Final = (_hosted_entry_flag(entry, lambda: provider) for entry in entries if isinstance(entry, Mapping)) return next((flag for flag in flags if flag is not None), None) @@ -537,6 +546,18 @@ class AnthropicCacheControlHook(CustomPromptManagement): return False return prompt_cache_options is not None or targets_openai_api(api_base) + @staticmethod + def catalog_targets_openai_prompt_cache_breakpoint(model: str, custom_llm_provider: str) -> bool: + return ( + _hosted_openai_dialect_flag( + model, + custom_llm_provider, + AnthropicCacheControlHook._resolve_provider, + GetModelCostMap.loaded_model_cost_map(), + ) + is True + ) + @staticmethod def _resolve_provider(model: str) -> str | None: return _served_provider_for_model(model) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index d382d653a53..5381447f28f 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -6,15 +6,19 @@ from typing import TYPE_CHECKING, Final, cast import httpx from httpx import Response +from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger +from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.litellm_core_utils.prompt_templates.common_utils import ( content_parts, convert_content_list_to_str, drop_tool_reference_parts_from_tool_messages, filter_value_from_dict, is_pdf_file_part, + prompt_cache_breakpoint_for_wire, + with_prompt_cache_breakpoint, ) from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.azure_ai.common_utils import ( @@ -27,7 +31,7 @@ from litellm.llms.openai.common_utils import drop_params_from_unprocessable_enti from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator, OpenAIConfig from litellm.llms.xai.chat.transformation import XAIChatConfig from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, PromptCacheBreakpoint from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ModelResponse, ModelResponseStream, ProviderField from litellm.utils import add_path_to_api_base, supports_tool_choice @@ -62,6 +66,72 @@ def _has_list_only_content_part(message: Mapping[str, object]) -> bool: return any(_is_list_only_content_part(part) for part in content_parts(message)) +_CONTENT_PARTS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) + + +def _text_parts(message: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + content: Final = message.get("content") + if not isinstance(content, list): + return () + try: + parts: Final = _CONTENT_PARTS_ADAPTER.validate_python(content) + except ValidationError: + return () + return tuple(part for part in parts if part.get("type") == "text") + + +def _with_only_text_parts(message: AllMessageValues) -> AllMessageValues: + if _has_list_only_content_part(message) or not (text_parts := _text_parts(message)): + return message + text_only: Final = {**message, "content": list(text_parts)} + return cast(AllMessageValues, text_only) # cast-ok: the same message, text parts only + + +_FIELDS: Final = TypeAdapter(Mapping[str, object]) + + +def _foundry_breakpoint(marker: object) -> PromptCacheBreakpoint | None: + if prompt_cache_breakpoint_for_wire(marker, drop_params=True) is None: + return None + return PromptCacheBreakpoint(mode="explicit") + + +def _part_with_foundry_breakpoint(part: object) -> object: + if not isinstance(part, Mapping): + return part + fields: Final = _FIELDS.validate_python(part) + if "prompt_cache_breakpoint" not in fields: + return fields + unmarked: Final = {key: value for key, value in fields.items() if key != "prompt_cache_breakpoint"} + return with_prompt_cache_breakpoint(unmarked, _foundry_breakpoint(fields["prompt_cache_breakpoint"])) + + +def _with_breakpoints_for_foundry(message: AllMessageValues, drop_params: bool) -> AllMessageValues: + renamed: Final = AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints(_with_only_text_parts(message)) + parts: Final = content_parts(renamed) if drop_params else () + if not parts: + return renamed + foundry_message: Final = {**renamed, "content": [_part_with_foundry_breakpoint(part) for part in parts]} + return cast(AllMessageValues, foundry_message) # cast-ok: the same message, its breakpoints in Foundry's shape + + +def _deployment_reads_breakpoints(model: str, litellm_params: Mapping[str, object]) -> bool: + base_model: Final = litellm_params.get("base_model") + named_models: Final = (base_model, model) if isinstance(base_model, str) and base_model else (model,) + model_info: Final = litellm_params.get("model_info") + if not isinstance(model_info, Mapping): + return any( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(name, "azure_ai") for name in named_models + ) + deployment_flag: Final = _FIELDS.validate_python(model_info).get("supports_prompt_cache_breakpoint") + if isinstance(deployment_flag, bool): + return deployment_flag + return any( + AnthropicCacheControlHook.catalog_targets_openai_prompt_cache_breakpoint(name, "azure_ai") + for name in named_models + ) + + class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator): def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract usage: Final = chunk.get("usage") @@ -251,6 +321,8 @@ class AzureAIStudioConfig(OpenAIConfig): self, messages: list[AllMessageValues], model: str, + reads_breakpoints: bool | None = None, + drop_params: bool = False, ) -> list: """ - Azure AI Studio doesn't support content as a list. This handles: @@ -260,13 +332,31 @@ class AzureAIStudioConfig(OpenAIConfig): Azure AI Foundry backends set additionalProperties=false and reject these with "Extra inputs are not permitted", which breaks multi-turn Anthropic-format clients that echo thinking blocks back as history. - 2. Transforms list content to a string. + 2. Transforms list content to a string, except on deployments that read OpenAI explicit + prompt cache breakpoints (the deployment's own model_info opt-in, else the cost-map entry + of its base_model or name, as loaded and without a sibling deployment's opt-in): only a + content part can carry the breakpoint (a message-level one is silently ignored), so those + keep their text parts, with each Anthropic `cache_control` mark renamed to a breakpoint the + way the openai provider renames it. Every other part is dropped, exactly as the string + conversion drops it, before the marks are renamed, so a message-level mark lands on the + last text part. Under drop_params each breakpoint goes out as `{"mode": "explicit"}`, the only + shape Foundry takes, and a malformed one is dropped. 3. If message contains an image, audio, or PDF file part, send as is (user-intended) Operates on a deep copy so the caller's messages keep their thinking blocks and provider metadata, which a fallback to another provider still needs. """ - stripped_messages: Final = copy.deepcopy(drop_tool_reference_parts_from_tool_messages(messages)) + deployment_reads_breakpoints: Final = ( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(model, "azure_ai") + if reads_breakpoints is None + else reads_breakpoints + ) + copied_messages: Final = copy.deepcopy(drop_tool_reference_parts_from_tool_messages(messages)) + stripped_messages: Final = ( + [_with_breakpoints_for_foundry(message, drop_params) for message in copied_messages] + if deployment_reads_breakpoints + else copied_messages + ) for message in stripped_messages: message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped on our copy for field in NON_OPENAI_SPEC_MESSAGE_FIELDS: @@ -275,6 +365,10 @@ class AzureAIStudioConfig(OpenAIConfig): if _has_list_only_content_part(message): continue + if deployment_reads_breakpoints and (text_parts := _text_parts(message)): + message_dict["content"] = list(text_parts) + continue + texts = convert_content_list_to_str(message=message) if texts: message["content"] = texts @@ -329,7 +423,16 @@ class AzureAIStudioConfig(OpenAIConfig): if extra_body and isinstance(extra_body, dict): optional_params.update(extra_body) optional_params.pop("max_retries", None) - return super().transform_request(model, messages, optional_params, litellm_params, headers) + return { + "model": model, + "messages": self._transform_messages( + messages, + model, + _deployment_reads_breakpoints(model, _FIELDS.validate_python(litellm_params)), + bool(litellm_params.get("drop_params") or litellm.drop_params), + ), + **optional_params, + } def transform_response( self, diff --git a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py index 3df73d3b767..d470605b1e5 100644 --- a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py +++ b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py @@ -153,15 +153,16 @@ def _foundry_deployment(rig: _Rig, scenario: Scenario, *, bare: bool) -> str: return name -def _assert_foundry_chat_wire(received: Request, prompt: str, *, options: JsonValue | None) -> None: +def _assert_foundry_chat_wire(received: Request, user_content: JsonValue, *, options: JsonValue | None) -> None: body: Final = body_of(received) + marked_system: Final = [{"type": "text", "text": SYSTEM, "prompt_cache_breakpoint": EXPLICIT}] assert urlsplit(received.target).netloc == _FOUNDRY_HOST, received.target assert urlsplit(received.target).path == _FOUNDRY_CHAT_PATH, received.target - assert "prompt_cache_breakpoint" not in received.body.decode(), received.body assert body["model"] == _FOUNDRY_MODEL, received.body - assert body["messages"] == [{"role": "system", "content": SYSTEM}, {"role": "user", "content": prompt}], ( - received.body - ) + assert body["messages"] == [ + {"role": "system", "content": marked_system}, + {"role": "user", "content": user_content}, + ], received.body assert body.get("prompt_cache_options") == options, received.body @@ -185,7 +186,7 @@ def test_b9_b12_foundry_gpt6_responses_reach_the_openai_dialect_on_the_foundry_h @pytest.mark.timeout(_CELL_TIMEOUT) -def test_b10_foundry_gpt6_messages_bridge_to_chat_without_a_breakpoint(rig: _Rig) -> None: +def test_b10_foundry_gpt6_messages_bridge_to_chat_keeps_the_injected_breakpoint(rig: _Rig) -> None: marker: Final = fresh_marker() with rig.gateway.scenario() as scenario: name: Final = _foundry_deployment(rig, scenario, bare=False) @@ -193,12 +194,12 @@ def test_b10_foundry_gpt6_messages_bridge_to_chat_without_a_breakpoint(rig: _Rig rig.gateway, rig.wire, "messages", request_body("messages", name, prompt_text(marker)) ) assert_answered(outcome, marker) - _assert_foundry_chat_wire(received, prompt_text(marker), options=IMPLICIT) + _assert_foundry_chat_wire(received, [{"type": "text", "text": prompt_text(marker)}], options=IMPLICIT) success_row(name, marker) @pytest.mark.timeout(_CELL_TIMEOUT) -def test_b11_foundry_gpt6_chat_stays_on_the_plain_chat_wire(rig: _Rig) -> None: +def test_b11_foundry_gpt6_chat_stays_on_the_chat_wire_with_the_injected_breakpoint(rig: _Rig) -> None: marker: Final = fresh_marker() with rig.gateway.scenario() as scenario: name: Final = _foundry_deployment(rig, scenario, bare=False) diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index 6269fd367fe..145d45317c3 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -1,6 +1,7 @@ import json import traceback from collections.abc import Callable +from types import MappingProxyType from typing import Final from unittest.mock import MagicMock, patch @@ -9,6 +10,7 @@ import respx from httpx import Response import litellm +from litellm import Router from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.llms.azure_ai.azure_model_router.transformation import ( @@ -880,6 +882,360 @@ def test_azure_ai_stripping_does_not_mutate_caller_messages(): } +def test_azure_ai_keeps_prompt_cache_breakpoint_for_a_model_that_supports_it(_local_model_cost_map): + """ + Foundry GPT-6 deployments take OpenAI explicit prompt caching, and the breakpoint only + counts on a content part, so a system message carrying one must reach the request body as + the content-part list it arrived in while the Anthropic-only fields are still stripped. + """ + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "stable instructions", + "prompt_cache_breakpoint": {"mode": "explicit"}, + "cache_control": {"type": "ephemeral"}, + } + ], + }, + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + ], + optional_params={"prompt_cache_options": {"mode": "explicit"}}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][0]["content"] == [ + {"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert request["messages"][1]["content"] == [{"type": "text", "text": "hi"}] + assert request["prompt_cache_options"] == {"mode": "explicit"} + + +def test_azure_ai_still_flattens_list_content_for_a_model_without_the_breakpoint_flag(_local_model_cost_map): + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-4o", + messages=[ + { + "role": "system", + "content": [ + {"type": "text", "text": "stable ", "prompt_cache_breakpoint": {"mode": "explicit"}}, + {"type": "text", "text": "instructions"}, + ], + }, + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][0]["content"] == "stable instructions" + assert request["messages"][1]["content"] == "hi" + assert not _find_key_anywhere(request["messages"], "prompt_cache_breakpoint") + + +def test_azure_ai_drops_thinking_parts_from_history_while_keeping_the_breakpoint(_local_model_cost_map): + """ + A multi-turn client can echo an assistant turn back with thinking parts inside its content + list. Foundry rejects those parts, and the string conversion always dropped them, so a model + that keeps content parts must still drop them while the text parts and the breakpoint stay. + """ + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[ + { + "role": "system", + "content": [ + {"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}} + ], + }, + {"role": "user", "content": "read the file"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "The user wants me to read a file.", "signature": "sig"}, + {"type": "redacted_thinking", "data": "opaque"}, + {"type": "text", "text": "Reading it now."}, + ], + }, + {"role": "user", "content": [{"type": "text", "text": "thanks"}]}, + ], + optional_params={"prompt_cache_options": {"mode": "explicit"}}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][0]["content"] == [ + {"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert request["messages"][2]["content"] == [{"type": "text", "text": "Reading it now."}] + assert request["messages"][3]["content"] == [{"type": "text", "text": "thanks"}] + assert not _find_key_anywhere(request["messages"], "thinking") + assert not _find_key_anywhere(request["messages"], "data") + + + +def test_azure_ai_renames_a_cache_control_mark_to_a_breakpoint_for_a_model_that_reads_them(_local_model_cost_map): + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[ + { + "role": "system", + "content": [{"type": "text", "text": "stable instructions", "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": "hi", "cache_control": {"type": "ephemeral"}}, + ], + optional_params={"prompt_cache_options": {"mode": "explicit"}}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][0]["content"] == [ + {"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert request["messages"][1]["content"] == [ + {"type": "text", "text": "hi", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert not _find_key_anywhere(request["messages"], "cache_control") + + +def test_azure_ai_drops_a_file_part_next_to_the_breakpoint_like_the_string_conversion(_local_model_cost_map): + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}}, + {"type": "file", "file": {"file_id": "file-held-by-another-provider"}}, + ], + } + ] + + kept: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", messages=messages, optional_params={}, litellm_params={}, headers={} + ) + flattened: Final = AzureAIStudioConfig().transform_request( + model="gpt-4o", messages=messages, optional_params={}, litellm_params={}, headers={} + ) + + assert kept["messages"][0]["content"] == [ + {"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert flattened["messages"][0]["content"] == "summarize" + + +def test_azure_ai_sends_a_pdf_part_next_to_the_breakpoint_as_is(_local_model_cost_map): + content: Final = [ + {"type": "file", "file": {"filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}}, + {"type": "text", "text": "What is the codeword?", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ] + + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[{"role": "user", "content": content}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"] == [{"role": "user", "content": content}] + + +def test_azure_ai_puts_a_message_level_mark_on_the_last_text_part_when_a_file_part_follows(_local_model_cost_map): + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[ + { + "role": "user", + "cache_control": {"type": "ephemeral"}, + "content": [ + {"type": "text", "text": "summarize"}, + {"type": "file", "file": {"file_id": "file-held-by-another-provider"}}, + ], + } + ], + optional_params={"prompt_cache_options": {"mode": "explicit"}}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][0]["content"] == [ + {"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + + +@pytest.mark.parametrize( + ("registered", "keeps_parts"), + [ + ({"azure_ai/qa-foundry-alias": {"litellm_provider": "azure_ai", "supports_prompt_cache_breakpoint": True}}, True), + ({"qa-foundry-alias": {"litellm_provider": "openai", "supports_prompt_cache_breakpoint": True}}, False), + ], + ids=["azure_ai-entry-opts-in", "openai-entry-does-not-speak-for-azure_ai"], +) +def test_azure_ai_keeps_parts_only_when_the_azure_ai_entry_reads_breakpoints( + _local_model_cost_map, registered: dict[str, dict[str, object]], keeps_parts: bool +): + litellm.register_model(model_cost=registered, persist_across_reloads=False, add_to_provider_model_sets=False) + breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}} + + request: Final = AzureAIStudioConfig().transform_request( + model="qa-foundry-alias", + messages=[{"role": "system", "content": [breakpoint_part]}], + optional_params={}, + litellm_params={"custom_llm_provider": "azure_ai", "drop_params": True}, + headers={}, + ) + + assert request["messages"][0]["content"] == ([breakpoint_part] if keeps_parts else "stable") + + +@pytest.mark.parametrize( + "registration_order", [("opted-in", "opted-out", "unflagged"), ("unflagged", "opted-out", "opted-in")] +) +def test_azure_ai_deployment_reads_only_its_own_breakpoint_flag_not_a_siblings_sharing_its_model_string( + _local_model_cost_map, respx_mock: respx.MockRouter, registration_order: tuple[str, ...] +) -> None: + breakpoint_flags: Final = MappingProxyType( + { + "opted-in": {"supports_prompt_cache_breakpoint": True}, + "opted-out": {"supports_prompt_cache_breakpoint": False}, + } + ) + def deployment(host: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": host, + "litellm_params": { + "model": "azure_ai/qa-shared-deployment", + "api_base": f"https://{host}.services.ai.azure.com/models", + "api_key": "azure-ai-test-key", + }, + "model_info": model_info, + } + + routes: Final = { + host: respx_mock.post(host=f"{host}.services.ai.azure.com", path="/models/chat/completions").mock( + return_value=Response( + 200, + json={ + "id": "chatcmpl-azure-ai", + "object": "chat.completion", + "created": 1700000000, + "model": "qa-shared-deployment", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, + }, + ) + ) + for host in registration_order + } + router: Final = Router( + model_list=[ + deployment(host, {"id": f"{host}-id", **breakpoint_flags.get(host, {})}) for host in registration_order + ] + ) + marked_system: Final = { + "role": "system", + "content": [{"type": "text", "text": "stable", "cache_control": {"type": "ephemeral"}}], + } + + for host in registration_order: + router.completion(model=host, messages=[marked_system, {"role": "user", "content": "hi"}]) + + def sent_system_content(host: str) -> object: + return json.loads(routes[host].calls.last.request.content)["messages"][0]["content"] + + assert sent_system_content("opted-in") == [ + {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert sent_system_content("opted-out") == "stable" + assert sent_system_content("unflagged") == "stable" + + +@pytest.mark.parametrize( + "litellm_params", + [ + {"base_model": "azure_ai/gpt-6-astra"}, + {"base_model": "azure_ai/gpt-6-astra", "model_info": {"id": "my-astra-deployment-id"}}, + ], +) +def test_azure_ai_keeps_parts_for_a_custom_deployment_name_whose_base_model_reads_breakpoints( + _local_model_cost_map, litellm_params: dict[str, object] +): + breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}} + + request: Final = AzureAIStudioConfig().transform_request( + model="my-astra-deployment", + messages=[{"role": "system", "content": [breakpoint_part]}], + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + + assert request["messages"][0]["content"] == [breakpoint_part] + + +@pytest.mark.parametrize( + "litellm_params", + [ + {"base_model": "azure/cost-tracking-row"}, + {"base_model": "azure/cost-tracking-row", "model_info": {"id": "astra-deployment-id"}}, + ], +) +def test_azure_ai_falls_back_to_the_model_name_when_the_base_model_row_does_not_read_breakpoints( + _local_model_cost_map, litellm_params: dict[str, object] +): + breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}} + + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[{"role": "system", "content": [breakpoint_part]}], + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + + assert request["messages"][0]["content"] == [breakpoint_part] + + +@pytest.mark.parametrize( + "sent_marker, litellm_params, wire_part", + [ + ( + {"mode": "explicit", "ttl": "30m"}, + {"drop_params": True}, + {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ), + ({"type": "ephemeral"}, {"drop_params": True}, {"type": "text", "text": "stable"}), + ( + {"mode": "explicit", "ttl": "30m"}, + {}, + {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}}, + ), + ], +) +def test_azure_ai_reshapes_breakpoints_for_foundry_only_under_drop_params( + _local_model_cost_map, + sent_marker: dict[str, object], + litellm_params: dict[str, object], + wire_part: dict[str, object], +): + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[ + {"role": "system", "content": [{"type": "text", "text": "stable", "prompt_cache_breakpoint": sent_marker}]} + ], + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + + assert request["messages"][0]["content"] == [wire_part] + + def test_azure_ai_reference_only_tool_result_is_sent_as_empty_text(): messages: Final = [ {"role": "user", "content": "load the WebFetch tool"}, From 953750ed41c69ea909fa4064a42317b04fa8a39b Mon Sep 17 00:00:00 2001 From: Devon Krisman <844724+dkrisman@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:06:57 -0400 Subject: [PATCH 09/64] fix(websearch-interception)!: stream native server-tool blocks to clients (#40967) * fix(websearch-interception): stream native server-tool blocks to clients Streaming native clients (Anthropic SDK, Claude Desktop, Claude Code) never received the server_tool_use and web_search_tool_result blocks the interception layer builds for them, so citation panels stayed empty and result counts read zero. Two gaps in the same flow: FakeAnthropicMessagesStreamIterator._create_content_block_chunks only knew text, thinking, redacted_thinking and tool_use; any other block type emitted a bare content_block_stop with no start and no payload, silently dropping the block from the fake stream. server_tool_use now streams as content_block_start plus input_json_delta and any other server-generated block rides whole in content_block_start, matching Anthropic's native streaming of these types. try_short_circuit_search only recognized a native client by finding a web_search_* tool in the tools list, but the pre-request and deployment hooks convert that tool to the LiteLLM standard shape before the short-circuit runs, recording the conversion in WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY. The flag now counts as native too, so the synthetic response carries the native blocks on the converted path instead of only when hook ordering happens to leave the tool untouched. Fixes #35333 * fix(anthropic): accept server_tool_use and web_search_tool_result blocks when costing a /v1/messages response The short-circuit now returns native web search blocks, and the cost path validated the served body against a content union that only allowed text, tool_use and thinking blocks, so non-streaming requests answered 500 * fix(anthropic): keep every server_tool_use field when re-streaming a /v1/messages reply * fix(anthropic): keep tool_use fields, redacted thinking data and text citations when re-streaming a /v1/messages reply * test(websearch-interception): build the cost test's logging object with a datetime start time * fix(websearch-interception): send search snippets once in the direct search answer * fix(guardrails): mask web search queries, titles and urls on /v1/messages and deliver stream rewrites past server_tool_use blocks Output guardrails now scan the native web search blocks the short-circuit emits, ahead of the reply text so guardrails that read the last text still scan the reply. A guardrail reply of a different length falls back to the reply-only write-back. Ended-stream tool call rewrites land on server_tool_use blocks too, the ended-stream scan key covers the result strings, and legacy-hook pipeline steps pass through the result strings their response view lacks. The wire test reads the streamed server_tool_use input from its input_json_delta. * fix(guardrails): deliver a legacy post-call hook's tool-call rewrites on streams instead of discarding the stream rewrite A legacy hook that masked a tool call's arguments on a buffered stream made the whole rewrite undeliverable, text included, so the original stream went out. With web search blocks in the stream that now hit any hook masking the search query. The adapter hands same-count tool-call rewrites to the translation's write-back, as a unified guardrail's already are, and still discards a rewrite that adds or drops a tool call. * fix(guardrails): show a legacy post-call hook the web search result blocks of a /v1/messages stream The legacy hook's view of a buffered Anthropic stream is rebuilt from the assembled ModelResponse, which keeps web_search_tool_result blocks only in provider-specific fields, so a hook masking result titles on a stream never saw them. The hook's view now carries those blocks, its rewrite of them goes back through the ended-stream write-back, and the adapter no longer needs to pass texts outside the hook's view through untouched. * test(guardrails): type the legacy hook stream test helpers * fix(guardrails): scan a streamed web search query as text like the non-streaming reply * test(guardrails): type the masking guardrail test helpers * fix(guardrails): keep the block id when scanning a streamed web search query --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../websearch_interception/handler.py | 15 +- .../websearch_interception/transformation.py | 10 +- .../chat/guardrail_translation/handler.py | 335 ++++++++++++++++-- .../messages/native/fake_stream_iterator.py | 56 +-- .../proxy/policy_engine/pipeline_executor.py | 36 +- litellm/types/llms/anthropic.py | 19 + .../test_websearch_interception_wire.py | 27 +- .../test_websearch_short_circuit.py | 122 +++++++ .../test_anthropic_guardrail_handler.py | 307 +++++++++++++++- .../native/test_fake_stream_iterator.py | 145 +++++++- .../policy_engine/test_pipeline_executor.py | 238 +++++++++++-- 11 files changed, 1191 insertions(+), 119 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 4739b3ee561..2bb89248be4 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -398,14 +398,17 @@ class WebSearchInterceptionLogger(CustomLogger): (t for t in tools if is_anthropic_native_web_search_tool(t)), None, ) + emit_native_blocks: Final = native_tool is not None or ( + kwargs is not None and bool(kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY)) + ) outcome: Final = await self._short_circuit_search_outcome(query, kwargs=kwargs) search_result_text: Final = WebSearchTransformation.search_outcome_text(outcome) content: Final[list[dict[str, object]]] = [] - if native_tool is not None: + if emit_native_blocks: tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}" - tool_name: Final = native_tool.get("name") or "web_search" + tool_name: Final = (native_tool.get("name") if native_tool is not None else None) or "web_search" content.append( { "type": "server_tool_use", @@ -415,10 +418,10 @@ class WebSearchInterceptionLogger(CustomLogger): } ) content.append( - WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome) + WebSearchTransformation.build_web_search_outcome_block( + tool_use_id=tool_use_id, outcome=outcome, with_snippets=False + ) ) - # Keep the text block so non-native short-circuit callers (Claude Code, - # github_copilot, etc.) see the same payload they always have. content.append({"type": "text", "text": search_result_text}) response: Final[dict[str, object]] = { @@ -435,7 +438,7 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.debug( "WebSearchInterception: Short-circuit search completed, returning synthetic response (%s chars, native_blocks=%s)", len(search_result_text), - native_tool is not None, + emit_native_blocks, ) return response diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index 58ae8bd6e48..55564dd1134 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -412,6 +412,7 @@ class WebSearchTransformation: def build_web_search_tool_result_block( tool_use_id: str, search_response: SearchResponse | None, + with_snippets: bool = True, ) -> dict[str, object]: """ Build an Anthropic-native ``web_search_tool_result`` content block. @@ -445,6 +446,8 @@ class WebSearchTransformation: ``litellm.asearch()``. If None or empty, the block is still emitted with an empty result list (signals "search ran, no results" rather than "search did not run"). + with_snippets: False when a text block in the same reply already + carries the snippets. """ items: Final[list[dict[str, object]]] = [] if search_response is not None: @@ -460,7 +463,7 @@ class WebSearchTransformation: "title": title, "page_age": page_age, "encrypted_content": "", - "snippet": getattr(r, "snippet", "") or "", + **({"snippet": getattr(r, "snippet", "") or ""} if with_snippets else {}), } ) return { @@ -481,12 +484,15 @@ class WebSearchTransformation: } @staticmethod - def build_web_search_outcome_block(tool_use_id: str, outcome: SearchOutcome) -> dict[str, object]: + def build_web_search_outcome_block( + tool_use_id: str, outcome: SearchOutcome, with_snippets: bool = True + ) -> dict[str, object]: match outcome: case SearchSucceeded(response=response): return WebSearchTransformation.build_web_search_tool_result_block( tool_use_id=tool_use_id, search_response=response, + with_snippets=with_snippets, ) case SearchFailed(error_code=error_code): return WebSearchTransformation.build_web_search_tool_result_error_block( diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e351542da01..7f688ae733b 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -24,7 +24,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.messages.adapters.transformation import ( OPENAI_MAX_TOOL_NAME_LENGTH, @@ -61,6 +61,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( unappliable_request_rewrite, ) from litellm.types.llms.anthropic import ( + ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX, AllAnthropicToolsValues, AnthropicMessagesRequest, ) @@ -153,6 +154,7 @@ class _TextAttrBlock(Protocol): text: str +@runtime_checkable class _WritableMessage(Protocol): @overload def get(self, key: str, /) -> object | None: ... @@ -195,6 +197,140 @@ class _ToolCallShape: arguments: str +@dataclass(frozen=True, slots=True) +class _WebSearchString: + block_idx: int + result_idx: int | None + field: str + value: str + + +def _streamed_web_search_result_blocks(response: ModelResponse) -> tuple[Mapping[str, object], ...]: + choice: Final = response.choices[0] if response.choices else None + fields: Final = as_str_mapping(choice.message.provider_specific_fields) if isinstance(choice, Choices) else None + blocks: Final = ( + (as_str_mapping(item) for item in iter_items(fields.get("web_search_results"))) + if fields is not None and isinstance(fields.get("web_search_results"), list) + else () + ) + return tuple( + deepcopy(block) for block in blocks if block is not None and block.get("type") == "web_search_tool_result" + ) + + +def _web_search_results(block: Mapping[str, object]) -> Iterable[object]: + return iter_items(block.get("content")) if isinstance(block.get("content"), list) else () + + +def _web_search_result_strings(block_idx: int, result_idx: int, result: object) -> tuple[_WebSearchString, ...]: + fields: Final = as_str_mapping(result) + if fields is None or fields.get("type") != "web_search_result": + return () + values: Final = ((field, fields.get(field)) for field in ("title", "url")) + return tuple( + _WebSearchString(block_idx, result_idx, field, value) + for field, value in values + if isinstance(value, str) and value + ) + + +def _is_server_tool_use(block: Mapping[str, object]) -> bool: + block_id: Final = block.get("id") + match block.get("type"): + case "server_tool_use": + return True + case "tool_use": + return isinstance(block_id, str) and block_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX) + case _: + return False + + +def _web_search_block_strings( + block_idx: int, block: Mapping[str, object], *, with_query: bool +) -> tuple[_WebSearchString, ...]: + """The strings a client reads off a native web search block: the ``server_tool_use`` query and each + ``web_search_tool_result`` result's title and url.""" + match block.get("type"): + case "server_tool_use" | "tool_use" if with_query and _is_server_tool_use(block): + tool_input: Final = as_str_mapping(block.get("input")) + query: Final = None if tool_input is None else tool_input.get("query") + return (_WebSearchString(block_idx, None, "query", query),) if isinstance(query, str) and query else () + case "web_search_tool_result": + return tuple( + chain.from_iterable( + _web_search_result_strings(block_idx, result_idx, result) + for result_idx, result in enumerate(_web_search_results(block)) + ) + ) + case _: + return () + + +def _rewritten_web_search_result( + result_idx: int, result: object, rewrites: Mapping[tuple[int | None, str], str] +) -> object: + fields: Final = as_str_mapping(result) + if fields is None: + return result + return { + **fields, + **{field: rewrites[(result_idx, field)] for field in ("title", "url") if (result_idx, field) in rewrites}, + } + + +def _rewritten_web_search_block( + block: Mapping[str, object], rewrites: Sequence[tuple[_WebSearchString, str]] +) -> Mapping[str, object]: + by_location: Final = MappingProxyType({(target.result_idx, target.field): value for target, value in rewrites}) + if _is_server_tool_use(block): + query: Final = by_location.get((None, "query")) + tool_input: Final = as_str_mapping(block.get("input")) or MappingProxyType({}) + return block if query is None else {**block, "input": {**tool_input, "query": query}} + if not isinstance(block.get("content"), list): + return block + return { + **block, + "content": [ + _rewritten_web_search_result(result_idx, result, by_location) + for result_idx, result in enumerate(_web_search_results(block)) + ], + } + + +def _write_back_web_search_block(content_block: object, rewrites: Sequence[tuple[_WebSearchString, str]]) -> None: + block: Final = as_str_mapping(content_block) + if block is None or not isinstance(content_block, _WritableMessage): + return + rewritten: Final = _rewritten_web_search_block(block, rewrites) + writable: Final = _as_writable(content_block) + for key in ("input", "content"): + if key in rewritten: + writable[key] = rewritten[key] + + +def _split_web_search_texts( + guardrailed: Sequence[str], web_search_count: int, reply_count: int +) -> tuple[Sequence[str], Sequence[str]]: + split_at: Final = web_search_count if len(guardrailed) == web_search_count + reply_count else 0 + return guardrailed[:split_at], guardrailed[split_at:] + + +def _web_search_rewrites_by_block( + strings: Sequence[_WebSearchString], guardrailed: Sequence[object] +) -> Mapping[int, tuple[tuple[_WebSearchString, str], ...]]: + rewrites: Final = tuple( + (target, value) + for target, value in zip(strings, guardrailed) + if isinstance(value, str) and value != target.value + ) + return MappingProxyType( + { + block_idx: tuple(rewrite for rewrite in rewrites if rewrite[0].block_idx == block_idx) + for block_idx in {target.block_idx for target, _ in rewrites} + } + ) + + def _is_client_tool_use(block: Mapping[str, object]) -> bool: return ( block.get("type") == "tool_use" @@ -289,6 +425,7 @@ def _write_back_message_text(message: _WritableMessage, target: MessageTextTarge _TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object]) _DECODED_SSE_DATA: Final = TypeAdapter(object) +_REBUILT_TOOL_CALL_BLOCK_TYPES: Final = frozenset({"tool_use", "server_tool_use"}) _RELEASED_TOOL_USE_STOP: Final = ( b"event: message_delta\n" b'data: {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": null}, ' @@ -315,6 +452,51 @@ def _write_back_tool_use( block["name"] = shape.name +def _streamed_query_strings( + tool_call_blocks: Sequence[tuple[int, Mapping[str, object]]], tool_calls: Sequence[_ToolCallShape] +) -> tuple[_WebSearchString, ...]: + if len(tool_call_blocks) != len(tool_calls): + return () + return tuple( + chain.from_iterable( + _web_search_block_strings( + index, {**block, "input": _rewritten_tool_use_input(tool_call.arguments)}, with_query=True + ) + for (index, block), tool_call in zip(tool_call_blocks, tool_calls) + ) + ) + + +def _query_rewrites(strings: Sequence[_WebSearchString], guardrailed: Sequence[object]) -> Mapping[int, str]: + return MappingProxyType( + { + target.block_idx: value + for target, value in zip(strings, guardrailed) + if isinstance(value, str) and value != target.value + } + ) + + +def _with_query(tool_call: _ToolCallShape, query: str) -> _ToolCallShape: + arguments: Final = _rewritten_tool_use_input(tool_call.arguments) + if arguments is None: + return tool_call + return _ToolCallShape(name=tool_call.name, arguments=json.dumps({**arguments, "query": query})) + + +def _with_query_rewrites( + tool_calls: tuple[_ToolCallShape, ...], + tool_call_blocks: Sequence[tuple[int, Mapping[str, object]]], + query_rewrites: Mapping[int, str], +) -> tuple[_ToolCallShape, ...]: + if not query_rewrites or len(tool_call_blocks) != len(tool_calls): + return tool_calls + return tuple( + _with_query(tool_call, query_rewrites[index]) if index in query_rewrites else tool_call + for (index, _), tool_call in zip(tool_call_blocks, tool_calls) + ) + + @dataclass(frozen=True, slots=True) class _SSEFieldRewrite: """One field of one nested section of a buffered SSE event, rewritten.""" @@ -402,7 +584,12 @@ class AnthropicMessagesHandler(BaseTranslation): def post_call_hook_response(self, response: object) -> object: if not isinstance(response, ModelResponse): return response - return self.adapter.translate_openai_response_to_anthropic(response) + translated: Final = self.adapter.translate_openai_response_to_anthropic(response) + result_blocks: Final = _streamed_web_search_result_blocks(response) + if not result_blocks: + return translated + content: Final = translated.get("content") or () + return {**translated, "content": [*result_blocks, *content]} @staticmethod def _build_streaming_usage_response( @@ -1273,9 +1460,17 @@ class AnthropicMessagesHandler(BaseTranslation): task_mappings, tool_calls_to_check, ) + block_fields: Final = tuple(map(self._output_block_fields, response_content)) + web_search_strings: Final = tuple( + chain.from_iterable( + _web_search_block_strings(content_idx, fields[1], with_query=True) + for content_idx, fields in enumerate(block_fields) + if fields is not None + ) + ) # Step 2: Apply guardrail to all texts in batch - if texts_to_check or tool_calls_to_check: + if texts_to_check or tool_calls_to_check or web_search_strings: request_data = self._prepare_request_data( request_data, response, @@ -1284,7 +1479,7 @@ class AnthropicMessagesHandler(BaseTranslation): ) inputs: Final = self._build_guardrail_inputs( - texts_to_check, + [*(target.value for target in web_search_strings), *texts_to_check], images_to_check, tool_calls_to_check, response, @@ -1297,14 +1492,19 @@ class AnthropicMessagesHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + web_search_rewrites, reply_rewrites = _split_web_search_texts( + guardrailed_inputs.get("texts", []), len(web_search_strings), len(texts_to_check) + ) # Step 3: Map guardrail responses back to original response structure await self._apply_guardrail_responses_to_output( response=response, - responses=guardrailed_texts, + responses=list(reply_rewrites), task_mappings=task_mappings, ) + self._write_back_web_search_strings( + response_content, _web_search_rewrites_by_block(web_search_strings, web_search_rewrites) + ) verbose_proxy_logger.debug("Anthropic Messages: Processed output response: %s", response) @@ -1350,8 +1550,14 @@ class AnthropicMessagesHandler(BaseTranslation): string_so_far = first_choice.message.content pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ()) guardrail_inputs: Final = GenericGuardrailAPIInputs() - if string_so_far: - guardrail_inputs["texts"] = [string_so_far] + leading_texts: Final = (string_so_far,) if isinstance(string_so_far, str) and string_so_far else () + tool_call_blocks: Final = self._streamed_tool_call_blocks(responses_so_far) + streamed_query_strings: Final = _streamed_query_strings(tool_call_blocks, pre_guardrail_tool_calls) + streamed_result_strings: Final = self._streamed_web_search_result_strings(responses_so_far) + streamed_web_search_strings: Final = (*streamed_result_strings, *streamed_query_strings) + texts_to_check: Final = [*(target.value for target in streamed_web_search_strings), *leading_texts] + if texts_to_check: + guardrail_inputs["texts"] = texts_to_check if tool_calls_list: guardrail_inputs["tool_calls"] = tool_calls_list @@ -1374,29 +1580,42 @@ class AnthropicMessagesHandler(BaseTranslation): responses_so_far, request_data ) raise - guardrailed_texts: Final = _guardrailed_inputs.get("texts") + web_search_rewrites, reply_rewrites = _split_web_search_texts( + _guardrailed_inputs.get("texts") or [], len(streamed_web_search_strings), len(leading_texts) + ) if ( deliver_ended_stream_rewrites and isinstance(string_so_far, str) and string_so_far - and guardrailed_texts - and guardrailed_texts[0] != string_so_far + and reply_rewrites + and reply_rewrites[0] != string_so_far ): self._write_ended_stream_text_rewrite( responses_so_far, - guardrailed_texts[0], + reply_rewrites[0], guardrail_name=guardrail_to_apply.guardrail_name or "unknown", ) + if deliver_ended_stream_rewrites: + self._write_ended_stream_web_search_rewrites( + responses_so_far, + _web_search_rewrites_by_block(streamed_result_strings, web_search_rewrites), + ) if deliver_ended_stream_rewrites: returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls") self._write_ended_stream_tool_call_rewrites( responses_so_far, pre_guardrail_tool_calls=pre_guardrail_tool_calls, - post_guardrail_tool_calls=_tool_call_shapes( - returned_tool_calls - if isinstance(returned_tool_calls, list) - and len(returned_tool_calls) == len(pre_guardrail_tool_calls) - else tool_calls_list or () + post_guardrail_tool_calls=_with_query_rewrites( + _tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) + and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + else tool_calls_list or () + ), + tool_call_blocks, + _query_rewrites( + streamed_query_strings, web_search_rewrites[len(streamed_result_strings) :] + ), ), guardrail_name=guardrail_to_apply.guardrail_name or "unknown", ) @@ -1451,6 +1670,14 @@ class AnthropicMessagesHandler(BaseTranslation): request_data["litellm_metadata"] = user_metadata return request_data + @staticmethod + def _write_back_web_search_strings( + response_content: Sequence[object], + rewrites_by_block: Mapping[int, tuple[tuple[_WebSearchString, str], ...]], + ) -> None: + for block_idx, rewrites in rewrites_by_block.items(): + _write_back_web_search_block(_content_block_at(response_content, block_idx), rewrites) + @staticmethod def _get_response_content(response: object) -> Sequence[object]: """Extract content list from a dict or object response.""" @@ -1573,30 +1800,22 @@ class AnthropicMessagesHandler(BaseTranslation): ) -> None: """Deliver ended-stream guardrail tool-call rewrites by rewriting the buffered chunks in place: the rebuilt response lists tool calls in the - order of the stream's ``tool_use`` blocks, so the nth rewritten call lands - on the nth block, its first ``input_json_delta`` carrying the full rewritten + order of the stream's ``tool_use`` and ``server_tool_use`` blocks, so the nth + rewritten call lands on the nth block, its first ``input_json_delta`` carrying the full rewritten arguments, every later one blanked, and ``content_block_start`` carrying the rewritten name. Blocks that do not line up with the rebuilt tool calls make the rewrite undeliverable, so the pipeline executor discards it and releases the original chunks.""" if post_guardrail_tool_calls == pre_guardrail_tool_calls: return - block_indices: Final = tuple( - index - for item in responses_so_far - for event in cls._iter_sse_events(item) - if event.get("type") == "content_block_start" - and isinstance(block := event.get("content_block"), Mapping) - and block.get("type") == "tool_use" - and isinstance(index := event.get("index"), int) - ) + block_indices: Final = tuple(index for index, _ in cls._streamed_tool_call_blocks(responses_so_far)) if len(block_indices) != len(post_guardrail_tool_calls): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite raise UndeliverableStreamRewrite( guardrail_name, f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried " - f"{len(block_indices)} tool_use blocks", + f"{len(block_indices)} tool_use or server_tool_use blocks", ) rewrites_by_block: Final = MappingProxyType( { @@ -1636,6 +1855,60 @@ class AnthropicMessagesHandler(BaseTranslation): ), ) + @classmethod + def _streamed_web_search_result_strings(cls, responses_so_far: Sequence[object]) -> tuple[_WebSearchString, ...]: + events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far) + result_blocks: Final = ( + (index, block) + for event in events + if event.get("type") == "content_block_start" + and isinstance(index := event.get("index"), int) + and (block := as_str_mapping(event.get("content_block"))) is not None + ) + return tuple( + chain.from_iterable( + _web_search_block_strings(index, block, with_query=False) for index, block in result_blocks + ) + ) + + @classmethod + def _write_ended_stream_web_search_rewrites( + cls, + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place + rewrites_by_block: Mapping[int, tuple[tuple[_WebSearchString, str], ...]], + ) -> None: + """Deliver ended-stream guardrail rewrites of web search result titles and urls into the + ``content_block_start`` that carries each whole ``web_search_tool_result`` block.""" + if not rewrites_by_block: + return + + def rewrite_result_block(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + index: Final = event.get("index") + block: Final = as_str_mapping(event.get("content_block")) + rewrites: Final = rewrites_by_block.get(index) if isinstance(index, int) else None + if event.get("type") != "content_block_start" or block is None or rewrites is None: + return None + return _SSEFieldRewrite( + "content_block", "content", _rewritten_web_search_block(block, rewrites).get("content") + ) + + cls._rewrite_ended_stream_events(responses_so_far, rewrite_result_block) + + @classmethod + def _streamed_tool_call_blocks( + cls, responses_so_far: Sequence[object] + ) -> tuple[tuple[int, Mapping[str, object]], ...]: + events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far) + return tuple( + (index, block) + for event in events + if event.get("type") == "content_block_start" + and (block := as_str_mapping(event.get("content_block"))) is not None + and isinstance(block_type := block.get("type"), str) + and block_type in _REBUILT_TOOL_CALL_BLOCK_TYPES + and isinstance(index := event.get("index"), int) + ) + @classmethod def _streamed_tool_use_ids(cls, responses_so_far: Sequence[object]) -> Mapping[int, str]: """The id of each buffered ``tool_use`` block, by its content block index.""" @@ -1738,8 +2011,10 @@ class AnthropicMessagesHandler(BaseTranslation): def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: stream_ended: Final = self._check_streaming_has_ended(responses_so_far) tool_use_fingerprints: Final = self._streamed_tool_use_fingerprints(responses_so_far) + streamed_text: Final = self.get_streaming_string_so_far(responses_so_far) + result_strings: Final = self._streamed_web_search_result_strings(responses_so_far) if stream_ended else () return StreamingScanKey( - texts=(self.get_streaming_string_so_far(responses_so_far),), + texts=(*(target.value for target in result_strings), streamed_text), tool_calls=tool_use_fingerprints if stream_ended else (), stream_ended=stream_ended, tool_calls_in_flight=bool(tool_use_fingerprints) and not stream_ended, diff --git a/litellm/llms/anthropic/messages/native/fake_stream_iterator.py b/litellm/llms/anthropic/messages/native/fake_stream_iterator.py index 39bdfbbc53a..3a26737a8ac 100644 --- a/litellm/llms/anthropic/messages/native/fake_stream_iterator.py +++ b/litellm/llms/anthropic/messages/native/fake_stream_iterator.py @@ -12,11 +12,20 @@ import json from collections.abc import Mapping, Sequence from typing import Final, cast +from pydantic import TypeAdapter + from litellm.llms.anthropic.messages.native.utils import anthropic_message_delta from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) +_CITATIONS: Final = TypeAdapter(tuple[object, ...]) + + +def _citations_of(block: Mapping[str, object]) -> tuple[object, ...]: + citations: Final = block.get("citations") + return _CITATIONS.validate_python(citations) if isinstance(citations, list) else () + class FakeAnthropicMessagesStreamIterator: """ @@ -47,16 +56,30 @@ class FakeAnthropicMessagesStreamIterator: def _create_content_block_chunks(self, block_dict: Mapping[str, object], index: int) -> list[bytes]: """Build SSE chunks for a single content block.""" - chunks: Final = [] + chunks: Final = list[bytes]() block_type: Final = block_dict.get("type") if block_type == "text": + citations: Final = _citations_of(block_dict) content_block_start = { "type": "content_block_start", "index": index, - "content_block": {"type": "text", "text": ""}, + "content_block": {"type": "text", "text": "", "citations": ()} + if citations + else {"type": "text", "text": ""}, } chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) + citation_deltas: Final = ( + { + "type": "content_block_delta", + "index": index, + "delta": {"type": "citations_delta", "citation": citation}, + } + for citation in citations + ) + chunks.extend( + f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n".encode() for delta in citation_deltas + ) text: Final = block_dict.get("text", "") content_block_delta = { "type": "content_block_delta", @@ -89,36 +112,23 @@ class FakeAnthropicMessagesStreamIterator: } chunks.append(f"event: content_block_delta\ndata: {json.dumps(signature_delta)}\n\n".encode()) - elif block_type == "redacted_thinking": - content_block_start = { + elif block_type in ("tool_use", "server_tool_use"): + empty_input: Final[Mapping[str, object]] = {} + tool_use_start: Final = { "type": "content_block_start", "index": index, - "content_block": {"type": "redacted_thinking"}, + "content_block": {**block_dict, "input": empty_input}, } - chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) - - elif block_type == "tool_use": - content_block_start = { - "type": "content_block_start", - "index": index, - "content_block": { - "type": "tool_use", - "id": block_dict.get("id"), - "name": block_dict.get("name"), - "input": {}, - }, - } - chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) - input_data: Final = block_dict.get("input", {}) - content_block_delta = { + chunks.append(f"event: content_block_start\ndata: {json.dumps(tool_use_start)}\n\n".encode()) + tool_use_delta: Final = { "type": "content_block_delta", "index": index, "delta": { "type": "input_json_delta", - "partial_json": json.dumps(input_data), + "partial_json": json.dumps(block_dict.get("input", {})), }, } - chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode()) + chunks.append(f"event: content_block_delta\ndata: {json.dumps(tool_use_delta)}\n\n".encode()) else: passthrough_start: Final = { diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 4ad29bd3478..47572dd7ceb 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -98,15 +98,11 @@ def _count_change(sent: tuple[object, ...] | None, returned: tuple[object, ...] return (len(sent), len(returned)) -def _tool_call_mismatch_reason( - sent: tuple[tuple[object, object], ...] | None, returned: tuple[tuple[object, object], ...] | None -) -> str | None: - if sent == returned: - return None +def _tool_call_count_mismatch_reason(sent: Sequence[object] | None, returned: Sequence[object] | None) -> str | None: sent_count: Final = len(sent or ()) returned_count: Final = len(returned or ()) if sent_count == returned_count: - return "the legacy hook changed a tool call's name or arguments, which this path cannot write back" + return None return f"the legacy hook returned {returned_count} tool calls for a stream that carried {sent_count}" @@ -205,11 +201,10 @@ class _LegacyHookStreamAdapter(CustomGuardrail): non-streaming hooks, an exception it raises ends the stream through the executor's fail/error classification, and the response it hands back, or the one it changed in place and returned ``None`` for, is re-scanned by the same translation so its texts reach the - client through the translation's ended-stream write-back. A - replacement whose scanned texts do not line up with the originals, or whose tool calls - differ from them, is undeliverable, so the executor releases the original chunks. A stream - that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as - long as the hook left the tool calls alone.""" + client through the translation's ended-stream write-back. Tool calls the hook rewrote go to the translation's + tool-call write-back like a unified guardrail's. A replacement whose scanned texts or tool + calls do not line up with the originals by count is undeliverable, so the executor releases + the original chunks.""" def __init__( self, @@ -258,14 +253,19 @@ class _LegacyHookStreamAdapter(CustomGuardrail): guardrail_name, f"the legacy hook returned {returned_text_count} texts for a stream that carried {sent_text_count}", ) - tool_call_mismatch: Final = _tool_call_mismatch_reason( - _tool_call_shapes(inputs.get("tool_calls")), _tool_call_shapes(rescanned.get("tool_calls")) + sent_tool_calls: Final = inputs.get("tool_calls") + rewritten_tool_calls: Final = rescanned.get("tool_calls") + tool_call_count_mismatch: Final = _tool_call_count_mismatch_reason(sent_tool_calls, rewritten_tool_calls) + if tool_call_count_mismatch is not None: + raise UndeliverableStreamRewrite(guardrail_name, tool_call_count_mismatch) + texts_update: Final[GenericGuardrailAPIInputs] = {"texts": rewritten} if rewritten else {} + tool_calls_update: Final[GenericGuardrailAPIInputs] = ( + {"tool_calls": rewritten_tool_calls} + if rewritten_tool_calls is not None + and _rewrote(_tool_call_shapes(sent_tool_calls), _tool_call_shapes(rewritten_tool_calls)) + else {} ) - if tool_call_mismatch is not None: - raise UndeliverableStreamRewrite(guardrail_name, tool_call_mismatch) - if not rewritten: - return inputs - rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} + rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, **texts_update, **tool_calls_update} return rewritten_inputs async def _rescan( diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 50c4a2340a3..a648376c542 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -713,6 +713,23 @@ class AnthropicResponseContentBlockRedactedThinking(LiteLLMBaseModel): data: str +class AnthropicResponseContentBlockServerToolUse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["server_tool_use"] + id: str + name: str + input: Mapping[str, object] + + +class AnthropicResponseContentBlockWebSearchToolResult(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["web_search_tool_result"] + tool_use_id: str + content: Sequence[Mapping[str, object]] | Mapping[str, object] + + class AnthropicResponseUsageBlock(LiteLLMBaseModel): model_config = ConfigDict(extra="allow") @@ -744,6 +761,8 @@ class AnthropicResponse(LiteLLMBaseModel): | AnthropicResponseContentBlockToolUse | AnthropicResponseContentBlockThinking | AnthropicResponseContentBlockRedactedThinking + | AnthropicResponseContentBlockServerToolUse + | AnthropicResponseContentBlockWebSearchToolResult ] """Content generated by the model.""" diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py index 5ec860910e5..1bd9f18253e 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py @@ -29,6 +29,24 @@ def sse_events(text: str) -> tuple[tuple[str, dict[str, object]], ...]: ) +def streamed_blocks(events: tuple[tuple[str, dict[str, object]], ...]) -> tuple[dict[str, object], ...]: + input_json: Final = tuple( + (event["index"], event["delta"]["partial_json"]) + for name, event in events + if name == "content_block_delta" and event["delta"]["type"] == "input_json_delta" + ) + return tuple( + { + **event["content_block"], + "input": json.loads("".join(part for index, part in input_json if index == event["index"])), + } + if any(index == event["index"] for index, _ in input_json) + else event["content_block"] + for name, event in events + if name == "content_block_start" + ) + + @pytest.mark.covers("other.provider_wire.bedrock.websearch_interception_streamed_capped_turn_ends_with_native_results") def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_snippets_and_ordered_blocks( gateway: Gateway, tmp_path: Path @@ -136,7 +154,7 @@ def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_sni for earlier, earlier_event in events[:position] if earlier == "content_block_start" }, response.text - started: Final = tuple(event["content_block"] for name, event in events if name == "content_block_start") + started: Final = streamed_blocks(events) search_ids: Final = tuple(block["id"] for block in started if block["type"] == "server_tool_use") assert search_ids and all(search_id.startswith("srvtoolu_") for search_id in search_ids), response.text assert started[-1] == {"type": "text", "text": ""}, response.text @@ -162,7 +180,12 @@ def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_sni ) ), response.text assert ( - "".join(event["delta"]["text"] for name, event in events if name == "content_block_delta") == "turn-1" + "".join( + event["delta"]["text"] + for name, event in events + if name == "content_block_delta" and event["delta"]["type"] == "text_delta" + ) + == "turn-1" ), response.text assert [event["delta"]["stop_reason"] for name, event in events if name == "message_delta"] == [ "end_turn" diff --git a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py index ef050d96dc5..10cb6ff12e9 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py @@ -5,7 +5,9 @@ Tests the short-circuit path that detects web-search-only /v1/messages requests and executes the search directly without routing through the backend LLM. """ +import datetime import json +from collections.abc import Iterator, Mapping, Sequence from typing import Final from unittest.mock import AsyncMock, patch @@ -16,6 +18,7 @@ import litellm from litellm.integrations.websearch_interception.handler import ( WebSearchInterceptionLogger, ) +from litellm.messages.dispatch import MessagesResult # --------------------------------------------------------------------------- # Detection tests @@ -73,6 +76,42 @@ class TestTryShortCircuitSearch: assert "Result" in text_block["text"] mock_search.assert_called_once_with("Search for Claude Code releases") + @pytest.mark.asyncio + async def test_emits_native_blocks_when_tools_already_converted(self, perplexity_search: respx.Route): + """Pre-request/deployment hooks convert the native tool to the LiteLLM + standard shape before the short-circuit runs and record that in + WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY; the flag alone must be enough for + native blocks to be emitted.""" + from litellm.integrations.websearch_interception.handler import ( + WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY, + ) + + logger: Final = WebSearchInterceptionLogger(enabled_providers=["hosted_vllm"]) + + result: Final = await logger.try_short_circuit_search( + model="hosted_vllm/some-model", + messages=[{"role": "user", "content": "Search for something"}], + tools=[ + { + "name": "litellm_web_search", + "description": "Search the web for information.", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + } + ], + custom_llm_provider="hosted_vllm", + kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + assert perplexity_search.call_count == 1 + assert result is not None + assert [b["type"] for b in result["content"]] == ["server_tool_use", "web_search_tool_result", "text"] + assert result["content"][0]["name"] == "web_search" + assert [r["url"] for r in result["content"][1]["content"]] == ["https://example.com/releases"] + @pytest.mark.asyncio async def test_does_not_short_circuit_mixed_tools(self): """Mix of web_search and other tools → NOT short-circuited""" @@ -468,6 +507,89 @@ class TestShortCircuitEntryPoint: assert text_block["text"] == "results" +class TestShortCircuitThroughMessagesEntryPoint: + """The pre-request hook renames the native tool before the short-circuit + runs, so only a request through anthropic_messages shows whether native + clients still get their server_tool_use and web_search_tool_result blocks.""" + + NATIVE_TOOL: Final = {"type": "web_search_20250305", "name": "web_search"} + + @pytest.fixture(autouse=True) + def _intercept_web_search(self, perplexity_search: respx.Route) -> Iterator[respx.Route]: + with patch("litellm.callbacks", [WebSearchInterceptionLogger(enabled_providers=["bedrock"])]): + yield perplexity_search + + @staticmethod + async def _send(*, tools: Sequence[Mapping[str, object]], stream: bool, **kwargs: object) -> MessagesResult: + return await litellm.anthropic_messages( + model="bedrock/converse/us.anthropic.claude-sonnet-5-5", + messages=[{"role": "user", "content": "latest python release"}], + tools=tools, + max_tokens=64, + stream=stream, + **kwargs, + ) + + @pytest.mark.asyncio + async def test_native_tool_gets_native_blocks(self, perplexity_search: respx.Route): + response: Final = await self._send(tools=[self.NATIVE_TOOL], stream=False) + + assert perplexity_search.call_count == 1 + assert [block["type"] for block in response["content"]] == [ + "server_tool_use", + "web_search_tool_result", + "text", + ] + server_tool_use, tool_result, text = response["content"] + assert server_tool_use["name"] == "web_search" + assert server_tool_use["input"] == {"query": "latest python release"} + assert tool_result["tool_use_id"] == server_tool_use["id"] + assert [result["url"] for result in tool_result["content"]] == ["https://example.com/releases"] + assert "Release notes" in text["text"] + assert "Release notes" not in json.dumps(tool_result) + + @pytest.mark.asyncio + async def test_native_blocks_reply_can_be_costed(self): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = Logging( + model="bedrock/converse/us.anthropic.claude-sonnet-5-5", + messages=[{"role": "user", "content": "latest python release"}], + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="short-circuit-cost", + function_id="short-circuit-cost", + ) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + response: Final = await self._send(tools=[self.NATIVE_TOOL], stream=False, litellm_logging_obj=logging_obj) + + assert logging_obj.response_cost_calculator(result=response) == 0.0 + + @pytest.mark.asyncio + async def test_native_tool_streams_native_blocks(self): + stream: Final = await self._send(tools=[self.NATIVE_TOOL], stream=True) + body: Final = b"".join([chunk async for chunk in stream]) + + starts: Final = [ + json.loads(line.removeprefix(b"data: "))["content_block"] + for line in body.splitlines() + if line.startswith(b"data: ") and b'"content_block_start"' in line + ] + assert [block["type"] for block in starts] == ["server_tool_use", "web_search_tool_result", "text"] + assert starts[1]["tool_use_id"] == starts[0]["id"] + + @pytest.mark.asyncio + async def test_litellm_standard_tool_stays_text_only(self): + from litellm.integrations.websearch_interception.tools import ( + get_litellm_web_search_tool, + ) + + response: Final = await self._send(tools=[get_litellm_web_search_tool()], stream=False) + + assert [block["type"] for block in response["content"]] == ["text"] + + class TestShortCircuitSafeguardResults: """A short-circuited reply never reaches the backend, so the proxy answers Claude Code's safeguards itself.""" diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 141ad8a0481..d815e1eedd3 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -6,6 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic """ import json +import re from typing import Any, Literal, Optional, Final from unittest.mock import MagicMock, patch @@ -552,13 +553,25 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert chunks == original @pytest.mark.asyncio - async def test_deliver_ended_stream_tool_use_rewrite_with_server_tool_use_block_fails_closed(self): - from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - + async def test_deliver_ended_stream_tool_use_rewrite_lands_on_server_tool_use_and_tool_use_blocks(self): handler = AnthropicMessagesHandler() server_tool_use = [ - ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}}), - ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}}), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}, + }, + ), ("content_block_stop", {"type": "content_block_stop", "index": 0}), ] tool_use = self._ended_tool_use_sse_chunks() @@ -568,6 +581,26 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: + [chunk.replace(b'"index": 0', b'"index": 1') for chunk in tool_use[1:]] ) + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', '{"fruit": "[MASKED]"}', "", ""] + raw = b"".join(chunks).decode() + assert '"id": "srvtoolu_1"' in raw and '"id": "toolu_1"' in raw + assert "persim" not in raw and '"query"' not in raw + + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_use_rewrite_on_an_unindexed_tool_block_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = AnthropicMessagesHandler() + tool_use = self._ended_tool_use_sse_chunks() + chunks = tool_use[:1] + [tool_use[1].replace(b'"index": 0', b'"index": "0"')] + tool_use[2:] + with pytest.raises(UndeliverableStreamRewrite): await handler.process_output_streaming_response( responses_so_far=chunks, @@ -576,6 +609,270 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: deliver_ended_stream_rewrites=True, ) + @staticmethod + def _python_masking_guardrail() -> CustomGuardrail: + class MaskPython(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + for tool_call in inputs.get("tool_calls", []): + if isinstance(tool_call, ChatCompletionMessageToolCall): + tool_call.function.arguments = re.sub( + "python", "[MASKED]", tool_call.function.arguments, flags=re.I + ) + return { + **inputs, + "texts": [re.sub("python", "[MASKED]", text, flags=re.I) for text in inputs.get("texts", [])], + } + + return MaskPython(guardrail_name="test") + + @staticmethod + def _python_texts_masking_guardrail() -> CustomGuardrail: + class MaskPythonTexts(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + return { + **inputs, + "texts": [re.sub("python", "[MASKED]", text, flags=re.I) for text in inputs.get("texts", [])], + } + + return MaskPythonTexts(guardrail_name="test") + + @staticmethod + def _web_search_result_block() -> dict[str, object]: + return { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.python.org/3.14/", + "title": "Python 3.14 docs", + "encrypted_content": "", + "page_age": None, + }, + { + "type": "web_search_result", + "url": "https://peps.example/745", + "title": "PEP 745", + "encrypted_content": "", + "page_age": None, + }, + ], + } + + @classmethod + def _web_search_stream_chunks(cls, server_tool_block_type: str = "server_tool_use") -> list[bytes]: + events = [ + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": server_tool_block_type, + "id": "srvtoolu_1", + "name": "web_search", + "input": {}, + }, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"query": "python 3.14 release date"}'}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "content_block_start", + {"type": "content_block_start", "index": 1, "content_block": cls._web_search_result_block()}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 1}), + ( + "content_block_start", + {"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 2, + "delta": {"type": "text_delta", "text": "Python 3.14 shipped"}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 2}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + @pytest.mark.parametrize("server_tool_block_type", ["server_tool_use", "tool_use"]) + @pytest.mark.parametrize( + "guardrail_factory", + [_python_masking_guardrail, _python_texts_masking_guardrail], + ids=["texts_and_tool_calls", "texts_only"], + ) + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_mask_the_web_search_query_results_and_text( + self, guardrail_factory, server_tool_block_type + ): + chunks = self._web_search_stream_chunks(server_tool_block_type) + + await AnthropicMessagesHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail_factory(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + raw = b"".join(chunks).decode() + assert "python" not in raw.lower() + result_block = next( + json.loads(line[len("data:") :])["content_block"] + for line in raw.split("\n") + if line.startswith("data:") and '"web_search_tool_result"' in line + ) + assert [(result["title"], result["url"]) for result in result_block["content"]] == [ + ("[MASKED] 3.14 docs", "https://docs.[MASKED].org/3.14/"), + ("PEP 745", "https://peps.example/745"), + ] + assert '"id": "srvtoolu_1"' in raw and '"tool_use_id": "srvtoolu_1"' in raw + + @staticmethod + def _last_text_redacting_guardrail() -> CustomGuardrail: + class RedactLastText(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": ["[REDACTED] " + inputs["texts"][-1]]} + + return RedactLastText(guardrail_name="test") + + @pytest.mark.asyncio + async def test_output_response_keeps_the_reply_as_the_last_text_a_guardrail_scans(self): + response = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "python 3.14 release date"}, + }, + self._web_search_result_block(), + {"type": "text", "text": "Python 3.14 shipped"}, + ], + "stop_reason": "end_turn", + } + + await AnthropicMessagesHandler().process_output_response( + response=response, + guardrail_to_apply=self._last_text_redacting_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert response["content"][0]["input"] == {"query": "python 3.14 release date"} + assert response["content"][1] == self._web_search_result_block() + assert response["content"][2] == {"type": "text", "text": "[REDACTED] Python 3.14 shipped"} + + @pytest.mark.asyncio + async def test_deliver_ended_stream_keeps_the_reply_as_the_last_text_a_guardrail_scans(self): + chunks = self._web_search_stream_chunks() + + await AnthropicMessagesHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._last_text_redacting_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + raw = b"".join(chunks).decode() + assert "[REDACTED] Python 3.14 shipped" in raw + assert raw.count("[REDACTED]") == 1 + assert json.dumps(self._web_search_result_block()) in raw + + def test_ended_stream_scan_key_differs_from_a_mid_stream_key_by_the_web_search_results(self): + handler = AnthropicMessagesHandler() + chunks = self._web_search_stream_chunks() + + mid_stream_key = handler.get_streaming_scan_key(chunks[:-2]) + ended_key = handler.get_streaming_scan_key(chunks) + + assert mid_stream_key.texts == ("Python 3.14 shipped",) + assert ended_key != mid_stream_key + assert ended_key.texts[-1] == "Python 3.14 shipped" + assert {"Python 3.14 docs", "https://docs.python.org/3.14/"} <= set(ended_key.texts) + + @pytest.mark.asyncio + async def test_output_response_masks_the_web_search_query_results_and_text(self): + response = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "python 3.14 release date"}, + }, + self._web_search_result_block(), + {"type": "text", "text": "Python 3.14 shipped"}, + ], + "stop_reason": "end_turn", + } + + await AnthropicMessagesHandler().process_output_response( + response=response, + guardrail_to_apply=self._python_masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert "python" not in json.dumps(response).lower() + assert response["content"][0]["input"] == {"query": "[MASKED] 3.14 release date"} + assert [(result["title"], result["url"]) for result in response["content"][1]["content"]] == [ + ("[MASKED] 3.14 docs", "https://docs.[MASKED].org/3.14/"), + ("PEP 745", "https://peps.example/745"), + ] + assert response["content"][2] == {"type": "text", "text": "[MASKED] 3.14 shipped"} + @pytest.mark.asyncio async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() diff --git a/tests/unit/llms/anthropic/messages/native/test_fake_stream_iterator.py b/tests/unit/llms/anthropic/messages/native/test_fake_stream_iterator.py index 233145f2706..ece5bb39225 100644 --- a/tests/unit/llms/anthropic/messages/native/test_fake_stream_iterator.py +++ b/tests/unit/llms/anthropic/messages/native/test_fake_stream_iterator.py @@ -1,6 +1,10 @@ import json -from collections.abc import Iterable, Mapping -from typing import Final +from collections.abc import Iterable, Mapping, Sequence +from typing import Final, cast + +import anthropic +import httpx +import pytest from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( FakeAnthropicMessagesStreamIterator, @@ -91,3 +95,140 @@ def test_message_delta_omits_absent_optional_fields(): "usage": {"input_tokens": 10, "output_tokens": 2}, } assert _message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None} + + +_WEB_SEARCH_RESULT_BLOCK: Final = { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://example.com/a", + "title": "Example A", + "encrypted_content": "abc123", + } + ], +} + + +def _content_events(content: Sequence[Mapping[str, object]]) -> tuple[tuple[str, dict[str, object]], ...]: + response: Final = { + "id": "msg_ws", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": list(content), + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + return tuple( + event for event in _decoded_events(FakeAnthropicMessagesStreamIterator(response=response)) + if event[0].startswith("content_block_") + ) + + +def test_server_tool_use_streams_empty_start_then_input_json_delta(): + """A real Anthropic stream opens server_tool_use with an empty input and + sends the query as an input_json_delta, which SDK accumulators read.""" + caller: Final = {"type": "code_execution_20260120", "tool_id": "srvtoolu_00"} + events: Final = _content_events( + [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "example search"}, + "caller": caller, + } + ] + ) + assert [event_type for event_type, _ in events] == [ + "content_block_start", + "content_block_delta", + "content_block_stop", + ] + assert events[0][1]["content_block"] == { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {}, + "caller": caller, + } + delta: Final = events[1][1]["delta"] + assert isinstance(delta, dict) + assert delta["type"] == "input_json_delta" + assert json.loads(delta["partial_json"]) == {"query": "example search"} + + +def test_web_search_tool_result_rides_whole_in_content_block_start(): + events: Final = _content_events([_WEB_SEARCH_RESULT_BLOCK]) + assert [event_type for event_type, _ in events] == ["content_block_start", "content_block_stop"] + assert events[0][1]["content_block"] == _WEB_SEARCH_RESULT_BLOCK + + +def test_native_search_blocks_keep_indices_aligned_with_text(): + events: Final = _content_events( + [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "q"}}, + _WEB_SEARCH_RESULT_BLOCK, + {"type": "text", "text": "Here is what I found."}, + ] + ) + starts: Final = tuple(payload for event_type, payload in events if event_type == "content_block_start") + stops: Final = tuple(payload["index"] for event_type, payload in events if event_type == "content_block_stop") + assert [(start["index"], cast(dict[str, object], start["content_block"])["type"]) for start in starts] == [ + (0, "server_tool_use"), + (1, "web_search_tool_result"), + (2, "text"), + ] + assert stops == (0, 1, 2) + + +_CALLER: Final = {"type": "code_execution_20260120", "tool_id": "srvtoolu_00"} + +_REPLAYED_BLOCKS: Final = ( + {"type": "thinking", "thinking": "Look up the release page.", "signature": "sig_01"}, + {"type": "redacted_thinking", "data": "EncryptedThinking=="}, + {"type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"city": "Paris"}, "caller": _CALLER}, + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_fetch", "input": {"url": "https://example.com/a"}, "caller": _CALLER}, + _WEB_SEARCH_RESULT_BLOCK, + { + "type": "text", + "text": "Example A is the newest release.", + "citations": [ + { + "type": "web_search_result_location", + "cited_text": "Example A was released today.", + "url": "https://example.com/a", + "title": "Example A", + "encrypted_index": "idx_01", + } + ], + }, +) + + +def _sdk_rebuilt_content(block: Mapping[str, object]) -> list[dict[str, object]]: + response: Final = { + "id": "msg_replay", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [block], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + body: Final = b"".join(FakeAnthropicMessagesStreamIterator(response=response)) + transport: Final = httpx.MockTransport( + lambda _request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"}) + ) + client: Final = anthropic.Anthropic(api_key="sk-ant-test", http_client=httpx.Client(transport=transport)) + with client.messages.stream( + model="claude-sonnet-5-5", max_tokens=64, messages=[{"role": "user", "content": "hi"}] + ) as stream: + return [content_block.model_dump(exclude_none=True) for content_block in stream.get_final_message().content] + + +@pytest.mark.parametrize("block", _REPLAYED_BLOCKS, ids=lambda block: str(block["type"])) +def test_anthropic_sdk_rebuilds_every_block_field_from_the_stream(block: Mapping[str, object]): + assert _sdk_rebuilt_content(block) == [block] diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index c42878bd4f7..da855753b38 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -5,9 +5,12 @@ Uses mock guardrails to validate pipeline execution without external services. """ import copy +import json import logging import pickle -from typing import Literal +from collections.abc import Callable, Iterable, Mapping, Sequence +from itertools import chain +from typing import Final, Literal from unittest.mock import MagicMock import pytest @@ -15,6 +18,7 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( @@ -1726,21 +1730,6 @@ async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up assert chunks == [_chunk()] -@pytest.mark.asyncio -async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(monkeypatch, caplog): - masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} - guardrail = _LegacyHookGuardrail( - replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[masked_tool_call]) - ) - chunks = [_chunk()] - - with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) - - _assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments") - assert chunks == [_chunk()] - - @pytest.mark.asyncio async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls(monkeypatch, caplog): guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[])) @@ -1772,21 +1761,6 @@ async def test_streaming_step_passes_a_tool_only_stream_the_legacy_hook_left_alo assert not any("discarded" in record.getMessage() for record in caplog.records) -@pytest.mark.asyncio -async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only_stream(monkeypatch, caplog): - masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} - guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement(tool_calls=[masked_tool_call])) - chunks = [_tool_only_chunk()] - - with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - result = await _run_legacy_streaming_step( - monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() - ) - - _assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments") - assert chunks == [_tool_only_chunk()] - - class _NoHooksGuardrail(CustomGuardrail): pass @@ -1825,6 +1799,208 @@ async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_r assert chunks == [_chunk()] +_AnthropicHookResponse = dict[str, list[dict[str, object]]] + + +class _AnthropicTextRewritingLegacyGuardrail(_LegacyHookGuardrail): + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: _AnthropicHookResponse + ) -> None: + self.calls.append({"response": copy.deepcopy(response)}) + for block in response["content"]: + if block["type"] == "text" and isinstance(text := block.get("text"), str): + block["text"] = "[REWRITTEN] " + text + + +def _anthropic_web_search_stream() -> list[bytes]: + result_block: Final = { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + {"type": "web_search_result", "url": "https://example.test/a", "title": "A", "encrypted_content": ""} + ], + } + events: Final = ( + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "m", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"query": "a"}'}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": result_block}, + {"type": "content_block_stop", "index": 1}, + {"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 2, "delta": {"type": "text_delta", "text": "hello world"}}, + {"type": "content_block_stop", "index": 2}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}}, + {"type": "message_stop"}, + ) + return _anthropic_sse(events) + + +def _anthropic_sse(events: Iterable[Mapping[str, object]]) -> list[bytes]: + return [f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events] + + +def _anthropic_tool_only_stream() -> list[bytes]: + return _anthropic_sse( + ( + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "m", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {}}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"ssn"'}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": ': "123"}'}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 2}}, + {"type": "message_stop"}, + ) + ) + + +def _streamed_events(chunks: Sequence[bytes]) -> list[dict[str, object]]: + return [ + json.loads(line[len("data: ") :]) + for line in b"".join(chunks).decode().splitlines() + if line.startswith("data: ") + ] + + +def _streamed_block_deltas(chunks: Sequence[bytes], index: int, field: str) -> str: + deltas: Final = ( + event["delta"] + for event in _streamed_events(chunks) + if event.get("index") == index and event["type"] == "content_block_delta" + ) + return "".join(str(delta[field]) for delta in deltas if isinstance(delta, dict)) + + +def _streamed_web_search_results(chunks: Sequence[bytes]) -> list[dict[str, str]]: + blocks: Final = ( + event["content_block"] for event in _streamed_events(chunks) if event["type"] == "content_block_start" + ) + results: Final = ( + block["content"] for block in blocks if isinstance(block, dict) and block["type"] == "web_search_tool_result" + ) + return list(chain.from_iterable(results)) + + +@pytest.mark.asyncio +async def test_streaming_step_delivers_a_legacy_hook_mask_of_web_search_result_titles( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + chunks: Final = _anthropic_web_search_stream() + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result: Final = await _run_legacy_streaming_step( + monkeypatch, _AnthropicMaskingLegacyGuardrail(), chunks, translation=AnthropicMessagesHandler() + ) + + assert [step.outcome for step in result.step_results] == ["pass"] + assert [(item["title"], item["url"]) for item in _streamed_web_search_results(chunks)] == [ + ("[MASKED]", "https://example.test/a") + ] + assert _streamed_block_deltas(chunks, 2, "text") == "[MASKED] world" + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +class _AnthropicMaskingLegacyGuardrail(_LegacyHookGuardrail): + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: _AnthropicHookResponse + ) -> None: + for block in response["content"]: + if block["type"] == "text" and isinstance(text := block.get("text"), str): + block["text"] = text.replace("hello", "[MASKED]") + if block["type"] == "tool_use" and isinstance(tool_input := block.get("input"), dict): + block["input"] = {key: "[MASKED]" for key in tool_input} + if block["type"] == "web_search_tool_result" and isinstance(results := block.get("content"), list): + block["content"] = [{**result, "title": "[MASKED]"} for result in results] + + +@pytest.mark.parametrize( + "stream, masked_key", + [(_anthropic_web_search_stream, "query"), (_anthropic_tool_only_stream, "ssn")], + ids=["web_search_query", "tool_only"], +) +@pytest.mark.asyncio +async def test_streaming_step_delivers_a_legacy_hook_mask_of_a_streamed_tool_input( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + stream: Callable[[], list[bytes]], + masked_key: str, +) -> None: + chunks: Final = stream() + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result: Final = await _run_legacy_streaming_step( + monkeypatch, _AnthropicMaskingLegacyGuardrail(), chunks, translation=AnthropicMessagesHandler() + ) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert json.loads(_streamed_block_deltas(chunks, 0, "partial_json")) == {masked_key: "[MASKED]"} + assert "hello" not in b"".join(chunks).decode() + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_delivers_a_legacy_text_rewrite_on_a_web_search_stream( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + guardrail: Final = _AnthropicTextRewritingLegacyGuardrail() + chunks: Final = _anthropic_web_search_stream() + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result: Final = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=AnthropicMessagesHandler() + ) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + raw: Final = b"".join(chunks).decode() + assert "[REWRITTEN] hello world" in raw + assert '"title": "A"' in raw and '"url": "https://example.test/a"' in raw + assert not any("discarded" in record.getMessage() for record in caplog.records) + + @pytest.mark.asyncio async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch): masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world")) From 56388fd13de12f13e7f49a43e0050ad590baa94d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 00:11:55 -0700 Subject: [PATCH 10/64] fix(count_tokens): count a history that carries a compaction block (#44784) * fix(count_tokens): count a compaction block instead of answering 500 A message history that replays Anthropic's signed compaction block made /v1/messages/count_tokens answer HTTP 500: neither the Anthropic nor the Vertex AI partner count-tokens request carried the compact-2026-09-04 beta, the provider rejected the block with a 400, the proxy fell back to the local tokenizer, and that raised on the unknown block type. Both provider count paths now add the compaction beta whenever the history carries a compaction block, and the local tokenizer counts the block's summary text like a text block, skipping the signature. * fix(token_counter): skip the tokenizer for empty thinking and compaction text * fix(count_tokens): forward context_management on Anthropic counts and count chat compaction blocks locally Vertex AI compaction support for inference and counting moves to its own PR, so the Vertex count handler is back to main here * fix(anthropic): map list-form context_management to edits for count_tokens and inference betas Count requests forwarded the OpenAI list form of context_management unchanged, while inference maps it to Anthropic edits, and the inference beta helper read the list as no edits. Both now map it once. A replayed compaction block with an empty summary no longer reaches the tokenizer. * fix(anthropic): forward count context_management only with edits, map Vertex list-form betas, and return provider context_management * fix(anthropic): drop an unmappable context_management list under drop_params --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/token_counter.py | 43 ++++-- litellm/llms/anthropic/common_utils.py | 74 ++++++++-- .../llms/anthropic/count_tokens/handler.py | 4 +- .../anthropic/count_tokens/token_counter.py | 1 + .../anthropic/count_tokens/transformation.py | 31 ++++- .../messages/native/transformation.py | 26 +--- .../anthropic/messages/transformation.py | 25 +--- .../proxy/anthropic_endpoints/endpoints.py | 18 ++- .../litellm_core_utils/test_token_counter.py | 91 ++++++++++++ .../test_anthropic_messages_transformation.py | 62 ++++++++- ...t_anthropic_count_tokens_transformation.py | 131 +++++++++++++++++- .../llms/anthropic/test_count_tokens_oauth.py | 2 +- .../vertex_ai/count_tokens/test_context.py | 3 +- ...artner_models_anthropic_messages_config.py | 23 +++ .../anthropic_endpoints/test_endpoints.py | 68 ++++++++- 15 files changed, 520 insertions(+), 82 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index ea3e645ffdf..800a0939cec 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -17,6 +17,7 @@ import anyio import anyio.lowlevel import httpx import tiktoken +from pydantic import TypeAdapter from tiktoken.model import MODEL_TO_ENCODING from typing_extensions import ParamSpec, TypeVar @@ -69,6 +70,7 @@ from litellm.types.llms.anthropic import ( AnthropicMessagesToolResultParam, AnthropicMessagesToolUseParam, AnthropicSearchResultParam, + CompactionBlock, ) from litellm.types.llms.openai import ( AllMessageValues, @@ -937,6 +939,8 @@ def _count_messages( default_token_count, params.pdf_pricing, ) + elif key == "provider_specific_fields": + num_tokens += sum(map(params.count_function, _compaction_summaries(value))) elif key == "search_results" and isinstance(value, list): from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_search_results_text, @@ -951,6 +955,25 @@ def _count_messages( return num_tokens +_OBJECT_LIST: Final = TypeAdapter(list[object]) + + +def _compaction_summaries(provider_specific_fields: object) -> tuple[str, ...]: + fields: Final = as_str_mapping(provider_specific_fields) + blocks: Final = fields.get("compaction_blocks") if fields is not None else None + block_fields: Final = ( + tuple(map(as_str_mapping, _OBJECT_LIST.validate_python(blocks))) if isinstance(blocks, list) else () + ) + return tuple( + content + for block in block_fields + if block is not None + and block.get("type") == "compaction" + and isinstance(content := block.get("content"), str) + and content + ) + + def _count_extra( count_function: TokenCounterFunction, tools: list[ChatCompletionToolParam] | None, @@ -1648,6 +1671,11 @@ def _count_anthropic_content( return tokens +_TEXT_FIELD_BY_BLOCK_TYPE: Final = MappingProxyType( + {"text": "text", "thinking": "thinking", "redacted_thinking": "thinking", "compaction": "content"} +) + + def _count_content_list( count_function: TokenCounterFunction, content_list: str @@ -1659,6 +1687,7 @@ def _count_content_list( | AnthropicMessagesImageParam | AnthropicMessagesDocumentParam | AnthropicSearchResultParam + | CompactionBlock ], use_default_image_token_count: bool, default_token_count: int | None, @@ -1670,8 +1699,10 @@ def _count_content_list( for c in content_list: if isinstance(c, str): num_tokens += count_function(c) - elif c["type"] == "text": - num_tokens += count_function(str(c.get("text", ""))) + elif c["type"] in _TEXT_FIELD_BY_BLOCK_TYPE: + block_text = str(c.get(_TEXT_FIELD_BY_BLOCK_TYPE[c["type"]]) or "") + if block_text: + num_tokens += count_function(block_text) elif c["type"] == "image_url": image_url = c.get("image_url") num_tokens += _count_image_tokens(image_url, use_default_image_token_count) @@ -1704,12 +1735,6 @@ def _count_content_list( default_token_count, pdf_pricing, ) - elif c["type"] in ("thinking", "redacted_thinking"): - # Claude extended thinking content block - # Count the thinking text and skip the opaque blobs (signature, redacted data) - thinking_text = str(c.get("thinking", "")) - if thinking_text: - num_tokens += count_function(thinking_text) elif c["type"] == "tool_reference": # Anthropic tool-search reference block: a lightweight pointer to # a deferred tool, e.g. {"type": "tool_reference", "tool_name": ...}. @@ -1727,7 +1752,7 @@ def _count_content_list( f"Invalid content item type: {content_type}. " f"Expected str or dict with 'type' field " f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, " - f"tool_reference, search_result)." + f"compaction, tool_reference, search_result)." ) return num_tokens except Exception as e: diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 0bc519187f4..82977f97854 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -8,12 +8,13 @@ import re from collections.abc import Container, Mapping, MutableMapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone +from itertools import chain from types import MappingProxyType from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar from urllib.parse import quote, urlparse import httpx -from pydantic import BeforeValidator, ConfigDict, Field, StrictBool, TypeAdapter, ValidationError +from pydantic import BeforeValidator, ConfigDict, Field, JsonValue, StrictBool, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -24,6 +25,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, ) from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, is_encrypted_reasoning_block, @@ -46,6 +48,7 @@ from litellm.llms.anthropic.wif import ( from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.anthropic import ( + ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_HOSTED_TOOLS, ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER, ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, @@ -118,6 +121,7 @@ _CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:" OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _DECODED_JSON: Final = TypeAdapter(object) +_JSON_VALUE: Final = TypeAdapter(JsonValue) _CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/") @@ -177,24 +181,68 @@ def is_anthropic_messages_url(url: str) -> bool: return parsed_url.hostname == "api.anthropic.com" or parsed_url.path.removesuffix("/").endswith("/v1/messages") +def _message_blocks(message: object) -> Sequence[object]: + fields: Final = as_str_mapping(message) + blocks: Final = validated_object_list(fields.get("content")) if fields is not None else None + return blocks or () + + +def _is_signed_compaction_block(block: object) -> bool: + fields: Final = as_str_mapping(block) + if fields is None or fields.get("type") != "compaction": + return False + signature: Final = fields.get("signature") + return isinstance(signature, str) and bool(signature) + + +def carries_signed_compaction_block(messages: Sequence[object]) -> bool: + return any(map(_is_signed_compaction_block, chain.from_iterable(map(_message_blocks, messages)))) + + def requires_native_compaction_beta( custom_llm_provider: str, optional_params: Mapping[str, object], messages: Sequence[object], ) -> bool: return custom_llm_provider == "anthropic" and ( - optional_params.get("compaction") is not None - or any( - isinstance(block, Mapping) - and block.get("type") == "compaction" - and isinstance(block.get("signature"), str) - and bool(block.get("signature")) - for message in messages - if isinstance(message, Mapping) - for content in (message.get("content"),) - if isinstance(content, (list, tuple)) - for block in content + optional_params.get("compaction") is not None or carries_signed_compaction_block(messages) + ) + + +def anthropic_context_management(context_management: JsonValue) -> JsonValue: + if not isinstance(context_management, list): + return context_management + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + edits: Final = AnthropicConfig.map_openai_context_management_to_anthropic( + [entry for entry in context_management if isinstance(entry, dict)] + ) + return context_management if edits is None else _JSON_VALUE.validate_python(edits) + + +def _context_management_edits(spec: JsonValue) -> Sequence[object]: + mapping: Final = as_str_mapping(spec) + return (validated_object_list(mapping.get("edits")) if mapping is not None else None) or () + + +def context_management_with_edits(context_management: JsonValue) -> JsonValue: + spec: Final = anthropic_context_management(context_management) + return spec if _context_management_edits(spec) else None + + +def context_management_betas(context_management: JsonValue) -> tuple[str, ...]: + edits: Final = _context_management_edits(anthropic_context_management(context_management)) + edit_types: Final = tuple(edit.get("type") for edit in map(as_str_mapping, edits) if edit is not None) + return tuple( + beta + for beta, used in ( + (ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, "compact_20260112" in edit_types), + ( + ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, + any(edit_type != "compact_20260112" for edit_type in edit_types), + ), ) + if used ) @@ -1367,8 +1415,6 @@ class AnthropicModelInfo(BaseLLMModelInfo): if is_vertex_request is True: # Vertex AI requires web search beta header for web search to work if web_search_tool_used: - from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES - headers["anthropic-beta"] = ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value elif len(all_betas) > 0: headers["anthropic-beta"] = ",".join(all_betas) diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index f8bf7232d9c..7a7d1d776d3 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -39,6 +39,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): tools: list[dict[str, JsonValue]] | None = None, system: JsonValue = None, optional_params: Mapping[str, JsonValue] | None = None, + context_management: JsonValue = None, drop_params: bool = False, ) -> dict[str, JsonValue]: """ @@ -70,6 +71,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): tools=tools, system=system, optional_params=optional_params, + context_management=context_management, drop_params=drop_params, ) @@ -81,7 +83,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): verbose_logger.debug("Making request to: %s", endpoint_url) # Get required headers - headers: Final = self.get_count_tokens_headers(auth_header) + headers: Final = self.get_count_tokens_headers(auth_header, messages, context_management) # Use LiteLLM's async httpx client async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 91fc77773d1..99b1cd7f241 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -132,6 +132,7 @@ class AnthropicTokenCounter(BaseTokenCounter): tools=tools, system=system, optional_params=count_options, + context_management=context_management.spec if context_management is not None else None, drop_params=drops_unsupported_params(litellm_params.get("drop_params")), ) diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index e60ffe67c8d..1c8f0c206cd 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -27,11 +27,15 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ) from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import ( + carries_signed_compaction_block, + context_management_betas, + context_management_with_edits, is_claude_model_target, merge_anthropic_beta_headers, tools_without_custom_field, ) from litellm.llms.anthropic.wif import resolve_anthropic_base +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject, ValidUserMessageContentTypes from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse @@ -376,6 +380,7 @@ class AnthropicCountTokensConfig: tools: list[dict[str, JsonValue]] | None = None, system: JsonValue = None, optional_params: Mapping[str, JsonValue] | None = None, + context_management: JsonValue = None, drop_params: bool = False, ) -> dict[str, JsonValue]: # mutable-ok: provider transport requires JSON dictionaries """ @@ -408,6 +413,7 @@ class AnthropicCountTokensConfig: drop_custom=drop_params, ), ), + ("context_management", context_management_with_edits(context_management)), ) if value is not None } @@ -419,18 +425,33 @@ class AnthropicCountTokensConfig: ) ) - def get_count_tokens_headers(self, auth_header: Mapping[str, str]) -> dict[str, str]: + def get_count_tokens_headers( + self, + auth_header: Mapping[str, str], + messages: Sequence[Mapping[str, JsonValue]], + context_management: JsonValue = None, + ) -> dict[str, str]: """The count-tokens headers around a resolved Anthropic auth header (``AnthropicModelInfo.get_auth_header``): x-api-key for a static key, an Authorization bearer for ``ANTHROPIC_AUTH_TOKEN`` and for sk-ant-oat tokens, whose mandatory oauth beta - merges with the token-counting beta instead of replacing it.""" + merges with the token-counting beta instead of replacing it. A history that carries a + signed compaction block also needs the compaction beta, or Anthropic rejects the block + as an unknown content type, and ``context_management`` edits need the betas inference sends.""" + compaction_betas: Final = ( + (ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value,) + if carries_signed_compaction_block(messages) + else () + ) + request_betas: Final = ( + ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, + *compaction_betas, + *context_management_betas(context_management), + ) return { "Content-Type": "application/json", "anthropic-version": "2023-06-01", **auth_header, - "anthropic-beta": merge_anthropic_beta_headers( - auth_header.get("anthropic-beta"), ANTHROPIC_TOKEN_COUNTING_BETA_VERSION - ), + "anthropic-beta": merge_anthropic_beta_headers(auth_header.get("anthropic-beta"), request_betas), } def validate_request( diff --git a/litellm/llms/anthropic/messages/native/transformation.py b/litellm/llms/anthropic/messages/native/transformation.py index c086618e829..ef855eff180 100644 --- a/litellm/llms/anthropic/messages/native/transformation.py +++ b/litellm/llms/anthropic/messages/native/transformation.py @@ -43,6 +43,7 @@ from ...common_utils import ( OBJECT_LIST_ADAPTER, AnthropicError, AnthropicModelInfo, + context_management_betas, is_claude_model_target, merge_anthropic_beta_headers, optionally_handle_anthropic_oauth, @@ -745,6 +746,8 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ) if transformed_context_management is not None: anthropic_messages_optional_request_params["context_management"] = transformed_context_management + elif drops_unsupported_params(litellm_params.drop_params): + anthropic_messages_optional_request_params.pop("context_management") ####### get required params for all anthropic messages requests ###### # Lazy %s: the f-string previously stringified the entire messages @@ -873,28 +876,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): if requires_native_compaction_beta(custom_llm_provider, optional_params, messages): beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value) - # Check for context management - context_management_param: Final = optional_params.get("context_management") - if context_management_param is not None: - # Check edits array for compact_20260112 type - edits: Final = context_management_param.get("edits", ()) - has_compact = False - has_other = False - - for edit in edits: - edit_type = edit.get("type", "") - if edit_type == "compact_20260112": - has_compact = True - else: - has_other = True - - # Add compact header if any compact edits exist - if has_compact: - beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value) - - # Add context management header if any other edits exist - if has_other: - beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value) + beta_values.update(context_management_betas(optional_params.get("context_management"))) # Check for structured outputs. Anthropic's newer request shape nests # the schema under output_config.format; the older top-level diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py index bfa94396ea7..da82766c2b9 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py @@ -3,7 +3,7 @@ from typing import Final from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.core_helpers import drops_unsupported_params -from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.anthropic.common_utils import AnthropicModelInfo, context_management_betas from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, messages_carry_output_config, @@ -171,28 +171,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert if existing_beta: beta_values.update(b.strip() for b in existing_beta.split(",")) - # Check for context management - context_management_param: Final = optional_params.get("context_management") - if context_management_param is not None: - # Check edits array for compact_20260112 type - edits: Final = context_management_param.get("edits", []) - has_compact = False - has_other = False - - for edit in edits: - edit_type = edit.get("type", "") - if edit_type == "compact_20260112": - has_compact = True - else: - has_other = True - - # Add compact header if any compact edits exist - if has_compact: - beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value) - - # Add context management header if any other edits exist - if has_other: - beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value) + beta_values.update(context_management_betas(optional_params.get("context_management"))) # Check for web search tool for tool in tools: diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1a2e610d5c4..0e204cf9211 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -8,7 +8,7 @@ from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import JSONResponse -from pydantic import JsonValue, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -54,6 +54,7 @@ from litellm.proxy.litellm_pre_call_utils import ( should_auto_drop_params_for_agentic_cli, ) from litellm.router import Router +from litellm.types.utils import TokenCountResponse router: Final = APIRouter() @@ -334,6 +335,15 @@ def _count_tokens_body_options(data: Mapping[str, JsonValue], auto_drop_params: } +_PROVIDER_COUNT_REPLY: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) + + +def _provider_context_management(response: TokenCountResponse) -> Mapping[str, object]: + original: Final = _PROVIDER_COUNT_REPLY.validate_python(response.original_response) + context_management: Final = original.get("context_management") if original is not None else None + return {} if context_management is None else {"context_management": context_management} + + @router.post( "/v1/messages/count_tokens", tags=["[beta] Anthropic Messages Token Counting"], @@ -413,7 +423,11 @@ async def count_tokens( log_llm_api_exception(e, litellm_call_id) return _anthropic_error_json_response(proxy_exception_from_route_error(e), request) else: - return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)} + return { + "input_tokens": token_response.total_tokens, + **_provider_context_management(token_response), + **count_source_fields(token_response), + } @router.post( diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index c2f8e12a726..697ac76b4e5 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -1962,6 +1962,97 @@ def test_token_counter_with_redacted_thinking_content(): assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block) + +def test_token_counter_with_compaction_block(): + """ + A replayed on-demand compaction block (Anthropic's signed summary of the earlier turns) counts its + summary text like a text block and nothing for its signature. It used to raise, which turned every + /v1/messages/count_tokens call on a compacted history into a 500 once the provider count fell back here. + """ + model = "anthropic/claude-sonnet-5-5" + summary = "The user is building a recipe app in Python and asked for one-sentence class descriptions." + follow_up = {"role": "user", "content": [{"type": "text", "text": "Now do the same for Ingredient."}]} + signed_block = {"type": "compaction", "content": summary, "signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40} + + as_compaction = [{"role": "assistant", "content": [signed_block]}, follow_up] + unsigned = [{"role": "assistant", "content": [{"type": "compaction", "content": summary}]}, follow_up] + as_text = [{"role": "assistant", "content": [{"type": "text", "text": summary}]}, follow_up] + + assert token_counter(model=model, messages=as_compaction) == token_counter(model=model, messages=as_text) + assert token_counter(model=model, messages=unsigned) == token_counter(model=model, messages=as_text) + + +def test_token_counter_counts_chat_compaction_blocks_like_their_summary_text(): + """ + Chat clients replay a compaction block in the assistant message's provider_specific_fields, and inference + sends it to Anthropic as a content block, so the local count includes its summary text and skips its signature. + """ + model = "anthropic/claude-sonnet-5-5" + summary = "The user is building a recipe app in Python and asked for one-sentence class descriptions." + follow_up = {"role": "user", "content": "Now do the same for Ingredient."} + signed_block = {"type": "compaction", "content": summary, "signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40} + + replayed = [ + {"role": "assistant", "content": None, "provider_specific_fields": {"compaction_blocks": [signed_block]}}, + follow_up, + ] + as_text = [{"role": "assistant", "content": summary}, follow_up] + + assert token_counter(model=model, messages=replayed) == token_counter(model=model, messages=as_text) + + +class _BosEncoder: + def encode_batch_fast(self, texts: list[str]) -> list[list[int]]: + return [[0] * (len(text.split()) + 1) for text in texts] + + +def test_token_counter_charges_nothing_for_a_chat_compaction_block_with_no_summary(): + """ + A tokenizer that prepends a BOS token charges one token for the empty string, so a replayed + compaction block with an empty summary must not reach it, like an empty thinking block. + """ + model = "anthropic/claude-sonnet-5-5" + tokenizer = {"type": "huggingface_tokenizer", "tokenizer": _BosEncoder()} + follow_up = {"role": "user", "content": "Now do the same for Ingredient."} + replayed = [ + { + "role": "assistant", + "content": None, + "provider_specific_fields": {"compaction_blocks": [{"type": "compaction", "content": ""}]}, + }, + follow_up, + ] + bare = [{"role": "assistant", "content": None}, follow_up] + + assert token_counter(model=model, messages=replayed, custom_tokenizer=tokenizer) == token_counter( + model=model, messages=bare, custom_tokenizer=tokenizer + ) + + +def test_count_content_list_skips_the_tokenizer_for_empty_block_text(): + """ + An empty thinking block, a redacted_thinking block (no thinking text at all) and a compaction + block with no summary carry no tokens, even under a tokenizer that charges for the empty string + (the llama-2 tokenizer prepends its BOS token to every encode call). Counting them through the + tokenizer added one token per block on such models. + """ + from litellm.litellm_core_utils.token_counter import _count_content_list + + def charges_for_empty_text(text: str) -> int: + return len(text.split()) + 1 + + empty_blocks = [ + {"type": "thinking", "thinking": "", "signature": "EqcLCkYICxgCKkCrqu6lP..."}, + {"type": "thinking", "thinking": None, "signature": "EqcLCkYICxgCKkCrqu6lP..."}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy3va3pzix/LafPsn4aDFIT2Xlxh0L5L8rLVyIw"}, + {"type": "compaction", "content": ""}, + {"type": "text", "text": ""}, + ] + + assert _count_content_list(charges_for_empty_text, empty_blocks, False, None) == 0 + assert _count_content_list(charges_for_empty_text, [{"type": "thinking", "thinking": "two words"}], False, None) == 3 + + def test_token_counter_with_tool_reference_block(): """ Regression test: a message containing an Anthropic tool-search diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py index 1c004d8c36e..8fa9d534a8c 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py @@ -24,7 +24,7 @@ from litellm.llms.tencent.messages.transformation import TencentAnthropicMessage from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) -from litellm.types.llms.anthropic import AnthropicFinishReason +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, AnthropicFinishReason from litellm.types.router import GenericLiteLLMParams from tests.unit.llms.anthropic.test_anthropic_wif import ( ScriptedPoster, @@ -321,6 +321,66 @@ def _derived_beta_header(config: AnthropicMessagesConfig) -> dict: ) +_COMPACT_EDIT_BETA: Final = ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value +_CLEARING_EDIT_BETA: Final = ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value + + +@pytest.mark.parametrize( + ("context_management", "betas"), + ( + pytest.param({"edits": [{"type": "compact_20260112"}]}, {_COMPACT_EDIT_BETA}, id="compact"), + pytest.param({"edits": [{"type": "clear_tool_uses_20250919"}]}, {_CLEARING_EDIT_BETA}, id="clearing"), + pytest.param( + {"edits": [{"type": "compact_20260112"}, {"type": "clear_thinking_20251015"}]}, + {_COMPACT_EDIT_BETA, _CLEARING_EDIT_BETA}, + id="both", + ), + pytest.param( + [{"type": "compaction", "compact_threshold": 150000}], {_COMPACT_EDIT_BETA}, id="openai_list_form" + ), + ), +) +def test_context_management_edits_add_the_beta_each_edit_needs( + context_management: JsonValue, betas: set[str], local_beta_headers_config: None +) -> None: + headers: Final = AnthropicMessagesConfig()._update_headers_with_anthropic_beta( + headers={}, optional_params={"context_management": context_management}, messages=[] + ) + + assert set(headers["anthropic-beta"].split(",")) == betas + + +_UNMAPPABLE_LIST_FORM: Final = [{"type": "clear_tool_uses_20250919"}] + + +@pytest.mark.parametrize("config_class", [AnthropicMessagesConfig, VertexAIPartnerModelsAnthropicMessagesConfig]) +@pytest.mark.parametrize( + ("context_management", "drop_params", "forwarded"), + ( + pytest.param(_UNMAPPABLE_LIST_FORM, True, None, id="unmappable_dropped_under_drop_params"), + pytest.param(_UNMAPPABLE_LIST_FORM, False, _UNMAPPABLE_LIST_FORM, id="unmappable_kept_for_the_provider_400"), + pytest.param( + [{"type": "compaction"}], + True, + {"edits": [{"type": "compact_20260112"}]}, + id="mappable_still_mapped_under_drop_params", + ), + ), +) +def test_context_management_the_list_form_cannot_map_is_dropped_only_under_drop_params( + config_class: type[AnthropicMessagesConfig], context_management: JsonValue, drop_params: bool, forwarded: JsonValue +) -> None: + request: Final = config_class().transform_anthropic_messages_request( + model="claude-haiku-5-5", + messages=[{"role": "user", "content": "Hi"}], + anthropic_messages_optional_request_params={"max_tokens": 50, "context_management": context_management}, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + + assert request.get("context_management") == forwarded + + @pytest.mark.parametrize("config", _classifier_configs(), ids=type) def test_routes_whose_platform_runs_the_classifier_keep_safeguards_and_the_beta( config: AnthropicMessagesConfig, local_beta_headers_config: None diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index a3e723eaca2..aa7fa2727b4 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -26,7 +26,24 @@ from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAn from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, ) -from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse +from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES + +COMPACTION_BETA = ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value +COMPACTED_HISTORY = [ + { + "role": "assistant", + "content": [ + { + "type": "compaction", + "content": "The user is building a recipe app and asked for one-sentence class descriptions.", + "signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40, + } + ], + }, + {"role": "user", "content": "Now do the same for Ingredient."}, +] @pytest.mark.parametrize( @@ -890,3 +907,115 @@ async def test_counter_sends_the_structured_output_tool_choice_and_thinking_chat inference["tool_choice"], inference["tools"], ) + + +def test_headers_add_the_compaction_beta_only_for_a_compacted_history(): + config = AnthropicCountTokensConfig() + auth_header = {"x-api-key": "sk-ant-api03-test-key"} + + with_block = config.get_count_tokens_headers(auth_header, COMPACTED_HISTORY)["anthropic-beta"].split(",") + plain = config.get_count_tokens_headers(auth_header, [{"role": "user", "content": "hi"}])["anthropic-beta"] + + assert COMPACTION_BETA in with_block + assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in with_block + assert plain == ANTHROPIC_TOKEN_COUNTING_BETA_VERSION + + +@pytest.mark.asyncio +async def test_handler_sends_the_compaction_beta_with_a_compacted_history(httpx_transport_clients): + with respx.mock: + route = respx.post("https://gateway.example/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 404}) + ) + result = await AnthropicCountTokensHandler().handle_count_tokens_request( + model="claude-sonnet-5-5", + messages=COMPACTED_HISTORY, + auth_header={"x-api-key": "sk-ant-api03-test-key"}, + api_base="https://gateway.example", + ) + + assert result == {"input_tokens": 404} + betas = route.calls.last.request.headers["anthropic-beta"].split(",") + assert COMPACTION_BETA in betas + assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in betas + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@pytest.mark.parametrize( + ("context_management", "forwarded", "sent_beta", "unsent_beta"), + ( + pytest.param( + {"edits": [{"type": "compact_20260112"}]}, + {"edits": [{"type": "compact_20260112"}]}, + ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, + ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, + id="compact", + ), + pytest.param( + {"edits": [{"type": "clear_tool_uses_20250919"}]}, + {"edits": [{"type": "clear_tool_uses_20250919"}]}, + ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, + ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, + id="clearing", + ), + pytest.param( + [{"type": "compaction", "compact_threshold": 150000}], + {"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]}, + ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, + ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, + id="openai_list_form", + ), + ), +) +async def test_counter_forwards_context_management_with_the_beta_inference_sends( + context_management: JsonValue, forwarded: JsonValue, sent_beta: str, unsent_beta: str +) -> None: + """Anthropic counts an unsigned compaction block only under the compact_20260112 edit that inference sends, + and a clearing edit changes the count, so the count carries the request's context_management as inference + sends it, the OpenAI list form mapped to edits, and its beta.""" + with respx.mock: + route: Final = respx.post("https://gateway.example/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 259}) + ) + result: Final = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-opus-5-5", + messages=[{"role": "assistant", "content": [{"type": "compaction", "content": "Recipe app summary."}]}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-ant-api03-test-key", "api_base": "https://gateway.example"}}, + context_management=CountContextManagement(spec=context_management, drop_invalid_edits=False), + ) + + sent: Final = route.calls.last.request + betas: Final = sent.headers["anthropic-beta"].split(",") + assert result is not None and result.total_tokens == 259 + assert TypeAdapter(dict[str, JsonValue]).validate_json(sent.content)["context_management"] == forwarded + assert (sent_beta in betas, unsent_beta in betas) == (True, False) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@pytest.mark.parametrize("context_management", ({}, {"edits": []}, []), ids=("empty_object", "empty_edits", "empty_list")) +async def test_counter_leaves_out_context_management_without_edits(context_management: JsonValue) -> None: + """Anthropic's count endpoint rejects a context_management with no edits, which inference accepts, so the count + leaves it out and adds no context-management beta.""" + with respx.mock: + route: Final = respx.post("https://gateway.example/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 11}) + ) + result: Final = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-opus-5-5", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-ant-api03-test-key", "api_base": "https://gateway.example"}}, + context_management=CountContextManagement(spec=context_management, drop_invalid_edits=False), + ) + + sent: Final = route.calls.last.request + betas: Final = frozenset(sent.headers.get("anthropic-beta", "").split(",")) + assert result is not None and result.total_tokens == 11 + assert "context_management" not in TypeAdapter(dict[str, JsonValue]).validate_json(sent.content) + assert not betas & { + ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, + ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, + } diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py index b8a7c001b37..3e3b408d825 100644 --- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py @@ -48,7 +48,7 @@ FEDERATED_DEPLOYMENT = { def count_tokens_headers_for(api_key: str) -> dict[str, str]: auth_header = AnthropicModelInfo.get_auth_header(api_key=api_key) assert auth_header is not None - return AnthropicCountTokensConfig().get_count_tokens_headers(auth_header) + return AnthropicCountTokensConfig().get_count_tokens_headers(auth_header, [{"role": "user", "content": "hi"}]) @pytest.fixture diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 38eb87debfd..67733884e12 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -666,7 +666,8 @@ async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_whe *_agent_turns(2, first_round=3), ) - view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact)) + with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")): + view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact)) assert isinstance(view, InferenceRequestView) assert _tool_outputs_kept(view) == (True, True) diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 1e7c3d52814..7f16957ddc4 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -526,6 +526,29 @@ def test_both_compact_and_context_management_headers_added(): ) +def test_openai_list_form_context_management_sends_the_compact_beta(): + config: Final = VertexAIPartnerModelsAnthropicMessagesConfig() + + with ( + patch.object(config, "_ensure_access_token", return_value=("token", "test-project")), + patch.object(config, "get_complete_vertex_url", return_value="https://mock-url"), + ): + updated_headers, _ = config.validate_anthropic_messages_environment( + headers={}, + model="claude-vertex-ai-opus-4-6", + messages=[], + optional_params={"context_management": [{"type": "compaction", "compact_threshold": 150000}]}, + litellm_params={ + "vertex_ai_project": "test-project", + "vertex_ai_location": "us-central1", + "vertex_credentials": "{}", + }, + api_base=None, + ) + + assert updated_headers["anthropic-beta"].split(",") == [ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value] + + def test_validate_environment_always_refreshes_token_ignoring_stale_bearer(): """Regression: stale Authorization in shared deployment extra_headers must not skip token refresh on /v1/messages — _ensure_access_token is always called.""" diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index f0317af915f..604fc84a096 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -6,6 +6,7 @@ import json import logging import unittest from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping +from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -587,6 +588,48 @@ class TestErrorLogCarriesCallId: assert call_id in record.getMessage() + +class TestCountTokensCompactedHistory: + """A history carrying Anthropic's signed compaction block used to fail /v1/messages/count_tokens once the count + ran on the local tokenizer, which rejected the block type.""" + + @pytest.mark.asyncio + async def test_count_tokens_counts_a_compaction_block_like_its_summary_text(self): + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + summary = "The user is building a recipe app and asked for one-sentence class descriptions." + follow_up = {"role": "user", "content": "Now do the same for Ingredient."} + compacted = { + "model": "claude-sonnet-5-5", + "messages": [ + { + "role": "assistant", + "content": [{"type": "compaction", "content": summary, "signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40}], + }, + follow_up, + ], + } + as_text = { + "model": "claude-sonnet-5-5", + "messages": [{"role": "assistant", "content": [{"type": "text", "text": summary}]}, follow_up], + } + request = MagicMock() + request.headers = {} + + async def count(body): + with ( + patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "llm_router", None), # test-quality-ok: module global; with no router the count runs on the local tokenizer, the path that raised + ): + return await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + counted = await count(compacted) + + assert counted["input_tokens"] > 0 + assert counted == await count(as_text) + def _claude_router() -> Router: return Router( model_list=[ @@ -759,7 +802,10 @@ async def test_count_tokens_drops_malformed_tools_the_counted_deployment_drops_e def _count_on_claude( - monkeypatch: pytest.MonkeyPatch, peer: Callable[[httpx.Request], httpx.Response], headers: Mapping[str, str] + monkeypatch: pytest.MonkeyPatch, + peer: Callable[[httpx.Request], httpx.Response], + headers: Mapping[str, str], + extra_body: Mapping[str, object] = MappingProxyType({}), ) -> httpx.Response: from fastapi import FastAPI @@ -799,7 +845,7 @@ def _count_on_claude( app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() return TestClient(app).post( "/v1/messages/count_tokens", - json={"model": "count-log-claude", "messages": [{"role": "user", "content": "Hello"}]}, + json={"model": "count-log-claude", "messages": [{"role": "user", "content": "Hello"}], **extra_body}, headers=dict(headers), ) @@ -827,6 +873,24 @@ def test_count_tokens_logs_one_line_with_the_callers_call_id_model_and_count( ) +def test_count_tokens_returns_the_context_management_the_provider_applied(monkeypatch: pytest.MonkeyPatch) -> None: + applied: Final = {"original_input_tokens": 640} + response: Final = _count_on_claude( + monkeypatch, + lambda request: httpx.Response( + 200, json={"input_tokens": 576, "context_management": applied}, request=request + ), + {}, + {"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}}, + ) + + assert response.json() == { + "input_tokens": 576, + "context_management": applied, + "litellm_count_source": "provider", + }, response.text + + def test_count_tokens_provider_rejection_logs_its_status_and_answers_the_same_error( monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] ) -> None: From cfb77ecc3ad72d9f3602dbda22d78d7e84739289 Mon Sep 17 00:00:00 2001 From: mika-lt <349045450+mika-lt@users.noreply.github.com> Date: Fri, 30 Oct 2026 00:14:49 -0700 Subject: [PATCH 11/64] fix(proxy): list only routable models for a wildcard alias over a provider org (#49152) * fix(proxy): list only routable models for a wildcard alias over a provider org A wildcard alias over an org-scoped deployment (together_ai/* over together_ai/nosuchorg/*) listed the provider's whole catalog on /v1/models and /model/info, and every listed id routed to a model the provider 404s. A skipped deployment path segment now counts as a routing prefix only when the routed model resolves to a catalog id through the get_model_info lookup order, so responses/, Bedrock converse/, and region paths list as before. Together's model list now also carries the cost map's together_ai rows, skipping the together-ai-* pricing tiers, so an alias over a real org lists the models Together serves today. * fix(proxy): keep flat catalog ids under routes the cost map can't parse, drop retired Together rows * fix(proxy): check skipped path segments only on Together so Hugging Face provider routes keep listing --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/__init__.py | 5 ++ litellm/constants.py | 2 +- litellm/proxy/auth/model_checks.py | 34 ++++++++- litellm/responses/main.py | 8 +- litellm/utils.py | 46 ++++++------ tests/unit/proxy/auth/test_model_checks.py | 87 +++++++++++++++++++++- 6 files changed, 150 insertions(+), 32 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index d59edcaff89..5f14c605cd0 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -872,6 +872,11 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje # ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params. if "-to-" not in key and "fireworks-ai-default" not in key: fireworks_ai_models.add(key) + elif value.get("litellm_provider") == "together_ai": + if _is_past_deprecation_date(value): + together_ai_models.discard(key) + elif not key.startswith("together-ai-"): + together_ai_models.add(key) elif value.get("litellm_provider") == "baseten": if _is_past_deprecation_date(value): baseten_models.discard(key) diff --git a/litellm/constants.py b/litellm/constants.py index cfee777f36a..e7c9c5dcbfb 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1274,7 +1274,7 @@ empower_models: Final = set( ] ) -together_ai_models: Final[set] = set( +together_ai_models: Final[set[str]] = set( [ # llama llms - chat "togethercomputer/llama-2-70b-chat", diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index c40ee190e12..b21315529eb 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -14,14 +14,16 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.responses.main import strip_responses_routing_prefix from litellm.router import Router from litellm.router_utils.fallback_event_handlers import get_fallback_model_group from litellm.router_utils.pattern_match_deployments import PatternMatchRouter from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params from litellm.types.utils import LlmProviders -from litellm.utils import ProviderConfigManager, get_valid_models +from litellm.utils import ProviderConfigManager, get_valid_models, potential_model_names_in_lookup_order LLM_PROVIDER_PREFIXES: Final = frozenset(llm_provider.value for llm_provider in LlmProviders) +_PROVIDERS_SENDING_THE_DEPLOYMENT_PATH_AS_THE_MODEL_ID: Final = frozenset({LlmProviders.TOGETHER_AI.value}) _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) @@ -331,6 +333,25 @@ def _routes_back_to( ) +def _resolves_to_catalog_model(model: str, provider: str, provider_model_ids: frozenset[str]) -> bool: + provider_id_prefix: Final = f"{provider}/" + resolved_model: Final = ( + f"{provider_id_prefix}{strip_responses_routing_prefix(model.removeprefix(provider_id_prefix))}" + ) + return any( + name.removeprefix(provider_id_prefix) in provider_model_ids + for name in potential_model_names_in_lookup_order(resolved_model, provider) + ) + + +def _routes_to_catalog_model( + routes_to_model: Callable[[str, str], bool], provider: str, provider_model_ids: frozenset[str] +) -> Callable[[str, str], bool]: + return lambda listed_id, model: ( + routes_to_model(listed_id, model) and _resolves_to_catalog_model(model, provider, provider_model_ids) + ) + + def _ids_routing_to_catalog_models( name_pattern: str, deployment_pattern: str, @@ -344,8 +365,15 @@ def _ids_routing_to_catalog_models( return () provider_id_prefix: Final = f"{provider}/" provider_model_ids: Final = tuple(catalog_id.removeprefix(provider_id_prefix) for catalog_id in catalog_ids) + deployment_model_pattern_without_provider: Final = deployment_pattern.removeprefix(provider_id_prefix) catalog_pattern: Final = _without_routing_segments( - deployment_pattern.removeprefix(provider_id_prefix), provider, name_pattern, provider_model_ids + deployment_model_pattern_without_provider, provider, name_pattern, provider_model_ids + ) + routing_prefix: Final = deployment_model_pattern_without_provider.removesuffix(catalog_pattern) + reaches_catalog_model: Final = ( + _routes_to_catalog_model(routes_to_model, provider, frozenset(provider_model_ids)) + if routing_prefix and provider in _PROVIDERS_SENDING_THE_DEPLOYMENT_PATH_AS_THE_MODEL_ID + else routes_to_model ) deployment_model_pattern: Final = f"{provider_id_prefix}{catalog_pattern}" deployment_regex: Final = re.compile(PatternMatchRouter().pattern_to_regex(deployment_model_pattern)) @@ -369,7 +397,7 @@ def _ids_routing_to_catalog_models( for listed_id, catalog_id in listed_ids if listed_id and _routes_back_to( - listed_id, name_regex, (deployment_model_pattern, catalog_id), deployment_pattern, routes_to_model + listed_id, name_regex, (deployment_model_pattern, catalog_id), deployment_pattern, reaches_catalog_model ) ) ) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 98356e9a5d4..063ef9d91ea 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -759,7 +759,7 @@ def _dispatched_model_name( from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata provider_model, _ = resolve_model_and_provider_for_metadata(model=model, custom_llm_provider=custom_llm_provider) - return _strip_responses_routing_prefix(provider_model) + return strip_responses_routing_prefix(provider_model) def _will_bridge_to_chat_completions( @@ -1245,7 +1245,7 @@ def _pop_use_chat_completions_api_kw(kwargs: dict[str, object]) -> bool: _RESPONSES_ROUTING_PREFIX: Final = "responses/" -def _strip_responses_routing_prefix(model: str) -> str: +def strip_responses_routing_prefix(model: str) -> str: if not model.startswith(_RESPONSES_ROUTING_PREFIX): return model return model[len(_RESPONSES_ROUTING_PREFIX) :] @@ -1273,7 +1273,7 @@ def _resolve_model_provider_for_responses( litellm_params.api_key = dynamic_api_key if dynamic_api_base is not None: litellm_params.api_base = dynamic_api_base - return _strip_responses_routing_prefix(provider_model), resolved_provider + return strip_responses_routing_prefix(provider_model), resolved_provider def _apply_managed_file_id_mapping( @@ -2893,7 +2893,7 @@ async def _aresponses_websocket( api_key=api_key, litellm_params=litellm_params, ) - resolved_model: Final = _strip_responses_routing_prefix(provider_model) + resolved_model: Final = strip_responses_routing_prefix(provider_model) litellm_params_dict["data_residency"] = infer_openai_data_residency( _custom_llm_provider, diff --git a/litellm/utils.py b/litellm/utils.py index 41a3fddf3a3..e85e089289a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6708,6 +6708,26 @@ def _first_registered_match( return next(matches, (None, None)) +def _model_cost_lookup_order( + model: str, potential_model_names: PotentialModelNamesAndCustomLLMProvider +) -> tuple[str, ...]: + return ( + potential_model_names["combined_model_name"], + model, + potential_model_names["region_free_combined_model_name"], + potential_model_names["split_model"], + potential_model_names["combined_stripped_model_name"], + potential_model_names["stripped_model_name"], + potential_model_names["provider_prefixed_model_name"], + ) + + +def potential_model_names_in_lookup_order(model: str, custom_llm_provider: str) -> tuple[str, ...]: + return _model_cost_lookup_order( + model, get_potential_model_names(model=model, custom_llm_provider=custom_llm_provider) + ) + + def _get_model_info_from_generalization( model: str, potential_model_names: PotentialModelNamesAndCustomLLMProvider, @@ -6726,15 +6746,7 @@ def _get_model_info_from_generalization( still had a priced exact name to try. O(number of rules); only call after the exact lookups have missed. """ - candidates: Final = ( - potential_model_names["combined_model_name"], - model, - potential_model_names["region_free_combined_model_name"], - potential_model_names["split_model"], - potential_model_names["combined_stripped_model_name"], - potential_model_names["stripped_model_name"], - potential_model_names["provider_prefixed_model_name"], - ) + candidates: Final = _model_cost_lookup_order(model, potential_model_names) if any(get_model_cost_key(candidate) is not None for candidate in candidates): return None for candidate in candidates: @@ -6973,11 +6985,6 @@ def get_model_info_helper( model=model, custom_llm_provider=custom_llm_provider or declared_authenticating_provider(model) ) - combined_model_name: Final = potential_model_names["combined_model_name"] - stripped_model_name: Final = potential_model_names["stripped_model_name"] - combined_stripped_model_name: Final = potential_model_names["combined_stripped_model_name"] - provider_prefixed_model_name: Final = potential_model_names["provider_prefixed_model_name"] - split_model: Final = potential_model_names["split_model"] custom_llm_provider = potential_model_names["custom_llm_provider"] model_cost_custom_llm_provider: Final = custom_llm_provider ######################### @@ -7044,16 +7051,9 @@ def get_model_info_helper( and strip it. Tried last so no model that already resolves through 1-6 can change. """ - lookup_order: Final = ( - combined_model_name, - model, - potential_model_names["region_free_combined_model_name"], - split_model, - combined_stripped_model_name, - stripped_model_name, - provider_prefixed_model_name, + lookup: Final = _first_registered_match( + _model_cost_lookup_order(model, potential_model_names), model_cost_custom_llm_provider ) - lookup: Final = _first_registered_match(lookup_order, model_cost_custom_llm_provider) key: str | None = lookup[0] _model_info: dict[str, Any] | None = lookup[1] diff --git a/tests/unit/proxy/auth/test_model_checks.py b/tests/unit/proxy/auth/test_model_checks.py index 3f0f97e63b0..c67610c24a1 100644 --- a/tests/unit/proxy/auth/test_model_checks.py +++ b/tests/unit/proxy/auth/test_model_checks.py @@ -1057,7 +1057,9 @@ def test_wildcard_lists_each_model_once_under_its_full_owner_id(provider: str): import litellm from litellm.proxy.auth.model_checks import get_known_models_from_wildcard - full_ids: Final = {f"{provider}/{model.removeprefix(f'{provider}/')}" for model in litellm.models_by_provider[provider]} + full_ids: Final = { + f"{provider}/{model.removeprefix(f'{provider}/')}" for model in litellm.models_by_provider[provider] + } assert any(model.count("/") > 1 for model in full_ids) assert sorted(get_known_models_from_wildcard(f"{provider}/*")) == sorted(full_ids) @@ -1372,12 +1374,57 @@ def test_partial_wildcard_lists_catalog_ids_stored_with_the_provider_prefix( assert sorted(listed) == expected +@pytest.mark.parametrize( + ("provider", "routing_segment"), [("bedrock", "converse"), ("bedrock", "us-east-1"), ("openai", "responses")] +) +def test_routing_segment_lists_the_same_catalog_ids_as_the_bare_provider_wildcard( + provider: str, routing_segment: str +) -> None: + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + bare_listing: Final = get_known_models_from_wildcard(f"{provider}/*", LiteLLM_Params(model=f"{provider}/*")) + routed_listing: Final = get_known_models_from_wildcard( + f"{provider}/*", LiteLLM_Params(model=f"{provider}/{routing_segment}/*") + ) + + assert bare_listing + assert sorted(routed_listing) == sorted(bare_listing) + + +def test_provider_wildcard_lists_cost_map_models_but_not_pricing_tiers() -> None: + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + model_id: Final = "together_ai/partial-wildcard-org/model-t" + pricing_tier: Final = "together-ai-partial-wildcard-tier" + with _registered_chat_models({model_id: "together_ai", pricing_tier: "together_ai"}): + listed: Final = get_known_models_from_wildcard("together_ai/*") + + assert model_id in listed + assert [listed_id for listed_id in listed if pricing_tier in listed_id] == [] + + +def test_provider_wildcard_stops_listing_a_cost_map_model_once_it_is_retired() -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + model_id: Final = "together_ai/partial-wildcard-org/model-retired" + with _registered_chat_models({model_id: "together_ai"}): + listed_while_served: Final = get_known_models_from_wildcard("together_ai/*") + litellm.add_known_models({model_id: {"litellm_provider": "together_ai", "deprecation_date": "2000-01-01"}}) + listed_once_retired: Final = get_known_models_from_wildcard("together_ai/*") + + assert model_id in listed_while_served + assert model_id not in listed_once_retired + + _ROUTE_CF_ID: Final = "cloudflare/@cf/partial-route-org/model-a-instruct" _ROUTE_CF_SUFFIXED_ID: Final = "cloudflare/@cf/partial-route-org/model-b-instruct-v2" _ROUTE_HF_ID: Final = "cloudflare/@hf/partial-route-org/model-c" _ROUTE_GEMINI_ID: Final = "gemini-partial-route-9-flash" _ROUTE_OPENAI_ID: Final = "gpt-partial-route-9-mini" _ROUTE_BEDROCK_ID: Final = "anthropic.partial-route-9-v1:0" +_ROUTE_TOGETHER_ID: Final = "together_ai/partial-route-org/model-t" _ROUTE_CF_IDS: Final = {_ROUTE_CF_ID, _ROUTE_CF_SUFFIXED_ID} @@ -1456,6 +1503,29 @@ _ROUTE_CF_IDS: Final = {_ROUTE_CF_ID, _ROUTE_CF_SUFFIXED_ID} set(), id="same-name-namespace-the-catalog-lacks", ), + pytest.param( + "together/*", "together_ai/partial-route-missing-org/*", None, set(), id="alias-over-org-the-catalog-lacks" + ), + pytest.param( + "openai/*", + "openai/chat_completions/*", + None, + {f"openai/chat_completions/{_ROUTE_OPENAI_ID}"}, + id="chat-completions-bridge-route", + ), + pytest.param( + "bedrock/*", "bedrock/llama/*", None, {f"bedrock/llama/{_ROUTE_BEDROCK_ID}"}, id="invoke-format-route" + ), + pytest.param( + "together/*", + "together_ai/partial-route-org/*", + None, + {_ROUTE_TOGETHER_ID}, + id="alias-over-org-a-cost-map-catalog-carries", + ), + pytest.param( + "bedrock/*", "bedrock/us-east-1/*", None, {f"bedrock/us-east-1/{_ROUTE_BEDROCK_ID}"}, id="region-route" + ), pytest.param( "*", "openai/*", "openai/gpt-partial-route-9*", set(), id="provider-prefixed-key-filter-over-catch-all" ), @@ -1501,6 +1571,7 @@ def test_every_listed_partial_wildcard_id_routes_to_the_catalog_model_it_names( _ROUTE_GEMINI_ID: "gemini", _ROUTE_OPENAI_ID: "openai", _ROUTE_BEDROCK_ID: "bedrock", + _ROUTE_TOGETHER_ID: "together_ai", } ): v1_models: Final = get_complete_model_list( @@ -1535,6 +1606,20 @@ def test_every_listed_partial_wildcard_id_routes_to_the_catalog_model_it_names( assert listed_model_info_names == sorted(reachable) +def test_inference_provider_route_lists_every_id_the_plain_provider_wildcard_lists() -> None: + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + plain: Final = get_known_models_from_wildcard("huggingface/*", LiteLLM_Params(model="huggingface/*")) + routed: Final = get_known_models_from_wildcard( + "huggingface/*", LiteLLM_Params(model="huggingface/featherless-ai/*") + ) + + assert plain + assert [model for model in plain if "/" not in model.removeprefix("huggingface/")] == [] + assert sorted(routed) == sorted(plain) + + def test_catch_all_lists_no_id_a_more_specific_sibling_wildcard_routes_elsewhere() -> None: from litellm import Router from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list From e81dd348a425c180c7a632c1eecea7a37682ad96 Mon Sep 17 00:00:00 2001 From: Filip Sinitsin Date: Fri, 30 Oct 2026 07:18:15 +0000 Subject: [PATCH 12/64] fix(proxy)!: log and bill Responses-style bodies sent to /cursor/chat/completions (#30277) * fix(proxy): restore spend logging for /cursor/chat/completions (#30126) The /cursor/chat/completions route processes requests through the Responses API pipeline, but _get_metadata_variable_name did not treat it as a litellm_metadata route. Proxy auth metadata therefore landed in the user-facing metadata param, so the spend tracking callback saw user_api_key=None and silently skipped the SpendLogs write. Add the route to LITELLM_METADATA_ROUTES and cover it with regression tests. * test(proxy): pin provider metadata and key billing on /cursor/chat/completions A Responses-format body on the Cursor route now sends OpenAI only the metadata the client sent and writes a spend row that bills the key. The chat-body branch, which this fix also moves onto litellm_metadata, still bills the key with its team and request tags. The scripted chat reply gets a unique id per test, since a shared id let only the first spend row per database land * fix(proxy): pick the Cursor route's metadata field from the body shape A Responses-style body on /cursor/chat/completions keeps proxy metadata in litellm_metadata, the same as /v1/responses, while a Chat Completions body keeps metadata, the same as /v1/chat/completions. Adding the whole route to the internal-metadata list let a client's chat-body tags replace the key's tags in spend logs and skip tag budgets. * fix(proxy): keep get_metadata_variable_name callable with only the request * fix(proxy): check client tags against tag budgets and read agent limits from litellm_metadata on Responses bodies A Responses-style body on /cursor/chat/completions now keeps its proxy fields in litellm_metadata, so the pre-call tag budget check and the v3 agent rate limiter stopped seeing the client's metadata.tags and the x-litellm-agent-id value. The pre-auth tag check now merges the caller's metadata.tags into the selected metadata field, sharing one helper with request setup, and the limiter reads agent_id from either field * fix(proxy): check client tags sent in a JSON-string metadata against tag budgets on Responses bodies * fix(proxy): answer 422, not 500, when a budget rejects a Responses body whose metadata is a JSON string * fix(proxy): keep key guardrails, metadata type errors and failure-row client details on Responses-style bodies --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 32 ++-- litellm/proxy/auth/auth_checks.py | 10 +- litellm/proxy/common_request_processing.py | 16 +- .../hooks/parallel_request_limiter_v3.py | 3 +- .../proxy/hooks/proxy_track_cost_callback.py | 37 +++- litellm/proxy/litellm_pre_call_utils.py | 88 ++++++++-- .../proxy/response_api_endpoints/endpoints.py | 12 +- litellm/proxy/utils.py | 10 +- .../providers/test_cursor_responses_wire.py | 163 +++++++++++++++--- .../integrations/test_custom_guardrail.py | 14 ++ .../hooks/test_parallel_request_limiter_v3.py | 5 +- .../hooks/test_proxy_track_cost_callback.py | 50 ++++++ .../response_api_endpoints/test_endpoints.py | 2 +- .../proxy/test_common_request_processing.py | 34 ++++ .../unit/proxy/test_litellm_pre_call_utils.py | 68 ++++++++ .../test_post_call_failure_hook.py | 29 ++++ 16 files changed, 482 insertions(+), 91 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..a9fb5c29bad 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -9,6 +9,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args import httpx +from pydantic import TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger @@ -218,6 +219,15 @@ def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = ) +_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) + + +def _guardrails_listed_in(metadata: object) -> list[str | dict[str, DynamicGuardrailParams]]: + if not isinstance(metadata, dict): + return [] + return metadata.get("guardrails") or [] + + class CustomGuardrail(CustomLogger): # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. use_native_during_call_hook: ClassVar[bool] = False @@ -791,28 +801,22 @@ class CustomGuardrail(CustomLogger): return True raise - def get_guardrail_from_metadata(self, data: dict) -> list[str] | list[dict[str, DynamicGuardrailParams]]: + def get_guardrail_from_metadata(self, data: dict) -> list[str | dict[str, DynamicGuardrailParams]]: """ - Returns the guardrail(s) to be run from the metadata or root + Returns the guardrail(s) to be run from the root, else from both metadata fields together """ if "guardrails" in data: return data["guardrails"] - # Check both metadata locations. For regular endpoints move_guardrails_to_metadata - # writes to "metadata"; for thread/assistant endpoints it writes to - # "litellm_metadata". We check the one that actually contains the "guardrails" - # key so that a non-empty litellm_metadata without guardrails does not shadow - # the merged list stored in metadata (which would cause team guardrails to be - # silently skipped while default_on=True policy guardrails still fire). - for meta_key in ("metadata", "litellm_metadata"): - meta = data.get(meta_key) or {} - if isinstance(meta, dict) and "guardrails" in meta: - return meta.get("guardrails") or [] - return [] + request_fields: Final = _REQUEST_FIELDS.validate_python(data) + return [ + *_guardrails_listed_in(request_fields.get("metadata")), + *_guardrails_listed_in(request_fields.get("litellm_metadata")), + ] def _guardrail_is_in_requested_guardrails( self, - requested_guardrails: list[str] | list[dict[str, DynamicGuardrailParams]], + requested_guardrails: Sequence[str | dict[str, DynamicGuardrailParams]], ) -> bool: for _guardrail in requested_guardrails: if isinstance(_guardrail, dict): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1ef5b1072ed..b8d9530d2eb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -446,7 +446,7 @@ def _request_has_trace_id(request_body: Mapping[str, object], request: Request, from litellm.proxy.litellm_pre_call_utils import ( _trace_id_from_traceparent, # pyright: ignore[reportPrivateUsage] # reuse W3C traceparent parser get_chain_id_from_headers, - metadata_variable_name_for_route, + metadata_variable_name_for_request, ) if get_chain_id_from_headers(dict(request.headers)): @@ -460,7 +460,7 @@ def _request_has_trace_id(request_body: Mapping[str, object], request: Request, traceparent: Final[str | None] = request.headers.get("traceparent") if isinstance(traceparent, str) and _trace_id_from_traceparent(traceparent): return True - metadata_variable_name: Final = metadata_variable_name_for_route(get_request_route(request)) + metadata_variable_name: Final = metadata_variable_name_for_request(get_request_route(request), request_body) selected_metadata: Final[object] = request_body.get(metadata_variable_name) if isinstance(selected_metadata, dict): selected_metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(selected_metadata) @@ -1367,11 +1367,15 @@ async def common_checks( if not skip_all_budget_checks: # Key metadata.tags are injected into request_body here so the tag budget # check can read them; this mutation must run before the gathered checks. - from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request if valid_token is not None: for tag_target in LiteLLMProxyRequestSetup.pre_auth_tag_merge_targets(request, request_body): LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=tag_target, route=route) + LiteLLMProxyRequestSetup.merge_caller_body_tags( + request_data=tag_target, + metadata_variable_name=metadata_variable_name_for_request(route, tag_target), + ) LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(request_data=tag_target, user_api_key_dict=valid_token) LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_body) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7cd752c2297..a79d8fe29db 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2401,9 +2401,10 @@ class ProxyBaseLLMRequestProcessing: version=version, proxy_config=proxy_config, ) + request_fields: Final = _REQUEST_FIELDS.validate_python(self.data) + metadata_variable_name: Final = get_metadata_variable_name(request, request_fields) if logging_auth is not None: - metadata_variable_name: Final = get_metadata_variable_name(request) - request_metadata: Final = self.data.get(metadata_variable_name) + request_metadata: Final = request_fields.get(metadata_variable_name) if isinstance(request_metadata, dict): self.data[metadata_variable_name] = { **TypeAdapter(dict[str, object]).validate_python(request_metadata), @@ -2431,12 +2432,11 @@ class ProxyBaseLLMRequestProcessing: # Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved if queue_time_seconds is not None: - _metadata_variable_name: Final = get_metadata_variable_name(request) - if _metadata_variable_name not in self.data: - self.data[_metadata_variable_name] = {} - if not isinstance(self.data[_metadata_variable_name], dict): - self.data[_metadata_variable_name] = {} - self.data[_metadata_variable_name]["queue_time_seconds"] = queue_time_seconds + if metadata_variable_name not in self.data: + self.data[metadata_variable_name] = {} + if not isinstance(self.data[metadata_variable_name], dict): + self.data[metadata_variable_name] = {} + self.data[metadata_variable_name]["queue_time_seconds"] = queue_time_seconds if isinstance(model, str): reject_url_valued_destination("model", model) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 14b08a09116..15282679622 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -3312,7 +3312,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if key_agent_id: return key_agent_id metadata: Final = data.get("metadata") or {} - return metadata.get("agent_id") + litellm_metadata: Final = data.get("litellm_metadata") or {} + return metadata.get("agent_id") or litellm_metadata.get("agent_id") def _get_session_id_from_data(self, data: dict) -> str | None: """Extract session_id from request metadata or litellm_session_id.""" diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 7d8b2e8329c..9c4a4bc775b 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -37,7 +37,7 @@ from litellm.proxy.db.db_spend_update_writer import ( get_llm_router, ) from litellm.proxy.hooks.search_free_tier import apply_search_free_tier -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope from litellm.proxy.spend_tracking.spend_event import ( @@ -137,19 +137,34 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( _OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) -def _proxy_stamped_used_client_oauth_token( - request_data: Mapping[str, object], request_route: str | None -) -> bool | None: +def _proxy_stamped_metadata(request_data: Mapping[str, object], request_route: str | None) -> Mapping[str, object]: proxy_bucket: Final = ( get_metadata_variable_name_from_kwargs(request_data) if request_route is None - else metadata_variable_name_for_route(request_route) + else metadata_variable_name_for_request(request_route, request_data) ) - proxy_metadata: Final = request_data.get(proxy_bucket) - stamped: Final = proxy_metadata.get("used_client_oauth_token") if isinstance(proxy_metadata, dict) else None + return _str_keyed_fields(request_data.get(proxy_bucket)) + + +def _proxy_stamped_used_client_oauth_token( + request_data: Mapping[str, object], request_route: str | None +) -> bool | None: + stamped: Final = _proxy_stamped_metadata(request_data, request_route).get("used_client_oauth_token") return stamped if isinstance(stamped, bool) else None +_PROXY_STAMPED_REQUEST_FIELDS: Final = ("requester_ip_address", "user_agent") + + +def _proxy_stamped_request_fields( + request_data: Mapping[str, object], request_route: str | None +) -> Mapping[str, object]: + if request_route is None: + return {} + proxy_metadata: Final = _proxy_stamped_metadata(request_data, request_route) + return {field: proxy_metadata[field] for field in _PROXY_STAMPED_REQUEST_FIELDS if field in proxy_metadata} + + def _str_keyed_fields(value: object) -> Mapping[str, object]: return value if is_str_keyed_mapping(value) else {} @@ -455,7 +470,13 @@ class _ProxyDBLogger(CustomLogger): _metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_fields, request_route) existing_metadata: Final[dict] = request_data.get("metadata", None) or {} - existing_metadata.update({**_routed_metadata_fallback(request_fields), **_metadata}) + existing_metadata.update( + { + **_routed_metadata_fallback(request_fields), + **_proxy_stamped_request_fields(request_fields, request_route), + **_metadata, + } + ) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") litellm_metadata: Final[ObjectMapping] = ( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index a87fc815402..c1e541d20a9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -379,6 +379,8 @@ LITELLM_METADATA_ROUTES: Final = ( "interactions", ) +CURSOR_CHAT_COMPLETIONS_ROUTE: Final = "/cursor/chat/completions" + LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset( { "mask_input", @@ -636,7 +638,7 @@ def _has_tag_budget(tag: "LiteLLM_TagTable | None") -> bool: return tag.litellm_budget_table.max_budget is not None -def _normalized_metadata_slot( +def normalized_metadata_slot( request_data: MutableMapping[str, object], metadata_variable_name: str ) -> dict[str, object]: """Return the request's metadata slot as a dict, normalising it in place first. @@ -653,6 +655,12 @@ def _normalized_metadata_slot( return normalized +def _holds_metadata_object(value: object) -> bool: + if value is None or isinstance(value, dict): + return True + return isinstance(value, str) and isinstance(_DECODED_JSON.validate_python(safe_json_loads(value)), dict) + + def _metadata_object_or_error(field: str, value: object) -> Mapping[str, object] | ProxyException: if value is None: return {} @@ -863,7 +871,7 @@ def _strip_router_reserved_metadata( ) -def get_metadata_variable_name(request: Request) -> str: +def get_metadata_variable_name(request: Request, data: Mapping[str, object] | None = None) -> str: """ Helper to return what the "metadata" field should be called in the request data @@ -874,7 +882,10 @@ def get_metadata_variable_name(request: Request) -> str: # Inline imports — auth_utils/route_checks participate in a proxy import cycle. from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 - return metadata_variable_name_for_route(get_request_route(request)) + route: Final = get_request_route(request) + if data is None: + return metadata_variable_name_for_route(route) + return metadata_variable_name_for_request(route, data) _get_metadata_variable_name: Final = get_metadata_variable_name @@ -890,6 +901,25 @@ def metadata_variable_name_for_route(route: str) -> Literal["metadata", "litellm return "metadata" +def is_chat_completions_body(data: Mapping[str, object]) -> bool: + messages: Final = data.get("messages") + if isinstance(messages, list) and messages: + return True + return "messages" in data and "input" not in data + + +def is_cursor_responses_request(route: str, body: Mapping[str, object]) -> bool: + return CURSOR_CHAT_COMPLETIONS_ROUTE in route and not is_chat_completions_body(body) + + +def metadata_variable_name_for_request( + route: str, body: Mapping[str, object] +) -> Literal["metadata", "litellm_metadata"]: + if is_cursor_responses_request(route, body): + return "litellm_metadata" + return metadata_variable_name_for_route(route) + + def _promoted_trace_control_fields( requester_metadata: Mapping[str, object], litellm_metadata: Mapping[str, object], @@ -2309,7 +2339,9 @@ class LiteLLMProxyRequestSetup: ``_tag_max_budget_check`` resolve to ``litellm_metadata``, leaving header tags invisible to per-tag budget enforcement. """ - if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES): + if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES) or is_cursor_responses_request( + route, request_data + ): request_data.setdefault("litellm_metadata", {}) @staticmethod @@ -2380,7 +2412,7 @@ class LiteLLMProxyRequestSetup: return tags: Final = _STRING_TAGS_ADAPTER.validate_python(raw_tags) slot: Final = metadata_key if metadata_key is not None else get_metadata_variable_name_from_kwargs(request_data) - metadata: Final = _normalized_metadata_slot(request_data, slot) + metadata: Final = normalized_metadata_slot(request_data, slot) existing: Final = metadata.get("tags") if existing and not isinstance(existing, list): return @@ -2393,6 +2425,34 @@ class LiteLLMProxyRequestSetup: tags=tuple(tag for tag in tags if tag not in sent_tags) ) + @staticmethod + def merge_caller_body_tags( + request_data: MutableMapping[str, object], # mutable-ok: the request body the tag merges write into + metadata_variable_name: str | None = None, + ) -> None: + slot: Final = ( + metadata_variable_name + if metadata_variable_name is not None + else get_metadata_variable_name_from_kwargs(request_data) + ) + if slot == "metadata" or not _holds_metadata_object(request_data.get(slot)): + return + raw_caller_metadata: Final = request_data.get("metadata") + caller_metadata: Final = ( + _DECODED_JSON.validate_python(safe_json_loads(raw_caller_metadata)) + if isinstance(raw_caller_metadata, str) + else raw_caller_metadata + ) + caller_tags: Final = caller_metadata.get("tags") if isinstance(caller_metadata, dict) else None + if not isinstance(caller_tags, list) or not caller_tags: + return + metadata: Final = normalized_metadata_slot(request_data, slot) + existing_tags: Final = metadata.get("tags") + metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags( + request_tags=existing_tags if isinstance(existing_tags, list) else None, + tags_to_add=caller_tags, + ) + @staticmethod def apply_key_tags_pre_auth( request_data: dict, @@ -2408,7 +2468,7 @@ class LiteLLMProxyRequestSetup: return _metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data) - metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name) + metadata: Final = normalized_metadata_slot(request_data, _metadata_variable_name) existing_tags: Final = metadata.get("tags") metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags( @@ -2460,7 +2520,7 @@ class LiteLLMProxyRequestSetup: # from (litellm_metadata vs metadata) so the merged tags are visible # to _tag_max_budget_check. _metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data) - metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name) + metadata: Final = normalized_metadata_slot(request_data, _metadata_variable_name) existing_tags: Final = metadata.get("tags") metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags( @@ -2637,7 +2697,7 @@ async def add_litellm_data_to_request( } safe_add_api_version_from_query_params(data, request) - _metadata_variable_name: Final = get_metadata_variable_name(request) + _metadata_variable_name: Final = get_metadata_variable_name(request, _REQUEST_FIELDS.validate_python(data)) if data.get(_metadata_variable_name, None) is None: data[_metadata_variable_name] = {} @@ -2974,17 +3034,7 @@ async def add_litellm_data_to_request( tags_to_add=tags, ) - _caller_body_metadata: Final = data.get("metadata") if _metadata_variable_name != "metadata" else None - _caller_body_tags: Final = ( - _caller_body_metadata.get("tags") - if isinstance(_caller_body_metadata, dict) and isinstance(_caller_body_metadata.get("tags"), list) - else None - ) - if _caller_body_tags: - data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup.merge_tags( # rebind-ok: matches file idiom - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=_caller_body_tags, - ) + LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data, metadata_variable_name=_metadata_variable_name) # caller_tags: the tags the request's headers and body carry (x-litellm-tags, body "tags", and body # "metadata.tags" / "litellm_metadata.tags" read before the mapped tags are added below), diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 9d0d89f681e..ef59deb2e5e 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -79,6 +79,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry from litellm.proxy.litellm_pre_call_utils import ( _CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list _UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list + is_chat_completions_body, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( request_tags_from_request_data, @@ -213,13 +214,6 @@ def _normalize_tool_dialect( return {**data, **{key: value for key, value in replaceable if key in data}} -def _is_chat_completions_body(data: Mapping[str, object]) -> bool: - messages: Final = data.get("messages") - if isinstance(messages, list) and messages: - return True - return "messages" in data and "input" not in data - - _CURSOR_THINKING_SEPARATOR: Final = "-thinking-" _CURSOR_FAST_SUFFIX: Final = "-fast" _CURSOR_THINKING_LEVELS: Final[frozenset[str]] = frozenset(get_args(REASONING_EFFORT)) @@ -261,7 +255,7 @@ def _resolve_cursor_model_variant( resolved: Final = {**data, "model": variant.base_model} if variant.reasoning_effort is None: return resolved - if _is_chat_completions_body(data): + if is_chat_completions_body(data): if "reasoning_effort" in data: return resolved return {**resolved, "reasoning_effort": variant.reasoning_effort} @@ -646,7 +640,7 @@ async def cursor_chat_completions( raw_body: Final = await read_request_body(request=request) - if _is_chat_completions_body(raw_body): + if is_chat_completions_body(raw_body): # Genuine chat completions body (Cursor sends these for models whose BYOK it # already fixed); delegate so behavior matches /chat/completions exactly. # Keyed on messages CONTENT, not key presence: Cursor can send a null or diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9e125b95a68..b1055bb9124 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -250,7 +250,11 @@ from litellm.proxy.hooks.sensitive_data_routing import ( # noqa: F401, RUF100 PROXY_SensitiveDataRoutingHandler, _PROXY_SensitiveDataRoutingHandler, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export ) -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata +from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + add_guardrails_from_auth_metadata, + normalized_metadata_slot, +) from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.policy_engine.pipeline_executor import ( PipelineExecutor, @@ -3973,9 +3977,7 @@ class ProxyLogging: ) request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP}) - if "metadata" not in request_data: - request_data["metadata"] = {} - request_data["metadata"].update(user_api_key_logged_metadata) + normalized_metadata_slot(request_data, "metadata").update(user_api_key_logged_metadata) if litellm_logging_obj is not None: ## UPDATE LOGGING INPUT diff --git a/tests/integration/providers/test_cursor_responses_wire.py b/tests/integration/providers/test_cursor_responses_wire.py index d77d35f2ff8..1bd79b5b575 100644 --- a/tests/integration/providers/test_cursor_responses_wire.py +++ b/tests/integration/providers/test_cursor_responses_wire.py @@ -3,11 +3,14 @@ from __future__ import annotations import json import os import uuid +from collections.abc import Mapping +from hashlib import sha256 from importlib.resources import files from typing import Final +import httpx import pytest -from integration._support.client import Gateway, Scenario +from integration._support.client import Gateway, Scenario, eventually, list_value, object_value from integration._support.responses_vendor import same_response from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.worker_readiness import wait_for_worker_readiness @@ -38,6 +41,7 @@ _CHAT_SUFFIXES: Final[tuple[tuple[str, str], ...]] = (("", ""),) + tuple( (f"-thinking-{level}", level) for level in _CHAT_THINKING_LEVELS ) _PROMPT: Final = "Inspect the requested change." +_REQUEST_TAG: Final = "cursor-integration" _INPUT: Final = [{"role": "user", "content": [{"type": "input_text", "text": _PROMPT}]}] _CUSTOM_TOOL: Final = { "type": "custom", @@ -107,16 +111,30 @@ def _stream(identity: str) -> tuple[bytes, ...]: ) + (b"data: [DONE]\n\n",) -def _chat_completion(stream: bool) -> bytes | tuple[bytes, ...]: +def _chat_completion(stream: bool, identity: str) -> bytes | tuple[bytes, ...]: if stream: - return ( - b'data: {"id":"chat_cursor_1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{"content":"Patch inspected."},"finish_reason":null}]}\n\n', - b'data: {"id":"chat_cursor_1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n', - b"data: [DONE]\n\n", - ) + return tuple( + f"data: {json.dumps(chunk)}\n\n".encode() + for chunk in ( + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _MODEL, + "choices": [{"index": 0, "delta": {"content": "Patch inspected."}, "finish_reason": None}], + }, + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _MODEL, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + ) + ) + (b"data: [DONE]\n\n",) return json.dumps( { - "id": "chat_cursor_1", + "id": identity, "object": "chat.completion", "created": 1, "model": _MODEL, @@ -133,11 +151,12 @@ def _chat_completion(stream: bool) -> bytes | tuple[bytes, ...]: ).encode() -def _deployment(scenario: Scenario, wire: Wire) -> str: +def _deployment(scenario: Scenario, wire: Wire, **pricing: JsonValue) -> str: model: Final = scenario.model( model=f"openai/{_MODEL}", api_key=_API_KEY, api_base=f"{wire.url}/v1", + **pricing, ) wait_for_worker_readiness(scenario.gateway, model, 1, _PROXY_WORKERS) return model @@ -167,7 +186,7 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output( assert request.method == "POST" assert request.headers["authorization"] == f"Bearer {_API_KEY}", dict(request.headers) if request.target == "/v1/chat/completions": - completion: Final = _chat_completion(stream) + completion: Final = _chat_completion(stream, raw_id) return Reply( content_type="text/event-stream" if stream else "application/json", chunks=completion if isinstance(completion, tuple) else None, @@ -175,7 +194,7 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output( ) assert request.target == "/v1/responses", request.target body: Final = _JSON_OBJECT.validate_json(request.body) - assert {key: value for key, value in body.items() if key != "metadata"} == { + assert body == { "model": _MODEL, "input": _INPUT, **reasoning, @@ -259,41 +278,69 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output( assert [request.target for request in drain_contract_requests(wire)] == ["/v1/responses"] -def test_cursor_responses_body_carries_no_proxy_metadata(gateway: Gateway) -> None: - pytest.skip( - "BUG: /cursor/chat/completions forwards proxy-internal metadata (key hash, team, user, client ip, user agent) " - "to the provider in the Responses body" - ) +@pytest.mark.parametrize( + "client_metadata", ({}, {"cursor_conversation": "integration"}), ids=("no_client_metadata", "client_metadata") +) +@pytest.mark.parametrize("stream", (False, True)) +def test_cursor_responses_body_sends_only_client_metadata_and_bills_the_key( + gateway: Gateway, + stream: bool, + client_metadata: Mapping[str, JsonValue], +) -> None: raw_id: Final = f"resp_cursor_metadata_{uuid.uuid4().hex}" + alias: Final = f"integration-cursor-{uuid.uuid4().hex}" + sent_metadata: Final = {"metadata": client_metadata} if client_metadata else {} def respond(request: Request) -> Reply: if request.method == "GET" and request.target == "/v1/models": return model_discovery_reply(_MODEL) assert request.method == "POST" and request.target == "/v1/responses", request.target + if stream: + return Reply(content_type="text/event-stream", chunks=_stream(raw_id)) return Reply(body=json.dumps(_response(raw_id)).encode()) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = _deployment(scenario, wire) + key: Final = scenario.key(key_alias=alias) + hashed_key: Final = sha256(key.encode()).hexdigest() response: Final = gateway.request( "POST", "/cursor/chat/completions", { - "model": f"{model}-thinking-high", + "model": model, "input": _INPUT, "messages": [], "tools": [_CUSTOM_TOOL], + "stream": stream, + **sent_metadata, }, + key=key, ) assert response.status_code == 200, response.text requests: Final = drain_contract_requests(wire) assert len(requests) == 1, requests - body: Final = _JSON_OBJECT.validate_json(requests[0].body) - assert body == { + assert _JSON_OBJECT.validate_json(requests[0].body) == { "model": _MODEL, "input": _INPUT, - "reasoning": {"effort": "high"}, "tools": [_CUSTOM_TOOL], + "stream": stream, + **sent_metadata, }, requests[0].body + (row,) = eventually( + lambda: gateway.rows( + 'SELECT call_type, spend, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (hashed_key,) + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert row["call_type"] == "aresponses", row + assert object_value(row["metadata"])["user_api_key_alias"] == alias, row + assert isinstance(row["spend"], float) and row["spend"] > 0, row + eventually( + lambda: gateway.rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed_key,)), + lambda tokens: len(tokens) == 1 and tokens[0]["spend"] == pytest.approx(row["spend"]), + seconds=70, + ) def test_cursor_stream_reports_usage_when_the_client_asks_for_it(gateway: Gateway) -> None: @@ -341,10 +388,12 @@ def test_cursor_messages_body_uses_chat_completions_wire( suffix: str, reasoning_effort: str, ) -> None: + chat_id: Final = f"chat_cursor_{uuid.uuid4().hex}" + def respond(request: Request) -> Reply: if request.method == "GET" and request.target == "/v1/models": return model_discovery_reply(_MODEL) - completion: Final = _chat_completion(stream) + completion: Final = _chat_completion(stream, chat_id) return Reply( content_type="text/event-stream" if stream else "application/json", chunks=completion if isinstance(completion, tuple) else None, @@ -367,6 +416,7 @@ def test_cursor_messages_body_uses_chat_completions_wire( "stream": stream, }, key=key, + headers={"x-litellm-tags": _REQUEST_TAG}, ) assert response.status_code == 200, response.text if stream: @@ -393,7 +443,7 @@ def test_cursor_messages_body_uses_chat_completions_wire( payload: Final = _JSON_OBJECT.validate_json(response.content) identity: Final = payload.get("id") created: Final = payload.get("created") - assert isinstance(identity, str) and same_response(identity, "chat_cursor_1"), response.text + assert isinstance(identity, str) and same_response(identity, chat_id), response.text assert isinstance(created, int) and created > 0, response.text assert payload == { "id": identity, @@ -456,3 +506,72 @@ def test_cursor_messages_body_uses_chat_completions_wire( **({"stream_options": {"include_usage": True}} if stream else {}), **({"reasoning_effort": reasoning_effort} if reasoning_effort else {}), }, requests[0].body + (row,) = eventually( + lambda: gateway.rows( + 'SELECT call_type, team_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (sha256(key.encode()).hexdigest(),), + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert (row["call_type"], row["team_id"]) == ("acompletion", team), row + assert _REQUEST_TAG in list_value(row["request_tags"]), row + + +@pytest.mark.parametrize( + ("body_shape", "metadata_as_json_string"), (("messages", False), ("input", False), ("input", True)) +) +def test_cursor_body_keeps_client_and_key_tags_and_enforces_the_client_tag_budget( + gateway: Gateway, body_shape: str, metadata_as_json_string: bool +) -> None: + chat_id: Final = f"chat_cursor_{uuid.uuid4().hex}" + response_id: Final = f"resp_cursor_tags_{uuid.uuid4().hex}" + budget_tag: Final = f"cursor-tag-budget-{uuid.uuid4().hex}" + key_tag: Final = f"cursor-key-tag-{uuid.uuid4().hex}" + client_metadata: Final = {"tags": [budget_tag]} + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.target == "/v1/responses": + return Reply(body=json.dumps(_response(response_id)).encode()) + completion: Final = _chat_completion(False, chat_id) + return Reply(body=completion if isinstance(completion, bytes) else b"") + + def delete_tag() -> None: + gateway.post("/tag/delete", {"name": budget_tag}) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, input_cost_per_token=0.01, output_cost_per_token=0.01) + key: Final = scenario.key(models=[model], metadata={"tags": [key_tag]}) + gateway.post("/tag/new", {"name": budget_tag, "max_budget": 0.0001}) + scenario.cleanups.callback(delete_tag) + + def tagged_chat(text: str) -> httpx.Response: + return gateway.request( + "POST", + "/cursor/chat/completions", + { + "model": model, + body_shape: [{"role": "user", "content": text}], + "metadata": json.dumps(client_metadata) if metadata_as_json_string else client_metadata, + }, + key=key, + ) + + first: Final = tagged_chat("first tagged chat") + assert first.status_code == 200, first.text + (row,) = eventually( + lambda: gateway.rows( + 'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (sha256(key.encode()).hexdigest(),), + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert {budget_tag, key_tag} <= set(list_value(row["request_tags"])), row + blocked: Final = eventually( + lambda: tagged_chat("tag budget probe"), lambda response: response.status_code != 200, seconds=70 + ) + assert blocked.status_code == 422, blocked.text + assert "budget" in blocked.text.lower(), blocked.text diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 614ff97553a..0483888533e 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3699,6 +3699,20 @@ def test_get_guardrail_from_metadata() -> None: ) == [{"test-guardrail": {"extra_body": {"key": "value"}}}] + +@pytest.mark.parametrize( + "caller_metadata", + [{"guardrails": []}, {"guardrails": ["not-a-guardrail"]}, {"guardrails": None}], +) +def test_caller_metadata_guardrails_never_hide_the_proxy_guardrails_in_litellm_metadata(caller_metadata: dict) -> None: + guardrail: Final = CustomGuardrail( + guardrail_name="key-guardrail", event_hook=GuardrailEventHooks.pre_call, default_on=False + ) + request_data: Final = {"metadata": caller_metadata, "litellm_metadata": {"guardrails": ["key-guardrail"]}} + + assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is True + + def test_guardrail_is_in_requested_guardrails() -> None: guardrail: Final = CustomGuardrail(guardrail_name="test-guardrail") diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index b66e7f3d091..3aad2bda9bf 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -2591,7 +2591,8 @@ async def test_agent_session_rate_limit_skipped_without_session_id(): @pytest.mark.asyncio -async def test_agent_rate_limit_from_metadata_agent_id(): +@pytest.mark.parametrize("metadata_field", ("metadata", "litellm_metadata")) +async def test_agent_rate_limit_from_metadata_agent_id(metadata_field): """ Test that agent rate limits work when agent_id comes from request metadata (header) rather than from the API key. @@ -2639,7 +2640,7 @@ async def test_agent_rate_limit_from_metadata_agent_id(): cache=local_cache, data={ "model": "gpt-4", - "metadata": {"agent_id": _agent_id}, + metadata_field: {"agent_id": _agent_id}, }, call_type="", ) diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index dc46b5faf4e..8d4f8cd87ad 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -623,6 +623,56 @@ async def test_async_post_call_failure_hook_reads_used_client_oauth_token_from_t assert json.loads(payload["metadata"])["used_client_oauth_token"] is expected + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_route, body", + [ + ( + "/cursor/chat/completions", + { + "input": "Hello", + "metadata": {"requester_ip_address": "203.0.113.9", "user_agent": "spoofed"}, + "litellm_metadata": {"requester_ip_address": "127.0.0.1", "user_agent": "Cursor/1.0"}, + }, + ), + ( + "/v1/chat/completions", + { + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"requester_ip_address": "127.0.0.1", "user_agent": "Cursor/1.0"}, + "litellm_metadata": {"requester_ip_address": "203.0.113.9", "user_agent": "spoofed"}, + }, + ), + ], +) +async def test_async_post_call_failure_hook_logs_the_proxy_stamped_ip_and_user_agent(request_route: str, body: dict): + request_data = { + "model": "gpt-4o", + "proxy_server_request": {"request_id": "test_request_id"}, + **body, + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await ProxyDBLogger().async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("invalid api key"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key", request_route=request_route), + ) + + payload = get_logging_payload( + kwargs=mock_update_database.call_args[1]["kwargs"], + response_obj={}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert payload["requester_ip_address"] == "127.0.0.1" + assert json.loads(payload["metadata"])["user_agent"] == "Cursor/1.0" + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_bills_guardrail_cost_on_blocked_request(): """LIT-5651: a request blocked by a guardrail never reaches the LLM, but the diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index a783ac70881..1e4dded801e 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3938,7 +3938,7 @@ class TestCursorInputArmFlattening: class TestChatCompletionsBodyDetection: def test_routing_matrix(self): - from litellm.proxy.response_api_endpoints.endpoints import _is_chat_completions_body + from litellm.proxy.litellm_pre_call_utils import is_chat_completions_body as _is_chat_completions_body assert _is_chat_completions_body({"messages": [{"role": "user", "content": "hi"}]}) is True assert _is_chat_completions_body({"messages": [{"role": "user", "content": "hi"}], "input": []}) is True diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index c867db2fd02..2d08609c93b 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -9711,6 +9711,40 @@ async def test_saved_weights_override_caller_input_and_preserve_key_precedence( assert logging.pre_call_hook.call_args.kwargs["data"].get("_router_weights") == expected +def _hook_data(**kwargs: object) -> object: + return kwargs["data"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prompt, proxy_field, client_field", [ + ({"input": "hi"}, "litellm_metadata", "metadata"), + ({"messages": [{"role": "user", "content": "hi"}]}, "metadata", "litellm_metadata"), +]) +async def test_cursor_route_writes_proxy_fields_only_to_the_field_its_body_shape_picks( + prompt: dict[str, str | list[dict[str, str]]], + proxy_field: str, + client_field: str, +) -> None: + from litellm.proxy import proxy_server + + logging: Final = MagicMock(spec=ProxyLogging) + logging.pre_call_hook = AsyncMock(side_effect=_hook_data) + processor: Final = ProxyBaseLLMRequestProcessing(data={"model": "gpt-x", "metadata": {"client": "cursor"}, **prompt}) + data, _ = await processor.common_processing_pre_call_logic( + request=Request({"type": "http", "method": "POST", "path": "/cursor/chat/completions", "headers": []}), + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash"), + proxy_logging_obj=logging, + proxy_config=proxy_server.ProxyConfig(), + route_type="acompletion", + logging_auth=ProxyUserAPIKeyAuth(api_key="logging-hash", key_alias="logging-alias"), + ) + assert data[proxy_field]["user_api_key_alias"] == "logging-alias" + assert "queue_time_seconds" in data[proxy_field] + assert "user_api_key_alias" not in data.get(client_field, {}) + assert "queue_time_seconds" not in data.get(client_field, {}) + assert data["metadata"]["client"] == "cursor" + class TestPerRequestModelGroupAlias: """``router_settings.model_group_alias`` on a key or team has to be resolved by the proxy: the Router resolves aliases from its own shared instance diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 1ce61a8e9b3..2edb032b048 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -121,6 +121,23 @@ class TestGetMetadataVariableName: request = self._make_request("/chat/completions") assert get_metadata_variable_name(request) == "metadata" + @pytest.mark.parametrize( + "body, expected", + [ + ({"messages": [{"role": "user", "content": "hi"}]}, "metadata"), + ({"messages": []}, "metadata"), + ({"input": [{"role": "user", "content": "hi"}]}, "litellm_metadata"), + ({"input": "hi", "messages": []}, "litellm_metadata"), + ], + ) + def test_cursor_route_picks_the_metadata_field_by_body_shape(self, body, expected): + request = self._make_request("/cursor/chat/completions") + assert get_metadata_variable_name(request, body) == expected + + def test_cursor_route_without_a_body_keeps_the_route_answer(self): + request = self._make_request("/cursor/chat/completions") + assert get_metadata_variable_name(request) == "metadata" + def test_returns_metadata_for_completions(self): request = self._make_request("/v1/completions") assert get_metadata_variable_name(request) == "metadata" @@ -681,6 +698,45 @@ async def test_add_litellm_data_to_request_rejects_non_object_metadata(field, va assert field not in data +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/responses", "/cursor/chat/completions"]) +@pytest.mark.parametrize("value, received_type", [("junk", "a string"), (5, "an integer"), ([1], "an array")]) +async def test_invalid_litellm_metadata_answers_400_when_metadata_carries_tags(route, value, received_type): + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.utils import ProxyLogging + + data = {"model": "gpt-4o", "input": "hi", "metadata": {"tags": ["client-tag"]}, "litellm_metadata": value} + request = _batches_request_mock() + await common_checks( + request_body=data, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(api_key="hashed-key"), + request=request, + ) + + with pytest.raises(ProxyException) as exc_info: + await add_litellm_data_to_request( + data=data, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert exc_info.value.code == "400" + assert exc_info.value.message == ( + f"Invalid type for 'litellm_metadata': expected an object, but got {received_type} instead." + ) + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_removes_every_invalid_metadata_field_before_raising(): """When both fields are invalid, the raise for the first must not leave the @@ -11098,3 +11154,15 @@ def test_mapped_baggage_tags_keep_non_string_client_tags_as_sent(): data: Final[dict[str, object]] = {"metadata": {"tags": [1, {"team": "x"}, "client-tag", "client-tag"]}} LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, data) assert data["metadata"]["tags"] == [1, {"team": "x"}, "client-tag", "client-tag", "tenant:a"] + + +@pytest.mark.parametrize("caller_metadata", ({"tags": ["client-tag"]}, json.dumps({"tags": ["client-tag"]}))) +def test_caller_tags_reach_litellm_metadata_whether_metadata_is_a_dict_or_a_json_string(caller_metadata): + data: Final[dict[str, object]] = { + "input": "hi", + "metadata": caller_metadata, + "litellm_metadata": {"tags": ["key-tag"]}, + } + LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data) + assert set(data["litellm_metadata"]["tags"]) == {"key-tag", "client-tag"} + assert data["metadata"] == caller_metadata diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index f880d3548f7..4386d5aec94 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -4,6 +4,7 @@ and ``_handle_logging_proxy_only_error``.""" from __future__ import annotations import asyncio +import json from collections.abc import Callable from datetime import datetime from types import MappingProxyType @@ -1192,3 +1193,31 @@ async def test_handle_logging_proxy_only_path_flags_requests_that_died_before_re blocked_by_guardrail.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is REJECTED_BEFORE_REQUEST_SETUP, blocked_by_guardrail["litellm_call_id"], ) == (True, "client-supplied", False, "call-1") + + +@pytest.mark.asyncio +async def test_auth_failure_on_a_responses_body_with_a_json_string_metadata_is_logged_against_the_key( + proxy_logging, monkeypatch +): + from litellm.proxy import proxy_server + + recorded: list[dict] = [] + + class _RecordingLogger(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + recorded.append(kwargs) + + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "callbacks", [_RecordingLogger()]) + proxy_logging.alert_types = [] + + await proxy_logging.post_call_failure_hook( + request_data={"model": "internal-model", "input": "hi", "metadata": json.dumps({"tags": ["checkout"]})}, + original_exception=litellm.BudgetExceededError(current_cost=1.0, max_budget=0.5), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-hashed", key_alias="checkout-key", request_route="/v1/responses"), + error_type=ProxyErrorTypes.auth_error, + route="/v1/responses", + ) + + assert len(recorded) == 1 + assert recorded[0]["standard_logging_object"]["metadata"]["user_api_key_alias"] == "checkout-key" From 73aaf9e3daec8bb174accdf56ecade90b34c27c8 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 01:16:16 -0700 Subject: [PATCH 13/64] fix(mistral): forward prompt_cache_key to the Mistral API (#47432) * fix(mistral): forward prompt_cache_key to the Mistral API * fix(mistral): send prompt_cache_key only to Mistral's API and accept reasoning_effort on devstral and code-agent models * fix(mistral): keep the upstream api_base out of the prompt_cache_key rejection message * fix(mistral): honor allowed_openai_params for prompt_cache_key and gate codestral async and env hosts * test(mistral): annotate the new adapter prompt_cache_key tests * fix(mistral): drop or reject a prompt_cache_key the Mistral API cannot take A non-ASCII prompt_cache_key gets a 500 from api.mistral.ai, and a non-string one a 422, so both now drop under drop_params and answer 400 without it, like any value the provider cannot take. The /v1/messages adapter derives no key from a non-ASCII Claude Code session for Mistral, and the 400 names the caller's provider, so a codestral deployment's error says codestral. * fix(mistral): send a reasoning_effort object as its level and keep UnsupportedParamsError on sync calls * fix(mistral): keep async codestral file uploads working by not rewriting the caller's file blocks --- .../litellm_core_utils/get_litellm_params.py | 1 + .../anthropic/messages/adapters/handler.py | 2 + .../messages/adapters/transformation.py | 30 +- litellm/llms/mistral/chat/transformation.py | 153 ++++++- litellm/llms/openai/openai.py | 5 +- litellm/main.py | 1 + litellm/types/llms/mistral.py | 19 + .../adapters/test_handler_prompt_cache_key.py | 68 ++++ .../test_mistral_chat_transformation.py | 372 +++++++++++++++--- 9 files changed, 576 insertions(+), 75 deletions(-) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 8792fd0f58a..29f96d61976 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -85,6 +85,7 @@ OPTIONAL_KWARGS_KEYS: Final = ( "github_copilot_user_session", "fireworks_forward_user_id", "responses_api_version", + "allowed_openai_params", PROVIDER_AFFINITY_HEADER_KWARG_KEY, } ) diff --git a/litellm/llms/anthropic/messages/adapters/handler.py b/litellm/llms/anthropic/messages/adapters/handler.py index 704bec9056f..b3b84f721c0 100644 --- a/litellm/llms/anthropic/messages/adapters/handler.py +++ b/litellm/llms/anthropic/messages/adapters/handler.py @@ -647,12 +647,14 @@ class LiteLLMMessagesToCompletionTransformationHandler: request_data["output_config"] = extra_kwargs["output_config"] custom_llm_provider: Final = extra_kwargs.get("custom_llm_provider") + api_base: Final = extra_kwargs.get("api_base") ( openai_request, tool_name_mapping, ) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping( request_data, custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + api_base=api_base if isinstance(api_base, str) else None, drop_params=drops_unsupported_params(extra_kwargs.get("drop_params")), ) diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index ddda759384f..8ea6293bbe6 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -24,6 +24,7 @@ OPENAI_MAX_TOOL_NAME_LENGTH: Final = 64 TOOL_NAME_HASH_LENGTH: Final = 8 TOOL_NAME_PREFIX_LENGTH: Final = OPENAI_MAX_TOOL_NAME_LENGTH - TOOL_NAME_HASH_LENGTH - 1 # 55 PROVIDERS_PROXYING_AN_UNKNOWN_BACKEND: Final = frozenset({"litellm_proxy"}) +PROVIDERS_SERVED_BY_THE_MISTRAL_CHAT_CONFIG: Final = frozenset({"mistral", "codestral"}) _COMPACTION_BLOCK: Final = TypeAdapter(dict[str, JsonValue]) @@ -255,7 +256,12 @@ class AnthropicAdapter: return result def translate_completion_input_params_with_tool_mapping( - self, kwargs, *, custom_llm_provider: str | None = None, drop_params: bool = True + self, + kwargs, + *, + custom_llm_provider: str | None = None, + api_base: str | None = None, + drop_params: bool = True, ) -> tuple[ChatCompletionRequest | None, dict[str, str]]: """ Translate Anthropic request params to OpenAI format, returning tool name mapping. @@ -289,6 +295,7 @@ class AnthropicAdapter: ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( anthropic_message_request=request_body, custom_llm_provider=custom_llm_provider, + api_base=api_base, drop_params=drop_params, ) @@ -1001,7 +1008,9 @@ class LiteLLMAnthropicMessagesAdapter: ) @staticmethod - def _supports_prompt_cache_key(model: str | None, custom_llm_provider: str | None) -> bool: + def _supports_prompt_cache_key( + prompt_cache_key: str, model: str | None, custom_llm_provider: str | None, api_base: str | None + ) -> bool: if not model or not custom_llm_provider: return False if custom_llm_provider in PROVIDERS_PROXYING_AN_UNKNOWN_BACKEND: @@ -1009,7 +1018,17 @@ class LiteLLMAnthropicMessagesAdapter: supported_params: Final = litellm.get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider ) - return "prompt_cache_key" in (supported_params or ()) + if "prompt_cache_key" not in (supported_params or ()): + return False + if custom_llm_provider not in PROVIDERS_SERVED_BY_THE_MISTRAL_CHAT_CONFIG: + return True + _, _, _, resolved_api_base = litellm.get_llm_provider( + model=model, custom_llm_provider=custom_llm_provider, api_base=api_base + ) + mistral_config: Final = litellm.MistralConfig() + return mistral_config.takes_prompt_cache_key_value( + prompt_cache_key + ) and mistral_config.accepts_prompt_cache_key_at(resolved_api_base) @staticmethod def _target_declares_reasoning_effort(model: str, custom_llm_provider: str | None) -> bool: @@ -1042,6 +1061,7 @@ class LiteLLMAnthropicMessagesAdapter: new_kwargs: ChatCompletionRequest, *, custom_llm_provider: str | None = None, + api_base: str | None = None, ) -> None: """Translate metadata fields from Anthropic request to OpenAI request.""" if "metadata" in anthropic_message_request: @@ -1050,7 +1070,7 @@ class LiteLLMAnthropicMessagesAdapter: new_kwargs["user"] = metadata["user_id"] prompt_cache_key: Final = prompt_cache_key_from_user_id(metadata["user_id"]) if prompt_cache_key is not None and self._supports_prompt_cache_key( - anthropic_message_request.get("model"), custom_llm_provider + prompt_cache_key, anthropic_message_request.get("model"), custom_llm_provider, api_base ): new_kwargs["prompt_cache_key"] = prompt_cache_key @@ -1244,6 +1264,7 @@ class LiteLLMAnthropicMessagesAdapter: anthropic_message_request: AnthropicMessagesRequest, *, custom_llm_provider: str | None = None, + api_base: str | None = None, preserve_midturn_system: bool = False, drop_params: bool = True, ) -> tuple[ChatCompletionRequest, dict[str, str]]: @@ -1282,6 +1303,7 @@ class LiteLLMAnthropicMessagesAdapter: anthropic_message_request=anthropic_message_request, new_kwargs=new_kwargs, custom_llm_provider=custom_llm_provider, + api_base=api_base, ) ## CONVERT TOOL CHOICE self._translate_tool_choice_to_openai( diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 2fb9402ed03..15e348453e0 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -9,11 +9,16 @@ Docs - https://docs.mistral.ai/api/ from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from itertools import accumulate, count from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload +from urllib.parse import urlparse import httpx +from pydantic import TypeAdapter from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 +import litellm from litellm._logging import verbose_logger +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, @@ -30,7 +35,15 @@ from litellm.router_utils.reasoning_effort_capability import ( nearest_declared_reasoning_effort, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.mistral import MistralTextBlock, MistralThinkingBlock, MistralToolCallMessage +from litellm.types.llms.mistral import ( + MistralCacheKeyParams, + MistralFilePart, + MistralReasoningEffortParams, + MistralRequestRouting, + MistralTextBlock, + MistralThinkingBlock, + MistralToolCallMessage, +) from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionTextObject from litellm.types.utils import ModelResponse, ModelResponseStream from litellm.utils import convert_to_model_response_object, supports_reasoning @@ -57,6 +70,17 @@ def _accepted_reasoning_effort(model: str, requested: str, custom_llm_provider: _NO_TOOL_CALL_TURN: Final[tuple[int, frozenset[str]]] = (-1, frozenset()) +_MISTRAL_REQUEST_ROUTING: Final = TypeAdapter(MistralRequestRouting) +_MISTRAL_CACHE_KEY_PARAMS: Final = TypeAdapter(MistralCacheKeyParams) +_MISTRAL_REASONING_EFFORT_PARAMS: Final = TypeAdapter(MistralReasoningEffortParams) +_MISTRAL_FILE_PART: Final = TypeAdapter(MistralFilePart) + + +def _reasoning_level_and_summary(reasoning_effort: object) -> tuple[object, object]: + requested: Final = as_str_mapping(reasoning_effort) + if requested is None: + return reasoning_effort, None + return requested.get("effort"), requested.get("summary") def _is_json_array(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a decoded JSON array is a list @@ -88,6 +112,8 @@ class MistralConfig(OpenAIGPTConfig): - `response_format` (object or null): An object specifying the format that the model must output. Setting to { "type": "json_object" } enables JSON mode, which guarantees the message the model generates is in JSON. When using JSON mode you MUST also instruct the model to produce JSON yourself with a system or a user message. """ + _is_base_class = False + temperature: int | None = None top_p: int | None = None max_tokens: int | None = None @@ -131,6 +157,7 @@ class MistralConfig(OpenAIGPTConfig): def get_supported_openai_params(self, model: str) -> list[str]: is_magistral: Final = "magistral" in model.lower() accepts_reasoning_effort: Final = self._accepts_reasoning(model) + accepts_prompt_cache_key: Final = self.custom_llm_provider == "mistral" return [ "stream", "temperature", @@ -145,6 +172,7 @@ class MistralConfig(OpenAIGPTConfig): "parallel_tool_calls", *(("thinking",) if is_magistral else ()), *(("reasoning_effort",) if accepts_reasoning_effort else ()), + *(("prompt_cache_key",) if accepts_prompt_cache_key else ()), ] def _map_tool_choice(self, tool_choice: str) -> str: @@ -211,15 +239,84 @@ class MistralConfig(OpenAIGPTConfig): optional_params["extra_body"] = {"random_seed": value} if param == "response_format": optional_params["response_format"] = value - if param == "reasoning_effort" and "magistral" not in model.lower(): - optional_params["reasoning_effort"] = _accepted_reasoning_effort(model, value, self.custom_llm_provider) + if ( + param == "reasoning_effort" + and "magistral" not in model.lower() + and ( + reasoning_level := self._reasoning_level_to_send( + model, _MISTRAL_REASONING_EFFORT_PARAMS.validate_python(non_default_params), drop_params + ) + ) + is not None + ): + optional_params["reasoning_effort"] = reasoning_level if param in ("reasoning_effort", "thinking") and "magistral" in model.lower(): # Flag that we need to add reasoning system prompt optional_params["_add_reasoning_prompt"] = True - if param == "parallel_tool_calls": - optional_params["parallel_tool_calls"] = value + if param in ("parallel_tool_calls", "prompt_cache_key"): + optional_params[param] = value return optional_params + def _reasoning_level_to_send( + self, model: str, requested: MistralReasoningEffortParams, drop_params: bool + ) -> str | None: + level, summary = _reasoning_level_and_summary(requested.get("reasoning_effort")) + sendable: Final = ( + _accepted_reasoning_effort(model, level, self.custom_llm_provider) if isinstance(level, str) else None + ) + if sendable is not None and summary is None: + return sendable + if drop_params or litellm.drop_params is True: + return sendable + raise UnsupportedParamsError( + message=( + f"{self.custom_llm_provider} does not support parameters: ['reasoning_effort'], for model={model}, " + "since the Mistral API takes it only as a level such as 'high', with no reasoning summary. " + "To drop it, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n" + ), + llm_provider=self.custom_llm_provider, + model=model, + ) + + def accepts_prompt_cache_key_at(self, api_base: str | None) -> bool: + resolved_api_base, _ = self._get_openai_compatible_provider_info(api_base, None) + hostname: Final = urlparse(resolved_api_base).hostname or "" + return hostname == "mistral.ai" or hostname.endswith(".mistral.ai") + + @staticmethod + def takes_prompt_cache_key_value(prompt_cache_key: object) -> bool: + return isinstance(prompt_cache_key, str) and prompt_cache_key.isascii() + + def _prompt_cache_key_rejection(self, prompt_cache_key: object, api_base: str | None) -> str | None: + if not self.accepts_prompt_cache_key_at(api_base): + return "since only the Mistral API (*.mistral.ai) takes it" + if not self.takes_prompt_cache_key_value(prompt_cache_key): + return "since the Mistral API takes it only as a string of ASCII characters" + return None + + def _prompt_cache_key_withheld(self, model: str, prompt_cache_key: object, routing: MistralRequestRouting) -> bool: + requested_api_base: Final = routing.get("api_base") + api_base: Final = requested_api_base if isinstance(requested_api_base, str) else None + allowed_openai_params: Final = routing.get("allowed_openai_params") + if isinstance(allowed_openai_params, (list, tuple)) and "prompt_cache_key" in allowed_openai_params: + return False + rejection: Final = self._prompt_cache_key_rejection(prompt_cache_key, api_base) + if rejection is None: + return False + if litellm.drop_params is True or routing.get("drop_params") is True: + return True + requested_provider: Final = routing.get("custom_llm_provider") + provider: Final = requested_provider if isinstance(requested_provider, str) else self.custom_llm_provider + raise UnsupportedParamsError( + message=( + f"{provider} does not support parameters: ['prompt_cache_key'], for model={model}, {rejection}. " + "To send it anyway, pass `allowed_openai_params=['prompt_cache_key']`. To drop it, set " + "`litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n" + ), + llm_provider=provider, + model=model, + ) + def _get_openai_compatible_provider_info(self, api_base: str | None, api_key: str | None) -> tuple[str, str | None]: # mistral is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.mistral.ai api_base = ( @@ -359,24 +456,31 @@ class MistralConfig(OpenAIGPTConfig): messages = self._handle_message_with_file(messages) return messages + @staticmethod + def _with_mistral_file_id(block: Mapping[str, object]) -> Mapping[str, object]: + file_part: Final = block.get("file") if block.get("type") == "file" else None + file_id: Final = ( + _MISTRAL_FILE_PART.validate_python(file_part).get("file_id") if isinstance(file_part, Mapping) else None + ) + if not file_id: + return block + return {**{key: value for key, value in block.items() if key != "file"}, "file_id": file_id} + + @classmethod + def _with_mistral_file_ids(cls, message: AllMessageValues) -> AllMessageValues: + content: Final = message.get("content") + if not isinstance(content, list): + return message + blocks: Final[tuple[Mapping[str, object], ...]] = tuple(content) + return cast( # cast-ok: only file blocks change, to the file_id shape the Mistral API takes + AllMessageValues, {**message, "content": [cls._with_mistral_file_id(block) for block in blocks]} + ) + def _handle_message_with_file(self, messages: list[AllMessageValues]) -> list[AllMessageValues]: """ Mistral API supports only 'file_id' in message content with type 'file'. """ - for m in messages: - _content_block = m.get("content") - if _content_block and isinstance(_content_block, list): - if any(c.get("type") == "file" for c in _content_block): - # If file content is present, we get file_id from 'file' attribute of content block - # then replace 'file' with 'file_id' and assign the value of 'file_id' attribute to it. - file_contents = [c for c in _content_block if c.get("type") == "file"] - for file_content in file_contents: - file_id = file_content.get("file", {}).get("file_id") - if file_id: - # Replace 'file' with 'file_id' - file_content["file_id"] = file_id # pyright: ignore[reportGeneralTypeIssues] # legacy in-place rewrite of the block shape - file_content.pop("file", None) - return messages + return [self._with_mistral_file_ids(m) for m in messages] def _add_reasoning_system_prompt_if_needed( self, messages: list[AllMessageValues], optional_params: dict @@ -794,7 +898,16 @@ class MistralConfig(OpenAIGPTConfig): if "magistral" in model.lower() and optional_params.get("_add_reasoning_prompt", False): messages = self._add_reasoning_system_prompt_if_needed(messages, optional_params) - upstream_params: Final = {key: value for key, value in optional_params.items() if key != "client_metadata"} + cache_key_params: Final = _MISTRAL_CACHE_KEY_PARAMS.validate_python(optional_params) + withheld_params: Final = ( + ("client_metadata", "prompt_cache_key") + if "prompt_cache_key" in cache_key_params + and self._prompt_cache_key_withheld( + model, cache_key_params["prompt_cache_key"], _MISTRAL_REQUEST_ROUTING.validate_python(litellm_params) + ) + else ("client_metadata",) + ) + upstream_params: Final = {key: value for key, value in optional_params.items() if key not in withheld_params} # Call parent transform_request which handles _transform_messages return super().transform_request( diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index e0a23d6c99f..bef99fd0719 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -27,6 +27,7 @@ import litellm from litellm import LlmProviders from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RETRIES +from litellm.exceptions import UnsupportedParamsError from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -971,8 +972,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): litellm.remove_index_from_tool_calls(messages=messages) else: raise e - except OpenAIError as e: - raise e + except (OpenAIError, UnsupportedParamsError): + raise except Exception as e: status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) diff --git a/litellm/main.py b/litellm/main.py index d469aa629fe..8d6e07d6a8f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6199,6 +6199,7 @@ def completion( PROVIDER_AFFINITY_HEADER_KWARG_KEY, "fireworks_forward_user_id", "responses_api_version", + "allowed_openai_params", ) if key in kwargs }, diff --git a/litellm/types/llms/mistral.py b/litellm/types/llms/mistral.py index cd61f32f430..68a9d26827b 100644 --- a/litellm/types/llms/mistral.py +++ b/litellm/types/llms/mistral.py @@ -23,3 +23,22 @@ class MistralThinkingBlock(TypedDict): type: Literal["thinking"] thinking: list[MistralTextBlock] closed: NotRequired[ReadOnly[bool]] + + +class MistralRequestRouting(TypedDict, total=False): + api_base: ReadOnly[object] + drop_params: ReadOnly[object] + allowed_openai_params: ReadOnly[object] + custom_llm_provider: ReadOnly[object] + + +class MistralCacheKeyParams(TypedDict, total=False): + prompt_cache_key: ReadOnly[object] + + +class MistralReasoningEffortParams(TypedDict, total=False): + reasoning_effort: ReadOnly[object] + + +class MistralFilePart(TypedDict, total=False): + file_id: ReadOnly[object] diff --git a/tests/unit/llms/anthropic/messages/adapters/test_handler_prompt_cache_key.py b/tests/unit/llms/anthropic/messages/adapters/test_handler_prompt_cache_key.py index ed6cd6134ff..77bdacc938f 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_handler_prompt_cache_key.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_handler_prompt_cache_key.py @@ -1,6 +1,7 @@ import json import os import sys +from typing import Final import pytest @@ -70,3 +71,70 @@ def test_prepare_completion_kwargs_keeps_prompt_cache_key_through_responses_rero ) assert completion_kwargs["model"] == "openai/responses/gpt-5.6-luna" assert completion_kwargs["prompt_cache_key"] == "session-abc" + + +@pytest.mark.parametrize( + ("api_base", "azure_api_base_env", "derives_key"), + [ + (None, None, True), + ("https://api.mistral.ai/v1", None, True), + ("https://my-resource.openai.azure.com/openai/v1", None, False), + (None, "https://my-resource.services.ai.azure.com/openai", False), + ], +) +def test_prepare_completion_kwargs_derives_prompt_cache_key_only_for_a_mistral_hosted_endpoint( + monkeypatch: pytest.MonkeyPatch, api_base: str | None, azure_api_base_env: str | None, derives_key: bool +) -> None: + if azure_api_base_env is None: + monkeypatch.delenv("MISTRAL_AZURE_API_BASE", raising=False) + else: + monkeypatch.setenv("MISTRAL_AZURE_API_BASE", azure_api_base_env) + + completion_kwargs: Final = _prepare( + "mistral/mistral-small-latest", {"custom_llm_provider": "mistral", "api_base": api_base} + ) + + assert completion_kwargs["user"] == CLAUDE_CODE_USER_ID + assert ("prompt_cache_key" in completion_kwargs) is derives_key + + +@pytest.mark.parametrize( + ("codestral_api_base_env", "derives_key"), + [(None, True), ("https://my-resource.services.ai.azure.com/openai/v1", False)], +) +def test_prepare_completion_kwargs_derives_prompt_cache_key_for_codestral_only_on_its_mistral_host( + monkeypatch: pytest.MonkeyPatch, codestral_api_base_env: str | None, derives_key: bool +) -> None: + monkeypatch.delenv("MISTRAL_AZURE_API_BASE", raising=False) + if codestral_api_base_env is None: + monkeypatch.delenv("CODESTRAL_API_BASE", raising=False) + else: + monkeypatch.setenv("CODESTRAL_API_BASE", codestral_api_base_env) + + completion_kwargs: Final = _prepare("codestral/codestral-latest", {"custom_llm_provider": "codestral"}) + + assert ("prompt_cache_key" in completion_kwargs) is derives_key + + +def test_prepare_completion_kwargs_derives_no_prompt_cache_key_mistral_cannot_take_from_a_non_ascii_session() -> None: + non_ascii_user_id: Final = json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": "sitzung-ü"}) + + mistral_kwargs: Final = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="mistral/mistral-small-latest", + metadata={"user_id": non_ascii_user_id}, + thinking=None, + extra_kwargs={"custom_llm_provider": "mistral"}, + )[0] + openai_kwargs: Final = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="openai/gpt-5.6-luna", + metadata={"user_id": non_ascii_user_id}, + thinking=None, + extra_kwargs={"custom_llm_provider": "openai"}, + )[0] + + assert "prompt_cache_key" not in mistral_kwargs + assert openai_kwargs["prompt_cache_key"] == "sitzung-ü" diff --git a/tests/unit/llms/mistral/test_mistral_chat_transformation.py b/tests/unit/llms/mistral/test_mistral_chat_transformation.py index b76e1707805..c340f77c317 100644 --- a/tests/unit/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/unit/llms/mistral/test_mistral_chat_transformation.py @@ -1,7 +1,13 @@ +import json +from collections.abc import Mapping from typing import Final, List, cast from unittest.mock import MagicMock, patch +from urllib.parse import urlparse +import httpx import pytest +import respx +from openai import AsyncOpenAI from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY from litellm.llms.mistral.chat.transformation import ( @@ -45,21 +51,15 @@ class TestMistralReasoningSupport: mistral_config = MistralConfig() # Test magistral model supports reasoning parameters - supported_params = mistral_config.get_supported_openai_params( - "mistral/magistral-medium-2506" - ) + supported_params = mistral_config.get_supported_openai_params("mistral/magistral-medium-2506") assert "reasoning_effort" in supported_params assert "thinking" in supported_params - supported_params_reasoning = mistral_config.get_supported_openai_params( - "mistral/mistral-medium-latest" - ) + supported_params_reasoning = mistral_config.get_supported_openai_params("mistral/mistral-medium-latest") assert "reasoning_effort" in supported_params_reasoning assert "thinking" not in supported_params_reasoning - supported_params_normal = mistral_config.get_supported_openai_params( - "mistral/mistral-large-latest" - ) + supported_params_normal = mistral_config.get_supported_openai_params("mistral/mistral-large-latest") assert "reasoning_effort" not in supported_params_normal assert "thinking" not in supported_params_normal @@ -117,6 +117,36 @@ class TestMistralReasoningSupport: ) assert optional_params["reasoning_effort"] == sent + @pytest.mark.parametrize( + ("requested", "drop_params", "sent"), + [ + ({"effort": "medium"}, False, {"reasoning_effort": "high"}), + ({"effort": "high", "summary": "auto"}, True, {"reasoning_effort": "high"}), + ({"summary": "auto"}, True, {}), + ], + ) + def test_reasoning_effort_object_is_sent_as_its_level(self, local_model_cost_map, requested, drop_params, sent): + import litellm + + optional_params = litellm.get_optional_params( + model="mistral-small-latest", + custom_llm_provider="mistral", + reasoning_effort=requested, + drop_params=drop_params, + ) + assert {key: optional_params[key] for key in ("reasoning_effort",) if key in optional_params} == sent + + @pytest.mark.parametrize("requested", [{"effort": "high", "summary": "auto"}, {"summary": "auto"}]) + def test_reasoning_effort_object_mistral_cannot_take_is_rejected(self, local_model_cost_map, requested): + import litellm + + with pytest.raises(litellm.UnsupportedParamsError, match=r"\['reasoning_effort'\]"): + litellm.get_optional_params( + model="devstral-latest", + custom_llm_provider="mistral", + reasoning_effort=requested, + ) + def test_reasoning_effort_is_forwarded_verbatim_when_the_map_declares_no_levels( self, local_model_cost_map, monkeypatch ): @@ -194,9 +224,7 @@ class TestMistralReasoningSupport: messages = [{"role": "user", "content": "What is 2+2?"}] optional_params = {"_add_reasoning_prompt": True} - result = mistral_config._add_reasoning_system_prompt_if_needed( - messages, optional_params - ) + result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params) # Should add a new system message at the beginning assert len(result) == 2 @@ -218,9 +246,7 @@ class TestMistralReasoningSupport: ] optional_params = {"_add_reasoning_prompt": True} - result = mistral_config._add_reasoning_system_prompt_if_needed( - messages, optional_params - ) + result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params) # Should modify existing system message assert len(result) == 2 @@ -251,9 +277,7 @@ class TestMistralReasoningSupport: ] optional_params = {"_add_reasoning_prompt": True} - result = mistral_config._add_reasoning_system_prompt_if_needed( - messages, optional_params - ) + result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params) # Should modify existing system message preserving list format assert len(result) == 2 @@ -266,10 +290,7 @@ class TestMistralReasoningSupport: # Original content should be preserved assert "You are a helpful assistant." in result[0]["content"][1]["text"] - assert ( - "You always provide detailed explanations." - in result[0]["content"][2]["text"] - ) + assert "You always provide detailed explanations." in result[0]["content"][2]["text"] assert result[1]["role"] == "user" @@ -287,9 +308,7 @@ class TestMistralReasoningSupport: ] string_params = {"_add_reasoning_prompt": True} - string_result = mistral_config._add_reasoning_system_prompt_if_needed( - string_messages, string_params - ) + string_result = mistral_config._add_reasoning_system_prompt_if_needed(string_messages, string_params) assert isinstance(string_result[0]["content"], str) assert "" in string_result[0]["content"] assert "You are helpful." in string_result[0]["content"] @@ -304,9 +323,7 @@ class TestMistralReasoningSupport: ] list_params = {"_add_reasoning_prompt": True} - list_result = mistral_config._add_reasoning_system_prompt_if_needed( - list_messages, list_params - ) + list_result = mistral_config._add_reasoning_system_prompt_if_needed(list_messages, list_params) assert isinstance(list_result[0]["content"], list) assert list_result[0]["content"][0]["type"] == "text" assert "" in list_result[0]["content"][0]["text"] @@ -319,9 +336,7 @@ class TestMistralReasoningSupport: messages = [{"role": "user", "content": "What is 2+2?"}] optional_params = {} - result = mistral_config._add_reasoning_system_prompt_if_needed( - messages, optional_params - ) + result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params) # Should return messages unchanged assert result == messages @@ -444,9 +459,7 @@ class TestMistralReasoningSupport: def test_mistral_streaming_chunk_preserves_thinking_blocks(): """Ensure streaming chunks keep magistral reasoning content.""" - iterator = MistralChatResponseIterator( - streaming_response=iter([]), sync_stream=True, json_mode=False - ) + iterator = MistralChatResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False) streamed_chunk = { "id": "chunk-1", @@ -523,9 +536,7 @@ class TestMistralParallelToolCalls: def test_get_supported_openai_params_includes_parallel_tool_calls(self): """Test that parallel_tool_calls is in supported parameters.""" mistral_config = MistralConfig() - supported_params = mistral_config.get_supported_openai_params( - "mistral/mistral-large-latest" - ) + supported_params = mistral_config.get_supported_openai_params("mistral/mistral-large-latest") assert "parallel_tool_calls" in supported_params def test_transform_request_preserves_parallel_tool_calls(self): @@ -691,9 +702,7 @@ class TestMistralEmptyContentHandling: result = MistralConfig._handle_empty_content_response(response_data) - assert ( - result["choices"][0]["message"]["content"] == "Hello, how can I help you?" - ) + assert result["choices"][0]["message"]["content"] == "Hello, how can I help you?" def test_handle_empty_content_response_handles_multiple_choices(self): """Test that only the first choice is processed for empty content.""" @@ -1431,9 +1440,7 @@ def test_mistral_transform_request_keeps_an_empty_tool_message_after_its_tool_ca { "role": "assistant", "content": "", - "tool_calls": [ - {"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}} - ], + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}], }, {"role": "tool", "tool_call_id": "call_1", "content": ""}, ], @@ -1457,9 +1464,7 @@ def test_mistral_transform_request_sends_a_null_tool_reply_as_empty_text(reply): { "role": "assistant", "content": "", - "tool_calls": [ - {"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}} - ], + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}], }, {"role": "tool", "tool_call_id": "call_1", **reply}, ], @@ -1514,9 +1519,7 @@ def test_mistral_transform_request_hoists_tool_message_image(): { "role": "assistant", "content": "", - "tool_calls": [ - {"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}} - ], + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}], }, { "role": "tool", @@ -1665,3 +1668,274 @@ def test_mistral_transform_request_drops_an_empty_tool_message_whose_call_id_is_ ) assert [m.get("role") for m in request["messages"]] == ["user", "assistant"] + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider"), + [ + ("mistral-large-latest", "mistral"), + ("mistral-large-4", "mistral"), + ("a-mistral-model-missing-from-the-cost-map", "mistral"), + ("codestral-latest", "codestral"), + ], +) +def test_prompt_cache_key_reaches_the_mistral_request_body( + local_model_cost_map: None, model: str, custom_llm_provider: str +) -> None: + import litellm + + optional_params: Final = litellm.get_optional_params( + model=model, custom_llm_provider=custom_llm_provider, prompt_cache_key="session-1" + ) + request: Final = MistralConfig().transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert request["prompt_cache_key"] == "session-1" + + +AZURE_HOSTED_MISTRAL_API_BASE: Final = "https://my-resource.services.ai.azure.com/openai/v1" + + +def _transform_with_prompt_cache_key( + litellm_params: dict[str, object], prompt_cache_key: object = "session-1" +) -> Mapping[str, object]: + return cast( + Mapping[str, object], + MistralConfig().transform_request( + model="mistral-small-latest", + messages=[{"role": "user", "content": "hi"}], + optional_params={"prompt_cache_key": prompt_cache_key, "temperature": 0.2}, + litellm_params=litellm_params, + headers={}, + ), + ) + + +@pytest.mark.parametrize("api_base", ["https://api.mistral.ai/v1", "https://codestral.mistral.ai/v1"]) +def test_prompt_cache_key_reaches_a_mistral_hosted_endpoint(api_base: str) -> None: + request: Final = _transform_with_prompt_cache_key({"api_base": api_base}) + + assert request["prompt_cache_key"] == "session-1" + + +@pytest.mark.parametrize("api_base", [AZURE_HOSTED_MISTRAL_API_BASE, "https://mistral.ai.example.com/v1"]) +def test_prompt_cache_key_for_another_host_is_rejected_without_drop_params(api_base: str) -> None: + import litellm + + with pytest.raises(litellm.UnsupportedParamsError, match="prompt_cache_key") as rejection: + _transform_with_prompt_cache_key({"api_base": api_base}) + + upstream_host: Final = urlparse(api_base).hostname or api_base + assert upstream_host not in str(rejection.value) + + +UNSENDABLE_PROMPT_CACHE_KEYS: Final = ("Büro-Sitzung-42", "日本", 123, ["session-1"]) + + +@pytest.mark.parametrize("prompt_cache_key", UNSENDABLE_PROMPT_CACHE_KEYS) +def test_prompt_cache_key_the_mistral_api_cannot_take_is_rejected_without_drop_params( + prompt_cache_key: object, +) -> None: + import litellm + + with pytest.raises(litellm.UnsupportedParamsError, match="ASCII"): + _transform_with_prompt_cache_key({"api_base": "https://api.mistral.ai/v1"}, prompt_cache_key) + + +@pytest.mark.parametrize("prompt_cache_key", UNSENDABLE_PROMPT_CACHE_KEYS) +def test_prompt_cache_key_the_mistral_api_cannot_take_is_dropped_under_drop_params(prompt_cache_key: object) -> None: + request: Final = _transform_with_prompt_cache_key( + {"api_base": "https://api.mistral.ai/v1", "drop_params": True}, prompt_cache_key + ) + + assert "prompt_cache_key" not in request + assert request["temperature"] == 0.2 + + +def test_prompt_cache_key_with_ascii_punctuation_reaches_the_mistral_api() -> None: + request: Final = _transform_with_prompt_cache_key({}, "team a/session:42") + + assert request["prompt_cache_key"] == "team a/session:42" + + +def test_prompt_cache_key_allowlisted_for_another_purpose_is_still_rejected_for_another_host() -> None: + import litellm + + with pytest.raises(litellm.UnsupportedParamsError, match="allowed_openai_params"): + _transform_with_prompt_cache_key( + {"api_base": AZURE_HOSTED_MISTRAL_API_BASE, "allowed_openai_params": ["reasoning_effort"]} + ) + + +def _chat_completion_reply() -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "mistral-small-latest", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + +def test_allowlisted_prompt_cache_key_reaches_a_gateway_in_front_of_mistral(respx_mock: respx.MockRouter) -> None: + import litellm + + gateway_route: Final = respx_mock.post("https://gateway.example.com/mistral/v1/chat/completions").mock( + return_value=_chat_completion_reply() + ) + + litellm.completion( + model="mistral/mistral-small-latest", + messages=[{"role": "user", "content": "hi"}], + api_base="https://gateway.example.com/mistral/v1", + api_key="gateway-key", + prompt_cache_key="session-1", + allowed_openai_params=["prompt_cache_key"], + max_retries=0, + ) + + assert json.loads(gateway_route.calls.last.request.content)["prompt_cache_key"] == "session-1" + + +@pytest.mark.asyncio +async def test_prompt_cache_key_is_dropped_from_an_async_codestral_call_to_another_host( + respx_mock: respx.MockRouter, +) -> None: + import litellm + + azure_route: Final = respx_mock.post(f"{AZURE_HOSTED_MISTRAL_API_BASE}/chat/completions").mock( + return_value=_chat_completion_reply() + ) + + await litellm.acompletion( + model="codestral/codestral-latest", + messages=[{"role": "user", "content": "hi"}], + api_base=AZURE_HOSTED_MISTRAL_API_BASE, + api_key="azure-key", + prompt_cache_key="session-1", + drop_params=True, + max_retries=0, + client=AsyncOpenAI( + api_key="azure-key", base_url=AZURE_HOSTED_MISTRAL_API_BASE, http_client=httpx.AsyncClient() + ), + ) + + assert "prompt_cache_key" not in json.loads(azure_route.calls.last.request.content) + + +@pytest.mark.asyncio +async def test_prompt_cache_key_rejection_for_an_async_codestral_call_names_codestral( + respx_mock: respx.MockRouter, +) -> None: + import litellm + + azure_route: Final = respx_mock.post(f"{AZURE_HOSTED_MISTRAL_API_BASE}/chat/completions").mock( + return_value=_chat_completion_reply() + ) + + with pytest.raises(litellm.BadRequestError, match=r"codestral does not support parameters: \['prompt_cache_key'\]"): + await litellm.acompletion( + model="codestral/codestral-latest", + messages=[{"role": "user", "content": "hi"}], + api_base=AZURE_HOSTED_MISTRAL_API_BASE, + api_key="azure-key", + prompt_cache_key="session-1", + max_retries=0, + client=AsyncOpenAI( + api_key="azure-key", base_url=AZURE_HOSTED_MISTRAL_API_BASE, http_client=httpx.AsyncClient() + ), + ) + + assert not azure_route.called + + +@pytest.mark.asyncio +async def test_async_codestral_call_sends_an_uploaded_file_by_its_id(respx_mock: respx.MockRouter) -> None: + import litellm + + codestral_api_base: Final = "https://codestral.mistral.ai/v1" + codestral_route: Final = respx_mock.post(f"{codestral_api_base}/chat/completions").mock( + return_value=_chat_completion_reply() + ) + + await litellm.acompletion( + model="codestral/codestral-latest", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Review this file"}, + {"type": "file", "file": {"file_id": "file-12345"}}, + ], + } + ], + api_base=codestral_api_base, + api_key="codestral-key", + max_retries=0, + client=AsyncOpenAI(api_key="codestral-key", base_url=codestral_api_base, http_client=httpx.AsyncClient()), + ) + + sent_content: Final = json.loads(codestral_route.calls.last.request.content)["messages"][0]["content"] + assert sent_content[1] == {"type": "file", "file_id": "file-12345"} + + +def test_prompt_cache_key_rejection_for_a_sync_codestral_call_is_an_unsupported_params_error( + respx_mock: respx.MockRouter, +) -> None: + import litellm + + azure_route: Final = respx_mock.post(f"{AZURE_HOSTED_MISTRAL_API_BASE}/chat/completions").mock( + return_value=_chat_completion_reply() + ) + + with pytest.raises( + litellm.UnsupportedParamsError, match=r"^litellm.UnsupportedParamsError: codestral does not support" + ): + litellm.completion( + model="codestral/codestral-latest", + messages=[{"role": "user", "content": "hi"}], + api_base=AZURE_HOSTED_MISTRAL_API_BASE, + api_key="azure-key", + prompt_cache_key="session-1", + max_retries=0, + ) + + assert not azure_route.called + + +def test_prompt_cache_key_for_another_host_is_dropped_under_drop_params() -> None: + request: Final = _transform_with_prompt_cache_key({"api_base": AZURE_HOSTED_MISTRAL_API_BASE, "drop_params": True}) + + assert "prompt_cache_key" not in request + assert request["temperature"] == 0.2 + + +def test_prompt_cache_key_for_the_azure_endpoint_set_in_the_environment_is_dropped_under_drop_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("MISTRAL_AZURE_API_BASE", "https://my-resource.services.ai.azure.com/openai") + + request: Final = _transform_with_prompt_cache_key({"drop_params": True}) + + assert "prompt_cache_key" not in request + + +def test_prompt_cache_key_stays_unsupported_for_vertex_hosted_mistral(local_model_cost_map: None) -> None: + import litellm + + with pytest.raises(litellm.UnsupportedParamsError): + litellm.get_optional_params(model="mistral-medium-3", custom_llm_provider="vertex_ai", prompt_cache_key="s") + + dropped: Final = litellm.get_optional_params( + model="mistral-medium-3", custom_llm_provider="vertex_ai", prompt_cache_key="s", drop_params=True + ) + assert "prompt_cache_key" not in dropped From 4c8b90945faf69eedfbdda3d4b7dad58ce74860d Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 01:31:39 -0700 Subject: [PATCH 14/64] fix(realtime)!: send the Gemini Live setup at the first content so session.update settings apply (#39864) * fix(realtime): send Gemini Live setup on the client's first message so session.update modalities apply Gemini and Vertex AI Live accept exactly one setup per session, and the proxy sent its default audio setup at connect, so a client's session.update (GA output_modalities or beta modalities) was dropped and every answer came back as audio. The setup now goes out on the client's first message: a session.update becomes that setup, anything else gets the model default. Bedrock Nova Sonic keeps its connect-time setup through a provider hook. Text setups label deltas response.output_text.delta, and downgrading TEXT to AUDIO on an audio-only model logs a warning * fix(realtime): keep one Gemini Live response modality and skip the "inf" token limit in the setup * fix(realtime): merge every session.update sent before the first content frame into the Gemini Live setup Gemini and Vertex Live accept one setup per connection, so the bridge used to turn only the first session.update into that setup and drop the rest without telling the client. Now every session.update that arrives before the first content frame is held, deep-merged (later values win per field), acknowledged with a session.updated built from the setup it would produce, and the merged session becomes the one setup sent with the first conversation item, audio frame, or tool result. The session events also surface the systemInstruction text as instructions, which the old string check never matched because the setup carries it as a Content object. * fix(realtime): cap the session.update merge depth so the recursive detector accepts it _merge_session_fields recurses once per nested mapping level, so CI's recursive_detector flagged it. Bound the descent at eight levels (the OpenAI session object nests four deep at most), replace anything deeper wholesale with the later session.update's value, and list the helper with the other depth-capped functions the detector allows. * fix(gemini-live): accept a string output modality and drop values Live cannot produce A bare string output_modalities is read as a one-element list and modality values Gemini Live cannot produce are dropped with a warning instead of reaching the backend, so a malformed session.update no longer costs the session. Transcription guardrails now fold their VAD disable into the first held session.update only, so the merged setup carries it exactly once. * fix(gemini realtime): answer an error for session updates Gemini Live cannot apply An output modality the model cannot produce, or a session.update sent after the first content, now gets an invalid_request_error event instead of a silent swap or drop, unless drop_params is set. The synthetic session.created reports the model's default modality, setupComplete no longer adds an extra session.updated, and merging nested session updates is linear in their size * fix(gemini realtime): send the setup before guardrail blocks and idle deadlines, mark push-to-talk turns A pre_call guardrail blocking the first typed message now goes out after the setup, so Gemini and Vertex no longer close the session with 1007. Push-to-talk sessions (turn_detection null or create_response false) wrap each audio turn in empty activityStart and activityEnd messages, which Gemini needs when automatic activity detection is off. A session with no content after 60 seconds sends its setup, since Vertex refuses a setup 300 seconds after connect. Held session.update frames past 10 MB answer session_update_too_large, and deferred mode now errors on a second session.update unless drop_params is set. * fix(realtime): ack unchanged Gemini Live session.update resends and honor null turn_detection An update that leaves the applied setup unchanged now gets session.updated instead of session_update_not_applied, so clients that resend their session with the first response keep working. Deferred setup mode merges held updates like the default mode, turn_detection null turns Gemini's activity detection off, drop_params drops a voice outside the documented Live voices, and the lazy setup deadline moves to 240 seconds * fix(realtime): keep Vertex activity detection off for null turn_detection and set up deferred sessions before text content * fix(realtime): refuse unscreenable Gemini audio, send held setup after buffered audio, reject unknown Vertex voices A realtime_input_transcription guardrail cannot screen Gemini Live audio before the answer starts, so audio input under one is refused with an audio_input_not_screenable error instead of reaching the model. In deferred setup mode a session.update that follows buffered audio is sent as setup at once. Only conversation content triggers the setup, so control events no longer lock in a default one. Vertex rejects an undocumented voice on the session.update, or drops it with drop_params, while Google AI Studio forwards any name. * fix(realtime): hold first content so a following session.update joins the Gemini setup Convert flat realtime tool schemas through Gemini's schema builder, and collapse a text and audio pair to what the model can produce on text-only models too * fix(realtime): check AI Studio Live voices like Vertex * fix(realtime): forward any voice to AI Studio Live models that accept it * fix(realtime): map OpenAI TTS voices before the Live voice check, keep outage spend rows, share the playground session config The Live voice gate now judges the Gemini voice an OpenAI TTS name maps to, so fable and onyx reach Vertex as Umbriel and Orus again. An upstream close before the held setup was sent is logged as a session, not a refusal. The playground sends one server_vad session config for typed and voice turns, and the e2e, integration, and dashboard tests follow lazy setup. * fix(realtime): take the lazy setup clock as a dependency so its tests run on a stepped clock The deadline and grace tests waited on real 50 ms timers, and the merge cost check compared wall-clock runs. The session now reads an injected clock, time.monotonic by default, the two tests step it, and the merge cost check moves to the CodSpeed benchmarks. The Nova Sonic style handler test passes the backend through wrap_backend, which main added after the test was written --- litellm/constants.py | 2 + .../litellm_core_utils/realtime_streaming.py | 330 ++++++++- .../llms/base_llm/realtime/transformation.py | 32 +- litellm/llms/bedrock/realtime/handler.py | 20 +- .../llms/bedrock/realtime/transformation.py | 19 +- litellm/llms/custom_httpx/llm_http_handler.py | 11 +- .../llms/gemini/realtime/transformation.py | 317 +++++++-- .../llms/vertex_ai/realtime/transformation.py | 17 +- litellm/types/llms/gemini.py | 16 +- litellm/utils.py | 2 +- .../test_realtime_session_merge_benchmarks.py | 21 + .../code_coverage_tests/recursive_detector.py | 1 + .../realtime/realtime_client.py | 3 + .../realtime/test_realtime_e2e.py | 4 +- ..._transcription_guardrail_session_update.py | 24 +- .../test_realtime_streaming.py | 670 +++++++++++++++++- .../custom_httpx/test_llm_http_handler.py | 189 +++++ .../test_gemini_realtime_transformation.py | 310 +++++++- .../test_vertex_ai_realtime_transformation.py | 110 +++ .../chat_ui/RealtimePlayground.test.tsx | 18 + .../components/chat_ui/RealtimePlayground.tsx | 65 +- 21 files changed, 1993 insertions(+), 188 deletions(-) create mode 100644 tests/benchmarks/test_realtime_session_merge_benchmarks.py diff --git a/litellm/constants.py b/litellm/constants.py index e7c9c5dcbfb..8e80c07b4ed 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -420,6 +420,8 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = get_env_float( "REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 20.0 ) BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS: Final = get_env_float("BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS", 5.0) +REALTIME_LAZY_SETUP_DEADLINE_SECONDS: Final = 240.0 +REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS: Final = 1.0 # RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 484bacdc72c..4db22e9d5bf 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,5 +1,6 @@ import asyncio import json +import time import traceback from collections.abc import Callable, Coroutine, Mapping, MutableMapping, Sequence from dataclasses import dataclass @@ -15,13 +16,19 @@ from litellm._logging import redact_internal_details_from_client_message, verbos from litellm._uuid import uuid from litellm.constants import ( BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, + REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS, + REALTIME_LAZY_SETUP_DEADLINE_SECONDS, REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError -from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend +from litellm.llms.base_llm.realtime.transformation import ( + BaseRealtimeConfig, + RealtimeBackend, + RealtimeSessionUpdateRejection, +) from litellm.types.llms.openai import ( OpenAIRealtimeEvents, OpenAIRealtimeOutputItemDone, @@ -125,6 +132,65 @@ class _LoggingWorker(Protocol): _DECODED_JSON: Final = TypeAdapter(object) +_SESSION_FIELDS: Final = TypeAdapter(dict[str, object]) +_NO_SESSION_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) +_SESSION_MERGE_MAX_DEPTH: Final = 8 + + +class _SessionUpdateFrame(TypedDict): + type: ReadOnly[str] + session: ReadOnly[Mapping[str, object]] + + +def _merge_session_fields( + base: Mapping[str, object], update: Mapping[str, object], depth: int = 0 +) -> Mapping[str, object]: + def merged_value(key: str) -> object: + if key not in update: + return base[key] + value: Final = update[key] + base_value: Final = base.get(key) + if depth < _SESSION_MERGE_MAX_DEPTH and isinstance(value, Mapping) and isinstance(base_value, Mapping): + return _merge_session_fields( + _SESSION_FIELDS.validate_python(base_value), _SESSION_FIELDS.validate_python(value), depth + 1 + ) + return update[key] + + return MappingProxyType({key: merged_value(key) for key in dict.fromkeys((*base, *update))}) + + +def _session_update_frame(session: Mapping[str, object]) -> str: + frame: Final[_SessionUpdateFrame] = {"type": "session.update", "session": session} + return json.dumps(frame, default=_SESSION_FIELDS.validate_python) + + +_HELD_SESSION_UPDATE_TOO_LARGE: Final = RealtimeSessionUpdateRejection( + code="session_update_too_large", + message=( + "The session.update frames sent before the first conversation content add up to more than 10 MB, " + "so this session.update was not applied" + ), + param="session", +) + + +def _unscreened_audio_error_event(client_event_id: object) -> Mapping[str, object]: + return { + "type": "error", + "event_id": f"event_{uuid.uuid4()}", + "error": { + "type": "invalid_request_error", + "code": "audio_input_not_screenable", + "message": ( + "A realtime_input_transcription guardrail is on, and Gemini Live only transcribes a turn after it " + "starts answering it, so audio input was not sent. Send text with conversation.item.create instead" + ), + "param": None, + "event_id": client_event_id if isinstance(client_event_id, str) else None, + }, + } + + def _decode_json_object(payload: str) -> Mapping[str, object]: """Decode a realtime frame into its top-level field mapping.""" return json.loads(payload) @@ -243,8 +309,17 @@ class RealTimeStreaming: event_normalizer: RealtimeEventNormalizer | None = None, usage_meter: RealtimeUsageMeter | None = None, logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, + lazy_setup_deadline_seconds: float = REALTIME_LAZY_SETUP_DEADLINE_SECONDS, + lazy_setup_content_grace_seconds: float = REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS, + clock: Callable[[], float] = time.monotonic, ): self.websocket: _ClientWebSocket = websocket + self._clock: Final = clock + self._lazy_setup_deadline: Final = clock() + lazy_setup_deadline_seconds + self._lazy_setup_content_grace_seconds: Final = lazy_setup_content_grace_seconds + self._held_client_content: tuple[str, ...] = () + self._held_client_content_bytes: int = 0 + self._held_client_content_deadline: float | None = None self._replay_scope: Final = client_websocket_scope(websocket) self.backend_ws = backend_ws self.logging_obj = logging_obj @@ -300,6 +375,10 @@ class RealTimeStreaming: # Gemini Live rejects a follow-up BidiGenerateContentSetup once any # content (realtimeInput / clientContent / toolResponse) has been sent. self._content_sent_after_setup: bool = False + self._held_session_update: Mapping[str, object] | None = None + self._applied_session_fields: Mapping[str, object] | None = None + self._held_session_update_bytes: int = 0 + self._unscreened_audio_refused: bool = False # Whether this is a transcription-only session (session.type == "transcription", # e.g. gpt-realtime-whisper). Such sessions must not be sent response.create and # their input_audio_transcription.completed usage drives duration-based cost. @@ -337,6 +416,17 @@ class RealTimeStreaming: ] ) _CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"]) + _CLIENT_CONTENT_TYPES = frozenset( + [ + "conversation.item.create", + "response.create", + "input_audio_buffer.append", + "input_audio_buffer.commit", + "input_audio_buffer.end", + ] + ) + _HELD_CONTENT_TYPES: Final = frozenset(["conversation.item.create", "input_audio_buffer.append"]) + _DEFERRED_SETUP_HELD_CONTENT_TYPES: Final = frozenset(["conversation.item.create"]) _AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = { "pcm16": {"type": "audio/pcm", "rate": 24000}, "g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000}, @@ -377,7 +467,7 @@ class RealTimeStreaming: message_obj: dict[str, Any] = cast(dict[str, Any], message) else: message_obj = cast(dict[str, Any], json.loads(message)) - self._collect_tool_calls_from_response_done(cast(dict, message_obj)) + self._collect_tool_calls_from_response_done(message_obj) if not self._should_store_message(message_obj): return try: @@ -502,7 +592,7 @@ class RealTimeStreaming: self.store_message(flush_event) self._capture_transcription_usage(flush_event) - def _collect_tool_calls_from_response_done(self, event_obj: dict | OpenAIRealtimeEvents) -> None: + def _collect_tool_calls_from_response_done(self, event_obj: Mapping[str, object] | OpenAIRealtimeEvents) -> None: """Extract function_call items from response.done events for spend logging.""" try: if event_obj.get("type") != "response.done": @@ -607,6 +697,167 @@ class RealTimeStreaming: self._usage_meter.observe_client_event(sent_event) return True + def _builds_setup_lazily(self) -> bool: + if self.provider_config is None: + return False + return self.provider_config.builds_setup_from_session_update() + + def _holds_session_updates_until_first_content(self) -> bool: + return self.session_configuration_request is None and self._builds_setup_lazily() + + async def _rejected_setup_session_update(self, message: str) -> bool: + provider_config: Final = self.provider_config + if provider_config is None or not provider_config.builds_setup_from_session_update(): + return False + client_event: Final = _decode_json_object(message) + session: Final = client_event.get("session") + update: Final = _SESSION_FIELDS.validate_python(session) if isinstance(session, Mapping) else _NO_SESSION_FIELDS + if self._leaves_applied_setup_unchanged(update): + await self._acknowledge_session_setup(self._applied_session_fields or _NO_SESSION_FIELDS) + return True + rejection: Final = provider_config.setup_session_update_rejection( + self.model, update, setup_sent=self.session_configuration_request is not None + ) + if rejection is None: + return False + await self.websocket.send_text(json.dumps(rejection.error_event(client_event.get("event_id")))) + return True + + def _held_session_update_message(self) -> str: + return _session_update_frame(self._held_session_update or _NO_SESSION_FIELDS) + + async def _hold_session_update(self, message: str) -> None: + client_event: Final = _decode_json_object(message) + if self._held_session_update_bytes + len(message.encode("utf-8")) > RealTimeStreaming._MAX_BUFFERED_BYTES: + await self.websocket.send_text( + json.dumps(_HELD_SESSION_UPDATE_TOO_LARGE.error_event(client_event.get("event_id"))) + ) + return + session: Final = client_event.get("session") + update: Final = _SESSION_FIELDS.validate_python(session) if isinstance(session, Mapping) else _NO_SESSION_FIELDS + held: Final = _merge_session_fields(self._held_session_update or _NO_SESSION_FIELDS, update) + self._held_session_update = held + self._held_session_update_bytes = len(_session_update_frame(held).encode("utf-8")) + await self._acknowledge_session_setup(held) + + def _setup_preview(self, session: Mapping[str, object]) -> str | None: + if self.provider_config is None: + return None + preview: Final = self.provider_config.transform_realtime_request( + _session_update_frame(session), self.model, None + ) + return next( + (frame for frame in preview if isinstance(frame, str) and "setup" in _decode_json_object(frame)), None + ) + + def _leaves_applied_setup_unchanged(self, update: Mapping[str, object]) -> bool: + applied: Final = self._applied_session_fields + if self.session_configuration_request is None or applied is None: + return False + applied_setup: Final = self._setup_preview(applied) + merged_setup: Final = self._setup_preview(_merge_session_fields(applied, update)) + return ( + applied_setup is not None + and merged_setup is not None + and _decode_json_object(applied_setup) == _decode_json_object(merged_setup) + ) + + async def _acknowledge_session_setup(self, session: Mapping[str, object]) -> None: + if self.provider_config is None: + return + setup: Final = self._setup_preview(session) + if setup is None: + return + acknowledged: Final = self.provider_config.transform_session_created_event( + self.model, str(self.logging_obj.litellm_trace_id), setup + ) + if acknowledged is None: + return + event: Final[Mapping[str, object]] = {**acknowledged, "type": "session.updated"} + event_str: Final = json.dumps(event) + self.store_message(event_str) + await self._send_event_to_client(event, event_str) + + async def _send_held_session_update_as_setup(self) -> None: + held: Final = self._held_session_update + if held is None: + return + message: Final = self._held_session_update_message() + self._held_session_update = None + self._held_session_update_bytes = 0 + if await self._send_to_backend(message): + self._applied_session_fields = held + await self._send_held_client_content() + + def _holds_client_content(self, msg_type: str | None, message: str) -> bool: + held_types: Final = ( + RealTimeStreaming._DEFERRED_SETUP_HELD_CONTENT_TYPES + if self._uses_deferred_backend_setup() + else RealTimeStreaming._HELD_CONTENT_TYPES + ) + return ( + msg_type in held_types + and self._held_session_update is not None + and self._holds_session_updates_until_first_content() + and not self._pending_messages_until_setup + and len(self._held_client_content) < RealTimeStreaming._MAX_BUFFERED_MESSAGES + and self._held_client_content_bytes + len(message.encode("utf-8")) <= RealTimeStreaming._MAX_BUFFERED_BYTES + ) + + def _hold_client_content(self, message: str) -> None: + if self._held_client_content_deadline is None: + self._held_client_content_deadline = self._clock() + self._lazy_setup_content_grace_seconds + self._held_client_content = (*self._held_client_content, message) + self._held_client_content_bytes += len(message.encode("utf-8")) + + async def _send_held_client_content(self) -> None: + held: Final = self._held_client_content + self._held_client_content = () + self._held_client_content_bytes = 0 + self._held_client_content_deadline = None + for message in held: + await self._send_to_backend(message) + + def _setup_deadline(self) -> float: + content_deadline: Final = self._held_client_content_deadline + return ( + self._lazy_setup_deadline if content_deadline is None else min(self._lazy_setup_deadline, content_deadline) + ) + + async def _send_setup_before_content(self, msg_type: str | None) -> None: + if self._holds_session_updates_until_first_content(): + await self._send_held_session_update_as_setup() + await self._send_default_setup_before_first_client_message(msg_type) + + async def _receive_client_message_before_setup_deadline(self) -> str | None: + if not self._holds_session_updates_until_first_content(): + return await self._receive_client_message() + receive: Final = asyncio.ensure_future(self._receive_client_message()) + try: + done, _ = await asyncio.wait((receive,), timeout=self._setup_deadline() - self._clock()) + if not done: + await self._send_setup_before_content(None) + return await receive + finally: + receive.cancel() + + async def _send_default_setup_before_first_client_message(self, msg_type: str | None) -> None: + if self.session_configuration_request is not None or self.provider_config is None: + return + if msg_type == "session.update" or ( + self._uses_deferred_backend_setup() + and (msg_type in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES or self._pending_messages_until_setup) + ): + return + default_setup: Final = self.provider_config.session_configuration_request(self.model) + if not default_setup: + return + setup: Final = self.maybe_inject_guardrail_auto_response_disable(default_setup) + await self.backend_ws.send(setup) + self._cache_session_configuration_request(setup) + self._applied_session_fields = _NO_SESSION_FIELDS + await self._send_held_client_content() + def _enforce_transcription_session_model(self, message: str) -> str: """Force client transcription session updates to the authorized model. @@ -851,6 +1102,20 @@ class RealTimeStreaming: self._guardrail_cancel_event_ids = self._guardrail_cancel_event_ids | {event_id} await self._send_to_backend(json.dumps({"type": "response.cancel", "event_id": event_id})) + async def _send_blocked_turn_to_backend(self, pre_block_backend_message: str | None, item_id: str | None) -> None: + if pre_block_backend_message is None and self._is_transcription_session: + return + await self._send_setup_before_content("conversation.item.create") + # Deliver any caller-supplied backend message FIRST so that + # protocol contracts requiring a specific ordering (e.g. + # Gemini Live's mandatory ``toolResponse`` after a + # ``toolCall``) are honored before the guardrail's own + # clientContent / cancel messages are sent. + if pre_block_backend_message is not None: + await self._send_to_backend(pre_block_backend_message) + if not self._is_transcription_session: + await self._withdraw_blocked_turn(item_id) + async def _withdraw_blocked_turn(self, item_id: str | None) -> None: await self._send_guardrail_cancel() if item_id is not None and self._deletes_conversation_items: @@ -958,6 +1223,21 @@ class RealTimeStreaming: """Transcription-only sessions have no assistant turn to gate and reject realtime session updates.""" return not self._is_transcription_session and self._has_audio_transcription_guardrails() + def _refuses_unscreened_audio(self, msg_type: str | None) -> bool: + return ( + msg_type in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES + and self._builds_setup_lazily() + and self._should_disable_vad_auto_response() + ) + + async def _refuse_unscreened_audio(self, msg_type: str | None, message: str) -> None: + if self._unscreened_audio_refused and msg_type not in RealTimeStreaming._CLIENT_AUDIO_BUFFER_COMMIT_TYPES: + return + self._unscreened_audio_refused = True + await self.websocket.send_text( + json.dumps(_unscreened_audio_error_event(_decode_json_object(message).get("event_id"))) + ) + async def run_realtime_guardrails( self, transcript: str, @@ -1030,15 +1310,7 @@ class RealTimeStreaming: # Use realtime_violation_message if configured; fall back to guardrail error text. error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg - # Deliver any caller-supplied backend message FIRST so that - # protocol contracts requiring a specific ordering (e.g. - # Gemini Live's mandatory ``toolResponse`` after a - # ``toolCall``) are honored before the guardrail's own - # clientContent / cancel messages are sent. - if pre_block_backend_message is not None: - await self._send_to_backend(pre_block_backend_message) - if not self._is_transcription_session: - await self._withdraw_blocked_turn(item_id) + await self._send_blocked_turn_to_backend(pre_block_backend_message, item_id) # Send the policy violation hint (shows as small gray status text in UI). await self.websocket.send_text( json.dumps( @@ -1136,6 +1408,9 @@ class RealTimeStreaming: break finally: self._flushing_pending_messages_until_setup = False + if self.session_created_sent_to_client and self._builds_setup_lazily(): + await self._maybe_send_guardrail_turn_detection_update() + continue if self.session_created_sent_to_client: # A synthetic session.created (with placeholder defaults) was # already forwarded to the client when we connected. The @@ -1295,7 +1570,11 @@ class RealTimeStreaming: return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket") def _backend_refused_session(self, close: BackendClose) -> bool: - return close.code != 1000 and not self._backend_sent_frames + return ( + close.code != 1000 + and not self._backend_sent_frames + and not self._holds_session_updates_until_first_content() + ) async def _log_session_end(self, transport_refusal: Exception | None) -> None: provider_refusal: Final = None if self.provider_config is None else self.provider_config.session_refusal() @@ -1488,7 +1767,7 @@ class RealTimeStreaming: client_event: _ClientEventFrame try: while True: - message = await self._receive_client_message() + message = await self._receive_client_message_before_setup_deadline() if message is None: return ClientLoopExit.CLIENT_DISCONNECTED @@ -1600,6 +1879,7 @@ class RealTimeStreaming: if ( msg_type == "session.update" and self.session_configuration_request is None + and self._held_session_update is None and not self._guardrail_turn_detection_update_sent and self._should_disable_vad_auto_response() and ( @@ -1663,12 +1943,25 @@ class RealTimeStreaming: except (json.JSONDecodeError, AttributeError): pass + if self._refuses_unscreened_audio(msg_type): + await self._refuse_unscreened_audio(msg_type, message) + continue + ## LOGGING # Log after any in-place modifications (GA remap, guardrail # turn_detection injection) so audit logs reflect what we # actually forward to the backend. self.store_input(message=message) + if self._holds_client_content(msg_type, message): + self._hold_client_content(message) + continue + if self._held_client_content and msg_type != "session.update": + await self._send_setup_before_content(msg_type) + + if msg_type in RealTimeStreaming._CLIENT_CONTENT_TYPES and self._uses_deferred_backend_setup(): + await self._send_setup_before_content(msg_type) + if self._should_buffer_client_message_until_setup(message): self._buffer_pending_message_until_setup(message) continue @@ -1697,6 +1990,15 @@ class RealTimeStreaming: # ``_maybe_send_guardrail_turn_detection_update`` would retry. if msg_type == "session.update" and self.provider_config is None: await self._send_session_gate_once() + if msg_type == "session.update" and await self._rejected_setup_session_update(message): + continue + if msg_type == "session.update" and self._holds_session_updates_until_first_content(): + await self._hold_session_update(message) + if self._pending_messages_until_setup: + await self._send_held_session_update_as_setup() + continue + if msg_type in RealTimeStreaming._CLIENT_CONTENT_TYPES: + await self._send_setup_before_content(msg_type) try: sent = await self._send_to_backend(message) except RealtimeTranscriptionProtocolError as e: diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index f1539d16b30..ff714bb908b 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -1,10 +1,12 @@ +import dataclasses from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from types import TracebackType -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, Literal, Protocol import httpx +from litellm._uuid import uuid from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents from litellm.types.realtime import ( RealtimeInputAudioTranscriptionUsage, @@ -22,6 +24,26 @@ else: LiteLLMLoggingObj = Any +@dataclasses.dataclass(frozen=True, slots=True) +class RealtimeSessionUpdateRejection: + code: Literal["invalid_type", "invalid_value", "session_update_not_applied", "session_update_too_large"] + message: str + param: str + + def error_event(self, client_event_id: object) -> dict[str, object]: + return { + "type": "error", + "event_id": f"event_{uuid.uuid4()}", + "error": { + "type": "invalid_request_error", + "code": self.code, + "message": self.message, + "param": self.param, + "event_id": client_event_id if isinstance(client_event_id, str) else None, + }, + } + + class RealtimeBackend(Protocol): async def __aenter__(self) -> object: ... @@ -93,6 +115,14 @@ class BaseRealtimeConfig(ABC): def session_configuration_request(self, model: str) -> str | None: # message sent to setup the realtime session return None + def builds_setup_from_session_update(self) -> bool: + return False + + def setup_session_update_rejection( + self, model: str, session: Mapping[str, object], *, setup_sent: bool + ) -> RealtimeSessionUpdateRejection | None: + return None + def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: return None diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index e9b2f26e977..00622a1e287 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -17,7 +17,6 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import redact_string, verbose_proxy_logger -from litellm._uuid import uuid from litellm.constants import ( BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS, BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY, @@ -31,13 +30,14 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes +from litellm.llms.base_llm.realtime.transformation import RealtimeSessionUpdateRejection from litellm.types.llms.bedrock import AwsAuthParams from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError -from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection, requested_modalities +from .transformation import BedrockRealtimeConfig, requested_modalities _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -256,23 +256,11 @@ async def _ack_session_update( async def _reject_session_update( client_ws: RealtimeClientWebSocket, parsed_client_message: Mapping[str, JsonValue], - rejection: BedrockSessionUpdateRejection, + rejection: RealtimeSessionUpdateRejection, ) -> bool: """Send an error event for a session.update Bedrock cannot apply; False means the client is gone.""" - client_event_id: Final = parsed_client_message.get("event_id") - error_event: Final = { - "type": "error", - "event_id": f"event_{uuid.uuid4()}", - "error": { - "type": "invalid_request_error", - "code": rejection.code, - "message": rejection.message, - "param": rejection.param, - "event_id": client_event_id if isinstance(client_event_id, str) else None, - }, - } try: - await client_ws.send_text(json.dumps(error_event)) + await client_ws.send_text(json.dumps(rejection.error_event(parsed_client_message.get("event_id")))) except Exception as e: # noqa: BLE001 # any send failure means the client is gone verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) return False diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 952d8858f78..25498fef46f 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -11,7 +11,7 @@ import math import uuid as uuid_lib from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final, Literal +from typing import Final import httpx from pydantic import ConfigDict, JsonValue, TypeAdapter @@ -19,7 +19,7 @@ from pydantic import ConfigDict, JsonValue, TypeAdapter from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig +from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeSessionUpdateRejection from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm from litellm.types.llms.base import LiteLLMBaseModel @@ -196,13 +196,6 @@ def _malformed_settings(session_config: Mapping[str, JsonValue]) -> tuple[tuple[ return tuple((name, expected) for name, expected, is_malformed in malformed if is_malformed) -@dataclasses.dataclass(frozen=True, slots=True) -class BedrockSessionUpdateRejection: - code: Literal["invalid_type", "invalid_value", "session_update_not_applied"] - message: str - param: str - - @dataclasses.dataclass(frozen=True, slots=True) class BedrockLockedSessionSettings: """The session settings Bedrock Nova Sonic locks in once the prompt starts.""" @@ -533,7 +526,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): max_tokens=max_tokens, ) - def session_update_rejection(self, json_message: Mapping[str, JsonValue]) -> BedrockSessionUpdateRejection | None: + def session_update_rejection(self, json_message: Mapping[str, JsonValue]) -> RealtimeSessionUpdateRejection | None: """The error for a session.update this session cannot apply, or None when it applies.""" session_config: Final = json_message.get("session") if not isinstance(session_config, Mapping): @@ -541,7 +534,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): malformed: Final = () if self.drop_params else _malformed_settings(session_config) if malformed: name, expected = malformed[0] - return BedrockSessionUpdateRejection( + return RealtimeSessionUpdateRejection( code="invalid_type", message=( f"Invalid type for 'session.{name}': Bedrock Nova Sonic needs {expected} there. " @@ -551,7 +544,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): ) unplayable_voice: Final = None if self.drop_params else _unplayable_voice(session_config) if unplayable_voice is not None: - return BedrockSessionUpdateRejection( + return RealtimeSessionUpdateRejection( code="invalid_value", message=( f"Invalid value for 'session.voice': '{unplayable_voice}' is not a Bedrock Nova Sonic voice. " @@ -562,7 +555,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if not self.prompt_started: return None if self._apply_locked_settings(session_config, self._started_settings) != self._started_settings: - return BedrockSessionUpdateRejection( + return RealtimeSessionUpdateRejection( code="session_update_not_applied", message="Bedrock Nova Sonic cannot change tools, instructions, voice, output audio format, temperature or max output tokens after the session has started, so this session.update was not applied", param="session", diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index c92315e53b3..61fc23589ac 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6973,14 +6973,19 @@ class BaseLLMHTTPHandler: ), ) - # Auto-send session setup if the provider requires it (e.g. - # Gemini/Vertex AI Live needs a `setup` before any realtime_input). + # Auto-send session setup at connect unless the provider builds it from the + # client's first session.update (Gemini/Vertex AI Live accept exactly one + # setup, so RealTimeStreaming sends theirs on the first client message). # Build the streaming handler first so a transcription guardrail's # auto-response disable can be folded into this one setup: Gemini # rejects a second setup, so a follow-up disable would be dropped # and the guardrail bypassed. _session_config: str | None = None - if provider_config.requires_session_configuration(): + sends_setup_at_connect: Final = ( + provider_config.requires_session_configuration() + and not provider_config.builds_setup_from_session_update() + ) + if sends_setup_at_connect: _session_config = provider_config.session_configuration_request(model) if _session_config: _session_config = realtime_streaming.maybe_inject_guardrail_auto_response_disable( diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 6f778ab2761..3134dfa1e4b 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -7,6 +7,7 @@ from collections import OrderedDict from collections.abc import Mapping, Sequence from typing import Final, cast +from pydantic import TypeAdapter from typing_extensions import ReadOnly, Required, TypedDict import litellm @@ -15,8 +16,9 @@ from litellm._uuid import uuid from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig +from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeSessionUpdateRejection from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + OPENAI_TO_GEMINI_TTS_VOICES, VertexGeminiConfig, ) from litellm.responses.litellm_completion_transformation.transformation import ( @@ -76,27 +78,82 @@ MAP_GEMINI_FIELD_TO_OPENAI_EVENT: Final[dict[str, OpenAIRealtimeEventTypes | Res # Keys the main transform loop handles; siblings like ``usageMetadata`` are skipped. _KNOWN_GEMINI_TOP_LEVEL_KEYS: Final[set] = {map_key.split(".", 1)[0] for map_key in MAP_GEMINI_FIELD_TO_OPENAI_EVENT} +_GEMINI_LIVE_RESPONSE_MODALITIES: Final = frozenset({"TEXT", "AUDIO"}) OPENAI_STOCK_REALTIME_VOICES: Final[frozenset[str]] = frozenset( {"alloy", "ash", "ballad", "cedar", "coral", "echo", "marin", "sage", "shimmer", "verse"} ) +# Per https://docs.cloud.google.com/vertex-ai/generative-ai/docs/live-api/configure-language-voice (2026-10-28) +GEMINI_LIVE_DOCUMENTED_VOICES: Final[frozenset[str]] = frozenset( + { + "achernar", + "achird", + "algenib", + "algieba", + "alnilam", + "aoede", + "autonoe", + "callirrhoe", + "charon", + "despina", + "enceladus", + "erinome", + "fenrir", + "gacrux", + "iapetus", + "kore", + "laomedeia", + "leda", + "orus", + "puck", + "pulcherrima", + "rasalgethi", + "sadachbia", + "sadaltager", + "schedar", + "sulafat", + "umbriel", + "vindemiatrix", + "zephyr", + "zubenelgenubi", + } +) -def _gemini_live_speech_config(voice: object) -> Mapping[str, object] | None: + +def _is_undocumented_live_voice(voice: object) -> bool: + return ( + isinstance(voice, str) + and voice.lower() not in OPENAI_STOCK_REALTIME_VOICES + and OPENAI_TO_GEMINI_TTS_VOICES.get(voice.lower(), voice).lower() not in GEMINI_LIVE_DOCUMENTED_VOICES + ) + + +def _gemini_live_speech_config(voice: object, drop_undocumented: bool) -> Mapping[str, object] | None: """Build the Gemini Live speechConfig for a client-requested voice. - OpenAI stock voice names have no Gemini equivalent and Gemini Live closes - the session on an unknown voice, so they are dropped with a warning and - the model keeps its default voice. Every other name is forwarded verbatim. + OpenAI stock voice names have no Gemini equivalent, so they are always + dropped, and so is any name outside Google's documented Live voices when + ``drop_undocumented`` is set, keeping the model's default voice. Every + other name is forwarded verbatim. """ - if isinstance(voice, str) and voice.lower() in OPENAI_STOCK_REALTIME_VOICES: + if not isinstance(voice, str): + return VertexGeminiConfig().map_audio_params({"voice": voice}) + if voice.lower() in OPENAI_STOCK_REALTIME_VOICES: verbose_logger.warning( "Gemini Realtime: voice %s is an OpenAI voice with no Gemini equivalent; " "dropping it so the session keeps the model's default voice.", voice, ) return None + if drop_undocumented and _is_undocumented_live_voice(voice): + verbose_logger.warning( + "Gemini Realtime: voice %s is not a documented Gemini Live voice; " + "drop_params is set, so it is dropped and the session keeps the model's default voice.", + voice, + ) + return None return VertexGeminiConfig().map_audio_params({"voice": voice}) @@ -142,16 +199,34 @@ def _base64_decoded_byte_count(data: str) -> int: return max(len(data) * 3 // 4 - padding, 0) +_JSON_ARRAY: Final = TypeAdapter(tuple[object, ...]) + + +def _chat_completion_tool(tool: object) -> dict[str, object]: + fields: Final = as_str_mapping(tool) or {} + if fields.get("type") != "function" or "function" in fields: + return {**fields} + return {"type": "function", "function": {key: value for key, value in fields.items() if key != "type"}} + + +def _system_instruction_text(system_instruction: HttpxContentType | None) -> str | None: + parts: Final = system_instruction.get("parts", []) if system_instruction else [] + texts: Final = tuple(text for part in parts if isinstance(text := part.get("text"), str)) + return "\n".join(texts) if texts else None + + class GeminiRealtimeConfig(BaseRealtimeConfig): _TOOL_CALL_ID_TO_NAME_MAX = 256 # LRU cap for call_id→name mapping - def __init__(self): + def __init__(self, drop_params: bool = False): super().__init__() + self.drop_params = drop_params self._tool_call_id_to_name: OrderedDict[str, str] = OrderedDict() # Gemini Live sometimes emits usageMetadata in a standalone frame between # turns; buffer it here so the next response.done carries the token counts. self._pending_usage_metadata: dict | None = None self._unbilled_input_audio_bytes: int = 0 + self._manual_activity_open: bool = False def is_setup_message(self, msg_obj: dict) -> bool: return "setup" in msg_obj @@ -167,6 +242,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): """Google AI Studio closes the socket (1007) on padding or silence tuning next to ``disabled: true``.""" return True + def _rejects_undocumented_voices(self, model: str) -> bool: + """Google AI Studio closes the socket (1007) on an unknown voice only for its native-audio Live models.""" + return bool(GeminiRealtimeConfig._model_cost_entry(model).get("gemini_native_audio")) + @staticmethod def _usage_detail_alias(details: object, defaults: dict[str, int]) -> dict[str, object]: detail_mapping: Final = as_str_mapping(details) @@ -272,17 +351,26 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): except (json.JSONDecodeError, TypeError, AttributeError): return False + def _start_manual_activity(self, session_configuration_request: str | None) -> tuple[str, ...]: + if self._manual_activity_open or not self._manual_turn_detection_enabled(session_configuration_request): + return () + self._manual_activity_open = True + activity_start: Final[BidiGenerateContentRealtimeInput] = {"activityStart": {}} + return (json.dumps({"realtimeInput": activity_start}),) + + def _end_manual_activity(self) -> tuple[str, ...]: + if not self._manual_activity_open: + return () + self._manual_activity_open = False + activity_end: Final[BidiGenerateContentRealtimeInput] = {"activityEnd": {}} + return (json.dumps({"realtimeInput": activity_end}),) + def _handle_input_audio_buffer_commit_or_end(self, session_configuration_request: str | None) -> list[str]: """Map OpenAI buffer commit/end to Gemini Live turn-boundary signals.""" if self._manual_turn_detection_enabled(session_configuration_request): - realtime_input_dict: BidiGenerateContentRealtimeInput = { - "activityEnd": True, - } - verbose_logger.debug("Gemini Realtime: Sending activityEnd realtimeInput to backend") - else: - realtime_input_dict = {"audioStreamEnd": True} - verbose_logger.debug("Gemini Realtime: Sending audioStreamEnd realtimeInput to backend") - return [json.dumps({"realtimeInput": realtime_input_dict})] + return list(self._end_manual_activity()) + audio_stream_end: Final[BidiGenerateContentRealtimeInput] = {"audioStreamEnd": True} + return [json.dumps({"realtimeInput": audio_stream_end})] def map_automatic_turn_detection(self, value: OpenAIRealtimeTurnDetection) -> AutomaticActivityDetection: """Map OpenAI ``server_vad`` to Gemini ``automaticActivityDetection``. @@ -325,7 +413,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): "voice", ] - def map_openai_params(self, optional_params: dict, non_default_params: dict) -> dict: + def map_openai_params(self, optional_params: dict, non_default_params: dict, *, model: str) -> dict: if "generationConfig" not in optional_params: optional_params["generationConfig"] = {} for key, value in non_default_params.items(): @@ -333,12 +421,12 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): optional_params["systemInstruction"] = HttpxContentType(role="user", parts=[{"text": value}]) elif key == "temperature": optional_params["generationConfig"]["temperature"] = value - elif key == "max_response_output_tokens": + elif key == "max_response_output_tokens" and isinstance(value, int): optional_params["generationConfig"]["maxOutputTokens"] = value elif key == "modalities": - optional_params["generationConfig"]["responseModalities"] = [ - modality.upper() for modality in cast(list[str], value) - ] + optional_params["generationConfig"]["responseModalities"] = list( + GeminiRealtimeConfig._requested_response_modalities(cast(object, value)) + ) elif key == "tools": from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, @@ -347,34 +435,43 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): vertex_gemini_config = VertexGeminiConfig() # Tools should be at the top level of setup, not inside generationConfig optional_params["tools"] = vertex_gemini_config.map_function( - value=value, optional_params=optional_params + value=[_chat_completion_tool(tool) for tool in _JSON_ARRAY.validate_python(value)], + optional_params=optional_params, ) elif key == "input_audio_transcription" and value is not None: optional_params["inputAudioTranscription"] = {} - elif key == "turn_detection" and value is not None: - value_typed = cast(OpenAIRealtimeTurnDetection, value) - if ( - isinstance(value_typed, dict) - and value_typed.get("type") == "semantic_vad" - and "create_response" not in value_typed - ): - # Pipecat/OpenAI GA semantic VAD — skip; Gemini uses its own VAD. - # Only skip when there is no create_response override so that - # a guardrail-injected create_response:false is not dropped. - continue - transformed_audio_activity_config = self.map_automatic_turn_detection(value_typed) - if transformed_audio_activity_config: - optional_params["realtimeInputConfig"] = BidiGenerateContentRealtimeInputConfig( - automaticActivityDetection=transformed_audio_activity_config - ) + elif key == "turn_detection": + realtime_input_config = self._realtime_input_config(cast(object, value)) + if realtime_input_config is not None: + optional_params["realtimeInputConfig"] = realtime_input_config elif key == "voice": - speech_config = _gemini_live_speech_config(value) + speech_config = _gemini_live_speech_config( + value, self.drop_params and self._rejects_undocumented_voices(model) + ) if speech_config: optional_params["generationConfig"]["speechConfig"] = speech_config if len(optional_params["generationConfig"]) == 0: optional_params.pop("generationConfig") return optional_params + def _realtime_input_config(self, turn_detection: object) -> BidiGenerateContentRealtimeInputConfig | None: + if turn_detection is None: + return BidiGenerateContentRealtimeInputConfig(automaticActivityDetection={"disabled": True}) + value_typed: Final = cast(OpenAIRealtimeTurnDetection, turn_detection) + if ( + isinstance(value_typed, dict) + and value_typed.get("type") == "semantic_vad" + and "create_response" not in value_typed + ): + # Pipecat/OpenAI GA semantic VAD — skip; Gemini uses its own VAD. + # Only skip when there is no create_response override so that + # a guardrail-injected create_response:false is not dropped. + return None + transformed_audio_activity_config: Final = self.map_automatic_turn_detection(value_typed) + if not transformed_audio_activity_config: + return None + return BidiGenerateContentRealtimeInputConfig(automaticActivityDetection=transformed_audio_activity_config) + @staticmethod def _extract_turn_detection(session: dict) -> dict | None: """Extract turn_detection from a session.update payload. @@ -423,6 +520,8 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if isinstance(input_cfg, dict): if "input_audio_transcription" not in normalized and "transcription" in input_cfg: normalized["input_audio_transcription"] = input_cfg["transcription"] + if "turn_detection" not in normalized and "turn_detection" in input_cfg: + normalized["turn_detection"] = input_cfg["turn_detection"] output_cfg: Final = audio.get("output") if isinstance(output_cfg, dict) and output_cfg.get("voice"): normalized["voice"] = output_cfg["voice"] @@ -434,7 +533,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return normalized @staticmethod - def _model_cost_entry(model: str) -> dict: + def _model_cost_entry(model: str) -> Mapping[str, object]: entry = litellm.model_cost.get(model) if entry is None: stripped: Final = model.split("/", 1)[-1] @@ -455,18 +554,122 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return "TEXT" if GeminiRealtimeConfig._is_text_only_live_model(model) else "AUDIO" @staticmethod - def _coerce_response_modalities(model: str, modalities: Sequence[object]) -> tuple[str, ...]: - """Swap responseModalities a Live model cannot produce: TEXT to AUDIO for - audio-only models, AUDIO to TEXT for text-only ones (e.g. transcribe-live).""" - normalized: Final = tuple( - modality.upper() if isinstance(modality, str) else str(modality).upper() for modality in modalities + def _requested_response_modalities(value: object) -> tuple[str, ...]: + if isinstance(value, str): + return (value.upper(),) + if not isinstance(value, (list, tuple)): + return () + return tuple(str(modality).upper() for modality in _JSON_ARRAY.validate_python(value)) + + @staticmethod + def _unsupported_response_modalities(model: str, requested: tuple[str, ...]) -> tuple[str, ...]: + unknown: Final = tuple(modality for modality in requested if modality not in _GEMINI_LIVE_RESPONSE_MODALITIES) + if unknown: + return unknown + if GeminiRealtimeConfig._is_audio_only_live_model(model) and frozenset(requested) == frozenset(("TEXT",)): + return ("TEXT",) + if GeminiRealtimeConfig._is_text_only_live_model(model) and frozenset(requested) == frozenset(("AUDIO",)): + return ("AUDIO",) + return () + + def setup_session_update_rejection( + self, model: str, session: Mapping[str, object], *, setup_sent: bool + ) -> RealtimeSessionUpdateRejection | None: + if self.drop_params: + return None + if setup_sent: + return RealtimeSessionUpdateRejection( + code="session_update_not_applied", + message=( + "Gemini Live accepts one session configuration per connection, and this session already sent " + "its configuration, so this session.update was not applied. Put every setting in session.update " + "before the first conversation content, or set drop_params to drop it instead" + ), + param="session", + ) + return self._modality_rejection(model, session) or self._voice_rejection(model, session) + + @staticmethod + def _modality_rejection(model: str, session: Mapping[str, object]) -> RealtimeSessionUpdateRejection | None: + param: Final = "output_modalities" if "output_modalities" in session else "modalities" + unsupported: Final = GeminiRealtimeConfig._unsupported_response_modalities( + model, GeminiRealtimeConfig._requested_response_modalities(session.get(param)) ) + if not unsupported: + return None + return RealtimeSessionUpdateRejection( + code="invalid_value", + message=( + f"Invalid value for 'session.{param}': {model} cannot produce " + f"{', '.join(modality.lower() for modality in unsupported)} output. Set drop_params to drop it instead" + ), + param=f"session.{param}", + ) + + def _voice_rejection(self, model: str, session: Mapping[str, object]) -> RealtimeSessionUpdateRejection | None: + audio_output: Final = as_str_mapping((as_str_mapping(session.get("audio")) or {}).get("output")) or {} + output_voice: Final = audio_output.get("voice") + voice: Final = output_voice or session.get("voice") + if not (self._rejects_undocumented_voices(model) and _is_undocumented_live_voice(voice)): + return None + param: Final = "session.audio.output.voice" if output_voice else "session.voice" + return RealtimeSessionUpdateRejection( + code="invalid_value", + message=( + f"Invalid value for '{param}': {voice} is not a Gemini Live voice {model} offers. " + "Set drop_params to drop it instead" + ), + param=param, + ) + + @staticmethod + def _coerce_response_modalities(model: str, modalities: Sequence[object]) -> tuple[str, ...]: + """Drop modalities Gemini Live cannot produce at all, swap the ones this + model cannot (TEXT to AUDIO for audio-only models, AUDIO to TEXT for + text-only ones such as transcribe-live), then keep a single modality + because Gemini Live rejects a setup naming two.""" + return GeminiRealtimeConfig._single_response_modality( + GeminiRealtimeConfig._swap_unsupported_response_modalities( + model, GeminiRealtimeConfig._known_response_modalities(model, modalities) + ) + ) + + @staticmethod + def _known_response_modalities(model: str, modalities: Sequence[object]) -> tuple[str, ...]: + normalized: Final = tuple(str(modality).upper() for modality in modalities) + known: Final = tuple(modality for modality in normalized if modality in _GEMINI_LIVE_RESPONSE_MODALITIES) + if len(known) < len(normalized): + verbose_logger.warning( + "Gemini Live: %s only produces TEXT or AUDIO; ignoring the requested %s", + model, + ", ".join(modality for modality in normalized if modality not in _GEMINI_LIVE_RESPONSE_MODALITIES), + ) + return known or (GeminiRealtimeConfig._default_response_modality(model),) + + @staticmethod + def _swap_unsupported_response_modalities(model: str, normalized: tuple[str, ...]) -> tuple[str, ...]: if GeminiRealtimeConfig._is_audio_only_live_model(model) and "TEXT" in normalized: + verbose_logger.warning( + "Gemini Live: %s only produces audio; downgrading the requested TEXT response modality to AUDIO", + model, + ) return tuple(modality for modality in normalized if modality != "TEXT") or ("AUDIO",) if GeminiRealtimeConfig._is_text_only_live_model(model) and "AUDIO" in normalized: return tuple(modality for modality in normalized if modality != "AUDIO") or ("TEXT",) return normalized + @staticmethod + def _single_response_modality(modalities: tuple[str, ...]) -> tuple[str, ...]: + if len(modalities) <= 1: + return modalities + chosen: Final = "AUDIO" if "AUDIO" in modalities else modalities[0] + verbose_logger.warning( + "Gemini Live: at most one response modality is allowed per session; using %s for the requested %s", + chosen, + ", ".join(modalities), + ) + return (chosen,) + @staticmethod def _finalize_gemini_live_setup(model: str, setup: dict[str, object]) -> dict[str, object]: generation_config: Final = setup.get("generationConfig") @@ -504,7 +707,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # would be silently dropped because ``map_openai_params`` only # recognises the flat OpenAI-beta key names. session_payload = self._normalize_session_payload_for_mapping(session_payload) - new_overrides: Final = self.map_openai_params(optional_params={}, non_default_params=session_payload) + new_overrides: Final = self.map_openai_params( + optional_params={}, non_default_params=session_payload, model=model + ) if session_configuration_request is None: generation_config: Final = new_overrides.setdefault("generationConfig", {}) @@ -618,14 +823,13 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): message_str = str(message) raise ValueError(f"Invalid JSON message: {message_str}") - messages: Final[list[str]] = [] msg_type: Final = json_message.get("type") if msg_type == "session.update": return self._handle_session_update(json_message, model, session_configuration_request) if msg_type == "response.create": - return [] # Gemini responds automatically; nothing to forward + return list(self._end_manual_activity()) if msg_type == "conversation.item.create": return self._handle_conversation_item(json_message) @@ -643,8 +847,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): gemini_msg: Final = json.dumps({"realtimeInput": realtime_input_dict}) verbose_logger.debug("Gemini Realtime: Sending audio realtimeInput to backend") - messages.append(gemini_msg) - return messages + return [*self._start_manual_activity(session_configuration_request), gemini_msg] if msg_type in ("input_audio_buffer.commit", "input_audio_buffer.end"): return self._handle_input_audio_buffer_commit_or_end(session_configuration_request) @@ -660,22 +863,21 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): logging_session_id: str, session_configuration_request: str | None = None, ) -> OpenAIRealtimeStreamSessionEvents: - if session_configuration_request: - session_configuration_request_dict: BidiGenerateContentSetup = _parse_setup(session_configuration_request) - else: - session_configuration_request_dict = {} + session_configuration_request_dict: Final = _parse_setup( + session_configuration_request or self.session_configuration_request(model) + ) _model: Final = session_configuration_request_dict.get("model") or model generation_config: Final = session_configuration_request_dict.get("generationConfig", {}) or {} gemini_modalities: Final = generation_config.get("responseModalities", ["AUDIO"]) _modalities: Final = [modality.lower() for modality in cast(list[str], gemini_modalities)] - _system_instruction: Final = session_configuration_request_dict.get("systemInstruction") + instructions: Final = _system_instruction_text(session_configuration_request_dict.get("systemInstruction")) session: Final = OpenAIRealtimeStreamSession( id=logging_session_id, modalities=_modalities, ) - if _system_instruction is not None and isinstance(_system_instruction, str): - session["instructions"] = _system_instruction + if instructions is not None: + session["instructions"] = instructions if _model is not None and isinstance(_model, str): # Strip Vertex/AI Studio path prefixes to expose the bare model name. if "/models/" in _model: @@ -1662,6 +1864,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # Deferred setup opt-in: litellm.gemini_live_defer_setup = True return not litellm.gemini_live_defer_setup + def builds_setup_from_session_update(self) -> bool: + return True + def session_configuration_request(self, model: str) -> str: """ diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index 13f547b3b92..7635cd32a06 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -39,8 +39,8 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): (they require async I/O) and injected at construction time. """ - def __init__(self, access_token: str, project: str, location: str) -> None: - super().__init__() + def __init__(self, access_token: str, project: str, location: str, drop_params: bool = False) -> None: + super().__init__(drop_params=drop_params) self._access_token = access_token self._project = project self._location = location @@ -51,6 +51,9 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): def _rejects_vad_tuning_with_detection_off(self) -> bool: return False + def _rejects_undocumented_voices(self, model: str) -> bool: + return True + # ------------------------------------------------------------------ # URL # ------------------------------------------------------------------ @@ -163,7 +166,7 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): # settings would be silently dropped because ``map_openai_params`` only # recognises the flat OpenAI-beta key names. session_params = self._normalize_session_payload_for_mapping(session_params) - setup_config: Final = self.map_openai_params(optional_params={}, non_default_params=session_params) + setup_config: Final = self.map_openai_params(optional_params={}, non_default_params=session_params, model=model) # Use full Vertex AI model path setup_config["model"] = self._vertex_model_path(model) @@ -177,15 +180,16 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): # ``silence_duration_ms``). ``map_automatic_turn_detection`` sets # ``disabled=True`` whenever ``create_response`` is absent or # ``False``. Force ``disabled=False`` only when the client did - # not explicitly request ``create_response: False`` — that path + # not explicitly request ``create_response: False`` or + # ``turn_detection: null`` — the first of those paths # is how transcription guardrails suppress automatic responses, # and overriding it here would silently bypass the guardrail. # Vertex Live has no "VAD on, no auto-response" mode, so callers # that need that behaviour must accept that VAD is off. client_turn_detection: Final = session_params.get("turn_detection") client_disabled_auto_response: Final = ( - isinstance(client_turn_detection, dict) and client_turn_detection.get("create_response") is False - ) + "turn_detection" in session_params and client_turn_detection is None + ) or (isinstance(client_turn_detection, dict) and client_turn_detection.get("create_response") is False) realtime_input_config: Final = setup_config.setdefault("realtimeInputConfig", {}) automatic_detection: Final = realtime_input_config.setdefault("automaticActivityDetection", {}) if not client_disabled_auto_response: @@ -261,4 +265,5 @@ def vertex_realtime_config( access_token=access_token, project=project, location=VertexBase.get_vertex_region(vertex_region=location, model=model), + drop_params=drop_params, ) diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index 2ad81e64e05..67d90eefab0 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -81,6 +81,14 @@ class BidiGenerateContentServerMessage(TypedDict, total=False): """Output only. The setup complete message.""" +class BidiGenerateContentActivityStart(TypedDict): + pass + + +class BidiGenerateContentActivityEnd(TypedDict): + pass + + class BidiGenerateContentRealtimeInput(TypedDict, total=False): text: str """The text to be sent to the model.""" @@ -94,11 +102,11 @@ class BidiGenerateContentRealtimeInput(TypedDict, total=False): audioStreamEnd: bool """Output only. If true, indicates that the audio stream has ended.""" - activityStart: bool - """Output only. If true, indicates that the activity has started.""" + activityStart: BidiGenerateContentActivityStart + """Marks the start of user activity when automatic activity detection is disabled.""" - activityEnd: bool - """Output only. If true, indicates that the activity has ended.""" + activityEnd: BidiGenerateContentActivityEnd + """Marks the end of user activity when automatic activity detection is disabled.""" StartOfSpeechSensitivityEnum = Literal[ diff --git a/litellm/utils.py b/litellm/utils.py index e85e089289a..6f56cc65f45 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10911,7 +10911,7 @@ class ProviderConfigManager: if LlmProviders.GEMINI == provider: from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig - return GeminiRealtimeConfig() + return GeminiRealtimeConfig(drop_params=drop_params) if LlmProviders.META == provider: from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig diff --git a/tests/benchmarks/test_realtime_session_merge_benchmarks.py b/tests/benchmarks/test_realtime_session_merge_benchmarks.py new file mode 100644 index 00000000000..29ee1b638ad --- /dev/null +++ b/tests/benchmarks/test_realtime_session_merge_benchmarks.py @@ -0,0 +1,21 @@ +from typing import Final + +import pytest + +from litellm.litellm_core_utils.realtime_streaming import _merge_session_fields + +NESTED_LEVELS: Final = 8 +LEAVES: Final = 2000 + + +def _nested_session(depth: int) -> dict[str, object]: + return {"child": _nested_session(depth - 1)} if depth else {f"k{index}": index for index in range(LEAVES)} + + +@pytest.mark.benchmark +def test_merging_a_session_update_that_repeats_every_nested_level() -> None: + session: Final = _nested_session(NESTED_LEVELS) + + merged: Final = _merge_session_fields(session, session) + + assert merged == session diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index c943ff2cbd0..9b61c1d39ae 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -85,6 +85,7 @@ IGNORE_FUNCTIONS = [ "_sort_processed_sets", # bounded by the nesting depth of the log-record extra it walks (a finite JSON tree, no cycles possible). "scrub_json_strings", # max depth set (MAX_SCRUB_DEPTH); fails closed by returning "[Filtered]" for anything nested past the cap. "_without_requester_ip_address", # walks only the fresh tree _sanitize_request_body_for_spend_logs_payload returns (mapping cycles already cut), so it never goes deeper than that sanitizer did. + "_merge_session_fields", # max depth set (_SESSION_MERGE_MAX_DEPTH); deeper mappings are replaced wholesale by the later session.update. ] diff --git a/tests/e2e/llm_translation/realtime/realtime_client.py b/tests/e2e/llm_translation/realtime/realtime_client.py index 4d8336775ef..6bd5d61d721 100644 --- a/tests/e2e/llm_translation/realtime/realtime_client.py +++ b/tests/e2e/llm_translation/realtime/realtime_client.py @@ -47,6 +47,7 @@ class RealtimeProvider: id: str alias: str litellm_params: LiteLLMParamsBody + modalities: tuple[str, ...] = ("text",) PROVIDERS = ( @@ -75,6 +76,7 @@ PROVIDERS = ( model="gemini/gemini-3.1-flash-live-preview", api_key="os.environ/GEMINI_API_KEY", ), + modalities=("audio",), ), RealtimeProvider( "vertex_ai", @@ -84,6 +86,7 @@ PROVIDERS = ( vertex_location="us-central1", vertex_credentials="os.environ/VERTEXAI_CREDENTIALS", ), + modalities=("audio",), ), RealtimeProvider( "bedrock", diff --git a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py index f5592abb34d..3b7ba3a2430 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py @@ -124,7 +124,8 @@ def test_text_conversation( session.send( SessionUpdate( session=SessionConfig( - instructions="You are a terse assistant. Reply in one short sentence." + modalities=list(provider.modalities), + instructions="You are a terse assistant. Reply in one short sentence.", ) ) ) @@ -159,6 +160,7 @@ def test_tool_call_round_trip( session.send( SessionUpdate( session=SessionConfig( + modalities=list(provider.modalities), tools=[WEATHER_TOOL], tool_choice="auto", instructions=( diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py index ad58a400cec..2712d6e30cd 100644 --- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py +++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py @@ -355,7 +355,13 @@ def _voice_scenario_created_after_the_first_frame(transcript: str) -> RealtimeRe def _live_scenario() -> RealtimeResponse: - return RealtimeResponse(content_type="application/x-realtime", events=()) + return RealtimeResponse( + content_type="application/x-realtime", + events=( + {"serverContent": {"modelTurn": {"parts": [{"text": "scripted $REQUEST_ID"}]}}}, + {"serverContent": {"turnComplete": True}}, + ), + ) def _realtime_update(turn_detection: JsonValue) -> dict[str, JsonValue]: @@ -1348,10 +1354,14 @@ def test_gemini_ai_studio_voice_session_with_vad_tuning_connects_under_the_guard _ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, - (_step(_update_frame(update), until="session.updated"),), + ( + _step(_update_frame(update), until="session.updated"), + _step(_user_text_item(CLEAN_TRANSCRIPT), RESPONSE_CREATE, until=RESPONSE_DONE), + ), ) observed: Final = _observed(upstream_url, handle.scenario_id) - assert session.types == ("session.created", "session.updated"), session + assert session.types[:2] == ("session.created", "session.updated"), session + assert session.types[-1] == RESPONSE_DONE, session assert _setup_detection(observed) == {"disabled": True}, _sent(observed) @@ -1365,10 +1375,14 @@ def test_vertex_live_voice_session_keeps_vad_tuning_under_the_guardrail(guardrai _ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, - (_step(_update_frame(_realtime_update(TUNED_SERVER_VAD)), until="session.updated"),), + ( + _step(_update_frame(_realtime_update(TUNED_SERVER_VAD)), until="session.updated"), + _step(_user_text_item(CLEAN_TRANSCRIPT), RESPONSE_CREATE, until=RESPONSE_DONE), + ), ) observed: Final = _observed(upstream_url, handle.scenario_id) - assert session.types == ("session.created", "session.updated"), session + assert session.types[:2] == ("session.created", "session.updated"), session + assert session.types[-1] == RESPONSE_DONE, session assert _setup_detection(observed) == GEMINI_TUNED_DETECTION, _sent(observed) diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index a8543ebcfc2..ba76138a268 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -1,7 +1,7 @@ import asyncio import base64 import json -from collections.abc import Coroutine, Mapping, Sequence +from collections.abc import Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -14,6 +14,7 @@ from websockets.frames import Close import litellm from litellm.constants import ( BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, + REALTIME_LAZY_SETUP_DEADLINE_SECONDS, REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -21,6 +22,7 @@ from litellm.constants import ( from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( RealTimeStreaming, + _merge_session_fields, client_sent_openai_beta_realtime_header, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -1255,11 +1257,15 @@ async def test_client_ack_caches_setup_to_prevent_duplicate_session_update_setup side_effect=[ json.dumps({"type": "session.update", "session": {"tools": []}}), json.dumps({"type": "session.update", "session": {"tools": []}}), + _TEXT_ITEM_CREATE, + _RESPONSE_CREATE, Exception("client done"), ] ) + websocket.send_text = AsyncMock() provider_config = MagicMock() + provider_config.setup_session_update_rejection = MagicMock(return_value=None) def _transform(message: str, model: str, session_configuration_request=None): if session_configuration_request is None: @@ -2285,6 +2291,7 @@ async def test_guardrail_turn_detection_injected_into_first_session_update_defer logging_obj.success_handler = MagicMock() provider_config = MagicMock() + provider_config.builds_setup_from_session_update = MagicMock(return_value=False) transformed_messages = [] def mock_transform(msg, model, session_config): @@ -2351,6 +2358,7 @@ async def test_guardrail_turn_detection_injection_tolerates_non_dict_value( logging_obj.success_handler = MagicMock() provider_config = MagicMock() + provider_config.builds_setup_from_session_update = MagicMock(return_value=False) transformed_messages = [] def mock_transform(msg, model, session_config): @@ -2413,6 +2421,7 @@ async def test_subsequent_session_update_cannot_reenable_vad_when_guardrails_act logging_obj.success_handler = MagicMock() provider_config = MagicMock() + provider_config.setup_session_update_rejection = MagicMock(return_value=None) transformed_messages = [] def mock_transform(msg, model, session_config): @@ -2465,6 +2474,7 @@ async def test_follow_up_setup_updates_cached_session_configuration_request(): logging_obj.success_handler = MagicMock() provider_config = MagicMock() + provider_config.setup_session_update_rejection = MagicMock(return_value=None) follow_up_setup = json.dumps( { "setup": { @@ -2498,6 +2508,636 @@ async def test_follow_up_setup_updates_cached_session_configuration_request(): assert streaming.session_configuration_request == follow_up_setup +def _gemini_default_mode_streaming( + monkeypatch: pytest.MonkeyPatch, + client_messages: Sequence[str], + model: str = "gemini-live-2.5-flash-native-audio", + drop_params: bool = False, + defer_setup: bool = False, +) -> tuple[RealTimeStreaming, MagicMock]: + monkeypatch.setattr(litellm, "gemini_live_defer_setup", defer_setup, raising=False) + from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig + + client_ws = MagicMock() + client_ws.receive_text = AsyncMock(side_effect=[*client_messages, ConnectionClosed(None, None)]) + client_ws.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + streaming = RealTimeStreaming( + client_ws, + backend_ws, + MagicMock(), + provider_config=GeminiRealtimeConfig(drop_params=drop_params), + model=model, + ) + return streaming, backend_ws + + +def _events_sent_to_client(streaming: RealTimeStreaming) -> list[dict]: + return [json.loads(call.args[0]) for call in streaming.websocket.send_text.await_args_list] + + +_TEXT_ITEM_CREATE = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + } +) + +_RESPONSE_CREATE = json.dumps({"type": "response.create"}) + + +@pytest.mark.asyncio +async def test_default_mode_sends_default_setup_before_first_non_session_client_message(monkeypatch): + """Gemini Live setup is no longer pre-sent at connect, so a client that never + sends session.update still gets the provider default setup ahead of its + first content frame.""" + streaming, backend_ws = _gemini_default_mode_streaming(monkeypatch, [_TEXT_ITEM_CREATE]) + + await streaming.client_ack_messages() + + frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["AUDIO"] + assert streaming.session_configuration_request == backend_ws.send.await_args_list[0].args[0] + + +@pytest.mark.asyncio +async def test_default_mode_session_update_first_builds_the_only_setup(monkeypatch): + session_update = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}}) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [session_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash" + ) + + await streaming.client_ack_messages() + + frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"] + + +@pytest.mark.asyncio +async def test_default_mode_merges_every_session_update_sent_before_the_first_content_frame(monkeypatch): + """Clients configure a session over several session.update messages. Gemini + Live takes one setup, so every update sent before the first content frame is + merged into that setup instead of only the first one applying.""" + instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}) + modalities_update = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}}) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [instructions_update, modalities_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash" + ) + + await streaming.client_ack_messages() + + frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"] + assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief." + + +@pytest.mark.asyncio +async def test_default_mode_merges_nested_audio_settings_across_session_updates(monkeypatch): + vad_update = json.dumps( + {"type": "session.update", "session": {"audio": {"input": {"turn_detection": {"silence_duration_ms": 1500}}}}} + ) + voice_update = json.dumps({"type": "session.update", "session": {"audio": {"output": {"voice": "Kore"}}}}) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [vad_update, voice_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash" + ) + + await streaming.client_ack_messages() + + setup = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"] + assert setup["realtimeInputConfig"]["automaticActivityDetection"]["silenceDurationMs"] == 1500 + assert setup["generationConfig"]["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + + +def test_merge_session_fields_replaces_nested_mappings_wholesale_past_the_depth_cap(): + def nested(levels: int, key: str) -> dict: + return {key: levels, **({"child": nested(levels - 1, key)} if levels else {})} + + def level(node, depth: int): + return node if depth == 0 else level(node["child"], depth - 1) + + merged = _merge_session_fields(nested(10, "keep"), nested(10, "new")) + + assert all({"keep", "new"} <= set(level(merged, depth)) for depth in range(9)) + assert "keep" not in level(merged, 9) + assert level(merged, 9)["new"] == 1 + + +@pytest.mark.asyncio +async def test_default_mode_acknowledges_each_held_session_update_before_any_backend_frame(monkeypatch): + """A client that waits for session.updated before sending content must not + hang while its updates are held, so each one is acknowledged right away + with the modalities the merged setup will carry.""" + instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}) + modalities_update = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}}) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [instructions_update, modalities_update], model="gemini-live-2.5-flash" + ) + + await streaming.client_ack_messages() + + acks = [json.loads(call.args[0]) for call in streaming.websocket.send_text.await_args_list] + assert [ack["type"] for ack in acks] == ["session.updated", "session.updated"] + assert [ack["session"].get("output_modalities") or ack["session"].get("modalities") for ack in acks] == [ + ["audio"], + ["text"], + ] + assert acks[1]["session"]["instructions"] == "Be brief." + backend_ws.send.assert_not_awaited() + + +_LATE_INSTRUCTIONS_UPDATE = json.dumps( + {"type": "session.update", "event_id": "evt_late", "session": {"instructions": "Be brief."}} +) + + +@pytest.mark.asyncio +async def test_default_mode_answers_a_session_update_after_the_first_content_with_an_error(monkeypatch): + streaming, backend_ws = _gemini_default_mode_streaming(monkeypatch, [_TEXT_ITEM_CREATE, _LATE_INSTRUCTIONS_UPDATE]) + + await streaming.client_ack_messages() + + errors: Final = [event["error"] for event in _events_sent_to_client(streaming) if event["type"] == "error"] + assert [(error["code"], error["event_id"]) for error in errors] == [("session_update_not_applied", "evt_late")] + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + + +@pytest.mark.asyncio +async def test_default_mode_drops_a_late_session_update_quietly_under_drop_params(monkeypatch): + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [_TEXT_ITEM_CREATE, _LATE_INSTRUCTIONS_UPDATE], drop_params=True + ) + + await streaming.client_ack_messages() + + assert _events_sent_to_client(streaming) == [] + assert backend_ws.send.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "requested_output"), + [ + ("gemini-live-2.5-flash-native-audio", {"modalities": ["text"]}), + ("gemini-live-2.5-flash", {"output_modalities": ["image"]}), + ], +) +async def test_default_mode_rejects_an_output_modality_the_model_cannot_produce(monkeypatch, model, requested_output): + update: Final = json.dumps( + {"type": "session.update", "session": {**requested_output, "instructions": "Be brief."}} + ) + streaming, backend_ws = _gemini_default_mode_streaming(monkeypatch, [update, _TEXT_ITEM_CREATE], model=model) + + await streaming.client_ack_messages() + + events: Final = _events_sent_to_client(streaming) + assert [(event["type"], event["error"]["code"]) for event in events] == [("error", "invalid_value")] + setup: Final = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"] + assert "systemInstruction" not in setup + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("modalities", "drop_params"), + [(["text", "audio"], False), (["text"], True)], + ids=["openai-default-pair", "text-dropped-under-drop-params"], +) +async def test_default_mode_applies_the_update_with_audio_output_on_an_audio_only_model( + monkeypatch, modalities, drop_params +): + update: Final = json.dumps( + {"type": "session.update", "session": {"modalities": modalities, "instructions": "Be brief."}} + ) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], drop_params=drop_params + ) + + await streaming.client_ack_messages() + + assert [event["type"] for event in _events_sent_to_client(streaming)] == ["session.updated"] + setup: Final = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"] + assert setup["generationConfig"]["responseModalities"] == ["AUDIO"] + assert setup["systemInstruction"]["parts"][0]["text"] == "Be brief." + + +@pytest.mark.asyncio +async def test_default_mode_sends_no_extra_session_updated_when_gemini_confirms_the_setup(monkeypatch): + """Each held session.update was already acknowledged, so Gemini's setupComplete + must not reach the client as one more session.updated nobody asked for.""" + monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False) + from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig + + model: Final = "gemini-live-2.5-flash-native-audio" + config: Final = GeminiRealtimeConfig() + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=[json.dumps({"setupComplete": {}}).encode(), ConnectionClosed(None, None)]) + logging_obj: Final = _mock_logging() + logging_obj.litellm_trace_id = "trace_lazy" + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj, provider_config=config, model=model) + streaming.session_created_sent_to_client = True + streaming.session_configuration_request = config.session_configuration_request(model) + + await streaming.backend_to_client_send_messages() + + assert [json.loads(call.args[0])["type"] for call in client_ws.send_text.await_args_list] == [] + + +@pytest.mark.asyncio +async def test_default_mode_folds_transcription_guardrail_vad_disable_into_the_merged_setup(monkeypatch): + """A transcription guardrail gates every turn, so its auto-response disable + has to land in the one setup built from the held session.updates, injected + once rather than into every update that gets merged.""" + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + debug = MagicMock() + monkeypatch.setattr("litellm.litellm_core_utils.realtime_streaming.verbose_logger.debug", debug) + instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}) + vad_update = json.dumps( + {"type": "session.update", "session": {"turn_detection": {"type": "server_vad", "silence_duration_ms": 900}}} + ) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [instructions_update, vad_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash" + ) + + await streaming.client_ack_messages() + + frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + assert frames[0]["setup"]["realtimeInputConfig"]["automaticActivityDetection"] == {"disabled": True} + assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief." + assert streaming._guardrail_turn_detection_update_sent is True + injected = [call for call in debug.call_args_list if str(call.args[0]).startswith("Injected turn_detection")] + assert len(injected) == 1 + + +def _blocking_text_guardrail() -> CustomGuardrail: + class _BlockingTextGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + if any("zebrafruit" in text for text in inputs.get("texts", [])): + raise ValueError("Content blocked") + return inputs + + return _BlockingTextGuardrail( + guardrail_name="test_text_guard", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + + +@pytest.mark.asyncio +async def test_default_mode_sends_the_setup_before_a_guardrail_blocks_the_first_message(monkeypatch): + """Gemini Live closes the socket when content arrives before the setup, so a + guardrail that blocks the first typed message must not reach the backend first.""" + monkeypatch.setattr(litellm, "callbacks", [_blocking_text_guardrail()]) + instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}) + blocked_item = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "zebrafruit"}]}, + } + ) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [instructions_update, blocked_item], model="gemini-live-2.5-flash" + ) + + await streaming.client_ack_messages() + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief." + assert "zebrafruit" not in json.dumps(frames[1]) + errors: Final = [event["error"]["type"] for event in _events_sent_to_client(streaming) if event["type"] == "error"] + assert errors == ["guardrail_violation"] + + +class _SteppedClock: + def __init__(self, epoch: float) -> None: + self.now = epoch + + def __call__(self) -> float: + return self.now + + +def _client_ws_receiving(*steps: Callable[[], Coroutine[object, object, str]]) -> MagicMock: + remaining: Final = iter(steps) + + async def receive_text() -> str: + return await next(remaining)() + + client_ws: Final = MagicMock() + client_ws.receive_text = receive_text + client_ws.send_text = AsyncMock() + return client_ws + + +@pytest.mark.asyncio +@pytest.mark.parametrize("epoch", [0.0, 1e9]) +async def test_default_mode_sends_the_held_setup_when_no_content_arrives_before_the_deadline(monkeypatch, epoch): + """Vertex AI refuses a setup sent minutes after connecting, so an idle session + sends what it holds at the deadline and answers any later session.update as late.""" + monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False) + from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig + + clock: Final = _SteppedClock(epoch) + setup_sent: Final = asyncio.Event() + + async def instructions_update_then_the_deadline_passes() -> str: + clock.now += REALTIME_LAZY_SETUP_DEADLINE_SECONDS + return json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}) + + async def late_update_once_the_setup_is_out() -> str: + await setup_sent.wait() + return json.dumps( + {"type": "session.update", "event_id": "evt_late", "session": {"instructions": "Answer in French."}} + ) + + async def disconnect() -> str: + raise ConnectionClosed(None, None) + + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock(side_effect=lambda frame: setup_sent.set()) + streaming: Final = RealTimeStreaming( + _client_ws_receiving(instructions_update_then_the_deadline_passes, late_update_once_the_setup_is_out, disconnect), + backend_ws, + MagicMock(), + provider_config=GeminiRealtimeConfig(), + model="gemini-live-2.5-flash", + clock=clock, + ) + + await asyncio.wait_for(streaming.client_ack_messages(), timeout=5) + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup"] + assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief." + events: Final = _events_sent_to_client(streaming) + assert [(event["type"], event.get("error", {}).get("code")) for event in events] == [ + ("session.updated", None), + ("error", "session_update_not_applied"), + ] + + +@pytest.mark.asyncio +async def test_default_mode_refuses_held_session_updates_past_the_size_cap(monkeypatch): + def instructions_update(text: str) -> str: + return json.dumps({"type": "session.update", "event_id": f"evt_{text[0]}", "session": {"instructions": text}}) + + half_the_cap: Final = RealTimeStreaming._MAX_BUFFERED_BYTES // 2 + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, + [instructions_update("a" * half_the_cap), instructions_update("b" * half_the_cap), _TEXT_ITEM_CREATE, _RESPONSE_CREATE], + model="gemini-live-2.5-flash", + ) + + await streaming.client_ack_messages() + + events: Final = _events_sent_to_client(streaming) + assert [(event["type"], event.get("error", {}).get("event_id")) for event in events] == [ + ("session.updated", None), + ("error", "evt_b"), + ] + assert events[1]["error"]["code"] == "session_update_too_large" + setup: Final = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"] + assert setup["systemInstruction"]["parts"][0]["text"].startswith("a") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("defer_setup", [False, True]) +async def test_a_session_update_between_the_first_items_and_response_create_joins_the_setup(monkeypatch, defer_setup): + modalities_update: Final = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}}) + instructions_update: Final = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, + [modalities_update, _TEXT_ITEM_CREATE, instructions_update, _RESPONSE_CREATE], + model="gemini-live-2.5-flash", + defer_setup=defer_setup, + ) + + await streaming.client_ack_messages() + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"] + assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief." + assert [event["type"] for event in _events_sent_to_client(streaming)] == ["session.updated", "session.updated"] + + +@pytest.mark.asyncio +async def test_default_mode_applies_a_session_update_that_arrives_while_audio_streams_in(monkeypatch): + modalities_update: Final = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}}) + append: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"}) + instructions_update: Final = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}) + commit: Final = json.dumps({"type": "input_audio_buffer.commit"}) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [modalities_update, append, instructions_update, commit], model="gemini-live-2.5-flash" + ) + + await streaming.client_ack_messages() + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames][:2] == ["setup", "realtimeInput"] + assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief." + assert frames[1]["realtimeInput"]["audio"]["data"] == "AAAA" + assert "error" not in [event["type"] for event in _events_sent_to_client(streaming)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("epoch", [0.0, 1e9]) +async def test_held_items_go_out_with_the_setup_when_no_response_create_follows(monkeypatch, epoch): + monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False) + from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig + + item_sent: Final = asyncio.Event() + + async def instructions_update() -> str: + return json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}) + + async def text_item() -> str: + return _TEXT_ITEM_CREATE + + async def disconnect_once_the_item_is_out() -> str: + await item_sent.wait() + raise ConnectionClosed(None, None) + + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock( + side_effect=lambda frame: item_sent.set() if "clientContent" in json.loads(frame) else None + ) + streaming: Final = RealTimeStreaming( + _client_ws_receiving(instructions_update, text_item, disconnect_once_the_item_is_out), + backend_ws, + MagicMock(), + provider_config=GeminiRealtimeConfig(), + model="gemini-live-2.5-flash", + lazy_setup_content_grace_seconds=0, + clock=_SteppedClock(epoch), + ) + + await asyncio.wait_for(streaming.client_ack_messages(), timeout=5) + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief." + + +@pytest.mark.asyncio +async def test_default_mode_acknowledges_an_unchanged_session_update_after_content(monkeypatch): + """Pipecat resends its unchanged session.update right after its first + conversation items and treats any error event as fatal.""" + session_update: Final = json.dumps( + {"type": "session.update", "session": {"modalities": ["text"], "instructions": "Be brief."}} + ) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [session_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE, session_update], model="gemini-live-2.5-flash" + ) + + await streaming.client_ack_messages() + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + events: Final = _events_sent_to_client(streaming) + assert [event["type"] for event in events] == ["session.updated", "session.updated"] + assert events[0]["session"] == events[1]["session"] + + +@pytest.mark.asyncio +async def test_deferred_setup_merges_session_updates_until_the_first_content(monkeypatch): + """Deferred setup merges every session.update sent before the first content, + the same as the default mode, and answers a changed one after it as late.""" + monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False) + from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig + + client_ws: Final = MagicMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps({"type": "session.update", "session": {"modalities": ["text"]}}), + json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}), + _TEXT_ITEM_CREATE, + _RESPONSE_CREATE, + json.dumps( + {"type": "session.update", "event_id": "evt_late", "session": {"instructions": "Answer in French."}} + ), + ConnectionClosed(None, None), + ] + ) + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, MagicMock(), provider_config=GeminiRealtimeConfig(), model="gemini-live-2.5-flash" + ) + + await streaming.client_ack_messages() + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"] + assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief." + errors: Final = [event["error"] for event in _events_sent_to_client(streaming) if event["type"] == "error"] + assert [(error["code"], error["event_id"]) for error in errors] == [("session_update_not_applied", "evt_late")] + + +@pytest.mark.asyncio +async def test_deferred_setup_sends_the_default_setup_before_text_content_that_comes_first(monkeypatch): + """Gemini Live closes the socket on content sent before any setup, so text that arrives + before a session.update in deferred mode gets the model's default setup first.""" + monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False) + from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig + + client_ws: Final = MagicMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + _TEXT_ITEM_CREATE, + json.dumps({"type": "session.update", "event_id": "evt_late", "session": {"modalities": ["text"]}}), + ConnectionClosed(None, None), + ] + ) + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, MagicMock(), provider_config=GeminiRealtimeConfig(), model="gemini-live-2.5-flash" + ) + + await streaming.client_ack_messages() + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + errors: Final = [event["error"] for event in _events_sent_to_client(streaming) if event["type"] == "error"] + assert [(error["code"], error["event_id"]) for error in errors] == [("session_update_not_applied", "evt_late")] + + +@pytest.mark.asyncio +async def test_deferred_setup_sends_a_session_update_that_follows_buffered_audio_at_once(monkeypatch): + append: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"}) + commit: Final = json.dumps({"type": "input_audio_buffer.commit"}) + session_update: Final = json.dumps({"type": "session.update", "session": {"instructions": "Be a pirate."}}) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, [append, commit, session_update], model="gemini-live-2.5-flash", defer_setup=True + ) + + await streaming.client_ack_messages() + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup"] + assert "Be a pirate." in json.dumps(frames[0]["setup"]["systemInstruction"]) + + +@pytest.mark.parametrize("defer_setup", [False, True]) +@pytest.mark.asyncio +async def test_a_transcription_guardrail_refuses_gemini_audio_it_cannot_screen(monkeypatch, defer_setup): + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + session_update: Final = json.dumps( + {"type": "session.update", "session": {"turn_detection": {"type": "server_vad", "create_response": False}}} + ) + append: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"}) + commit: Final = json.dumps({"type": "input_audio_buffer.commit", "event_id": "commit_1"}) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, + [session_update, append, append, commit, json.dumps({"type": "response.create"}), _TEXT_ITEM_CREATE], + model="gemini-live-2.5-flash", + defer_setup=defer_setup, + ) + + await streaming.client_ack_messages() + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + refusals: Final = [ + event["error"]["event_id"] + for event in _events_sent_to_client(streaming) + if event["type"] == "error" and event["error"]["code"] == "audio_input_not_screenable" + ] + assert refusals == [None, "commit_1"] + + +@pytest.mark.parametrize("defer_setup", [False, True]) +@pytest.mark.parametrize("control_type", ["response.cancel", "conversation.item.delete", "input_audio_buffer.clear"]) +@pytest.mark.asyncio +async def test_a_control_event_before_session_update_leaves_the_setup_open(monkeypatch, defer_setup, control_type): + control_event: Final = json.dumps({"type": control_type, "item_id": "item_x"}) + session_update: Final = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}}) + streaming, backend_ws = _gemini_default_mode_streaming( + monkeypatch, + [control_event, session_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], + model="gemini-live-2.5-flash", + defer_setup=defer_setup, + ) + + await streaming.client_ack_messages() + + frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"] + assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"] + assert [event["type"] for event in _events_sent_to_client(streaming) if event["type"] == "error"] == [] + + @pytest.mark.asyncio async def test_deferred_setup_buffers_audio_until_backend_setup_complete(monkeypatch): """Pipecat may send audio before session.update when setup is deferred.""" @@ -2533,14 +3173,18 @@ async def test_deferred_setup_buffers_audio_until_backend_setup_complete(monkeyp @pytest.mark.asyncio -async def test_deferred_setup_sends_session_update_before_buffered_audio(monkeypatch): +async def test_deferred_setup_sends_the_held_session_update_when_more_audio_arrives(monkeypatch): monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False) from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig client_ws = MagicMock() + client_ws.send_text = AsyncMock() audio_msg = json.dumps({"type": "input_audio_buffer.append", "audio": "AA=="}) session_update = json.dumps({"type": "session.update", "session": {"modalities": ["audio"]}}) - client_ws.receive_text = AsyncMock(side_effect=[audio_msg, session_update, ConnectionClosed(None, None)]) + later_audio_msg = json.dumps({"type": "input_audio_buffer.append", "audio": "AQ=="}) + client_ws.receive_text = AsyncMock( + side_effect=[audio_msg, session_update, later_audio_msg, ConnectionClosed(None, None)] + ) backend_ws = MagicMock() backend_ws.send = AsyncMock() logging_obj: Final = _mock_logging() @@ -2560,7 +3204,7 @@ async def test_deferred_setup_sends_session_update_before_buffered_audio(monkeyp sent_payload = json.loads(backend_ws.send.await_args_list[0].args[0]) assert "setup" in sent_payload assert "realtimeInput" not in sent_payload - assert streaming._pending_messages_until_setup == [audio_msg] + assert streaming._pending_messages_until_setup == [audio_msg, later_audio_msg] @pytest.mark.asyncio @@ -3523,6 +4167,24 @@ async def test_upstream_refusal_after_a_synthetic_session_created_still_logs_a_f assert session.logging.logged_sessions == () +class _LazySetupRealtimeConfig(_PassthroughRealtimeConfig): + def builds_setup_from_session_update(self) -> bool: + return True + + +@pytest.mark.asyncio +async def test_upstream_outage_before_a_held_setup_was_sent_logs_the_session_not_a_refusal(): + upstream_close: Final = ConnectionClosed(Close(1011, "upstream outage"), None) + session: Final = _relay_session( + _client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close), _LazySetupRealtimeConfig() + ) + + await session.run() + + assert session.logging.logged_failures == () + assert len(session.logging.logged_sessions) == 1 + + @pytest.mark.asyncio async def test_upstream_close_after_relayed_events_still_logs_the_session_as_success(): client_ws: Final = _client_ws_that_never_sends() diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 794ac2d26d5..23d7feb2557 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -14,6 +14,7 @@ import httpcore import httpx import pytest import respx +import websockets.exceptions from botocore.credentials import RefreshableCredentials from h2.events import ConnectionTerminated @@ -2787,6 +2788,194 @@ async def test_async_realtime_error_event_send_failure_still_closes(): assert client_ws.events == [("close", (1011, "Internal server error: vertex token refresh exploded"))] +_GEMINI_LIVE_SETUP_COMPLETE = json.dumps({"setupComplete": {}}) +_GEMINI_LIVE_TEXT_TURN = json.dumps({"serverContent": {"modelTurn": {"parts": [{"text": "Blue."}]}}}) +_GEMINI_LIVE_TURN_COMPLETE = json.dumps({"serverContent": {"turnComplete": True}}) + + +class _ScriptedRealtimeBackend: + """Backend socket that answers the first setup frame with the scripted replies + and closes once the client is done, so both bridge loops run to completion.""" + + def __init__(self, replies_after_setup): + self.sent = [] + self.closed = asyncio.Event() + self._replies_after_setup = list(replies_after_setup) + self._inbox = asyncio.Queue() + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return None + + async def send(self, payload): + self.sent.append(payload) + if "setup" in json.loads(payload): + for reply in self._replies_after_setup: + self._inbox.put_nowait(reply) + + def close_after_replies(self): + self._inbox.put_nowait(None) + + async def recv(self, decode=True): + reply = await self._inbox.get() + if reply is None: + self.closed.set() + raise websockets.exceptions.ConnectionClosedOK(None, None) + return reply + + +class _ScriptedRealtimeClient: + def __init__(self, messages, backend): + self.scope = {"headers": []} + self.exceptions = websockets.exceptions + self.sent = [] + self._messages = list(messages) + self._backend = backend + + async def send_text(self, payload): + self.sent.append(payload) + + async def receive_text(self): + if self._messages: + return self._messages.pop(0) + self._backend.close_after_replies() + await self._backend.closed.wait() + raise websockets.exceptions.ConnectionClosedOK(None, None) + + async def close(self, code=None, reason=None): + return None + + +def _realtime_logging_obj(): + logging_obj = Mock() + logging_obj.litellm_trace_id = "trace-lit-6974" + logging_obj.async_success_handler = AsyncMock() + logging_obj.dispatch_failure_handlers = AsyncMock() + return logging_obj + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model, session, backend_modality, client_modality, backend_turn, drop_params", + [ + ( + "gemini-live-2.5-flash", + {"type": "realtime", "output_modalities": ["text"]}, + "TEXT", + "text", + [_GEMINI_LIVE_TEXT_TURN], + False, + ), + ("gemini-live-2.5-flash", {"modalities": ["text"]}, "TEXT", "text", [_GEMINI_LIVE_TEXT_TURN], False), + ( + "gemini-live-2.5-flash-native-audio", + {"type": "realtime", "output_modalities": ["text"]}, + "AUDIO", + "audio", + [], + True, + ), + ], +) +async def test_async_realtime_vertex_first_session_update_builds_the_only_setup( + model, session, backend_modality, client_modality, backend_turn, drop_params: bool, monkeypatch +): + """Regression for text-only Vertex Live sessions answered in audio: the + connect-time default setup (AUDIO) used to go out before the client's + session.update, and Gemini accepts only one setup per connection, so the + requested modality was dropped and billed at the audio rate. The setup now + goes out on the client's first message, built from session.update when that + comes first. Under drop_params an audio-only model still gets AUDIO, and the + client learns the downgrade from session.updated.""" + from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig + + monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False) + backend = _ScriptedRealtimeBackend([_GEMINI_LIVE_SETUP_COMPLETE, *backend_turn, _GEMINI_LIVE_TURN_COMPLETE]) + client_ws = _ScriptedRealtimeClient( + [ + json.dumps({"type": "session.update", "session": session}), + json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "What color is the sky? One word."}], + }, + } + ), + json.dumps({"type": "response.create"}), + ], + backend, + ) + handler = BaseLLMHTTPHandler() + + with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(return_value=backend)): + await asyncio.wait_for( + handler.async_realtime( + model=model, + websocket=client_ws, + logging_obj=_realtime_logging_obj(), + provider_config=VertexAIRealtimeConfig( + access_token="fake-token", project="fake-project", location="us-central1", drop_params=drop_params + ), + headers={}, + ), + timeout=10, + ) + + backend_frames = [json.loads(frame) for frame in backend.sent] + setups = [frame["setup"] for frame in backend_frames if "setup" in frame] + assert len(setups) == 1 + assert "setup" in backend_frames[0] + assert setups[0]["generationConfig"]["responseModalities"] == [backend_modality] + assert any("clientContent" in frame for frame in backend_frames[1:]) + + client_events = [json.loads(frame) for frame in client_ws.sent] + event_types = [event["type"] for event in client_events] + assert event_types[:2] == ["session.created", "session.updated"] + assert client_events[1]["session"]["modalities"] == [client_modality] + if backend_turn: + assert "response.output_text.delta" in event_types + assert "response.output_audio_transcript.delta" not in event_types + + +@pytest.mark.asyncio +async def test_async_realtime_sends_setup_at_connect_when_provider_cannot_build_it_from_session_update(): + """Bedrock Nova Sonic style providers keep the connect-time setup and get no + synthetic session.created.""" + setup = json.dumps({"event": {"sessionStart": {}}}) + backend = _ScriptedRealtimeBackend([]) + client_ws = _ScriptedRealtimeClient([], backend) + provider_config = Mock() + provider_config.get_complete_url.return_value = "wss://backend.example/live" + provider_config.validate_environment.return_value = {} + provider_config.requires_session_configuration.return_value = True + provider_config.builds_setup_from_session_update.return_value = False + provider_config.session_configuration_request.return_value = setup + provider_config.unbilled_usage_on_session_close.return_value = None + provider_config.open_backend = AsyncMock(return_value=None) + provider_config.wrap_backend.side_effect = lambda opened: opened + handler = BaseLLMHTTPHandler() + + with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(return_value=backend)): + await asyncio.wait_for( + handler.async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=client_ws, + logging_obj=_realtime_logging_obj(), + provider_config=provider_config, + headers={}, + ), + timeout=10, + ) + + assert backend.sent == [setup] + assert client_ws.sent == [] + + class _JSONBodyAudioTranscriptionConfig(BaseAudioTranscriptionConfig): def get_supported_openai_params(self, model): return [] diff --git a/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py index 017f279d6c4..460a50ed61a 100644 --- a/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,3 +1,4 @@ +import base64 import json from collections.abc import Mapping from typing import Final, cast @@ -7,7 +8,7 @@ import pytest import litellm -from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig +from litellm.llms.gemini.realtime.transformation import GEMINI_LIVE_DOCUMENTED_VOICES, GeminiRealtimeConfig from litellm.types.llms.gemini import BidiGenerateContentServerMessage @@ -439,6 +440,51 @@ def test_gemini_realtime_session_update_with_tools(): assert "parameters" in function_decl +def _schema_nodes(schema: object): + if isinstance(schema, dict): + yield schema + for value in schema.values(): + yield from _schema_nodes(value) + elif isinstance(schema, list): + for value in schema: + yield from _schema_nodes(value) + + +def test_gemini_session_update_converts_flat_realtime_tool_schemas_to_gemini_schemas(): + session_update: Final = { + "type": "session.update", + "session": { + "tools": [ + { + "type": "function", + "name": "lookup_weather", + "description": "Look up the weather in a city.", + "parameters": { + "type": "object", + "$defs": { + "City": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]} + }, + "properties": {"city": {"$ref": "#/$defs/City"}, "note": {"type": ["string", "null"]}}, + "required": ["city"], + }, + } + ] + }, + } + + messages: Final = GeminiRealtimeConfig().transform_realtime_request( + json.dumps(session_update), "gemini-2.5-flash", session_configuration_request=None + ) + + declaration: Final = json.loads(messages[0])["setup"]["tools"][0]["function_declarations"][0] + assert declaration["name"] == "lookup_weather" + parameters: Final = declaration["parameters"] + assert parameters["properties"]["city"]["properties"]["name"]["type"] == "string" + nodes: Final = tuple(_schema_nodes(parameters)) + assert not any("$defs" in node or "$ref" in node for node in nodes) + assert all(isinstance(node["type"], str) for node in nodes if "type" in node) + + def test_gemini_session_update_defaults_to_audio_modality(): config = GeminiRealtimeConfig() @@ -613,6 +659,35 @@ def test_gemini_audio_only_live_models_drop_text_from_text_audio_combo(patch_gem assert setup["generationConfig"]["responseModalities"] == ["AUDIO"] +@pytest.mark.parametrize("output_modalities", ["text", ["text"]]) +def test_gemini_session_update_accepts_a_single_string_output_modality(output_modalities): + """Regression: a bare string was iterated per character, so output_modalities + "text" reached Vertex as responseModalities ["T"], which it rejects with 1007.""" + config = GeminiRealtimeConfig() + session_update = {"type": "session.update", "session": {"type": "realtime", "output_modalities": output_modalities}} + + messages = config.transform_realtime_request( + json.dumps(session_update), "gemini-live-2.5-flash", session_configuration_request=None + ) + + assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == ["TEXT"] + + +@pytest.mark.parametrize( + ("output_modalities", "expected"), + [(["video", "text"], ["TEXT"]), (["video"], ["AUDIO"]), (None, ["AUDIO"])], +) +def test_gemini_session_update_drops_modalities_live_cannot_produce(output_modalities, expected): + config = GeminiRealtimeConfig() + session_update = {"type": "session.update", "session": {"type": "realtime", "output_modalities": output_modalities}} + + messages = config.transform_realtime_request( + json.dumps(session_update), "gemini-live-2.5-flash", session_configuration_request=None + ) + + assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == expected + + def test_gemini_non_live_model_preserves_text_modality(): config = GeminiRealtimeConfig() session_update = { @@ -1396,22 +1471,71 @@ def test_gemini_input_audio_buffer_clear_is_local_noop(): assert messages == [] -def test_gemini_input_audio_buffer_commit_maps_to_activity_end_when_manual_vad(): +_MANUAL_VAD_SETUP = json.dumps({"setup": {"realtimeInputConfig": {"automaticActivityDetection": {"disabled": True}}}}) +_APPEND = json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(b"\x00\x01" * 80).decode()}) + + +def _realtime_input_keys(config: GeminiRealtimeConfig, client_event: str, setup: str | None) -> list[str]: + frames = config.transform_realtime_request(client_event, "gemini-live-2.5-flash-native-audio", setup) + return [key for frame in frames for key in json.loads(frame)["realtimeInput"]] + + +@pytest.mark.parametrize("turn_end", ["input_audio_buffer.commit", "response.create"]) +def test_gemini_manual_vad_wraps_each_audio_turn_in_activity_start_and_end(turn_end): + """With activity detection off, Gemini Live answers only after activityEnd closes an + activityStart, and it closes the socket on any other shape of those markers.""" config = GeminiRealtimeConfig() - setup = { - "setup": { - "realtimeInputConfig": { - "automaticActivityDetection": {"disabled": True}, - } - } - } - messages = config.transform_realtime_request( - json.dumps({"type": "input_audio_buffer.commit"}), - "gemini-live-2.5-flash-native-audio", - session_configuration_request=json.dumps(setup), + end = json.dumps({"type": turn_end}) + + turns = [ + [_realtime_input_keys(config, event, _MANUAL_VAD_SETUP) for event in (_APPEND, _APPEND, end, end)] + for _ in range(2) + ] + + assert turns == [[["activityStart", "audio"], ["audio"], ["activityEnd"], []]] * 2 + + +def test_gemini_manual_vad_activity_markers_are_empty_messages(): + config = GeminiRealtimeConfig() + commit = json.dumps({"type": "input_audio_buffer.commit"}) + + first_append = config.transform_realtime_request(_APPEND, "gemini-live-2.5-flash", _MANUAL_VAD_SETUP) + turn_end = config.transform_realtime_request(commit, "gemini-live-2.5-flash", _MANUAL_VAD_SETUP) + + assert json.loads(first_append[0]) == {"realtimeInput": {"activityStart": {}}} + assert [json.loads(frame) for frame in turn_end] == [{"realtimeInput": {"activityEnd": {}}}] + + +@pytest.mark.parametrize( + ("session", "expected_keys"), + [ + ({"turn_detection": None}, [["activityStart", "audio"], ["activityEnd"]]), + ({"type": "realtime", "audio": {"input": {"turn_detection": None}}}, [["activityStart", "audio"], ["activityEnd"]]), + ({"type": "realtime", "audio": {"input": {}}}, [["audio"], ["audioStreamEnd"]]), + ], +) +def test_gemini_null_turn_detection_makes_the_client_mark_each_turn(session, expected_keys): + """OpenAI's turn_detection null turns server VAD off, so the client commits each turn itself.""" + config = GeminiRealtimeConfig() + setup = config.transform_realtime_request( + json.dumps({"type": "session.update", "session": session}), "gemini-live-2.5-flash-native-audio", None ) - assert len(messages) == 1 - assert json.loads(messages[0]) == {"realtimeInput": {"activityEnd": True}} + commit = json.dumps({"type": "input_audio_buffer.commit"}) + + keys = [_realtime_input_keys(config, event, setup[0]) for event in (_APPEND, commit)] + + assert keys == expected_keys + + +def test_gemini_automatic_vad_sends_no_activity_markers(): + config = GeminiRealtimeConfig() + commit = json.dumps({"type": "input_audio_buffer.commit"}) + response_create = json.dumps({"type": "response.create"}) + + keys = [_realtime_input_keys(config, event, None) for event in (_APPEND, commit)] + + assert keys == [["audio"], ["audioStreamEnd"]] + assert config.transform_realtime_request(response_create, "gemini-live-2.5-flash", None) == [] def test_gemini_tool_call_id_to_name_evicts_oldest_when_capped(): @@ -1883,13 +2007,75 @@ def test_map_openai_params_drops_stock_voice_case_insensitively(): """ cfg = GeminiRealtimeConfig() - dropped = cfg.map_openai_params(optional_params={}, non_default_params={"voice": "Alloy"}) + dropped = cfg.map_openai_params(optional_params={}, non_default_params={"voice": "Alloy"}, model="gemini-3.8-live") assert "speechConfig" not in dropped.get("generationConfig", {}) - passthrough = cfg.map_openai_params(optional_params={}, non_default_params={"voice": "Kore"}) + passthrough = cfg.map_openai_params( + optional_params={}, non_default_params={"voice": "Kore"}, model="gemini-3.8-live" + ) assert passthrough["generationConfig"]["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" +_DOCUMENTED_VOICE: Final = sorted(GEMINI_LIVE_DOCUMENTED_VOICES)[0].capitalize() + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_google_ai_studio_forwards_a_documented_voice(drop_params): + config: Final = GeminiRealtimeConfig(drop_params=drop_params) + mapped: Final = config.map_openai_params( + optional_params={}, non_default_params={"voice": _DOCUMENTED_VOICE}, model="gemini-3.8-live" + ) + + speech_config: Final = mapped["generationConfig"]["speechConfig"] + assert speech_config["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == _DOCUMENTED_VOICE + assert config.setup_session_update_rejection("gemini-3.8-live", {"voice": _DOCUMENTED_VOICE}, setup_sent=False) is None + + +@pytest.mark.parametrize( + ("session", "param"), + [ + ({"voice": "Nova"}, "session.voice"), + ({"type": "realtime", "audio": {"output": {"voice": "Nova"}}}, "session.audio.output.voice"), + ], +) +def test_google_ai_studio_rejects_an_undocumented_voice_on_the_session_update(session, param): + rejection: Final = GeminiRealtimeConfig().setup_session_update_rejection( + "gemini-2.5-flash-native-audio-preview-12-2025", session, setup_sent=False + ) + + assert rejection is not None + assert (rejection.code, rejection.param) == ("invalid_value", param) + + +def _setup_speech_config(config: GeminiRealtimeConfig, model: str, session: dict[str, object]) -> object: + messages: Final = config.transform_realtime_request( + json.dumps({"type": "session.update", "session": session}), model, session_configuration_request=None + ) + return json.loads(messages[0])["setup"].get("generationConfig", {}).get("speechConfig") + + +@pytest.mark.parametrize( + "session", + [{"voice": "Nova"}, {"type": "realtime", "audio": {"output": {"voice": "Nova"}}}], +) +def test_google_ai_studio_drops_an_undocumented_voice_with_drop_params(session): + config: Final = GeminiRealtimeConfig(drop_params=True) + model: Final = "gemini-2.5-flash-native-audio-preview-12-2025" + + assert config.setup_session_update_rejection(model, session, setup_sent=False) is None + assert _setup_speech_config(config, model, session) is None + + +@pytest.mark.parametrize("drop_params", [False, True]) +@pytest.mark.parametrize("model", ["gemini-3.8-live", "gemini-3.1-flash-live-preview"]) +def test_google_ai_studio_forwards_any_voice_to_models_that_accept_it(model, drop_params): + config: Final = GeminiRealtimeConfig(drop_params=drop_params) + + assert config.setup_session_update_rejection(model, {"voice": "Nova"}, setup_sent=False) is None + speech_config: Final = _setup_speech_config(config, model, {"voice": "Nova"}) + assert speech_config == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Nova"}}} + + @pytest.fixture(autouse=False) def patch_gemini_transcribe_live_cost_map_entry(monkeypatch): """Inject the gemini-3.5-transcribe-live registry entry locally. @@ -1954,6 +2140,55 @@ def test_gemini_transcribe_live_coerces_audio_modality_to_text(modalities, patch assert setup["generationConfig"]["responseModalities"] == ["TEXT"] +@pytest.mark.parametrize("model", ["gemini-3.5-transcribe-live", "gemini/gemini-3.5-transcribe-live"]) +def test_gemini_transcribe_live_synthetic_session_created_reports_text_output( + model, patch_gemini_transcribe_live_cost_map_entry +): + session = GeminiRealtimeConfig().transform_session_created_event(model, "sess_transcribe")["session"] + + assert session["modalities"] == ["text"] + + +@pytest.mark.parametrize( + ("model", "modalities", "producible"), + [ + ("gemini-3.5-transcribe-live", ["audio"], False), + ("gemini-3.5-transcribe-live", ["audio", "text"], True), + ("gemini-live-2.5-flash-native-audio", ["text"], False), + ("gemini-live-2.5-flash-native-audio", ["text", "audio"], True), + ], +) +@pytest.mark.parametrize("drop_params", [False, True]) +def test_gemini_live_rejects_only_output_modalities_the_model_cannot_produce_at_all( + model, + modalities, + producible, + drop_params, + patch_gemini_transcribe_live_cost_map_entry, + patch_gemini_audio_cost_map_entries, +): + rejection = GeminiRealtimeConfig(drop_params=drop_params).setup_session_update_rejection( + model, {"output_modalities": modalities}, setup_sent=False + ) + + assert (rejection is None) is (drop_params or producible) + assert rejection is None or rejection.param == "session.output_modalities" + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_gemini_provider_realtime_config_carries_drop_params(drop_params): + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + model = "gemini-live-2.5-flash-native-audio" + config = ProviderConfigManager.get_provider_realtime_config( + model=model, provider=LlmProviders.GEMINI, drop_params=drop_params + ) + + rejection = config.setup_session_update_rejection(model, {"modalities": ["text"]}, setup_sent=False) + assert (rejection is None) is drop_params + + def test_gemini_chat_model_with_text_output_modalities_keeps_audio_eager_setup( patch_gemini_transcribe_live_cost_map_entry, ): @@ -2169,6 +2404,47 @@ def test_unbilled_usage_on_session_close_flushes_trailing_audio(patch_gemini_tra assert config.unbilled_usage_on_session_close("gemini-3.5-transcribe-live") is None +def test_gemini_realtime_two_output_modalities_collapse_to_audio(): + """Gemini Live closes the socket with 1007 when the setup names two response + modalities, so a GA session.update asking for ["text", "audio"] must reach the + backend as AUDIO only, the same choice the beta shim makes.""" + config = GeminiRealtimeConfig() + session_update = { + "type": "session.update", + "session": {"output_modalities": ["text", "audio"], "instructions": "Be brief."}, + } + + messages = config.transform_realtime_request( + json.dumps(session_update), + "gemini-live-2.5-flash", + session_configuration_request=None, + ) + + assert len(messages) == 1 + assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == ["AUDIO"] + + +def test_gemini_realtime_inf_max_response_output_tokens_is_not_forwarded(): + """OpenAI's default max_response_output_tokens is the string "inf", which Gemini + rejects with 1007 (TYPE_INT32), so only integer limits map to maxOutputTokens.""" + config = GeminiRealtimeConfig() + + def generation_config_for(limit): + session_update = { + "type": "session.update", + "session": {"modalities": ["text"], "max_response_output_tokens": limit}, + } + messages = config.transform_realtime_request( + json.dumps(session_update), + "gemini-live-2.5-flash", + session_configuration_request=None, + ) + return json.loads(messages[0])["setup"]["generationConfig"] + + assert "maxOutputTokens" not in generation_config_for("inf") + assert generation_config_for(200)["maxOutputTokens"] == 200 + + def _grounded_live_frame(grounding_metadata: Mapping[str, object] | None) -> Mapping[str, object]: """One Live server frame. Grounding metadata and usageMetadata arrive together, as Vertex sends them.""" from typing import Final diff --git a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index 0bfff870d15..7f0106616ce 100644 --- a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -18,6 +18,7 @@ import websockets.exceptions # registers websockets.exceptions on the websocket import litellm from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import OPENAI_TO_GEMINI_TTS_VOICES from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig # --------------------------------------------------------------------------- @@ -179,6 +180,25 @@ def test_vertex_audio_only_live_model_coerces_text_modality_to_audio( assert setup["generationConfig"]["responseModalities"] == ["AUDIO"] +@pytest.mark.parametrize("drop_params", [False, True]) +def test_vertex_realtime_config_carries_drop_params(drop_params, patch_native_audio_cost_map_entry): + from litellm.llms.vertex_ai.realtime.transformation import vertex_realtime_config + + cfg: Final = vertex_realtime_config( + _NATIVE_AUDIO_MODEL, + access_token="tok", + resolve_access_token=AsyncMock(return_value="tok"), + project="my-proj", + location="us-central1", + drop_params=drop_params, + ) + + rejection: Final = cfg.setup_session_update_rejection( + _NATIVE_AUDIO_MODEL, {"modalities": ["text"]}, setup_sent=False + ) + assert (rejection is None) is drop_params + + def test_vertex_session_update_normalizes_ga_remapped_fields( patch_native_audio_cost_map_entry, ): @@ -253,6 +273,26 @@ def test_vertex_keeps_turn_detection_tuning_when_auto_response_is_off(): } + +@pytest.mark.parametrize( + ("session", "disabled"), + [ + ({"turn_detection": None}, True), + ({"type": "realtime", "audio": {"input": {"turn_detection": None}}}, True), + ({"type": "realtime", "audio": {"input": {}}}, False), + ], +) +def test_vertex_turns_activity_detection_off_for_null_turn_detection(session, disabled): + """OpenAI's turn_detection null is push-to-talk, so Vertex must wait for the client's activity markers.""" + cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1") + messages: Final = cfg.transform_realtime_request( + json.dumps({"type": "session.update", "session": session}), + _NATIVE_AUDIO_MODEL, + session_configuration_request=None, + ) + setup: Final = json.loads(messages[0])["setup"] + assert setup["realtimeInputConfig"]["automaticActivityDetection"]["disabled"] is disabled + # --------------------------------------------------------------------------- # Round-trip test: text-in / text-out via RealTimeStreaming # --------------------------------------------------------------------------- @@ -607,3 +647,73 @@ def test_vertex_native_audio_unmapped_voice_passes_through(patch_native_audio_co generation_config = json.loads(messages[0])["setup"]["generationConfig"] assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + + +def test_vertex_session_update_two_output_modalities_collapse_to_audio(): + """Vertex Live rejects a setup naming two response modalities, so a GA + ["text", "audio"] request must be sent as AUDIO only.""" + cfg = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="global") + session_update = { + "type": "session.update", + "session": {"output_modalities": ["text", "audio"], "instructions": "Be brief."}, + } + + messages = cfg.transform_realtime_request( + json.dumps(session_update), + "gemini-live-2.5-flash", + session_configuration_request=None, + ) + + assert len(messages) == 1 + assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == ["AUDIO"] + + +@pytest.mark.parametrize( + ("session", "param"), + [ + ({"voice": "Nova"}, "session.voice"), + ({"type": "realtime", "audio": {"output": {"voice": "Nova"}}}, "session.audio.output.voice"), + ], +) +def test_vertex_rejects_an_undocumented_voice_on_the_session_update(session, param): + cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1") + + rejection: Final = cfg.setup_session_update_rejection("gemini-live-2.5-flash", session, setup_sent=False) + + assert rejection is not None + assert (rejection.code, rejection.param) == ("invalid_value", param) + + +@pytest.mark.parametrize("voice", ["Kore", "alloy"]) +def test_vertex_accepts_documented_and_openai_stock_voices_on_the_session_update(voice): + cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1") + + assert cfg.setup_session_update_rejection("gemini-live-2.5-flash", {"voice": voice}, setup_sent=False) is None + + +def test_vertex_drops_an_undocumented_voice_with_drop_params(): + cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1", drop_params=True) + session_update: Final = {"type": "session.update", "session": {"voice": "Nova", "instructions": "Be brief."}} + + assert cfg.setup_session_update_rejection("gemini-live-2.5-flash", session_update["session"], setup_sent=False) is None + messages: Final = cfg.transform_realtime_request( + json.dumps(session_update), "gemini-live-2.5-flash", session_configuration_request=None + ) + assert "speechConfig" not in json.loads(messages[0])["setup"].get("generationConfig", {}) + + +@pytest.mark.parametrize("drop_params", [False, True]) +@pytest.mark.parametrize("voice", ["fable", "onyx"]) +def test_vertex_sends_the_documented_voice_an_openai_tts_voice_maps_to(voice, drop_params): + cfg: Final = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1", drop_params=drop_params + ) + session_update: Final = {"type": "session.update", "session": {"voice": voice}} + + assert cfg.setup_session_update_rejection("gemini-live-2.5-flash", session_update["session"], setup_sent=False) is None + messages: Final = cfg.transform_realtime_request( + json.dumps(session_update), "gemini-live-2.5-flash", session_configuration_request=None + ) + assert json.loads(messages[0])["setup"]["generationConfig"]["speechConfig"] == { + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES[voice]}} + } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.test.tsx index 117796d658e..56135c779e9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.test.tsx @@ -173,6 +173,24 @@ describe("RealtimePlayground", () => { expect(update.session.type).toBe("realtime"); }); + it("starts recording on the configuration the session was created with", async () => { + const user = userEvent.setup(); + render(); + + await connect(user); + await act(async () => { + latestSocket().emit({ type: "session.created" }); + }); + await user.click(screen.getByTitle("Start recording")); + + const updates = latestSocket() + .sent.map((raw) => JSON.parse(raw)) + .filter((payload) => payload.type === "session.update"); + expect(updates).toHaveLength(2); + expect(updates[1]).toEqual(updates[0]); + expect(updates[0].session.turn_detection).toEqual({ type: "server_vad" }); + }); + it("sends what was typed and then asks for a response", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx index 01151203e65..408cbd58593 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx @@ -22,6 +22,20 @@ interface RealtimePlaygroundProps { selectedGuardrails?: string[]; } +const sessionUpdate = (voice: string) => + JSON.stringify({ + type: "session.update", + session: { + type: "realtime", + modalities: ["text", "audio"], + voice, + input_audio_format: "pcm16", + output_audio_format: "pcm16", + input_audio_transcription: { model: "gpt-4o-mini-transcribe" }, + turn_detection: { type: "server_vad" }, + }, + }); + const RealtimePlayground: React.FC = ({ accessToken, selectedModel, @@ -125,21 +139,7 @@ const RealtimePlayground: React.FC = ({ const type = data.type; if (type === "session.created") { - // GA: session.type is required ("realtime" | "transcription") - ws.send( - JSON.stringify({ - type: "session.update", - session: { - type: "realtime", - modalities: ["text", "audio"], - voice: selectedVoice, - input_audio_format: "pcm16", - output_audio_format: "pcm16", - input_audio_transcription: { model: "gpt-4o-mini-transcribe" }, - turn_detection: null, - }, - }), - ); + ws.send(sessionUpdate(selectedVoice)); } else if (type === "session.updated") { // session configured } else if ( @@ -231,22 +231,7 @@ const RealtimePlayground: React.FC = ({ const startRecording = useCallback(async () => { if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return; - // Switch to server VAD mode for voice input - // GA: session.type is required - wsRef.current.send( - JSON.stringify({ - type: "session.update", - session: { - type: "realtime", - modalities: ["text", "audio"], - voice: selectedVoice, - input_audio_format: "pcm16", - output_audio_format: "pcm16", - input_audio_transcription: { model: "gpt-4o-mini-transcribe" }, - turn_detection: { type: "server_vad" }, - }, - }), - ); + wsRef.current.send(sessionUpdate(selectedVoice)); try { const stream = await navigator.mediaDevices.getUserMedia({ audio: true }); @@ -301,7 +286,7 @@ const RealtimePlayground: React.FC = ({ } catch (err: any) { addMessage("status", `Microphone error: ${err.message}`); } - }, [addMessage]); + }, [addMessage, selectedVoice]); const stopRecording = useCallback(() => { processorRef.current?.disconnect(); @@ -317,21 +302,7 @@ const RealtimePlayground: React.FC = ({ if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return; if (configureSessionRef.current) return; configureSessionRef.current = true; - // GA: session.type is required - wsRef.current.send( - JSON.stringify({ - type: "session.update", - session: { - type: "realtime", - modalities: ["text", "audio"], - voice: selectedVoice, - input_audio_format: "pcm16", - output_audio_format: "pcm16", - input_audio_transcription: { model: "gpt-4o-mini-transcribe" }, - turn_detection: null, - }, - }), - ); + wsRef.current.send(sessionUpdate(selectedVoice)); }, [selectedVoice]); const sendTextMessage = useCallback(() => { From 373f1752d30c70f8d2c87022e677f4c0ec58e14e Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 02:21:39 -0700 Subject: [PATCH 15/64] fix(anthropic): send PDF image_url content to Claude as document blocks (#46984) * fix(anthropic): convert PDF image_url content to documents Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(prompt_templates): strip content-type parameters from fetched media Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(prompt_templates): format MIME normalization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../prompt_templates/common_utils.py | 6 +- .../prompt_templates/factory.py | 29 +++----- ...llm_core_utils_prompt_templates_factory.py | 74 +++++++++++++++++++ .../litellm_core_utils/test_image_handling.py | 31 ++++++++ 4 files changed, 119 insertions(+), 21 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index b3f89b33d1d..862e61cb631 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2023,12 +2023,14 @@ def infer_content_type_from_url_and_content( """ from litellm.litellm_core_utils.token_counter import get_image_type + media_type: Final = current_content_type.split(";", 1)[0].strip().lower() if current_content_type else None + # If we have a valid content type that's not generic, use it - if current_content_type and current_content_type not in [ + if media_type and media_type not in [ "binary/octet-stream", "application/octet-stream", ]: - return current_content_type + return media_type # Extension to MIME type mapping # Supports images, documents, and other common file types diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 834d26b4fe8..de896cde4dc 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -912,9 +912,9 @@ def create_anthropic_image_param( image_url_input: str | dict, format: str | None = None, is_bedrock_invoke: bool = False, -) -> AnthropicMessagesImageParam: +) -> AnthropicMessagesImageParam | AnthropicMessagesDocumentParam: """ - Create an AnthropicMessagesImageParam from an image URL input. + Create an Anthropic content block from an image URL input. Supports both URL references (for HTTP/HTTPS URLs) and base64 encoding. """ @@ -933,14 +933,7 @@ def create_anthropic_image_param( if is_bedrock_invoke or image_url.startswith("http://"): base64_url: Final = convert_url_to_base64(url=image_url) image_chunk = convert_to_anthropic_image_obj(openai_image_url=base64_url, format=format) - return AnthropicMessagesImageParam( - type="image", - source=AnthropicContentParamSource( - type="base64", - media_type=image_chunk["media_type"], - data=image_chunk["data"], - ), - ) + return _anthropic_content_element_factory(image_chunk) else: # HTTPS URL - pass directly for regular Anthropic return AnthropicMessagesImageParam( @@ -953,14 +946,7 @@ def create_anthropic_image_param( else: # Convert to base64 for data URIs or other formats image_chunk = convert_to_anthropic_image_obj(openai_image_url=image_url, format=format) - return AnthropicMessagesImageParam( - type="image", - source=AnthropicContentParamSource( - type="base64", - media_type=image_chunk["media_type"], - data=image_chunk["data"], - ), - ) + return _anthropic_content_element_factory(image_chunk) # The following XML functions will be deprecated once JSON schema support is available on Bedrock and Vertex @@ -1910,7 +1896,12 @@ def convert_to_anthropic_tool_result( anthropic_content_element=_anthropic_image_param, original_content_element=content, ) - anthropic_content_list.append(cast(AnthropicMessagesImageParam, _anthropic_image_param)) + anthropic_content_list.append( + cast( + AnthropicMessagesImageParam | AnthropicMessagesDocumentParam, + _anthropic_image_param, + ) + ) elif content["type"] == "tool_reference": anthropic_content_list.append(ToolReference(type="tool_reference", tool_name=content["tool_name"])) elif content["type"] == "file": diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 53be3608979..778c9d68be5 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1246,6 +1246,80 @@ def test_anthropic_messages_pt_keeps_signed_thinking_block(): assert thinking_blocks[0]["thinking"] == "genuine anthropic reasoning" +@pytest.mark.parametrize("llm_provider", ("anthropic", "vertex_ai")) +def test_anthropic_messages_pt_converts_pdf_image_url_data_uri_to_document(llm_provider: str) -> None: + pdf_data = "JVBERi0xLjQK" + result: Final = anthropic_messages_pt( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize"}, + {"type": "image_url", "image_url": {"url": f"data:application/pdf;base64,{pdf_data}"}}, + ], + } + ], + model="claude-opus-5-5", + llm_provider=llm_provider, + ) + + assert result[0]["content"][1] == { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": pdf_data}, + } + + +@pytest.mark.parametrize("llm_provider", ("anthropic", "vertex_ai")) +def test_anthropic_messages_pt_uses_pdf_format_for_generic_mime_data_uri(llm_provider: str) -> None: + pdf_data = "JVBERi0xLjQK" + result: Final = anthropic_messages_pt( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize"}, + { + "type": "image_url", + "image_url": { + "url": f"data:application/octet-stream;base64,{pdf_data}", + "format": "application/pdf", + }, + }, + ], + } + ], + model="claude-opus-5-5", + llm_provider=llm_provider, + ) + + assert result[0]["content"][1] == { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": pdf_data}, + } + + +def test_anthropic_messages_pt_keeps_png_image_url_as_image() -> None: + image_data = "iVBORw0KGgo=" + result: Final = anthropic_messages_pt( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": f"data:image/png;base64,{image_data}"}, + ], + } + ], + model="claude-opus-5-5", + llm_provider="anthropic", + ) + + assert result[0]["content"][1] == { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": image_data}, + } + + EMPTY_USER_MESSAGES: Final = pytest.mark.parametrize( "empty_user_message", [{"role": "user", "content": None}, {"role": "user", "content": []}, {"role": "user"}], diff --git a/tests/unit/litellm_core_utils/test_image_handling.py b/tests/unit/litellm_core_utils/test_image_handling.py index 9d9aef0c22a..d1bebde9917 100644 --- a/tests/unit/litellm_core_utils/test_image_handling.py +++ b/tests/unit/litellm_core_utils/test_image_handling.py @@ -288,6 +288,37 @@ def test_convert_url_to_base64_keeps_a_real_content_type(monkeypatch): assert result.startswith("data:image/jpeg;base64,") +@pytest.mark.parametrize( + ("url", "content_type", "expected"), + [ + ( + "https://docs.example/document.pdf", + "application/pdf; qs=0.001", + "data:application/pdf;base64,cGRm", + ), + ( + "https://img.example/image.png", + "image/png; charset=binary", + "data:image/png;base64,cGRm", + ), + ( + "https://docs.example/document.pdf", + "application/octet-stream; charset=binary", + "data:application/pdf;base64,cGRm", + ), + ], +) +def test_convert_url_to_base64_uses_bare_content_type( + monkeypatch: pytest.MonkeyPatch, + url: str, + content_type: str, + expected: str, +) -> None: + monkeypatch.setattr(litellm, "module_level_client", _stub_image_client(b"pdf", content_type)) + + assert convert_url_to_base64(url) == expected + + def test_convert_url_to_base64_raises_when_no_content_type_is_determinable(monkeypatch): monkeypatch.setattr( litellm, From d1758cda91a663b4e00b35af0921389bc6e2c95e Mon Sep 17 00:00:00 2001 From: simonp9 <349053457+simonp9@users.noreply.github.com> Date: Fri, 30 Oct 2026 02:40:01 -0700 Subject: [PATCH 16/64] fix(router)!: keep a deployment's own model_id over the request body's (#49045) * fix(router): keep a deployment's own model_id over the request body's A request body model_id overrode the deployment's pinned model_id, so a client could move SageMaker chat and embedding calls to another inference component. The router now drops the request's model_id when the deployment sets one. * fix(router): let deployments opt back in to a client model_id and stop client component headers on SageMaker chat Drop the request's model_id after the router defaults are applied, so a default_litellm_params model_id cannot win over the deployment either, and apply the same precedence on the sync text_completion path. A deployment that lists model_id in configurable_clientside_auth_params keeps the client's value. SageMaker chat drops any case variant of a client's X-Amzn-SageMaker-Inference-Component header when the deployment sets a component * fix(router): keep a SageMaker deployment's component when it pins it through extra_headers * fix(router): read the SageMaker component header as a pin only on SageMaker deployments * fix(sagemaker): send the inference component only as a header, never as a body model_id SageMaker picks the component from X-Amzn-SageMaker-Inference-Component, so the chat container body and legacy parameters no longer carry model_id. Nova endpoints accept a fixed body field list, so a pinned Nova deployment could reject the extra field. * fix(health): probe header-pinned SageMaker deployments at their component Health checks bypass the router, so a legacy sagemaker/ deployment pinned only through its inference component header sent no component on /health while routed requests now send it. Also update the two SageMaker integration tests that still asserted the old precedence and the dropped body field. * fix(health): apply the router's model_id pin rule to health probes Health probes now share the router's pin function, so a deployment that lets clients set model_id gets no header pin and its own model_id wins without reading the header. Header keys that are not strings are skipped when reading the inference component header instead of failing the call. * fix(sagemaker): keep header keys that aren't strings and report a string opt-in in health checks The chat transformation now keeps extra_headers keys that aren't strings, so such a deployment fails the way it did before instead of answering 500, and the model_id pin checks configurable_clientside_auth_params with isinstance, so a string there reaches the health probe's own error capture again. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/sagemaker/chat/transformation.py | 14 +- litellm/llms/sagemaker/common_utils.py | 33 ++- .../sagemaker/completion/transformation.py | 6 +- litellm/proxy/health_check.py | 8 +- litellm/router.py | 7 + litellm/router_utils/common_utils.py | 11 + .../providers/test_sagemaker_chat_wire.py | 2 +- .../test_sagemaker_embedding_wire.py | 4 +- .../test_sagemaker_chat_transformation.py | 20 ++ .../test_sagemaker_completion_handler.py | 21 ++ .../proxy/test_health_check_max_tokens.py | 82 +++++++ tests/unit/test_router/test_router.py | 219 ++++++++++++++++++ 12 files changed, 414 insertions(+), 13 deletions(-) diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py index 4a59d5f1802..76eafcd6bb6 100644 --- a/litellm/llms/sagemaker/chat/transformation.py +++ b/litellm/llms/sagemaker/chat/transformation.py @@ -28,7 +28,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import LlmProviders from ...openai.chat.gpt_transformation import OpenAIGPTConfig -from ..common_utils import AWSEventStreamDecoder, SagemakerError +from ..common_utils import INFERENCE_COMPONENT_HEADER, OBJECT_KEYED_MAPPING, AWSEventStreamDecoder, SagemakerError if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -59,7 +59,12 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): inference_component_name: Final = optional_params.get("model_id") if not isinstance(inference_component_name, str): return headers - return {**headers, "X-Amzn-SageMaker-Inference-Component": inference_component_name} + other_headers: Final = { + name: value + for name, value in OBJECT_KEYED_MAPPING.validate_python(headers).items() + if not isinstance(name, str) or name.lower() != INFERENCE_COMPONENT_HEADER.lower() + } + return {**other_headers, INFERENCE_COMPONENT_HEADER: inference_component_name} def transform_request( self, @@ -76,9 +81,8 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): litellm_params=litellm_params, headers=headers, ) - container_request: Final = { - key: value for key, value in request.items() if key not in self.aws_authentication_params - } + header_only_params: Final = frozenset((*self.aws_authentication_params, "model_id")) + container_request: Final = {key: value for key, value in request.items() if key not in header_only_params} served_model_name: Final = litellm_params.get("hf_model_name") if not isinstance(served_model_name, str): return container_request diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index eaa805e7769..bb7719a19a1 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -14,9 +14,10 @@ from typing_extensions import ReadOnly, Required, TypedDict from litellm import verbose_logger from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import GenericStreamingChunk as GChunk -from litellm.types.utils import StreamingChatCompletionChunk +from litellm.types.utils import LlmProviders, StreamingChatCompletionChunk STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +OBJECT_KEYED_MAPPING: Final = TypeAdapter(Mapping[object, object], config=ConfigDict(hide_input_in_errors=True)) @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) @@ -167,6 +168,36 @@ def get_sagemaker_response_stream_shape(): return _load_sagemaker_response_stream_shape() +INFERENCE_COMPONENT_HEADER: Final = "X-Amzn-SageMaker-Inference-Component" + + +_SAGEMAKER_PROVIDERS: Final = frozenset( + provider.value for provider in (LlmProviders.SAGEMAKER, LlmProviders.SAGEMAKER_CHAT, LlmProviders.SAGEMAKER_NOVA) +) + + +def _deployment_provider(litellm_params: Mapping[str, object]) -> str | None: + custom_llm_provider: Final = litellm_params.get("custom_llm_provider") + if isinstance(custom_llm_provider, str) and custom_llm_provider: + return custom_llm_provider + model: Final = litellm_params.get("model") + return model.split("/", 1)[0] if isinstance(model, str) else None + + +def deployment_inference_component(litellm_params: Mapping[str, object]) -> str | None: + headers: Final = litellm_params.get("extra_headers") + if _deployment_provider(litellm_params) not in _SAGEMAKER_PROVIDERS or not isinstance(headers, Mapping): + return None + return next( + ( + value + for name, value in OBJECT_KEYED_MAPPING.validate_python(headers).items() + if isinstance(name, str) and name.lower() == INFERENCE_COMPONENT_HEADER.lower() and isinstance(value, str) + ), + None, + ) + + def with_inference_component_hint(error_message: str) -> str: if "Inference Component Name header is required" in error_message: return f"{error_message}\n pass in via `litellm.embedding(..., model_id={{InferenceComponentName}})`" diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py index a69846c816f..5f6bef4246e 100644 --- a/litellm/llms/sagemaker/completion/transformation.py +++ b/litellm/llms/sagemaker/completion/transformation.py @@ -23,7 +23,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, Usage from litellm.utils import token_counter -from ..common_utils import SagemakerError +from ..common_utils import STR_KEYED_MAPPING, SagemakerError if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -153,7 +153,9 @@ class SagemakerConfig(BaseConfig): litellm_params: dict, headers: dict, ) -> dict: - inference_params: Final = optional_params.copy() + inference_params: Final = { + key: value for key, value in STR_KEYED_MAPPING.validate_python(optional_params).items() if key != "model_id" + } stream: Final = inference_params.pop("stream", False) data: Final[dict] = {"parameters": inference_params} if stream is True: diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 73c2fc35492..20d920f1104 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -32,7 +32,7 @@ from litellm.router_strategy.auto_router.auto_router_model_naming import ( classify_strategy_router_model, strategy_router_dependencies, ) -from litellm.router_utils.common_utils import is_model_capability_field +from litellm.router_utils.common_utils import deployment_model_id_pin, is_model_capability_field from litellm.types.utils import secret_bearing_wif_litellm_params, server_owned_wif_litellm_params # Provider routing and server-owned federation or OAuth token-exchange fields. Allowed for proxy @@ -855,6 +855,8 @@ def update_litellm_params_for_health_check( `deployment_model` (default: the incoming `litellm_params` model), so a `health_check_model` or a `model` in `health_check_params` is judged on its own. Keys already set through `health_check_params` are kept, and `modalities` set there without "audio" turns the audio fields off + - sends a SageMaker deployment's `X-Amzn-SageMaker-Inference-Component` extra header as `model_id` when the + deployment sets none, so legacy `sagemaker/` probes reach that component the way routed requests do """ incoming_model: Final = _PROBE_PARAMS_ADAPTER.validate_python(litellm_params).get("model") base_model_owner: Final = deployment_model or (incoming_model if isinstance(incoming_model, str) else None) @@ -935,9 +937,11 @@ def update_litellm_params_for_health_check( audio_output_params: Final = _health_check_audio_output_params( _MODEL_INFO_ADAPTER.validate_python(model_info), probe_params, mode, base_model_owner ) + model_id_pin: Final = deployment_model_id_pin(probe_params) + model_id_params: Final = {} if model_id_pin is None else {"model_id": model_id_pin} return { key: value - for key, value in {**audio_output_params, **probe_params}.items() + for key, value in {**audio_output_params, **model_id_params, **probe_params}.items() if not is_model_capability_field(key) } diff --git a/litellm/router.py b/litellm/router.py index 20af2e5649e..af2181cedb5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -255,6 +255,7 @@ from litellm.router_utils.clientside_credential_handler import ( from litellm.router_utils.common_utils import ( MID_STREAM_CONTINUATION_KWARG, as_retry_skipped_deployment_ids, + deployment_model_id_pin, deployments_the_caller_may_use, filter_team_based_models, filter_web_search_deployments, @@ -5316,6 +5317,9 @@ class Router(DeprecatedRouterMembers): ): kwargs[STREAM_USAGE_INJECTED_FIELD] = True self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name) + pinned_model_id: Final = deployment_model_id_pin(deployment_context) + if pinned_model_id is not None: + kwargs["model_id"] = pinned_model_id # rebind-ok: this method's contract is filling the caller's kwargs hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs) if _is_fallback_hop(hop_kwargs): EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( @@ -6242,6 +6246,9 @@ class Router(DeprecatedRouterMembers): kwargs[k] = v elif k == "metadata": kwargs[k].update(v) + pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data)) + if pinned_model_id is not None: + kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults # call via litellm.completion() return litellm.text_completion( diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 2451422b5e6..a06860e7786 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -16,6 +16,7 @@ from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure +from litellm.llms.sagemaker.common_utils import deployment_inference_component from litellm.types.router import ( CredentialLiteLLMParams, Deployment, @@ -522,6 +523,16 @@ PROVIDER_SPECIFIC_CAPABILITY_FIELDS: Final = frozenset(ProviderSpecificModelInfo _LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +def deployment_model_id_pin(deployment_litellm_params: Mapping[str, object]) -> object | None: + clientside_params: Final = deployment_litellm_params.get("configurable_clientside_auth_params") + if isinstance(clientside_params, list | tuple) and "model_id" in clientside_params: + return None + own_model_id: Final = deployment_litellm_params.get("model_id") + if own_model_id is not None: + return own_model_id + return deployment_inference_component(deployment_litellm_params) + + def is_model_capability_field(name: str) -> bool: return name in PROVIDER_SPECIFIC_CAPABILITY_FIELDS or ( name.startswith("supports_") and name in GetModelCostMap.loaded_model_cost_map_fields() diff --git a/tests/integration/providers/test_sagemaker_chat_wire.py b/tests/integration/providers/test_sagemaker_chat_wire.py index fdf0d94f7c7..784e43f85e2 100644 --- a/tests/integration/providers/test_sagemaker_chat_wire.py +++ b/tests/integration/providers/test_sagemaker_chat_wire.py @@ -243,7 +243,7 @@ def test_sagemaker_chat_signs_the_inference_component_header_and_sends_hf_model_ ) assert "x-amzn-sagemaker-inference-component" in signed_headers.removeprefix("SignedHeaders=").split(";") assert JSON_OBJECT.validate_json(request.body) == { - **container_body(wire.url, PROMPT, stream=False, model_id=_INFERENCE_COMPONENT), + **container_body(wire.url, PROMPT, stream=False), "model": SERVED_MODEL, }, request.body diff --git a/tests/integration/providers/test_sagemaker_embedding_wire.py b/tests/integration/providers/test_sagemaker_embedding_wire.py index 74f3c8c8a00..332f473aeaf 100644 --- a/tests/integration/providers/test_sagemaker_embedding_wire.py +++ b/tests/integration/providers/test_sagemaker_embedding_wire.py @@ -668,12 +668,12 @@ def test_an_extra_body_that_is_not_a_populated_object_adds_nothing_to_the_body( assert _body(request) == {"input": [text]}, request.body -def test_a_caller_sent_model_id_overrides_the_deployments_inference_component(proxy: Gateway, wire: Wire) -> None: +def test_a_caller_sent_model_id_cannot_override_the_deployments_inference_component(proxy: Gateway, wire: Wire) -> None: text: Final = _text() response: Final = _embed(proxy, "sm-openai", [text], model_id="caller-chosen-component") assert response.status_code == 200, response.text request: Final = _only_request(wire) - assert request.headers[COMPONENT_HEADER] == "caller-chosen-component", dict(request.headers) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body diff --git a/tests/unit/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/unit/llms/sagemaker/test_sagemaker_chat_transformation.py index f8caa8d5fed..356528d8909 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_chat_transformation.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_chat_transformation.py @@ -343,6 +343,26 @@ def test_body_carries_no_aws_credentials( assert client.request_headers["X-Amz-Security-Token"] == "test-session-token" +@pytest.mark.parametrize("model", ["sagemaker_chat/my-endpoint", "sagemaker_nova/my-endpoint"]) +def test_inference_component_travels_only_in_the_header(monkeypatch: pytest.MonkeyPatch, model: str): + client = _invoke_sagemaker_chat(monkeypatch, model=model, model_id="my-inference-component") + + assert client.request_headers[_INFERENCE_COMPONENT_HEADER] == "my-inference-component" + assert "model_id" not in client.request_body + + +def test_inference_component_header_keeps_header_keys_that_are_not_strings(): + headers = SagemakerChatConfig().validate_environment( + headers={7: "other", "x-amzn-sagemaker-inference-component": "client-component"}, + model="my-endpoint", + messages=[], + optional_params={"model_id": "my-inference-component"}, + litellm_params={}, + ) + + assert headers == {7: "other", _INFERENCE_COMPONENT_HEADER: "my-inference-component"} + + def test_body_model_stays_the_endpoint_name_when_hf_model_name_is_unset(monkeypatch): """Without `hf_model_name` the body must keep the model it has today.""" client = _invoke_sagemaker_chat(monkeypatch) diff --git a/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py index 0980c9f5c66..5e40c2ebb60 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py @@ -372,6 +372,27 @@ def test_completion_sagemaker_non_stream( assert response.choices[0].message.content == "SageMaker reply" +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.usefixtures("fake_provider_credentials") +def test_completion_sagemaker_sends_the_inference_component_only_as_a_header(respx_mock: respx.MockRouter) -> None: + url: Final = "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/my-endpoint/invocations" + route: Final = respx_mock.post(url).mock( + return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"}) + ) + litellm.completion( + model="sagemaker/my-endpoint", + messages=[{"role": "user", "content": "hi"}], + temperature=0.2, + model_id="my-inference-component", + aws_region_name="us-west-2", + num_retries=0, + ) + + request: Final = route.calls[0].request + assert request.headers["X-Amzn-SageMaker-Inference-Component"] == "my-inference-component" + assert json.loads(request.content) == {"inputs": "hi", "parameters": {"temperature": 0.2}} + + @pytest.mark.usefixtures("httpx_transport") @pytest.mark.usefixtures("fake_provider_credentials") def test_completion_sagemaker_non_stream_with_aws_params(respx_mock: respx.MockRouter) -> None: diff --git a/tests/unit/proxy/test_health_check_max_tokens.py b/tests/unit/proxy/test_health_check_max_tokens.py index df0429316ed..26d9fe2fb3d 100644 --- a/tests/unit/proxy/test_health_check_max_tokens.py +++ b/tests/unit/proxy/test_health_check_max_tokens.py @@ -1,5 +1,6 @@ import asyncio import copy +import io import json import logging from typing import Final @@ -1564,3 +1565,84 @@ async def test_a_failing_health_probe_leaves_the_deployment_serving_requests(mon assert "error" in probe assert reply.choices[0].message.content == "ok" + + +@pytest.mark.parametrize( + ("component_params", "probed_component"), + [ + ({"extra_headers": {"X-Amzn-SageMaker-Inference-Component": "header-component"}}, "header-component"), + ( + {"model_id": "own-component", "extra_headers": {"X-Amzn-SageMaker-Inference-Component": "header-component"}}, + "own-component", + ), + ( + {"model_id": "own-component", "extra_headers": {1: "other", "X-Amzn-SageMaker-Inference-Component": "x"}}, + "own-component", + ), + ( + {"extra_headers": {1: "other", "X-Amzn-SageMaker-Inference-Component": "header-component"}}, + "header-component", + ), + ( + { + "configurable_clientside_auth_params": ["model_id"], + "extra_headers": {"X-Amzn-SageMaker-Inference-Component": "header-component"}, + }, + None, + ), + ], + ids=[ + "header_pinned", + "own_model_id_beats_the_header", + "own_model_id_beside_a_non_string_header_key", + "header_pinned_beside_a_non_string_header_key", + "client_configurable_model_id_takes_no_header_pin", + ], +) +@pytest.mark.asyncio +async def test_run_model_health_check_probes_a_legacy_sagemaker_embedding_at_the_deployments_component( + component_params: dict[str, object], probed_component: str | None +) -> None: + embedding_reply: Final = {"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}]} + invoke_endpoint: Final = MagicMock(return_value={"Body": io.BytesIO(json.dumps(embedding_reply).encode())}) + client: Final = MagicMock(invoke_endpoint=invoke_endpoint) + deployment: Final = { + "litellm_params": { + "model": "sagemaker/openai/embed-endpoint", + "aws_region_name": "us-east-1", + "aws_access_key_id": "test-key", + "aws_secret_access_key": "test-secret", + **component_params, + }, + "model_info": {"mode": "embedding"}, + } + + with patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=client))): + result: Final = await hc_module._run_model_health_check(deployment) + + assert "error" not in result, result + assert invoke_endpoint.call_count == 1 + assert invoke_endpoint.call_args.kwargs.get("InferenceComponentName") == probed_component + + +@pytest.mark.asyncio +async def test_run_model_health_check_reports_why_a_string_configurable_clientside_auth_params_fails() -> None: + invoke_endpoint: Final = MagicMock() + deployment: Final = { + "litellm_params": { + "model": "sagemaker/openai/embed-endpoint", + "model_id": "own-component", + "aws_region_name": "us-east-1", + "aws_access_key_id": "test-key", + "aws_secret_access_key": "test-secret", + }, + "model_info": {"mode": "embedding", "health_check_params": {"configurable_clientside_auth_params": "model_id"}}, + } + + client: Final = MagicMock(invoke_endpoint=invoke_endpoint) + + with patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=client))): + result: Final = await hc_module._run_model_health_check(deployment) + + assert "configurable_clientside_auth_params" in result["error"] + invoke_endpoint.assert_not_called() diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index a0a7dd78193..74959155380 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -4,6 +4,7 @@ import copy import functools import gc import itertools +import io import json import linecache import logging @@ -31523,6 +31524,224 @@ def test_router_keeps_supported_fallback_shapes( assert router.fallbacks == expected_fallbacks +_SAGEMAKER_ROUTER_CREDENTIALS: Final = MappingProxyType( + {"aws_region_name": "us-east-1", "aws_access_key_id": "test-key", "aws_secret_access_key": "test-secret"} +) + +_PINNED_COMPONENT: Final = MappingProxyType({"model_id": "deployment-component"}) +_HEADER_PINNED_COMPONENT: Final = MappingProxyType( + {"extra_headers": {"x-amzn-sagemaker-inference-component": "deployment-component"}} +) + +_INFERENCE_COMPONENT_PRECEDENCE: Final = pytest.mark.parametrize( + ("deployment_params", "sent_component"), + [ + (_PINNED_COMPONENT, "deployment-component"), + (_HEADER_PINNED_COMPONENT, "deployment-component"), + (MappingProxyType({}), "client-component"), + ( + MappingProxyType({**_PINNED_COMPONENT, "configurable_clientside_auth_params": ["model_id"]}), + "client-component", + ), + ], + ids=[ + "deployment_pins_the_component", + "deployment_pins_the_component_with_a_header", + "deployment_leaves_the_component_to_the_client", + "deployment_lets_the_client_override_its_component", + ], +) + + +def _sagemaker_router( + model: str, + deployment_params: Mapping[str, object], + mode: str, + router_defaults: Mapping[str, str] = MappingProxyType({}), + model_name: str = "sagemaker-model", +) -> Router: + return Router( + model_list=[ + { + "model_name": model_name, + "litellm_params": {"model": model, **_SAGEMAKER_ROUTER_CREDENTIALS, **deployment_params}, + "model_info": {"mode": mode}, + } + ], + default_litellm_params=dict(router_defaults), + ) + + +_CHAT_COMPLETION_REPLY: Final = MappingProxyType( + { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "served-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } +) + + +def _sagemaker_chat_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post( + "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/chat-endpoint/invocations" + ).respond(json=dict(_CHAT_COMPLETION_REPLY)) + + +@_INFERENCE_COMPONENT_PRECEDENCE +@pytest.mark.asyncio +async def test_a_client_model_id_cannot_move_a_sagemaker_chat_call_off_the_deployments_component( + respx_mock: respx.MockRouter, + _router_httpx_transport: None, + deployment_params: Mapping[str, object], + sent_component: str, +) -> None: + route: Final = _sagemaker_chat_route(respx_mock) + router: Final = _sagemaker_router("sagemaker_chat/chat-endpoint", deployment_params, "chat") + + await router.acompletion( + model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], model_id="client-component" + ) + + assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == sent_component + + +@_INFERENCE_COMPONENT_PRECEDENCE +def test_a_client_model_id_cannot_move_a_sync_sagemaker_text_completion_off_the_deployments_component( + respx_mock: respx.MockRouter, deployment_params: Mapping[str, object], sent_component: str +) -> None: + route: Final = _sagemaker_chat_route(respx_mock) + router: Final = _sagemaker_router( + "sagemaker_chat/chat-endpoint", deployment_params, "completion", model_name="sagemaker_chat/chat-endpoint" + ) + + router.text_completion(model="sagemaker_chat/chat-endpoint", prompt="hi", model_id="client-component") + + assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == sent_component + + +@pytest.mark.parametrize( + "client_header", + [ + MappingProxyType({"extra_headers": {"x-amzn-sagemaker-inference-component": "client-component"}}), + MappingProxyType({"headers": {"x-amzn-sagemaker-inference-component": "client-component"}}), + MappingProxyType( + { + "provider_specific_header": { + "custom_llm_provider": "sagemaker_chat", + "extra_headers": {"x-amzn-sagemaker-inference-component": "client-component"}, + } + } + ), + ], + ids=["extra_headers", "headers", "provider_specific_header"], +) +@pytest.mark.parametrize( + "deployment_pin", [_PINNED_COMPONENT, _HEADER_PINNED_COMPONENT], ids=["model_id_pin", "header_pin"] +) +@pytest.mark.asyncio +async def test_a_client_inference_component_header_cannot_move_a_pinned_sagemaker_chat_call( + respx_mock: respx.MockRouter, + _router_httpx_transport: None, + client_header: Mapping[str, object], + deployment_pin: Mapping[str, object], +) -> None: + route: Final = _sagemaker_chat_route(respx_mock) + router: Final = _sagemaker_router("sagemaker_chat/chat-endpoint", deployment_pin, "chat") + + await router.acompletion(model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], **client_header) + + assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == "deployment-component" + + +@pytest.mark.parametrize( + "request_component", + [MappingProxyType({}), MappingProxyType({"model_id": "deployment-component"})], + ids=["request_without_model_id", "request_with_the_pinned_model_id"], +) +@pytest.mark.asyncio +async def test_a_router_default_model_id_cannot_move_a_sagemaker_call_off_the_deployments_component( + respx_mock: respx.MockRouter, _router_httpx_transport: None, request_component: Mapping[str, str] +) -> None: + route: Final = _sagemaker_chat_route(respx_mock) + router: Final = _sagemaker_router( + "sagemaker_chat/chat-endpoint", + _PINNED_COMPONENT, + "chat", + router_defaults=MappingProxyType({"model_id": "default-component"}), + ) + + await router.acompletion(model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], **request_component) + + assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == "deployment-component" + + +@_INFERENCE_COMPONENT_PRECEDENCE +@pytest.mark.asyncio +async def test_a_client_model_id_cannot_move_a_sagemaker_embedding_call_off_the_deployments_component( + deployment_params: Mapping[str, object], sent_component: str +) -> None: + component_embeddings: Final = MappingProxyType({"deployment-component": [0.1], "client-component": [0.2]}) + + def invoke_endpoint(**request: str | bytes) -> Mapping[str, io.BytesIO]: + served: Final = component_embeddings[str(request["InferenceComponentName"])] + return { + "Body": io.BytesIO( + json.dumps({"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]}).encode() + ) + } + + client: Final = MagicMock(invoke_endpoint=invoke_endpoint) + router: Final = _sagemaker_router("sagemaker/openai/embed-endpoint", deployment_params, "embedding") + + with patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=client))): + response: Final = await router.aembedding(model="sagemaker-model", input=["hello"], model_id="client-component") + + assert response.data[0]["embedding"] == component_embeddings[sent_component] + + +@pytest.mark.asyncio +async def test_a_header_pin_holds_when_the_deployment_names_sagemaker_through_custom_llm_provider( + respx_mock: respx.MockRouter, _router_httpx_transport: None +) -> None: + route: Final = _sagemaker_chat_route(respx_mock) + router: Final = _sagemaker_router( + "chat-endpoint", MappingProxyType({**_HEADER_PINNED_COMPONENT, "custom_llm_provider": "sagemaker_chat"}), "chat" + ) + + await router.acompletion( + model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], model_id="client-component" + ) + + assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == "deployment-component" + + +@pytest.mark.asyncio +async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clients_model_id_alone( + respx_mock: respx.MockRouter, _router_httpx_transport: None +) -> None: + route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond(json=dict(_CHAT_COMPLETION_REPLY)) + router: Final = Router( + model_list=[ + { + "model_name": "gateway-model", + "litellm_params": { + "model": "openai/served-model", + "api_base": "https://gateway.local/v1", + "api_key": "test-key", + "extra_headers": {"X-Amzn-SageMaker-Inference-Component": "gateway-component"}, + }, + } + ] + ) + + await router.acompletion(model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value") + + assert json.loads(route.calls.last.request.content)["model_id"] == "client-value" + + def test_router_warns_when_a_deployment_enforces_endpoints_without_a_list(caplog: pytest.LogCaptureFixture) -> None: with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): router: Final = Router( From 080e643edc2db94cf34824fe851721c320099bd7 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:19:25 -0700 Subject: [PATCH 17/64] fix(proxy): apply new model rows on every worker as soon as they are written (#49034) * fix(proxy): apply published model rows on every worker before the debounced resync Config sync messages now carry the ids of model rows a create, update, or upsert wrote. Every worker that receives one loads those rows by id right away, including while its debounced full resync is still waiting, so a group added one deployment at a time is complete on every worker as soon as the write returns. Deletes, bulk writes, and messages from older pods carry no ids and keep the full resync path * fix(proxy): load published model rows from the primary under the reconcile lock Read the by-id rows through the writer so a lagging read replica cannot put an old row back, and read them inside MODEL_RECONCILE_LOCK so a concurrent delete or reload stays the last word. Look the ids up in bounded chunks, reload DB credentials only when a row names one this worker lacks, and leave auto-router rows to the debounced full resync. Add an integration test that boots two proxies and needs the peer to fall back to a just-added backup within 4 s. * test(proxy): cover a published model row that names no credential * fix(proxy): keep the unknown model name read-through on the replica and outside the lock Only the by-id load that follows a published write reads the primary inside the reconcile lock. A request for a model name a worker does not know reads through the replica routing before taking the lock again, as on main, so unknown names put no load on the primary and do not wait behind a reload. * fix(proxy): keep access group edits from dropping staying deployments and batch published config ids * fix(proxy): read access group rows from the primary so quick edits are not lost to a lagging replica * fix(proxy): install ordinary rows queued with an auto router and check access group edits before writing A batch of published model ids that held one auto router skipped every row in it, so ordinary rows written at the same moment waited for the full reload. The by-id load now leaves only the auto router rows to that reload. An access group update now reads every deployment from the primary, answers 404 for an unknown group and 400 for an unknown model id or name before it writes anything, then removes the group from the deployments that leave before it adds it to the ones that join, so a failed write leaves the group narrower than asked, never wider. * fix(proxy): write an access group update in one transaction after a fresh read The update now checks model names first, reads every deployment from the primary, and commits every join and leave in one batch right after that read, so a failed write changes nothing and a field another admin set during validation is kept. Other workers hear about the rows only after the commit. * fix(proxy): lock the model rows while an access group update reads and rewrites them * test(proxy): run the model resync tests on the real installer and a real router The by-id resync tests now install rows through ProxyConfig's real installer and credential loader on a real Router, faking only the database, and the config sync test row model is frozen. * fix(proxy): lock only the rows an access group update touches and answer 404 before 400 An access group update locked every model row, so a lock on any unrelated row stalled it for 30 s and then failed it with a 500, and unrelated model edits on other proxies waited behind it. It now takes a per-group advisory lock, locks only the rows that carry the group or that it selects, and reads them again once locked. The group existence check runs before the model name check again, so a missing group answers 404 as on main. * fix(proxy): never wait on a locked row in an access group update The update now takes its row locks with NOWAIT. When another edit holds one of the group's rows, it rolls back, holding nothing, and tries again for about 10 s before answering 409. Holding rows while waiting deadlocked against team deletes and let a concurrent model edit undo the update's change. The per-group lock also moves to the two-key advisory space, apart from team ids. * fix(proxy): retry access group row locks for 30 s with SKIP LOCKED and lock only the rows model_ids select * fix(proxy): take the access group lock with pg_try_advisory_xact_lock and retry while another edit holds it --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../proxy/common_utils/config_sync_pubsub.py | 136 +++- .../common_utils/registry_read_through.py | 126 ++- ...model_access_group_management_endpoints.py | 315 ++++++-- .../model_management_endpoints.py | 6 +- litellm/proxy/proxy_server.py | 2 + ...backup_deployment_serves_peer_fallbacks.py | 109 +++ .../common_utils/test_config_sync_pubsub.py | 482 +++++++++--- .../test_registry_read_through.py | 347 +++++++++ .../test_access_group_management.py | 735 +++++++++++++++++- .../test_model_management_endpoints.py | 6 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 12 files changed, 2049 insertions(+), 222 deletions(-) create mode 100644 tests/integration/management/test_added_backup_deployment_serves_peer_fallbacks.py diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index bc64a4c7ef0..c76b9be7c33 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -1648,7 +1648,7 @@ }, "/access_group/{access_group}/update": { "put": { - "description": "Update an access group's model names.\n\nThis will:\n1. Remove the access group from all current deployments\n2. Add the access group to all deployments for the new model_names list\n\nExample:\n```bash\ncurl -X PUT 'http://localhost:4000/access_group/production-models/update' \\\n -H \"Authorization: Bearer $LITELLM_MASTER_KEY\" \\\n -H 'Content-Type: application/json' \\\n -d '{\n \"model_names\": [\"gpt-4\", \"claude-3-sonnet\"]\n }'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n- model_names: List[str] - New list of model groups to include\n\nReturns:\n- NewModelGroupResponse with the updated access group details\n\nRaises:\n- HTTPException 400: If any model names don't exist\n- HTTPException 404: If access group not found", + "description": "Update an access group's model names.\n\nThis will:\n1. Add the access group to all deployments for the new model_names list\n2. Remove the access group from every other deployment\n\nExample:\n```bash\ncurl -X PUT 'http://localhost:4000/access_group/production-models/update' \\\n -H \"Authorization: Bearer $LITELLM_MASTER_KEY\" \\\n -H 'Content-Type: application/json' \\\n -d '{\n \"model_names\": [\"gpt-4\", \"claude-3-sonnet\"]\n }'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n- model_names: List[str] - New list of model groups to include\n\nReturns:\n- NewModelGroupResponse with the updated access group details\n\nRaises:\n- HTTPException 400: If any model names don't exist\n- HTTPException 404: If access group not found\n- HTTPException 409: If another edit holds the group or one of its models for 30 s", "operationId": "update_access_group_access_group__access_group__update_put", "parameters": [ { diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py index da74851b2d8..e9aaa0e81e4 100644 --- a/litellm/proxy/common_utils/config_sync_pubsub.py +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -2,9 +2,12 @@ import asyncio import json import random import time -from collections.abc import Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # untyped prisma/redis boundary needs cast +from itertools import chain +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, cast # noqa: TID251 # untyped prisma/redis boundary + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.repositories.prisma_protocols import RowT_co, TableActions @@ -106,24 +109,69 @@ _pubsub_capable_client: Final = pubsub_capable_client @dataclass(frozen=True, slots=True) class _ConfigChangeMessage: object_type: str + model_ids: tuple[str, ...] = () -def _config_change_message_json(object_type: str) -> str: - return json.dumps(asdict(_ConfigChangeMessage(object_type=object_type))) +@dataclass(frozen=True, slots=True) +class _PubSubEnvelope: + data: str | bytes -async def publish_config_change(redis_cache: "RedisCache | None", object_type: str) -> None: +class _WrittenModelRow(BaseModel): + model_config = ConfigDict(from_attributes=True, frozen=True) + + model_id: str + + +_CONFIG_CHANGE_MESSAGE: Final = TypeAdapter(_ConfigChangeMessage) +_PUBSUB_ENVELOPE: Final = TypeAdapter(_PubSubEnvelope) +_MODEL_TABLE_NAME: Final = "litellm_proxymodeltable" +_TARGETED_WRITE_ACTION_NAMES: Final[frozenset[str]] = frozenset({"create", "update", "upsert"}) + +PublishConfigChange: TypeAlias = Callable[[str, tuple[str, ...]], Awaitable[None]] + + +def _config_change_message_json(object_type: str, model_ids: tuple[str, ...]) -> str: + return json.dumps(asdict(_ConfigChangeMessage(object_type=object_type, model_ids=model_ids))) + + +def parse_config_change(message: object) -> _ConfigChangeMessage | None: + try: + payload: Final = ( + message if isinstance(message, (str, bytes)) else _PUBSUB_ENVELOPE.validate_python(message).data + ) + return _CONFIG_CHANGE_MESSAGE.validate_json(payload) + except ValidationError: + return None + + +def written_model_ids(object_type: str, action_name: str, result: object) -> tuple[str, ...]: + if object_type != _MODEL_TABLE_NAME or action_name not in _TARGETED_WRITE_ACTION_NAMES: + return () + try: + return (_WrittenModelRow.model_validate(result).model_id,) + except ValidationError: + return () + + +async def publish_config_change( + redis_cache: "RedisCache | None", object_type: str, model_ids: tuple[str, ...] = () +) -> None: if redis_cache is None: return try: client: Final = pubsub_capable_client(redis_cache) - await client.publish(config_sync_channel(redis_cache), _config_change_message_json(object_type)) + await client.publish(config_sync_channel(redis_cache), _config_change_message_json(object_type, model_ids)) except Exception as e: # noqa: BLE001 # best-effort publish; writes must never fail on redis errors verbose_proxy_logger.warning("config sync publish for %s failed: %s", object_type, e) -async def publish_config_change_for_object_type(object_type: str) -> None: - await publish_config_change(redis_cache=coordination_redis_cache(), object_type=object_type) +async def publish_config_change_for_object_type(object_type: str, model_ids: tuple[str, ...] = ()) -> None: + await publish_config_change(redis_cache=coordination_redis_cache(), object_type=object_type, model_ids=model_ids) + + +async def publish_model_rows_written(model_ids: tuple[str, ...]) -> None: + await publish_config_change_for_object_type(_MODEL_TABLE_NAME, model_ids) async def publish_config_param_change(param_name: str) -> None: @@ -139,7 +187,7 @@ async def publish_config_param_change(param_name: str) -> None: class _PublishOnWriteActions: __slots__ = ("_actions", "_object_type", "_publish") - def __init__(self, actions: object, object_type: str, publish: Callable[[str], Awaitable[None]]) -> None: + def __init__(self, actions: object, object_type: str, publish: PublishConfigChange) -> None: self._actions = actions self._object_type = object_type self._publish = publish @@ -157,7 +205,7 @@ class _PublishOnWriteActions: **kwargs: object, # kwargs-ok: transparent passthrough to untyped prisma action ) -> object: result: Final = await write_action(*args, **kwargs) - await publish(object_type) + await publish(object_type, written_model_ids(object_type, name, result)) return result return _write_then_publish @@ -166,7 +214,7 @@ class _PublishOnWriteActions: def wrap_table_actions_for_config_sync( actions: "TableActions[RowT_co]", table_name: str, - publish: Callable[[str], Awaitable[None]] = publish_config_change_for_object_type, + publish: PublishConfigChange = publish_config_change_for_object_type, ) -> "TableActions[RowT_co]": if table_name not in _CONFIG_SYNCED_TABLE_NAMES: return actions @@ -174,8 +222,14 @@ def wrap_table_actions_for_config_sync( return cast("TableActions[RowT_co]", wrapped) # cast-ok: dynamic write-through proxy keeps the wrapped row type +def _published_model_ids(messages: tuple[object, ...]) -> tuple[str, ...]: + changes: Final = (parse_config_change(message) for message in messages) + return tuple(dict.fromkeys(chain.from_iterable(change.model_ids for change in changes if change is not None))) + + class ConfigSyncSubscriber: __slots__ = ( + "_apply_model_ids", "_backoff_initial_seconds", "_backoff_max_seconds", "_debounce_seconds", @@ -195,6 +249,7 @@ class ConfigSyncSubscriber: self, redis_cache: "RedisCache", resync_callbacks: tuple[Callable[[], Awaitable[None]], ...], + apply_model_ids: Callable[[tuple[str, ...]], Awaitable[None]] | None = None, debounce_seconds: float = CONFIG_SYNC_DEBOUNCE_SECONDS, jitter_max_seconds: float = CONFIG_SYNC_JITTER_MAX_SECONDS, min_resync_interval_seconds: float = CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS, @@ -206,6 +261,7 @@ class ConfigSyncSubscriber: ) -> None: self._redis_cache = redis_cache self._resync_callbacks = resync_callbacks + self._apply_model_ids = apply_model_ids self._debounce_seconds = debounce_seconds self._jitter_max_seconds = jitter_max_seconds self._min_resync_interval_seconds = min_resync_interval_seconds @@ -266,32 +322,64 @@ class ConfigSyncSubscriber: message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS) if message is None: continue + await self._apply(pubsub, message) await self._debounce_and_resync(pubsub) - async def _debounce_and_resync(self, pubsub: _ConfigSyncPubSub) -> None: - await self._sleep(self._debounce_seconds + self._rng.uniform(0.0, self._jitter_max_seconds)) - await self._wait_for_min_resync_interval() + async def _debounce_and_resync(self, pubsub: ConfigSyncPubSub) -> None: + await self._apply_until( + pubsub, self._monotonic() + self._debounce_seconds + self._rng.uniform(0.0, self._jitter_max_seconds) + ) + await self._apply_until(pubsub, self._next_resync_at()) await self._drain_pending(pubsub) await self._run_resync_callbacks() self._last_resync_at = self._monotonic() - async def _wait_for_min_resync_interval(self) -> None: + def _next_resync_at(self) -> float: if self._last_resync_at is None: - return - seconds_until_next_resync = self._min_resync_interval_seconds - (self._monotonic() - self._last_resync_at) - if seconds_until_next_resync <= 0: - return - verbose_proxy_logger.debug( - "config sync resync throttled for %.1fs to cap fleet-wide reload rate", - seconds_until_next_resync, - ) - await self._sleep(seconds_until_next_resync) + return self._monotonic() + next_resync_at: Final = self._last_resync_at + self._min_resync_interval_seconds + throttled_for: Final = next_resync_at - self._monotonic() + if throttled_for > 0: + verbose_proxy_logger.debug( + "config sync resync throttled for %.1fs to cap fleet-wide reload rate", + throttled_for, + ) + return next_resync_at + + async def _apply_until(self, pubsub: ConfigSyncPubSub, deadline: float) -> None: + while True: + remaining = deadline - self._monotonic() + if remaining <= 0: + return + message = await pubsub.get_message( + ignore_subscribe_messages=True, timeout=min(remaining, _POLL_TIMEOUT_SECONDS) + ) + if message is not None: + await self._apply(pubsub, message) @staticmethod async def _drain_pending(pubsub: ConfigSyncPubSub) -> None: while await pubsub.get_message(ignore_subscribe_messages=True, timeout=0) is not None: pass + @staticmethod + async def _ready_messages(pubsub: ConfigSyncPubSub) -> AsyncIterator[object]: + while (message := await pubsub.get_message(ignore_subscribe_messages=True, timeout=0)) is not None: + yield message + + async def _apply(self, pubsub: ConfigSyncPubSub, message: object) -> None: + apply: Final = self._apply_model_ids + if apply is None: + return + ready: Final = (message, *[ready_message async for ready_message in self._ready_messages(pubsub)]) + model_ids: Final = _published_model_ids(ready) + if not model_ids: + return + try: + await apply(model_ids) + except Exception as e: # noqa: BLE001 # the debounced full resync retries what a targeted apply could not + verbose_proxy_logger.warning("config sync targeted apply of %s failed: %s", model_ids, e) + async def _run_resync_callbacks(self) -> None: for callback in self._resync_callbacks: try: diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 22df11429c0..487c6feb68d 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -15,12 +15,13 @@ or the first routed request on this replica would pick a tier it has not loaded. import asyncio import time from collections.abc import Awaitable, Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper if TYPE_CHECKING: @@ -33,8 +34,14 @@ if TYPE_CHECKING: ) from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.repositories.prisma_protocols import TableActions from litellm.types.agents import AgentResponse + _ModelTable: TypeAlias = TableActions[prisma_models.LiteLLM_ProxyModelTable] + _ModelRows: TypeAlias = Sequence[prisma_models.LiteLLM_ProxyModelTable] + +ModelLoadOutcome: TypeAlias = Literal["loaded", "missing", "no_router", "auto_router"] + READ_THROUGH_MISS_TTL_SECONDS: Final = 2.0 READ_THROUGH_RESYNC_WINDOW_SECONDS: Final = 5.0 READ_THROUGH_MAX_RESYNCS_PER_WINDOW: Final = 20 @@ -123,47 +130,114 @@ async def _resync_model_deployments(model_name: str) -> bool: return False prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None - table: Final = ModelRepository(prisma_client).table - name_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_name": model_name} - id_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_id": model_name} - rows: Final = await table.find_many(where=name_filter) or await table.find_many(where=id_filter) + rows: Final = await _find_rows_by_name_or_id(model_name, ModelRepository(prisma_client).table) if not rows: return False - router: Final = proxy_server.llm_router - if router is None: - await proxy_server.proxy_config.add_deployment( - prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj - ) - return proxy_server.llm_router is not None - if _holds_an_auto_router(rows): - await proxy_server.proxy_config.add_deployment( - prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj - ) - return _model_is_loaded(model_name) async with proxy_server.MODEL_RECONCILE_LOCK: + outcome: Final = await _install_model_rows(rows, _always) + if outcome in ("no_router", "auto_router"): + await _reload_every_model() + return outcome == "loaded" or _model_is_loaded(model_name) + + +async def resync_model_deployments_by_ids(model_ids: tuple[str, ...]) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient + from litellm.repositories.model_repository import ModelRepository + + if not _db_backed_registries_enabled("models"): + return + prisma_client: Final = proxy_server.prisma_client + assert prisma_client is not None + writer_table: Final = ModelRepository(WriterPinnedClient(prisma_client.db)).table + async with proxy_server.MODEL_RECONCILE_LOCK: + rows: Final = await _find_rows_by_ids(model_ids, writer_table) + _ = await _install_model_rows(_without_auto_routers(rows), _names_an_unloaded_credential) + + +async def _find_rows_by_name_or_id(model_name: str, table: "_ModelTable") -> "_ModelRows": + name_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_name": model_name} + id_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_id": model_name} + return await table.find_many(where=name_filter) or await table.find_many(where=id_filter) + + +async def _find_rows_by_ids(model_ids: tuple[str, ...], table: "_ModelTable") -> "_ModelRows": + from litellm.repositories.chunked_in import find_many_in + + return list(await find_many_in(table, "model_id", model_ids)) + + +async def _install_model_rows( + rows: "_ModelRows", + needs_credentials: Callable[["_ModelRows"], bool], +) -> ModelLoadOutcome: + from litellm.proxy import proxy_server + + router: Final = proxy_server.llm_router + if not rows: + return "missing" + if router is None: + return "no_router" + if _holds_an_auto_router(rows): + return "auto_router" + if needs_credentials(rows): + prisma_client: Final = proxy_server.prisma_client + assert prisma_client is not None await proxy_server.proxy_config.get_credentials(prisma_client=prisma_client) - proxy_server.proxy_config._add_deployment(db_models=rows) - proxy_server.llm_model_list = router.registry().deployments() + proxy_server.proxy_config._add_deployment(db_models=rows) + proxy_server.llm_model_list = router.registry().deployments() + return "loaded" + + +async def _reload_every_model() -> None: + from litellm.proxy import proxy_server + + prisma_client: Final = proxy_server.prisma_client + assert prisma_client is not None + await proxy_server.proxy_config.add_deployment( + prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj + ) + + +def _always(rows: "_ModelRows") -> bool: return True -def _holds_an_auto_router(rows: Sequence["prisma_models.LiteLLM_ProxyModelTable"]) -> bool: +def _holds_an_auto_router(rows: "_ModelRows") -> bool: return any(_names_an_auto_router(row.litellm_params) for row in rows) +def _without_auto_routers(rows: "_ModelRows") -> "_ModelRows": + return list(row for row in rows if not _names_an_auto_router(row.litellm_params)) + + +def _names_an_unloaded_credential(rows: "_ModelRows") -> bool: + return any(_names_an_unloaded_credential_param(row.litellm_params) for row in rows) + + _STORED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) -def _names_an_auto_router(stored_litellm_params: object) -> bool: +def _decrypted_param(stored_litellm_params: object, key: str) -> str | None: try: params: Final = _STORED_LITELLM_PARAMS.validate_python(stored_litellm_params) except ValidationError: - return False - stored_model: Final = params.get("model") - if not isinstance(stored_model, str): - return False - model: Final = decrypt_value_helper(value=stored_model, key="model", return_original_value=True) - return isinstance(model, str) and model.startswith("auto_router/") + return None + stored_value: Final = params.get(key) + if not isinstance(stored_value, str): + return None + value: Final = decrypt_value_helper(value=stored_value, key=key, return_original_value=True) + return value if isinstance(value, str) else None + + +def _names_an_auto_router(stored_litellm_params: object) -> bool: + model: Final = _decrypted_param(stored_litellm_params, "model") + return model is not None and model.startswith("auto_router/") + + +def _names_an_unloaded_credential_param(stored_litellm_params: object) -> bool: + credential_name: Final = _decrypted_param(stored_litellm_params, "litellm_credential_name") + return credential_name is not None and CredentialAccessor.find_credential(credential_name) is None async def _resync_guardrails(guardrail_name: str) -> bool: diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index 2d88882a2ac..3a2de12d108 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -12,23 +12,30 @@ Endpoints here: - DELETE /access_group/{access_group}/budget - Clear an access group's shared budget """ +import asyncio +import contextlib +import itertools import json -from collections.abc import Mapping, Sequence -from datetime import datetime +import time +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from datetime import datetime, timedelta from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel, ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.config_sync_pubsub import publish_model_rows_written from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, model_access_group_cache_key, model_access_group_registry_cache_key, ) +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient, writer_wrapper from litellm.proxy.management_endpoints.common_utils import validate_budget_duration # Clear cache and reload models to pick up the access group changes @@ -116,6 +123,28 @@ class _ModelTableClient(Protocol): async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... +class _ModelTableBatchActions(Protocol): + def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... + + +class _ModelTableBatch(Protocol): + @property + def litellm_proxymodeltable(self) -> _ModelTableBatchActions: ... + + +class _ModelTableTxActions(Protocol): + async def find_many(self) -> Sequence[_DeploymentRow]: ... + + +class _ModelTableTx(Protocol): + @property + def litellm_proxymodeltable(self) -> _ModelTableTxActions: ... + + async def query_raw(self, query: str, *args: object) -> object: ... + + def batch_(self) -> contextlib.AbstractAsyncContextManager[_ModelTableBatch]: ... + + class _BudgetRow(Protocol): @property def budget_id(self) -> str: ... @@ -168,7 +197,42 @@ class _ModelAccessGroupBudgetTableClient(Protocol): def _model_table(prisma_client: PrismaClient) -> _ModelTableClient: - return ModelRepository(prisma_client).table + return ModelRepository(WriterPinnedClient(prisma_client.db)).table + + +_TRY_LOCK_ACCESS_GROUP: Final = ( + "SELECT pg_try_advisory_xact_lock(hashtext('model_access_group'), hashtext($1)) AS locked" +) +_LOCK_MODEL_ROWS: Final = ( + 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = ANY($1::text[]) ' + "ORDER BY model_id FOR UPDATE SKIP LOCKED" +) +_MEMBERSHIP_TX_TIMEOUT: Final = timedelta(seconds=30) +_LOCK_RETRY_BUDGET_SECONDS: Final = _MEMBERSHIP_TX_TIMEOUT.total_seconds() +_LOCK_RETRY_BACKOFF_SECONDS: Final = (0.05, 0.1, 0.2, 0.4, 0.8) +_LOCK_RETRY_LONGEST_DELAY_SECONDS: Final = 1.0 + + +class _LockedModelRow(BaseModel): + model_config = ConfigDict(frozen=True) + + model_id: str + + +_LOCKED_MODEL_ROWS: Final = TypeAdapter(tuple[_LockedModelRow, ...]) + + +class _AccessGroupLock(BaseModel): + model_config = ConfigDict(frozen=True) + + locked: bool + + +_ACCESS_GROUP_LOCK: Final = TypeAdapter(tuple[_AccessGroupLock]) + + +def _model_table_tx(prisma_client: PrismaClient) -> contextlib.AbstractAsyncContextManager[_ModelTableTx]: + return writer_wrapper(prisma_client.db).tx(timeout=_MEMBERSHIP_TX_TIMEOUT) def _model_access_group_budget_table(prisma_client: PrismaClient) -> _ModelAccessGroupBudgetTableClient: @@ -379,11 +443,10 @@ async def _tag_deployment_with_access_group( """Write `access_group` into one deployment's model_info; returns the (model_id, updated model_info) pair when a write happened, None when the deployment already carried the group.""" - updated_model_info, was_modified = add_access_group_to_deployment( - model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)), - access_group=access_group, + updated_model_info: Final = _model_info_with_access_group( + model_id=model_id, model_info=model_info, access_group=access_group ) - if not was_modified: + if updated_model_info is None: return None await _model_table(prisma_client).update( where={"model_id": model_id}, @@ -393,6 +456,141 @@ async def _tag_deployment_with_access_group( return (model_id, updated_model_info) +def _model_info_with_access_group(model_id: str, model_info: object, access_group: str) -> Mapping[str, object] | None: + updated_model_info, was_modified = add_access_group_to_deployment( + model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)), + access_group=access_group, + ) + return updated_model_info if was_modified else None + + +def _model_info_without_access_group( + model_id: str, model_info: object, access_group: str +) -> Mapping[str, object] | None: + updated_model_info, was_modified = remove_access_group_from_deployment( + model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)), + access_group=access_group, + ) + return updated_model_info if was_modified else None + + +def _membership_rewrite(deployment: _DeploymentRow, selected: bool, access_group: str) -> Mapping[str, object] | None: + if selected: + return _model_info_with_access_group( + model_id=deployment.model_id, model_info=deployment.model_info, access_group=access_group + ) + return _model_info_without_access_group( + model_id=deployment.model_id, model_info=deployment.model_info, access_group=access_group + ) + + +def _membership_rewrites( + deployments: Sequence[_DeploymentRow], selected_model_ids: frozenset[str], access_group: str +) -> tuple[tuple[str, Mapping[str, object]], ...]: + rewrites: Final = ( + ( + deployment.model_id, + _membership_rewrite( + deployment=deployment, + selected=deployment.model_id in selected_model_ids, + access_group=access_group, + ), + ) + for deployment in deployments + ) + return tuple((model_id, model_info) for model_id, model_info in rewrites if model_info is not None) + + +def _carries_access_group(model_info: object, access_group: str) -> bool: + access_groups: Final = (model_info_as_mapping(model_info) or {}).get("access_groups") + return isinstance(access_groups, list) and access_group in access_groups + + +def _model_ids_to_lock( + deployments: Sequence[_DeploymentRow], access_group: str, data: UpdateModelGroupRequest +) -> frozenset[str]: + selected_model_ids: Final = frozenset(data.model_ids or ()) + selected_model_names: Final = frozenset(() if data.model_ids else data.model_names or ()) + return frozenset( + deployment.model_id + for deployment in deployments + if deployment.model_id in selected_model_ids + or deployment.model_name in selected_model_names + or _carries_access_group(deployment.model_info, access_group) + ) + + +async def _raise_404_unless_the_primary_carries(access_group: str, prisma_client: PrismaClient) -> None: + if access_group not in _access_groups_of_deployments(await _model_table(prisma_client).find_many()): + raise _http_error(404, f"Access group '{access_group}' not found") + + +async def _rewrite_membership_unless_locked( + prisma_client: PrismaClient, access_group: str, data: UpdateModelGroupRequest +) -> tuple[tuple[_DeploymentRow, ...], tuple[tuple[str, Mapping[str, object]], ...]] | None: + async with _model_table_tx(prisma_client) as tx: + (group_lock,) = _ACCESS_GROUP_LOCK.validate_python(await tx.query_raw(_TRY_LOCK_ACCESS_GROUP, access_group)) + if not group_lock.locked: + return None + model_ids_to_lock: Final = _model_ids_to_lock( + deployments=await tx.litellm_proxymodeltable.find_many(), access_group=access_group, data=data + ) + locked_rows: Final = _LOCKED_MODEL_ROWS.validate_python( + await tx.query_raw(_LOCK_MODEL_ROWS, sorted(model_ids_to_lock)) + ) + if frozenset(row.model_id for row in locked_rows) != model_ids_to_lock: + return None + deployments: Final = tuple( + deployment + for deployment in await tx.litellm_proxymodeltable.find_many() + if deployment.model_id in model_ids_to_lock + ) + if access_group not in _access_groups_of_deployments(deployments): + raise HTTPException( + status_code=404, + detail={"error": f"Access group '{access_group}' not found"}, + ) + selected_deployments: Final = _selected_deployments_or_400(deployments=deployments, data=data) + rewrites: Final = _membership_rewrites( + deployments=deployments, + selected_model_ids=frozenset(deployment.model_id for deployment in selected_deployments), + access_group=access_group, + ) + if rewrites: + async with tx.batch_() as batch: + for model_id, model_info in rewrites: + batch.litellm_proxymodeltable.update( + where={"model_id": model_id}, data={"model_info": json.dumps(model_info)} + ) + return selected_deployments, rewrites + + +def _lock_retry_delays() -> Iterator[float]: + return itertools.chain(_LOCK_RETRY_BACKOFF_SECONDS, itertools.repeat(_LOCK_RETRY_LONGEST_DELAY_SECONDS)) + + +async def _rewrite_membership_once_unlocked( + prisma_client: PrismaClient, + access_group: str, + data: UpdateModelGroupRequest, + retry_budget_seconds: float = _LOCK_RETRY_BUDGET_SECONDS, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + clock: Callable[[], float] = time.monotonic, +) -> tuple[tuple[_DeploymentRow, ...], tuple[tuple[str, Mapping[str, object]], ...]]: + deadline: Final = clock() + retry_budget_seconds + for delay in _lock_retry_delays(): + if ( + rewritten := await _rewrite_membership_unless_locked( + prisma_client=prisma_client, access_group=access_group, data=data + ) + ) is not None: + return rewritten + if clock() + delay > deadline: + break + await sleep(delay) + raise _http_error(409, f"Another edit holds access group '{access_group}' or one of its models, try again") + + def _readable_model_info_or_raise(model_id: str, model_info: object) -> Mapping[str, object]: """These helpers rewrite the model_info column wholesale, so a present-but-unreadable value must refuse loudly rather than be silently replaced with a fresh object; an @@ -412,11 +610,10 @@ async def _strip_access_group_from_deployment( """Remove `access_group` from one deployment's model_info; returns the (model_id, updated model_info) pair when a write happened, None when the deployment did not carry the group.""" - updated_model_info, was_modified = remove_access_group_from_deployment( - model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)), - access_group=access_group, + updated_model_info: Final = _model_info_without_access_group( + model_id=model_id, model_info=model_info, access_group=access_group ) - if not was_modified: + if updated_model_info is None: return None await _model_table(prisma_client).update( where={"model_id": model_id}, @@ -493,6 +690,35 @@ async def update_specific_deployments_with_access_group( return tuple(pair for pair in tagged if pair is not None) +def _selected_deployments_or_400( + deployments: Sequence[_DeploymentRow], data: UpdateModelGroupRequest +) -> tuple[_DeploymentRow, ...]: + if data.model_ids: + known_model_ids: Final = frozenset(deployment.model_id for deployment in deployments) + unknown_model_id: Final = next( + (model_id for model_id in data.model_ids if model_id not in known_model_ids), None + ) + if unknown_model_id is not None: + raise HTTPException( + status_code=400, + detail={"error": f"Deployment with model_id '{unknown_model_id}' not found in Database."}, + ) + selected_model_ids: Final = frozenset(data.model_ids) + return tuple(deployment for deployment in deployments if deployment.model_id in selected_model_ids) + model_names: Final = data.model_names or () + known_model_names: Final = frozenset(deployment.model_name for deployment in deployments) + unknown_model_name: Final = next((name for name in model_names if name not in known_model_names), None) + if unknown_model_name is not None: + raise HTTPException( + status_code=400, + detail={ + "error": f"Can't find model '{unknown_model_name}' in Database. Access group management is only supported for database models." + }, + ) + selected_model_names: Final = frozenset(model_names) + return tuple(deployment for deployment in deployments if deployment.model_name in selected_model_names) + + async def _find_deployment_or_400(model_id: str, prisma_client: PrismaClient) -> object: deployment: Final = await _model_table(prisma_client).find_unique(where={"model_id": model_id}) if deployment is None: @@ -536,10 +762,10 @@ async def get_all_access_groups_from_db( Returns: Dict[str, AccessGroupInfo]: Dictionary mapping access_group name to info """ - # Get all deployments - deployments: Final = await ModelRepository(prisma_client).table.find_many() + return dict(_access_groups_of_deployments(await ModelRepository(prisma_client).table.find_many())) - # Build access group map + +def _access_groups_of_deployments(deployments: Sequence[_DeploymentRow]) -> Mapping[str, AccessGroupInfo]: model_names_by_group: Final[dict[str, list[str]]] = {} for deployment in deployments: @@ -835,8 +1061,8 @@ async def update_access_group( Update an access group's model names. This will: - 1. Remove the access group from all current deployments - 2. Add the access group to all deployments for the new model_names list + 1. Add the access group to all deployments for the new model_names list + 2. Remove the access group from every other deployment Example: ```bash @@ -858,6 +1084,7 @@ async def update_access_group( Raises: - HTTPException 400: If any model names don't exist - HTTPException 404: If access group not found + - HTTPException 409: If another edit holds the group or one of its models for 30 s """ from litellm.proxy.proxy_server import llm_router, prisma_client @@ -879,26 +1106,14 @@ async def update_access_group( detail={"error": "Either model_names or model_ids must be provided and non-empty"}, ) - use_model_ids: Final = has_model_ids - - # Validation: Check if access group exists try: - access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client) - if access_group not in access_groups_map: - raise HTTPException( - status_code=404, - detail={"error": f"Access group '{access_group}' not found"}, - ) + await _raise_404_unless_the_primary_carries(access_group=access_group, prisma_client=prisma_client) except HTTPException: raise except Exception as e: - raise HTTPException( - status_code=500, - detail={"error": f"Failed to check access group existence: {e}"}, - ) + raise _http_error(500, f"Failed to check access group existence: {e}") - # Validation: Check if all new models exist (only if using model_names path) - if not use_model_ids and has_model_names: + if not has_model_ids and has_model_names: assert data.model_names is not None missing_models: Final = await _missing_models_after_read_through( model_names=data.model_names, @@ -912,43 +1127,19 @@ async def update_access_group( ) try: - # Step 1: Remove access group from ALL DB deployments (skip config models) - all_deployments: Final = await _model_table(prisma_client).find_many() - - stripped: Final = [ - await _strip_access_group_from_deployment( - model_id=deployment.model_id, - model_info=deployment.model_info, - access_group=access_group, - prisma_client=prisma_client, - ) - for deployment in all_deployments - ] - stripped_pairs: Final = tuple(pair for pair in stripped if pair is not None) - - # Step 2: Add access group using the appropriate method - if use_model_ids: - assert data.model_ids is not None - updated_pairs = await update_specific_deployments_with_access_group( - model_ids=data.model_ids, - access_group=access_group, - prisma_client=prisma_client, - ) - else: - assert data.model_names is not None - updated_pairs = await update_deployments_with_access_group( - model_names=data.model_names, - access_group=access_group, - prisma_client=prisma_client, - ) - models_updated: Final = len(updated_pairs) + selected_deployments, rewrites = await _rewrite_membership_once_unlocked( + prisma_client=prisma_client, access_group=access_group, data=data + ) + if rewrites: + await publish_model_rows_written(tuple(model_id for model_id, _ in rewrites)) + models_updated: Final = len(selected_deployments) # Clear cache and reload models to pick up the access group changes live_before_reload: Final = live_model_ids_snapshot() reload_outcome: Final = await clear_cache() _raise_http_if_reload_degraded_serving( before=live_before_reload, - written_models=list({**dict(stripped_pairs), **dict(updated_pairs)}.items()), + written_models=rewrites, access_group=access_group, still_desired=reload_outcome.still_desired, live_after=reload_outcome.live_after, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index cdfe3c90182..74fa8e3b3c1 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -779,7 +779,11 @@ async def _auto_router_capability_slot( limit=limit, ) yield tables.litellm_proxymodeltable - await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") + await publish_config_change( + redis_cache=coordination_redis_cache(), + object_type="litellm_proxymodeltable", + model_ids=() if model_id is None else (model_id,), + ) ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b0ccf74c63b..3e004f1ced7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -549,6 +549,7 @@ from litellm.proxy.common_utils.periodic_reload_schedule import ( write_reload_interval, ) from litellm.proxy.common_utils.proxy_state import ProxyState +from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState from litellm.proxy.common_utils.scheduled_job_stagger import ( @@ -8750,6 +8751,7 @@ class ProxyConfig: subscriber: Final = ConfigSyncSubscriber( redis_cache=redis_cache, resync_callbacks=(_resync_config_from_db,), + apply_model_ids=resync_model_deployments_by_ids, ) self.config_sync_subscriber = subscriber subscriber.start() diff --git a/tests/integration/management/test_added_backup_deployment_serves_peer_fallbacks.py b/tests/integration/management/test_added_backup_deployment_serves_peer_fallbacks.py new file mode 100644 index 00000000000..942bed6c165 --- /dev/null +++ b/tests/integration/management/test_added_backup_deployment_serves_peer_fallbacks.py @@ -0,0 +1,109 @@ +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Final + +import httpx +import pytest +from pydantic import TypeAdapter +from redis import Redis + +from tests.integration._support.client import Gateway, eventually, list_value, object_value, string_value +from tests.integration._support.database import scratch_database +from tests.integration._support.process import graceful_stop_seconds, owned_proxy +from tests.integration._support.redis_process import owned_redis + +pytestmark: Final = [pytest.mark.timeout(int(4 * graceful_stop_seconds() + 120))] + +COORDINATION_CONFIG: Final = Path(__file__).resolve().parents[1] / "coordination_redis_proxy_config.yaml" +CONFIG_SYNC_CHANNEL: Final = "litellm_proxy.config_change" +SUBSCRIBER_COUNTS: Final = TypeAdapter(tuple[tuple[str, int], ...]) +PERIODIC_RELOAD_SECONDS: Final = 3600 +PEER_PICKUP_SECONDS: Final = 4 +PROVIDER_KEY: Final = "integration-provider-key" + + +@contextmanager +def _proxy_pair(gateway: Gateway, directory: Path) -> Iterator[tuple[Gateway, Gateway]]: + with scratch_database("integration_model_add_sync") as database, owned_redis(directory) as coordination: + shared_stores: Final = { + "DATABASE_URL": database, + "REDIS_HOST": coordination.host, + "REDIS_PORT": str(coordination.port), + "PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": str(PERIODIC_RELOAD_SECONDS), + } + with ( + owned_proxy(gateway, directory / "writer", shared_stores, config=COORDINATION_CONFIG) as writer, + owned_proxy(gateway, directory / "peer", shared_stores, config=COORDINATION_CONFIG) as peer, + Redis(host=coordination.host, port=coordination.port, socket_timeout=2) as redis, + ): + eventually( + lambda: SUBSCRIBER_COUNTS.validate_python(redis.pubsub_numsub(CONFIG_SYNC_CHANNEL))[0][1], + lambda count: count >= 2, + seconds=30, + ) + yield writer, peer + + +@contextmanager +def _failing_upstream_model(gateway: Gateway, upstream_model: str) -> Iterator[None]: + script: Final = f"{gateway.upstream_url}/__scripts/{upstream_model}" + httpx.post(script, json={"statuses": [500]}, timeout=5, trust_env=False).raise_for_status() + try: + yield + finally: + httpx.delete(script, timeout=5, trust_env=False).raise_for_status() + + +def _add_deployment(writer: Gateway, group: str, upstream_model: str, order: int) -> None: + writer.post( + "/model/new", + { + "model_name": group, + "litellm_params": { + "model": f"openai/{upstream_model}", + "api_key": PROVIDER_KEY, + "api_base": f"{writer.upstream_url}/v1", + "order": order, + "num_retries": 0, + }, + }, + ) + + +def _listed_models(proxy: Gateway) -> frozenset[str]: + return frozenset(string_value(object_value(entry)["id"]) for entry in list_value(proxy.get("/v1/models")["data"])) + + +def _chat(proxy: Gateway, group: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": group, "messages": [{"role": "user", "content": f"order fallback {uuid.uuid4().hex}"}]}, + ) + + +def test_a_backup_added_on_one_proxy_catches_the_peers_failing_primary_right_away( + gateway: Gateway, tmp_path: Path +) -> None: + group: Final = f"integration-{uuid.uuid4().hex}" + primary: Final = f"integration-primary-{uuid.uuid4().hex}" + backup: Final = f"integration-backup-{uuid.uuid4().hex}" + with _proxy_pair(gateway, tmp_path) as (writer, peer), _failing_upstream_model(gateway, primary): + _add_deployment(writer, group, primary, order=1) + eventually(lambda: _listed_models(peer), lambda listed: group in listed, seconds=30) + + _add_deployment(writer, group, backup, order=2) + on_writer: Final = _chat(writer, group) + on_peer: Final = eventually( + lambda: _chat(peer, group), + lambda response: response.status_code == 200, + seconds=PEER_PICKUP_SECONDS, + return_last_on_timeout=True, + ) + + assert on_writer.status_code == 200, on_writer.text + assert object_value(on_writer.json())["model"] == backup, on_writer.text + assert on_peer.status_code == 200, on_peer.text + assert object_value(on_peer.json())["model"] == backup, on_peer.text diff --git a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py index 92cfb6a71e5..cbb37c964c6 100644 --- a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py @@ -5,6 +5,7 @@ from typing import Callable, Coroutine, Final, Iterable, List, Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pydantic import BaseModel, ConfigDict from redis.asyncio import Redis import litellm @@ -102,8 +103,8 @@ class _ScriptedPubSubClient: class _QueuePubSub: - def __init__(self, initial_messages: Iterable[str] = ()) -> None: - self.queue: "asyncio.Queue[str]" = asyncio.Queue() + def __init__(self, initial_messages: Iterable[object] = ()) -> None: + self.queue: "asyncio.Queue[object]" = asyncio.Queue() for message in initial_messages: self.queue.put_nowait(message) self.subscribed_channels: List[str] = [] @@ -112,7 +113,7 @@ class _QueuePubSub: async def subscribe(self, *channels: str) -> None: self.subscribed_channels.extend(channels) - async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]: if timeout == 0: try: return self.queue.get_nowait() @@ -128,7 +129,7 @@ class _QueuePubSub: class _BrokenPubSub(_QueuePubSub): - async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]: raise ConnectionError("connection lost") @@ -138,11 +139,11 @@ class _CloseFailingBrokenPubSub(_BrokenPubSub): class _EmptyPollsThenMessagePubSub(_QueuePubSub): - def __init__(self, empty_polls: int, initial_messages: Iterable[str] = ()) -> None: + def __init__(self, empty_polls: int, initial_messages: Iterable[object] = ()) -> None: super().__init__(initial_messages=initial_messages) self.remaining_empty_polls = empty_polls - async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]: if timeout != 0 and self.remaining_empty_polls > 0: self.remaining_empty_polls -= 1 return None @@ -157,6 +158,38 @@ class _FakeClock: return self.now +class _VirtualTimePubSub(_QueuePubSub): + """Timed polls move the injected clock forward instead of waiting on the wall clock. + + A delivery scheduled at a clock time is returned by the first poll whose window + reaches it, with the clock set to that time, so a test asserts when each resync + and targeted apply ran in subscriber time. + """ + + def __init__(self, clock: _FakeClock, deliveries: Iterable[Tuple[float, object]] = ()) -> None: + super().__init__() + self._clock = clock + self.pending: List[Tuple[float, object]] = sorted(deliveries, key=lambda delivery: delivery[0]) + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]: + await asyncio.sleep(0) + due_by = self._clock.now + timeout + if self.pending and self.pending[0][0] <= due_by: + delivered_at, message = self.pending.pop(0) + self._clock.now = max(self._clock.now, delivered_at) + return message + self._clock.now = due_by + return None + + +def _change(*model_ids: str) -> str: + return json.dumps({"object_type": "litellm_proxymodeltable", "model_ids": list(model_ids)}) + + +def _redis_delivery(message: str) -> dict: + return {"type": "message", "pattern": None, "channel": b"litellm_proxy.config_change", "data": message.encode()} + + class _ScriptedPubSubRedisClient(Redis): def __init__(self, pubsubs: Iterable[_QueuePubSub]) -> None: self._scripted_pubsubs = iter(pubsubs) @@ -204,7 +237,16 @@ async def test_publish_sends_object_type_json_on_channel() -> None: assert len(client.published) == 1 channel, message = client.published[0] assert channel == "litellm_proxy.config_change" - assert json.loads(message) == {"object_type": "litellm_proxymodeltable"} + assert json.loads(message) == {"object_type": "litellm_proxymodeltable", "model_ids": []} + + +async def test_publish_carries_the_written_model_ids() -> None: + client = _RecordingRedisClient() + cache = _FakeRedisCache(client) + + await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable", model_ids=("m-1", "m-2")) + + assert json.loads(client.published[0][1]) == {"object_type": "litellm_proxymodeltable", "model_ids": ["m-1", "m-2"]} async def test_publish_uses_namespaced_channel() -> None: @@ -235,7 +277,7 @@ async def test_publish_reaches_cluster_derived_pubsub_clients() -> None: await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable") assert client.published == [ - (CONFIG_SYNC_CHANNEL, json.dumps({"object_type": "litellm_proxymodeltable"})) + (CONFIG_SYNC_CHANNEL, json.dumps({"object_type": "litellm_proxymodeltable", "model_ids": []})) ] @@ -312,29 +354,210 @@ class _MaxJitterRandom(random.Random): return b -async def test_debounce_sleep_adds_jitter_from_injected_rng() -> None: - pubsub = _QueuePubSub(initial_messages=[json.dumps({"object_type": "litellm_proxymodeltable"})]) - cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) - sleeps: List[float] = [] - fired = asyncio.Event() +def _virtual_time_subscriber( + pubsub: _VirtualTimePubSub, + clock: _FakeClock, + events: List[str], + resynced: asyncio.Event, + *, + debounce_seconds: float = 0.0, + jitter_max_seconds: float = 0.0, + min_resync_interval_seconds: float = 10.0, + rng: Optional[random.Random] = None, + expected_resyncs: int = 1, + apply_error: Optional[Exception] = None, +) -> ConfigSyncSubscriber: + async def resync() -> None: + events.append(f"resync@{clock.now}") + if sum(event.startswith("resync@") for event in events) >= expected_resyncs: + resynced.set() - async def recording_sleep(seconds: float) -> None: - sleeps.append(seconds) + async def apply(model_ids: Tuple[str, ...]) -> None: + events.append(f"apply:{','.join(model_ids)}@{clock.now}") + if apply_error is not None: + raise apply_error - subscriber = ConfigSyncSubscriber( - redis_cache=cache, - resync_callbacks=(_recording_callback([], "resync", fired),), - debounce_seconds=1.0, - jitter_max_seconds=4.0, - rng=_MaxJitterRandom(), - sleep=recording_sleep, + return ConfigSyncSubscriber( + redis_cache=_FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])), + resync_callbacks=(resync,), + apply_model_ids=apply, + debounce_seconds=debounce_seconds, + jitter_max_seconds=jitter_max_seconds, + min_resync_interval_seconds=min_resync_interval_seconds, + rng=rng, + monotonic=clock, ) - subscriber.start() - await asyncio.wait_for(fired.wait(), timeout=5) - await subscriber.stop() - assert sleeps == [5.0] +async def _run_until_resynced(subscriber: ConfigSyncSubscriber, resynced: asyncio.Event) -> None: + subscriber.start() + try: + await asyncio.wait_for(resynced.wait(), timeout=5) + finally: + await subscriber.stop() + + +async def test_debounce_waits_out_the_jitter_from_injected_rng() -> None: + clock = _FakeClock() + pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, _change())]) + events: List[str] = [] + resynced = asyncio.Event() + subscriber = _virtual_time_subscriber( + pubsub, clock, events, resynced, debounce_seconds=1.0, jitter_max_seconds=4.0, rng=_MaxJitterRandom() + ) + + await _run_until_resynced(subscriber, resynced) + + assert events == ["resync@1005.0"] + + +async def test_published_model_ids_apply_on_receipt_while_the_full_resync_still_debounces() -> None: + clock = _FakeClock() + pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, _change("m-1"))]) + events: List[str] = [] + resynced = asyncio.Event() + subscriber = _virtual_time_subscriber( + pubsub, clock, events, resynced, debounce_seconds=1.0, jitter_max_seconds=4.0, rng=_MaxJitterRandom() + ) + + await _run_until_resynced(subscriber, resynced) + + assert events == ["apply:m-1@1000.0", "resync@1005.0"] + + +async def test_model_ids_published_inside_the_throttle_wait_apply_without_waiting() -> None: + clock = _FakeClock() + pubsub = _VirtualTimePubSub( + clock, deliveries=[(1000.0, _change()), (1004.0, _change("m-2")), (1006.0, _change("m-3", "m-4"))] + ) + events: List[str] = [] + resynced = asyncio.Event() + subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, expected_resyncs=2) + + await _run_until_resynced(subscriber, resynced) + + assert events == ["resync@1000.0", "apply:m-2@1004.0", "apply:m-3,m-4@1006.0", "resync@1010.0"] + + +async def test_model_ids_already_waiting_apply_together_in_one_call() -> None: + clock = _FakeClock() + pubsub = _VirtualTimePubSub( + clock, deliveries=[(1000.0, _change("m-1")), (1000.0, _change("m-2", "m-1")), (1000.0, _change("m-3"))] + ) + events: List[str] = [] + resynced = asyncio.Event() + subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced) + + await _run_until_resynced(subscriber, resynced) + + assert events == ["apply:m-1,m-2,m-3@1000.0", "resync@1000.0"] + + +async def test_redis_shaped_message_applies_its_model_ids() -> None: + clock = _FakeClock() + pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, _redis_delivery(_change("m-1")))]) + events: List[str] = [] + resynced = asyncio.Event() + subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced) + + await _run_until_resynced(subscriber, resynced) + + assert events == ["apply:m-1@1000.0", "resync@1000.0"] + + +async def test_messages_without_usable_model_ids_only_schedule_the_full_resync() -> None: + clock = _FakeClock() + pubsub = _VirtualTimePubSub( + clock, + deliveries=[ + (1000.0, "change"), + (1004.0, json.dumps({"object_type": "litellm_proxymodeltable"})), + (1005.0, json.dumps({"object_type": "litellm_proxymodeltable", "model_ids": "m-1"})), + (1006.0, {"type": "message", "pattern": None, "channel": b"c", "data": 7}), + (1007.0, None), + ], + ) + events: List[str] = [] + resynced = asyncio.Event() + subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, expected_resyncs=2) + + await _run_until_resynced(subscriber, resynced) + + assert events == ["resync@1000.0", "resync@1010.0"] + + +async def test_failing_targeted_apply_still_runs_the_full_resync() -> None: + clock = _FakeClock() + pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, _change("m-1"))]) + events: List[str] = [] + resynced = asyncio.Event() + subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, apply_error=RuntimeError("db down")) + + await _run_until_resynced(subscriber, resynced) + + assert events == ["apply:m-1@1000.0", "resync@1000.0"] + + +async def test_subscriber_without_an_apply_callback_only_resyncs() -> None: + clock = _FakeClock() + pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, _change("m-1"))]) + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=_FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])), + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.0, + jitter_max_seconds=0.0, + monotonic=clock, + ) + + await _run_until_resynced(subscriber, fired) + + assert resyncs == ["resync"] + + +async def test_peer_holds_every_published_row_before_its_next_full_resync() -> None: + publisher = _RecordingRedisClient() + cache = _FakeRedisCache(publisher) + await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable", model_ids=("m-order-1",)) + await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable", model_ids=("m-order-2",)) + on_the_wire = [_redis_delivery(message) for _, message in publisher.published] + clock = _FakeClock() + pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, on_the_wire[0]), (1007.0, on_the_wire[1])]) + db_rows_written_at = {"m-order-1": 1000.0, "m-order-2": 1007.0} + registry: set = set() + snapshots: List[Tuple[str, float, frozenset]] = [] + resynced = asyncio.Event() + + async def apply(model_ids: Tuple[str, ...]) -> None: + registry.update(model_ids) + snapshots.append(("apply", clock.now, frozenset(registry))) + + async def resync() -> None: + registry.update(row for row, written_at in db_rows_written_at.items() if written_at <= clock.now) + snapshots.append(("resync", clock.now, frozenset(registry))) + if sum(kind == "resync" for kind, _, _ in snapshots) == 2: + resynced.set() + + subscriber = ConfigSyncSubscriber( + redis_cache=_FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])), + resync_callbacks=(resync,), + apply_model_ids=apply, + debounce_seconds=1.0, + jitter_max_seconds=5.0, + min_resync_interval_seconds=10.0, + rng=_MaxJitterRandom(), + monotonic=clock, + ) + + await _run_until_resynced(subscriber, resynced) + + assert snapshots == [ + ("apply", 1000.0, frozenset({"m-order-1"})), + ("resync", 1006.0, frozenset({"m-order-1"})), + ("apply", 1007.0, frozenset({"m-order-1", "m-order-2"})), + ("resync", 1016.0, frozenset({"m-order-1", "m-order-2"})), + ] def test_default_jitter_window_is_nonzero() -> None: @@ -345,92 +568,43 @@ def test_default_min_resync_interval_caps_reload_rate() -> None: assert CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS > CONFIG_SYNC_JITTER_MAX_SECONDS -def _throttled_subscriber( - cache: object, - events: List[str], - fired: asyncio.Event, - clock: _FakeClock, - min_resync_interval_seconds: float = 10.0, -) -> ConfigSyncSubscriber: - async def recording_sleep(seconds: float) -> None: - events.append(f"sleep:{seconds}") - await asyncio.sleep(0) - - async def resync() -> None: - events.append("resync") - fired.set() - - return ConfigSyncSubscriber( - redis_cache=cache, - resync_callbacks=(resync,), - debounce_seconds=0.0, - jitter_max_seconds=0.0, - min_resync_interval_seconds=min_resync_interval_seconds, - sleep=recording_sleep, - monotonic=clock, - ) - - async def test_resync_arriving_inside_min_interval_waits_out_the_remainder() -> None: - pubsub = _QueuePubSub() - cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) - events: List[str] = [] - fired = asyncio.Event() clock = _FakeClock() - subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock) + pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, "change"), (1004.0, "change")]) + events: List[str] = [] + resynced = asyncio.Event() + subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, expected_resyncs=2) - subscriber.start() - pubsub.queue.put_nowait("change") - await asyncio.wait_for(fired.wait(), timeout=5) - fired.clear() - clock.now += 4.0 - pubsub.queue.put_nowait("change") - await asyncio.wait_for(fired.wait(), timeout=5) - await subscriber.stop() + await _run_until_resynced(subscriber, resynced) - assert events == ["sleep:0.0", "resync", "sleep:0.0", "sleep:6.0", "resync"] + assert events == ["resync@1000.0", "resync@1010.0"] async def test_resync_after_min_interval_elapsed_is_not_throttled() -> None: - pubsub = _QueuePubSub() - cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) - events: List[str] = [] - fired = asyncio.Event() clock = _FakeClock() - subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock) + pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, "change"), (1030.0, "change")]) + events: List[str] = [] + resynced = asyncio.Event() + subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, expected_resyncs=2) - subscriber.start() - pubsub.queue.put_nowait("change") - await asyncio.wait_for(fired.wait(), timeout=5) - fired.clear() - clock.now += 30.0 - pubsub.queue.put_nowait("change") - await asyncio.wait_for(fired.wait(), timeout=5) - await subscriber.stop() + await _run_until_resynced(subscriber, resynced) - assert events == ["sleep:0.0", "resync", "sleep:0.0", "resync"] + assert events == ["resync@1000.0", "resync@1030.0"] async def test_writes_during_the_throttle_wait_collapse_into_the_next_resync() -> None: - pubsub = _QueuePubSub() - cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) - events: List[str] = [] - fired = asyncio.Event() clock = _FakeClock() - subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock) + pubsub = _VirtualTimePubSub( + clock, deliveries=[(1000.0, "change")] + [(1004.0 + offset, "change") for offset in range(5)] + ) + events: List[str] = [] + resynced = asyncio.Event() + subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, expected_resyncs=2) - subscriber.start() - pubsub.queue.put_nowait("change") - await asyncio.wait_for(fired.wait(), timeout=5) - fired.clear() - for _ in range(5): - pubsub.queue.put_nowait("change") - await asyncio.wait_for(fired.wait(), timeout=5) - await asyncio.sleep(0.1) - await subscriber.stop() + await _run_until_resynced(subscriber, resynced) - assert events.count("resync") == 2 - assert pubsub.queue.empty() + assert events == ["resync@1000.0", "resync@1010.0"] + assert pubsub.pending == [] async def test_polls_without_messages_do_not_trigger_resyncs() -> None: @@ -646,13 +820,39 @@ class _AllWritesTableActions: return action -def _recording_publish(calls: List[Tuple[str, str]]) -> Callable[[str], Coroutine[None, None, None]]: - async def publish(object_type: str) -> None: +def _recording_publish(calls: List[Tuple[str, str]]) -> Callable[[str, Tuple[str, ...]], Coroutine[None, None, None]]: + async def publish(object_type: str, model_ids: Tuple[str, ...]) -> None: calls.append(("publish", object_type)) return publish +def _recording_publish_with_ids( + calls: List[Tuple[str, Tuple[str, ...]]], +) -> Callable[[str, Tuple[str, ...]], Coroutine[None, None, None]]: + async def publish(object_type: str, model_ids: Tuple[str, ...]) -> None: + calls.append((object_type, model_ids)) + + return publish + + +class _ModelRow(BaseModel): + model_config = ConfigDict(frozen=True) + + model_id: str + model_name: str + + +class _ModelRowTableActions: + def __getattr__(self, name: str) -> Callable[..., Coroutine[None, None, object]]: + async def action(*args: object, **kwargs: object) -> object: + if name.endswith("_many"): + return 2 + return _ModelRow(model_id="m-1", model_name="gpt-5.2") + + return action + + def test_wrapper_passes_through_unsynced_tables() -> None: actions = object() @@ -722,6 +922,46 @@ async def test_wrapper_publishes_for_every_write_action(action_name: str) -> Non assert publish_calls == [("publish", "litellm_guardrailstable")] +@pytest.mark.parametrize( + "action_name, expected_ids", + [ + ("create", ("m-1",)), + ("update", ("m-1",)), + ("upsert", ("m-1",)), + ("delete", ()), + ("create_many", ()), + ("update_many", ()), + ("delete_many", ()), + ], +) +async def test_wrapper_publishes_the_written_model_id_for_row_writes( + action_name: str, expected_ids: Tuple[str, ...] +) -> None: + publish_calls: List[Tuple[str, Tuple[str, ...]]] = [] + wrapped = wrap_table_actions_for_config_sync( + actions=_ModelRowTableActions(), + table_name="litellm_proxymodeltable", + publish=_recording_publish_with_ids(publish_calls), + ) + + await getattr(wrapped, action_name)(data={}) + + assert publish_calls == [("litellm_proxymodeltable", expected_ids)] + + +async def test_wrapper_publishes_no_model_ids_for_other_tables() -> None: + publish_calls: List[Tuple[str, Tuple[str, ...]]] = [] + wrapped = wrap_table_actions_for_config_sync( + actions=_ModelRowTableActions(), + table_name="litellm_guardrailstable", + publish=_recording_publish_with_ids(publish_calls), + ) + + await wrapped.create(data={}) + + assert publish_calls == [("litellm_guardrailstable", ())] + + async def test_model_repository_write_publishes_via_live_coordination_cache() -> None: from litellm.proxy import proxy_server from litellm.proxy.proxy_server import _set_redis_usage_cache @@ -747,7 +987,7 @@ async def test_model_repository_write_publishes_via_live_coordination_cache() -> assert len(client.published) == 1 channel, message = client.published[0] assert channel == CONFIG_SYNC_CHANNEL - assert json.loads(message) == {"object_type": "litellm_proxymodeltable"} + assert json.loads(message) == {"object_type": "litellm_proxymodeltable", "model_ids": ["m-1"]} async def test_ui_settings_write_publishes_via_live_coordination_cache() -> None: @@ -774,7 +1014,7 @@ async def test_ui_settings_write_publishes_via_live_coordination_cache() -> None assert len(client.published) == 1 channel, message = client.published[0] assert channel == CONFIG_SYNC_CHANNEL - assert json.loads(message) == {"object_type": "litellm_uisettings"} + assert json.loads(message) == {"object_type": "litellm_uisettings", "model_ids": []} async def _publish_calls_for_invalidated_param(param_name: str) -> List[Tuple[str, str]]: @@ -799,7 +1039,7 @@ async def test_invalidate_config_param_publishes_params_a_resync_applies(param_n assert len(published) == 1 channel, message = published[0] assert channel == CONFIG_SYNC_CHANNEL - assert json.loads(message) == {"object_type": param_name} + assert json.loads(message) == {"object_type": param_name, "model_ids": []} @pytest.mark.parametrize("param_name", _STARTUP_ONLY_CONFIG_PARAM_NAMES) @@ -932,6 +1172,46 @@ async def test_proxy_config_subscriber_resyncs_deployments_only() -> None: assert subscriber._task is None +async def test_proxy_config_subscriber_applies_published_model_ids_from_the_db( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import ProxyConfig + + rows = [MagicMock()] + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=rows) + router = MagicMock() + router.registry.return_value.deployments.return_value = ["deployment"] + installed: List[object] = [] + + def install(db_models: object) -> None: + installed.append(db_models) + + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + monkeypatch.setattr(proxy_server.proxy_config, "get_credentials", AsyncMock()) + monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", install) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()])) + config = ProxyConfig() + + config.start_config_sync_subscriber(prisma_client=prisma_client, proxy_logging_obj=MagicMock(), redis_cache=cache) + subscriber = config.config_sync_subscriber + assert subscriber is not None + apply = subscriber._apply_model_ids + assert apply is not None + await apply(("m-1", "m-2")) + await config.stop_config_sync_subscriber() + + prisma_client.db.litellm_proxymodeltable.find_many.assert_awaited_once_with( + where={"model_id": {"in": ["m-1", "m-2"]}} + ) + assert installed == [rows] + assert proxy_server.llm_model_list == ["deployment"] + + async def test_proxy_config_does_not_start_subscriber_without_coordination_redis() -> None: from litellm.proxy.proxy_server import ProxyConfig diff --git a/tests/unit/proxy/common_utils/test_registry_read_through.py b/tests/unit/proxy/common_utils/test_registry_read_through.py index 04729167707..da7b1f497c9 100644 --- a/tests/unit/proxy/common_utils/test_registry_read_through.py +++ b/tests/unit/proxy/common_utils/test_registry_read_through.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final import pytest @@ -6,7 +7,10 @@ import pytest from litellm.proxy.common_utils.registry_read_through import RegistryReadThrough if TYPE_CHECKING: + from litellm import Router + from litellm.proxy._types import LiteLLM_ProxyModelTable from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry + from litellm.types.utils import CredentialItem def nothing_loaded(_key: str) -> bool: @@ -487,6 +491,178 @@ async def test_resync_model_deployments_loads_db_credentials_before_reconciling_ installed.assert_called_once_with(db_models=rows, credential={"api_key": "sk-from-db"}) +@pytest.fixture +def stored_params_salt(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-stored-model-rows") + + +def _stored_model_row(model_id: str, stored_params: Mapping[str, str]) -> "LiteLLM_ProxyModelTable": + from litellm.proxy._types import LiteLLM_ProxyModelTable + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + return LiteLLM_ProxyModelTable( + model_id=model_id, + model_name="gpt-failover-group", + litellm_params={key: encrypt_value_helper(value) for key, value in stored_params.items()}, + model_info={"id": model_id}, + ) + + +def _backup_row(model_id: str, model: str = "openai/gpt-backup") -> "LiteLLM_ProxyModelTable": + return _stored_model_row(model_id, {"model": model, "api_key": "sk-backup"}) + + +def _router_serving(*model_ids: str) -> "Router": + from litellm import Router + + return Router( + model_list=[ + { + "model_name": "gpt-failover-group", + "litellm_params": {"model": "openai/gpt-primary", "api_key": "sk-primary"}, + "model_info": {"id": model_id}, + } + for model_id in model_ids + ] + ) + + +def _served_ids(deployments: "Sequence[Mapping[str, Mapping[str, object]]] | None") -> frozenset[object]: + return frozenset(deployment["model_info"]["id"] for deployment in deployments or ()) + + +class _ProxyModelTable: + def __init__(self, rows: "tuple[LiteLLM_ProxyModelTable, ...]") -> None: + self.rows = rows + self.reads = 0 + + async def find_many( + self, where: Mapping[str, Mapping[str, Sequence[str]]] + ) -> "tuple[LiteLLM_ProxyModelTable, ...]": + self.reads += 1 + return tuple(row for row in self.rows if row.model_id in where["model_id"]["in"]) + + async def create(self, data: Mapping[str, object]) -> None: + raise AssertionError(f"a model load wrote a row: {data}") + + +class _CredentialsTable: + def __init__(self, rows: "tuple[CredentialItem, ...]") -> None: + self.rows = rows + + async def find_many(self) -> "tuple[CredentialItem, ...]": + return self.rows + + +class _Db: + def __init__( + self, litellm_proxymodeltable: _ProxyModelTable, litellm_credentialstable: _CredentialsTable | None = None + ) -> None: + self.litellm_proxymodeltable = litellm_proxymodeltable + self.litellm_credentialstable = litellm_credentialstable or _CredentialsTable(()) + self.other_tables_read: list[str] = [] + + def __getattr__(self, table_name: str) -> object: + self.other_tables_read.append(table_name) + raise AttributeError(table_name) + + +class _PrismaClientWithoutConfigRows: + def __init__(self, db: object) -> None: + self.db = db + + async def get_generic_data(self, key: str, value: str, table_name: str) -> None: + return None + + +def _serve_db_models(monkeypatch: pytest.MonkeyPatch, db: object, router: "Router | None") -> None: + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", _PrismaClientWithoutConfigRows(db)) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("stored_params_salt") +async def test_resync_model_deployments_by_ids_installs_only_the_published_rows( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids + + router: Final = _router_serving("primary-id") + _serve_db_models( + monkeypatch, _Db(_ProxyModelTable((_backup_row("m-1"), _backup_row("m-2"), _backup_row("m-3")))), router + ) + + await resync_model_deployments_by_ids(("m-1", "m-2")) + + assert _served_ids(router.registry().deployments()) == {"primary-id", "m-1", "m-2"} + assert _served_ids(proxy_server.llm_model_list) == {"primary-id", "m-1", "m-2"} + assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("stored_params_salt") +async def test_resync_model_deployments_by_ids_leaves_the_router_alone_for_deleted_rows( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids + + router: Final = _router_serving("primary-id") + _serve_db_models(monkeypatch, _Db(_ProxyModelTable(())), router) + + await resync_model_deployments_by_ids(("deleted-id",)) + + assert _served_ids(router.registry().deployments()) == {"primary-id"} + assert proxy_server.llm_model_list is None + assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("stored_params_salt") +async def test_resync_model_deployments_by_ids_without_a_router_leaves_the_reload_to_the_debounced_resync( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids + + table: Final = _ProxyModelTable((_backup_row("m-1"),)) + db: Final = _Db(table) + _serve_db_models(monkeypatch, db, None) + + await resync_model_deployments_by_ids(("m-1",)) + + assert proxy_server.llm_router is None + assert table.reads == 1 + assert db.other_tables_read == [] + + +@pytest.mark.asyncio +async def test_resync_model_deployments_by_ids_respects_supported_db_objects( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock( + side_effect=AssertionError("db hit for an object type this replica does not load") + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["guardrails"]}) + + await resync_model_deployments_by_ids(("gated-out-id",)) + + prisma_client.db.litellm_proxymodeltable.find_many.assert_not_awaited() + + @pytest.mark.asyncio async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch): from unittest.mock import AsyncMock, MagicMock @@ -786,3 +962,174 @@ async def test_resync_model_deployments_reconciles_every_db_row_when_the_missed_ assert await _resync_model_deployments(router_name) is router_holds_the_name_after_reconcile reconciled.assert_called_once_with(prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("stored_params_salt") +async def test_resync_model_deployments_by_ids_reads_the_writer_not_a_lagging_replica( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unittest.mock import MagicMock + + from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer: Final = MagicMock(name="writer_prisma") + writer.litellm_proxymodeltable = _ProxyModelTable((_backup_row("backup-id"),)) + replica_table: Final = _ProxyModelTable(()) + replica: Final = MagicMock(name="lagging_replica_prisma") + replica.litellm_proxymodeltable = replica_table + router: Final = _router_serving("primary-id") + _serve_db_models( + monkeypatch, + RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=replica, iam_token_db_auth=False), + ), + router, + ) + + await resync_model_deployments_by_ids(("backup-id",)) + + assert router.registry().has_deployment("backup-id") + assert replica_table.reads == 0 + + +@pytest.mark.asyncio +async def test_unknown_model_name_read_through_reads_the_replica_without_waiting_for_a_reload( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer: Final = MagicMock(name="writer_prisma") + writer.litellm_proxymodeltable = SimpleNamespace(find_many=AsyncMock(return_value=[]), create=AsyncMock()) + replica: Final = MagicMock(name="replica_prisma") + replica.litellm_proxymodeltable = SimpleNamespace(find_many=AsyncMock(return_value=[]), create=AsyncMock()) + prisma_client: Final = MagicMock() + prisma_client.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=replica, iam_token_db_auth=False), + ) + reload_in_progress: Final = asyncio.Lock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) + monkeypatch.setattr(proxy_server, "MODEL_RECONCILE_LOCK", reload_in_progress) + + async with reload_in_progress: + found: Final = await asyncio.wait_for(_resync_model_deployments("no-such-model"), timeout=2) + + assert found is False + writer.litellm_proxymodeltable.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("stored_params_salt") +async def test_resync_model_deployments_by_ids_installs_the_row_a_concurrent_reload_left_current( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids + + table: Final = _ProxyModelTable((_backup_row("m-1", model="openai/gpt-before-update"),)) + router: Final = _router_serving("primary-id") + _serve_db_models(monkeypatch, _Db(table), router) + reload_in_progress: Final = asyncio.Lock() + monkeypatch.setattr(proxy_server, "MODEL_RECONCILE_LOCK", reload_in_progress) + + async with reload_in_progress: + targeted_load: Final = asyncio.create_task(resync_model_deployments_by_ids(("m-1",))) + for _ in range(5): + await asyncio.sleep(0) + assert not router.registry().has_deployment("m-1") + table.rows = (_backup_row("m-1", model="openai/gpt-after-update"),) + await targeted_load + + installed: Final = router.registry().deployment("m-1") + assert installed is not None + assert installed.litellm_params.model == "openai/gpt-after-update" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("stored_params_salt") +@pytest.mark.parametrize( + ("stored_credential_name", "credentials_on_this_worker", "api_key_this_worker_uses"), + [ + ("shared-openai", (), "sk-from-db"), + ("shared-openai", ("another-openai",), "sk-from-db"), + ("shared-openai", ("shared-openai",), "sk-on-this-worker"), + (None, (), None), + ], +) +async def test_resync_model_deployments_by_ids_reloads_credentials_only_for_one_this_worker_lacks( + monkeypatch: pytest.MonkeyPatch, + stored_credential_name: str | None, + credentials_on_this_worker: tuple[str, ...], + api_key_this_worker_uses: str | None, +) -> None: + import litellm + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids + from litellm.types.utils import CredentialItem + + model_params: Final = {"model": "openai/gpt-backup"} + row: Final = _stored_model_row( + "backup-id", + model_params if stored_credential_name is None else model_params | {"litellm_credential_name": stored_credential_name}, + ) + stored_credential: Final = CredentialItem( + credential_name="shared-openai", + credential_values={"api_key": encrypt_value_helper("sk-from-db")}, + credential_info={}, + ) + router: Final = _router_serving("primary-id") + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem(credential_name=name, credential_values={"api_key": "sk-on-this-worker"}, credential_info={}) + for name in credentials_on_this_worker + ], + ) + _serve_db_models(monkeypatch, _Db(_ProxyModelTable((row,)), _CredentialsTable((stored_credential,))), router) + + await resync_model_deployments_by_ids(("backup-id",)) + + credential: Final = CredentialAccessor.find_credential("shared-openai") + assert (credential.credential_values.get("api_key") if credential else None) == api_key_this_worker_uses + assert router.registry().has_deployment("backup-id") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("stored_params_salt") +async def test_resync_model_deployments_by_ids_installs_the_ordinary_rows_queued_with_an_auto_router( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids + + table: Final = _ProxyModelTable( + ( + _stored_model_row("auto-router-id", {"model": "auto_router/complexity-router"}), + _backup_row("backup-id"), + ) + ) + db: Final = _Db(table) + router: Final = _router_serving("primary-id") + _serve_db_models(monkeypatch, db, router) + + await resync_model_deployments_by_ids(("auto-router-id", "backup-id")) + + assert _served_ids(router.registry().deployments()) == {"primary-id", "backup-id"} + assert _served_ids(proxy_server.llm_model_list) == {"primary-id", "backup-id"} + assert table.reads == 1 + assert db.other_tables_read == [] + assert not proxy_server.MODEL_RECONCILE_LOCK.locked() diff --git a/tests/unit/proxy/management_endpoints/test_access_group_management.py b/tests/unit/proxy/management_endpoints/test_access_group_management.py index c36576beaa9..d79af0af519 100644 --- a/tests/unit/proxy/management_endpoints/test_access_group_management.py +++ b/tests/unit/proxy/management_endpoints/test_access_group_management.py @@ -687,33 +687,152 @@ class _FakeAccessGroupBudgetTable: return self.rows.pop(access_group_name, None) +class _FakeClock: + def __init__(self) -> None: + self.now = 0.0 + self.slept: list[float] = [] + + def __call__(self) -> float: + return self.now + + async def sleep(self, seconds: float) -> None: + self.slept.append(seconds) + self.now += seconds + + class _FakeModelTable: - def __init__(self, journal: list[str], deployments) -> None: + def __init__(self, journal: list[str], deployments, failing_model_ids=frozenset()) -> None: self.journal = journal self.deployments = list(deployments) + self.failing_model_ids = failing_model_ids self.updates: list[tuple[dict, dict]] = [] + self.rows_locked_by_another_session: frozenset[str] = frozenset() + self.group_locked_by_another_session = False + self.edits_the_other_session_commits: dict[str, dict] = {} + self.other_session_commits_after_refusals: int | None = None + self.lock_refusals = 0 async def find_many(self, where=None, **kwargs): - return list(self.deployments) + model_names = (where or {}).get("model_name", {}).get("in") + return [d for d in self.deployments if model_names is None or d.model_name in model_names] + + def lock_group(self, wait): + if not self.group_locked_by_another_session: + return True + if wait: + raise TimeoutError("lock_timeout expired waiting for the access group lock") + self._refuse_lock() + return False + + def lock_rows(self, model_ids, skip_locked): + held = self.rows_locked_by_another_session.intersection(model_ids) + if not held: + return list(model_ids) + if not skip_locked: + raise TimeoutError(f"timed out waiting for the row lock on {sorted(held)}") + self._refuse_lock() + return [model_id for model_id in model_ids if model_id not in held] + + def _refuse_lock(self): + self.lock_refusals += 1 + if self.lock_refusals == self.other_session_commits_after_refusals: + self._other_session_commits() + + def _other_session_commits(self): + self.deployments = [ + _FakeDeployment( + model_id=d.model_id, + model_name=d.model_name, + model_info=self.edits_the_other_session_commits.get(d.model_id, d.model_info), + ) + for d in self.deployments + ] + self.rows_locked_by_another_session = frozenset() + self.group_locked_by_another_session = False async def find_unique(self, where, include=None): return next((d for d in self.deployments if d.model_id == where["model_id"]), None) async def update(self, where, data, include=None): + if where["model_id"] in self.failing_model_ids: + raise RuntimeError(f"database write to {where['model_id']} failed") self.journal.append(f"model_table.update:{where['model_id']}") self.updates.append((dict(where), dict(data))) return None + async def create(self, data, **kwargs): + raise AssertionError("access group endpoints never create deployments") + + +class _FakeModelTableBatch: + """Queues model-table updates and applies them all or none when the batch closes, like a Prisma batch.""" + + def __init__(self, model_table: _FakeModelTable) -> None: + self.model_table = model_table + self.queued: list[tuple[dict, dict]] = [] + self.litellm_proxymodeltable = SimpleNamespace(update=self._queue_update) + + def _queue_update(self, where, data): + self.queued.append((dict(where), dict(data))) + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + if exc_type is not None: + return False + failing = [where["model_id"] for where, _ in self.queued if where["model_id"] in self.model_table.failing_model_ids] + if failing: + raise RuntimeError(f"batch write to {failing} failed, nothing applied") + for where, data in self.queued: + await self.model_table.update(where=where, data=data) + return False + + +class _FakeModelTableTx: + """An interactive transaction on the primary that journals the row lock and the read it guards.""" + + def __init__(self, journal: list[str], model_table: _FakeModelTable) -> None: + self.journal = journal + self.model_table = model_table + self.litellm_proxymodeltable = SimpleNamespace(find_many=self._find_many) + + async def query_raw(self, query, *args): + if "advisory_xact_lock" in query: + self.journal.append(f"access_group.lock:{args[0]}") + return [{"locked": self.model_table.lock_group(wait="pg_try_" not in query)}] + if "FOR UPDATE" in query: + model_ids = sorted(args[0] if args else (d.model_id for d in self.model_table.deployments)) + self.journal.append(f"model_table.lock:{','.join(model_ids)}") + locked = self.model_table.lock_rows(model_ids, skip_locked="SKIP LOCKED" in query) + return [{"model_id": model_id} for model_id in locked] + self.journal.append(f"model_table.query:{query}") + return [] + + async def _find_many(self): + self.journal.append("model_table.read") + return await self.model_table.find_many() + + def batch_(self): + return _FakeModelTableBatch(self.model_table) + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + class _FakePrismaClient: - def __init__(self, journal: list[str], deployments=()) -> None: + def __init__(self, journal: list[str], deployments=(), failing_model_ids=frozenset()) -> None: self.budget_table = _FakeBudgetTable(journal) self.access_group_budget_table = _FakeAccessGroupBudgetTable(journal, self.budget_table) - self.model_table = _FakeModelTable(journal, deployments) + self.model_table = _FakeModelTable(journal, deployments, failing_model_ids) self.db = SimpleNamespace( litellm_budgettable=self.budget_table, litellm_modelaccessgroupbudgettable=self.access_group_budget_table, litellm_proxymodeltable=self.model_table, + tx=lambda timeout=None: _FakeModelTableTx(journal, self.model_table), ) def jsonify_object(self, data): @@ -784,6 +903,12 @@ def _proxy_with_stubbed_reload(prisma): yield +def _serve_model_groups(router, model_groups): + router.registry.return_value.names.side_effect = lambda model_name=None, *, kind="groups", **_: ( + list(model_groups) if kind == "groups" else [] + ) + + def _eviction_journal(access_group): """Both auth cache keys, in the order a write path has to evict them.""" from litellm.proxy.common_utils.user_api_key_cache import ( @@ -1290,3 +1415,605 @@ async def test_deleting_an_access_group_that_never_had_a_budget_still_evicts(): assert response.models_updated == 1 assert prisma.access_group_budget_table.rows == {} _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "selection", + [{"model_ids": ["stays", "joins"]}, {"model_names": ["gpt-stays", "gpt-joins"]}], +) +async def test_updating_an_access_group_writes_only_the_deployments_that_join_or_leave_it(selection): + import json + + import litellm.proxy.proxy_server as proxy_server + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient( + [], + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="leaves", model_name="gpt-leaves"), + _deployment(model_id="joins", model_name="gpt-joins", access_groups=()), + ], + ) + + with _proxy_with_stubbed_reload(prisma): + _serve_model_groups(proxy_server.llm_router, ["gpt-stays", "gpt-leaves", "gpt-joins"]) + response = await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(**selection), + user_api_key_dict=_admin(), + ) + + written_access_groups = { + where["model_id"]: json.loads(data["model_info"])["access_groups"] + for where, data in prisma.model_table.updates + } + assert written_access_groups == {"joins": ["prod-models"], "leaves": []} + assert response.models_updated == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "selection", + [{"model_ids": ["stays", "joins"]}, {"model_names": ["gpt-stays", "gpt-joins"]}], +) +async def test_updating_an_access_group_reads_the_primary_not_a_lagging_replica(selection): + import json + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + primary = _FakePrismaClient( + [], + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="leaves", model_name="gpt-leaves"), + _deployment(model_id="joins", model_name="gpt-joins", access_groups=()), + ], + ) + replica_before_the_last_edit = _FakePrismaClient( + [], + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="leaves", model_name="gpt-leaves", access_groups=()), + _deployment(model_id="joins", model_name="gpt-joins"), + ], + ) + primary.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=primary.db, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=replica_before_the_last_edit.db, iam_token_db_auth=False), + ) + + with _proxy_with_stubbed_reload(primary): + _serve_model_groups(proxy_server.llm_router, ["gpt-stays", "gpt-leaves", "gpt-joins"]) + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(**selection), + user_api_key_dict=_admin(), + ) + + written_access_groups = { + where["model_id"]: json.loads(data["model_info"])["access_groups"] + for where, data in primary.model_table.updates + } + assert written_access_groups == {"joins": ["prod-models"], "leaves": []} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failing_model_id", ["leaves", "joins"]) +async def test_a_failed_access_group_update_leaves_every_deployment_unchanged(failing_model_id): + from fastapi import HTTPException + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + deployments = [ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="leaves", model_name="gpt-leaves"), + _deployment(model_id="joins", model_name="gpt-joins", access_groups=()), + ] + prisma = _FakePrismaClient([], deployments=deployments, failing_model_ids=frozenset({failing_model_id})) + + with _proxy_with_stubbed_reload(prisma): + _serve_model_groups(proxy_server.llm_router, ["gpt-stays", "gpt-leaves", "gpt-joins"]) + with pytest.raises(HTTPException) as raised: + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=["stays", "joins"]), + user_api_key_dict=_admin(), + ) + + assert raised.value.status_code == 500 + assert prisma.model_table.updates == [] + + +@pytest.mark.asyncio +async def test_an_access_group_update_keeps_a_model_info_field_another_admin_set_while_names_were_validated(): + import json + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient( + [], + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="joins", model_name="gpt-joins", access_groups=()), + ], + ) + + def another_admin_edits_joins_while_names_are_checked(model_name=None, *, kind="groups", **_): + if kind != "groups": + return [] + prisma.model_table.deployments = [ + deployment + if deployment.model_id != "joins" + else _FakeDeployment( + model_id="joins", + model_name="gpt-joins", + model_info={"access_groups": [], "description": "set by another admin"}, + ) + for deployment in prisma.model_table.deployments + ] + return ["gpt-stays", "gpt-joins"] + + with _proxy_with_stubbed_reload(prisma): + proxy_server.llm_router.registry.return_value.names.side_effect = ( + another_admin_edits_joins_while_names_are_checked + ) + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_names=["gpt-stays", "gpt-joins"]), + user_api_key_dict=_admin(), + ) + + written = {where["model_id"]: json.loads(data["model_info"]) for where, data in prisma.model_table.updates} + assert written["joins"]["description"] == "set by another admin" + assert written["joins"]["access_groups"] == ["prod-models"] + + +@pytest.mark.asyncio +async def test_an_access_group_update_locks_the_rows_it_reads_and_tells_other_workers_after_the_commit(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient( + journal, + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="leaves", model_name="gpt-leaves"), + _deployment(model_id="joins", model_name="gpt-joins", access_groups=()), + ], + ) + + async def record_publish(model_ids): + journal.append(f"publish:{','.join(sorted(model_ids))}") + + with ( + _proxy_with_stubbed_reload(prisma), + patch( # test-quality-ok: the endpoint publishes through its module-level import; Redis is the real seam + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.publish_model_rows_written", + new=record_publish, + ), + ): + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=["stays", "joins"]), + user_api_key_dict=_admin(), + ) + + assert journal == [ + "access_group.lock:prod-models", + "model_table.read", + "model_table.lock:joins,leaves,stays", + "model_table.read", + "model_table.update:leaves", + "model_table.update:joins", + "publish:joins,leaves", + ] + + +@pytest.mark.asyncio +async def test_an_access_group_update_that_changes_no_membership_writes_and_announces_nothing(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient( + journal, + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="also-stays", model_name="gpt-also-stays"), + _deployment(model_id="outside", model_name="gpt-outside", access_groups=()), + ], + ) + + async def record_publish(model_ids): + journal.append(f"publish:{','.join(sorted(model_ids))}") + + with ( + _proxy_with_stubbed_reload(prisma), + patch( # test-quality-ok: the endpoint publishes through its module-level import; Redis is the real seam + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.publish_model_rows_written", + new=record_publish, + ), + ): + response = await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=["stays", "also-stays"]), + user_api_key_dict=_admin(), + ) + + assert response.models_updated == 2 + assert journal == [ + "access_group.lock:prod-models", + "model_table.read", + "model_table.lock:also-stays,stays", + "model_table.read", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ignored_model_names", [None, ["gpt-unrelated"]]) +async def test_an_access_group_update_goes_through_while_another_session_holds_a_model_row_outside_the_group( + ignored_model_names, +): + import json + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient( + [], + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="leaves", model_name="gpt-leaves"), + _deployment(model_id="joins", model_name="gpt-joins", access_groups=()), + _deployment(model_id="unrelated", model_name="gpt-unrelated", access_groups=("other-group",)), + ], + ) + prisma.model_table.rows_locked_by_another_session = frozenset({"unrelated"}) + + with _proxy_with_stubbed_reload(prisma): + response = await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=["stays", "joins"], model_names=ignored_model_names), + user_api_key_dict=_admin(), + ) + + written_access_groups = { + where["model_id"]: json.loads(data["model_info"])["access_groups"] + for where, data in prisma.model_table.updates + } + assert written_access_groups == {"joins": ["prod-models"], "leaves": []} + assert response.models_updated == 2 + + +@pytest.mark.asyncio +async def test_an_access_group_update_retries_a_locked_row_and_keeps_the_edit_committed_there(): + import json + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient( + [], + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="joins", model_name="gpt-joins", access_groups=()), + ], + ) + prisma.model_table.rows_locked_by_another_session = frozenset({"joins"}) + prisma.model_table.edits_the_other_session_commits = { + "joins": {"access_groups": [], "description": "saved while the update waited"} + } + prisma.model_table.other_session_commits_after_refusals = 1 + + with _proxy_with_stubbed_reload(prisma): + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=["stays", "joins"]), + user_api_key_dict=_admin(), + ) + + written = {where["model_id"]: json.loads(data["model_info"]) for where, data in prisma.model_table.updates} + assert written == {"joins": {"access_groups": ["prod-models"], "description": "saved while the update waited"}} + + +@pytest.mark.asyncio +async def test_an_access_group_emptied_while_the_update_retried_a_locked_row_answers_404_and_writes_nothing(): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient( + [], + deployments=[ + _deployment(model_id="leaves", model_name="gpt-leaves"), + _deployment(model_id="joins", model_name="gpt-joins", access_groups=()), + ], + ) + prisma.model_table.rows_locked_by_another_session = frozenset({"leaves"}) + prisma.model_table.edits_the_other_session_commits = {"leaves": {"access_groups": []}} + prisma.model_table.other_session_commits_after_refusals = 1 + + with _proxy_with_stubbed_reload(prisma), pytest.raises(HTTPException) as raised: + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=["joins"]), + user_api_key_dict=_admin(), + ) + + assert raised.value.status_code == 404 + assert prisma.model_table.updates == [] + + +def _another_session_holds_a_member_row(model_table): + model_table.rows_locked_by_another_session = frozenset({"leaves"}) + + +def _another_edit_holds_the_group(model_table): + model_table.group_locked_by_another_session = True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "hold", [_another_session_holds_a_member_row, _another_edit_holds_the_group], ids=["member_row", "group"] +) +async def test_an_access_group_update_answers_409_and_writes_nothing_while_another_session_keeps_a_lock_it_needs( + hold, +): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + _rewrite_membership_once_unlocked, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient( + [], + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="leaves", model_name="gpt-leaves"), + ], + ) + hold(prisma.model_table) + clock = _FakeClock() + + with pytest.raises(HTTPException) as raised: + await _rewrite_membership_once_unlocked( + prisma_client=prisma, + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=["stays"]), + retry_budget_seconds=0.5, + sleep=clock.sleep, + clock=clock, + ) + + assert raised.value.status_code == 409 + assert 0 < clock.now <= 0.5 + assert prisma.model_table.lock_refusals == len(clock.slept) + 1 + assert prisma.model_table.updates == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "hold", [_another_session_holds_a_member_row, _another_edit_holds_the_group], ids=["member_row", "group"] +) +async def test_an_access_group_update_waits_out_a_lock_held_for_over_ten_seconds_and_writes_the_edit(hold): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + _rewrite_membership_once_unlocked, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient( + [], + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="leaves", model_name="gpt-leaves"), + ], + ) + hold(prisma.model_table) + prisma.model_table.other_session_commits_after_refusals = 20 + clock = _FakeClock() + + _, rewrites = await _rewrite_membership_once_unlocked( + prisma_client=prisma, + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=["stays"]), + sleep=clock.sleep, + clock=clock, + ) + + assert clock.now > 10 + assert [model_id for model_id, _ in rewrites] == ["leaves"] + assert [where["model_id"] for where, _ in prisma.model_table.updates] == ["leaves"] + + +@pytest.mark.asyncio +async def test_updating_a_missing_access_group_answers_404_before_checking_the_model_names(): + from fastapi import HTTPException + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient( + [], deployments=[_deployment(model_id="other", model_name="gpt-other", access_groups=("other-group",))] + ) + + with _proxy_with_stubbed_reload(prisma): + _serve_model_groups(proxy_server.llm_router, ["gpt-other"]) + with pytest.raises(HTTPException) as raised: + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_names=["no-such-model"]), + user_api_key_dict=_admin(), + ) + + assert raised.value.status_code == 404 + assert prisma.model_table.updates == [] + + +@pytest.mark.asyncio +async def test_an_access_group_update_answers_500_when_the_deployments_cannot_be_read(): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment(model_id="stays")]) + prisma.model_table.find_many = AsyncMock(side_effect=RuntimeError("database unreachable")) + + with _proxy_with_stubbed_reload(prisma), pytest.raises(HTTPException) as raised: + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=["stays"]), + user_api_key_dict=_admin(), + ) + + assert raised.value.status_code == 500 + assert prisma.model_table.updates == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_ids", [["no-such-deployment"], ["joins", "no-such-deployment"]]) +async def test_updating_an_access_group_with_an_unknown_model_id_writes_no_deployment(model_ids): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient( + [], + deployments=[_deployment(model_id="stays"), _deployment(model_id="joins", access_groups=())], + ) + + with _proxy_with_stubbed_reload(prisma), pytest.raises(HTTPException) as raised: + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=model_ids), + user_api_key_dict=_admin(), + ) + + assert raised.value.status_code == 400 + assert prisma.model_table.updates == [] + + +@pytest.mark.asyncio +async def test_updating_an_access_group_no_deployment_carries_answers_404_and_writes_nothing(): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment(model_id="other", access_groups=("other-group",))]) + + with _proxy_with_stubbed_reload(prisma), pytest.raises(HTTPException) as raised: + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_ids=["other"]), + user_api_key_dict=_admin(), + ) + + assert raised.value.status_code == 404 + assert prisma.model_table.updates == [] + + +@pytest.mark.asyncio +async def test_updating_an_access_group_with_a_model_name_only_the_config_serves_writes_no_deployment(): + from fastapi import HTTPException + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateModelGroupRequest, + ) + + prisma = _FakePrismaClient( + [], + deployments=[ + _deployment(model_id="stays", model_name="gpt-stays"), + _deployment(model_id="joins", model_name="gpt-joins", access_groups=()), + ], + ) + + with _proxy_with_stubbed_reload(prisma): + _serve_model_groups(proxy_server.llm_router, ["gpt-stays", "gpt-joins", "config-only-model"]) + with pytest.raises(HTTPException) as raised: + await update_access_group( + access_group="prod-models", + data=UpdateModelGroupRequest(model_names=["gpt-joins", "config-only-model"]), + user_api_key_dict=_admin(), + ) + + assert raised.value.status_code == 400 + assert prisma.model_table.updates == [] diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 14d93f368c1..872d1607351 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -6508,7 +6508,11 @@ class TestStrategyRouterWriteValidation: assert fake.tx_obj.raw_calls == [] return assert handle is fake.tx_obj.litellm_proxymodeltable - published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable") + published.assert_awaited_once_with( + redis_cache=None, + object_type="litellm_proxymodeltable", + model_ids=(model_id,) if model_id else (), + ) (lock_sql, lock_params), (count_sql, count_params) = fake.tx_obj.raw_calls assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql assert lock_params == (AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY,) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 13a064517a4..7d9e9839a5f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -699,8 +699,8 @@ export interface paths { * @description Update an access group's model names. * * This will: - * 1. Remove the access group from all current deployments - * 2. Add the access group to all deployments for the new model_names list + * 1. Add the access group to all deployments for the new model_names list + * 2. Remove the access group from every other deployment * * Example: * ```bash @@ -722,6 +722,7 @@ export interface paths { * Raises: * - HTTPException 400: If any model names don't exist * - HTTPException 404: If access group not found + * - HTTPException 409: If another edit holds the group or one of its models for 30 s */ put: operations["update_access_group_access_group__access_group__update_put"]; post?: never; From c0ded7b2824e7af9d8e5e8008fdf35d1afec0c3c Mon Sep 17 00:00:00 2001 From: mika-lt <349045450+mika-lt@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:47:59 -0700 Subject: [PATCH 18/64] fix(responses)!: send a dict reasoning_effort the same way on every Responses bridge (#49082) * fix(responses): send a dict reasoning_effort the same way on every Responses bridge Azure and ChatGPT bridges dropped reasoning mode and context unless the dict also carried summary. Every bridge now runs the reasoning step the OpenAI bridge uses, and each Responses config declares the reasoning values its API accepts, so drop_params drops what a backend refuses. Perplexity ignores mode and context, so it answers 400 for them without drop_params. * fix(responses): refuse ignored reasoning fields on xAI and Fireworks, resolve openai/ names for OpenRouter * fix(responses): refuse a reasoning mode Copilot ignores, keep catalog modes, return refusals as values Apply additional_drop_params before the refusal check, keep a null reasoning field instead of refusing it, let a cost-map row's reasoning_modes and reasoning_contexts override a backend's declared values, and declare GitHub Copilot's accepted values from a live probe * fix(responses): read Azure base_model and deployment rows, accept standard mode, refuse only mode and context * fix(responses): send only effort and summary to an Azure deployment litellm cannot identify under drop_params Azure and Azure AI deployment names are user-chosen, so when neither the name, model_info.base_model, nor a cost map row resolves the model, drop_params now drops mode and context instead of forwarding a pro the model may refuse. Names resolve by their last path segment, so a regional base_model like azure/eu/gpt-5.4 follows gpt-5.4's rows * fix(responses): resolve deployment names by model family and keep reasoning_modes per deployment A custom Azure deployment name now resolves to the longest OpenAI catalog prefix of its base_model or name, an unresolved name forwards mode and context, and a deployment's model_info reasoning_modes and reasoning_contexts apply only to that deployment instead of every sibling sharing its model name. litellm_proxy forwards every reasoning field. * fix(responses): find the model family anywhere in an Azure deployment name and keep only effort and summary for unknown gpt-5 or gpt-6 names under drop_params * fix(responses): look up OpenAI model names inside deployment names only on Azure Databricks endpoints like databricks-gpt-5-5-pro spell versions with hyphens and other backends name local models freely, so only Azure and Azure AI search inside the name, and a run that cuts a hyphen-spelled version never matches * fix(responses): judge Azure AI reasoning fields by the request api_base --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../llms/azure/responses/transformation.py | 7 + .../llms/base_llm/responses/transformation.py | 28 +- .../llms/chatgpt/responses/transformation.py | 4 + .../fireworks_ai/responses/transformation.py | 10 + .../responses/transformation.py | 7 + .../litellm_proxy/responses/transformation.py | 4 + .../perplexity/responses/transformation.py | 11 + litellm/llms/xai/responses/transformation.py | 10 + litellm/main.py | 26 +- litellm/responses/utils.py | 186 +++++-- litellm/types/utils.py | 2 + tests/unit/test_main.py | 457 ++++++++++++++++++ 12 files changed, 716 insertions(+), 36 deletions(-) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 09777da8669..2b757bb10d9 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -14,6 +14,7 @@ from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.common_utils import BaseAzureLLM +from litellm.llms.base_llm.responses.transformation import SUMMARY_ONLY_REASONING_FIELD_VALUES, ReasoningFieldValues from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import * from litellm.types.responses.main import * @@ -92,6 +93,12 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def _effort_resolves_to_none(model: str, effort: str | None) -> bool: return AzureOpenAIGPT5Config.effort_resolves_to_none(model, effort) + def deployment_names_wrap_openai_models(self) -> bool: + return True + + def unresolved_model_reasoning_field_values(self, model: str) -> ReasoningFieldValues | None: + return SUMMARY_ONLY_REASONING_FIELD_VALUES if self._is_gpt_5_model(model) else None + def get_supported_openai_params(self, model: str) -> list: """ Azure Responses API does not support context_management (compaction). diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 792b14090b8..32d0e408aba 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -1,7 +1,7 @@ import types from abc import ABC, abstractmethod from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast, get_args import httpx @@ -27,6 +27,20 @@ else: BaseLLMException = Any ReasoningSummaryLevel: TypeAlias = Literal["auto", "concise", "detailed"] +ReasoningFieldValues: TypeAlias = Mapping[str, frozenset[str | None] | None] +REASONING_SUMMARY_VALUES: Final = frozenset((*get_args(ReasoningSummaryLevel), None)) +SUMMARY_ONLY_REASONING_FIELD_VALUES: Final[ReasoningFieldValues] = types.MappingProxyType( + {"summary": REASONING_SUMMARY_VALUES} +) +STANDARD_MODE_WITHOUT_CONTEXT_REASONING_FIELD_VALUES: Final[ReasoningFieldValues] = types.MappingProxyType( + {**SUMMARY_ONLY_REASONING_FIELD_VALUES, "mode": frozenset({"standard"})} +) +STANDARD_MODE_REASONING_FIELD_VALUES: Final[ReasoningFieldValues] = types.MappingProxyType( + {**STANDARD_MODE_WITHOUT_CONTEXT_REASONING_FIELD_VALUES, "context": None} +) +ANY_REASONING_FIELD_VALUES: Final[ReasoningFieldValues] = types.MappingProxyType( + {**SUMMARY_ONLY_REASONING_FIELD_VALUES, "mode": None, "context": None} +) class BaseResponsesAPIConfig(ABC): @@ -100,6 +114,18 @@ class BaseResponsesAPIConfig(ABC): """The ``reasoning.summary`` LiteLLM's opt-in auto summary asks this provider for.""" return "detailed" + def reasoning_field_values(self) -> ReasoningFieldValues | None: + return None + + def deployment_names_wrap_openai_models(self) -> bool: + return False + + def unresolved_model_reasoning_field_values(self, model: str) -> ReasoningFieldValues | None: + return None + + def ignores_unsupported_reasoning_fields(self) -> bool: + return False + @abstractmethod def map_openai_params( self, diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 72b5cde21d1..26a5a3f150b 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -11,6 +11,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo safe_convert_created_field, ) from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.responses.transformation import STANDARD_MODE_REASONING_FIELD_VALUES, ReasoningFieldValues from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig @@ -295,3 +296,6 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): def supports_native_websocket(self) -> bool: """ChatGPT does not support native WebSocket for Responses API""" return False + + def reasoning_field_values(self) -> ReasoningFieldValues: + return STANDARD_MODE_REASONING_FIELD_VALUES diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py index 86d9df9618a..463b9cbfc9a 100644 --- a/litellm/llms/fireworks_ai/responses/transformation.py +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -6,6 +6,10 @@ from urllib.parse import unquote import httpx from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.responses.transformation import ( + STANDARD_MODE_WITHOUT_CONTEXT_REASONING_FIELD_VALUES, + ReasoningFieldValues, +) from litellm.llms.fireworks_ai.common_utils import ( FIREWORKS_FORWARD_USER_ID_PARAM, get_fireworks_forwarded_user_id, @@ -129,3 +133,9 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def supports_native_file_search(self) -> bool: return False + + def reasoning_field_values(self) -> ReasoningFieldValues: + return STANDARD_MODE_WITHOUT_CONTEXT_REASONING_FIELD_VALUES + + def ignores_unsupported_reasoning_fields(self) -> bool: + return True diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index b055d388dd5..b87a3d543e7 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -24,6 +24,7 @@ from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.responses.transformation import STANDARD_MODE_REASONING_FIELD_VALUES, ReasoningFieldValues from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import ( ResponseInputParam, @@ -437,6 +438,12 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): return False + def reasoning_field_values(self) -> ReasoningFieldValues: + return STANDARD_MODE_REASONING_FIELD_VALUES + + def ignores_unsupported_reasoning_fields(self) -> bool: + return True + def supports_native_websocket(self) -> bool: """GitHub Copilot does not support native WebSocket for Responses API""" return False diff --git a/litellm/llms/litellm_proxy/responses/transformation.py b/litellm/llms/litellm_proxy/responses/transformation.py index fc94eb30504..ed62e132c99 100644 --- a/litellm/llms/litellm_proxy/responses/transformation.py +++ b/litellm/llms/litellm_proxy/responses/transformation.py @@ -12,6 +12,7 @@ from pydantic import TypeAdapter, ValidationError from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.responses.transformation import ANY_REASONING_FIELD_VALUES, ReasoningFieldValues from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.utils import LlmProviders @@ -42,6 +43,9 @@ class LiteLLMProxyResponsesAPIConfig(OpenAIResponsesAPIConfig): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.LITELLM_PROXY + def reasoning_field_values(self) -> ReasoningFieldValues: + return ANY_REASONING_FIELD_VALUES + def get_complete_url( self, api_base: str | None, diff --git a/litellm/llms/perplexity/responses/transformation.py b/litellm/llms/perplexity/responses/transformation.py index 58f71028451..afd27b0f696 100644 --- a/litellm/llms/perplexity/responses/transformation.py +++ b/litellm/llms/perplexity/responses/transformation.py @@ -9,6 +9,7 @@ The only provider quirks: Ref: https://docs.perplexity.ai/api-reference/responses-post """ +from types import MappingProxyType from typing import Any, Final import httpx @@ -17,6 +18,7 @@ from pydantic import TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.responses.transformation import STANDARD_MODE_REASONING_FIELD_VALUES, ReasoningFieldValues from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig, drops_params from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ResponseInputParam, ResponsesAPIResponse @@ -24,6 +26,9 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders _DECODED_JSON: Final = TypeAdapter(object) +PERPLEXITY_REASONING_FIELD_VALUES: Final[ReasoningFieldValues] = MappingProxyType( + {**STANDARD_MODE_REASONING_FIELD_VALUES, "context": frozenset({"auto", "current_turn", "all_turns"})} +) class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): @@ -127,3 +132,9 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): def supports_native_websocket(self) -> bool: """Perplexity does not support native WebSocket for Responses API""" return False + + def reasoning_field_values(self) -> ReasoningFieldValues: + return PERPLEXITY_REASONING_FIELD_VALUES + + def ignores_unsupported_reasoning_fields(self) -> bool: + return True diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index b6ea516fb79..c907fde9926 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -10,6 +10,10 @@ from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.responses.transformation import ( + STANDARD_MODE_WITHOUT_CONTEXT_REASONING_FIELD_VALUES, + ReasoningFieldValues, +) from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd from litellm.secret_managers.main import get_secret_str @@ -290,3 +294,9 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def supports_native_websocket(self) -> bool: """XAI does not support native WebSocket for Responses API""" return False + + def reasoning_field_values(self) -> ReasoningFieldValues: + return STANDARD_MODE_WITHOUT_CONTEXT_REASONING_FIELD_VALUES + + def ignores_unsupported_reasoning_fields(self) -> bool: + return True diff --git a/litellm/main.py b/litellm/main.py index 8d6e07d6a8f..54afe654cec 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6267,6 +6267,9 @@ def completion( # reasoningSummary/reasoning_summary without tools (AI SDK) that the first # (early) check doesn't cover. _reasoning_summary_for_bridge: Final = peek_reasoning_summary_aliases(optional_params) + typed_deployment_model_info: Final = str_keyed_mapping( + cast(object, model_info) # cast-ok: model_info is the untyped dict kwarg + ) if responses_api_model_info.get("mode") != "responses": responses_api_model_info, model = responses_api_bridge_check( model=model, @@ -6276,9 +6279,7 @@ def completion( reasoning_effort=bridge_reasoning_effort(resolved_reasoning_effort, optional_params.get("extra_body")), reasoning_summary=_reasoning_summary_for_bridge, api_base=api_base, - deployment_model_info=str_keyed_mapping( - cast(object, model_info) # cast-ok: model_info is the untyped dict kwarg - ), + deployment_model_info=typed_deployment_model_info, ) # Use base_model (the true underlying model) for Azure model-type @@ -6290,19 +6291,28 @@ def completion( optional_params, rs_val = strip_reasoning_summary_aliases_from_optional_params(optional_params) - if custom_llm_provider == "openai" and isinstance(resolved_reasoning_effort, dict): - from litellm.responses.utils import openai_bridge_reasoning + if isinstance(resolved_reasoning_effort, dict): + from litellm.responses.utils import ( + RefusedReasoningFields, + bridge_reasoning, + raise_refused_reasoning_fields, + ) - optional_params["reasoning_effort"] = openai_bridge_reasoning( + bridged_reasoning: Final = bridge_reasoning( requested=resolved_reasoning_effort, effort_kept="reasoning_effort" in optional_params, summary_alias=rs_val, model=model, + base_model=base_model if isinstance(base_model, str) else None, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + deployment_model_info=typed_deployment_model_info, drop_params=drop_unsupported_reasoning, additional_drop_params=routing_kwargs.get("additional_drop_params"), ) - elif isinstance(resolved_reasoning_effort, dict) and "summary" in resolved_reasoning_effort: - optional_params["reasoning_effort"] = resolved_reasoning_effort + if isinstance(bridged_reasoning, RefusedReasoningFields): + raise_refused_reasoning_fields(bridged_reasoning, model, custom_llm_provider) + optional_params["reasoning_effort"] = bridged_reasoning elif rs_val is not None: eff: Final = optional_params.get("reasoning_effort", resolved_reasoning_effort) if isinstance(eff, dict): diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index b0db0516658..1a97c1cb48c 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1,10 +1,11 @@ import base64 import os import re -from collections.abc import Callable, Collection, Iterable, Mapping, Sequence +from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass from functools import reduce from types import MappingProxyType -from typing import Any, Final, Optional, TypeVar, Union, cast, get_args, get_type_hints, overload +from typing import Any, Final, NoReturn, Optional, TypeVar, Union, cast, get_type_hints, overload from pydantic import BaseModel, TypeAdapter from typing_extensions import TypedDict as ExtraItemsTypedDict @@ -15,7 +16,13 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, iter_items from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils -from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig, ReasoningSummaryLevel +from litellm.llms.base_llm.responses.transformation import ( + ANY_REASONING_FIELD_VALUES, + REASONING_SUMMARY_VALUES, + BaseResponsesAPIConfig, + ReasoningFieldValues, + ReasoningSummaryLevel, +) from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.openai import ( AllMessageValues, @@ -1411,9 +1418,6 @@ def reasoning_with_summary_default( return {**reasoning} if summary is None else {**reasoning, "summary": summary} -_SUMMARY_VALUES: Final = frozenset((*get_args(ReasoningSummaryLevel), None)) - - def _openai_model_info(model: str) -> ModelInfo | None: from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration @@ -1425,24 +1429,115 @@ def _declared_values(values: Sequence[str] | None) -> frozenset[str | None] | No return None if values is None else frozenset(values) -def _openai_reasoning_field_values(model: str) -> Mapping[str, frozenset[str | None] | None]: +_REFUSABLE_REASONING_FIELDS: Final = frozenset({"mode", "context"}) + + +def _openai_reasoning_field_values(model: str) -> ReasoningFieldValues: model_info: Final = _openai_model_info(model) if model_info is None: - return MappingProxyType({"summary": _SUMMARY_VALUES, "mode": None, "context": None}) + return ANY_REASONING_FIELD_VALUES if not model_info.get("supports_reasoning"): return MappingProxyType({}) return MappingProxyType( { - "summary": _SUMMARY_VALUES, + "summary": REASONING_SUMMARY_VALUES, "mode": _declared_values(model_info.get("reasoning_modes")), "context": _declared_values(model_info.get("reasoning_contexts")), } ) -def _openai_accepts_reasoning_field( - key: str, value: object, values_by_field: Mapping[str, frozenset[str | None] | None] -) -> bool: +_CATALOG_REASONING_FIELDS: Final = (("mode", "reasoning_modes"), ("context", "reasoning_contexts")) + + +def _catalog_values(values: object) -> frozenset[str | None]: + return frozenset(value for value in iter_items(values) if isinstance(value, str)) + + +def _declared_reasoning_values(row: Mapping[str, object] | None) -> Mapping[str, frozenset[str | None]]: + if row is None: + return {} + return { + field: _catalog_values(row[key]) for field, key in _CATALOG_REASONING_FIELDS if isinstance(row.get(key), list) + } + + +def _with_declared_reasoning_values( + defaults: ReasoningFieldValues | None, + model: str, + custom_llm_provider: str, + deployment_model_info: Mapping[str, object] | None, +) -> ReasoningFieldValues | None: + catalog: Final = cast(Mapping[str, object], litellm.model_cost) # cast-ok: read-only view of the legacy registry + catalog_row: Final = as_str_mapping(catalog.get(f"{custom_llm_provider}/{model.removeprefix('responses/')}")) + overrides: Final = { + **_declared_reasoning_values(catalog_row), + **_declared_reasoning_values(deployment_model_info), + } + if not overrides: + return defaults + return MappingProxyType({**(ANY_REASONING_FIELD_VALUES if defaults is None else defaults), **overrides}) + + +def _is_version_number(part: str) -> bool: + return part.isdigit() and len(part) <= 2 + + +def _splits_a_version(parts: Sequence[str], end: int) -> bool: + return end < len(parts) and _is_version_number(parts[end - 1]) and _is_version_number(parts[end]) + + +def _hyphen_runs_longest_first(name: str) -> Iterator[str]: + parts: Final = name.rsplit("/", 1)[-1].split("-") + for length in range(len(parts), 0, -1): + yield from ( + "-".join(parts[start : start + length]) + for start in range(len(parts) - length + 1) + if not _splits_a_version(parts, start + length) + ) + + +def _openai_catalog_family(name: str) -> str | None: + runs: Final = _hyphen_runs_longest_first(name) + return next((run for run in runs if OpenAIGPTConfig.is_openai_catalog_model(run)), None) + + +def _openai_catalog_name(name: str) -> str | None: + last_segment: Final = name.rsplit("/", 1)[-1] + return last_segment if OpenAIGPTConfig.is_openai_catalog_model(last_segment) else None + + +def _openai_catalog_model(model: str, base_model: str | None, names_wrap_models: bool) -> str | None: + find: Final = _openai_catalog_family if names_wrap_models else _openai_catalog_name + families: Final = (find(name) for name in (base_model, model) if name) + return next((family for family in families if family is not None), None) + + +def _default_reasoning_field_values( + model: str, base_model: str | None, responses_config: BaseResponsesAPIConfig | None +) -> ReasoningFieldValues | None: + declared: Final = None if responses_config is None else responses_config.reasoning_field_values() + if declared is not None: + return declared + names_wrap_models: Final = responses_config is not None and responses_config.deployment_names_wrap_openai_models() + catalog_model: Final = _openai_catalog_model(model, base_model, names_wrap_models) + if catalog_model is not None: + return _openai_reasoning_field_values(catalog_model) + return None if responses_config is None else responses_config.unresolved_model_reasoning_field_values(model) + + +def _reasoning_field_values( + model: str, + base_model: str | None, + custom_llm_provider: str, + responses_config: BaseResponsesAPIConfig | None, + deployment_model_info: Mapping[str, object] | None, +) -> ReasoningFieldValues | None: + defaults: Final = _default_reasoning_field_values(model, base_model, responses_config) + return _with_declared_reasoning_values(defaults, model, custom_llm_provider, deployment_model_info) + + +def _accepts_reasoning_field(key: str, value: object, values_by_field: ReasoningFieldValues) -> bool: if key == "effort": return True if key not in values_by_field: @@ -1451,11 +1546,8 @@ def _openai_accepts_reasoning_field( return allowed is None or (isinstance(value, str | None) and value in allowed) -def _openai_accepted_reasoning(reasoning: Mapping[str, object], model: str) -> Mapping[str, object]: - values_by_field: Final = _openai_reasoning_field_values(model) - return { - key: value for key, value in reasoning.items() if _openai_accepts_reasoning_field(key, value, values_by_field) - } +def _accepted_reasoning(reasoning: Mapping[str, object], values_by_field: ReasoningFieldValues) -> Mapping[str, object]: + return {key: value for key, value in reasoning.items() if _accepts_reasoning_field(key, value, values_by_field)} def _additional_drop_paths(additional_drop_params: object) -> tuple[str, ...]: @@ -1464,26 +1556,66 @@ def _additional_drop_paths(additional_drop_params: object) -> tuple[str, ...]: return tuple(path for path in iter_items(additional_drop_params) if isinstance(path, str)) -def openai_bridge_reasoning( +@dataclass(frozen=True, slots=True) +class RefusedReasoningFields: + refused: tuple[str, ...] + + +def raise_refused_reasoning_fields(refusal: RefusedReasoningFields, model: str, custom_llm_provider: str) -> NoReturn: + raise litellm.UnsupportedParamsError( + message=( + f"{custom_llm_provider} does not support parameters: " + f"{[f'reasoning_effort.{key}' for key in refusal.refused]}, for model={model}. " + "To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n" + ), + llm_provider=custom_llm_provider, + model=model, + ) + + +def bridge_reasoning( *, requested: Mapping[str, object], effort_kept: bool, summary_alias: object, model: str, + base_model: str | None, + custom_llm_provider: str, + api_base: str | None, + deployment_model_info: Mapping[str, object] | None, drop_params: bool, additional_drop_params: object, -) -> Mapping[str, object] | None: +) -> Mapping[str, object] | RefusedReasoningFields | None: + from litellm.utils import ProviderConfigManager + drop_paths: Final = _additional_drop_paths(additional_drop_params) if "reasoning_effort" in drop_paths: return None kept: Final = requested if effort_kept else {key: value for key, value in requested.items() if key != "effort"} - with_summary: Final = reasoning_with_summary_default(kept, summary_alias, model, "openai") - accepted: Final = ( - _openai_accepted_reasoning(with_summary, model) - if drop_params and OpenAIGPTConfig.is_openai_catalog_model(model) - else {**with_summary} + with_summary: Final = reasoning_with_summary_default(kept, summary_alias, model, custom_llm_provider) + sent: Final = ( + as_str_mapping( + _apply_nested_drop_params({"reasoning_effort": {**with_summary}}, list(drop_paths)).get("reasoning_effort") + ) + or {} ) - remaining: Final = as_str_mapping( - _apply_nested_drop_params({"reasoning_effort": accepted}, list(drop_paths)).get("reasoning_effort") + responses_config: Final = ProviderConfigManager.get_provider_responses_api_config( + custom_llm_provider, model, api_base ) - return remaining or None + values_by_field: Final = _reasoning_field_values( + model, base_model, custom_llm_provider, responses_config, deployment_model_info + ) + accepted: Final = sent if values_by_field is None else _accepted_reasoning(sent, values_by_field) + refused: Final = tuple( + key + for key, value in sent.items() + if key in _REFUSABLE_REASONING_FIELDS and key not in accepted and value is not None + ) + if ( + refused + and not drop_params + and responses_config is not None + and responses_config.ignores_unsupported_reasoning_fields() + ): + return RefusedReasoningFields(refused) + return (accepted if drop_params else sent) or None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5905aaa3380..6a758961c7b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4155,6 +4155,8 @@ DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( "supports_assistant_prefill", "supports_bedrock_native_structured_output", "supports_pdf_input", + "reasoning_modes", + "reasoning_contexts", } ) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index b2ae6b43d26..b0bf1cb879e 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -2463,6 +2463,463 @@ _BRIDGED_RESPONSE_ANSWER: Final = json.dumps( ).encode() +@pytest.fixture +def non_openai_responses_bridge( + restore_model_registry, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> Callable[[str], tuple[str, Mapping[str, str], respx.Route]]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("LITELLM_REASONING_AUTO_SUMMARY", raising=False) + for chatgpt_env in ("CHATGPT_AUTH_FILE", "CHATGPT_API_BASE", "OPENAI_CHATGPT_API_BASE"): + monkeypatch.delenv(chatgpt_env, raising=False) + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + (tmp_path / "auth.json").write_text( + json.dumps({"access_token": "token", "account_id": "acct", "expires_at": 4102444800}) + ) + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + (tmp_path / "api-key.json").write_text( + json.dumps({"token": "token", "expires_at": 4102444800, "endpoints": {"api": "https://api.githubcopilot.com"}}) + ) + litellm.register_model( + { + **{ + model_name: { + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "reasoning_modes": ["standard"], + "reasoning_contexts": ["auto", "current_turn", "all_turns"], + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + for model_name in ("gpt-6", "gpt-6.1-standard-only") + }, + **{ + f"{provider}/{model_name}": { + "litellm_provider": provider, + "mode": "responses", + "supports_reasoning": True, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + for provider, model_name in ( + ("chatgpt", "gpt-6.1-bridge-test"), + ("perplexity", "bridge-test-model"), + ("github_copilot", "gpt-6.1-bridge-test"), + ) + }, + }, + persist_across_reloads=False, + ) + backends: Final = { + "azure": ( + "azure/responses/gpt-6.1-standard-only", + "https://test.openai.azure.com/", + {"api_base": "https://test.openai.azure.com", "api_key": "fake-key"}, + ), + "azure_ai": ( + "azure_ai/responses/gpt-6.1-standard-only-prod", + _FOUNDRY_API_BASE, + {"api_base": _FOUNDRY_API_BASE, "api_key": "fake-key"}, + ), + "chatgpt": ("chatgpt/gpt-6.1-bridge-test", "https://chatgpt.com/backend-api/codex/responses", {}), + "github_copilot": ("github_copilot/gpt-6.1-bridge-test", "https://api.githubcopilot.com/responses", {}), + "databricks": ( + "databricks/responses/databricks-gpt-bridge-test", + "https://dbx.example/serving-endpoints/responses", + {"api_base": "https://dbx.example/serving-endpoints", "api_key": "fake-key"}, + ), + "perplexity": ( + "perplexity/bridge-test-model", + "https://api.perplexity.ai/v1/responses", + {"api_key": "fake-key"}, + ), + "xai": ("xai/responses/grok-bridge-test", "https://api.x.ai/v1/responses", {"api_key": "fake-key"}), + "fireworks_ai": ( + "fireworks_ai/responses/accounts/fireworks/models/bridge-test", + "https://api.fireworks.ai/inference/v1/responses", + {"api_key": "fake-key"}, + ), + "openrouter": ( + "openrouter/responses/openai/gpt-6.1-standard-only", + "https://openrouter.ai/api/v1/responses", + {"api_key": "fake-key"}, + ), + "litellm_proxy": ( + "litellm_proxy/responses/gpt-6.1-standard-only", + "https://proxy.example/v1/responses", + {"api_base": "https://proxy.example/v1", "api_key": "fake-key"}, + ), + } + + def route_backend(backend: str) -> tuple[str, Mapping[str, str], respx.Route]: + model, url, kwargs = backends[backend] + route: Final = respx_mock.post(url__startswith=url).respond( + content=_BRIDGED_RESPONSE_ANSWER, headers={"content-type": "application/json"} + ) + return model, kwargs, route + + return route_backend + + +@pytest.mark.parametrize( + ("backend", "reasoning_effort"), + [ + ("azure", {"effort": "low", "mode": "pro", "context": "all_turns"}), + ("chatgpt", {"effort": "low", "mode": "pro", "context": "all_turns"}), + ("github_copilot", {"effort": "low", "mode": "standard", "context": "all_turns"}), + ("databricks", {"effort": "low", "mode": "pro", "context": "all_turns"}), + ("openrouter", {"effort": "low", "mode": "pro", "context": "all_turns"}), + ("perplexity", {"effort": "low", "mode": "standard", "context": "all_turns", "generate_summary": "auto"}), + ("xai", {"effort": "low", "mode": "standard", "generate_summary": "auto"}), + ("fireworks_ai", {"effort": "low", "mode": "standard", "generate_summary": "auto"}), + ], +) +@pytest.mark.parametrize("summary", [{}, {"summary": "auto"}]) +def test_non_openai_responses_bridge_sends_the_whole_reasoning_effort_dict_with_or_without_summary( + backend: str, + reasoning_effort: dict[str, str], + summary: dict[str, str], + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + model, kwargs, route = non_openai_responses_bridge(backend) + requested: Final = {**reasoning_effort, **summary} + + litellm.completion( + model=model, messages=[{"role": "user", "content": "hello"}], reasoning_effort=requested, **kwargs + ) + + assert json.loads(route.calls[0].request.content)["reasoning"] == requested + + +@pytest.mark.parametrize( + ("backend", "reasoning_effort", "expected_reasoning"), + [ + ("azure", {"effort": "low", "mode": "pro", "context": "all_turns"}, {"effort": "low", "context": "all_turns"}), + ( + "chatgpt", + {"effort": "low", "mode": "pro", "context": "all_turns"}, + {"effort": "low", "context": "all_turns"}, + ), + ( + "chatgpt", + {"effort": "low", "mode": "standard", "context": "current_turn"}, + {"effort": "low", "mode": "standard", "context": "current_turn"}, + ), + ( + "github_copilot", + {"effort": "low", "mode": "pro", "context": "all_turns"}, + {"effort": "low", "context": "all_turns"}, + ), + ( + "openrouter", + {"effort": "low", "mode": "pro", "context": "all_turns"}, + {"effort": "low", "context": "all_turns"}, + ), + ( + "perplexity", + {"effort": "low", "mode": "pro", "context": "all_turns"}, + {"effort": "low", "context": "all_turns"}, + ), + ("perplexity", {"effort": "low", "mode": "standard", "context": "every_turn"}, {"effort": "low", "mode": "standard"}), + ("xai", {"effort": "low", "mode": "pro", "context": "all_turns"}, {"effort": "low"}), + ("fireworks_ai", {"effort": "low", "mode": "pro", "context": "all_turns"}, {"effort": "low"}), + ], +) +@pytest.mark.parametrize("summary", [{}, {"summary": "auto"}]) +def test_non_openai_responses_bridge_drops_reasoning_fields_the_backend_refuses_under_drop_params( + backend: str, + reasoning_effort: dict[str, str], + expected_reasoning: dict[str, str], + summary: dict[str, str], + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + model, kwargs, route = non_openai_responses_bridge(backend) + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={**reasoning_effort, **summary}, + drop_params=True, + **kwargs, + ) + + assert json.loads(route.calls[0].request.content)["reasoning"] == {**expected_reasoning, **summary} + + +@pytest.mark.parametrize( + ("backend", "ignored"), + [ + ("perplexity", {"mode": "pro"}), + ("perplexity", {"context": "every_turn"}), + ("xai", {"context": "all_turns"}), + ("fireworks_ai", {"context": "all_turns"}), + ("github_copilot", {"mode": "pro"}), + ], +) +@pytest.mark.parametrize("summary", [{}, {"summary": "auto"}]) +def test_responses_bridge_rejects_reasoning_fields_the_backend_ignores_without_drop_params( + backend: str, + ignored: dict[str, str], + summary: dict[str, str], + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + model, kwargs, route = non_openai_responses_bridge(backend) + + with pytest.raises(litellm.UnsupportedParamsError, match=rf"reasoning_effort\.{next(iter(ignored))}") as raised: + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", **ignored, **summary}, + **kwargs, + ) + + assert raised.value.status_code == 400 + assert route.called is False + + +@pytest.mark.parametrize("backend", ["perplexity", "xai", "fireworks_ai"]) +def test_responses_bridge_sends_the_rest_when_additional_drop_params_removes_the_refused_fields( + backend: str, + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + model, kwargs, route = non_openai_responses_bridge(backend) + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "all_turns", "summary": "auto"}, + additional_drop_params=["reasoning_effort.mode", "reasoning_effort.context"], + **kwargs, + ) + + assert json.loads(route.calls[0].request.content)["reasoning"] == {"effort": "low", "summary": "auto"} + + +@pytest.mark.parametrize("backend", ["perplexity", "xai", "fireworks_ai"]) +def test_responses_bridge_forwards_a_null_reasoning_field_instead_of_rejecting_it( + backend: str, + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + model, kwargs, route = non_openai_responses_bridge(backend) + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": None}, + **kwargs, + ) + + assert json.loads(route.calls[0].request.content)["reasoning"] == {"effort": "low", "mode": None} + + +def test_chatgpt_bridge_keeps_a_reasoning_mode_its_catalog_row_declares_under_drop_params( + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + _, kwargs, route = non_openai_responses_bridge("chatgpt") + litellm.register_model( + { + "chatgpt/gpt-6.1-pro-bridge-test": { + "litellm_provider": "chatgpt", + "mode": "responses", + "supports_reasoning": True, + "reasoning_modes": ["pro"], + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + + litellm.completion( + model="chatgpt/gpt-6.1-pro-bridge-test", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro"}, + drop_params=True, + **kwargs, + ) + + assert json.loads(route.calls[0].request.content)["reasoning"] == {"effort": "low", "mode": "pro"} + + +@pytest.mark.parametrize( + ("model", "model_info"), + [ + ("azure/responses/custom-deployment-name", {"base_model": "azure/gpt-6.1-standard-only"}), + ("azure/responses/custom-deployment-name", {"base_model": "azure/eu/gpt-6.1-standard-only"}), + ("azure/responses/custom-deployment-name", {"base_model": "azure/gpt-6.1-standard-only-2026-07-09"}), + ("azure/responses/gpt-6.1-standard-only-prod", {}), + ("azure/responses/eu-gpt-6.1-standard-only-prod", {}), + ("azure/responses/gpt-6-2026-07-09", {}), + ], +) +def test_azure_bridge_reads_reasoning_modes_from_the_deployment_model_family_under_drop_params( + model: str, + model_info: dict[str, str], + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + _, kwargs, route = non_openai_responses_bridge("azure") + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "all_turns"}, + model_info=model_info, + drop_params=True, + **kwargs, + ) + + assert json.loads(route.calls[0].request.content)["reasoning"] == {"effort": "low", "context": "all_turns"} + + +def test_azure_ai_bridge_judges_reasoning_fields_by_the_request_api_base_under_drop_params( + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://legacy-models.eastus2.inference.ai.azure.com") + model, kwargs, route = non_openai_responses_bridge("azure_ai") + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "all_turns"}, + drop_params=True, + **kwargs, + ) + + assert json.loads(route.calls[0].request.content)["reasoning"] == {"effort": "low", "context": "all_turns"} + + +@pytest.mark.parametrize("model", ["azure/responses/custom-gpt-6x", "azure/responses/gpt-6-1-sol"]) +@pytest.mark.parametrize("summary", [{}, {"summary": "auto"}]) +def test_azure_bridge_sends_only_effort_and_summary_to_an_unknown_gpt_series_deployment_under_drop_params( + model: str, + summary: dict[str, str], + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + _, kwargs, route = non_openai_responses_bridge("azure") + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "all_turns", **summary}, + drop_params=True, + **kwargs, + ) + + assert json.loads(route.calls[0].request.content)["reasoning"] == {"effort": "low", **summary} + + +@pytest.mark.parametrize( + ("backend", "model"), + [ + ("azure", "azure/responses/terra-prod"), + ("databricks", "databricks/responses/databricks-gpt-bridge-test"), + ("databricks", "databricks/responses/databricks-gpt-6.1-standard-only"), + ("litellm_proxy", "litellm_proxy/responses/gpt-6.1-standard-only"), + ], +) +@pytest.mark.parametrize("summary", [{}, {"summary": "auto"}]) +def test_responses_bridge_forwards_mode_and_context_it_cannot_judge_under_drop_params( + backend: str, + model: str, + summary: dict[str, str], + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + _, kwargs, route = non_openai_responses_bridge(backend) + requested: Final = {"effort": "low", "mode": "pro", "context": "all_turns", **summary} + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=requested, + drop_params=True, + **kwargs, + ) + + assert json.loads(route.calls[0].request.content)["reasoning"] == requested + + +def test_azure_bridge_applies_a_deployment_reasoning_modes_only_to_that_deployment_under_drop_params( + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + _, kwargs, route = non_openai_responses_bridge("azure") + litellm.register_model( + { + "gpt-6.1-pro-capable": { + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "reasoning_modes": ["standard", "pro"], + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + litellm_params: Final = {"model": "azure/responses/gpt-6.1-pro-capable", **kwargs} + router: Final = litellm.Router( + model_list=[ + { + "model_name": "standard-only", + "litellm_params": litellm_params, + "model_info": {"reasoning_modes": ["standard"]}, + }, + {"model_name": "pro-capable", "litellm_params": litellm_params}, + ] + ) + requested: Final = {"effort": "low", "mode": "pro"} + + for model_name in ("standard-only", "pro-capable"): + router.completion( + model=model_name, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=requested, + drop_params=True, + ) + litellm.completion( + model="azure/responses/gpt-6.1-pro-capable", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=requested, + drop_params=True, + **kwargs, + ) + + assert [json.loads(call.request.content)["reasoning"] for call in route.calls] == [ + {"effort": "low"}, + requested, + requested, + ] + + +def test_azure_bridge_reads_reasoning_modes_from_the_deployment_cost_map_row_under_drop_params( + non_openai_responses_bridge: Callable[[str], tuple[str, Mapping[str, str], respx.Route]], +): + _, kwargs, route = non_openai_responses_bridge("azure") + litellm.register_model( + { + "azure/custom-standard-deployment": { + "litellm_provider": "azure", + "mode": "responses", + "supports_reasoning": True, + "reasoning_modes": ["standard"], + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + + litellm.completion( + model="azure/responses/custom-standard-deployment", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "all_turns"}, + drop_params=True, + **kwargs, + ) + + assert json.loads(route.calls[0].request.content)["reasoning"] == {"effort": "low", "context": "all_turns"} + + def test_responses_bridge_fake_stream_is_bounded_by_the_request_timeout_not_the_stream_timeout( restore_model_registry: None, ) -> None: From b8283d8b61f9d15bdc734d58a7f8a0a38cff5bd4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:58:36 -0700 Subject: [PATCH 19/64] feat(usage): show subscription-covered providers and account fees separately (#44772) * feat(usage): show subscription-covered providers and account fees separately * fix(usage): write only the named fee columns and sync the root schema The fee update now writes just the columns the request names, so two admins changing different fields at once both land, and the dashboard edit form sends only what changed. The root schema.prisma carries the new table like the two copies, and a billing period anchored in the calendar's last month ends on December 31, 9999 instead of failing. * fix(usage): register the subscription table span and keep provider logic under llms * fix(subscription_accounts): register the fee routes as management routes and reject an empty account id * fix(usage): run the new tests in CI and close subscription card review gaps Adds the subscription test file to its unit shard, sets the migration lock timeout, rejects infinite fees, keeps the card from rewriting login files, lists each deployment once without aliases, shows the provider badge on entity views, and asks before removing a fee * fix(usage): resolve named-credential ChatGPT accounts and keep admin viewers read-only on the fee card * fix(usage): count a subscription fee only for accounts behind a current deployment Usage listed a fee row for an account no deployment uses while another ChatGPT deployment existed and dropped it when none did. It now lists only accounts behind a current deployment in both cases. Adds a real proxy and Postgres integration test for the fee API and for ChatGPT Responses rows on a deployment's own login, and names the card's multi-component test with the integration suffix. * fix(usage): keep a stored subscription fee listed after its account's deployments are removed The usage route lists every account behind a subscription deployment plus every account carrying a stored fee, and reads the fees on every proxy, so a fee stays visible, counted, and editable for past ranges after its last deployment goes. The integration test now polls every day the calls spanned, so a UTC midnight crossing no longer fails it. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + backend/routes/allowlist.py | 1 + .../migration.sql | 21 + .../litellm_proxy_extras/schema.prisma | 16 + litellm/constants.py | 1 + litellm/integrations/otel/model/spans.py | 1 + litellm/llms/chatgpt/authenticator.py | 8 +- litellm/llms/subscription_accounts.py | 37 ++ litellm/models/__init__.py | 2 + litellm/models/subscription_account.py | 28 ++ litellm/proxy/_types.py | 4 + .../common_daily_activity.py | 35 +- .../subscription_accounts.py | 335 ++++++++++++++ litellm/proxy/proxy_server.py | 4 + litellm/proxy/schema.prisma | 16 + .../subscription_billing_periods.py | 54 +++ .../subscription_account_repository.py | 134 ++++++ litellm/responses/main.py | 6 +- .../subscription_accounts.py | 62 +++ schema.prisma | 16 + .../spend/test_chatgpt_subscription_usage.py | 322 +++++++++++++ tests/unit/llms/test_subscription_accounts.py | 110 +++++ .../proxy/db/test_db_spend_update_writer.py | 31 ++ .../test_common_daily_activity.py | 65 +++ .../test_subscription_accounts.py | 438 ++++++++++++++++++ .../test_subscription_billing_periods.py | 73 +++ tests/unit/repositories/test_repositories.py | 43 ++ .../EntityUsage/EntityUsage.test.tsx | 28 ++ .../EntityUsage/SpendByProvider.test.tsx | 28 ++ .../EntityUsage/SpendByProvider.tsx | 18 +- .../EntityUsage/entityUsageAggregations.ts | 6 +- .../SubscriptionCostCard.test.tsx | 185 ++++++++ .../SubscriptionCost/SubscriptionCostCard.tsx | 221 +++++++++ .../SubscriptionCosts.integration.test.tsx | 69 +++ .../SubscriptionCost/SubscriptionCosts.tsx | 38 ++ .../SubscriptionFeeModal.test.tsx | 83 ++++ .../SubscriptionCost/SubscriptionFeeModal.tsx | 114 +++++ .../components/SubscriptionCost/formatting.ts | 34 ++ .../components/SubscriptionCost/types.ts | 74 +++ .../useSubscriptionFeeActions.test.tsx | 79 ++++ .../useSubscriptionFeeActions.ts | 90 ++++ .../components/UsagePageView.test.tsx | 99 +++- .../_components/components/UsagePageView.tsx | 24 +- .../components/overview/UsageOverview.tsx | 4 + .../components/overview/overviewData.test.ts | 29 ++ .../components/overview/overviewData.ts | 10 + .../_components/hooks/useSubscriptionUsage.ts | 47 ++ .../src/components/networking.tsx | 35 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 328 +++++++++++++ 49 files changed, 3485 insertions(+), 22 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261028000000_subscription_accounts/migration.sql create mode 100644 litellm/llms/subscription_accounts.py create mode 100644 litellm/models/subscription_account.py create mode 100644 litellm/proxy/management_endpoints/subscription_accounts.py create mode 100644 litellm/proxy/spend_tracking/subscription_billing_periods.py create mode 100644 litellm/repositories/subscription_account_repository.py create mode 100644 litellm/types/proxy/management_endpoints/subscription_accounts.py create mode 100644 tests/integration/spend/test_chatgpt_subscription_usage.py create mode 100644 tests/unit/llms/test_subscription_accounts.py create mode 100644 tests/unit/proxy/management_endpoints/test_subscription_accounts.py create mode 100644 tests/unit/proxy/spend_tracking/test_subscription_billing_periods.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCostCard.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCostCard.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCosts.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCosts.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionFeeModal.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionFeeModal.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/formatting.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/types.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/useSubscriptionFeeActions.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/useSubscriptionFeeActions.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useSubscriptionUsage.ts diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 9570b4d8cc8..68823d99519 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -227,6 +227,7 @@ jobs: tests/unit/llms/test_oss_decision.py tests/unit/llms/test_polling_url_origin_match.py tests/unit/llms/test_predibase_transformation.py + tests/unit/llms/test_subscription_accounts.py tests/unit/llms/a2a tests/unit/llms/aiml tests/unit/llms/aiohttp_openai diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 2a7522cf34e..2d3ccf40175 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -55,6 +55,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/cache_settings", "/coordination_redis/", "/cost_tracking", + "/subscription_accounts/", "/cost_optimization/", "/cost/", "/credentials", diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261028000000_subscription_accounts/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261028000000_subscription_accounts/migration.sql new file mode 100644 index 00000000000..e2788c85be9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261028000000_subscription_accounts/migration.sql @@ -0,0 +1,21 @@ +SET LOCAL lock_timeout = '10s'; + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_SubscriptionAccountTable" ( + "subscription_account_id" TEXT NOT NULL, + "custom_llm_provider" TEXT NOT NULL, + "account_id" TEXT NOT NULL, + "label" TEXT, + "monthly_fee" DOUBLE PRECISION NOT NULL, + "currency" TEXT NOT NULL, + "billing_period_start" TEXT NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_SubscriptionAccountTable_pkey" PRIMARY KEY ("subscription_account_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_SubscriptionAccountTable_provider_account_key" ON "LiteLLM_SubscriptionAccountTable"("custom_llm_provider", "account_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 2965174114b..179b741fe7c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -2131,3 +2131,19 @@ model LiteLLM_AccountingDelivery { @@id([event_id, target, version]) @@index([target, state, next_attempt_at]) } + +model LiteLLM_SubscriptionAccountTable { + subscription_account_id String @id @default(uuid()) + custom_llm_provider String + account_id String + label String? + monthly_fee Float + currency String + billing_period_start String + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String + + @@unique([custom_llm_provider, account_id], map: "LiteLLM_SubscriptionAccountTable_provider_account_key") +} diff --git a/litellm/constants.py b/litellm/constants.py index 8e80c07b4ed..97e93309ceb 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -899,6 +899,7 @@ PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset( } ) +SUBSCRIPTION_BACKED_PROVIDERS: Final = frozenset({"chatgpt"}) A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS: Final = frozenset({"databricks_agent"}) LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [ diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 4907bd4368d..320231430d2 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -361,6 +361,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_LensReview", "LiteLLM_LensIngestionKey", "LiteLLM_LensWorker", + "LiteLLM_SubscriptionAccountTable", "LiteLLM_LensSignalConfig", "LiteLLM_LensTraceSignal", ) diff --git a/litellm/llms/chatgpt/authenticator.py b/litellm/llms/chatgpt/authenticator.py index a5fb1c9c6b9..b237116c5f1 100644 --- a/litellm/llms/chatgpt/authenticator.py +++ b/litellm/llms/chatgpt/authenticator.py @@ -122,6 +122,12 @@ class Authenticator: return tokens["access_token"] def get_account_id(self, access_token: str | None = None) -> str | None: + return self._account_id(access_token, persist_derived=True) + + def read_account_id(self, access_token: str | None = None) -> str | None: + return self._account_id(access_token, persist_derived=False) + + def _account_id(self, access_token: str | None, persist_derived: bool) -> str | None: auth_data: Final = self._read_auth_file() if access_token and (not auth_data or auth_data.get("access_token") != access_token): return extract_chatgpt_account_id(access_token) @@ -133,7 +139,7 @@ class Authenticator: id_token: Final = auth_data.get("id_token") saved_access_token: Final = auth_data.get("access_token") derived: Final = extract_chatgpt_account_id(_optional_str(id_token or saved_access_token)) - if derived: + if derived and persist_derived: self._write_auth_file({**auth_data, "account_id": derived}) return derived diff --git a/litellm/llms/subscription_accounts.py b/litellm/llms/subscription_accounts.py new file mode 100644 index 00000000000..79c186ca124 --- /dev/null +++ b/litellm/llms/subscription_accounts.py @@ -0,0 +1,37 @@ +""" +Providers whose requests are covered by a flat subscription, each with the way to read +the account id of the login a deployment is signed into. +""" + +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final + +from litellm.llms.chatgpt.authenticator import Authenticator, get_cached_authenticator, get_chatgpt_auth_file +from litellm.secret_managers.main import UnresolvableSecretReference, resolve_local_oidc_reference +from litellm.utils import with_named_credential + + +def _chatgpt_account_id(litellm_params: Mapping[str, object]) -> str | None: + params: Final = with_named_credential(litellm_params) + auth_file: Final = get_chatgpt_auth_file(params) + if auth_file: + return get_cached_authenticator(auth_file).read_account_id() + api_key: Final = params.get("api_key") + if not isinstance(api_key, str) or not api_key: + return Authenticator().read_account_id() + match resolve_local_oidc_reference(api_key): + case UnresolvableSecretReference(): + return None + case str() as access_token: + return Authenticator().read_account_id(access_token) + + +SUBSCRIPTION_ACCOUNT_ID_RESOLVERS: Final[Mapping[str, Callable[[Mapping[str, object]], str | None]]] = MappingProxyType( + {"chatgpt": _chatgpt_account_id} +) + + +def resolve_subscription_account_id(custom_llm_provider: str, litellm_params: Mapping[str, object]) -> str | None: + resolver: Final = SUBSCRIPTION_ACCOUNT_ID_RESOLVERS.get(custom_llm_provider) + return None if resolver is None else resolver(litellm_params) diff --git a/litellm/models/__init__.py b/litellm/models/__init__.py index 50eb6f8af4f..bb47bd52c2d 100644 --- a/litellm/models/__init__.py +++ b/litellm/models/__init__.py @@ -30,6 +30,7 @@ from litellm.models.organization_membership import LiteLLM_OrganizationMembershi from litellm.models.project import LiteLLM_ProjectTable from litellm.models.skills import LiteLLM_SkillsTable from litellm.models.spend_logs import LiteLLM_ErrorLogs, LiteLLM_SpendLogs +from litellm.models.subscription_account import LiteLLM_SubscriptionAccountTable from litellm.models.tag import LiteLLM_TagTable from litellm.models.team import LiteLLM_TeamTable from litellm.models.team_membership import LiteLLM_TeamMembership @@ -59,6 +60,7 @@ __all__ = [ "LiteLLM_ProxyModelTable", "LiteLLM_SkillsTable", "LiteLLM_SpendLogs", + "LiteLLM_SubscriptionAccountTable", "LiteLLM_TagTable", "LiteLLM_TeamMemberTable", "LiteLLM_TeamMembership", diff --git a/litellm/models/subscription_account.py b/litellm/models/subscription_account.py new file mode 100644 index 00000000000..36eaa0ffc37 --- /dev/null +++ b/litellm/models/subscription_account.py @@ -0,0 +1,28 @@ +""" +Subscription account table model. + +Canonical definition for ``litellm_subscriptionaccounttable``: the operator-entered +recurring fee of a provider account whose requests are covered by a flat subscription. +""" + +from datetime import datetime + +from pydantic import ConfigDict + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_SubscriptionAccountTable(LiteLLMPydanticObjectBase): + model_config = ConfigDict(protected_namespaces=(), frozen=True) + + subscription_account_id: str + custom_llm_provider: str + account_id: str + label: str | None = None + monthly_fee: float + currency: str + billing_period_start: str + created_at: datetime + created_by: str + updated_at: datetime + updated_by: str diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0a09fe4a855..f165355da5a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -667,6 +667,10 @@ class LiteLLMRoutes(enum.Enum): "/team/spend/by_user", # gateway request counts (SGR); deployment-wide, admin-only "/gateway/daily/activity", + "/subscription_accounts/usage", + "/subscription_accounts/new", + "/subscription_accounts/update", + "/subscription_accounts/delete", # model "/model/new", "/model/update", diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 98372b64015..6e4ea0bd0cb 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -12,7 +12,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never from litellm import constants from litellm._logging import verbose_proxy_logger -from litellm.constants import PTU_SENTINEL_API_KEY +from litellm.constants import PTU_SENTINEL_API_KEY, SUBSCRIPTION_BACKED_PROVIDERS from litellm.proxy._types import CommonProxyErrors from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_details, @@ -295,6 +295,12 @@ def _entity_metadata( return stored if stored is not None else {} +def provider_bucket_metadata(provider: str, base: Mapping[str, object]) -> Mapping[str, object]: + if provider in SUBSCRIPTION_BACKED_PROVIDERS: + return {**base, "subscription_covered": True} + return dict(base) + + def update_breakdown_metrics( breakdown: BreakdownMetrics, record: DailySpendRecord, @@ -395,7 +401,7 @@ def update_breakdown_metrics( if provider not in breakdown.providers: breakdown.providers[provider] = MetricWithMetadata( metrics=SpendMetrics(), - metadata=provider_metadata.get(provider, {}), # Add any provider-specific metadata here + metadata=provider_bucket_metadata(provider, provider_metadata.get(provider, {})), ) breakdown.providers[provider].metrics = update_metrics(breakdown.providers[provider].metrics, record) @@ -795,12 +801,23 @@ def _aggregate_grouping_sets_records_sync( grouped_data[date_str] = bucket return bucket - def assign_metric_with_metadata(target: dict[str, MetricWithMetadata], key: str, metrics: SpendMetrics) -> None: + def assign_metric_with_metadata( + target: dict[str, MetricWithMetadata], + key: str, + metrics: SpendMetrics, + metadata: Mapping[str, object] | None = None, + ) -> None: existing: Final = target.get(key) - if existing is None: - target[key] = MetricWithMetadata(metrics=metrics, metadata={}) - else: - existing.metrics = sum_spend_metrics((existing.metrics, metrics)) + target[key] = ( + MetricWithMetadata(metrics=metrics, metadata=dict(metadata or {})) + if existing is None + else existing.model_copy( + update={ + "metrics": sum_spend_metrics((existing.metrics, metrics)), + "metadata": {**existing.metadata, **(metadata or {})}, + } + ) + ) def assign_api_key_breakdown( target: dict[str, MetricWithMetadata], @@ -864,7 +881,9 @@ def _aggregate_grouping_sets_records_sync( # provider the base build reported. provider_metrics = metrics.model_copy(update={"flat_cost": 0.0}) provider = record.custom_llm_provider or "unknown" - assign_metric_with_metadata(breakdown.providers, provider, provider_metrics) + assign_metric_with_metadata( + breakdown.providers, provider, provider_metrics, provider_bucket_metadata(provider, {}) + ) elif level == _GROUP_DATE_PROVIDER_API_KEY: if not is_ptu_sentinel: provider = record.custom_llm_provider or "unknown" diff --git a/litellm/proxy/management_endpoints/subscription_accounts.py b/litellm/proxy/management_endpoints/subscription_accounts.py new file mode 100644 index 00000000000..363d6cfba5d --- /dev/null +++ b/litellm/proxy/management_endpoints/subscription_accounts.py @@ -0,0 +1,335 @@ +""" +Subscription accounts: the operator-entered recurring fee of a provider account whose +requests are covered by a flat subscription, shown beside metered spend on the Usage page. +""" + +import re +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +from datetime import date +from itertools import chain +from typing import TYPE_CHECKING, Annotated, Final, TypeAlias + +from fastapi import APIRouter, Depends, HTTPException, Query, status +from pydantic import TypeAdapter + +from litellm.constants import LITELLM_PROXY_ADMIN_NAME, SUBSCRIPTION_BACKED_PROVIDERS +from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider +from litellm.llms.subscription_accounts import resolve_subscription_account_id +from litellm.models.subscription_account import LiteLLM_SubscriptionAccountTable +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth, user_api_key_has_admin_view +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.spend_tracking.subscription_billing_periods import billing_periods_starting_within +from litellm.proxy.utils import get_prisma_client_or_throw +from litellm.repositories.subscription_account_repository import ( + DuplicateSubscriptionAccount, + KeepLabel, + SetLabel, + SubscriptionAccountFeeChanges, + SubscriptionAccountFeeFields, + SubscriptionAccountRepository, + SubscriptionAccountStore, +) +from litellm.types.proxy.management_endpoints.subscription_accounts import ( + BillingPeriod, + DeleteSubscriptionAccountRequest, + SubscriptionAccountFeeRequest, + SubscriptionAccountUsage, + SubscriptionUsageResponse, + UpdateSubscriptionAccountFeeRequest, +) +from litellm.types.router import DeploymentTypedDict + +if TYPE_CHECKING: + from litellm.router import Router + +router: Final = APIRouter(prefix="/subscription_accounts", tags=["subscription accounts"]) +Auth: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)] + + +@dataclass(frozen=True, slots=True) +class SubscriptionDeployment: + model_name: str + custom_llm_provider: str + account_id: str | None + + +AccountIdResolver: TypeAlias = Callable[[str, Mapping[str, object]], str | None] +_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _declared_provider(litellm_params: Mapping[str, object]) -> str | None: + model: Final = litellm_params.get("model") + configured_provider: Final = litellm_params.get("custom_llm_provider") + return declared_authenticating_provider( + model if isinstance(model, str) else None, + configured_provider if isinstance(configured_provider, str) else None, + ) + + +def _model_info(deployment: Mapping[str, object]) -> Mapping[str, object]: + return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) + + +def _public_model_name(entry: DeploymentTypedDict) -> str: + team_public_model_name: Final = _model_info(entry).get("team_public_model_name") + return team_public_model_name if isinstance(team_public_model_name, str) else entry["model_name"] + + +def _subscription_deployment( + entry: DeploymentTypedDict, account_id_resolver: AccountIdResolver +) -> SubscriptionDeployment | None: + litellm_params: Final = entry["litellm_params"] + provider: Final = _declared_provider(litellm_params) + if provider is None or provider not in SUBSCRIPTION_BACKED_PROVIDERS: + return None + return SubscriptionDeployment( + model_name=_public_model_name(entry), + custom_llm_provider=provider, + account_id=account_id_resolver(provider, litellm_params), + ) + + +def subscription_deployments_from_model_list( + model_list: Iterable[DeploymentTypedDict], account_id_resolver: AccountIdResolver +) -> tuple[SubscriptionDeployment, ...]: + candidates: Final = (_subscription_deployment(entry, account_id_resolver) for entry in model_list) + return tuple(candidate for candidate in candidates if candidate is not None) + + +def subscription_deployments_of_router( + llm_router: "Router | None", account_id_resolver: AccountIdResolver +) -> tuple[SubscriptionDeployment, ...]: + if llm_router is None: + return () + registry: Final = llm_router.registry() + indexed_names: Final = tuple(registry.index("model_name")) + deployments: Final = chain.from_iterable(registry.deployments(name, source="named") or () for name in indexed_names) + return subscription_deployments_from_model_list(deployments, account_id_resolver) + + +def get_subscription_deployments() -> tuple[SubscriptionDeployment, ...]: + from litellm.proxy.proxy_server import llm_router + + return subscription_deployments_of_router(llm_router, resolve_subscription_account_id) + + +def get_subscription_account_repository() -> SubscriptionAccountStore: + return SubscriptionAccountRepository(get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value)) + + +Store: TypeAlias = Annotated[SubscriptionAccountStore, Depends(get_subscription_account_repository)] +Deployments: TypeAlias = Annotated[tuple[SubscriptionDeployment, ...], Depends(get_subscription_deployments)] + + +def _forbidden(user_api_key_dict: UserAPIKeyAuth) -> HTTPException: + return HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": f"{CommonProxyErrors.not_allowed_access.value}, your role={user_api_key_dict.user_role}"}, + ) + + +def _account_usage( + custom_llm_provider: str, + account_id: str | None, + deployments: Iterable[SubscriptionDeployment], + fee_row: LiteLLM_SubscriptionAccountTable | None, + start_date: date, + end_date: date, +) -> SubscriptionAccountUsage: + model_names: Final = tuple(sorted({deployment.model_name for deployment in deployments})) + if fee_row is None: + return SubscriptionAccountUsage( + custom_llm_provider=custom_llm_provider, + account_id=account_id, + deployments=model_names, + fee=None, + billing_periods=(), + fixed_cost=None, + ) + periods: Final = billing_periods_starting_within( + date.fromisoformat(fee_row.billing_period_start), start_date, end_date + ) + return SubscriptionAccountUsage( + custom_llm_provider=custom_llm_provider, + account_id=account_id, + deployments=model_names, + fee=fee_row, + billing_periods=tuple(BillingPeriod(start=period.start, end=period.end) for period in periods), + fixed_cost=len(periods) * fee_row.monthly_fee, + ) + + +def build_subscription_usage( + deployments: Iterable[SubscriptionDeployment], + fee_rows: Iterable[LiteLLM_SubscriptionAccountTable], + start_date: date, + end_date: date, +) -> SubscriptionUsageResponse: + deployment_tuple: Final = tuple(deployments) + fees_by_account: Final = {(row.custom_llm_provider, row.account_id): row for row in fee_rows} + deployment_keys: Final = {(d.custom_llm_provider, d.account_id) for d in deployment_tuple} + account_keys: Final = sorted(deployment_keys | set(fees_by_account), key=lambda key: (key[0], key[1] or "~")) + accounts: Final = tuple( + _account_usage( + provider, + account_id, + (d for d in deployment_tuple if (d.custom_llm_provider, d.account_id) == (provider, account_id)), + fees_by_account.get((provider, account_id)) if account_id is not None else None, + start_date, + end_date, + ) + for provider, account_id in account_keys + ) + return SubscriptionUsageResponse( + start_date=start_date, + end_date=end_date, + subscription_providers=tuple(sorted(SUBSCRIPTION_BACKED_PROVIDERS)), + accounts=accounts, + ) + + +_DAY_PATTERN: Final = re.compile(r"\d{4}-\d{2}-\d{2}") + + +def _parse_day(value: str) -> date | None: + if _DAY_PATTERN.fullmatch(value) is None: + return None + try: + return date.fromisoformat(value) + except ValueError: + return None + + +def _parse_date(value: str, name: str) -> date: + parsed: Final = _parse_day(value) + if parsed is None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail={"error": f"{name} must be YYYY-MM-DD, got {value!r}"} + ) + return parsed + + +@router.get("/usage", response_model=SubscriptionUsageResponse, dependencies=[Depends(user_api_key_auth)]) +async def get_subscription_usage( + start_date: Annotated[str, Query(description="First day of the Usage page filter, YYYY-MM-DD")], + end_date: Annotated[str, Query(description="Last day of the Usage page filter, YYYY-MM-DD")], + user_api_key_dict: Auth, + repository: Store, + deployments: Deployments, +) -> SubscriptionUsageResponse: + """ + Subscription-backed accounts behind the proxy's deployments or carrying a stored fee, each with + its configured monthly fee and the fixed cost attributed to the date range: the fee once per billing + period whose start date falls inside the range, however many deployments share the account. + """ + if not user_api_key_has_admin_view(user_api_key_dict): + raise _forbidden(user_api_key_dict) + start: Final = _parse_date(start_date, "start_date") + end: Final = _parse_date(end_date, "end_date") + if end < start: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "end_date must not be before start_date"} + ) + return build_subscription_usage(deployments, await repository.list_all(), start, end) + + +def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise _forbidden(user_api_key_dict) + + +def _actor(user_api_key_dict: UserAPIKeyAuth) -> str: + return user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME + + +@router.post("/new", response_model=LiteLLM_SubscriptionAccountTable, dependencies=[Depends(user_api_key_auth)]) +async def new_subscription_account( + request: SubscriptionAccountFeeRequest, + user_api_key_dict: Auth, + repository: Store, +) -> LiteLLM_SubscriptionAccountTable: + """Record the monthly fee of one subscription-backed provider account.""" + _require_proxy_admin(user_api_key_dict) + if request.custom_llm_provider not in SUBSCRIPTION_BACKED_PROVIDERS: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": f"{request.custom_llm_provider!r} is not a subscription-backed provider; " + f"expected one of {sorted(SUBSCRIPTION_BACKED_PROVIDERS)}" + }, + ) + created: Final = await repository.create_account( + request.custom_llm_provider, + request.account_id, + SubscriptionAccountFeeFields( + monthly_fee=request.monthly_fee, + currency=request.currency, + billing_period_start=request.billing_period_start.isoformat(), + label=request.label, + ), + _actor(user_api_key_dict), + ) + if isinstance(created, DuplicateSubscriptionAccount): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": f"A fee is already configured for {created.custom_llm_provider} account " + f"{created.account_id}; use /subscription_accounts/update to change it" + }, + ) + return created + + +def _fee_changes(request: UpdateSubscriptionAccountFeeRequest) -> SubscriptionAccountFeeChanges: + return SubscriptionAccountFeeChanges( + monthly_fee=request.monthly_fee, + currency=request.currency, + billing_period_start=( + request.billing_period_start.isoformat() if request.billing_period_start is not None else None + ), + label=SetLabel(request.label) if "label" in request.model_fields_set else KeepLabel(), + ) + + +@router.post("/update", response_model=LiteLLM_SubscriptionAccountTable, dependencies=[Depends(user_api_key_auth)]) +async def update_subscription_account( + request: UpdateSubscriptionAccountFeeRequest, + user_api_key_dict: Auth, + repository: Store, +) -> LiteLLM_SubscriptionAccountTable: + """ + Change the fee, currency, billing period start, or label of a subscription account. + Only the fields named in the request are written, so two admins changing different + fields at the same time both land. + """ + _require_proxy_admin(user_api_key_dict) + updated: Final = await repository.update_fee( + request.subscription_account_id, _fee_changes(request), _actor(user_api_key_dict) + ) + if updated is None: + raise _not_found(request.subscription_account_id) + return updated + + +@router.post("/delete", response_model=LiteLLM_SubscriptionAccountTable, dependencies=[Depends(user_api_key_auth)]) +async def delete_subscription_account( + request: DeleteSubscriptionAccountRequest, + user_api_key_dict: Auth, + repository: Store, +) -> LiteLLM_SubscriptionAccountTable: + """Stop tracking a subscription account's fee; its requests stay marked as subscription-covered.""" + _require_proxy_admin(user_api_key_dict) + if await repository.find_by_id(request.subscription_account_id) is None: + raise _not_found(request.subscription_account_id) + deleted: Final = await repository.delete_account(request.subscription_account_id) + if deleted is None: + raise _not_found(request.subscription_account_id) + return deleted + + +def _not_found(subscription_account_id: str) -> HTTPException: + return HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"No subscription account with id {subscription_account_id}"}, + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3e004f1ced7..ae0e7fce667 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -777,6 +777,9 @@ from litellm.proxy.management_endpoints.session_endpoints import ( router as session_management_router, ) from litellm.proxy.management_endpoints.sso.saml_sso import SAML_AUTHN_REQUEST_TTL_SECONDS +from litellm.proxy.management_endpoints.subscription_accounts import ( + router as subscription_accounts_router, +) from litellm.proxy.management_endpoints.tag_management_endpoints import ( router as tag_management_router, ) @@ -21769,6 +21772,7 @@ app.include_router(memory_router) app.include_router(lens_router) app.include_router(plugin_router) app.include_router(cost_tracking_settings_router) +app.include_router(subscription_accounts_router) app.include_router(prompt_caching_requests_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2965174114b..179b741fe7c 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -2131,3 +2131,19 @@ model LiteLLM_AccountingDelivery { @@id([event_id, target, version]) @@index([target, state, next_attempt_at]) } + +model LiteLLM_SubscriptionAccountTable { + subscription_account_id String @id @default(uuid()) + custom_llm_provider String + account_id String + label String? + monthly_fee Float + currency String + billing_period_start String + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String + + @@unique([custom_llm_provider, account_id], map: "LiteLLM_SubscriptionAccountTable_provider_account_key") +} diff --git a/litellm/proxy/spend_tracking/subscription_billing_periods.py b/litellm/proxy/spend_tracking/subscription_billing_periods.py new file mode 100644 index 00000000000..a7282cc942c --- /dev/null +++ b/litellm/proxy/spend_tracking/subscription_billing_periods.py @@ -0,0 +1,54 @@ +from calendar import monthrange +from dataclasses import dataclass +from datetime import date, timedelta +from typing import Final + + +@dataclass(frozen=True, slots=True) +class BillingPeriodRange: + start: date + end: date + + +def _month_index(anchor: date, months_after_anchor: int) -> int: + return anchor.year * 12 + (anchor.month - 1) + months_after_anchor + + +_LAST_MONTH_INDEX: Final = _month_index(date.max, 0) + + +def _month_start(anchor: date, months_after_anchor: int) -> date: + month_index: Final = _month_index(anchor, months_after_anchor) + year: Final = month_index // 12 + month: Final = month_index % 12 + 1 + return date(year, month, min(anchor.day, monthrange(year, month)[1])) + + +def _period_end(anchor: date, months_after_anchor: int) -> date: + if _month_index(anchor, months_after_anchor + 1) > _LAST_MONTH_INDEX: + return date.max + return _month_start(anchor, months_after_anchor + 1) - timedelta(days=1) + + +def _first_candidate_index(anchor: date, start_date: date) -> int: + months_apart: Final = (start_date.year - anchor.year) * 12 + (start_date.month - anchor.month) + return max(0, months_apart - 1) + + +def billing_periods_starting_within(anchor: date, start_date: date, end_date: date) -> tuple[BillingPeriodRange, ...]: + """Monthly periods anchored on `anchor`'s day of month whose start lies inside [start_date, end_date].""" + if end_date < start_date or end_date < anchor: + return () + first_index: Final = _first_candidate_index(anchor, start_date) + last_index: Final = min( + first_index + _month_span(start_date, end_date) + 1, _LAST_MONTH_INDEX - _month_index(anchor, 0) + ) + candidates: Final = ( + BillingPeriodRange(start=_month_start(anchor, index), end=_period_end(anchor, index)) + for index in range(first_index, last_index + 1) + ) + return tuple(period for period in candidates if start_date <= period.start <= end_date) + + +def _month_span(start_date: date, end_date: date) -> int: + return (end_date.year - start_date.year) * 12 + (end_date.month - start_date.month) diff --git a/litellm/repositories/subscription_account_repository.py b/litellm/repositories/subscription_account_repository.py new file mode 100644 index 00000000000..ddf441a9191 --- /dev/null +++ b/litellm/repositories/subscription_account_repository.py @@ -0,0 +1,134 @@ +""" +Subscription account repository for database operations on LiteLLM_SubscriptionAccountTable. +""" + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias + +from litellm.models.subscription_account import LiteLLM_SubscriptionAccountTable +from litellm.repositories.base_repository import BaseRepository, is_unique_violation +from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository + +if TYPE_CHECKING: + from prisma import models as prisma_models + + +@dataclass(frozen=True, slots=True) +class DuplicateSubscriptionAccount: + custom_llm_provider: str + account_id: str + + +@dataclass(frozen=True, slots=True) +class SubscriptionAccountFeeFields: + monthly_fee: float + currency: str + billing_period_start: str + label: str | None + + +@dataclass(frozen=True, slots=True) +class KeepLabel: + pass + + +@dataclass(frozen=True, slots=True) +class SetLabel: + value: str | None + + +LabelChange: TypeAlias = KeepLabel | SetLabel + + +@dataclass(frozen=True, slots=True) +class SubscriptionAccountFeeChanges: + monthly_fee: float | None = None + currency: str | None = None + billing_period_start: str | None = None + label: LabelChange = KeepLabel() + + def columns(self) -> Mapping[str, float | str | None]: + named: Final = { + "monthly_fee": self.monthly_fee, + "currency": self.currency, + "billing_period_start": self.billing_period_start, + } + present: Final = {column: value for column, value in named.items() if value is not None} + label_column: Final = {"label": self.label.value} if isinstance(self.label, SetLabel) else {} + return MappingProxyType({**present, **label_column}) + + +class SubscriptionAccountStore(Protocol): + async def list_all(self) -> tuple[LiteLLM_SubscriptionAccountTable, ...]: ... + + async def find_by_id(self, id_value: str) -> LiteLLM_SubscriptionAccountTable | None: ... + + async def create_account( + self, custom_llm_provider: str, account_id: str, fee: SubscriptionAccountFeeFields, created_by: str + ) -> LiteLLM_SubscriptionAccountTable | DuplicateSubscriptionAccount: ... + + async def update_fee( + self, subscription_account_id: str, changes: SubscriptionAccountFeeChanges, updated_by: str + ) -> LiteLLM_SubscriptionAccountTable | None: ... + + async def delete_account(self, subscription_account_id: str) -> LiteLLM_SubscriptionAccountTable | None: ... + + +class _SubscriptionAccountTable(PrismaTableRepository["prisma_models.LiteLLM_SubscriptionAccountTable"]): + table_name = "litellm_subscriptionaccounttable" + + +class SubscriptionAccountRepository(BaseRepository[LiteLLM_SubscriptionAccountTable]): + def __init__(self, prisma_client: object) -> None: + super().__init__(prisma_client) + self._table: Final = _SubscriptionAccountTable(prisma_client) + + @property + def table(self) -> TableActions["prisma_models.LiteLLM_SubscriptionAccountTable"]: + return self._table.table + + @property + def model_class(self) -> type[LiteLLM_SubscriptionAccountTable]: + return LiteLLM_SubscriptionAccountTable + + async def list_all(self) -> tuple[LiteLLM_SubscriptionAccountTable, ...]: + return tuple(await self.find_many(order={"created_at": "asc"})) + + async def find_by_id( + self, id_value: str, id_field: str = "subscription_account_id" + ) -> LiteLLM_SubscriptionAccountTable | None: + return await super().find_by_id(id_value, id_field) + + async def create_account( + self, custom_llm_provider: str, account_id: str, fee: SubscriptionAccountFeeFields, created_by: str + ) -> LiteLLM_SubscriptionAccountTable | DuplicateSubscriptionAccount: + try: + return await self.create( + { + "custom_llm_provider": custom_llm_provider, + "account_id": account_id, + "monthly_fee": fee.monthly_fee, + "currency": fee.currency, + "billing_period_start": fee.billing_period_start, + "label": fee.label, + "created_by": created_by, + "updated_by": created_by, + } + ) + except Exception as exc: + if is_unique_violation(exc): + return DuplicateSubscriptionAccount(custom_llm_provider=custom_llm_provider, account_id=account_id) + raise + + async def update_fee( + self, subscription_account_id: str, changes: SubscriptionAccountFeeChanges, updated_by: str + ) -> LiteLLM_SubscriptionAccountTable | None: + return await self.update( + subscription_account_id, {**changes.columns(), "updated_by": updated_by}, id_field="subscription_account_id" + ) + + async def delete_account(self, subscription_account_id: str) -> LiteLLM_SubscriptionAccountTable | None: + return await self.delete(subscription_account_id, id_field="subscription_account_id") diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 063ef9d91ea..5c8350589b8 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -31,7 +31,10 @@ from litellm.integrations.anthropic_cache_control_hook import ( from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params -from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header +from litellm.litellm_core_utils.get_llm_provider_logic import ( + declared_authenticating_provider, + get_provider_affinity_header, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -805,6 +808,7 @@ def provider_stores_the_conversation(model: str, custom_llm_provider: str | None return not _will_bridge_to_chat_completions( model, custom_llm_provider + or declared_authenticating_provider(model) or litellm.get_llm_provider(model=model, api_base=api_base, litellm_params=litellm_params)[1], bool(kwargs.get("use_chat_completions_api")), kwargs.get("model_info"), diff --git a/litellm/types/proxy/management_endpoints/subscription_accounts.py b/litellm/types/proxy/management_endpoints/subscription_accounts.py new file mode 100644 index 00000000000..a924b97034b --- /dev/null +++ b/litellm/types/proxy/management_endpoints/subscription_accounts.py @@ -0,0 +1,62 @@ +from datetime import date +from typing import Final + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.models.subscription_account import LiteLLM_SubscriptionAccountTable + +CURRENCY_CODE_PATTERN: Final = r"^[A-Z]{3}$" + + +class SubscriptionAccountFeeRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + custom_llm_provider: str = Field(min_length=1) + account_id: str = Field(min_length=1) + monthly_fee: float = Field(gt=0, allow_inf_nan=False) + currency: str = Field(pattern=CURRENCY_CODE_PATTERN) + billing_period_start: date + label: str | None = None + + +class UpdateSubscriptionAccountFeeRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + subscription_account_id: str + monthly_fee: float | None = Field(default=None, gt=0, allow_inf_nan=False) + currency: str | None = Field(default=None, pattern=CURRENCY_CODE_PATTERN) + billing_period_start: date | None = None + label: str | None = None + + +class DeleteSubscriptionAccountRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + subscription_account_id: str + + +class BillingPeriod(BaseModel): + model_config = ConfigDict(frozen=True) + + start: date + end: date + + +class SubscriptionAccountUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + custom_llm_provider: str + account_id: str | None + deployments: tuple[str, ...] + fee: LiteLLM_SubscriptionAccountTable | None + billing_periods: tuple[BillingPeriod, ...] + fixed_cost: float | None + + +class SubscriptionUsageResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + start_date: date + end_date: date + subscription_providers: tuple[str, ...] + accounts: tuple[SubscriptionAccountUsage, ...] diff --git a/schema.prisma b/schema.prisma index 2965174114b..179b741fe7c 100644 --- a/schema.prisma +++ b/schema.prisma @@ -2131,3 +2131,19 @@ model LiteLLM_AccountingDelivery { @@id([event_id, target, version]) @@index([target, state, next_attempt_at]) } + +model LiteLLM_SubscriptionAccountTable { + subscription_account_id String @id @default(uuid()) + custom_llm_provider String + account_id String + label String? + monthly_fee Float + currency String + billing_period_start String + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String + + @@unique([custom_llm_provider, account_id], map: "LiteLLM_SubscriptionAccountTable_provider_account_key") +} diff --git a/tests/integration/spend/test_chatgpt_subscription_usage.py b/tests/integration/spend/test_chatgpt_subscription_usage.py new file mode 100644 index 00000000000..8f1fb2d0f74 --- /dev/null +++ b/tests/integration/spend/test_chatgpt_subscription_usage.py @@ -0,0 +1,322 @@ +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Final + +import httpx +import pytest +from integration._support import codex_vendor as cv +from integration._support import responses_vendor as rv +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, string_value +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +_MODEL: Final = "chatgpt/gpt-5.5" +_PROVIDER: Final = "chatgpt" +_FEE_TABLE_QUERY: Final = ( + 'SELECT subscription_account_id FROM "LiteLLM_SubscriptionAccountTable" WHERE custom_llm_provider = %s' + " AND account_id = %s" +) +_FEE_ROW_QUERY: Final = ( + 'SELECT subscription_account_id FROM "LiteLLM_SubscriptionAccountTable" WHERE subscription_account_id = %s' +) +_MONTHLY_FEE: Final = 200.0 +_BILLING_PERIOD_START: Final = "2026-01-15" +_RANGE: Final = {"start_date": "2026-03-01", "end_date": "2026-03-31"} +_PERIOD_IN_RANGE: Final = ("2026-03-15", "2026-04-14") +_TOKENS: Final = TypeAdapter(int) + + +class _Metrics(BaseModel): + model_config = ConfigDict(frozen=True) + + spend: float + prompt_tokens: int + completion_tokens: int + api_requests: int + successful_requests: int + failed_requests: int + + +class _ProviderUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + metrics: _Metrics + metadata: Mapping[str, JsonValue] + + +class _Breakdown(BaseModel): + model_config = ConfigDict(frozen=True) + + providers: Mapping[str, _ProviderUsage] + + +class _Day(BaseModel): + model_config = ConfigDict(frozen=True) + + breakdown: _Breakdown + + +class _DailyActivity(BaseModel): + model_config = ConfigDict(frozen=True) + + results: tuple[_Day, ...] + + +class _Fee(BaseModel): + model_config = ConfigDict(frozen=True) + + subscription_account_id: str + custom_llm_provider: str + account_id: str + monthly_fee: float + currency: str + billing_period_start: str + + +class _Period(BaseModel): + model_config = ConfigDict(frozen=True) + + start: str + end: str + + +class _Account(BaseModel): + model_config = ConfigDict(frozen=True) + + custom_llm_provider: str + account_id: str | None + deployments: tuple[str, ...] + fee: _Fee | None + billing_periods: tuple[_Period, ...] + fixed_cost: float | None + + +class _SubscriptionUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + accounts: tuple[_Account, ...] + + +@dataclass(frozen=True, slots=True) +class _Rig: + wire: Wire + proxy: OwnedProxy + model: str + + @property + def gateway(self) -> Gateway: + return self.proxy.gateway + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("chatgpt-subscription-usage") + no_default_login: Final = directory / "no-default-login" + no_default_login.mkdir() + deployment_login: Final = cv.login(directory) / "auth.json" + with gateway_from_environment() as gateway, wire_server(cv.CodexVendor().respond) as wire: + overrides: Final = {"CHATGPT_TOKEN_DIR": str(no_default_login), "CHATGPT_API_BASE": wire.url} + config: Final = cv.proxy_config(directory, probe=False) + with ( + owned_proxy_process(gateway, directory, overrides, config=config, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model=_MODEL, api_base=wire.url, api_key=None, chatgpt_auth_file=str(deployment_login) + ) + yield _Rig(wire, owned, model) + + +def _responses_body(model: str, marker: str, *, stream: bool) -> Mapping[str, JsonValue]: + return { + "model": model, + "input": [{"role": "user", "content": f"Say marker-{marker}"}], + "stream": stream, + "cache": {"no-cache": True}, + } + + +def _assert_answered(response: httpx.Response, marker: str, content_type: str) -> None: + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith(content_type), response.headers + assert rv.answer(marker) in response.text, response.text + + +def _today() -> str: + return datetime.now(UTC).date().isoformat() + + +def _chatgpt_usage_between(gateway: Gateway, first_day: str, last_day: str) -> tuple[_ProviderUsage, ...]: + response: Final = gateway.request( + "GET", "/user/daily/activity", params={"start_date": first_day, "end_date": last_day} + ) + assert response.status_code == 200, response.text + activity: Final = _DailyActivity.model_validate_json(response.content) + return tuple( + result.breakdown.providers[_PROVIDER] for result in activity.results if _PROVIDER in result.breakdown.providers + ) + + +def _summed(buckets: tuple[_ProviderUsage, ...]) -> _Metrics: + return _Metrics( + spend=sum(bucket.metrics.spend for bucket in buckets), + prompt_tokens=sum(bucket.metrics.prompt_tokens for bucket in buckets), + completion_tokens=sum(bucket.metrics.completion_tokens for bucket in buckets), + api_requests=sum(bucket.metrics.api_requests for bucket in buckets), + successful_requests=sum(bucket.metrics.successful_requests for bucket in buckets), + failed_requests=sum(bucket.metrics.failed_requests for bucket in buckets), + ) + + +def _usage(gateway: Gateway, *, key: str | None = None) -> httpx.Response: + return gateway.request("GET", "/subscription_accounts/usage", params=_RANGE, key=key) + + +def _accounts(gateway: Gateway) -> tuple[_Account, ...]: + response: Final = _usage(gateway) + assert response.status_code == 200, response.text + return _SubscriptionUsage.model_validate_json(response.content).accounts + + +def _deployment_account(rig: _Rig) -> _Account: + (account,) = _accounts(rig.gateway) + assert (account.custom_llm_provider, account.account_id, account.deployments) == ( + _PROVIDER, + cv.ACCOUNT, + (rig.model,), + ), account + return account + + +def _assert_no_fee(account: _Account) -> None: + assert (account.fee, account.billing_periods, account.fixed_cost) == (None, (), None), account + + +def _fee_body(account_id: str, monthly_fee: float = _MONTHLY_FEE) -> Mapping[str, JsonValue]: + return { + "custom_llm_provider": _PROVIDER, + "account_id": account_id, + "monthly_fee": monthly_fee, + "currency": "USD", + "billing_period_start": _BILLING_PERIOD_START, + "label": "integration", + } + + +def _delete_fee_if_present(gateway: Gateway, identity: str) -> None: + if gateway.rows(_FEE_ROW_QUERY, (identity,)): + gateway.post("/subscription_accounts/delete", {"subscription_account_id": identity}) + assert gateway.rows(_FEE_ROW_QUERY, (identity,)) == [] + + +def _set_fee(scenario: Scenario, account_id: str) -> str: + created: Final = scenario.gateway.post("/subscription_accounts/new", _fee_body(account_id)) + identity: Final = string_value(created["subscription_account_id"]) + scenario.cleanups.callback(_delete_fee_if_present, scenario.gateway, identity) + return identity + + +def _assert_fee_reported(account: _Account, identity: str) -> None: + assert account.fee is not None, account + assert ( + account.fee.subscription_account_id, + account.fee.account_id, + account.fee.monthly_fee, + account.fee.currency, + account.fee.billing_period_start, + ) == (identity, cv.ACCOUNT, _MONTHLY_FEE, "USD", _BILLING_PERIOD_START), account + assert tuple((period.start, period.end) for period in account.billing_periods) == (_PERIOD_IN_RANGE,), account + assert account.fixed_cost == _MONTHLY_FEE, account + + +def test_responses_calls_on_a_deployment_login_count_as_subscription_covered_chatgpt_usage(rig: _Rig) -> None: + first_day: Final = _today() + rig.wire.drain() + plain_marker: Final = uuid.uuid4().hex + stream_marker: Final = uuid.uuid4().hex + plain: Final = rig.gateway.request("POST", "/v1/responses", _responses_body(rig.model, plain_marker, stream=False)) + _assert_answered(plain, plain_marker, "application/json") + streamed: Final = rig.gateway.request( + "POST", "/v1/responses", _responses_body(rig.model, stream_marker, stream=True) + ) + _assert_answered(streamed, stream_marker, "text/event-stream") + forwarded: Final = rig.wire.drain() + assert len(forwarded) == 2, [request.target for request in forwarded] + for request, marker in zip(forwarded, (plain_marker, stream_marker), strict=True): + cv.forwarded(request, marker) + last_day: Final = _today() + buckets: Final = eventually( + lambda: _chatgpt_usage_between(rig.gateway, first_day, last_day), + lambda found: _summed(found).api_requests >= 2, + seconds=70, + ) + totals: Final = _summed(buckets) + assert (totals.api_requests, totals.successful_requests, totals.failed_requests, totals.spend) == ( + 2, + 2, + 0, + 0, + ), buckets + assert (totals.prompt_tokens, totals.completion_tokens) == ( + 2 * _TOKENS.validate_python(cv.USAGE["input_tokens"]), + 2 * _TOKENS.validate_python(cv.USAGE["output_tokens"]), + ), buckets + assert all(bucket.metadata.get("subscription_covered") is True for bucket in buckets), buckets + + +def test_fee_on_the_deployment_account_reports_its_fixed_cost_until_deleted(rig: _Rig) -> None: + _assert_no_fee(_deployment_account(rig)) + with rig.gateway.scenario() as scenario: + identity: Final = _set_fee(scenario, cv.ACCOUNT) + _assert_fee_reported(_deployment_account(rig), identity) + deleted: Final = rig.gateway.post("/subscription_accounts/delete", {"subscription_account_id": identity}) + assert deleted["subscription_account_id"] == identity, deleted + assert rig.gateway.rows(_FEE_ROW_QUERY, (identity,)) == [] + _assert_no_fee(_deployment_account(rig)) + + +def test_second_fee_for_the_same_account_answers_400_and_keeps_the_first(rig: _Rig) -> None: + with rig.gateway.scenario() as scenario: + identity: Final = _set_fee(scenario, cv.ACCOUNT) + duplicate: Final = rig.gateway.request( + "POST", "/subscription_accounts/new", _fee_body(cv.ACCOUNT, monthly_fee=3 * _MONTHLY_FEE) + ) + assert duplicate.status_code == 400, duplicate.text + assert rig.gateway.rows(_FEE_TABLE_QUERY, (_PROVIDER, cv.ACCOUNT)) == [{"subscription_account_id": identity}] + _assert_fee_reported(_deployment_account(rig), identity) + + +def test_update_and_delete_of_an_unknown_fee_answer_404(rig: _Rig) -> None: + unknown: Final = str(uuid.uuid4()) + updated: Final = rig.gateway.request( + "POST", "/subscription_accounts/update", {"subscription_account_id": unknown, "monthly_fee": _MONTHLY_FEE} + ) + assert updated.status_code == 404, updated.text + deleted: Final = rig.gateway.request("POST", "/subscription_accounts/delete", {"subscription_account_id": unknown}) + assert deleted.status_code == 404, deleted.text + assert rig.gateway.rows(_FEE_ROW_QUERY, (unknown,)) == [] + + +def test_internal_user_is_refused_setting_a_fee_and_reading_usage(rig: _Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(user_id=scenario.user(user_role="internal_user")) + created: Final = rig.gateway.request("POST", "/subscription_accounts/new", _fee_body(cv.ACCOUNT), key=key) + assert created.status_code == 401, created.text + assert rig.gateway.rows(_FEE_TABLE_QUERY, (_PROVIDER, cv.ACCOUNT)) == [] + usage: Final = _usage(rig.gateway, key=key) + assert usage.status_code == 401, usage.text + + +def test_fee_on_an_account_no_deployment_uses_stays_listed_and_counted(rig: _Rig) -> None: + unused: Final = f"acct-unused-{uuid.uuid4().hex}" + with rig.gateway.scenario() as scenario: + identity: Final = _set_fee(scenario, unused) + accounts: Final = _accounts(rig.gateway) + assert frozenset(account.account_id for account in accounts) == {cv.ACCOUNT, unused}, accounts + (unused_account,) = (account for account in accounts if account.account_id == unused) + assert unused_account.fee is not None, unused_account + assert (unused_account.deployments, unused_account.fee.subscription_account_id) == ((), identity) + assert unused_account.fixed_cost == _MONTHLY_FEE, unused_account diff --git a/tests/unit/llms/test_subscription_accounts.py b/tests/unit/llms/test_subscription_accounts.py new file mode 100644 index 00000000000..6995ee5d4df --- /dev/null +++ b/tests/unit/llms/test_subscription_accounts.py @@ -0,0 +1,110 @@ +import base64 +import json +from pathlib import Path +from typing import Final + +import pytest + +import litellm +from litellm.constants import SUBSCRIPTION_BACKED_PROVIDERS +from litellm.llms.chatgpt.common_utils import OPENAI_AUTH_CLAIM_KEY +from litellm.llms.subscription_accounts import SUBSCRIPTION_ACCOUNT_ID_RESOLVERS, resolve_subscription_account_id +from litellm.types.utils import CredentialItem + + +def _write_login(path: Path, account_id: str) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({"account_id": account_id})) + return path + + +def _access_token_for(account_id: str) -> str: + claims: Final = json.dumps({OPENAI_AUTH_CLAIM_KEY: {"chatgpt_account_id": account_id}}).encode() + payload: Final = base64.urlsafe_b64encode(claims).decode().rstrip("=") + return f"header.{payload}.signature" + + +@pytest.fixture +def signed_in_chatgpt_token_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + _write_login(tmp_path / "auth.json", "acct-from-token-file") + return tmp_path + + +def test_every_subscription_backed_provider_has_an_account_id_resolver() -> None: + assert frozenset(SUBSCRIPTION_ACCOUNT_ID_RESOLVERS) == SUBSCRIPTION_BACKED_PROVIDERS + + +def test_chatgpt_resolves_to_the_account_in_the_signed_in_token_file(signed_in_chatgpt_token_dir: Path) -> None: + assert resolve_subscription_account_id("chatgpt", {"model": "chatgpt/gpt-5.6-terra"}) == "acct-from-token-file" + + +def test_each_deployment_takes_the_account_of_its_own_auth_file(signed_in_chatgpt_token_dir: Path) -> None: + first: Final = _write_login(signed_in_chatgpt_token_dir / "team-a" / "auth.json", "acct-team-a") + second: Final = _write_login(signed_in_chatgpt_token_dir / "team-b" / "auth.json", "acct-team-b") + + assert resolve_subscription_account_id("chatgpt", {"chatgpt_auth_file": str(first)}) == "acct-team-a" + assert resolve_subscription_account_id("chatgpt", {"chatgpt_auth_file": str(second)}) == "acct-team-b" + + +def test_a_configured_access_token_resolves_to_its_own_account(signed_in_chatgpt_token_dir: Path) -> None: + litellm_params: Final = {"model": "chatgpt/gpt-5.6-terra", "api_key": _access_token_for("acct-from-api-key")} + + assert resolve_subscription_account_id("chatgpt", litellm_params) == "acct-from-api-key" + + +def test_a_metered_provider_resolves_to_no_account_even_with_a_chatgpt_login( + signed_in_chatgpt_token_dir: Path, +) -> None: + assert resolve_subscription_account_id("openai", {"model": "openai/gpt-5.6-terra"}) is None + + +def test_chatgpt_without_a_login_resolves_to_no_account(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + empty_dir: Final = tmp_path / "no-login" + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(empty_dir)) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + assert resolve_subscription_account_id("chatgpt", {"model": "chatgpt/gpt-5.6-terra"}) is None + + +def test_resolving_an_account_leaves_the_login_file_untouched(signed_in_chatgpt_token_dir: Path) -> None: + login: Final = signed_in_chatgpt_token_dir / "no-saved-account" / "auth.json" + login.parent.mkdir(parents=True) + saved: Final = json.dumps({"access_token": _access_token_for("acct-derived"), "refresh_token": "rt-current"}) + login.write_text(saved) + + assert resolve_subscription_account_id("chatgpt", {"chatgpt_auth_file": str(login)}) == "acct-derived" + assert login.read_text() == saved + + +@pytest.mark.parametrize("credential_field", ["chatgpt_auth_file", "api_key"]) +def test_a_named_credential_resolves_to_the_account_it_signs_in_as( + credential_field: str, signed_in_chatgpt_token_dir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + login: Final = _write_login(signed_in_chatgpt_token_dir / "named" / "auth.json", "acct-named-credential") + values: Final = { + "chatgpt_auth_file": {"chatgpt_auth_file": str(login)}, + "api_key": {"api_key": _access_token_for("acct-named-credential")}, + }[credential_field] + monkeypatch.setattr( + litellm, + "credential_list", + [CredentialItem(credential_name="chatgpt-team", credential_values=values, credential_info={})], + ) + litellm_params: Final = {"model": "chatgpt/gpt-5.6-terra", "litellm_credential_name": "chatgpt-team"} + + assert resolve_subscription_account_id("chatgpt", litellm_params) == "acct-named-credential" + + +@pytest.mark.parametrize( + ("api_key", "expected_account"), + [("oidc/env/CHATGPT_TEAM_TOKEN", "acct-from-oidc-env"), ("oidc/env/CHATGPT_UNSET_TOKEN", None)], +) +def test_an_oidc_env_api_key_resolves_to_the_account_of_the_token_it_points_to( + api_key: str, expected_account: str | None, signed_in_chatgpt_token_dir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("CHATGPT_TEAM_TOKEN", _access_token_for("acct-from-oidc-env")) + monkeypatch.delenv("CHATGPT_UNSET_TOKEN", raising=False) + litellm_params: Final = {"model": "chatgpt/gpt-5.6-terra", "api_key": api_key} + + assert resolve_subscription_account_id("chatgpt", litellm_params) == expected_account diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index c161aa5fe50..c3a93fd89be 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -6,6 +6,7 @@ import re from collections.abc import AsyncIterator, Callable from contextlib import AbstractAsyncContextManager, asynccontextmanager from datetime import datetime, timedelta, timezone +from pathlib import Path from types import SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, call, patch @@ -4325,6 +4326,36 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_only_on_rows_the_bri assert prisma.spend_log_flush_requested.is_set() is expects_flush + +@pytest.mark.asyncio +@pytest.mark.timeout(30) +async def test_a_chatgpt_responses_row_queues_without_a_default_chatgpt_login( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "no-default-login")) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + own_login: Final = tmp_path / "team" / "auth.json" + own_login.parent.mkdir() + own_login.write_text(json.dumps({"access_token": "at-team", "account_id": "acct-team"})) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "chatgpt-team", + "litellm_params": {"model": "chatgpt/gpt-5.6-terra", "chatgpt_auth_file": str(own_login)}, + "model_info": {"id": "dep-chatgpt"}, + } + ] + ) + prisma: Final = _tool_usage_prisma() + prisma.spend_log_flush_requested = asyncio.Event() + payload: Final = {"request_id": "req-1", "call_type": "aresponses", "model_id": "dep-chatgpt"} + + await DBSpendUpdateWriter()._insert_spend_log_to_db(payload=payload, prisma_client=prisma, llm_router=router) + + assert prisma.spend_log_transactions == [payload] + assert prisma.spend_log_flush_requested.is_set() is False + + def _batch_cost_payload() -> dict: return { **_minimal_spend_payload(), diff --git a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py index c80fa3572fd..5fe4cd2edcd 100644 --- a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py @@ -2553,6 +2553,71 @@ async def test_get_daily_activity_rejects_non_canonical_dates_before_querying(st mock_table.find_many.assert_not_awaited() +def _provider_record(provider: str, *, model: str, spend: float): + base = vars(_spend_record("real-key", model=model, spend=spend)) + return SimpleNamespace(**{**base, "custom_llm_provider": provider, "api_requests": 1, "successful_requests": 1}) + + +def test_per_row_provider_bucket_marks_subscription_backed_providers_only(): + from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + breakdown = BreakdownMetrics() + update_breakdown_metrics(breakdown, _provider_record("chatgpt", model="gpt-5.6-terra", spend=0.0), {}, {}, {}) + update_breakdown_metrics(breakdown, _provider_record("openai", model="gpt-5.6", spend=0.0), {}, {}, {}) + update_breakdown_metrics(breakdown, _provider_record("chatgpt", model="gpt-5.6-sol", spend=0.0), {}, {}, {}) + + assert breakdown.providers["chatgpt"].metadata == {"subscription_covered": True} + assert "subscription_covered" not in breakdown.providers["openai"].metadata + assert breakdown.providers["chatgpt"].metrics.api_requests == 2 + assert "subscription_covered" not in breakdown.models["gpt-5.6-terra"].metadata + + +def test_per_row_provider_marker_keeps_existing_provider_metadata(): + from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + breakdown = BreakdownMetrics() + update_breakdown_metrics( + breakdown, + _provider_record("chatgpt", model="gpt-5.6-terra", spend=0.0), + {}, + {"chatgpt": {"region": "us"}}, + {}, + ) + + assert breakdown.providers["chatgpt"].metadata == {"region": "us", "subscription_covered": True} + + +@pytest.mark.parametrize("provider_row_first", [True, False]) +def test_grouping_sets_provider_bucket_marks_subscription_backed_providers_in_any_row_order(provider_row_first): + from litellm.proxy.management_endpoints.common_daily_activity import ( + _GROUP_DATE_PROVIDER, + _GROUP_DATE_PROVIDER_API_KEY, + _GROUP_GRAND_TOTAL, + _aggregate_grouping_sets_records_sync, + ) + + provider_rows = [ + _grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="chatgpt", spend=0.0), + _grouping_row(_GROUP_DATE_PROVIDER_API_KEY, custom_llm_provider="chatgpt", api_key="real-key", spend=0.0), + ] + ordered_rows = provider_rows if provider_row_first else list(reversed(provider_rows)) + records = [ + *ordered_rows, + _grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="openai", spend=0.0), + _grouping_row(_GROUP_DATE_PROVIDER_API_KEY, custom_llm_provider="openai", api_key="real-key", spend=0.0), + _grouping_row(_GROUP_GRAND_TOTAL, spend=0.0), + ] + + aggregated = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={}) + + providers = aggregated["results"][0].breakdown.providers + assert providers["chatgpt"].metadata == {"subscription_covered": True} + assert "real-key" in providers["chatgpt"].api_key_breakdown + assert "subscription_covered" not in providers["openai"].metadata + + @pytest.mark.parametrize("missing", [None, ""]) def test_missing_dimensions_preserve_all_paginated_metrics(missing: str | None) -> None: records: Final = ( diff --git a/tests/unit/proxy/management_endpoints/test_subscription_accounts.py b/tests/unit/proxy/management_endpoints/test_subscription_accounts.py new file mode 100644 index 00000000000..bcd31efd4b7 --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_subscription_accounts.py @@ -0,0 +1,438 @@ +from collections.abc import Iterator +from datetime import datetime, timezone +from typing import Final +from uuid import uuid4 + +from unittest.mock import MagicMock + +import pytest +from fastapi import FastAPI, HTTPException, Request +from fastapi.testclient import TestClient +from pydantic import ValidationError + +from litellm.models.subscription_account import LiteLLM_SubscriptionAccountTable +from litellm.proxy._types import LiteLLM_JWTAuth, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import allowed_routes_check +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.subscription_accounts import ( + SubscriptionDeployment, + get_subscription_account_repository, + get_subscription_deployments, + router, + subscription_deployments_from_model_list, + subscription_deployments_of_router, +) +from litellm.router import Router +from litellm.types.proxy.management_endpoints.subscription_accounts import ( + SubscriptionAccountFeeRequest, + UpdateSubscriptionAccountFeeRequest, +) +from litellm.repositories.subscription_account_repository import ( + DuplicateSubscriptionAccount, + SubscriptionAccountFeeChanges, + SubscriptionAccountFeeFields, +) + +_ACCOUNT: Final = "org-team-workspace" +_TEST_FEE: Final = { + "custom_llm_provider": "chatgpt", + "account_id": _ACCOUNT, + "monthly_fee": 25.0, + "currency": "USD", + "billing_period_start": "2026-10-01", + "label": "Team plan (test value)", +} +_TWO_DEPLOYMENTS_ONE_ACCOUNT: Final = ( + SubscriptionDeployment(model_name="gpt-5.6-terra", custom_llm_provider="chatgpt", account_id=_ACCOUNT), + SubscriptionDeployment(model_name="gpt-5.6-sol", custom_llm_provider="chatgpt", account_id=_ACCOUNT), +) + + +class _FakeStore: + def __init__(self) -> None: + self.rows: dict[str, LiteLLM_SubscriptionAccountTable] = {} + self.written_columns: list[dict[str, float | str | None]] = [] + + async def list_all(self) -> list[LiteLLM_SubscriptionAccountTable]: + return sorted(self.rows.values(), key=lambda row: row.created_at) + + async def find_by_id(self, id_value: str) -> LiteLLM_SubscriptionAccountTable | None: + return self.rows.get(id_value) + + async def create_account( + self, custom_llm_provider: str, account_id: str, fee: SubscriptionAccountFeeFields, created_by: str + ) -> LiteLLM_SubscriptionAccountTable | DuplicateSubscriptionAccount: + for row in self.rows.values(): + if (row.custom_llm_provider, row.account_id) == (custom_llm_provider, account_id): + return DuplicateSubscriptionAccount(custom_llm_provider=custom_llm_provider, account_id=account_id) + now = datetime.now(timezone.utc) + row = LiteLLM_SubscriptionAccountTable( + subscription_account_id=str(uuid4()), + custom_llm_provider=custom_llm_provider, + account_id=account_id, + label=fee.label, + monthly_fee=fee.monthly_fee, + currency=fee.currency, + billing_period_start=fee.billing_period_start, + created_at=now, + created_by=created_by, + updated_at=now, + updated_by=created_by, + ) + self.rows[row.subscription_account_id] = row + return row + + async def update_fee( + self, subscription_account_id: str, changes: SubscriptionAccountFeeChanges, updated_by: str + ) -> LiteLLM_SubscriptionAccountTable | None: + columns = dict(changes.columns()) + self.written_columns.append(columns) + existing = self.rows.get(subscription_account_id) + if existing is None: + return None + updated = existing.model_copy( + update={**columns, "updated_by": updated_by, "updated_at": datetime.now(timezone.utc)} + ) + self.rows[subscription_account_id] = updated + return updated + + async def delete_account(self, subscription_account_id: str) -> LiteLLM_SubscriptionAccountTable | None: + return self.rows.pop(subscription_account_id, None) + + +def _resolve_auth(request: Request) -> UserAPIKeyAuth: + role: Final = LitellmUserRoles(request.headers.get("x-user-role", LitellmUserRoles.PROXY_ADMIN.value)) + return UserAPIKeyAuth(user_id="admin", user_role=role, api_key="sk-test") + + +@pytest.fixture +def deployments() -> tuple[SubscriptionDeployment, ...]: + return _TWO_DEPLOYMENTS_ONE_ACCOUNT + + +@pytest.fixture +def client(deployments: tuple[SubscriptionDeployment, ...]) -> Iterator[tuple[TestClient, _FakeStore]]: + store: Final = _FakeStore() + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[get_subscription_account_repository] = lambda: store + app.dependency_overrides[get_subscription_deployments] = lambda: deployments + app.dependency_overrides[user_api_key_auth] = _resolve_auth + with TestClient(app) as test_client: + yield test_client, store + + +def _usage(test_client: TestClient, start_date: str, end_date: str, **headers: str) -> dict: + response = test_client.get( + "/subscription_accounts/usage", params={"start_date": start_date, "end_date": end_date}, headers=headers + ) + assert response.status_code == 200, response.text + return response.json() + + +def test_untracked_account_lists_its_deployments_with_no_fee(client): + test_client, _ = client + + body = _usage(test_client, "2026-10-01", "2026-10-07") + + assert body["subscription_providers"] == ["chatgpt"] + assert body["accounts"] == [ + { + "custom_llm_provider": "chatgpt", + "account_id": _ACCOUNT, + "deployments": ["gpt-5.6-sol", "gpt-5.6-terra"], + "fee": None, + "billing_periods": [], + "fixed_cost": None, + } + ] + + +def test_fee_counts_once_per_billing_period_however_many_deployments_share_the_account(client): + test_client, _ = client + created = test_client.post("/subscription_accounts/new", json=_TEST_FEE) + assert created.status_code == 200, created.text + + one_period = _usage(test_client, "2026-10-01", "2026-10-07") + two_periods = _usage(test_client, "2026-09-15", "2026-11-15") + between_starts = _usage(test_client, "2026-10-02", "2026-10-31") + + (account,) = one_period["accounts"] + assert account["deployments"] == ["gpt-5.6-sol", "gpt-5.6-terra"] + assert account["fee"]["monthly_fee"] == 25.0 + assert account["fee"]["currency"] == "USD" + assert account["fee"]["label"] == "Team plan (test value)" + assert account["billing_periods"] == [{"start": "2026-10-01", "end": "2026-10-31"}] + assert account["fixed_cost"] == 25.0 + assert two_periods["accounts"][0]["fixed_cost"] == 50.0 + assert [period["start"] for period in two_periods["accounts"][0]["billing_periods"]] == [ + "2026-10-01", + "2026-11-01", + ] + assert between_starts["accounts"][0]["fixed_cost"] == 0.0 + assert between_starts["accounts"][0]["billing_periods"] == [] + + +def test_fee_of_an_account_no_deployment_uses_stays_listed_and_counted(client): + test_client, _ = client + other = {**_TEST_FEE, "account_id": "org-other", "monthly_fee": 200.0, "currency": "EUR"} + assert test_client.post("/subscription_accounts/new", json=other).status_code == 200 + assert test_client.post("/subscription_accounts/new", json=_TEST_FEE).status_code == 200 + + body = _usage(test_client, "2026-10-01", "2026-10-07") + + unused, used = body["accounts"] + assert (unused["account_id"], unused["deployments"], unused["fee"]["currency"]) == ("org-other", [], "EUR") + assert unused["fixed_cost"] == 200.0 + assert (used["account_id"], used["fixed_cost"]) == (_ACCOUNT, 25.0) + + +@pytest.mark.parametrize( + "deployments", + [(SubscriptionDeployment(model_name="gpt-5.6-terra", custom_llm_provider="chatgpt", account_id=None),)], +) +def test_account_not_signed_in_shows_no_account_and_takes_no_fee(client): + test_client, _ = client + assert test_client.post("/subscription_accounts/new", json=_TEST_FEE).status_code == 200 + + body = _usage(test_client, "2026-10-01", "2026-10-07") + + unsigned = next(account for account in body["accounts"] if account["account_id"] is None) + assert unsigned["deployments"] == ["gpt-5.6-terra"] + assert unsigned["fee"] is None + assert unsigned["fixed_cost"] is None + + +def test_update_changes_the_fee_and_delete_returns_the_account_to_untracked(client): + test_client, store = client + created = test_client.post("/subscription_accounts/new", json=_TEST_FEE).json() + account_id = created["subscription_account_id"] + + updated = test_client.post( + "/subscription_accounts/update", + json={"subscription_account_id": account_id, "monthly_fee": 30.0, "billing_period_start": "2026-10-05"}, + ) + assert updated.status_code == 200, updated.text + assert updated.json()["monthly_fee"] == 30.0 + assert updated.json()["currency"] == "USD" + assert updated.json()["billing_period_start"] == "2026-10-05" + assert updated.json()["label"] == _TEST_FEE["label"] + assert _usage(test_client, "2026-10-01", "2026-10-07")["accounts"][0]["fixed_cost"] == 30.0 + + cleared = test_client.post( + "/subscription_accounts/update", json={"subscription_account_id": account_id, "label": None} + ) + assert cleared.status_code == 200, cleared.text + assert cleared.json()["label"] is None + assert cleared.json()["monthly_fee"] == 30.0 + assert store.written_columns == [{"monthly_fee": 30.0, "billing_period_start": "2026-10-05"}, {"label": None}] + + unknown = test_client.post( + "/subscription_accounts/update", json={"subscription_account_id": "no-such-account", "monthly_fee": 1.0} + ) + assert unknown.status_code == 404 + + deleted = test_client.post("/subscription_accounts/delete", json={"subscription_account_id": account_id}) + assert deleted.status_code == 200, deleted.text + assert store.rows == {} + assert _usage(test_client, "2026-10-01", "2026-10-07")["accounts"][0]["fixed_cost"] is None + + missing = test_client.post("/subscription_accounts/delete", json={"subscription_account_id": account_id}) + assert missing.status_code == 404 + + +def test_rejects_a_second_fee_for_the_same_account(client): + test_client, _ = client + assert test_client.post("/subscription_accounts/new", json=_TEST_FEE).status_code == 200 + + duplicate = test_client.post("/subscription_accounts/new", json={**_TEST_FEE, "monthly_fee": 99.0}) + + assert duplicate.status_code == 400 + assert "/subscription_accounts/update" in duplicate.json()["detail"]["error"] + + +def test_rejects_a_fee_on_a_metered_provider(client): + test_client, store = client + + response = test_client.post("/subscription_accounts/new", json={**_TEST_FEE, "custom_llm_provider": "openai"}) + + assert response.status_code == 400 + assert store.rows == {} + + +@pytest.mark.parametrize( + "payload", + [ + {**_TEST_FEE, "monthly_fee": 0}, + {**_TEST_FEE, "account_id": ""}, + {**_TEST_FEE, "currency": "usd"}, + {**_TEST_FEE, "currency": "$"}, + {**_TEST_FEE, "billing_period_start": "10/01/2026"}, + ], +) +def test_rejects_malformed_fee_fields(client, payload): + test_client, store = client + + assert test_client.post("/subscription_accounts/new", json=payload).status_code == 422 + assert store.rows == {} + + +def test_rejects_an_infinite_monthly_fee(): + with pytest.raises(ValidationError): + SubscriptionAccountFeeRequest.model_validate({**_TEST_FEE, "monthly_fee": float("inf")}) + with pytest.raises(ValidationError): + UpdateSubscriptionAccountFeeRequest.model_validate( + {"subscription_account_id": "fee-1", "monthly_fee": float("inf")} + ) + + +@pytest.mark.parametrize("deployments", [()]) +def test_fee_stays_listed_after_the_last_subscription_deployment_is_removed(client): + test_client, _ = client + assert test_client.post("/subscription_accounts/new", json=_TEST_FEE).status_code == 200 + + body = _usage(test_client, "2026-10-01", "2026-10-07") + + (account,) = body["accounts"] + assert (account["account_id"], account["deployments"], account["fixed_cost"]) == (_ACCOUNT, [], 25.0) + + +def test_usage_rejects_a_malformed_or_inverted_range(client): + test_client, _ = client + + malformed = [ + test_client.get("/subscription_accounts/usage", params={"start_date": start, "end_date": "2026-10-07"}) + for start in ("2026-10-1", "20261001", "2026-W40-1", "2026-13-01") + ] + inverted = test_client.get( + "/subscription_accounts/usage", params={"start_date": "2026-10-07", "end_date": "2026-10-01"} + ) + + assert [response.status_code for response in malformed] == [400, 400, 400, 400] + assert inverted.status_code == 400 + + +def test_viewer_reads_usage_but_cannot_write(client): + test_client, store = client + viewer = {"x-user-role": LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value} + + assert _usage(test_client, "2026-10-01", "2026-10-07", **viewer)["accounts"][0]["account_id"] == _ACCOUNT + assert test_client.post("/subscription_accounts/new", json=_TEST_FEE, headers=viewer).status_code == 403 + assert store.rows == {} + + +_SUBSCRIPTION_ROUTES: Final = ( + "/subscription_accounts/usage", + "/subscription_accounts/new", + "/subscription_accounts/update", + "/subscription_accounts/delete", +) + + +def test_routes_follow_the_admin_route_policy(): + jwt_admin_defaults = LiteLLM_JWTAuth() + + for route in _SUBSCRIPTION_ROUTES: + assert allowed_routes_check( + user_role=LitellmUserRoles.PROXY_ADMIN, user_route=route, litellm_proxy_roles=jwt_admin_defaults + ) + assert RouteChecks.is_management_route(route) + + +def test_viewer_route_gate_reads_usage_and_blocks_writes(): + viewer_role: Final = LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value + get_usage = MagicMock(spec=Request, method="GET", query_params={}) + post_new = MagicMock(spec=Request, method="POST", query_params={}) + + RouteChecks._check_proxy_admin_viewer_access( + route="/subscription_accounts/usage", _user_role=viewer_role, request_data={}, request=get_usage + ) + with pytest.raises(HTTPException) as denied: + RouteChecks._check_proxy_admin_viewer_access( + route="/subscription_accounts/new", _user_role=viewer_role, request_data=_TEST_FEE, request=post_new + ) + assert denied.value.status_code == 403 + + +def test_internal_user_cannot_read_subscription_accounts(client): + test_client, _ = client + + response = test_client.get( + "/subscription_accounts/usage", + params={"start_date": "2026-10-01", "end_date": "2026-10-07"}, + headers={"x-user-role": LitellmUserRoles.INTERNAL_USER.value}, + ) + + assert response.status_code == 403 + + +def test_deployments_take_the_account_of_their_own_login(): + model_list = [ + {"model_name": "gpt-5.6-terra", "litellm_params": {"model": "chatgpt/gpt-5.6-terra"}}, + {"model_name": "gpt-5.6-sol", "litellm_params": {"model": "chatgpt/gpt-5.6-sol"}}, + { + "model_name": "gpt-5.6-luna", + "litellm_params": {"model": "chatgpt/gpt-5.6-luna", "chatgpt_auth_file": "/logins/other.json"}, + }, + {"model_name": "sonnet", "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-x"}}, + {"model_name": "broken", "litellm_params": {"model": "no-such-provider-model"}}, + ] + lookups: list[tuple[str, object]] = [] + + def resolver(provider: str, litellm_params) -> str | None: + lookups.append((provider, litellm_params.get("model"))) + return "org-other" if litellm_params.get("chatgpt_auth_file") == "/logins/other.json" else _ACCOUNT + + deployments = subscription_deployments_from_model_list(model_list, resolver) + + assert deployments == ( + *_TWO_DEPLOYMENTS_ONE_ACCOUNT, + SubscriptionDeployment(model_name="gpt-5.6-luna", custom_llm_provider="chatgpt", account_id="org-other"), + ) + assert lookups == [ + ("chatgpt", "chatgpt/gpt-5.6-terra"), + ("chatgpt", "chatgpt/gpt-5.6-sol"), + ("chatgpt", "chatgpt/gpt-5.6-luna"), + ] + + +def test_deployments_list_each_router_deployment_once_by_the_name_callers_use(): + llm_router: Final = Router( + model_list=[ + {"model_name": "gpt-5.6-sol", "litellm_params": {"model": "chatgpt/gpt-5.6-sol"}}, + { + "model_name": "model_name_team-1_7f3a", + "litellm_params": {"model": "chatgpt/gpt-5.6-terra"}, + "model_info": {"team_id": "team-1", "team_public_model_name": "team-terra"}, + }, + ], + model_group_alias={"sol-alias": "gpt-5.6-sol"}, + ) + + deployments = subscription_deployments_of_router(llm_router, lambda provider, litellm_params: _ACCOUNT) + + assert sorted(deployment.model_name for deployment in deployments) == ["gpt-5.6-sol", "team-terra"] + + +@pytest.mark.parametrize( + "deployments", + [ + ( + *_TWO_DEPLOYMENTS_ONE_ACCOUNT, + SubscriptionDeployment(model_name="gpt-5.6-luna", custom_llm_provider="chatgpt", account_id="org-other"), + ) + ], +) +def test_two_logins_each_count_their_own_fee_once(client): + test_client, _ = client + other = {**_TEST_FEE, "account_id": "org-other", "monthly_fee": 40.0} + assert test_client.post("/subscription_accounts/new", json=_TEST_FEE).status_code == 200 + assert test_client.post("/subscription_accounts/new", json=other).status_code == 200 + + body = _usage(test_client, "2026-10-01", "2026-10-07") + + assert [(a["account_id"], a["deployments"], a["fixed_cost"]) for a in body["accounts"]] == [ + ("org-other", ["gpt-5.6-luna"], 40.0), + (_ACCOUNT, ["gpt-5.6-sol", "gpt-5.6-terra"], 25.0), + ] diff --git a/tests/unit/proxy/spend_tracking/test_subscription_billing_periods.py b/tests/unit/proxy/spend_tracking/test_subscription_billing_periods.py new file mode 100644 index 00000000000..c2c98a05b3f --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_subscription_billing_periods.py @@ -0,0 +1,73 @@ +from datetime import date + +from litellm.proxy.spend_tracking.subscription_billing_periods import ( + BillingPeriodRange, + billing_periods_starting_within, +) + + +def test_one_period_per_month_start_inside_the_filter(): + periods = billing_periods_starting_within(date(2026, 8, 5), date(2026, 9, 1), date(2026, 10, 31)) + + assert periods == ( + BillingPeriodRange(start=date(2026, 9, 5), end=date(2026, 10, 4)), + BillingPeriodRange(start=date(2026, 10, 5), end=date(2026, 11, 4)), + ) + + +def test_filter_boundaries_are_inclusive(): + anchor = date(2026, 10, 1) + + assert billing_periods_starting_within(anchor, date(2026, 10, 1), date(2026, 10, 7)) == ( + BillingPeriodRange(start=date(2026, 10, 1), end=date(2026, 10, 31)), + ) + assert billing_periods_starting_within(anchor, date(2026, 9, 25), date(2026, 10, 1)) == ( + BillingPeriodRange(start=date(2026, 10, 1), end=date(2026, 10, 31)), + ) + + +def test_filter_between_two_period_starts_attributes_nothing(): + assert billing_periods_starting_within(date(2026, 10, 15), date(2026, 10, 1), date(2026, 10, 7)) == () + + +def test_no_period_before_the_anchor(): + assert billing_periods_starting_within(date(2026, 10, 1), date(2026, 8, 1), date(2026, 9, 30)) == () + assert billing_periods_starting_within(date(2026, 10, 1), date(2026, 8, 1), date(2026, 10, 1)) == ( + BillingPeriodRange(start=date(2026, 10, 1), end=date(2026, 10, 31)), + ) + + +def test_anchor_day_clamps_to_short_months_without_drifting(): + periods = billing_periods_starting_within(date(2026, 1, 31), date(2026, 2, 1), date(2026, 5, 31)) + + assert [period.start for period in periods] == [ + date(2026, 2, 28), + date(2026, 3, 31), + date(2026, 4, 30), + date(2026, 5, 31), + ] + assert periods[0].end == date(2026, 3, 30) + assert periods[-1].end == date(2026, 6, 29) + + +def test_period_ends_the_day_before_the_next_start_across_a_year_boundary(): + periods = billing_periods_starting_within(date(2025, 12, 20), date(2025, 12, 1), date(2026, 1, 31)) + + assert periods == ( + BillingPeriodRange(start=date(2025, 12, 20), end=date(2026, 1, 19)), + BillingPeriodRange(start=date(2026, 1, 20), end=date(2026, 2, 19)), + ) + + +def test_inverted_filter_is_empty(): + assert billing_periods_starting_within(date(2026, 1, 1), date(2026, 3, 1), date(2026, 2, 1)) == () + + +def test_the_calendar_s_last_month_ends_on_its_last_day_instead_of_failing(): + assert billing_periods_starting_within(date(9999, 11, 15), date(9999, 11, 1), date(9999, 12, 31)) == ( + BillingPeriodRange(start=date(9999, 11, 15), end=date(9999, 12, 14)), + BillingPeriodRange(start=date(9999, 12, 15), end=date(9999, 12, 31)), + ) + assert billing_periods_starting_within(date(9999, 12, 1), date(9999, 12, 1), date(9999, 12, 31)) == ( + BillingPeriodRange(start=date(9999, 12, 1), end=date(9999, 12, 31)), + ) diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py index cbcd0b1cd28..4e410e22b7f 100644 --- a/tests/unit/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -30,6 +30,11 @@ from litellm.repositories.object_permission_repository import ( ) from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.subscription_account_repository import ( + SetLabel, + SubscriptionAccountFeeChanges, + SubscriptionAccountRepository, +) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import ( @@ -140,6 +145,7 @@ class MockPrismaClient: self.db.litellm_projecttable = MockTable(pk_field="project_id") self.db.litellm_objectpermissiontable = MockTable(pk_field="object_permission_id") self.db.litellm_credentialstable = MockTable() + self.db.litellm_subscriptionaccounttable = MockTable(pk_field="subscription_account_id") class TestBaseRepository: @@ -2461,3 +2467,40 @@ class TestAutoRouterSessionRepository: assert AutoRouterSessionRepository(client).table is client.db.litellm_autoroutersession with pytest.raises(RuntimeError, match="No DB Connected"): _ = AutoRouterSessionRepository(None).table + + +class TestSubscriptionAccountRepository: + @pytest.mark.asyncio + async def test_update_fee_writes_only_the_named_columns(self): + client = MockPrismaClient() + stored_at = datetime(2026, 10, 1) + await client.db.litellm_subscriptionaccounttable.create( + data={ + "subscription_account_id": "sub-1", + "custom_llm_provider": "chatgpt", + "account_id": "acct_123", + "label": "Team plan (test value)", + "monthly_fee": 25.0, + "currency": "USD", + "billing_period_start": "2026-10-01", + "created_at": stored_at, + "created_by": "admin-a", + "updated_at": stored_at, + "updated_by": "admin-a", + } + ) + repo = SubscriptionAccountRepository(client) + + fee_only = await repo.update_fee("sub-1", SubscriptionAccountFeeChanges(monthly_fee=30.0), "admin-b") + label_cleared = await repo.update_fee("sub-1", SubscriptionAccountFeeChanges(label=SetLabel(None)), "admin-c") + missing = await repo.update_fee("sub-2", SubscriptionAccountFeeChanges(monthly_fee=1.0), "admin-b") + + assert fee_only is not None and label_cleared is not None + assert (fee_only.monthly_fee, fee_only.currency, fee_only.label, fee_only.updated_by) == ( + 30.0, + "USD", + "Team plan (test value)", + "admin-b", + ) + assert (label_cleared.monthly_fee, label_cleared.label, label_cleared.updated_by) == (30.0, None, "admin-c") + assert missing is None diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx index 660168fa6ab..5c8b284c72f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx @@ -1160,6 +1160,34 @@ describe("EntityUsage", () => { expect(screen.getByTestId("provider-spend-total")).toHaveTextContent("$100.50"); }); + it("lists a subscription-covered provider on the team view as included in the subscription at zero spend", async () => { + const day = mockSpendData.results[0]; + mockTeamDailyActivityCall.mockResolvedValue({ + ...mockSpendData, + results: [ + { + ...day, + breakdown: { + ...day.breakdown, + providers: { + ...day.breakdown.providers, + chatgpt: { + metrics: { ...day.breakdown.providers.openai.metrics, spend: 0, api_requests: 4 }, + metadata: { subscription_covered: true }, + }, + }, + }, + }, + ], + }); + + render(); + + expect(await screen.findByText("Included in subscription")).toBeInTheDocument(); + expect(screen.getByText("chatgpt")).toBeInTheDocument(); + expect(screen.getAllByTestId("provider-share-segment")).toHaveLength(1); + }); + it("should label the chart with user_email metadata instead of the raw UUID (LIT-3889)", async () => { const userUuid = "c0e68be8-057e-4e2f-9d3a-000000000000"; const spendDataForUser = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.test.tsx index 8da636e68a7..dca5ff91fa6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.test.tsx @@ -258,6 +258,34 @@ describe("SpendByProvider", () => { expect(screen.getByTestId("provider-spend-total")).toHaveTextContent("$200.50"); }); + it("always shows a subscription-covered provider and marks it included in the subscription", () => { + const providerSpendWithSubscription = [ + { + provider: "chatgpt", + spend: 0, + requests: 7, + successful_requests: 7, + failed_requests: 0, + tokens: 11608, + subscription_covered: true, + }, + { + provider: "google", + spend: 0, + requests: 0, + successful_requests: 0, + failed_requests: 0, + tokens: 0, + }, + ]; + render(); + expect(screen.getAllByText("chatgpt").length).toBeGreaterThan(0); + expect(screen.getByText("Included in subscription")).toBeInTheDocument(); + expect(screen.getByText("11,608")).toBeInTheDocument(); + expect(screen.queryByText("google")).not.toBeInTheDocument(); + expect(screen.getByText("Included in subscription").closest("td")).not.toHaveTextContent("$"); + }); + it("should include all providers with spend greater than zero by default", () => { const providerSpendWithMixed = [ { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx index 13d9a812a12..befb7d13b55 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx @@ -4,6 +4,7 @@ import { MoneyCell } from "@/components/shared/table_cells"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { Info, Server } from "lucide-react"; import type { ColumnDef } from "@tanstack/react-table"; +import { Badge } from "@/components/ui/badge"; import { Switch } from "@/components/ui/switch"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import React, { useId, useState } from "react"; @@ -11,13 +12,14 @@ import { ProviderLogo } from "@/components/molecules/models/ProviderLogo"; import { ChartLoader } from "@/components/shared/chart_loader"; import { Panel } from "../overview/Primitives"; -type ProviderSpendData = { +export type ProviderSpendData = { provider: string; spend: number; requests: number; successful_requests: number; failed_requests: number; tokens: number; + subscription_covered?: boolean; }; interface SpendByProviderProps { @@ -48,7 +50,15 @@ const columns: ColumnDef[] = [ header: "Spend", accessorKey: "spend", meta: { numeric: true, ...QUIET_HEADER }, - cell: ({ row }) => , + cell: ({ row }) => + row.original.subscription_covered ? ( +
+ {row.original.spend > 0 && } + Included in subscription +
+ ) : ( + + ), }, { header: "Successful", @@ -164,6 +174,10 @@ const SpendByProvider: React.FC = ({ loading, isDateChangi return includeUnknown; } + if (provider.subscription_covered) { + return true; + } + // If includeZeroSpend is true, include all providers (including those with 0 spend) // Otherwise, only include providers with spend > 0 if (includeZeroSpend) { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts index 54569608b0e..be9f44d0786 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts @@ -1,6 +1,7 @@ import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel"; import type { TopKeyItem } from "@/components/UsagePage/components/EntityUsage/TopKeyView"; import { BreakdownMetrics, DailyData, KeyMetricWithMetadata, TagUsage } from "@/components/UsagePage/types"; +import { subscriptionCoveredProviders } from "../SubscriptionCost/types"; export type ExtendedDailyData = DailyData & { breakdown: BreakdownMetrics; @@ -15,6 +16,7 @@ export interface ProviderSpendRow extends Record { successful_requests: number; failed_requests: number; tokens: number; + subscription_covered: boolean; } export const getTopModels = ( @@ -205,6 +207,7 @@ export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number }; export const getProviderSpend = (results: ExtendedDailyData[]): ProviderSpendRow[] => { + const covered = subscriptionCoveredProviders(results); const providerSpend: Record = {}; results.forEach((day) => { Object.entries(day.breakdown.providers || {}).forEach(([provider, metrics]) => { @@ -216,6 +219,7 @@ export const getProviderSpend = (results: ExtendedDailyData[]): ProviderSpendRow successful_requests: 0, failed_requests: 0, tokens: 0, + subscription_covered: covered.has(provider), }; } try { @@ -231,6 +235,6 @@ export const getProviderSpend = (results: ExtendedDailyData[]): ProviderSpendRow }); return Object.values(providerSpend) - .filter((provider) => provider.spend > 0) + .filter((provider) => provider.spend > 0 || provider.subscription_covered) .sort((a, b) => b.spend - a.spend); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCostCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCostCard.test.tsx new file mode 100644 index 00000000000..37f3bf4abf4 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCostCard.test.tsx @@ -0,0 +1,185 @@ +import { fireEvent, render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import SubscriptionCostCard from "./SubscriptionCostCard"; +import type { SubscriptionAccountUsage, SubscriptionFee } from "./types"; + +vi.mock("@/components/shared/chart_loader", () => ({ + ChartLoader: () =>
Loading chart data...
, +})); + +vi.mock("@/components/molecules/models/ProviderLogo", () => ({ + ProviderLogo: ({ provider }: { provider: string }) =>
, +})); + +const fee = (overrides: Partial = {}): SubscriptionFee => ({ + subscription_account_id: "fee-1", + custom_llm_provider: "chatgpt", + account_id: "acct_123", + label: null, + monthly_fee: 25, + currency: "USD", + billing_period_start: "2026-10-01", + created_at: "2026-10-06T00:00:00Z", + updated_at: "2026-10-06T00:00:00Z", + created_by: "admin", + updated_by: "admin", + ...overrides, +}); + +const account = (overrides: Partial = {}): SubscriptionAccountUsage => ({ + custom_llm_provider: "chatgpt", + account_id: "acct_123", + deployments: ["gpt-5.6-terra", "gpt-5.6-sol"], + fee: null, + billing_periods: [], + fixed_cost: null, + ...overrides, +}); + +const handlers = () => ({ onSetFee: vi.fn(), onEditFee: vi.fn(), onRemoveFee: vi.fn() }); + +const renderCard = (accounts: SubscriptionAccountUsage[], props: { canEdit?: boolean; meteredSpend?: number } = {}) => { + const actions = handlers(); + const view = render( + , + ); + return { ...view, ...actions }; +}; + +describe("SubscriptionCostCard", () => { + it("renders nothing when there are no subscription accounts", () => { + const { container } = render( + , + ); + expect(container).toBeEmptyDOMElement(); + expect(screen.queryByText("Subscription fixed costs")).not.toBeInTheDocument(); + }); + + it("calls out an untracked fixed cost and offers to set the fee only to an editor", () => { + const { onSetFee, unmount } = renderCard([account()]); + expect(screen.getByText("Fixed cost not tracked")).toBeInTheDocument(); + expect(screen.getByText("gpt-5.6-terra, gpt-5.6-sol")).toBeInTheDocument(); + expect(screen.getByText("None in range")).toBeInTheDocument(); + expect(screen.getByTestId("fixed-cost-total")).toHaveTextContent("Fixed cost in range: none"); + fireEvent.click(screen.getByRole("button", { name: "Set monthly fee" })); + expect(onSetFee).toHaveBeenCalledWith(account()); + unmount(); + + renderCard([account()], { canEdit: false }); + expect(screen.getByText("Fixed cost not tracked")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Set monthly fee" })).not.toBeInTheDocument(); + }); + + it("shows a configured fee once for two deployments sharing the account and the fixed cost in range", () => { + const covered = account({ + fee: fee({ label: "Team workspace" }), + billing_periods: [{ start: "2026-10-01", end: "2026-10-31" }], + fixed_cost: 25, + }); + const { onEditFee, onRemoveFee } = renderCard([covered], { meteredSpend: 0.0838 }); + + expect(screen.getAllByText("25.00 USD")).toHaveLength(2); + expect(screen.getByText("Team workspace")).toBeInTheDocument(); + expect(screen.getByText("Oct 1 - Oct 31, 2026")).toBeInTheDocument(); + expect(screen.getByTestId("fixed-cost-total")).toHaveTextContent("Fixed cost in range: 25.00 USD"); + expect(screen.getByTestId("metered-spend")).toHaveTextContent("Usage-billed spend in range, all providers: $0.08"); + expect(screen.queryByText(/25\.08/)).not.toBeInTheDocument(); + + fireEvent.click(screen.getByRole("button", { name: "Edit" })); + expect(onEditFee).toHaveBeenCalledWith(covered); + }); + + it("removes a fee only after the admin confirms it", async () => { + const covered = account({ fee: fee(), billing_periods: [], fixed_cost: 0 }); + const { onRemoveFee } = renderCard([covered]); + + fireEvent.click(screen.getByRole("button", { name: "Remove" })); + expect(onRemoveFee).not.toHaveBeenCalled(); + fireEvent.click(await screen.findByRole("button", { name: "Cancel" })); + expect(onRemoveFee).not.toHaveBeenCalled(); + + fireEvent.click(screen.getByRole("button", { name: "Remove" })); + fireEvent.click(await screen.findByRole("button", { name: "Remove fee" })); + expect(onRemoveFee).toHaveBeenCalledWith(covered); + }); + + it("sums fixed costs per currency and keeps a second currency separate", () => { + const october = { start: "2026-10-01", end: "2026-10-31" }; + const twoPeriods = { + fee: fee(), + billing_periods: [{ start: "2026-09-01", end: "2026-09-30" }, october], + fixed_cost: 50, + }; + const secondUsd = { subscription_account_id: "fee-2", account_id: "acct_456", monthly_fee: 10 }; + const noDeployments = { + account_id: "acct_456", + deployments: [], + fee: fee(secondUsd), + billing_periods: [october], + fixed_cost: 10, + }; + const euro = { subscription_account_id: "fee-3", account_id: "acct_789", monthly_fee: 20, currency: "EUR" }; + const inEuros = { account_id: "acct_789", fee: fee(euro), billing_periods: [october], fixed_cost: 20 }; + renderCard([account(twoPeriods), account(noDeployments), account(inEuros)]); + + const totals = screen.getAllByTestId("fixed-cost-total").map((node) => node.textContent); + expect(totals).toEqual(["Fixed cost in range: 60.00 USD", "Fixed cost in range: 20.00 EUR"]); + expect(screen.getByText("None")).toBeInTheDocument(); + expect(screen.getByText("Sep 1 - Sep 30, 2026")).toBeInTheDocument(); + }); + + it("explains an unresolved account and offers no action for it", () => { + renderCard([account({ account_id: null })]); + expect(screen.getByText("Account not signed in on the proxy")).toBeInTheDocument(); + expect(screen.getByText("Fixed cost not tracked")).toBeInTheDocument(); + expect(screen.queryByRole("button")).not.toBeInTheDocument(); + }); + + it("shows the loader instead of the table while a range is loading", () => { + render( + , + ); + expect(screen.getByTestId("chart-loader")).toBeInTheDocument(); + expect(screen.queryByText("25.00 USD")).not.toBeInTheDocument(); + }); + + it("says so instead of disappearing when the usage request fails", () => { + render( + , + ); + expect(screen.getByText("Could not load subscription accounts for this range.")).toBeInTheDocument(); + expect(screen.queryByTestId("chart-loader")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCostCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCostCard.tsx new file mode 100644 index 00000000000..deee30a2db5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCostCard.tsx @@ -0,0 +1,221 @@ +import type { ColumnDef } from "@tanstack/react-table"; +import { Info } from "lucide-react"; +import React from "react"; +import { ProviderLogo } from "@/components/molecules/models/ProviderLogo"; +import { ChartLoader } from "@/components/shared/chart_loader"; +import { DataTable } from "@/components/shared/DataTable"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, + AlertDialogTrigger, +} from "@/components/ui/alert-dialog"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { ATTRIBUTION_RULE, fixedCostByCurrency, formatBillingPeriod, formatFee } from "./formatting"; +import type { SubscriptionAccountUsage, SubscriptionUsageResponse } from "./types"; + +interface SubscriptionCostCardProps { + loading: boolean; + failed: boolean; + isDateChanging: boolean; + usage: SubscriptionUsageResponse | null; + meteredSpend: number; + canEdit: boolean; + onSetFee: (account: SubscriptionAccountUsage) => void; + onEditFee: (account: SubscriptionAccountUsage) => void; + onRemoveFee: (account: SubscriptionAccountUsage) => void; +} + +type FeeActions = Pick; + +const accountRowId = (account: SubscriptionAccountUsage): string => + `${account.custom_llm_provider}:${account.account_id ?? "unresolved"}`; + +const AccountCell = ({ account }: { account: SubscriptionAccountUsage }) => { + if (account.account_id === null) { + return Account not signed in on the proxy; + } + return ( +
+ {account.account_id} + {account.fee?.label && {account.fee.label}} +
+ ); +}; + +const ActionsCell = ({ + account, + onSetFee, + onEditFee, + onRemoveFee, +}: { account: SubscriptionAccountUsage } & FeeActions) => { + if (account.account_id === null) return null; + if (account.fee === null) { + return ( + + ); + } + return ( +
+ + + }>Remove + + + Remove this monthly fee? + + {account.account_id} goes back to "Fixed cost not tracked". Its requests stay included in the + subscription + + + + Cancel + onRemoveFee(account)}> + Remove fee + + + + +
+ ); +}; + +const baseColumns: ColumnDef[] = [ + { + header: "Provider", + accessorKey: "custom_llm_provider", + cell: ({ row }) => ( +
+ + {row.original.custom_llm_provider} +
+ ), + }, + { + header: "Account", + accessorKey: "account_id", + cell: ({ row }) => , + }, + { + header: "Deployments", + id: "deployments", + cell: ({ row }) => (row.original.deployments.length === 0 ? "None" : row.original.deployments.join(", ")), + }, + { + header: "Monthly fee", + id: "monthly_fee", + meta: { numeric: true }, + cell: ({ row }) => + row.original.fee === null ? ( + Fixed cost not tracked + ) : ( + formatFee(row.original.fee.monthly_fee, row.original.fee.currency) + ), + }, + { + header: "Billing periods in range", + id: "billing_periods", + cell: ({ row }) => + row.original.billing_periods.length === 0 ? ( + None in range + ) : ( +
+ {row.original.billing_periods.map((period) => ( + {formatBillingPeriod(period)} + ))} +
+ ), + }, + { + header: "Fixed cost in range", + id: "fixed_cost", + meta: { numeric: true }, + cell: ({ row }) => + row.original.fee === null || row.original.fixed_cost === null + ? "-" + : formatFee(row.original.fixed_cost, row.original.fee.currency), + }, +]; + +const actionsColumn = (actions: FeeActions): ColumnDef => ({ + header: "", + id: "actions", + meta: { numeric: true }, + cell: ({ row }) => , +}); + +const SubscriptionCostCard: React.FC = ({ + loading, + failed, + isDateChanging, + usage, + meteredSpend, + canEdit, + onSetFee, + onEditFee, + onRemoveFee, +}) => { + if (!failed && (usage === null || usage.accounts.length === 0)) return null; + const columns = canEdit ? [...baseColumns, actionsColumn({ onSetFee, onEditFee, onRemoveFee })] : baseColumns; + const accounts = usage?.accounts ?? []; + const totals = fixedCostByCurrency(accounts); + + return ( + + + + Subscription fixed costs + + } /> + {ATTRIBUTION_RULE} + + + + + {failed &&

Could not load subscription accounts for this range.

} + {!failed && loading && } + {!failed && !loading && ( +
+ +
+ {totals.length === 0 ? ( + Fixed cost in range: none + ) : ( + totals.map(({ currency, total }) => ( + + Fixed cost in range: {formatFee(total, currency)} + + )) + )} + + Usage-billed spend in range, all providers: ${formatNumberWithCommas(meteredSpend, 2)} + +
+

{ATTRIBUTION_RULE}

+
+ )} + + + ); +}; + +export default SubscriptionCostCard; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCosts.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCosts.integration.test.tsx new file mode 100644 index 00000000000..f057e98dd70 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCosts.integration.test.tsx @@ -0,0 +1,69 @@ +import { fireEvent, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders, testQueryClient } from "@/../tests/test-utils"; +import { subscriptionAccountCreateCall, subscriptionUsageCall } from "@/components/networking"; +import SubscriptionCosts from "./SubscriptionCosts"; +import type { SubscriptionUsageResponse } from "./types"; + +vi.mock("@/components/networking", () => ({ + formatDate: (date: Date) => date.toISOString().slice(0, 10), + subscriptionUsageCall: vi.fn(), + subscriptionAccountCreateCall: vi.fn().mockResolvedValue({}), + subscriptionAccountUpdateCall: vi.fn().mockResolvedValue({}), + subscriptionAccountDeleteCall: vi.fn().mockResolvedValue({}), +})); +vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), fromError: vi.fn() } })); + +const untracked: SubscriptionUsageResponse = { + start_date: "2026-10-01", + end_date: "2026-10-07", + subscription_providers: ["chatgpt"], + accounts: [ + { + custom_llm_provider: "chatgpt", + account_id: "acct_123", + deployments: ["gpt-5.6-sol", "gpt-5.6-terra"], + fee: null, + billing_periods: [], + fixed_cost: null, + }, + ], +}; + +const fill = (label: string, value: string) => fireEvent.change(screen.getByLabelText(label), { target: { value } }); + +describe("SubscriptionCosts", () => { + beforeEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + vi.mocked(subscriptionUsageCall).mockResolvedValue(untracked); + }); + + it("saves a fee set from the card on that card's account and reloads the card", async () => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(await screen.findByRole("button", { name: "Set monthly fee" })); + expect(screen.getByText("chatgpt account acct_123")).toBeInTheDocument(); + fill("Monthly fee", "25"); + fill("Currency", "USD"); + fill("Billing period start", "2026-10-01"); + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => expect(subscriptionUsageCall).toHaveBeenCalledTimes(2)); + expect(subscriptionAccountCreateCall).toHaveBeenCalledWith( + "sk-admin", + expect.objectContaining({ custom_llm_provider: "chatgpt", account_id: "acct_123", monthly_fee: 25 }), + ); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCosts.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCosts.tsx new file mode 100644 index 00000000000..e326fa1ff01 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionCosts.tsx @@ -0,0 +1,38 @@ +import React from "react"; +import { useSubscriptionUsage, type SubscriptionUsageOptions } from "../../hooks/useSubscriptionUsage"; +import SubscriptionCostCard from "./SubscriptionCostCard"; +import SubscriptionFeeModal from "./SubscriptionFeeModal"; +import { useSubscriptionFeeActions } from "./useSubscriptionFeeActions"; + +interface SubscriptionCostsProps extends SubscriptionUsageOptions { + meteredSpend: number; + canEdit: boolean; +} + +const SubscriptionCosts: React.FC = ({ meteredSpend, canEdit, ...usageOptions }) => { + const usage = useSubscriptionUsage(usageOptions); + const feeActions = useSubscriptionFeeActions({ accessToken: usageOptions.accessToken, refetch: usage.refetch }); + return ( + <> + + + + ); +}; + +export default SubscriptionCosts; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionFeeModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionFeeModal.test.tsx new file mode 100644 index 00000000000..a920f6ec883 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionFeeModal.test.tsx @@ -0,0 +1,83 @@ +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import SubscriptionFeeModal from "./SubscriptionFeeModal"; +import type { SubscriptionAccountUsage } from "./types"; + +const account: SubscriptionAccountUsage = { + custom_llm_provider: "chatgpt", + account_id: "acct_123", + deployments: ["gpt-5.6-terra"], + fee: null, + billing_periods: [], + fixed_cost: null, +}; + +const fill = (label: string, value: string) => fireEvent.change(screen.getByLabelText(label), { target: { value } }); + +describe("SubscriptionFeeModal", () => { + it("submits normalized values for a new fee", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + render(); + + expect(screen.getByText("Set monthly fee")).toBeInTheDocument(); + expect(screen.getByText("chatgpt account acct_123")).toBeInTheDocument(); + fill("Monthly fee", "25"); + fill("Currency", "usd"); + fill("Billing period start", "2026-10-01"); + fill("Label (optional)", "Team workspace"); + await userEvent.setup().click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); + const normalized = { + monthly_fee: 25, + currency: "USD", + billing_period_start: "2026-10-01", + label: "Team workspace", + }; + expect(onSubmit).toHaveBeenCalledWith(normalized); + }); + + it("rejects a zero fee and a two-letter currency", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + const user = userEvent.setup(); + render(); + + fill("Monthly fee", "0"); + fill("Billing period start", "2026-10-01"); + await user.click(screen.getByRole("button", { name: "Save" })); + expect(await screen.findByText("Monthly fee must be greater than 0")).toBeInTheDocument(); + + fill("Monthly fee", "25"); + fill("Currency", "US"); + await user.click(screen.getByRole("button", { name: "Save" })); + expect(await screen.findByText("Currency must be a three-letter code")).toBeInTheDocument(); + expect(onSubmit).not.toHaveBeenCalled(); + }); + + it("prefills the existing fee when editing", () => { + const existing: SubscriptionAccountUsage = { + ...account, + fee: { + subscription_account_id: "fee-1", + custom_llm_provider: "chatgpt", + account_id: "acct_123", + label: "Team workspace", + monthly_fee: 25, + currency: "EUR", + billing_period_start: "2026-09-15", + created_at: "2026-10-06T00:00:00Z", + updated_at: "2026-10-06T00:00:00Z", + created_by: "admin", + updated_by: "admin", + }, + }; + render(); + + expect(screen.getByText("Edit monthly fee")).toBeInTheDocument(); + expect(screen.getByLabelText("Monthly fee")).toHaveValue(25); + expect(screen.getByLabelText("Currency")).toHaveValue("EUR"); + expect(screen.getByLabelText("Billing period start")).toHaveValue("2026-09-15"); + expect(screen.getByLabelText("Label (optional)")).toHaveValue("Team workspace"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionFeeModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionFeeModal.tsx new file mode 100644 index 00000000000..1c211a5280c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/SubscriptionFeeModal.tsx @@ -0,0 +1,114 @@ +"use client"; + +import React, { useEffect } from "react"; +import { z } from "zod"; +import { formatDate } from "@/components/networking"; +import { FormField } from "@/components/shared/form/FormField"; +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { FieldGroup } from "@/components/ui/field"; +import { Input } from "@/components/ui/input"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import type { SubscriptionAccountUsage, SubscriptionFeeFormValues } from "./types"; + +const feeShape = { + monthly_fee: z + .string() + .min(1, "Enter a monthly fee") + .refine((value) => Number(value) > 0, "Monthly fee must be greater than 0"), + currency: z + .string() + .trim() + .toUpperCase() + .regex(/^[A-Z]{3}$/, "Currency must be a three-letter code"), + billing_period_start: z.string().regex(/^\d{4}-\d{2}-\d{2}$/, "Pick a billing period start date"), + label: z.string().trim(), +}; + +const feeSchema = z.object(feeShape); + +type FeeFormInput = z.input; + +const formInputFor = (account: SubscriptionAccountUsage | null): FeeFormInput => ({ + monthly_fee: account?.fee ? String(account.fee.monthly_fee) : "", + currency: account?.fee?.currency ?? "USD", + billing_period_start: account?.fee?.billing_period_start ?? formatDate(new Date()), + label: account?.fee?.label ?? "", +}); + +interface SubscriptionFeeModalProps { + visible: boolean; + account: SubscriptionAccountUsage | null; + onCancel: () => void; + onSubmit: (values: SubscriptionFeeFormValues) => Promise; +} + +const SubscriptionFeeModal: React.FC = ({ visible, account, onCancel, onSubmit }) => { + const form = useZodForm(feeSchema, { defaultValues: formInputFor(account) }); + const { reset } = form; + + useEffect(() => { + if (visible) reset(formInputFor(account)); + }, [visible, account, reset]); + + const handleFinish = async (values: z.output) => { + const submitted: SubscriptionFeeFormValues = { + monthly_fee: Number(values.monthly_fee), + currency: values.currency, + billing_period_start: values.billing_period_start, + label: values.label === "" ? null : values.label, + }; + await onSubmit(submitted); + }; + + const handleCancel = () => { + reset(formInputFor(null)); + onCancel(); + }; + + return ( + !open && handleCancel()}> + + + {account?.fee ? "Edit monthly fee" : "Set monthly fee"} + + {account ? `${account.custom_llm_provider} account ${account.account_id ?? ""}` : ""} + + +
+ + + {({ ref, ...field }) => } + + + {({ ref, ...field }) => } + + + {({ ref, ...field }) => } + + + {({ ref, ...field }) => } + + + + + + +
+ + + ); +}; + +export default SubscriptionFeeModal; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/formatting.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/formatting.ts new file mode 100644 index 00000000000..b1d4e0a48ef --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/formatting.ts @@ -0,0 +1,34 @@ +import type { BillingPeriod } from "./types"; + +export const ATTRIBUTION_RULE = + "A subscription fee is counted once for each monthly billing period whose start date falls inside the selected dates. Fixed costs sit beside metered spend and are never added to it, to key or team spend, or to budgets."; + +export const formatFee = (amount: number, currency: string): string => `${amount.toFixed(2)} ${currency}`; + +const parseLocalDay = (day: string): Date => { + const [year, month, date] = day.split("-").map(Number); + return new Date(year, month - 1, date); +}; + +export const formatBillingPeriod = ({ start, end }: BillingPeriod): string => { + const startDate = parseLocalDay(start); + const endDate = parseLocalDay(end); + const startLabel = startDate.toLocaleDateString("en-US", { + month: "short", + day: "numeric", + year: startDate.getFullYear() !== endDate.getFullYear() ? "numeric" : undefined, + }); + const endLabel = endDate.toLocaleDateString("en-US", { month: "short", day: "numeric", year: "numeric" }); + return `${startLabel} - ${endLabel}`; +}; + +export const fixedCostByCurrency = ( + accounts: readonly { fee: { currency: string } | null; fixed_cost: number | null }[], +): readonly { currency: string; total: number }[] => { + const totals = new Map(); + accounts.forEach((account) => { + if (account.fee === null || account.fixed_cost === null) return; + totals.set(account.fee.currency, (totals.get(account.fee.currency) ?? 0) + account.fixed_cost); + }); + return Array.from(totals, ([currency, total]) => ({ currency, total })); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/types.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/types.ts new file mode 100644 index 00000000000..6af5e0cdf3b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/types.ts @@ -0,0 +1,74 @@ +import type { DailyData } from "@/components/UsagePage/types"; + +export interface SubscriptionFee { + subscription_account_id: string; + custom_llm_provider: string; + account_id: string; + label: string | null; + monthly_fee: number; + currency: string; + billing_period_start: string; + created_at: string; + updated_at: string; + created_by: string; + updated_by: string; +} + +export interface BillingPeriod { + start: string; + end: string; +} + +export interface SubscriptionAccountUsage { + custom_llm_provider: string; + account_id: string | null; + deployments: string[]; + fee: SubscriptionFee | null; + billing_periods: BillingPeriod[]; + fixed_cost: number | null; +} + +export interface SubscriptionUsageResponse { + subscription_providers: string[]; + accounts: SubscriptionAccountUsage[]; +} + +export interface SubscriptionAccountCreateRequest { + custom_llm_provider: string; + account_id: string; + monthly_fee: number; + currency: string; + billing_period_start: string; + label: string | null; +} + +export interface SubscriptionAccountUpdateRequest { + subscription_account_id: string; + monthly_fee?: number; + currency?: string; + billing_period_start?: string; + label?: string | null; +} + +export interface SubscriptionAccountDeleteRequest { + subscription_account_id: string; +} + +export interface SubscriptionFeeFormValues { + monthly_fee: number; + currency: string; + billing_period_start: string; + label: string | null; +} + +export const isSubscriptionCovered = (metadata: object | undefined): boolean => + metadata !== undefined && "subscription_covered" in metadata && metadata.subscription_covered === true; + +export const subscriptionCoveredProviders = (results: readonly DailyData[]): ReadonlySet => + new Set( + results.flatMap((day) => + Object.entries(day.breakdown.providers ?? {}) + .filter(([, entry]) => isSubscriptionCovered(entry.metadata)) + .map(([provider]) => provider), + ), + ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/useSubscriptionFeeActions.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/useSubscriptionFeeActions.test.tsx new file mode 100644 index 00000000000..9406e07b497 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/useSubscriptionFeeActions.test.tsx @@ -0,0 +1,79 @@ +import { act, renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { subscriptionAccountCreateCall, subscriptionAccountUpdateCall } from "@/components/networking"; +import type { SubscriptionAccountUsage, SubscriptionFeeFormValues } from "./types"; +import { useSubscriptionFeeActions } from "./useSubscriptionFeeActions"; + +vi.mock("@/components/networking", () => ({ + subscriptionAccountCreateCall: vi.fn().mockResolvedValue({}), + subscriptionAccountUpdateCall: vi.fn().mockResolvedValue({}), + subscriptionAccountDeleteCall: vi.fn(), +})); +vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), fromError: vi.fn() } })); + +const stored: SubscriptionFeeFormValues = { + monthly_fee: 25, + currency: "USD", + billing_period_start: "2026-10-01", + label: "Team plan (test value)", +}; + +const tracked: SubscriptionAccountUsage = { + custom_llm_provider: "chatgpt", + account_id: "acct_123", + deployments: ["gpt-5.6-terra"], + fee: { + subscription_account_id: "sub-1", + custom_llm_provider: "chatgpt", + account_id: "acct_123", + ...stored, + created_at: "2026-10-01T00:00:00Z", + updated_at: "2026-10-01T00:00:00Z", + created_by: "admin", + updated_by: "admin", + }, + billing_periods: [{ start: "2026-10-01", end: "2026-10-31" }], + fixed_cost: 25, +}; + +const submit = async (account: SubscriptionAccountUsage, values: SubscriptionFeeFormValues) => { + const refetch = vi.fn().mockResolvedValue(undefined); + const { result } = renderHook(() => useSubscriptionFeeActions({ accessToken: "sk-admin", refetch })); + act(() => result.current.openFeeModal(account)); + await act(() => result.current.submitFee(values)); + return { refetch, result }; +}; + +describe("useSubscriptionFeeActions", () => { + beforeEach(() => vi.clearAllMocks()); + + it("sends only the fields the admin changed when editing a fee", async () => { + const { refetch, result } = await submit(tracked, { ...stored, monthly_fee: 30 }); + + expect(subscriptionAccountUpdateCall).toHaveBeenCalledWith("sk-admin", { + subscription_account_id: "sub-1", + monthly_fee: 30, + }); + expect(refetch).toHaveBeenCalledTimes(1); + expect(result.current.feeModalAccount).toBeNull(); + }); + + it("clears a label with an explicit null and nothing else", async () => { + await submit(tracked, { ...stored, label: null }); + + expect(subscriptionAccountUpdateCall).toHaveBeenCalledWith("sk-admin", { + subscription_account_id: "sub-1", + label: null, + }); + }); + + it("creates a fee with every field for an untracked account", async () => { + const untracked: SubscriptionAccountUsage = { ...tracked, fee: null, billing_periods: [], fixed_cost: null }; + const createRequest = { custom_llm_provider: "chatgpt", account_id: "acct_123", ...stored }; + + await submit(untracked, stored); + + expect(subscriptionAccountCreateCall).toHaveBeenCalledWith("sk-admin", createRequest); + expect(subscriptionAccountUpdateCall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/useSubscriptionFeeActions.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/useSubscriptionFeeActions.ts new file mode 100644 index 00000000000..dd2b2f01a30 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/SubscriptionCost/useSubscriptionFeeActions.ts @@ -0,0 +1,90 @@ +import { useCallback, useState } from "react"; +import { + subscriptionAccountCreateCall, + subscriptionAccountDeleteCall, + subscriptionAccountUpdateCall, +} from "@/components/networking"; +import { toast } from "@/lib/toast"; +import type { + SubscriptionAccountUpdateRequest, + SubscriptionAccountUsage, + SubscriptionFee, + SubscriptionFeeFormValues, +} from "./types"; + +interface Options { + accessToken: string | null; + refetch: () => Promise; +} + +interface Result { + feeModalAccount: SubscriptionAccountUsage | null; + openFeeModal: (account: SubscriptionAccountUsage) => void; + closeFeeModal: () => void; + submitFee: (values: SubscriptionFeeFormValues) => Promise; + removeFee: (account: SubscriptionAccountUsage) => Promise; +} + +const changedFields = (fee: SubscriptionFee, values: SubscriptionFeeFormValues): SubscriptionAccountUpdateRequest => ({ + subscription_account_id: fee.subscription_account_id, + ...(values.monthly_fee !== fee.monthly_fee && { monthly_fee: values.monthly_fee }), + ...(values.currency !== fee.currency && { currency: values.currency }), + ...(values.billing_period_start !== fee.billing_period_start && { + billing_period_start: values.billing_period_start, + }), + ...(values.label !== fee.label && { label: values.label }), +}); + +const saveFee = ( + accessToken: string, + account: SubscriptionAccountUsage, + values: SubscriptionFeeFormValues, +): Promise | null => { + if (account.fee) { + return subscriptionAccountUpdateCall(accessToken, changedFields(account.fee, values)); + } + if (account.account_id === null) return null; + return subscriptionAccountCreateCall(accessToken, { + custom_llm_provider: account.custom_llm_provider, + account_id: account.account_id, + ...values, + }); +}; + +export function useSubscriptionFeeActions({ accessToken, refetch }: Options): Result { + const [feeModalAccount, setFeeModalAccount] = useState(null); + const closeFeeModal = useCallback(() => setFeeModalAccount(null), []); + + const submitFee = useCallback( + async (values: SubscriptionFeeFormValues) => { + if (!accessToken || feeModalAccount === null) return; + const request = saveFee(accessToken, feeModalAccount, values); + if (request === null) return; + try { + await request; + toast.success("Monthly fee saved"); + setFeeModalAccount(null); + await refetch(); + } catch (error) { + toast.fromError(error); + } + }, + [accessToken, feeModalAccount, refetch], + ); + + const removeFee = useCallback( + async (account: SubscriptionAccountUsage) => { + if (!accessToken || account.fee === null) return; + try { + await subscriptionAccountDeleteCall(accessToken, account.fee.subscription_account_id); + toast.success("Monthly fee removed"); + await refetch(); + } catch (error) { + toast.fromError(error); + } + }, + [accessToken, refetch], + ); + + return { feeModalAccount, openFeeModal: setFeeModalAccount, closeFeeModal, submitFee, removeFee }; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx index f46a4569369..9d39e9c080c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx @@ -8,7 +8,7 @@ import type { ComponentProps } from "react"; import userEvent from "@testing-library/user-event"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import type { UrlUpdateEvent } from "nuqs/adapters/testing"; -import { chooseSelectOption, renderWithProviders } from "@/../tests/test-utils"; +import { chooseSelectOption, renderWithProviders, testQueryClient } from "@/../tests/test-utils"; import { toast } from "@/lib/toast"; import { processActivityData } from "@/components/activity_metrics"; import type { Organization } from "@/components/networking"; @@ -74,6 +74,11 @@ beforeAll(() => { // Mock the networking module vi.mock("@/components/networking", () => ({ + formatDate: (date: Date) => date.toISOString().slice(0, 10), + subscriptionUsageCall: vi.fn(), + subscriptionAccountCreateCall: vi.fn(), + subscriptionAccountUpdateCall: vi.fn(), + subscriptionAccountDeleteCall: vi.fn(), dailyActivityAggregatedCall: vi.fn(), dailyActivityKeyPageCall: vi.fn(), userDailyActivityUserPageCall: vi.fn(), @@ -142,6 +147,14 @@ vi.mock("./EntityUsage/SpendByProvider", () => ({ default: () =>
Spend By Provider
, })); +vi.mock("./SubscriptionCost/SubscriptionCostCard", () => ({ + default: ({ usage, canEdit }: { usage: unknown; canEdit: boolean }) => ( +
+ Subscription Cost +
+ ), +})); + vi.mock("./EndpointUsage/EndpointUsage", () => ({ default: () =>
Endpoint Usage
, })); @@ -278,6 +291,7 @@ describe("UsagePage", () => { const mockTagListCall = vi.mocked(networking.tagListCall); const mockGatewayDailyActivityCall = vi.mocked(networking.gatewayDailyActivityCall); const mockRequestErrorActivityCall = vi.mocked(networking.requestErrorActivityCall); + const mockSubscriptionUsageCall = vi.mocked(networking.subscriptionUsageCall); const mockUseAgents = vi.mocked(useAgents); const mockUseAuthorized = vi.mocked(useAuthorized); const mockUseCurrentUser = vi.mocked(useCurrentUser); @@ -496,7 +510,24 @@ describe("UsagePage", () => { organizations: [], }; + const mockSubscriptionUsage = { + subscription_providers: ["chatgpt"], + accounts: [ + { + custom_llm_provider: "chatgpt", + account_id: "acct_123", + deployments: ["gpt-5.6-terra", "gpt-5.6-sol"], + fee: null, + billing_periods: [], + fixed_cost: null, + }, + ], + }; + beforeEach(() => { + testQueryClient.clear(); + mockSubscriptionUsageCall.mockReset(); + mockSubscriptionUsageCall.mockResolvedValue(mockSubscriptionUsage); mockClientUsagePanel.mockClear(); vi.mocked(processActivityData).mockReset(); vi.mocked(processActivityData).mockImplementation((_data, key) => ({ @@ -578,6 +609,72 @@ describe("UsagePage", () => { } as any); }); + it("hands the subscription usage to the fixed cost card for an admin and never fetches it otherwise", async () => { + const { unmount } = renderWithProviders(); + await waitFor(() => { + expect(mockSubscriptionUsageCall).toHaveBeenCalledWith("test-token", expect.any(Date), expect.any(Date)); + }); + await waitFor(() => { + expect(screen.getByTestId("subscription-cost-card")).toHaveAttribute( + "data-usage", + JSON.stringify(mockSubscriptionUsage), + ); + }); + unmount(); + + for (const session of [nonAdminSession, { ...nonAdminSession, userRole: "org_admin" }]) { + mockSubscriptionUsageCall.mockClear(); + mockUserDailyActivityAggregatedCall.mockClear(); + mockUseAuthorized.mockReturnValue(session); + const rendered = renderWithProviders(); + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + expect(mockSubscriptionUsageCall).not.toHaveBeenCalled(); + expect(screen.getByTestId("subscription-cost-card")).toHaveAttribute("data-usage", "null"); + rendered.unmount(); + } + }); + + it("lets a proxy admin edit subscription fees but shows an admin viewer the same card read-only", async () => { + const adminRender = renderWithProviders(); + await waitFor(() => { + expect(screen.getByTestId("subscription-cost-card")).toHaveAttribute( + "data-usage", + JSON.stringify(mockSubscriptionUsage), + ); + }); + expect(screen.getByTestId("subscription-cost-card")).toHaveAttribute("data-can-edit", "true"); + adminRender.unmount(); + + const adminViewerSession = { + ...nonAdminSession, + userRole: "Admin", + userRoleLabel: "Admin Viewer", + isViewOnly: true, + }; + mockUseAuthorized.mockReturnValue(adminViewerSession); + renderWithProviders(); + await waitFor(() => { + expect(screen.getByTestId("subscription-cost-card")).toHaveAttribute( + "data-usage", + JSON.stringify(mockSubscriptionUsage), + ); + }); + expect(screen.getByTestId("subscription-cost-card")).toHaveAttribute("data-can-edit", "false"); + }); + + it("never sets an account-wide subscription fee beside one user's spend when an admin filters by user", async () => { + renderWithProviders(, { searchParams: "?user=user-001" }); + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledWith( + expect.objectContaining({ entityIds: ["user-001"] }), + ); + }); + expect(mockSubscriptionUsageCall).not.toHaveBeenCalled(); + expect(screen.getByTestId("subscription-cost-card")).toHaveAttribute("data-usage", "null"); + }); + it("should render and fetch usage data on mount", async () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index d5ec3ef7675..7c891754195 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -26,7 +26,7 @@ import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrgan import { toast } from "@/lib/toast"; import { hasCapability } from "@/utils/capabilities"; import { copyToClipboard } from "@/utils/dataUtils"; -import { hasProxyWideSpendView, internalUserRoles } from "@/utils/roles"; +import { hasProxyWideSpendView, internalUserRoles, isProxyAdminRole } from "@/utils/roles"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; import CloudZeroExportModal from "@/components/cloudzero_export_modal"; import UserDropdown from "@/components/common_components/UserDropdown"; @@ -92,7 +92,7 @@ import { usageViewFromParams, usageViewPatch, } from "../usageUrlState"; -import { overviewTotals, rollUpBreakdown } from "./overview/overviewData"; +import { overviewTotals, providerSpendRows } from "./overview/overviewData"; import TopAgents from "./overview/TopAgents"; import { parseUsageClient, topAgents, type AgentRow } from "./overview/agentCatalog"; import ClientUsagePanel from "./ClientUsagePanel"; @@ -100,6 +100,7 @@ import { resolveClientUsageScope } from "../hooks/clientUsageScope"; import { useUsagePanelContext } from "../hooks/useUsagePanelContext"; import ClientUsageSelect from "./ClientUsageSelect"; import SpendByProvider from "./EntityUsage/SpendByProvider"; +import SubscriptionCosts from "./SubscriptionCost/SubscriptionCosts"; import { useTagSummary } from "@/app/(dashboard)/hooks/tags/useTagSummary"; import { Panel } from "./overview/Primitives"; import UsageOverview from "./overview/UsageOverview"; @@ -112,7 +113,7 @@ interface UsagePageProps { } const UsagePage: React.FC = ({ teams, organizations, teamsLoading = false, defaultTimeRange }) => { - const { accessToken, userRole, userId: userID, premiumUser, isLoading: authorizationLoading } = useAuthorized(); + const { accessToken, userRole, userId: userID, premiumUser, isViewOnly, isLoading: authLoading } = useAuthorized(); const [gatewayActivityData, setGatewayActivityData] = useState(null); const defaultRange = useMemo(() => defaultUsageDateRange(defaultTimeRange), [defaultTimeRange]); @@ -187,7 +188,7 @@ const UsagePage: React.FC = ({ teams, organizations, teamsLoadin }; const { isLoading: organizationsLoading } = useOrganizations(); - const authorizationSettled = !authorizationLoading && userRole !== null; + const authorizationSettled = !authLoading && userRole !== null; const accessSettled = authorizationSettled && !organizationsLoading && !teamsLoading; useEffect(() => { if (!accessSettled) return; @@ -374,10 +375,7 @@ const UsagePage: React.FC = ({ teams, organizations, teamsLoadin () => overviewTotals(userSpendData.metadata, gatewayActivity), [userSpendData.metadata, gatewayActivity], ); - const providerSpend = useMemo( - () => rollUpBreakdown(userSpendData.results, "providers").map(({ key, ...row }) => ({ provider: key, ...row })), - [userSpendData.results], - ); + const providerSpend = useMemo(() => providerSpendRows(userSpendData.results), [userSpendData.results]); const { data: tagSummary, isLoading: tagSummaryLoading, @@ -546,6 +544,16 @@ const UsagePage: React.FC = ({ teams, organizations, teamsLoadin loading={loading} requestCountsPending={requestCountsPending} budget={currentUser?.max_budget ?? null} + subscriptionCosts={ + + } topKeys={ + {subscriptionCosts} + {topAgents} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.test.ts index f00d53ea1bc..1dfdfd1da78 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.test.ts @@ -7,6 +7,7 @@ import { labelForDate, modelColorMap, OTHER_COLOR, + providerSpendRows, rollUpBreakdown, seriesBy, } from "./overviewData"; @@ -88,6 +89,34 @@ describe("bucketSeries", () => { }); }); +const providerDay = (date: string, providers: Record>): DailyData => + ({ + ...day(date, 0, {}), + breakdown: { + ...day(date, 0, {}).breakdown, + providers: Object.fromEntries( + Object.entries(providers).map(([name, metadata]) => [ + name, + { metrics: metrics(0), metadata, api_key_breakdown: {} }, + ]), + ), + }, + }) as unknown as DailyData; + +describe("providerSpendRows", () => { + it("marks a provider covered by a subscription on any day of the range and leaves metered ones unmarked", () => { + const rows = providerSpendRows([ + providerDay("2026-10-01", { chatgpt: {}, openai: {} }), + providerDay("2026-10-02", { chatgpt: { subscription_covered: true }, openai: { subscription_covered: false } }), + ]); + expect(Object.fromEntries(rows.map((row) => [row.provider, row.subscription_covered]))).toEqual({ + chatgpt: true, + openai: false, + }); + expect(rows.find((row) => row.provider === "chatgpt")).toEqual(expect.objectContaining({ requests: 2 })); + }); +}); + describe("rollUpBreakdown", () => { it("sums a dimension across days without changing the input", () => { const days = [day("2026-10-01", 3, { a: 3 }), day("2026-10-02", 2, { a: 2 })]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.ts index c103c2d7c60..9d9e8f9fa32 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.ts @@ -2,6 +2,7 @@ import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; import { stackedUsageColor } from "@/components/shared/charts"; import { EMPTY_DAILY_ACTIVITY_METADATA, type DailyActivityMetadata } from "@/components/UsagePage/dailyActivityApi"; import type { GatewayActivity } from "../gatewayActivity"; +import { subscriptionCoveredProviders } from "../SubscriptionCost/types"; export type BreakdownDimension = "models" | "model_groups" | "providers"; @@ -40,6 +41,15 @@ export const rollUpBreakdown = (results: readonly DailyData[], dimension: Breakd return [...totals.values()].sort((a, b) => b.spend - a.spend); }; +export const providerSpendRows = (results: readonly DailyData[]) => { + const covered = subscriptionCoveredProviders(results); + return rollUpBreakdown(results, "providers").map(({ key, ...row }) => ({ + provider: key, + ...row, + subscription_covered: covered.has(key), + })); +}; + export const modelColorMap = (results: readonly DailyData[]): ReadonlyMap => { const groups = rollUpBreakdown(results, "model_groups"); const models = rollUpBreakdown(results, "models"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useSubscriptionUsage.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useSubscriptionUsage.ts new file mode 100644 index 00000000000..afda5ec57b8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useSubscriptionUsage.ts @@ -0,0 +1,47 @@ +import { keepPreviousData, queryOptions, skipToken, useQuery } from "@tanstack/react-query"; +import { formatDate, subscriptionUsageCall } from "@/components/networking"; +import type { SubscriptionUsageResponse } from "../components/SubscriptionCost/types"; + +export interface SubscriptionUsageOptions { + accessToken: string | null; + startTime: Date | null; + endTime: Date | null; + enabled: boolean; +} + +interface Result { + data: SubscriptionUsageResponse | null; + loading: boolean; + failed: boolean; + refetch: () => Promise; +} + +interface UsageRequest { + accessToken: string; + startTime: Date; + endTime: Date; +} + +const subscriptionUsageQuery = (request: UsageRequest | null) => { + const queryKey = [ + "subscriptionUsage", + request ? formatDate(request.startTime) : null, + request ? formatDate(request.endTime) : null, + ]; + const queryFn = request + ? () => subscriptionUsageCall(request.accessToken, request.startTime, request.endTime) + : skipToken; + return queryOptions({ queryKey, queryFn, placeholderData: keepPreviousData }); +}; + +export function useSubscriptionUsage({ accessToken, startTime, endTime, enabled }: SubscriptionUsageOptions): Result { + const range = startTime && endTime ? { startTime, endTime } : null; + const request = enabled && accessToken && range ? { accessToken, ...range } : null; + const query = useQuery(subscriptionUsageQuery(request)); + return { + data: request && query.data ? query.data : null, + loading: request !== null && (query.isPending || query.isPlaceholderData), + failed: query.isError, + refetch: query.refetch, + }; +} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 99a9dc7a707..c13052be245 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -119,6 +119,13 @@ import { MCP_TOOLS_PREVIEW_FORBIDDEN_MESSAGE } from "./mcp_tools/constants"; import type { ComplexityRouterConfigPayload } from "./add_model/build_complexity_router_config"; import type { AutoRouterPresetsResponse } from "@/lib/autorouter_presets"; import type { VectorStoreIndex } from "@/app/(dashboard)/vector-stores/_components/IndexesTab"; +import type { + SubscriptionAccountCreateRequest, + SubscriptionAccountDeleteRequest, + SubscriptionAccountUpdateRequest, + SubscriptionFee, + SubscriptionUsageResponse, +} from "@/app/(dashboard)/usage/_components/components/SubscriptionCost/types"; import type { RoutingDecision } from "./logs/detail/RoutingDecisionCard"; import type { SpanDetail, @@ -1395,6 +1402,34 @@ export const dailyActivityAggregatedCall = ( query: dailyActivityQuery(entity, req, req.apiKeyLimit === undefined ? {} : { api_key_limit: req.apiKeyLimit }), }); +export const subscriptionUsageCall = ( + accessToken: string, + startTime: Date, + endTime: Date, +): Promise => + apiClient.get("/subscription_accounts/usage", { + accessToken, + query: { start_date: formatDate(startTime), end_date: formatDate(endTime) }, + }); + +export const subscriptionAccountCreateCall = ( + accessToken: string, + body: SubscriptionAccountCreateRequest, +): Promise => apiClient.post("/subscription_accounts/new", { accessToken, body }); + +export const subscriptionAccountUpdateCall = ( + accessToken: string, + body: SubscriptionAccountUpdateRequest, +): Promise => apiClient.post("/subscription_accounts/update", { accessToken, body }); + +export const subscriptionAccountDeleteCall = ( + accessToken: string, + subscriptionAccountId: string, +): Promise => { + const body: SubscriptionAccountDeleteRequest = { subscription_account_id: subscriptionAccountId }; + return apiClient.post("/subscription_accounts/delete", { accessToken, body }); +}; + export const dailyActivityKeyPageCall = ( entity: DailyActivityEntity, req: DailyActivityRequest, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7d9e9839a5f..f53e57c2f18 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16080,6 +16080,90 @@ export interface paths { patch?: never; trace?: never; }; + "/subscription_accounts/delete": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Delete Subscription Account + * @description Stop tracking a subscription account's fee; its requests stay marked as subscription-covered. + */ + post: operations["delete_subscription_account_subscription_accounts_delete_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/subscription_accounts/new": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * New Subscription Account + * @description Record the monthly fee of one subscription-backed provider account. + */ + post: operations["new_subscription_account_subscription_accounts_new_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/subscription_accounts/update": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Update Subscription Account + * @description Change the fee, currency, billing period start, or label of a subscription account. + * Only the fields named in the request are written, so two admins changing different + * fields at the same time both land. + */ + post: operations["update_subscription_account_subscription_accounts_update_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/subscription_accounts/usage": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Subscription Usage + * @description Subscription-backed accounts behind the proxy's deployments or carrying a stored fee, each with + * its configured monthly fee and the fixed cost attributed to the date range: the fee once per billing + * period whose start date falls inside the range, however many deployments share the account. + */ + get: operations["get_subscription_usage_subscription_accounts_usage_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/systemone": { parameters: { query?: never; @@ -27402,6 +27486,19 @@ export interface components { /** Entities */ entities: components["schemas"]["BedrockChecksSensitiveInformationEntityItem"][]; }; + /** BillingPeriod */ + BillingPeriod: { + /** + * End + * Format: date + */ + end: string; + /** + * Start + * Format: date + */ + start: string; + }; /** BlockKeyRequest */ BlockKeyRequest: { /** Key */ @@ -31882,6 +31979,11 @@ export interface components { */ type: string; }; + /** DeleteSubscriptionAccountRequest */ + DeleteSubscriptionAccountRequest: { + /** Subscription Account Id */ + subscription_account_id: string; + }; /** DeleteTeamRequest */ DeleteTeamRequest: { /** Team Ids */ @@ -36017,6 +36119,37 @@ export interface components { */ user: string | null; }; + /** LiteLLM_SubscriptionAccountTable */ + LiteLLM_SubscriptionAccountTable: { + /** Account Id */ + account_id: string; + /** Billing Period Start */ + billing_period_start: string; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Created By */ + created_by: string; + /** Currency */ + currency: string; + /** Custom Llm Provider */ + custom_llm_provider: string; + /** Label */ + label?: string | null; + /** Monthly Fee */ + monthly_fee: number; + /** Subscription Account Id */ + subscription_account_id: string; + /** + * Updated At + * Format: date-time + */ + updated_at: string; + /** Updated By */ + updated_by: string; + }; /** LiteLLM_TeamMembership */ LiteLLM_TeamMembership: { /** Budget Id */ @@ -46225,6 +46358,55 @@ export interface components { /** Window Start */ window_start: number; }; + /** SubscriptionAccountFeeRequest */ + SubscriptionAccountFeeRequest: { + /** Account Id */ + account_id: string; + /** + * Billing Period Start + * Format: date + */ + billing_period_start: string; + /** Currency */ + currency: string; + /** Custom Llm Provider */ + custom_llm_provider: string; + /** Label */ + label?: string | null; + /** Monthly Fee */ + monthly_fee: number; + }; + /** SubscriptionAccountUsage */ + SubscriptionAccountUsage: { + /** Account Id */ + account_id: string | null; + /** Billing Periods */ + billing_periods: components["schemas"]["BillingPeriod"][]; + /** Custom Llm Provider */ + custom_llm_provider: string; + /** Deployments */ + deployments: string[]; + fee: components["schemas"]["LiteLLM_SubscriptionAccountTable"] | null; + /** Fixed Cost */ + fixed_cost: number | null; + }; + /** SubscriptionUsageResponse */ + SubscriptionUsageResponse: { + /** Accounts */ + accounts: components["schemas"]["SubscriptionAccountUsage"][]; + /** + * End Date + * Format: date + */ + end_date: string; + /** + * Start Date + * Format: date + */ + start_date: string; + /** Subscription Providers */ + subscription_providers: string[]; + }; /** * SuccessfulKeyUpdate * @description Successfully updated key with its updated information @@ -48898,6 +49080,19 @@ export interface components { UpdateSearchToolRequest: { search_tool: components["schemas"]["SearchTool"]; }; + /** UpdateSubscriptionAccountFeeRequest */ + UpdateSubscriptionAccountFeeRequest: { + /** Billing Period Start */ + billing_period_start?: string | null; + /** Currency */ + currency?: string | null; + /** Label */ + label?: string | null; + /** Monthly Fee */ + monthly_fee?: number | null; + /** Subscription Account Id */ + subscription_account_id: string; + }; /** * UpdateTeamMemberPermissionsRequest * @description Request to update the team member permissions for a team @@ -72988,6 +73183,139 @@ export interface operations { }; }; }; + delete_subscription_account_subscription_accounts_delete_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["DeleteSubscriptionAccountRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SubscriptionAccountTable"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + new_subscription_account_subscription_accounts_new_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["SubscriptionAccountFeeRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SubscriptionAccountTable"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + update_subscription_account_subscription_accounts_update_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["UpdateSubscriptionAccountFeeRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SubscriptionAccountTable"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_subscription_usage_subscription_accounts_usage_get: { + parameters: { + query: { + /** @description First day of the Usage page filter, YYYY-MM-DD */ + start_date: string; + /** @description Last day of the Usage page filter, YYYY-MM-DD */ + end_date: string; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["SubscriptionUsageResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; systemone_systemone_post: { parameters: { query?: never; From 0070ce23b311375c0e2b34a51f193e90ec5f6c65 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 04:27:11 -0700 Subject: [PATCH 20/64] fix(router)!: apply content_policy_fallbacks when router settings update at runtime (#45428) * fix(router): apply content_policy_fallbacks when router settings update at runtime * test(router): annotate the fallback regression test's locals as Final * fix(proxy): validate router_settings.content_policy_fallbacks on /config/update like its siblings * test(proxy): annotate the content_policy_fallbacks config update test locals as Final * fix(proxy): keep litellm_settings content policy fallbacks, accept dict targets, skip invalid stored rules A runtime content_policy_fallbacks rule now layers over the litellm_settings rules instead of replacing them, so models the runtime rules do not name keep their file fallbacks. Fallback targets may be deployment dicts as well as model names. A stored rule that fails validation is logged and skipped at sync, so the other router settings still apply. Per-user token discovery now reads every fallback key from key and team router_settings. * fix(proxy): validate content policy fallback rules one by one and keep a config-owned list alone An invalid stored rule is now skipped on its own instead of dropping every runtime rule, a router_settings list from the config file stays the router's only rules across DB syncs, and GET /fallback names the model of a dict target. * fix(proxy): answer an error for a string-valued fallback rule on GET /fallback instead of its characters * fix(proxy): keep config-file content-policy rules as Router init accepted them and log a skipped rule once per change * fix(proxy): log a skipped content-policy rule whenever the saved rules change * fix(proxy): keep general fallbacks when no content-policy rule is saved, name a model-less target's model, and take only model names in key rules * fix(proxy): reject an empty dict target in a key's content-policy rule --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/litellm_pre_call_utils.py | 17 +- .../fallback_management_endpoints.py | 95 ++++++- litellm/proxy/proxy_server.py | 33 ++- .../management_endpoints/key_requests.py | 24 +- litellm/types/router.py | 6 + .../test_key_management_endpoints.py | 16 +- .../test_fallback_management_endpoints.py | 100 ++++++- .../unit/proxy/test_litellm_pre_call_utils.py | 11 +- ...test_proxy_server_endpoints_and_startup.py | 243 +++++++++++++++++- tests/unit/test_router/test_router.py | 33 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 12 files changed, 553 insertions(+), 32 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 97e93309ceb..c586a5f767a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -68,6 +68,7 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( "retry_after", "fallbacks", "context_window_fallbacks", + "content_policy_fallbacks", "retry_policy", "model_group_retry_policy", "model_group_alias", diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index c1e541d20a9..9b52f2f8a8d 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -3188,9 +3188,9 @@ def _fallback_lists( ) -> tuple[Sequence[object], ...]: """Every fallback list the router can route this request through: the three router-level lists, the same three keys from the request body (which replace - the router's when present, so the union is the safe superset), and the - key/team ``router_settings.fallbacks`` the proxy prefers over router-level - fallbacks in ``_configured_fallbacks``.""" + the router's when present, so the union is the safe superset), and the same + three keys from key/team ``router_settings``, which the proxy forwards into + the request after this discovery runs.""" lists: Final[list[Sequence[object]]] = [] # mutable-ok: accumulates each fallback list found configured: Final = llm_router.settings() for router_list in ( @@ -3205,11 +3205,12 @@ def _fallback_lists( if isinstance(entries, list): lists.append(cast("Sequence[object]", entries)) # cast-ok: request fallback lists hold mixed entry shapes for settings in (router_settings, team_router_settings): - key_fallbacks: object = settings.get("fallbacks") if settings is not None else None - if isinstance(key_fallbacks, list): - lists.append( - cast("Sequence[object]", key_fallbacks) # cast-ok: settings fallback lists hold mixed entry shapes - ) + for key in _REQUEST_FALLBACK_KEYS: + key_fallbacks: object = settings.get(key) if settings is not None else None + if isinstance(key_fallbacks, list): + lists.append( + cast("Sequence[object]", key_fallbacks) # cast-ok: settings fallback lists hold mixed entry shapes + ) return tuple(lists) diff --git a/litellm/proxy/management_endpoints/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index d688aac0ed8..477cf44073f 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -14,10 +14,12 @@ import json from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Annotated, Final, Literal -from pydantic import Field, JsonValue, TypeAdapter +from pydantic import Field, JsonValue, TypeAdapter, ValidationError +import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_utils import fallback_target_model_name from litellm.proxy.auth.model_checks import get_all_fallbacks from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.utils import PrismaClient, invalidate_config_param @@ -26,6 +28,7 @@ if TYPE_CHECKING: from fastapi import APIRouter, Depends, HTTPException, status from litellm.proxy.proxy_server import ProxyConfig + from litellm.router import Router else: try: from fastapi import APIRouter, Depends, HTTPException, status @@ -41,6 +44,7 @@ from litellm.types.management_endpoints.router_settings_endpoints import ( FallbackGetResponse, FallbackResponse, ) +from litellm.types.router import ContentPolicyFallbackRule, ContentPolicyFallbackRules, FallbackTargets router: Final = APIRouter() @@ -50,12 +54,83 @@ StoredFallback = Annotated[FallbackRule | dict[str, object] | str, Field(union_m _STORED_FALLBACKS: Final = TypeAdapter(list[StoredFallback]) _JSON_RULES: Final = TypeAdapter(list[JsonValue]) _STORED_ROUTER_SETTINGS: Final = TypeAdapter(dict[str, JsonValue]) +_OPTIONAL_JSON_RULES: Final[TypeAdapter[Sequence[JsonValue] | None]] = TypeAdapter(Sequence[JsonValue] | None) +_CONTENT_POLICY_FALLBACK_RULE: Final[TypeAdapter[ContentPolicyFallbackRule]] = TypeAdapter(ContentPolicyFallbackRule) +_FALLBACK_TARGETS: Final[TypeAdapter[FallbackTargets]] = TypeAdapter(FallbackTargets) def _rule_covers(entry: StoredFallback, model: str) -> bool: return isinstance(entry, dict) and model in entry +def _layered_over_config_rules( + runtime_rules: ContentPolicyFallbackRules, config_rules: Sequence[Mapping[str, JsonValue]] +) -> Sequence[Mapping[str, object]]: + covered: Final = frozenset[str]().union(*runtime_rules) + return [*runtime_rules, *(rule for rule in config_rules if covered.isdisjoint(rule))] + + +def _config_file_rules(rules: object) -> Sequence[Mapping[str, JsonValue]]: + try: + candidates: Final = _OPTIONAL_JSON_RULES.validate_python(rules) or () + except ValidationError: + return () + return tuple(rule for rule in candidates if isinstance(rule, Mapping)) + + +def _log_ignored_content_policy_rules(rules: object, error: ValidationError) -> None: + verbose_proxy_logger.error( + "Ignoring invalid content_policy_fallbacks %r from the DB, every other rule still applies. " + "Fix it through POST /fallback or /config/update: %s", + rules, + error, + ) + + +def _checked_content_policy_rule(rule: JsonValue) -> ContentPolicyFallbackRule | ValidationError: + try: + return _CONTENT_POLICY_FALLBACK_RULE.validate_python(rule) + except ValidationError as invalid_rule: + return invalid_rule + + +def _checked_content_policy_rules( + rules: object, +) -> Sequence[tuple[object, ContentPolicyFallbackRule | ValidationError]]: + try: + candidates: Final = _OPTIONAL_JSON_RULES.validate_python(rules) or () + except ValidationError as not_a_rule_list: + return ((rules, not_a_rule_list),) + return tuple((candidate, _checked_content_policy_rule(candidate)) for candidate in candidates) + + +def apply_content_policy_fallbacks(llm_router: "Router", rules: object, *, log_skipped_rules: bool) -> None: + checked_rules: Final = _checked_content_policy_rules(rules) + runtime_rules: Final = tuple(result for _, result in checked_rules if not isinstance(result, ValidationError)) + layered_rules: Final = ( + _layered_over_config_rules( + runtime_rules, + _config_file_rules( + litellm.content_policy_fallbacks # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped module list + ), + ) + or None + ) + for rule, result in checked_rules: + if log_skipped_rules and isinstance(result, ValidationError): + _log_ignored_content_policy_rules(rule, result) + if layered_rules == llm_router.settings().content_policy_fallbacks: # pyright: ignore[reportUnknownMemberType] # untyped RouterSettings list + return + llm_router.apply(ChangeSet(settings={"content_policy_fallbacks": layered_rules})) + + +def _apply_fallback_rules(llm_router: "Router", fallback_key: str, rules: Sequence[StoredFallback]) -> None: + if fallback_key == "content_policy_fallbacks": + apply_content_policy_fallbacks(llm_router, rules, log_skipped_rules=True) + return + llm_router.apply(ChangeSet(assign={fallback_key: rules})) + + async def _reload_config_file_ownership(proxy_config: "ProxyConfig") -> None: await proxy_config.get_config() @@ -83,6 +158,7 @@ def _editable_rules( async def _write_rules( prisma_client: PrismaClient, proxy_config: "ProxyConfig", + llm_router: "Router", stored_settings: Mapping[str, JsonValue], fallback_key: str, rules: Sequence[StoredFallback], @@ -91,8 +167,10 @@ async def _write_rules( section_name=ROUTER_SETTINGS_PARAM, changed_keys={fallback_key: _JSON_RULES.validate_python(rules)} ) if proxy_config.router_settings.owned_by_config(fallback_key): + llm_router.apply(ChangeSet(assign={fallback_key: rules})) return await _persist_router_settings(prisma_client, {**stored_settings, fallback_key: rules}) + _apply_fallback_rules(llm_router, fallback_key, rules) async def _persist_router_settings(prisma_client: PrismaClient, router_settings: Mapping[str, object]) -> None: @@ -218,10 +296,7 @@ async def create_fallback( # Add new fallback existing_fallbacks.append({data.model: data.fallback_models}) - await _write_rules(prisma_client, proxy_config, router_settings, fallback_key, existing_fallbacks) - - # Update the in-memory router configuration - llm_router.apply(ChangeSet(assign={fallback_key: existing_fallbacks})) + await _write_rules(prisma_client, proxy_config, llm_router, router_settings, fallback_key, existing_fallbacks) verbose_proxy_logger.info( "Fallback configured: %s -> %s (type: %s)", data.model, data.fallback_models, data.fallback_type @@ -277,7 +352,10 @@ async def get_fallback( ) # Get fallbacks using the existing utility function - fallback_models: Final = get_all_fallbacks(model=model, llm_router=llm_router, fallback_type=fallback_type) + fallback_targets: Final = _FALLBACK_TARGETS.validate_python( + get_all_fallbacks(model=model, llm_router=llm_router, fallback_type=fallback_type) + ) + fallback_models: Final = [fallback_target_model_name(target) or model for target in fallback_targets] if not fallback_models: raise HTTPException( @@ -367,10 +445,7 @@ async def delete_fallback( detail={"error": f"No {fallback_type} fallbacks configured for model '{model}'"}, ) - await _write_rules(prisma_client, proxy_config, router_settings, fallback_key, updated_fallbacks) - - # Update the in-memory router configuration - llm_router.apply(ChangeSet(assign={fallback_key: updated_fallbacks})) + await _write_rules(prisma_client, proxy_config, llm_router, router_settings, fallback_key, updated_fallbacks) verbose_proxy_logger.info("Fallback deleted: %s (type: %s)", model, fallback_type) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ae0e7fce667..f078267c6b2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -718,6 +718,9 @@ from litellm.proxy.management_endpoints.customer_endpoints import ( router as customer_router, ) from litellm.proxy.management_endpoints.daily_activity_routes import router as daily_activity_router +from litellm.proxy.management_endpoints.fallback_management_endpoints import ( + apply_content_policy_fallbacks, +) from litellm.proxy.management_endpoints.fallback_management_endpoints import ( router as fallback_management_router, ) @@ -8163,14 +8166,38 @@ class ProxyConfig: if db_router_settings is not None and db_router_settings.param_value is not None else _EMPTY_SETTINGS_MAPPING ) + previous_content_policy_rules: Final = self.router_settings.resolved().get("content_policy_fallbacks") self.router_settings.apply_db_row("router_settings", db_values) combined_router_settings: Final = self.router_settings.resolved() if combined_router_settings: - self._apply_router_settings(llm_router, combined_router_settings) + self._apply_router_settings( + llm_router, + combined_router_settings, + content_policy_owned_by_config=self.router_settings.owned_by_config("content_policy_fallbacks"), + content_policy_rules_changed=( + combined_router_settings.get("content_policy_fallbacks") != previous_content_policy_rules + ), + ) @staticmethod - def _apply_router_settings(llm_router: Router, router_settings: Mapping[str, object]) -> None: - llm_router.apply(ChangeSet(settings={k: v for k, v in router_settings.items() if k != "routing_groups"})) + def _apply_router_settings( + llm_router: Router, + router_settings: Mapping[str, object], + *, + content_policy_owned_by_config: bool, + content_policy_rules_changed: bool, + ) -> None: + llm_router.apply( + ChangeSet( + settings={ + k: v for k, v in router_settings.items() if k not in ("routing_groups", "content_policy_fallbacks") + } + ) + ) + if "content_policy_fallbacks" in router_settings and not content_policy_owned_by_config: + apply_content_policy_fallbacks( + llm_router, router_settings["content_policy_fallbacks"], log_skipped_rules=content_policy_rules_changed + ) if "routing_groups" not in router_settings: return try: diff --git a/litellm/types/proxy/management_endpoints/key_requests.py b/litellm/types/proxy/management_endpoints/key_requests.py index 172357d8ac4..b49cb63c8e0 100644 --- a/litellm/types/proxy/management_endpoints/key_requests.py +++ b/litellm/types/proxy/management_endpoints/key_requests.py @@ -1,14 +1,14 @@ import enum -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from datetime import datetime from types import MappingProxyType -from typing import Any, Literal # noqa: TID251 # allowed_routes and metadata stay Any for existing callers +from typing import Any, Final, Literal # noqa: TID251 # allowed_routes and metadata stay Any for existing callers from pydantic import Field, PositiveInt, field_validator, model_validator from typing_extensions import TypeIs # noqa: TID251 # narrows the raw payload a before-validator receives from litellm.types.proxy.management_endpoints.request_base import GenerateRequestBase -from litellm.types.router import UpdateRouterConfig +from litellm.types.router import ContentPolicyFallbackRules, UpdateRouterConfig from litellm.types.utils import LiteLLMPydanticObjectBase @@ -17,6 +17,12 @@ class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): index_permissions: list[Literal["read", "write"]] +def _fallback_targets(rules: ContentPolicyFallbackRules) -> Iterator[object]: + for rule in rules: + for targets in rule.values(): + yield from targets + + class KeyRequestBase(GenerateRequestBase): key: str | None = None tpd_limit: int | None = None @@ -42,6 +48,18 @@ class KeyRequestBase(GenerateRequestBase): router_settings: UpdateRouterConfig | None = None access_group_ids: list[str] | None = None + @field_validator("router_settings") + @classmethod + def content_policy_fallbacks_name_models( + cls, router_settings: UpdateRouterConfig | None + ) -> UpdateRouterConfig | None: + rules: Final = None if router_settings is None else router_settings.content_policy_fallbacks + if rules is not None and not all(isinstance(target, str) for target in _fallback_targets(rules)): + raise ValueError( + 'A key\'s content_policy_fallbacks targets must be model names, e.g. [{"gpt-4": ["gpt-4o"]}]' + ) + return router_settings + class LiteLLMKeyType(str, enum.Enum): """ diff --git a/litellm/types/router.py b/litellm/types/router.py index f927ae0475b..726a328e913 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -192,6 +192,11 @@ OptionalPreCallChecks = list[ ] +FallbackTargets: TypeAlias = Sequence[str | Mapping[str, JsonValue]] +ContentPolicyFallbackRule: TypeAlias = Annotated[Mapping[str, FallbackTargets], Field(min_length=1, max_length=1)] +ContentPolicyFallbackRules: TypeAlias = Sequence[ContentPolicyFallbackRule] + + class UpdateRouterConfig(LiteLLMBaseModel): """ Set of params that you can modify via `router.update_settings()`. @@ -211,6 +216,7 @@ class UpdateRouterConfig(LiteLLMBaseModel): retry_after: float | None = None fallbacks: list[dict] | None = None context_window_fallbacks: list[dict] | None = None + content_policy_fallbacks: ContentPolicyFallbackRules | None = None model_group_alias: dict[str, str | dict] | None = Field(default={}) enable_tag_filtering: bool | None = None weights: RouterWeights | None = None diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index b0ef7e64e60..a0c32f2235f 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -13,7 +13,7 @@ import pytest import yaml from fastapi import HTTPException from fastapi.testclient import TestClient -from pydantic import JsonValue +from pydantic import JsonValue, ValidationError import litellm from litellm.constants import PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS @@ -7597,6 +7597,7 @@ async def test_generate_key_with_router_settings(monkeypatch): "num_retries": 3, "model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}, "weights": {"gpt-4": {"weighted-id": 1}}, + "content_policy_fallbacks": [{"gpt-4": ["gpt-4o", "gpt-4o-mini"]}], } request_data = GenerateKeyRequest( @@ -7736,6 +7737,19 @@ async def test_update_key_with_router_settings( assert json.loads(cleared["router_settings"]) == {} +@pytest.mark.parametrize("request_type", [GenerateKeyRequest, UpdateKeyRequest]) +@pytest.mark.parametrize("dict_target", [{"model": "gpt-4o-mini"}, {}]) +def test_a_key_content_policy_rule_can_only_name_fallback_models( + request_type: type[GenerateKeyRequest] | type[UpdateKeyRequest], + dict_target: dict[str, str], +): + with pytest.raises(ValidationError, match="must be model names"): + request_type( + key="sk-key", + router_settings={"content_policy_fallbacks": [{"gpt-4": ["gpt-4o", dict_target]}]}, + ) + + @pytest.mark.asyncio async def test_validate_max_budget(): """ diff --git a/tests/unit/proxy/test_fallback_management_endpoints.py b/tests/unit/proxy/test_fallback_management_endpoints.py index d28d61af72b..e3d0d072e0a 100644 --- a/tests/unit/proxy/test_fallback_management_endpoints.py +++ b/tests/unit/proxy/test_fallback_management_endpoints.py @@ -12,7 +12,7 @@ import copy import json from collections.abc import AsyncIterator, Mapping from types import MappingProxyType, SimpleNamespace -from typing import Final, get_args +from typing import Final, Literal, get_args from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -20,6 +20,7 @@ from fastapi import HTTPException from pydantic import JsonValue, TypeAdapter from redis.exceptions import NoPermissionError +import litellm from litellm import Router from litellm.proxy.config_resolvers.settings_rules import Section from litellm.proxy.management_endpoints.fallback_management_endpoints import ( @@ -388,6 +389,88 @@ class TestFallbackWritesSeeTheLatestStoredRules: assert stored.stored_fallbacks() == [GATEWAY_RULE, TEAM_RULE] + async def test_content_policy_writes_keep_the_litellm_settings_rules_for_other_models( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "content_policy_fallbacks", [TEAM_RULE]) + router: Final = _team_router() + stored: Final = _StoredRouterSettings(None) + with ( + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + await create_fallback( + FallbackCreateRequest( + model="gpt-5.4-mini", fallback_models=["team-fallback"], fallback_type="content_policy" + ), + MagicMock(), + ) + after_create: Final = await get_fallback("team-primary", "content_policy", MagicMock()) + await delete_fallback("gpt-5.4-mini", "content_policy", MagicMock()) + after_delete: Final = await get_fallback("team-primary", "content_policy", MagicMock()) + + assert after_create.fallback_models == ["team-fallback"] + assert after_delete.fallback_models == ["team-fallback"] + + async def test_a_content_policy_write_applies_beside_an_invalid_stored_rule(self) -> None: + router: Final = _team_router() + stored: Final = _StoredRouterSettings({"content_policy_fallbacks": [MALFORMED_GATEWAY_RULE]}) + with ( + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + await create_fallback( + FallbackCreateRequest( + model="team-primary", fallback_models=["team-fallback"], fallback_type="content_policy" + ), + MagicMock(), + ) + after_create: Final = await get_fallback("team-primary", "content_policy", MagicMock()) + + assert after_create.fallback_models == ["team-fallback"] + + @pytest.mark.parametrize( + ("fallback_type", "router_setting"), + [("general", "fallbacks"), ("content_policy", "content_policy_fallbacks")], + ) + async def test_get_names_the_model_of_a_dict_target( + self, fallback_type: Literal["general", "content_policy"], router_setting: str + ) -> None: + router: Final = _team_router() + router.update_settings( + **{router_setting: [{"team-primary": [{"model": "gpt-5.4-mini", "api_key": "sk-test"}, "team-fallback"]}]} + ) + with patch("litellm.proxy.proxy_server.llm_router", router): + response: Final = await get_fallback("team-primary", fallback_type, MagicMock()) + + assert response.fallback_models == ["gpt-5.4-mini", "team-fallback"] + + @pytest.mark.parametrize( + ("fallback_type", "router_setting"), + [("general", "fallbacks"), ("content_policy", "content_policy_fallbacks")], + ) + async def test_get_names_the_rules_own_model_for_a_dict_target_that_names_none( + self, fallback_type: Literal["general", "content_policy"], router_setting: str + ) -> None: + router: Final = _team_router() + router.update_settings(**{router_setting: [{"team-primary": [{"temperature": 0}, "team-fallback"]}]}) + with patch("litellm.proxy.proxy_server.llm_router", router): + response: Final = await get_fallback("team-primary", fallback_type, MagicMock()) + + assert response.fallback_models == ["team-primary", "team-fallback"] + + async def test_get_answers_an_error_for_a_string_valued_rule_instead_of_its_characters(self) -> None: + router: Final = _team_router() + router.update_settings(fallbacks=[{"team-primary": "team-fallback"}]) + with patch("litellm.proxy.proxy_server.llm_router", router), pytest.raises(HTTPException) as raised: + await get_fallback("team-primary", "general", MagicMock()) + + assert raised.value.status_code == 500 + async def test_a_create_keeps_a_rule_another_instance_stored_since_this_one_last_read(self) -> None: stored: Final = _StoredRouterSettings({}) await get_config_param(stored.prisma_client, "router_settings") @@ -568,6 +651,21 @@ class TestFallbackWritesHonourTheConfigFile: assert stored.row is None assert router.fallbacks == [CONFIG_RULE, CONFIG_WINDOW_RULE] + async def test_a_content_policy_restate_keeps_the_config_list_without_the_litellm_settings_rules( + self, router: Router, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "content_policy_fallbacks", [TEAM_RULE]) + stored: Final = _StoredRouterSettings(None, config_file={"content_policy_fallbacks": [CONFIG_RULE]}) + request: Final = FallbackCreateRequest( + model="gpt-5.4-mini", fallback_models=["team-fallback"], fallback_type="content_policy" + ) + + response: Final = await self._create(request, router, stored) + + assert response.fallback_models == ["team-fallback"] + assert stored.row is None + assert router.content_policy_fallbacks == [CONFIG_RULE] + async def test_a_create_that_restates_a_config_rule_next_to_an_env_reference_writes_nothing( self, router: Router, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 2edb032b048..c5136ff6e4a 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -10143,9 +10143,12 @@ class TestResolveUserProviderCredentials: assert data["secret_fields"]["user_provider_credentials_user_id"] == "user-a" @pytest.mark.asyncio - async def test_key_router_settings_fallback_loads_per_user_tokens(self, monkeypatch): - """The proxy prefers key/team router_settings.fallbacks over router-level - fallbacks (_configured_fallbacks), so discovery must read them too.""" + @pytest.mark.parametrize("fallback_key", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) + async def test_key_router_settings_fallback_loads_per_user_tokens( + self, monkeypatch: pytest.MonkeyPatch, fallback_key: str + ): + """The proxy forwards key/team router_settings fallback lists into the + request after discovery, so discovery must read every one of them.""" from litellm.proxy.litellm_pre_call_utils import _resolve_user_provider_credentials_for_request table = self._env(monkeypatch) @@ -10175,7 +10178,7 @@ class TestResolveUserProviderCredentials: authenticated_user_id="user-a", team_id=None, llm_router=router, - router_settings={"fallbacks": [{"gpt-4o": [{"model": "copilot-chat"}]}]}, + router_settings={fallback_key: [{"gpt-4o": [{"model": "copilot-chat"}]}]}, ) table.find_many.assert_awaited() assert data["secret_fields"]["user_provider_credentials_user_id"] == "user-a" diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 62a90109d26..bc3a228321a 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -10,6 +10,7 @@ import subprocess import sys import time import types +from collections.abc import Callable from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Final @@ -5583,6 +5584,215 @@ async def test_valid_db_routing_groups_still_replace_router_groups(): assert router._get_routing_context("m2", None)[0] == "least-busy" +def _content_policy_router(**router_kwargs: object) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake", "mock_response": "primary answered"}, + }, + { + "model_name": "secondary", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake", "mock_response": "secondary answered"}, + }, + { + "model_name": "backup", + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "fake", "mock_response": "backup answered"}, + }, + ], + num_retries=0, + **router_kwargs, + ) + + +def _router_settings_row(content_policy_fallbacks: object) -> MagicMock: + row: Final = MagicMock() + row.param_value = {"allowed_fails": 7, "content_policy_fallbacks": content_policy_fallbacks} + prisma: Final = MagicMock() + prisma.db.litellm_config.find_first = AsyncMock(return_value=row) + return prisma + + +async def _answer_after_content_policy_refusal(router: litellm.Router, model: str = "primary") -> str | None: + response: Final = await router.acompletion( + model=model, messages=[{"role": "user", "content": "hi"}], mock_testing_content_policy_fallbacks=True + ) + return response.choices[0].message.content + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored_target", ["backup", {"model": "backup"}]) +async def test_db_content_policy_fallbacks_reach_the_router(stored_target: str | dict[str, str]): + router: Final = _content_policy_router() + + await proxy_server_module.ProxyConfig()._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row([{"primary": [stored_target]}]) + ) + + assert await _answer_after_content_policy_refusal(router) == "backup answered" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored_rules", [[], [{"primary": "backup"}]]) +async def test_a_saved_content_policy_list_with_no_valid_rule_leaves_refusals_to_the_general_fallbacks( + monkeypatch: pytest.MonkeyPatch, stored_rules: list[object] +): + monkeypatch.setattr(litellm, "content_policy_fallbacks", None) + router: Final = _content_policy_router(fallbacks=[{"primary": ["backup"]}]) + + await proxy_server_module.ProxyConfig()._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row(stored_rules) + ) + + assert await _answer_after_content_policy_refusal(router) == "backup answered" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("litellm_settings_rules", "stored_rules", "refused_model"), + [ + ([{"secondary": ["backup"]}], [{"primary": ["backup"]}], "secondary"), + ([{"*": ["secondary"]}], [{"*": ["backup"]}], "primary"), + ([{"primary": ["backup"]}], [], "primary"), + ], +) +async def test_db_content_policy_fallbacks_layer_over_the_litellm_settings_rules( + monkeypatch: pytest.MonkeyPatch, + litellm_settings_rules: list[dict[str, list[str]]], + stored_rules: list[dict[str, list[str]]], + refused_model: str, +): + monkeypatch.setattr(litellm, "content_policy_fallbacks", litellm_settings_rules) + router: Final = _content_policy_router() + + await proxy_server_module.ProxyConfig()._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row(stored_rules) + ) + + assert await _answer_after_content_policy_refusal(router, refused_model) == "backup answered" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "invalid_rule", + [{"secondary": "backup"}, ["backup"], dict[str, list[str]](), {"secondary": ["backup"], "primary": ["secondary"]}], +) +async def test_an_invalid_db_content_policy_rule_skips_only_itself(monkeypatch: pytest.MonkeyPatch, invalid_rule: object): + monkeypatch.setattr(litellm, "content_policy_fallbacks", [{"primary": ["backup"]}]) + router: Final = _content_policy_router() + + await proxy_server_module.ProxyConfig()._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row([invalid_rule, {"secondary": ["backup"]}]) + ) + + assert router.allowed_fails == 7 + assert await _answer_after_content_policy_refusal(router) == "backup answered" + assert await _answer_after_content_policy_refusal(router, "secondary") == "backup answered" + + +@pytest.mark.asyncio +async def test_a_db_content_policy_value_that_is_not_a_list_keeps_the_litellm_settings_rules( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "content_policy_fallbacks", [{"primary": ["backup"]}]) + router: Final = _content_policy_router() + + await proxy_server_module.ProxyConfig()._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row("oops") + ) + + assert router.allowed_fails == 7 + assert await _answer_after_content_policy_refusal(router) == "backup answered" + + +@pytest.mark.asyncio +async def test_a_db_sync_keeps_a_litellm_settings_rule_with_a_target_router_init_accepted( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "content_policy_fallbacks", [{"primary": ["backup", 5]}]) + router: Final = _content_policy_router() + + await proxy_server_module.ProxyConfig()._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row([{"secondary": ["backup"]}]) + ) + + assert await _answer_after_content_policy_refusal(router) == "backup answered" + assert await _answer_after_content_policy_refusal(router, "secondary") == "backup answered" + + +@pytest.mark.asyncio +async def test_an_invalid_db_content_policy_rule_is_logged_once_per_change_of_the_rules( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + monkeypatch.setattr(litellm, "content_policy_fallbacks", []) + router: Final = _content_policy_router() + proxy_config: Final = proxy_server_module.ProxyConfig() + stored_rules: Final = [{"secondary": "backup"}, {"primary": ["backup"]}] + changed_rules: Final = [{"secondary": "backup"}, {"secondary": ["backup"]}] + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + await proxy_config._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row(stored_rules) + ) + await proxy_config._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row(stored_rules) + ) + await proxy_config._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row(changed_rules) + ) + + notices: Final = [ + record for record in caplog.records if "Ignoring invalid content_policy_fallbacks" in record.getMessage() + ] + assert len(notices) == 2 + assert await _answer_after_content_policy_refusal(router, "secondary") == "backup answered" + + +@pytest.mark.asyncio +async def test_a_db_sync_logs_a_new_invalid_content_policy_rule_when_no_valid_rule_changes( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + monkeypatch.setattr(litellm, "content_policy_fallbacks", []) + router: Final = _content_policy_router() + proxy_config: Final = proxy_server_module.ProxyConfig() + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + await proxy_config._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row([{"secondary": "backup"}]) + ) + await proxy_config._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row([{"secondary": "backup"}, {"primary": "backup"}]) + ) + + notices: Final = [ + record.getMessage() + for record in caplog.records + if "Ignoring invalid content_policy_fallbacks" in record.getMessage() + ] + assert len(notices) == 3 + assert "{'primary': 'backup'}" in notices[-1] + + +@pytest.mark.asyncio +async def test_config_file_content_policy_fallbacks_stay_the_only_rules_after_a_db_sync( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "content_policy_fallbacks", [{"secondary": ["backup"]}]) + file_rules: Final = [{"primary": ["backup"]}] + router: Final = _content_policy_router(content_policy_fallbacks=file_rules) + proxy_config: Final = proxy_server_module.ProxyConfig() + proxy_config.router_settings.load_yaml({"content_policy_fallbacks": file_rules}) + + await proxy_config._add_router_settings_from_db_config( + llm_router=router, prisma_client=_router_settings_row([{"secondary": ["backup"]}]) + ) + + assert router.allowed_fails == 7 + assert await _answer_after_content_policy_refusal(router) == "backup answered" + with pytest.raises(litellm.ContentPolicyViolationError): + await _answer_after_content_policy_refusal(router, "secondary") + + @pytest.mark.asyncio async def test_add_router_settings_from_db_config_empty_db_lists_do_not_clobber_config_fallbacks(): """ @@ -5591,6 +5801,7 @@ async def test_add_router_settings_from_db_config_empty_db_lists_do_not_clobber_ empty lists are "no value" and must not clobber config.yaml fallbacks, matching _deep_merge_dicts semantics. Non-empty DB lists still win. """ + from collections import ChainMap from unittest.mock import AsyncMock, MagicMock from litellm.proxy.proxy_server import ProxyConfig @@ -5625,10 +5836,10 @@ async def test_add_router_settings_from_db_config_empty_db_lists_do_not_clobber_ prisma_client=mock_prisma_client, ) - combined_settings = mock_router.update_settings.call_args.kwargs + combined_settings = ChainMap(*(call.kwargs for call in reversed(mock_router.update_settings.call_args_list))) assert combined_settings["fallbacks"] == [{"gpt-oss-120b": ["granite-4-h-small"]}] assert combined_settings["context_window_fallbacks"] == [{"gpt-oss-120b": ["granite-4-h-small"]}] - assert combined_settings["content_policy_fallbacks"] == [{"gpt-oss-120b": ["granite-4-h-small"]}] + assert combined_settings.get("content_policy_fallbacks") != [{"gpt-oss-120b": ["other-model"]}] assert combined_settings["num_retries"] == 3 @@ -9997,6 +10208,34 @@ def test_update_config_accepts_disjoint_routing_groups(_update_config_setup): restore() +def test_update_config_validates_content_policy_fallbacks_like_fallbacks( + _update_config_setup: Callable[..., tuple[TestClient, _FakePrismaClient, Callable[[], None]]], +): + client, prisma, restore = _update_config_setup(initial_rows={"router_settings": {"num_retries": 2}}) + rule: Final = [{"gpt-5.4-mini": ["gpt-5.4-nano", {"model": "gpt-5.4-nano", "temperature": 1}]}] + malformed: Final = ( + "oops", + 5, + [1, 2], + [{"gpt-5.4-mini": "gpt-5.4-nano"}], + [{"gpt-5.4-mini": ["gpt-5.4-nano", 7]}], + [dict[str, list[str]]()], + [{"gpt-5.4-mini": ["gpt-5.4-nano"], "gpt-5.4": ["gpt-5.4-nano"]}], + ) + try: + rejected: Final = tuple( + client.post("/config/update", json={"router_settings": {"content_policy_fallbacks": value}}).status_code + for value in malformed + ) + assert rejected == (422,) * len(malformed) + assert prisma.db.litellm_config.rows["router_settings"] == {"num_retries": 2} + accepted: Final = client.post("/config/update", json={"router_settings": {"content_policy_fallbacks": rule}}) + assert accepted.status_code == 200 + assert prisma.db.litellm_config.rows["router_settings"] == {"num_retries": 2, "content_policy_fallbacks": rule} + finally: + restore() + + def test_update_config_env_var_round_trip_not_double_encrypted(_update_config_setup, monkeypatch): """Endpoint-level regression for the /config/update double-encryption bug. diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 74959155380..1ac44b5976d 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2731,6 +2731,39 @@ def test_update_settings_model_group_alias_drops_cached_group_info(): assert after.input_cost_per_token is not None and after.input_cost_per_token > 0 +@pytest.mark.parametrize( + "fallback_kind, mock_testing_param", + [ + ("fallbacks", "mock_testing_fallbacks"), + ("context_window_fallbacks", "mock_testing_context_fallbacks"), + ("content_policy_fallbacks", "mock_testing_content_policy_fallbacks"), + ], +) +async def test_update_settings_fallback_rule_of_every_kind_is_used_on_its_error( + fallback_kind: str, mock_testing_param: str +): + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": "gpt-5.4-mini", "api_key": "fake", "mock_response": "primary answered"}, + }, + { + "model_name": "backup", + "litellm_params": {"model": "gpt-5.4-nano", "api_key": "fake", "mock_response": "backup answered"}, + }, + ], + num_retries=0, + ) + + router.update_settings(**{fallback_kind: [{"primary": ["backup"]}]}) + + response: Final = await router.acompletion( + model="primary", messages=[{"role": "user", "content": "hi"}], **{mock_testing_param: True} + ) + assert response.choices[0].message.content == "backup answered" + + _PAID_INPUT_COST_PER_TOKEN: Final = 3e-06 _PAID_OUTPUT_COST_PER_TOKEN: Final = 1.5e-05 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f53e57c2f18..1aec7d5b761 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -49017,6 +49017,12 @@ export interface components { UpdateRouterConfig: { /** Allowed Fails */ allowed_fails?: number | null; + /** Content Policy Fallbacks */ + content_policy_fallbacks?: { + [key: string]: (string | { + [key: string]: components["schemas"]["JsonValue"]; + })[]; + }[] | null; /** Context Window Fallbacks */ context_window_fallbacks?: { [key: string]: unknown; From 02ef44c3885b2b0500cf267abf57064fbd206dcd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 04:41:02 -0700 Subject: [PATCH 21/64] fix(streaming): convert upstream logprobs so streams with logprobs end with their finish chunk (#33549) * fix(streaming): normalize provider logprobs to prevent MockValSer serialization crash * fix(streaming): convert upstream chunk logprobs to litellm ChoiceLogprobs so the finish chunk serializes A streaming chat completion with logprobs on the OpenAI SDK path attached the SDK's own ChoiceLogprobs object to the outgoing chunk when the chunk had no content to rebuild from, which is the provider's finish chunk (empty delta, logprobs for the end token, finish_reason). The SDK parses stream chunks without building that nested class, so its pydantic serializer is still a MockValSer and model_dump_json fails. The proxy then sent an error frame in place of the finish chunk, and when the same chunk also carried usage the SDK wrapper itself raised MidStreamFallbackError. The attach site now validates upstream logprobs through a TypeAdapter that dumps any non-litellm pydantic model before validating into litellm's ChoiceLogprobs, and passes litellm instances through unchanged. * fix(streaming): keep upstream logprobs that don't fit ChoiceLogprobs as plain JSON so the stream still finishes * fix(streaming): build ChoiceLogprobs' serializer up front and keep nonconforming logprobs as JSON on chunk rebuilds ChoiceLogprobs inherits the openai SDK's deferred build, so a fresh worker whose first logprobs arrive on the finish chunk still had no serializer. Foreign logprobs models now dump in full, so a missing top_logprobs becomes [] and SDK callers keep attribute access, and chunk rebuilds for content and usage keep logprobs that don't fit as plain JSON instead of raising. * fix(streaming): keep only the logprobs keys upstream sent and type refusal logprobs * fix(types): pass nonconforming non-streaming logprobs through instead of failing the response --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../litellm_core_utils/streaming_handler.py | 8 +- litellm/types/utils.py | 30 +- .../test_convert_dict_to_response.py | 39 ++ .../test_streaming_handler.py | 354 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 5 files changed, 428 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 5542249d2c7..682b2a2f538 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -40,6 +40,7 @@ from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + CHOICE_LOGPROBS_OR_JSON, CacheCreationTokenDetails, CompletionTokensDetailsWrapper, Delta, @@ -1201,7 +1202,10 @@ class CustomStreamWrapper: choice_json.pop( "finish_reason", None ) # for mistral etc. which return a value in their last chunk (not-openai compatible). - choices.append(StreamingChoices(**choice_json)) + choice_json.pop("logprobs", None) + choices.append( + StreamingChoices(**choice_json, logprobs=getattr(choice, "logprobs", None)) + ) except Exception: choices.append(StreamingChoices()) setattr(model_response, "choices", choices) @@ -1605,7 +1609,7 @@ class CustomStreamWrapper: model_response.system_fingerprint = response_obj["original_chunk"].system_fingerprint self.system_fingerprint = response_obj["original_chunk"].system_fingerprint if response_obj["logprobs"] is not None: - model_response.choices[0].logprobs = response_obj["logprobs"] + model_response.choices[0].logprobs = CHOICE_LOGPROBS_OR_JSON.validate_python(response_obj["logprobs"]) if response_obj["usage"] is not None: if isinstance(response_obj["usage"], dict): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6a758961c7b..780c9ef6824 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -7,6 +7,7 @@ from enum import Enum from types import MappingProxyType from typing import ( TYPE_CHECKING, + Annotated, Any, Final, Literal, @@ -38,6 +39,7 @@ from openai.types.moderation_create_response import ( ) from pydantic import ( BaseModel, + BeforeValidator, ConfigDict, Field, FieldSerializationInfo, @@ -45,6 +47,7 @@ from pydantic import ( JsonValue, PrivateAttr, SkipValidation, + TypeAdapter, ValidationError, field_serializer, field_validator, @@ -1199,7 +1202,7 @@ class ChatCompletionTokenLogprob(OpenAIObject): unlikely. """ - top_logprobs: list[TopLogprob] + top_logprobs: list[TopLogprob] = Field(default_factory=list) """List of the most likely tokens and their log probability, at this token position. @@ -1236,9 +1239,13 @@ class ChatCompletionTokenLogprob(OpenAIObject): class ChoiceLogprobs(OpenAIObject): + model_config = ConfigDict(defer_build=False) + content: list[ChatCompletionTokenLogprob] | None = None """A list of message content tokens with log probability information.""" + refusal: Sequence[ChatCompletionTokenLogprob] | None = None + def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator return hasattr(self, key) @@ -1252,6 +1259,21 @@ class ChoiceLogprobs(OpenAIObject): return getattr(self, key) +def _sdk_logprobs_as_dict(logprobs: object) -> object: + if isinstance(logprobs, ChoiceLogprobs) or not isinstance(logprobs, BaseModel): + return logprobs + return logprobs.model_dump(exclude_unset=True) + + +CHOICE_LOGPROBS_OR_JSON: Final[TypeAdapter[ChoiceLogprobs | JsonValue]] = TypeAdapter( + Annotated[ + ChoiceLogprobs | JsonValue, + Field(union_mode="left_to_right"), + BeforeValidator(_sdk_logprobs_as_dict), + ] +) + + class FunctionCall(OpenAIObject): arguments: str name: str | None = None @@ -1765,7 +1787,7 @@ class Choices(SafeAttributeModel, OpenAIObject): params["message"] = Message(**dump) if logprobs is not None: if isinstance(logprobs, dict): - params["logprobs"] = ChoiceLogprobs(**logprobs) + params["logprobs"] = CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs) else: params["logprobs"] = logprobs else: @@ -2131,8 +2153,8 @@ class StreamingChoices(OpenAIObject): if enhancements is not None: self.enhancements = enhancements - if logprobs is not None and isinstance(logprobs, dict): - self.logprobs = ChoiceLogprobs(**logprobs) + if isinstance(logprobs, (dict, BaseModel)): + self.logprobs = CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs) else: self.logprobs = logprobs diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index faecb8d8ef6..de8d6a7a8c5 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -1,6 +1,7 @@ import asyncio import importlib import json +from collections.abc import Mapping from datetime import datetime, timedelta from typing import Final @@ -1111,6 +1112,44 @@ def test_convert_to_model_response_object_with_logprobs(): assert result.system_fingerprint is None assert result._hidden_params is not None + +@pytest.mark.parametrize( + "upstream_logprobs", + [ + pytest.param({"content": [{"token": "Hi", "logprob": None}]}, id="null-content-logprob"), + pytest.param( + { + "content": [{"token": "Hi", "logprob": -0.1, "top_logprobs": []}], + "refusal": [{"token": "I", "logprob": None}], + }, + id="null-refusal-logprob", + ), + ], +) +def test_convert_to_model_response_object_passes_nonconforming_logprobs_through( + upstream_logprobs: Mapping[str, object], +) -> None: + result: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1702685778, + "model": "test-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hi"}, + "logprobs": upstream_logprobs, + "finish_reason": "stop", + } + ], + }, + model_response_object=ModelResponse(), + ) + + assert json.loads(result.model_dump_json())["choices"][0]["logprobs"] == upstream_logprobs + + def test_convert_to_model_response_object_error(): """Test error handling for None response object.""" with pytest.raises(Exception, match="Error in response object format"): diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 41e8c50fa85..b9c8feb223d 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -2,12 +2,15 @@ import asyncio import datetime import itertools import json +import os import queue +import subprocess import sys import threading import time from collections.abc import Mapping, Sequence from dataclasses import dataclass +from pathlib import Path from types import CodeType, FrameType, MappingProxyType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -33,6 +36,7 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( + ChoiceLogprobs, CompletionTokensDetailsWrapper, Delta, ModelResponse, @@ -1624,6 +1628,356 @@ async def test_stream_with_null_choices_chunk_finishes_normally( assert [choice.finish_reason for choice in streamed_choices if choice.finish_reason] == ["stop"] +def _token_logprobs(token: str) -> dict[str, object]: + return { + "content": [{"token": token, "logprob": -0.5, "bytes": list(token.encode()), "top_logprobs": []}], + "refusal": None, + } + + +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("transport", ["sdk", "http_handler"]) +@pytest.mark.parametrize( + ("finish_chunk_extra", "stream_options", "expected_usage_totals"), + [ + pytest.param({}, None, [], id="finish-chunk-with-logprobs"), + pytest.param( + {"usage": {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6}}, + {"include_usage": True}, + [6], + id="finish-chunk-with-logprobs-and-usage", + ), + ], +) +async def test_stream_finish_chunk_with_logprobs_reaches_caller_as_litellm_logprobs( + async_stream: bool, + transport: str, + finish_chunk_extra: dict[str, object], + stream_options: dict[str, bool] | None, + expected_usage_totals: list[int], + logging_obj: Logging, +) -> None: + upstream_tokens: Final = ("Hel", "lo", "<|return|>") + upstream_chunks: Final = ( + { + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "Hel"}, + "logprobs": _token_logprobs("Hel"), + "finish_reason": None, + } + ] + }, + { + "choices": [ + {"index": 0, "delta": {"content": "lo"}, "logprobs": _token_logprobs("lo"), "finish_reason": None} + ] + }, + { + "choices": [{"index": 0, "delta": {}, "logprobs": _token_logprobs("<|return|>"), "finish_reason": "stop"}], + **finish_chunk_extra, + }, + ) + body: Final = ( + "".join( + "data: " + + json.dumps( + { + "id": "chatcmpl-logprobs", + "object": "chat.completion.chunk", + "created": 1, + "model": "test-model", + **chunk, + } + ) + + "\n\n" + for chunk in upstream_chunks + ) + + "data: [DONE]\n\n" + ) + raw_response: Final = httpx.Response( + 200, + content=body, + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), + ) + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + stream: Final = ( + ( + AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + ) + if transport == "sdk" + else OpenAIChatCompletionStreamingHandler( + streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(), + sync_stream=not async_stream, + ) + ) + response: Final = CustomStreamWrapper( + completion_stream=stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm", + stream_options=stream_options, + ) + + try: + chunks: Final = [chunk async for chunk in response] if async_stream else list(response) + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + wire_choices: Final = tuple( + itertools.chain.from_iterable( + json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True))["choices"] for chunk in chunks + ) + ) + finish_choices: Final = [choice for choice in wire_choices if choice.get("finish_reason")] + streamed_logprobs: Final = tuple( + choice.logprobs + for choice in itertools.chain.from_iterable(chunk.choices for chunk in chunks) + if choice.logprobs is not None + ) + streamed_tokens: Final = tuple( + entry.token for entry in itertools.chain.from_iterable(logprobs.content or () for logprobs in streamed_logprobs) + ) + + assert [choice["finish_reason"] for choice in finish_choices] == ["stop"] + assert [entry["token"] for entry in finish_choices[0]["logprobs"]["content"]] == [upstream_tokens[-1]] + assert all(isinstance(logprobs, ChoiceLogprobs) for logprobs in streamed_logprobs) + assert streamed_tokens == upstream_tokens + assert [chunk.usage.total_tokens for chunk in chunks if getattr(chunk, "usage", None)] == expected_usage_totals + + +async def _sdk_stream_chunks( + upstream_chunks: Sequence[Mapping[str, object]], + async_stream: bool, + logging_obj: Logging, + stream_options: dict[str, bool] | None = None, +) -> list[ModelResponseStream]: + body: Final = ( + "".join( + "data: " + + json.dumps( + { + "id": "chatcmpl-logprobs", + "object": "chat.completion.chunk", + "created": 1, + "model": "test-model", + **chunk, + } + ) + + "\n\n" + for chunk in upstream_chunks + ) + + "data: [DONE]\n\n" + ) + raw_response: Final = httpx.Response( + 200, + content=body, + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), + ) + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + stream: Final = ( + AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + ) + response: Final = CustomStreamWrapper( + completion_stream=stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai", + stream_options=stream_options, + ) + try: + return [chunk async for chunk in response] if async_stream else list(response) + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize( + ("upstream_logprobs", "finish_chunk_extra", "stream_options", "expected_usage_totals"), + [ + pytest.param( + {"content": [{"token": "<|return|>", "bytes": None, "logprob": None, "top_logprobs": None}], "refusal": None}, + {}, + None, + [], + id="null-logprob", + ), + pytest.param( + {"content": [{"token": "<|return|>", "logprob": None}]}, + {}, + None, + [], + id="null-logprob-keeps-only-the-keys-upstream-sent", + ), + pytest.param(-0.25, {}, None, [], id="scalar-logprobs"), + pytest.param( + {"content": [{"token": "<|return|>", "bytes": None, "logprob": None, "top_logprobs": None}], "refusal": None}, + {"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}}, + {"include_usage": True}, + [5], + id="null-logprob-and-usage", + ), + ], +) +async def test_stream_with_nonstandard_logprobs_passes_them_through_and_finishes( + async_stream: bool, + upstream_logprobs: object, + finish_chunk_extra: dict[str, object], + stream_options: dict[str, bool] | None, + expected_usage_totals: list[int], + logging_obj: Logging, +) -> None: + chunks: Final = await _sdk_stream_chunks( + ( + { + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "hi"}, + "logprobs": upstream_logprobs, + "finish_reason": None, + } + ] + }, + { + "choices": [{"index": 0, "delta": {}, "logprobs": upstream_logprobs, "finish_reason": "stop"}], + **finish_chunk_extra, + }, + ), + async_stream, + logging_obj, + stream_options, + ) + + wire_choices: Final = tuple( + itertools.chain.from_iterable( + json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True))["choices"] for chunk in chunks + ) + ) + logprobs_choices: Final = [choice for choice in wire_choices if "logprobs" in choice] + + assert [choice.get("finish_reason") for choice in logprobs_choices] == [None, "stop"] + assert [choice["logprobs"] for choice in logprobs_choices] == [upstream_logprobs, upstream_logprobs] + assert [chunk.usage.total_tokens for chunk in chunks if getattr(chunk, "usage", None)] == expected_usage_totals + + +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("logprobs_field", ["content", "refusal"]) +async def test_stream_logprobs_without_top_logprobs_reach_caller_as_litellm_logprobs( + async_stream: bool, logprobs_field: Literal["content", "refusal"], logging_obj: Logging +) -> None: + chunks: Final = await _sdk_stream_chunks( + tuple( + { + "choices": [ + { + "index": 0, + "delta": delta, + "logprobs": {logprobs_field: [{"token": token, "logprob": -0.5}]}, + "finish_reason": finish_reason, + } + ] + } + for token, delta, finish_reason in ( + ("hi", {"role": "assistant", "content": "hi"}, None), + ("<|return|>", {}, "stop"), + ) + ), + async_stream, + logging_obj, + ) + + streamed_logprobs: Final = tuple( + choice.logprobs + for choice in itertools.chain.from_iterable(chunk.choices for chunk in chunks) + if choice.logprobs is not None + ) + streamed_entries: Final = tuple( + itertools.chain.from_iterable( + (logprobs.content if logprobs_field == "content" else logprobs.refusal) or () + for logprobs in streamed_logprobs + ) + ) + + assert [(entry.token, entry.logprob, entry.top_logprobs) for entry in streamed_entries] == [ + ("hi", -0.5, []), + ("<|return|>", -0.5, []), + ] + + +FINISH_ONLY_LOGPROBS_STREAM_SCRIPT: Final = """ +import json +import sys + +sys.path.insert(0, sys.argv[1]) + +import httpx +from openai import OpenAI + +import litellm + +upstream_choices = ( + {"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}, + { + "index": 0, + "delta": {}, + "logprobs": {"content": [{"token": "<|return|>", "logprob": -0.5, "top_logprobs": []}]}, + "finish_reason": "stop", + }, +) +body = "".join( + "data: " + + json.dumps({"id": "c", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [choice]}) + + "\\n\\n" + for choice in upstream_choices +) + "data: [DONE]\\n\\n" +client = OpenAI( + api_key="unused", + base_url="https://provider.example/v1", + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"}) + ) + ), +) +stream = litellm.completion( + model="openai/m", messages=[{"role": "user", "content": "hi"}], stream=True, logprobs=True, client=client +) +print(json.dumps([json.loads(chunk.model_dump_json(exclude_none=True)) for chunk in stream])) +""" + + +def test_stream_with_logprobs_only_on_finish_chunk_serializes_in_a_fresh_process() -> None: + result: Final = subprocess.run( + [sys.executable, "-I", "-c", FINISH_ONLY_LOGPROBS_STREAM_SCRIPT, str(Path(litellm.__file__).parent.parent)], + capture_output=True, + text=True, + timeout=120, + env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + ) + + assert result.returncode == 0, result.stdout + result.stderr + wire_choices: Final = tuple( + itertools.chain.from_iterable(chunk["choices"] for chunk in json.loads(result.stdout.strip().splitlines()[-1])) + ) + finish_choices: Final = [choice for choice in wire_choices if choice.get("finish_reason")] + assert [choice["finish_reason"] for choice in finish_choices] == ["stop"] + assert [entry["token"] for entry in finish_choices[0]["logprobs"]["content"]] == ["<|return|>"] + + def _hosted_vllm_stream_wrapper(logging_obj: Logging, error_payload: dict) -> CustomStreamWrapper: """A CustomStreamWrapper over the real OpenAI-compatible line iterator, fed an HTTP 200 SSE body that carries an in-body error payload the way diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 1aec7d5b761..960616e2631 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29161,7 +29161,7 @@ export interface components { /** Token */ token: string; /** Top Logprobs */ - top_logprobs: components["schemas"]["TopLogprob"][]; + top_logprobs?: components["schemas"]["TopLogprob"][]; } & { [key: string]: unknown; }; @@ -29292,6 +29292,8 @@ export interface components { ChoiceLogprobs: { /** Content */ content?: components["schemas"]["ChatCompletionTokenLogprob"][] | null; + /** Refusal */ + refusal?: components["schemas"]["ChatCompletionTokenLogprob"][] | null; } & { [key: string]: unknown; }; From 87fec5467fb7a965a2cf3998fb542c3beab17234 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Fri, 30 Oct 2026 05:07:00 -0700 Subject: [PATCH 22/64] fix(sso): keep CLI SSO and Claude Code gateway team lists in the user record's order (#49033) * fix(sso): bind CLI SSO and Claude Code gateway logins to the first team on the user record * fix(proxy): keep SCIM and JWT team sync in record order, and look up SSO groups once SCIM PATCH and JWT team sync saved a user's teams from a set, so the first team on the record, the one a Claude Code gateway sign-in binds, could change on any sync. Both now keep the record's order and append new teams in the order the IdP sent them. SSO group sync now drops claim groups that are not LiteLLM teams with one lookup before adding the missing teams one at a time, so a claim full of non-team groups no longer costs one failed add per group on every sign-in. * test(sso): expect the gateway team dropdown in the user record's order * fix(sso): read team existence from the writer, keep new-user and SCIM replace team order * fix(scim): keep the teams a user already has in record order on PATCH and PUT * fix(scim): keep record team order when a SCIM POST matches an existing email * fix(scim): join new teams in request order so a failed roster write keeps it * fix(scim): let a groups replace win over later group ops in a PATCH again Applying group adds and removes after a replace made a PATCH that deactivates a user and adds a team that does not exist answer 404 and leave the user active, where it answered 200 before. Keep the record's team order, but ignore group ops after the last replace as the base did. * test(sso): type the CLI team detail row helper --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: rj-cairn <349020452+rj-cairn@users.noreply.github.com> --- litellm/proxy/auth/handle_jwt.py | 23 +- .../internal_user_endpoints.py | 36 +-- .../management_endpoints/scim/scim_v2.py | 83 +++-- litellm/proxy/management_endpoints/ui_sso.py | 55 +++- .../test_claude_code_gateway_oauth.py | 52 +++ tests/unit/proxy/auth/test_handle_jwt.py | 25 ++ .../scim/test_scim_patch_user.py | 73 ++++- .../scim/test_scim_v2_endpoints.py | 146 +++++++++ .../test_internal_user_endpoints.py | 59 ++++ .../proxy/management_endpoints/test_ui_sso.py | 300 +++++++++++------- 10 files changed, 661 insertions(+), 191 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 7bd7b1226b1..9e11179c686 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -2199,15 +2199,19 @@ class JWTAuthManager: # not mistaken for claimless and left with stale DB memberships the fallback # could later attribute. With the flag off, keep the upstream plural-only # reconciliation so existing deployments are unchanged. - jwt_team_ids: Final = set( - jwt_handler.get_all_jwt_team_ids(jwt_valid_token) - if jwt_handler.litellm_jwtauth.fallback_to_db_teams - else jwt_handler.get_team_ids_from_jwt(jwt_valid_token) + claim_team_ids: Final = tuple( + dict.fromkeys( + jwt_handler.get_all_jwt_team_ids(jwt_valid_token) + if jwt_handler.litellm_jwtauth.fallback_to_db_teams + else jwt_handler.get_team_ids_from_jwt(jwt_valid_token) + ) ) - existing_teams: Final = set(user_object.teams or []) - teams_to_add: Final = jwt_team_ids - existing_teams + jwt_team_ids: Final = frozenset(claim_team_ids) + existing_team_order: Final = tuple(dict.fromkeys(user_object.teams or [])) + existing_teams: Final = frozenset(existing_team_order) + teams_to_add: Final = tuple(team_id for team_id in claim_team_ids if team_id not in existing_teams) preserve_db_teams_without_claims: Final = jwt_handler.litellm_jwtauth.fallback_to_db_teams and not jwt_team_ids - teams_to_remove: Final = set() if preserve_db_teams_without_claims else existing_teams - jwt_team_ids + teams_to_remove: Final = frozenset() if preserve_db_teams_without_claims else existing_teams - jwt_team_ids if teams_to_add or teams_to_remove: from litellm.proxy.management_endpoints.scim.scim_v2 import ( patch_team_membership, @@ -2218,7 +2222,10 @@ class JWTAuthManager: teams_ids_to_add_user_to=list(teams_to_add), teams_ids_to_remove_user_from=list(teams_to_remove), ) - user_object.teams = list(jwt_team_ids) + user_object.teams = [ + *(team_id for team_id in existing_team_order if team_id not in teams_to_remove), + *teams_to_add, + ] if user_api_key_cache is not None: await user_api_key_cache.async_set_cache( key=user_object.user_id, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 6ceab27befa..6950cf08a25 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -13,6 +13,7 @@ These are members of a Team on LiteLLM """ import asyncio +import contextlib import json import traceback from collections.abc import Awaitable, Mapping, Sequence @@ -443,30 +444,25 @@ async def add_new_user_to_default_team( teams: Sequence[str | NewUserRequestTeam], prisma_client: "PrismaClient", ): - tasks: Final[list[Awaitable[object]]] = [] - for team in teams: - user_role: Literal["user", "admin"] = "user" - max_budget_in_team: float | None = None - if isinstance(team, str): - team_id = team - elif isinstance(team, NewUserRequestTeam): - team_id = team.team_id - user_role = team.user_role - max_budget_in_team = team.max_budget_in_team - else: - raise ValueError(f"Invalid team type: {type(team)}") - - tasks.append( - _add_user_to_team( + memberships: Final = tuple(_default_team_membership(team) for team in teams) + for membership in memberships: + with contextlib.suppress(Exception): + await _add_user_to_team( user_id=user_id, - team_id=team_id, + team_id=membership.team_id, user_email=user_email, user_api_key_dict=user_api_key_dict, - max_budget_in_team=max_budget_in_team, - user_role=user_role, + max_budget_in_team=membership.max_budget_in_team, + user_role=membership.user_role, ) - ) - await asyncio.gather(*tasks, return_exceptions=True) + + +def _default_team_membership(team: str | NewUserRequestTeam) -> NewUserRequestTeam: + if isinstance(team, str): + return NewUserRequestTeam(team_id=team) + if isinstance(team, NewUserRequestTeam): + return team + raise ValueError(f"Invalid team type: {type(team)}") async def _fetch_user_team_ids(user_id: str, prisma_client: "PrismaClient") -> tuple[str, ...]: diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 668bd64227b..e48f68a2d0a 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -229,7 +229,11 @@ class UserProvisionerHelpers: return None requested_teams: Final = list(dict.fromkeys(new_user_request.teams or [])) - new_teams: Final = requested_teams if requested_teams else list(existing_user.teams or []) + new_teams: Final = ( + list(_keep_record_team_order(existing_user.teams or [], requested_teams)) + if requested_teams + else list(existing_user.teams or []) + ) if new_user_request.user_id != existing_user.user_id: verbose_proxy_logger.info( @@ -1165,7 +1169,7 @@ async def _handle_team_membership_changes( existing_teams_set: Final = set(existing_teams) new_teams_set: Final = set(new_teams) - teams_to_add: Final = new_teams_set - existing_teams_set + teams_to_add: Final = tuple(team for team in dict.fromkeys(new_teams) if team not in existing_teams_set) teams_to_remove: Final = existing_teams_set - new_teams_set if teams_to_add or teams_to_remove: @@ -1948,7 +1952,11 @@ async def update_user( # SCIM User.groups is readOnly (RFC 7643 4.1.2): IdPs sync membership via /Groups and send # no groups or `groups: []` on profile PUTs, so empty means unspecified, not "remove from every team" - target_teams: Final = user_data["teams"] or existing_user.teams + target_teams: Final = ( + list(_keep_record_team_order(existing_user.teams, user_data["teams"])) + if user_data["teams"] + else existing_user.teams + ) await _handle_team_membership_changes( user_id=user_id, existing_teams=existing_user.teams, @@ -2241,19 +2249,29 @@ def _handle_name_update(path: str, op_type: str, value: object, scim_metadata: d scim_metadata["familyName"] = str(value) -def _handle_group_operations(op_type: str, value: object, teams_set: set[str], path: str | None) -> set[str] | None: - """Handle group/team membership operations.""" - group_values = _extract_group_values(value) - if not group_values and value is None: - group_values = _extract_ids_from_path_filter(path, "groups") - if op_type == "replace": - return set(group_values) - elif op_type == "add": - teams_set.update(group_values) - elif op_type == "remove": - for gid in group_values: - teams_set.discard(gid) - return None +def _is_group_path(path: str | None) -> bool: + return isinstance(path, str) and path.lower().startswith("groups") + + +def _apply_group_operation(teams: tuple[str, ...], op: SCIMPatchOperation) -> tuple[str, ...]: + listed: Final = _extract_group_values(op.value) + group_values: Final = listed if listed or op.value is not None else _extract_ids_from_path_filter(op.path, "groups") + if op.op == "replace": + return tuple(dict.fromkeys(group_values)) + if op.op == "add": + return tuple(dict.fromkeys((*teams, *group_values))) + removed: Final = frozenset(group_values) + return tuple(team for team in teams if team not in removed) + + +def _keep_record_team_order(record_teams: Sequence[str], target_teams: Sequence[str]) -> tuple[str, ...]: + target: Final = frozenset(target_teams) + record: Final = tuple(dict.fromkeys(record_teams)) + known: Final = frozenset(record) + return ( + *(team for team in record if team in target), + *(team for team in dict.fromkeys(target_teams) if team not in known), + ) def _multi_valued_attribute_base(path: str) -> str: @@ -2520,8 +2538,8 @@ def _handle_generic_metadata(path: str, op_type: str, value: object, metadata: d def _apply_patch_ops( existing_user: LiteLLM_UserTable, patch_ops: SCIMPatchOp, -) -> tuple[dict[str, object], set[str]]: - """Apply patch operations and return update data and final team set.""" +) -> tuple[dict[str, object], tuple[str, ...]]: + """Apply patch operations and return update data and the final teams in order.""" update_data: Final[dict[str, object]] = {} stored_metadata: Final[object] = TypeAdapter(object).validate_python(existing_user.model_dump().get("metadata")) metadata: Final = _JSON_OBJECT_ADAPTER.validate_python(stored_metadata) if isinstance(stored_metadata, dict) else {} @@ -2541,9 +2559,6 @@ def _apply_patch_ops( if isinstance(enterprise, _InvalidEnterpriseValue): _raise_invalid_enterprise_value(enterprise) - teams_set: Final[set[str]] = set(existing_user.teams or []) - replace_team_set: set[str] | None = None - for op in patch_ops.Operations: path = op.path.lower() if isinstance(op.path, str) else "" value = op.value @@ -2585,16 +2600,20 @@ def _apply_patch_ops( _handle_name_update(path, op_type, value, scim_metadata) elif _multi_valued_attribute_base(path) in SCIM_MULTI_VALUED_ATTRIBUTE_METADATA_KEYS: _handle_multi_valued_attribute_update(path, op_type, value, metadata) - elif path.startswith("groups"): - new_replace_set = _handle_group_operations(op_type, value, teams_set, op.path) - if new_replace_set is not None: - replace_team_set = new_replace_set - elif _enterprise_extension_sub_path(path) is not None: + elif _is_group_path(path) or _enterprise_extension_sub_path(path) is not None: continue else: _handle_generic_metadata(path, op_type, value, metadata) - final_team_set: Final = replace_team_set if replace_team_set is not None else teams_set + record_teams: Final = tuple(existing_user.teams or []) + group_ops: Final = tuple(op for op in patch_ops.Operations if _is_group_path(op.path)) + last_replace: Final = next((op for op in reversed(group_ops) if op.op == "replace"), None) + final_teams: Final = _keep_record_team_order( + record_teams, + _apply_group_operation(record_teams, last_replace) + if last_replace is not None + else reduce(_apply_group_operation, group_ops, record_teams), + ) final_metadata: Final = ( { **{key: metadata_value for key, metadata_value in metadata.items() if key != SCIM_ENTERPRISE_METADATA_KEY}, @@ -2605,7 +2624,7 @@ def _apply_patch_ops( else {**metadata, "scim_metadata": scim_metadata} ) update_data["metadata"] = final_metadata - return update_data, final_team_set + return update_data, final_teams def _is_user_not_in_team_error(exc: HTTPException) -> bool: @@ -2758,7 +2777,7 @@ async def patch_user( prev_active: Final = _user_scim_active(existing_user) - update_data, final_team_set = _apply_patch_ops( + update_data, final_teams = _apply_patch_ops( existing_user=existing_user, patch_ops=patch_ops, ) @@ -2770,15 +2789,15 @@ async def patch_user( await _handle_team_membership_changes( user_id=user_id, existing_teams=existing_user.teams or [], - new_teams=list(final_team_set), + new_teams=list(final_teams), ) - update_data["teams"] = list(final_team_set) + update_data["teams"] = list(final_teams) admin_group: Final = await _get_scim_admin_group() if admin_group is not None: update_data["user_role"] = _resolve_scim_user_role( - await _scim_groups_from_team_ids(prisma_client, list(final_team_set)), + await _scim_groups_from_team_ids(prisma_client, list(final_teams)), admin_group, _default_scim_user_role(), ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index a980b9c33ff..59c2df91d77 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -8,7 +8,6 @@ Has all /sso/* routes /sso/debug/callback - returns the OpenID object returned by the SSO provider """ -import asyncio import base64 import hashlib import inspect @@ -151,6 +150,7 @@ from litellm.proxy.utils import ( get_custom_url, get_server_root_path, ) +from litellm.repositories.chunked_in import find_many_in from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import SSOConfigRepository from litellm.repositories.team_repository import TeamRepository @@ -1855,7 +1855,7 @@ async def warn_if_id_jag_capture_gap(*, retention_enabled: RetentionCheck | None verbose_proxy_logger.warning("SSO debug callback ran with an oauth2_id_jag capture gap: %s", gap) -async def create_team_member_add_task(team_id, user_info): +async def create_team_member_add_task(team_id: str, user_info: NewUserResponse | LiteLLM_UserTable): """Create a task for adding a member to a team.""" try: member: Final = Member(user_id=user_info.user_id, role="user") @@ -1871,22 +1871,43 @@ async def create_team_member_add_task(team_id, user_info): verbose_proxy_logger.debug("[Non-Blocking] Error trying to add sso user to db: %s", e) -async def add_missing_team_member(user_info: NewUserResponse | LiteLLM_UserTable, sso_teams: list[str]): +async def find_existing_team_ids(team_ids: Sequence[str]) -> frozenset[str] | None: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return None + try: + rows: Final = await find_many_in(TeamRepository(prisma_client, use_writer=True).table, "team_id", team_ids) + except Exception as e: + verbose_proxy_logger.debug("[Non-Blocking] Error looking up SSO teams: %s", e) + return None + return frozenset(row.team_id for row in rows) + + +async def add_missing_team_member( + user_info: NewUserResponse | LiteLLM_UserTable, + sso_teams: Sequence[str], + *, + add_team_member: Callable[ + [str, NewUserResponse | LiteLLM_UserTable], Awaitable[object] + ] = create_team_member_add_task, + existing_team_ids: Callable[[Sequence[str]], Awaitable[frozenset[str] | None]] = find_existing_team_ids, +) -> None: """ - Get missing teams (diff b/w user_info.team_ids and sso_teams) - - Add missing user to missing teams + - Drop the IdP groups that are not teams with one lookup, since group claims often carry many of them + - Add missing user to missing teams, one at a time in sso_teams order, so user.teams keeps the claim's order """ - # Handle None as empty list for new users - user_teams: Final = user_info.teams if user_info.teams is not None else [] - missing_teams: Final = set(sso_teams) - set(user_teams) - missing_teams_list: Final = list(missing_teams) - tasks = [] - tasks = [create_team_member_add_task(team_id, user_info) for team_id in missing_teams_list] - - try: - await asyncio.gather(*tasks) - except Exception as e: - verbose_proxy_logger.debug("[Non-Blocking] Error trying to add sso user to db: %s", e) + held_teams: Final = frozenset(user_info.teams or ()) + missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(sso_teams) if team_id not in held_teams) + if not missing_teams: + return + existing: Final = await existing_team_ids(missing_teams) + teams_to_add: Final = ( + missing_teams if existing is None else tuple(team_id for team_id in missing_teams if team_id in existing) + ) + for team_id in teams_to_add: + await add_team_member(team_id, user_info) def get_disabled_non_admin_personal_key_creation(): @@ -2451,7 +2472,9 @@ async def fetch_cli_sso_team_details( except Exception as e: verbose_proxy_logger.error("Error fetching team details for CLI SSO session: %s", e) return None - return tuple(_cli_sso_team_detail(team_row.model_dump()) for team_row in prisma_teams) + fetched_details: Final = tuple(_cli_sso_team_detail(team_row.model_dump()) for team_row in prisma_teams) + details_by_team_id: Final = MappingProxyType({detail.team_id: detail for detail in fetched_details}) + return tuple(details_by_team_id[team_id] for team_id in dict.fromkeys(teams) if team_id in details_by_team_id) def _cli_sso_session_teams(team_details: Sequence[CliSsoTeamDetail]) -> list[str]: diff --git a/tests/integration/authorization/test_claude_code_gateway_oauth.py b/tests/integration/authorization/test_claude_code_gateway_oauth.py index 4dedbdcd8e4..23476c439de 100644 --- a/tests/integration/authorization/test_claude_code_gateway_oauth.py +++ b/tests/integration/authorization/test_claude_code_gateway_oauth.py @@ -636,3 +636,55 @@ def test_device_flow_mints_the_session_for_the_team_the_developer_picks_in_the_b served_models: Final = {model.id for model in sdk.models.list()} assert served_models == {team_models[chosen_team]} assert len(identity_provider.drain()) == 3 + + +@pytest.mark.timeout(300) +def test_gateway_login_lists_and_defaults_to_the_first_team_on_the_user_record_after_a_write_to_its_row( + enabled_gateway: Gateway, identity_provider: Wire +) -> None: + run_id: Final = uuid.uuid4().hex + with enabled_gateway.scenario() as scenario: + user_id: Final = scenario.user(user_id=_USER_ID, user_email=_USER_EMAIL, user_role="internal_user") + first_team: Final = scenario.team(team_id=f"zeta-{run_id}") + second_team: Final = scenario.team(team_id=f"alpha-{run_id}") + for team_id in (first_team, second_team): + enabled_gateway.post( + "/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}} + ) + enabled_gateway.post("/team/update", {"team_id": first_team, "metadata": {"updated_by_run": run_id}}) + assert enabled_gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (user_id,)) == [ + {"teams": [first_team, second_team]} + ] + + authorization_response: Final = enabled_gateway.client.post( + "/claude_code_gateway/oauth/device_authorization", data={"client_id": _CLIENT_ID} + ) + assert authorization_response.status_code == 200, authorization_response.text + authorization: Final = _DeviceAuthorization.model_validate_json(authorization_response.content) + verification: Final = enabled_gateway.client.get(authorization.verification_uri, follow_redirects=False) + assert verification.status_code == 303, verification.text + identity_provider_redirect: Final = enabled_gateway.client.get( + verification.headers["location"], follow_redirects=False + ) + assert identity_provider_redirect.status_code == 302, identity_provider_redirect.text + callback: Final = enabled_gateway.client.get( + identity_provider_redirect.headers["location"], follow_redirects=False + ) + assert callback.status_code == 200, callback.text + assert tuple(_TEAM_OPTION.findall(callback.text)) == (first_team, second_team), callback.text + browser_token: Final = _BROWSER_TOKEN.search(callback.text) + assert browser_token is not None, callback.text + completion_without_a_pick: Final = enabled_gateway.client.post( + f"/sso/cli/complete/{authorization.device_code.partition('.')[0]}", + data={ + "user_code": authorization.user_code, + "browser_complete_token": html.unescape(browser_token.group(1)), + }, + ) + assert completion_without_a_pick.status_code == 200, completion_without_a_pick.text + token_response: Final = _device_code_grant(enabled_gateway, authorization.device_code) + assert token_response.status_code == 200, token_response.text + token: Final = _AccessToken.model_validate_json(token_response.content) + identity_provider.drain() + + assert (token.user_id, token.team_id) == (user_id, first_team) diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index 51cf15996cd..4c223edc20d 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -913,6 +913,31 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_team_change(): assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable +@pytest.mark.asyncio +async def test_sync_user_role_and_teams_keeps_record_order_and_adds_claim_teams_in_claim_order(): + jwt_handler: Final = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=AsyncMock(), + litellm_jwtauth=LiteLLM_JWTAuth(team_ids_jwt_field="groups", sync_user_role_and_teams=True), + ) + cache: Final = AsyncMock() + user: Final = LiteLLM_UserTable(user_id="u1", teams=["team-zeta", "team-kappa", "team-alpha"]) + token: Final = {"groups": ["team-mu", "team-alpha", "team-beta", "team-mu", "team-zeta", "team-omega"]} + + with patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + new_callable=AsyncMock, + ) as mock_patch: + await JWTAuthManager.sync_user_role_and_teams(jwt_handler, token, user, AsyncMock(), user_api_key_cache=cache) + + assert mock_patch.await_args.kwargs["teams_ids_to_add_user_to"] == ["team-mu", "team-beta", "team-omega"] + assert mock_patch.await_args.kwargs["teams_ids_to_remove_user_from"] == ["team-kappa"] + expected_teams: Final = ["team-zeta", "team-alpha", "team-mu", "team-beta", "team-omega"] + assert user.teams == expected_teams + assert cache.async_set_cache.await_args.kwargs["value"].teams == expected_teams + + @pytest.mark.asyncio async def test_sync_user_role_and_teams_no_cache_write_when_nothing_changes(): """Test that cache is NOT written when role and teams already match.""" diff --git a/tests/unit/proxy/management_endpoints/scim/test_scim_patch_user.py b/tests/unit/proxy/management_endpoints/scim/test_scim_patch_user.py index 31c5a54a753..1901e0efd4b 100644 --- a/tests/unit/proxy/management_endpoints/scim/test_scim_patch_user.py +++ b/tests/unit/proxy/management_endpoints/scim/test_scim_patch_user.py @@ -502,9 +502,9 @@ def test_apply_patch_ops_remove_group_filtered_path_without_value(): Operations=[SCIMPatchOperation(op="remove", path='groups[value eq "team-1"]')] ) - _, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) + _, final_teams = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) - assert final_team_set == {"team-2"} + assert final_teams == ("team-2",) def test_apply_patch_ops_add_group_filtered_path_without_value(): @@ -519,9 +519,9 @@ def test_apply_patch_ops_add_group_filtered_path_without_value(): Operations=[SCIMPatchOperation(op="add", path="groups[value eq 'team-3']")] ) - _, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) + _, final_teams = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) - assert final_team_set == {"team-1", "team-3"} + assert final_teams == ("team-1", "team-3") def test_apply_patch_ops_replace_groups_empty_value_does_not_use_path_filter(): @@ -539,9 +539,70 @@ def test_apply_patch_ops_replace_groups_empty_value_does_not_use_path_filter(): ] ) - _, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) + _, final_teams = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) - assert final_team_set == set() + assert final_teams == () + + +@pytest.mark.parametrize( + "operations,expected_teams", + [ + pytest.param( + [ + SCIMPatchOperation(op="add", path="groups", value=[{"value": "team-mu"}, {"value": "team-beta"}]), + SCIMPatchOperation(op="remove", path='groups[value eq "team-kappa"]'), + SCIMPatchOperation(op="add", path="groups", value=[{"value": "team-alpha"}, {"value": "team-omega"}]), + ], + ("team-zeta", "team-alpha", "team-delta", "team-mu", "team-beta", "team-omega"), + id="adds_follow_the_record_in_op_order", + ), + pytest.param( + [ + SCIMPatchOperation( + op="replace", + path="groups", + value=[{"value": "team-omega"}, {"value": "team-delta"}, {"value": "team-beta"}], + ) + ], + ("team-delta", "team-omega", "team-beta"), + id="replace_keeps_teams_the_user_has_in_record_order_and_adds_new_ones_in_request_order", + ), + pytest.param( + [ + SCIMPatchOperation( + op="replace", + path="groups", + value=[{"value": "team-omega"}, {"value": "team-delta"}, {"value": "team-beta"}], + ), + SCIMPatchOperation(op="add", path="groups", value=[{"value": "team-mu"}]), + SCIMPatchOperation(op="remove", path='groups[value eq "team-delta"]'), + ], + ("team-delta", "team-omega", "team-beta"), + id="a_replace_wins_over_later_group_ops", + ), + pytest.param( + [ + SCIMPatchOperation(op="remove", path='groups[value eq "team-zeta"]'), + SCIMPatchOperation(op="add", path="groups", value=[{"value": "team-zeta"}]), + ], + ("team-zeta", "team-alpha", "team-kappa", "team-delta"), + id="removing_and_re_adding_a_team_keeps_its_place", + ), + ], +) +def test_apply_patch_ops_keeps_team_order_so_the_first_team_stays_first( + operations: list[SCIMPatchOperation], expected_teams: tuple[str, ...] +): + user: Final = LiteLLM_UserTable( + user_id="user-order", + user_email="order@example.com", + teams=["team-zeta", "team-alpha", "team-kappa", "team-delta"], + metadata={}, + ) + + _, final_teams = _apply_patch_ops(existing_user=user, patch_ops=SCIMPatchOp(Operations=operations)) + + assert final_teams == expected_teams @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 6014c9229f6..180362f45f1 100644 --- a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -25,6 +25,7 @@ from litellm.proxy._types import ( NewUserResponse, ProxyErrorTypes, ProxyException, + TeamMemberAddRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.scim.scim_v2 import ( @@ -941,6 +942,57 @@ async def test_handle_existing_user_by_email_syncs_roster_and_dedups_teams(mocke assert update_calls[0].kwargs["data"]["teams"] == ["team-a", "team-b"] +class _EmailMatchedUserRow: + def __init__(self, user: LiteLLM_UserTable) -> None: + self.user = user + + async def find_first(self, *, where: dict[str, object]) -> LiteLLM_UserTable: + return self.user + + async def update(self, *, where: dict[str, object], data: dict[str, object]) -> LiteLLM_UserTable: + self.user = self.user.model_copy(update=data) + return self.user + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_keeps_record_team_order(mocker: MockerFixture): + user_rows: Final = _EmailMatchedUserRow( + LiteLLM_UserTable( + user_id="same-id", + user_email="member@example.com", + user_alias="Member", + teams=["team-zeta", "team-alpha", "team-kappa"], + metadata={}, + ) + ) + mock_prisma_client: Final = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = user_rows + + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=None), + ) + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, + new_user_request=NewUserRequest( + user_id="same-id", + user_email="member@example.com", + user_alias="Member", + teams=["team-mu", "team-kappa", "team-beta", "team-zeta"], + metadata={}, + auto_create_key=False, + ), + ) + + assert user_rows.user.teams == ["team-zeta", "team-kappa", "team-mu", "team-beta"] + + @pytest.mark.asyncio async def test_handle_existing_user_by_email_without_teams_preserves_memberships(mocker): """Adoption via POST /Users without ``groups`` must keep the user's existing teams. @@ -1290,6 +1342,49 @@ async def test_handle_team_membership_changes_add_and_remove(mocker): assert call_args[1]["teams_ids_to_remove_user_from"] == ["team1"] +class _TeamRosterRecord: + def __init__(self, missing_team: str) -> None: + self.missing_team = missing_team + self.teams: tuple[str, ...] = () + + async def team_member_add(self, *, data: TeamMemberAddRequest, user_api_key_dict: UserAPIKeyAuth) -> None: + if data.team_id == self.missing_team: + raise HTTPException(status_code=404, detail={"error": "Team not found"}) + self.teams = (*self.teams, data.team_id) + + +@pytest.mark.asyncio +async def test_handle_team_membership_changes_adds_teams_in_request_order_when_one_fails(mocker: MockerFixture): + requested_teams: Final = ( + "team-zeta", + "team-alpha", + "team-kappa", + "team-kept", + "team-mu", + "team-missing", + "team-beta", + "team-omega", + "team-delta", + "team-sigma", + "team-gamma", + "team-tau", + ) + roster: Final = _TeamRosterRecord(missing_team="team-missing") + mocker.patch( # test-quality-ok: team_member_add is module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + roster.team_member_add, + ) + + with pytest.raises(SCIMRosterSyncError): + await _handle_team_membership_changes( + user_id="test-user", + existing_teams=["team-kept"], + new_teams=list(requested_teams), + ) + + assert roster.teams == tuple(team for team in requested_teams if team not in {"team-kept", "team-missing"}) + + @pytest.mark.asyncio async def test_update_user_success(mocker): """Should successfully update user with PUT request""" @@ -1492,6 +1587,57 @@ async def test_update_user_without_groups_preserves_memberships_and_role(mocker, assert "user_role" not in update_data +@pytest.mark.asyncio +async def test_update_user_put_keeps_the_record_order_of_teams_the_user_already_has( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.proxy_server import proxy_config + + async def no_scim_admin_group() -> dict[str, object]: + return {} + + monkeypatch.setattr(proxy_config, "get_config", no_scim_admin_group) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + existing_user: Final = mocker.MagicMock() + existing_user.teams = ["team-zeta", "team-alpha", "team-kappa"] + existing_user.metadata = {} + scim_user: Final = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="order-user", + emails=[SCIMUserEmail(value="order@example.com")], + groups=[ + SCIMUserGroup(value="team-mu"), + SCIMUserGroup(value="team-kappa"), + SCIMUserGroup(value="team-beta"), + SCIMUserGroup(value="team-zeta"), + ], + ) + mock_prisma_client: Final = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "order-user"}) + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + mocker.patch( # test-quality-ok: roster writes are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + + await update_user(user_id="order-user", user=scim_user) + + update_data: Final = mock_prisma_client.db.litellm_usertable.update.call_args.kwargs["data"] + assert update_data["teams"] == ["team-zeta", "team-kappa", "team-mu", "team-beta"] + + @pytest.mark.asyncio async def test_update_user_not_found(mocker): """Should raise 404 when user doesn't exist""" diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 42162087c9b..39586a1a204 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4195,6 +4195,65 @@ async def test_add_new_user_to_default_team_string_teams_have_no_member_budget(m assert mock_add.call_args.kwargs["team_id"] == "string-team" +@pytest.mark.asyncio +async def test_add_new_user_to_default_team_joins_teams_in_the_order_given_and_past_a_failing_team( + mocker: MockerFixture, +) -> None: + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + add_new_user_to_default_team, + ) + + requested: Final = ("team-zeta", "team-broken", "team-mu", "team-alpha") + seconds_per_team: Final = {"team-zeta": 0.03, "team-broken": 0.0, "team-mu": 0.02, "team-alpha": 0.0} + finished: Final = mocker.MagicMock() + + async def join_team(*, team_id: str, **_: object) -> None: + await asyncio.sleep(seconds_per_team[team_id]) + if team_id == "team-broken": + raise RuntimeError("team row locked") + finished(team_id) + + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._add_user_to_team", + side_effect=join_team, + ) + + await add_new_user_to_default_team( + user_id="scim-user", + user_email=None, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + teams=list(requested), + prisma_client=mocker.MagicMock(), + ) + + assert [call.args[0] for call in finished.call_args_list] == ["team-zeta", "team-mu", "team-alpha"] + + +@pytest.mark.asyncio +async def test_add_new_user_to_default_team_rejects_an_invalid_team_before_joining_any( + mocker: MockerFixture, +) -> None: + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + add_new_user_to_default_team, + ) + + join_team: Final = mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._add_user_to_team", + new_callable=mocker.AsyncMock, + ) + + with pytest.raises(ValueError, match="Invalid team type"): + await add_new_user_to_default_team( + user_id="new-user", + user_email=None, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + teams=["team-zeta", {"team_id": "team-mu"}], + prisma_client=mocker.MagicMock(), + ) + + join_team.assert_not_awaited() + + @pytest.mark.asyncio async def test_add_user_to_team_logs_unknown_team_at_error(mocker, caplog): """A default team that no longer exists makes every membership write 404. diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index cfb64a85359..136cc86e67a 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -4049,8 +4049,8 @@ class TestCLIKeyRegenerationFlow: assert ( '" ) in page @@ -4530,6 +4530,39 @@ class TestCLIKeyRegenerationFlow: == () ) + @pytest.mark.asyncio + async def test_fetch_cli_sso_team_details_follows_the_user_record_order(self): + """The CLI login binds to the first detail, so it must be the first team on the user record. + + Postgres returns the rows in whatever order the table holds them, which every write to a + team row changes, so the bound team used to flip between sign-ins. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + fetch_cli_sso_team_details, + ) + + def team_row(team_id: str) -> MagicMock: + row: Final = MagicMock() + row.model_dump.return_value = { + "team_id": team_id, + "team_alias": None, + "models": [], + "litellm_model_table": None, + } + return row + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[team_row("team-alpha"), team_row("team-zeta")] + ) + + details: Final = await fetch_cli_sso_team_details( + prisma_client=prisma_client, + teams=["team-zeta", "team-gone", "team-alpha", "team-zeta"], + ) + + assert [detail.team_id for detail in details] == ["team-zeta", "team-alpha"] + @pytest.mark.asyncio async def test_cli_poll_key_mints_jwt_with_selected_team_grants(self): """The selected team's grants must reach the mint, not just its alias.""" @@ -6692,72 +6725,57 @@ class TestPKCEFunctionality: class TestAddMissingTeamMember: """Tests for the add_missing_team_member function""" + @staticmethod + def _recording_add() -> tuple[ + list[tuple[str, str | None]], + Callable[[str, NewUserResponse | LiteLLM_UserTable], Awaitable[None]], + ]: + added: Final[list[tuple[str, str | None]]] = [] + + async def add(team_id: str, user_info: NewUserResponse | LiteLLM_UserTable) -> None: + added.append((team_id, user_info.user_id)) + + return added, add + @pytest.mark.asyncio async def test_add_missing_team_member_with_new_user_response_teams_none(self): - """ - Bug reproduction: When a NewUserResponse has teams=None (new SSO user), - add_missing_team_member() should still add the user to the SSO teams. - - Currently FAILS: The function returns early when teams is None. - """ + """A new SSO user's NewUserResponse has teams=None, and the user must still be added to the SSO teams.""" from litellm.proxy._types import NewUserResponse from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member - # Simulate a new SSO user - NewUserResponse has teams=None by default - new_user = NewUserResponse( - user_id="new-sso-user-123", - key="sk-xxxxx", - teams=None, # This is the default for NewUserResponse + new_user = NewUserResponse(user_id="new-sso-user-123", key="sk-xxxxx", teams=None) + added, add = self._recording_add() + + await add_missing_team_member( + user_info=new_user, + sso_teams=["team-from-entra-1", "team-from-entra-2"], + add_team_member=add, ) - sso_teams = ["team-from-entra-1", "team-from-entra-2"] - - with patch( - "litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task" - ) as mock_add_task: - mock_add_task.return_value = AsyncMock() - - await add_missing_team_member(user_info=new_user, sso_teams=sso_teams) - - # Bug: This assertion currently FAILS - no teams are added - # because function returns early when teams is None - assert ( - mock_add_task.call_count == 2 - ), f"Expected 2 calls to add user to teams, but got {mock_add_task.call_count}" - called_team_ids = [call.args[0] for call in mock_add_task.call_args_list] - assert set(called_team_ids) == { - "team-from-entra-1", - "team-from-entra-2", - } + assert added == [ + ("team-from-entra-1", "new-sso-user-123"), + ("team-from-entra-2", "new-sso-user-123"), + ] @pytest.mark.asyncio async def test_add_missing_team_member_with_litellm_user_table_empty_teams(self): - """ - Control test: When a LiteLLM_UserTable has teams=[] (existing user, no teams), - add_missing_team_member() should add the user to SSO teams. - - This test PASSES because LiteLLM_UserTable defaults teams to [] not None. - """ + """An existing user with teams=[] is added to every SSO team.""" from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member - # Existing user has teams=[] by default (not None) - existing_user = LiteLLM_UserTable( - user_id="existing-user-456", - teams=[], # Empty list, not None + existing_user = LiteLLM_UserTable(user_id="existing-user-456", teams=[]) + added, add = self._recording_add() + + await add_missing_team_member( + user_info=existing_user, + sso_teams=["team-from-entra-1", "team-from-entra-2"], + add_team_member=add, ) - sso_teams = ["team-from-entra-1", "team-from-entra-2"] - - with patch( - "litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task" - ) as mock_add_task: - mock_add_task.return_value = AsyncMock() - - await add_missing_team_member(user_info=existing_user, sso_teams=sso_teams) - - # This PASSES - teams are added because teams=[] not None - assert mock_add_task.call_count == 2 + assert added == [ + ("team-from-entra-1", "existing-user-456"), + ("team-from-entra-2", "existing-user-456"), + ] @pytest.mark.asyncio async def test_add_user_to_teams_from_sso_response_new_user(self): @@ -6800,96 +6818,160 @@ class TestAddMissingTeamMember: @pytest.mark.asyncio async def test_sso_first_login_full_flow_adds_user_to_teams(self): - """ - End-to-end test: Simulates complete first-time SSO login with Entra groups. - Verifies teams are created AND user is added as a member. - """ + """A first-time SSO login with Entra groups adds the new user as a member of every group's team.""" from litellm.proxy._types import NewUserResponse from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member - team_member_calls = [] + new_user = NewUserResponse(user_id="first-time-sso-user", key="sk-xxxxx", teams=None) + added, add = self._recording_add() - async def track_team_member_add(team_id, user_info): - team_member_calls.append({"team_id": team_id, "user_id": user_info.user_id}) - - # New SSO user with Entra groups - new_user = NewUserResponse( - user_id="first-time-sso-user", - key="sk-xxxxx", - teams=None, # The problematic default + await add_missing_team_member( + user_info=new_user, + sso_teams=["entra-team-alpha", "entra-team-beta"], + add_team_member=add, ) - sso_teams = ["entra-team-alpha", "entra-team-beta"] - - with patch( - "litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task", - side_effect=track_team_member_add, - ): - await add_missing_team_member(user_info=new_user, sso_teams=sso_teams) - - # Bug: With current code, team_member_calls will be empty - # After fix: Should have 2 entries - assert ( - len(team_member_calls) == 2 - ), f"Expected 2 teams to be added, but got {len(team_member_calls)}" - assert {c["team_id"] for c in team_member_calls} == { - "entra-team-alpha", - "entra-team-beta", - } - assert all(c["user_id"] == "first-time-sso-user" for c in team_member_calls) + assert added == [ + ("entra-team-alpha", "first-time-sso-user"), + ("entra-team-beta", "first-time-sso-user"), + ] @pytest.mark.asyncio @pytest.mark.parametrize( - "user_info_factory,teams_value,expected_teams_added", + "user_info_factory,expected_teams_added", [ - # Bug case: NewUserResponse with teams=None pytest.param( lambda uid: NewUserResponse(user_id=uid, key="sk-xxx", teams=None), - None, - ["team-1", "team-2"], # Should still add teams + ["team-1", "team-2"], id="new_user_teams_none", ), - # Working case: LiteLLM_UserTable with teams=[] pytest.param( lambda uid: LiteLLM_UserTable(user_id=uid, teams=[]), - [], ["team-1", "team-2"], id="existing_user_empty_teams", ), - # Existing user with some teams already pytest.param( lambda uid: LiteLLM_UserTable(user_id=uid, teams=["team-1"]), - ["team-1"], - ["team-2"], # Only missing team should be added + ["team-2"], id="existing_user_partial_teams", ), + pytest.param( + lambda uid: LiteLLM_UserTable(user_id=uid, teams=["team-2", "team-1"]), + [], + id="existing_user_all_teams", + ), ], ) async def test_add_missing_team_member_handles_all_user_types( - self, user_info_factory, teams_value, expected_teams_added + self, user_info_factory, expected_teams_added ): - """ - Parametrized test ensuring add_missing_team_member works for all user types. - """ + """Only the SSO teams the user is not already on are added.""" from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member - user_info = user_info_factory("test-user-id") - sso_teams = ["team-1", "team-2"] + added, add = self._recording_add() - added_teams = [] + await add_missing_team_member( + user_info=user_info_factory("test-user-id"), + sso_teams=["team-1", "team-2"], + add_team_member=add, + ) - async def mock_create_task(team_id, user): - added_teams.append(team_id) + assert added == [(team_id, "test-user-id") for team_id in expected_teams_added] - with patch( - "litellm.proxy.management_endpoints.ui_sso.create_team_member_add_task", - side_effect=mock_create_task, - ): - await add_missing_team_member(user_info=user_info, sso_teams=sso_teams) + @pytest.mark.asyncio + async def test_add_missing_team_member_lands_teams_in_claim_order(self): + """user.teams is appended to as each add finishes, and the first team on it is the one a CLI + login binds to, so the adds must finish in the claim's order even when an earlier add is slower. - assert set(added_teams) == set( - expected_teams_added - ), f"Expected teams {expected_teams_added}, but got {added_teams}" + The adds used to run concurrently from a set, so they finished in an arbitrary order. + """ + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + claim: Final = ["team-zeta", "team-alpha", "team-held", "team-zeta", "team-mu"] + yields_before_finishing: Final = {"team-zeta": 3, "team-alpha": 2, "team-mu": 1} + finished: Final[list[str]] = [] + + async def slower_when_earlier(team_id: str, user_info: NewUserResponse | LiteLLM_UserTable) -> None: + for _ in range(yields_before_finishing[team_id]): + await asyncio.sleep(0) + finished.append(team_id) + + await add_missing_team_member( + user_info=LiteLLM_UserTable(user_id="group-sync-user", teams=["team-held"]), + sso_teams=claim, + add_team_member=slower_when_earlier, + ) + + assert finished == ["team-zeta", "team-alpha", "team-mu"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "existing,expected_teams_added", + [ + pytest.param(frozenset({"team-zeta", "team-alpha"}), ["team-zeta", "team-alpha"], id="only_teams"), + pytest.param(None, ["grp-everyone", "team-zeta", "grp-vpn", "team-alpha"], id="lookup_failed"), + ], + ) + async def test_add_missing_team_member_looks_up_missing_groups_once_and_adds_only_teams( + self, existing: frozenset[str] | None, expected_teams_added: list[str] + ): + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + added, add = self._recording_add() + lookups: Final[list[tuple[str, ...]]] = [] + + async def existing_team_ids(team_ids: Sequence[str]) -> frozenset[str] | None: + lookups.append(tuple(team_ids)) + return existing + + await add_missing_team_member( + user_info=LiteLLM_UserTable(user_id="group-sync-user", teams=["team-held"]), + sso_teams=["grp-everyone", "team-zeta", "team-held", "grp-vpn", "team-alpha"], + add_team_member=add, + existing_team_ids=existing_team_ids, + ) + + assert lookups == [("grp-everyone", "team-zeta", "grp-vpn", "team-alpha")] + assert added == [(team_id, "group-sync-user") for team_id in expected_teams_added] + + @pytest.mark.asyncio + async def test_add_missing_team_member_skips_the_lookup_when_every_group_is_held(self): + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + added, add = self._recording_add() + lookup: Final = AsyncMock(return_value=frozenset()) + + await add_missing_team_member( + user_info=LiteLLM_UserTable(user_id="group-sync-user", teams=["team-alpha", "team-zeta"]), + sso_teams=["team-zeta", "team-alpha"], + add_team_member=add, + existing_team_ids=lookup, + ) + + lookup.assert_not_awaited() + assert added == [] + + @pytest.mark.asyncio + async def test_find_existing_team_ids_reads_the_writer_and_returns_none_when_the_db_cannot_answer(self): + from litellm.proxy.management_endpoints.ui_sso import find_existing_team_ids + + stored_on_writer: Final = ("team-zeta", "team-alpha") + + async def find_many(*, where: Mapping[str, Mapping[str, Sequence[str]]]) -> list[SimpleNamespace]: + return [ + SimpleNamespace(team_id=team_id) for team_id in stored_on_writer if team_id in where["team_id"]["in"] + ] + + prisma: Final = MagicMock() + prisma.writer_db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many) + prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + assert await find_existing_team_ids(("grp-vpn", "team-alpha")) == frozenset({"team-alpha"}) + prisma.writer_db.litellm_teamtable.find_many.side_effect = RuntimeError("db unavailable") + assert await find_existing_team_ids(("team-alpha",)) is None + with patch("litellm.proxy.proxy_server.prisma_client", None): + assert await find_existing_team_ids(("team-alpha",)) is None @pytest.mark.asyncio From a44e44b98059a9f740fd721fc475d139b721f89f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 05:11:42 -0700 Subject: [PATCH 23/64] fix(cache_control): pick the chat-path dialect from the deployment's provider prefix (#45295) * fix(cache_control): resolve the chat-path dialect with the deployment api_base * test(cache_control): type the chat-path resolver helpers * fix(cache_control): pick the chat-path dialect from the deployment's provider prefix Resolving with the deployment api_base followed the dispatch's azure_ai to azure flip, so chat on a non-Foundry Azure host sent Anthropic cache_control while /v1/messages, told azure_ai, sent the OpenAI breakpoint. Use the model's provider prefix when the caller does not name a provider, so chat matches the told path whatever AZURE_AI_API_BASE holds, and keep name-only resolution everywhere else. * fix(cache_control): keep litellm_proxy routing when the chat path reads the provider prefix With USE_LITELLM_PROXY or litellm.use_litellm_proxy on, every call goes to the upstream proxy, so the provider prefix no longer decides the cache format there and resolution matches main. The plain OpenAI check reads the same provider value as the hosted check, so one function no longer mixes two lookups * fix(cache_control): resolve the chat-path dialect with the deployment api_base, as the call is routed The hosted cache-format check now resolves the provider with the deployment's own api_base when the caller names none, the same provider the call dispatches to and the Responses path passes, so AZURE_AI_API_BASE no longer decides the chat format. The plain OpenAI check keeps main's name-only lookup. * fix(cache_control): send acompletion fallbacks the caller's request, not the primary's hooked kwargs acompletion handed async_completion_with_fallbacks the kwargs the cache seed and prompt hook had already rewritten for the first model, so a Claude fallback behind a Foundry GPT model received prompt_cache_options and got a 400. Fallbacks now get a copy of the caller's model, messages, tools, and kwargs taken before the hooks, as completion already does with its args. Both cache seeds also read base_url before api_base, since completion sends the call to base_url when both are set. * fix(cache_control): keep the caller's tools list intact for acompletion fallbacks The vector-store hook pops recognized file_search tools from the request's tools list in place, so the fallback request built from the pre-hook snapshot lost the retrieved context. Snapshot a copy of the list instead. * fix(cache_control): skip acompletion's prompt hooks when fallbacks will run them Each fallback attempt runs the prompt hooks for its own model, so running them first in the outer call only cost a vector-store search whose result the hand-off threw away. * fix(cache_control): run acompletion's prompt hooks for fallbacks and leave only the cache hook to each attempt --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../anthropic_cache_control_hook.py | 10 +- litellm/main.py | 40 +++- .../test_anthropic_cache_control_hook.py | 71 +++++++ .../litellm_core_utils/test_fallback_utils.py | 187 ++++++++++++++++++ 4 files changed, 294 insertions(+), 14 deletions(-) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 45c82589d2d..d816bc661f3 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -116,12 +116,12 @@ def configured_injection_points(value: object) -> Sequence[CacheControlInjection return tuple(cast(CacheControlInjectionPoint, entry) for entry in value if isinstance(entry, dict)) -def _served_provider_for_model(model: str) -> str | None: +def _served_provider_for_model(model: str, api_base: object = None) -> str | None: from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider try: - _, provider, _, _ = get_llm_provider(model=model) + _, provider, _, _ = get_llm_provider(model=model, api_base=api_base if isinstance(api_base, str) else None) except BadRequestError: return None return provider @@ -536,7 +536,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): if model is None: return False hosted_flag: Final = _hosted_openai_dialect_flag( - model, custom_llm_provider, AnthropicCacheControlHook._resolve_provider + model, custom_llm_provider, lambda name: AnthropicCacheControlHook._resolve_provider(name, api_base) ) if hosted_flag is not None: return hosted_flag @@ -559,8 +559,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): ) @staticmethod - def _resolve_provider(model: str) -> str | None: - return _served_provider_for_model(model) + def _resolve_provider(model: str, api_base: object = None) -> str | None: + return _served_provider_for_model(model, api_base) @staticmethod def count_request_cache_breakpoints( diff --git a/litellm/main.py b/litellm/main.py index 54afe654cec..5346741e4d1 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -681,7 +681,7 @@ async def acompletion( - The `completion` function is called using `run_in_executor` to execute synchronously in the event loop. - If `stream` is True, the function returns an async generator that yields completion lines. """ - fallbacks = kwargs.get("fallbacks", None) + fallbacks: Final = kwargs.get("fallbacks", None) or litellm.model_fallbacks mock_timeout = kwargs.get("mock_timeout", None) if mock_timeout is True: @@ -697,6 +697,7 @@ async def acompletion( loop: Final = asyncio.get_event_loop() custom_llm_provider = kwargs.get("custom_llm_provider", None) _ = _resolve_control_options(kwargs, model) + caller_kwargs: Final = cast("Mapping[str, object]", {**kwargs}) # cast-ok: acompletion takes untyped kwargs ## PROMPT MANAGEMENT HOOKS ## ######################################################### @@ -715,15 +716,34 @@ async def acompletion( custom_llm_provider=cast(str | None, custom_llm_provider), # cast-ok: read from untyped kwargs tools=tools, enable_prompt_caching=cast(bool | None, kwargs.get("enable_prompt_caching")), # cast-ok: untyped kwargs - api_base=kwargs.get("api_base") or base_url, + api_base=base_url or kwargs.get("api_base"), ) - if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( - litellm_logging_obj.should_run_prompt_management_hooks( - prompt_id=kwargs.get("prompt_id", None), + prompt_id: Final = cast(str | None, kwargs.get("prompt_id", None)) # cast-ok: read from untyped kwargs + prompt_management_logger: Final = ( + litellm_logging_obj.get_custom_logger_for_prompt_management( + model=model, + non_default_params=kwargs, + tools=tools, + prompt_id=prompt_id, + dynamic_callback_params=litellm_logging_obj.standard_callback_dynamic_params, + ) + if isinstance(litellm_logging_obj, LiteLLMLoggingObj) + and litellm_logging_obj.should_run_prompt_management_hooks( + prompt_id=prompt_id, non_default_params=kwargs, tools=tools, ) + else None + ) + each_fallback_writes_its_own_cache_dialect: Final = fallbacks is not None and isinstance( + prompt_management_logger, AnthropicCacheControlHook + ) + + if ( + isinstance(litellm_logging_obj, LiteLLMLoggingObj) + and prompt_management_logger is not None + and not each_fallback_writes_its_own_cache_dialect ): ( model, @@ -733,8 +753,9 @@ async def acompletion( model=model, messages=messages, non_default_params=kwargs, - prompt_id=kwargs.get("prompt_id", None), + prompt_id=prompt_id, prompt_variables=kwargs.get("prompt_variables", None), + prompt_management_logger=prompt_management_logger, tools=tools, prompt_label=kwargs.get("prompt_label", None), prompt_version=kwargs.get("prompt_version", None), @@ -826,9 +847,10 @@ async def acompletion( litellm_params=(GenericLiteLLMParams.model_validate(_supplemental) if _supplemental else None), ) - fallbacks = fallbacks or litellm.model_fallbacks if fallbacks is not None: - response = await async_completion_with_fallbacks(**completion_kwargs, kwargs={"fallbacks": fallbacks, **kwargs}) + response = await async_completion_with_fallbacks( + **completion_kwargs, kwargs={"fallbacks": fallbacks, **caller_kwargs} + ) if response is None: raise Exception( "No response from fallbacks. Got none. Turn on `litellm.set_verbose=True` to see more details." @@ -5765,7 +5787,7 @@ def completion( custom_llm_provider=cast(str | None, kwargs.get("custom_llm_provider")), # cast-ok: untyped kwargs tools=tools, enable_prompt_caching=cast(bool | None, kwargs.get("enable_prompt_caching")), # cast-ok: untyped kwargs - api_base=kwargs.get("api_base") or base_url, + api_base=base_url or kwargs.get("api_base"), ) if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index c70c7080ab8..fa03879c05a 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -3998,6 +3998,77 @@ class TestChatPathProviderStamp: assert params["prompt_cache_options"] == {"mode": "explicit"} +class TestChatPathResolvesTheDeploymentProvider: + MODEL = "azure_ai/gpt-6-astra" + AMBIENT_API_BASE = "https://other-deployment.openai.azure.com" + FOUNDRY_API_BASE = "https://foundry.services.ai.azure.com" + UPSTREAM_PROXY_API_BASE = "https://upstream-proxy.example.com" + POINTS = [{"location": "message", "role": "system"}] + MESSAGES = [{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}] + + def _pin_ambient_api_base(self, monkeypatch: pytest.MonkeyPatch, pinned: bool) -> None: + if pinned: + monkeypatch.setenv("AZURE_AI_API_BASE", self.AMBIENT_API_BASE) + else: + monkeypatch.delenv("AZURE_AI_API_BASE", raising=False) + + def _chat( + self, custom_llm_provider: str | None, api_base: str, model: str = MODEL + ) -> tuple[list[AllMessageValues], dict[str, object]]: + seeded: Final[dict[str, object]] = {"cache_control_injection_points": copy.deepcopy(self.POINTS)} + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=seeded, + messages=copy.deepcopy(self.MESSAGES), + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + ) + _, out, params = AnthropicCacheControlHook().get_chat_completion_prompt( + model=model, + messages=copy.deepcopy(self.MESSAGES), + non_default_params=seeded, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + return out, params + + def test_chat_path_stamps_the_breakpoint_on_a_foundry_deployment_while_the_ambient_base_is_azure_openai( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + self._pin_ambient_api_base(monkeypatch, True) + messages, params = self._chat(None, self.FOUNDRY_API_BASE) + assert (messages, params) == self._chat("azure_ai", self.FOUNDRY_API_BASE) + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 1 + assert ("prompt_cache_options" in params) == supports_openai_prompt_cache_breakpoint(self.MODEL, "azure_ai") + + @pytest.mark.parametrize("pinned", [True, False]) + @pytest.mark.parametrize( + "api_base", + [FOUNDRY_API_BASE, "https://resource.cognitiveservices.azure.com", "https://resource.openai.azure.com"], + ) + def test_chat_path_marks_the_deployment_like_the_provider_its_api_base_dispatches_to( + self, monkeypatch: pytest.MonkeyPatch, api_base: str, pinned: bool + ) -> None: + self._pin_ambient_api_base(monkeypatch, pinned) + dispatched_provider: Final = litellm.get_llm_provider(model=self.MODEL, api_base=api_base)[1] + assert self._chat(None, api_base) == self._chat(dispatched_provider, api_base) + + @pytest.mark.parametrize("model", [MODEL, "bedrock_mantle/openai.gpt-5.6-sol"]) + @pytest.mark.parametrize("set_globally", ["env", "module"]) + def test_chat_path_marks_the_deployment_like_litellm_proxy_when_calls_route_to_one( + self, monkeypatch: pytest.MonkeyPatch, model: str, set_globally: str + ) -> None: + self._pin_ambient_api_base(monkeypatch, False) + if set_globally == "env": + monkeypatch.setenv("USE_LITELLM_PROXY", "True") + else: + monkeypatch.setattr(litellm, "use_litellm_proxy", True) + assert self._chat(None, self.UPSTREAM_PROXY_API_BASE, model) == self._chat( + "litellm_proxy", self.UPSTREAM_PROXY_API_BASE, model + ) + + class TestClientBreakpointsCountedOnce: def test_client_message_breakpoints_are_not_double_counted(self): messages = [ diff --git a/tests/unit/litellm_core_utils/test_fallback_utils.py b/tests/unit/litellm_core_utils/test_fallback_utils.py index 1e3ca0ba90d..c0a344cc6c3 100644 --- a/tests/unit/litellm_core_utils/test_fallback_utils.py +++ b/tests/unit/litellm_core_utils/test_fallback_utils.py @@ -1,6 +1,7 @@ """Tests for litellm.litellm_core_utils.fallback_utils.""" import json +import pathlib from typing import Final import httpx @@ -8,12 +9,14 @@ import pytest import respx import litellm +from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.fallback_utils import ( async_completion_with_fallbacks, ) from litellm.main import acompletion as python_acompletion from litellm.main import completion as python_completion +from litellm.vector_stores.vector_store_registry import LiteLLM_ManagedVectorStore, VectorStoreRegistry @pytest.mark.asyncio @@ -444,3 +447,187 @@ def test_completion_fallbacks_sync( assert response.choices[0].message.content == "sync fallback answer" assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("first", "second") + + +CACHE_POINTS: Final = ({"location": "message", "role": "system"},) +CACHED_MESSAGES: Final = ( + {"role": "system", "content": "You are a careful assistant."}, + {"role": "user", "content": "hello"}, +) +OPENAI_ANSWER: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "gpt-6.1-sol", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +ANTHROPIC_ANSWER: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} + + +def _cached_request(**kwargs: object) -> dict[str, object]: + return { + "messages": [dict(message) for message in CACHED_MESSAGES], + "cache_control_injection_points": [dict(point) for point in CACHE_POINTS], + **kwargs, + } + + +@pytest.mark.asyncio +async def test_acompletion_fallback_sends_the_fallback_model_the_request_a_direct_call_sends( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + openai_route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response( + 401, + json={"error": {"message": "invalid key", "type": "authentication_error", "param": None, "code": None}}, + ) + ) + anthropic_route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=ANTHROPIC_ANSWER) + ) + + await python_acompletion(**_cached_request(model="anthropic/claude-haiku-4-5", api_key="anthropic-key")) + await python_acompletion( + **_cached_request( + model="openai/gpt-6.1-sol", + api_key="openai-key", + fallbacks=[{"model": "anthropic/claude-haiku-4-5", "api_key": "anthropic-key"}], + ) + ) + + assert "prompt_cache_options" in json.loads(openai_route.calls[0].request.content) + direct_body, fallback_body = (json.loads(call.request.content) for call in anthropic_route.calls) + assert fallback_body == direct_body + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_cache_dialect_follows_base_url_over_an_api_base_the_call_does_not_use( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, is_async: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, json=OPENAI_ANSWER) + ) + + async def send(**kwargs: object) -> None: + request: Final = _cached_request( + model="openai/gpt-6.1-sol", api_key="openai-key", base_url="https://api.openai.com/v1", **kwargs + ) + if is_async: + await python_acompletion(**request) + else: + python_completion(**request) + + await send() + await send(api_base="https://unused-gateway.example.com/v1") + + base_url_only, with_unused_api_base = (json.loads(call.request.content) for call in route.calls) + assert with_unused_api_base == base_url_only + + +@pytest.mark.asyncio +async def test_acompletion_fallback_sends_the_context_a_vector_store_tool_retrieves( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + vector_stores=[ + LiteLLM_ManagedVectorStore( + vector_store_id="vs_docs", custom_llm_provider="openai", litellm_params={"api_key": "openai-key"} + ) + ] + ), + ) + search_route: Final = respx_mock.post("https://api.openai.com/v1/vector_stores/vs_docs/search").mock( + return_value=httpx.Response( + 200, + json={ + "object": "vector_store.search_results.page", + "search_query": "what is litellm?", + "data": [{"score": 1.0, "content": [{"type": "text", "text": "litellm is an LLM gateway"}]}], + }, + ) + ) + openai_route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response( + 401, + json={"error": {"message": "invalid key", "type": "authentication_error", "param": None, "code": None}}, + ) + ) + anthropic_route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=ANTHROPIC_ANSWER) + ) + + async def send(**kwargs: object) -> None: + await python_acompletion( + messages=[{"role": "user", "content": "what is litellm?"}], + tools=[{"type": "file_search", "vector_store_ids": ["vs_docs"]}], + **kwargs, + ) + + await send(model="anthropic/claude-haiku-4-5", api_key="anthropic-key") + await send( + model="openai/gpt-6.1-sol", + api_key="openai-key", + fallbacks=[{"model": "anthropic/claude-haiku-4-5", "api_key": "anthropic-key"}], + ) + + direct_body, fallback_body = (json.loads(call.request.content) for call in anthropic_route.calls) + assert "litellm is an LLM gateway" in json.dumps(direct_body) + assert fallback_body == direct_body + assert search_route.call_count <= openai_route.call_count + anthropic_route.call_count + + +@pytest.mark.asyncio +async def test_acompletion_fallback_sends_the_prompt_the_primary_model_prefix_template_renders( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path +) -> None: + _ = (tmp_path / "brief.prompt").write_text( + "---\ninput:\n schema:\n question: string\n---\nSystem: Answer in one word.\n\nUser: {{question}}\n" + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "global_prompt_directory", str(tmp_path)) + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", []) + _ = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response( + 401, + json={"error": {"message": "invalid key", "type": "authentication_error", "param": None, "code": None}}, + ) + ) + anthropic_route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=ANTHROPIC_ANSWER) + ) + + async def send(**kwargs: object) -> None: + await python_acompletion( + prompt_id="brief", + prompt_variables={"question": "What color is the sky?"}, + messages=[{"role": "user", "content": "placeholder"}], + **kwargs, + ) + + await send(model="dotprompt/anthropic/claude-haiku-4-5", api_key="anthropic-key") + await send( + model="dotprompt/openai/gpt-6.1-sol", + api_key="openai-key", + fallbacks=[{"model": "anthropic/claude-haiku-4-5", "api_key": "anthropic-key"}], + ) + + direct_body, fallback_body = (json.loads(call.request.content) for call in anthropic_route.calls) + assert "What color is the sky?" in json.dumps(direct_body) + assert fallback_body == direct_body From b71491c6c910c15bc1580e9469b6d3ecb3b376b1 Mon Sep 17 00:00:00 2001 From: Siddhant Badola Date: Fri, 30 Oct 2026 09:27:59 -0400 Subject: [PATCH 24/64] fix(bedrock): route a Bedrock model id with no Invoke provider to Converse (#43290) * fix(bedrock): route unmapped converse-only family model ids to converse #42713 stopped Bedrock-native OpenAI ids that are not in the cost map yet from falling to the invoke route. The same gap exists for the other families that have no invoke transformation at all: minimax, nvidia, writer, zai, google and xai. An id from one of them that is newer than the cost map hits get_bedrock_route -> "invoke" -> get_bedrock_invoke_provider() -> None -> BedrockError, when Converse would have served it Add a second family regex for those six and default them to converse, the same way the openai family is handled. _OPENAI_FAMILY_MODEL_RE is left alone because the prompt-caching rule is openai-specific Every shipped id of these families already resolves to a converse entry via its base model (checked against the cost map), so this only changes behaviour for ids that are not in the map yet Tests: six unmapped ids route to converse, and an unmapped meta.* id still routes to invoke since meta has an invoke handler * fix(bedrock): route a model id with no Invoke provider to Converse A bedrock/ id missing from the cost map fell to the Invoke route, where an id that names no Invoke provider (a new vendor's model, a bare provisioned-model, imported-model, or custom-model-deployment ARN) fails locally with "Unknown provider=None" before any request is sent. Such ids now go to Converse, which replaces the per-family regex with the general rule. Ids that name an Invoke provider and explicit route prefixes are unchanged. * test(bedrock): expect Converse for ids with no Invoke provider in the nova ARN and not-found tests * fix(bedrock): accept tools on Converse for a model with no cost-map row or Invoke provider, and keep a 404 for an invalid id * fix(bedrock): answer 404 for an invalid model id instead of AWS's 400 AWS answers a mistyped Bedrock model id with a 400 ValidationException. The mapping already raised NotFoundError, but passing AWS's 400 response through made the openai base class report 400, so the proxy answered 400 and the router skipped its cooldown * fix(bedrock): drop lookaround regexes from tool schemas for a model the cost map does not list * fix(bedrock): apply the unknown-model defaults only to bare ids LiteLLM knows nothing about A bare Bedrock id with no cost-map row and no Invoke vendor gets Converse's open params, its reasoning_effort as sent, lookaround-free tool schemas, and a 404 for AWS's invalid-identifier reply, which keeps AWS's response headers. An explicit route like converse/, an ARN, or an id LiteLLM recognizes keeps main's params, regexes and error status, so drop_params still drops tools there * fix(bedrock): judge lookaround support by model id on every route, ignore price-only rows The regex check now strips route and region prefixes from the model and its base_model before asking whether LiteLLM knows the model, so bedrock/converse/ gets the same answer in a direct transform call as on the completion path. A cost-map row counts only when it names a provider, so a price-only register_model row leaves an unlisted id unknown, as on the proxy. AmazonInvokeConfig now inherits the Invoke provider lookup instead of keeping an identical copy. * fix(bedrock): treat a region path like the bare id when the cost map does not list it A deployment like bedrock/us-east-1/ now gets the same tools, reasoning_effort and typo 404 as bedrock/, so Claude Code works on it and a mistyped one still answers 404. * fix(bedrock): keep today's 404 for ids that reach Converse only by default, and keep an empty id local * fix(bedrock): answer 404 for an invalid id only when Converse rejected it, and keep a region path with no id local * fix(bedrock): keep a path-shaped id local, and map /v1/completions errors with the model's provider An id with a slash that is not an ARN, like bedrock/../../converse or an unknown region path, keeps Invoke and its local 404 instead of reaching AWS Converse. atext_completion mapped every error as OpenAI, so a Bedrock typo on /v1/completions answered 400 while chat answered 404. It now resolves the provider from the request's model, as the other async entry points do. * fix(text completion): keep /v1/completions error mapping as it is on main Mapping async /v1/completions errors with each provider's rules is wider than this Bedrock routing change. It moves classes, statuses, cooldown and fallbacks for most providers, and some lose context-window fallbacks. So this restores main's mapping here and leaves that change to its own PR. A Bedrock typo on /v1/completions answers AWS's 400, while chat keeps the 404. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../exception_mapping_utils.py | 32 +- .../bedrock/chat/converse_transformation.py | 8 +- .../base_invoke_transformation.py | 62 +--- litellm/llms/bedrock/common_utils.py | 68 +++- .../test_bedrock_kimi_k3_cache_point_wire.py | 2 +- .../test_exception_mapping_utils.py | 70 +++- .../chat/test_converse_transformation.py | 312 +++++++++++++++++- .../llms/bedrock/test_bedrock_common_utils.py | 107 +++++- .../llms/bedrock/test_nova_imported_models.py | 6 +- 9 files changed, 559 insertions(+), 108 deletions(-) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 5a27f5446ad..33b6d4e1857 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.bug_report import ( ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.secret_redaction import redact_string as redact_secret_string +from litellm.llms.bedrock.common_utils import bedrock_model_reaches_converse_by_default from litellm.llms.bigmodel.common_utils import raise_bigmodel_business_error from litellm.types.utils import LlmProviders @@ -938,13 +939,22 @@ _BEDROCK_MANTLE_CONTEXT_WINDOW_GENERIC_MESSAGE: Final = ( ) -def _bedrock_mantle_context_window_response(original_exception: _ProviderHTTPException) -> httpx.Response: +def _bedrock_response_with_status(original_exception: _ProviderHTTPException, status_code: int) -> httpx.Response: original_response: Final = getattr(original_exception, "response", None) if not isinstance(original_response, httpx.Response) or getattr(original_response, "_request", None) is None: - return httpx.Response(status_code=400, request=httpx.Request(method="POST", url="https://api.openai.com/v1/")) - if original_response.status_code == 400: + return httpx.Response( + status_code=status_code, request=httpx.Request(method="POST", url="https://api.openai.com/v1/") + ) + if original_response.status_code == status_code: return original_response - return httpx.Response(status_code=400, headers=original_response.headers, request=original_response.request) + return httpx.Response(status_code=status_code, headers=original_response.headers, request=original_response.request) + + +def _came_from_bedrock_converse(original_exception: _ProviderHTTPException) -> bool: + original_response: Final = getattr(original_exception, "response", None) + if not isinstance(original_response, httpx.Response) or getattr(original_response, "_request", None) is None: + return False + return original_response.request.url.path.endswith(("/converse", "/converse-stream")) def _get_bedrock_mantle_context_window_message(error_str: str) -> str | None: @@ -980,7 +990,7 @@ def _map_bedrock_exception( message=mantle_context_window_message, model=model, llm_provider=custom_llm_provider, - response=_bedrock_mantle_context_window_response(original_exception), + response=_bedrock_response_with_status(original_exception, 400), ) if ( "too many tokens" in error_str @@ -1003,6 +1013,18 @@ def _map_bedrock_exception( llm_provider="bedrock", response=getattr(original_exception, "response", None), ) + elif ( + "The provided model identifier is invalid" in error_str + and bool(model) + and _came_from_bedrock_converse(original_exception) + and bedrock_model_reaches_converse_by_default(model) + ): + raise NotFoundError( + message=f"BedrockException - {error_str}", + model=model, + llm_provider="bedrock", + response=_bedrock_response_with_status(original_exception, 404), + ) elif "Malformed input request" in error_str: raise BadRequestError( message=f"BedrockException - {error_str}", diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 443e722168d..5e621a5f322 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -63,6 +63,7 @@ from litellm.llms.bedrock.chat.tool_result_images import place_tool_result_image from litellm.llms.bedrock.common_utils import ( bedrock_chat_uses_native_structured_output, bedrock_cleaned_output_schema, + bedrock_model_is_unknown, bedrock_model_supports_regex_lookaround, bedrock_native_output_format, bedrock_native_output_is_opt_in, @@ -529,13 +530,14 @@ class AmazonConverseConfig(BaseConfig): """ Handle the reasoning_effort parameter based on the model type. - - GPT-OSS and DeepSeek V3 models: passed through unchanged via additionalModelRequestFields. + - GPT-OSS, DeepSeek V3 and unknown models (``bedrock_model_is_unknown``): passed through unchanged + via additionalModelRequestFields. - OpenAI GPT-5.x and GPT-6 models: mapped to ``reasoning.effort`` via additionalModelRequestFields. - Nova 2 models: transformed to reasoningConfig. - Anthropic models: mapped to ``thinking`` (and ``output_config.effort`` on adaptive Claude 4.6 / 4.7). """ - if "gpt-oss" in model or "deepseek" in model: + if "gpt-oss" in model or "deepseek" in model or bedrock_model_is_unknown(model): optional_params["reasoning_effort"] = reasoning_effort elif self._is_openai_gpt_reasoning_model(model): reasoning: Final[BedrockConverseGptReasoningEffortBlock] = {"effort": reasoning_effort} @@ -681,7 +683,7 @@ class AmazonConverseConfig(BaseConfig): "parallel_tool_calls", ] - if "arn" in model: # we can't infer the model from the arn, so just add all params + if "arn" in model or bedrock_model_is_unknown(model): supported_params.append("tools") supported_params.append("tool_choice") supported_params.append("thinking") diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 852179e8caa..877fb4ca692 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -2,7 +2,7 @@ import copy import json import time from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, cast, get_args +from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config @@ -74,10 +74,6 @@ _GUARDRAIL_CONFIG_EXPECTED_FORMAT: Final = ( "{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}" ) -_INVOKE_PROVIDERS: Final = TypeAdapter(tuple[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, ...]).validate_python( - get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL) -) - def _response_header(headers: Mapping[str, str], name: str) -> str | None: return headers.get(name, None) @@ -562,62 +558,6 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): """ return False - @staticmethod - def get_bedrock_invoke_provider( - model: str, - ) -> litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None: - """ - Helper function to get the bedrock provider from the model - - handles 4 scenarios: - 1. model=invoke/anthropic.claude-3-5-sonnet-20240620-v1:0 -> Returns `anthropic` - 2. model=anthropic.claude-3-5-sonnet-20240620-v1:0 -> Returns `anthropic` - 3. model=llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n -> Returns `llama` - 4. model=us.amazon.nova-pro-v1:0 -> Returns `nova` - """ - if model.startswith("invoke/"): - model = model.replace("invoke/", "", 1) - - # Special case: Check for "nova" in model name first (before "amazon") - # This handles amazon.nova-* models which would otherwise match "amazon" (Titan) - if "nova" in model.lower(): - if "nova" in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): - return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, "nova") - - _split_model: Final = model.split(".")[0] - if _split_model in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): - return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, _split_model) - - # If not a known provider, check for pattern with two slashes - provider = AmazonInvokeConfig._get_provider_from_model_path(model) - if provider is not None: - return provider - - for provider in _INVOKE_PROVIDERS: - if provider in model: - return provider - return None - - @staticmethod - def _get_provider_from_model_path( - model_path: str, - ) -> litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None: - """ - Helper function to get the provider from a model path with format: provider/model-name - - Args: - model_path (str): The model path (e.g., 'llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n' or 'anthropic/model-name') - - Returns: - Optional[str]: The provider name, or None if no valid provider found - """ - parts: Final = model_path.split("/") - if len(parts) >= 1: - provider: Final = parts[0] - if provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): - return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider) - return None - def convert_messages_to_prompt(self, model: str, messages, provider, custom_prompt_dict) -> tuple[str, list | None]: # handle anthropic prompts and amazon titan prompts prompt = "" diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index d88d7e6280c..d6f5d5b6bd7 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1098,6 +1098,11 @@ def _model_cost_entry(key: str) -> Mapping[str, object] | None: return None if raw is None else _MODEL_COST_ENTRY_ADAPTER.validate_python(raw) +def _cost_map_row_names_a_provider(key: str) -> bool: + entry: Final = _model_cost_entry(key) + return entry is not None and entry.get("litellm_provider") is not None + + def _bedrock_price_map_entries(model: str) -> tuple[Mapping[str, object] | None, ...]: return tuple( _model_cost_entry(key) @@ -1567,7 +1572,8 @@ def bedrock_model_supports_regex_lookaround(model: str, litellm_params: Mapping[ The deployment's ``model_info.supports_regex_lookaround`` wins, then the ``model_prices_and_context_window.json`` entry of its ``base_model``, then the - entry of ``model`` itself. A model nobody flagged keeps its schema as sent. + entry of ``model`` itself. A model nobody flagged keeps its schema as sent, unless + no candidate names a model LiteLLM knows, whatever route prefix it carries. """ params: Final = litellm_params or {} deployment_flag: Final = deployment_model_info(params).get("supports_regex_lookaround") @@ -1576,7 +1582,41 @@ def bedrock_model_supports_regex_lookaround(model: str, litellm_params: Mapping[ base_model: Final = params.get("base_model") candidates: Final = (*((base_model,) if isinstance(base_model, str) else ()), model) flags: Final = (_bedrock_converse_model_flag(candidate, "supports_regex_lookaround") for candidate in candidates) - return next((flag for flag in flags if flag is not None), True) + flag: Final = next((flag for flag in flags if flag is not None), None) + if flag is not None: + return flag + return not all(bedrock_model_is_unknown(split_bedrock_region_path(candidate)[1]) for candidate in candidates) + + +def _bedrock_model_id_without_region(model: str) -> str: + routed_id: Final = model.removeprefix("bedrock/") + region, separator, regional_id = routed_id.partition("/") + return regional_id if separator and region in get_all_bedrock_regions() else routed_id + + +def bedrock_model_is_unknown(model: str) -> bool: + model_id: Final = _bedrock_model_id_without_region(model) + if "/" in model_id: + return False + base: Final = get_bedrock_base_model(model_id) + regional_ids: Final = (f"{region}.{base}" for region in get_bedrock_cross_region_inference_regions()) + if any(_cost_map_row_names_a_provider(key) for key in (*_bedrock_model_info_candidates(model_id), *regional_ids)): + return False + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + return BaseAWSLLM.get_bedrock_invoke_provider(model_id) is None + + +def _names_no_invoke_provider(model: str) -> bool: + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + model_id: Final = _bedrock_model_id_without_region(model) + names_one_model: Final = bool(model_id) and ("/" not in model_id or model_id.startswith("arn:")) + return names_one_model and BaseAWSLLM.get_bedrock_invoke_provider(model) is None + + +def bedrock_model_reaches_converse_by_default(model: str) -> bool: + return BedrockModelInfo.named_bedrock_route(model, None) is None and _names_no_invoke_provider(model) _BedrockConverseModelFlag: TypeAlias = Literal[ @@ -1585,10 +1625,13 @@ _BedrockConverseModelFlag: TypeAlias = Literal[ ] -def _bedrock_converse_model_flag(model: str, key: _BedrockConverseModelFlag) -> bool | None: +def _bedrock_model_info_candidates(model: str) -> tuple[str, ...]: base: Final = get_bedrock_base_model(model) - candidates: Final = dict.fromkeys((model, base, _BEDROCK_MODEL_VERSION_SUFFIX_RE.sub("", base))) - for candidate in candidates: + return tuple(dict.fromkeys((model, base, _BEDROCK_MODEL_VERSION_SUFFIX_RE.sub("", base)))) + + +def _bedrock_converse_model_flag(model: str, key: _BedrockConverseModelFlag) -> bool | None: + for candidate in _bedrock_model_info_candidates(model): with contextlib.suppress(Exception): model_info = get_cached_model_info()( model=candidate, @@ -1766,8 +1809,19 @@ class BedrockModelInfo(BaseLLMModelInfo): default (``bedrock_runtime_chat_completions_is_default``) and ``chat_completions/`` opts any other model in; ``request_params`` (the caller's chat params) sends such a request to Converse when it needs a feature only Converse serves, and ``converse/`` pins a model to Converse. Every other OpenAI-family - model stays on Converse without the prefix. + model stays on Converse without the prefix. A model the cost map does not list as Converse goes to + Invoke only when its id names an Invoke provider, and to Converse otherwise. """ + named_route: Final = BedrockModelInfo.named_bedrock_route(model, request_params) + if named_route is not None: + return named_route + return "converse" if _names_no_invoke_provider(model) else "invoke" + + @staticmethod + def named_bedrock_route( + model: str, + request_params: Mapping[str, object] | None, + ) -> BedrockRoute | None: route_mappings: dict[ str, Literal[ @@ -1822,7 +1876,7 @@ class BedrockModelInfo(BaseLLMModelInfo): return "converse" if _OPENAI_FAMILY_MODEL_RE.search(base_model): return "converse" - return "invoke" + return None @staticmethod def _explicit_converse_route(model: str) -> bool: diff --git a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py index 854da6fce0b..0772e928b35 100644 --- a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py +++ b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py @@ -1037,7 +1037,7 @@ def test_s07_a_5kb_caller_model_under_the_wildcard_is_a_bounded_4xx_and_liveline error: Final = _JSON.validate_python(_JSON.validate_python(json.loads(outcome.raw))["error"]) assert "a" * 5120 in str(error["message"]), outcome.raw[:300] assert_within_factor_of_control(oversized.seconds, control=control.seconds, factor=20) - assert rig.wire.drain() == () + assert len(rig.wire.drain()) == 1 _failure_row(outcome.call_id) diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 88501a23235..efb9648cf13 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -60,7 +60,6 @@ def test_openai_cyber_policy_uses_the_structured_error_code(shape: str, code: st assert caught.value.response.headers["x-request-id"] == "cyber-policy-request" - def _debug_info_for_api_base(configured_api_base: str) -> str: request: Final = httpx.Request("POST", "https://gateway.example/v1/chat/completions") original: Final = OpenAIError( @@ -1602,6 +1601,75 @@ def test_bedrock_mantle_400_maps_to_bad_request(): assert type(excinfo.value) is litellm.BadRequestError +class _ErrorWithoutResponse(Exception): + def __init__(self, message: str, status_code: int) -> None: + super().__init__(message) + self.message = message + self.status_code = status_code + + +def test_bedrock_mantle_context_overflow_without_a_provider_response_answers_400(): + with pytest.raises(litellm.ContextWindowExceededError) as excinfo: + exception_type( + model="openai.gpt-5.6-sol", + original_exception=_ErrorWithoutResponse( + '{"error":{"code":"context_length_exceeded","message":"too long"}}', 400 + ), + custom_llm_provider="bedrock_mantle", + ) + + assert excinfo.value.response.status_code == 400 + + +def _invalid_model_identifier_from(url: str) -> BedrockError: + return BedrockError( + status_code=400, + message='{"message":"The provided model identifier is invalid."}', + response=httpx.Response(status_code=400, request=httpx.Request(method="POST", url=url)), + ) + + +@pytest.mark.parametrize( + "path", + ["/model/us.newvendor.typo-1/converse", "/model/us.newvendor.typo-1/converse-stream"], +) +def test_bedrock_converse_rejecting_an_id_with_no_invoke_provider_answers_404(path): + with pytest.raises(litellm.NotFoundError) as excinfo: + exception_type( + model="us.newvendor.typo-1", + original_exception=_invalid_model_identifier_from(f"https://bedrock-runtime.us-east-1.amazonaws.com{path}"), + custom_llm_provider="bedrock", + ) + + assert excinfo.value.response.status_code == 404 + + +@pytest.mark.parametrize( + "original_exception", + [ + _invalid_model_identifier_from("https://bedrock-runtime.us-east-1.amazonaws.com/model/claude-typo-9/invoke"), + _ErrorWithoutResponse("The provided model identifier is invalid.", 400), + ], + ids=["invoke", "no-response"], +) +def test_bedrock_invalid_model_identifier_not_from_converse_stays_a_bad_request(original_exception): + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type(model="claude-typo-9", original_exception=original_exception, custom_llm_provider="bedrock") + + assert excinfo.value.status_code == 400 + + +def test_bedrock_invalid_model_identifier_without_a_model_stays_a_bad_request(): + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model=None, + original_exception=_ErrorWithoutResponse("The provided model identifier is invalid.", 400), + custom_llm_provider="bedrock", + ) + + assert excinfo.value.status_code == 400 + + def test_bedrock_mantle_context_overflow_maps_to_context_window_exceeded(): from litellm.llms.base_llm.chat.transformation import BaseLLMException diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index d82ba192b1f..743e7e74431 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -1221,6 +1221,97 @@ def test_transform_request_drops_lookaround_regex_when_the_deployment_model_info assert properties == _LOOKAROUND_FREE_PROPERTIES +@pytest.mark.parametrize( + ("model", "litellm_params"), + [ + ("us.newvendor.model-1", None), + ("converse/us.newvendor.model-1", None), + ("us-east-1/us.newvendor.model-1", None), + ("us.newvendor.model-1", {"base_model": "bedrock/converse/us.newvendor.model-1"}), + ], +) +def test_transform_request_drops_lookaround_regex_for_a_model_the_cost_map_does_not_list(model, litellm_params): + assert "us.newvendor.model-1" not in litellm.model_cost + + properties: Final = _tool_schema_properties(model, _ARTIFACT_DATA_ANTHROPIC_TOOL, litellm_params) + + assert properties == _LOOKAROUND_FREE_PROPERTIES + + +@pytest.mark.parametrize( + "model", + [ + "arn:aws:bedrock:us-east-1:123456789012:custom-model-deployment/abc123", + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + "us.anthropic.claude-unreleased-v1:0", + "eu.openai.gpt-unreleased", + "apac.newvendor.listed-1", + ], +) +def test_transform_request_keeps_lookaround_regex_for_an_unlisted_model_litellm_still_recognizes(model, monkeypatch): + monkeypatch.setitem(litellm.model_cost, "us.newvendor.listed-1", {"litellm_provider": "bedrock_converse"}) + + properties = _tool_schema_properties(model, _ARTIFACT_DATA_ANTHROPIC_TOOL) + + assert properties == _ARTIFACT_DATA_INPUT_SCHEMA["properties"] + + +def test_transform_request_keeps_lookaround_regex_for_an_unlisted_model_whose_deployment_opts_out(): + properties = _tool_schema_properties( + "us.newvendor.model-1", _ARTIFACT_DATA_ANTHROPIC_TOOL, {"model_info": {"supports_regex_lookaround": True}} + ) + + assert properties == _ARTIFACT_DATA_INPUT_SCHEMA["properties"] + + +@pytest.mark.parametrize( + "deployment_params", + [ + {"model": "bedrock/us.newvendor.model-1"}, + {"model": "bedrock/converse/us.newvendor.model-1", "allowed_openai_params": ["tools"]}, + ], +) +def test_a_router_deployment_of_a_model_the_cost_map_does_not_list_drops_lookaround_regex( + respx_mock, deployment_params +): + route: Final = respx_mock.post( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.newvendor.model-1/converse" + ).mock( + return_value=httpx.Response( + 200, + json={ + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + }, + ) + ) + router: Final = Router( + model_list=[ + { + "model_name": "unlisted", + "litellm_params": { + **deployment_params, + "aws_access_key_id": "fake-access-key", + "aws_secret_access_key": "fake-secret-key", + "aws_region_name": "us-east-1", + }, + "model_info": {"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6}, + } + ] + ) + + router.completion( + model="unlisted", + messages=[{"role": "user", "content": "hi"}], + tools=[_ARTIFACT_DATA_OPENAI_TOOL], + client=HTTPHandler(), + ) + + tool: Final = json.loads(route.calls[0].request.content)["toolConfig"]["tools"][0] + assert tool["toolSpec"]["inputSchema"]["json"]["properties"] == _LOOKAROUND_FREE_PROPERTIES + + def test_transform_request_keeps_lookaround_regex_when_the_deployment_model_info_opts_out(): properties = _tool_schema_properties( "global.moonshotai.kimi-k3", _ARTIFACT_DATA_ANTHROPIC_TOOL, {"model_info": {"supports_regex_lookaround": True}} @@ -9189,18 +9280,223 @@ def test_bedrock_tools_transformation_valid_params(): assert "test" in result[0]["toolSpec"]["inputSchema"]["json"]["required"] -def test_not_found_error(): +def _invalid_model_identifier_reply(sent: list[httpx.Request]): + def respond(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response( + 400, + json={"message": "The provided model identifier is invalid."}, + headers={"x-amzn-ErrorType": "ValidationException", "x-amzn-RequestId": "req-123"}, + ) + + return respond + + +@pytest.mark.parametrize( + "model, path", + [ + ("bedrock/bad_model", b"/model/bad_model/converse"), + ("bedrock/us-east-1/bad_model", b"/model/bad_model/converse"), + ("bedrock/newvendor.listed-v1:0", b"/model/newvendor.listed-v1%3A0/converse"), + ("bedrock/us-east-1/newvendor.listed-v1:0", b"/model/newvendor.listed-v1%3A0/converse"), + ], +) +def test_not_found_error(model, path, monkeypatch): + monkeypatch.setitem(litellm.model_cost, "newvendor.listed-v1:0", {"litellm_provider": "bedrock", "mode": "chat"}) + sent: Final[list[httpx.Request]] = [] + + with pytest.raises(litellm.NotFoundError) as exc_info: + completion( + model=model, + messages=[{"role": "user", "content": "What is the meaning of life"}], + client=HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_invalid_model_identifier_reply(sent))) + ), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert exc_info.value.status_code == 404 + assert exc_info.value.response.status_code == 404 + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-123" + assert [request.url.raw_path for request in sent] == [path] + + +@pytest.mark.parametrize( + "model", ["bedrock/", "bedrock/us-east-1/", "bedrock/../../converse", "bedrock/xx-fake-1/us.xai.grok-4.7"] +) +def test_a_model_id_that_is_empty_or_a_path_fails_locally_without_calling_aws(model): + sent: Final[list[httpx.Request]] = [] + with pytest.raises(litellm.NotFoundError): completion( - model="bedrock/bad_model", - messages=[ - { - "role": "user", - "content": "What is the meaning of life", - } - ], + model=model, + messages=[{"role": "user", "content": "What is the meaning of life"}], + client=HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_invalid_model_identifier_reply(sent))) + ), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", ) + assert sent == [] + + +@pytest.mark.parametrize( + "model, path", + [ + ("bedrock/converse/bad_model", b"/model/bad_model/converse"), + ("bedrock/anthropic.claude-typo-v1:0", b"/model/anthropic.claude-typo-v1:0/invoke"), + ], +) +def test_invalid_model_identifier_stays_a_bad_request_for_a_model_litellm_recognizes(model, path): + sent: Final[list[httpx.Request]] = [] + + with pytest.raises(litellm.BadRequestError) as exc_info: + completion( + model=model, + messages=[{"role": "user", "content": "What is the meaning of life"}], + client=HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_invalid_model_identifier_reply(sent))) + ), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert exc_info.value.status_code == 400 + assert [request.url.raw_path for request in sent] == [path] + + +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +def _converse_tool_call_reply(sent: list[httpx.Request]): + def respond(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response( + 200, + json={ + "output": { + "message": { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "t1", "name": "get_weather", "input": {"city": "Paris"}}} + ], + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + }, + ) + + return respond + + +@pytest.mark.parametrize("model", ["bedrock/us.newvendor.model-1", "bedrock/us-east-1/us.newvendor.model-1"]) +def test_model_missing_from_cost_map_sends_tools_to_converse(model): + assert "us.newvendor.model-1" not in litellm.model_cost + sent: Final[list[httpx.Request]] = [] + + response: Final = completion( + model=model, + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_WEATHER_TOOL], + tool_choice="auto", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_converse_tool_call_reply(sent)))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert response.choices[0].message.tool_calls[0].function.name == "get_weather" + assert [request.url.raw_path for request in sent] == [b"/model/us.newvendor.model-1/converse"] + tool_config: Final = json.loads(sent[0].content)["toolConfig"] + assert [tool["toolSpec"]["name"] for tool in tool_config["tools"]] == ["get_weather"] + assert tool_config["toolChoice"] == {"auto": {}} + + +def test_a_price_only_row_for_an_unlisted_model_still_sends_tools_to_converse(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, "us.newvendor.priced-1", {"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6} + ) + sent: Final[list[httpx.Request]] = [] + + completion( + model="bedrock/us.newvendor.priced-1", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_WEATHER_TOOL], + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_converse_tool_call_reply(sent)))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + tool_config: Final = json.loads(sent[0].content)["toolConfig"] + assert [tool["toolSpec"]["name"] for tool in tool_config["tools"]] == ["get_weather"] + + +def test_explicit_converse_route_to_an_unlisted_model_still_drops_tools_under_drop_params(): + sent: Final[list[httpx.Request]] = [] + + completion( + model="bedrock/converse/us.newvendor.model-1", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_WEATHER_TOOL], + drop_params=True, + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_converse_tool_call_reply(sent)))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert "toolConfig" not in json.loads(sent[0].content) + + +def test_unknown_model_gets_reasoning_effort_as_sent(): + sent: Final[list[httpx.Request]] = [] + + completion( + model="bedrock/us.newvendor.model-1", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + reasoning_effort="high", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_converse_tool_call_reply(sent)))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert json.loads(sent[0].content)["additionalModelRequestFields"] == {"reasoning_effort": "high"} + + +def test_mapped_model_without_function_calling_still_rejects_tools(): + model: Final = "bedrock/converse/meta.llama3-8b-instruct-v1:0" + assert "meta.llama3-8b-instruct-v1:0" in litellm.model_cost + assert not litellm.supports_function_calling(model=model) + sent: Final[list[httpx.Request]] = [] + + with pytest.raises(litellm.UnsupportedParamsError): + completion( + model=model, + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_WEATHER_TOOL], + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_converse_tool_call_reply(sent)))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert sent == [] + @pytest.mark.parametrize( "model, expected_base_model", diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index dd89d9c46b6..20c0e10cb81 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -214,24 +214,17 @@ def test_system_defined_inference_profile_arn_still_routes_to_converse(): assert route == "converse" -def test_other_opaque_arn_types_still_route_to_invoke(): - """ - Only application-inference-profile ARNs are auto-routed to converse. Other - opaque ARNs (provisioned-model, imported-model, custom-model-deployment) - also yield no invoke provider, but they are frequently invoke-only with - provider-specific body formats, so routing them to converse could break - them. Guards the deliberate scope against an over-broad "any opaque ARN -> - converse" generalization. - """ - for arn_segment in ( - "provisioned-model/abcdefgh1234", - "imported-model/abcdefgh1234", - "custom-model-deployment/abcdefgh1234", - ): - route = BedrockModelInfo.get_bedrock_route( - model=f"bedrock/arn:aws:bedrock:us-east-1:123412341234:{arn_segment}" - ) - assert route == "invoke", f"{arn_segment} should stay on invoke route" +@pytest.mark.parametrize( + "arn_segment", + ["provisioned-model/abcdefgh1234", "imported-model/abcdefgh1234", "custom-model-deployment/abcdefgh1234"], +) +def test_bare_opaque_arn_with_no_invoke_provider_routes_to_converse(arn_segment: str): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + model: Final = f"bedrock/arn:aws:bedrock:us-east-1:123412341234:{arn_segment}" + assert BaseAWSLLM.get_bedrock_invoke_provider(model) is None + assert BedrockModelInfo.get_bedrock_route(model) == "converse" + assert BedrockModelInfo.get_bedrock_route(f"bedrock/llama/arn:aws:bedrock:us-east-1:123412341234:{arn_segment}") == "invoke" def test_govcloud_cross_region_inference_prefix(): @@ -1098,6 +1091,84 @@ def test_unmapped_openai_family_model_routes_to_converse(): assert BedrockModelInfo.get_bedrock_route(imported) == "openai" +@pytest.mark.parametrize( + "unmapped", + [ + "bedrock/minimax.minimax-m99-unmapped", + "bedrock/us.nvidia.nemotron-99b-unmapped", + "bedrock/writer.palmyra-x99-unmapped", + "bedrock/zai.glm-99-unmapped", + "bedrock/google.gemma-99-unmapped", + "bedrock/global.xai.grok-99-unmapped", + ], +) +def test_unmapped_converse_only_family_model_routes_to_converse(unmapped: str): + import litellm + + assert BedrockModelInfo.get_base_model(unmapped) not in litellm.bedrock_converse_models + assert BedrockModelInfo.get_bedrock_route(unmapped) == "converse" + + +def test_unmapped_model_of_family_with_invoke_handler_still_routes_to_invoke(): + from typing import Final + + import litellm + + unmapped: Final = "bedrock/meta.llama99-unmapped-v1:0" + assert BedrockModelInfo.get_base_model(unmapped) not in litellm.bedrock_converse_models + assert BedrockModelInfo.get_bedrock_route(unmapped) == "invoke" + + +@pytest.mark.parametrize( + ("model", "model_path", "response_json"), + [ + ( + "bedrock/us.newvendor.model-1", + b"/model/us.newvendor.model-1/converse", + { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + }, + ), + ( + "bedrock/mistral.mistral-99-unmapped-v1:0", + b"/model/mistral.mistral-99-unmapped-v1:0/invoke", + {"outputs": [{"text": "ok", "stop_reason": "stop"}]}, + ), + ], +) +def test_completion_on_unmapped_model_posts_converse_unless_its_id_names_an_invoke_provider( + model: str, model_path: bytes, response_json: dict[str, object] +): + import httpx + + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + assert BedrockModelInfo.get_base_model(model) not in litellm.bedrock_converse_models + sent: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response( + 200, + json=response_json, + headers={"x-amzn-bedrock-input-token-count": "1", "x-amzn-bedrock-output-token-count": "1"}, + ) + + response: Final = litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert response.choices[0].message.content == "ok" + assert [request.url.raw_path for request in sent] == [model_path] + + @pytest.mark.parametrize( ("model", "expected"), [ diff --git a/tests/unit/llms/bedrock/test_nova_imported_models.py b/tests/unit/llms/bedrock/test_nova_imported_models.py index e3677aaf9e6..231f3097dcd 100644 --- a/tests/unit/llms/bedrock/test_nova_imported_models.py +++ b/tests/unit/llms/bedrock/test_nova_imported_models.py @@ -24,10 +24,8 @@ class TestGetBedrockRoute: def test_nova2_prefix_routes_to_converse(self): assert BedrockModelInfo.get_bedrock_route(NOVA2_MODEL) == "converse" - def test_plain_arn_routes_to_invoke(self): - # Without spec prefix, ARN doesn't match converse models - result = BedrockModelInfo.get_bedrock_route(f"bedrock/{NOVA_ARN}") - assert result == "invoke" + def test_plain_arn_without_an_invoke_provider_routes_to_converse(self): + assert BedrockModelInfo.get_bedrock_route(f"bedrock/{NOVA_ARN}") == "converse" class TestGetBedrockBaseModel: From cbfb1abdb3fe86a139431360ea7d561579dbfb7f Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 06:58:05 -0700 Subject: [PATCH 25/64] fix(proxy)!: label shared-handler errors the way the route catch-alls do (#39555) * fix(proxy): name shared-handler errors by the status they answer handle_exception_on_proxy pinned type=internal_server_error on every exception it wrapped, so a management route rejecting a caller's own argument answered {"type":"internal_server_error","code":"400"} and an SDK branching on the type read its own bad request as a gateway outage. Derive the type from the status the response actually carries, keeping a type the exception already names. * test(proxy): pin a 500 for exceptions whose status is None * fix(proxy): keep internal_server_error for timeouts and 5xx, redact hosts in shared handler errors Provider 5xx and timeouts keep internal_server_error so retry loops still retry them, the shared handler redacts internal hosts from raw exception messages, and /key/list routes its HTTPException through the shared handler instead of its own copy of the old label. * fix(proxy): label shared-handler errors like the chat route, pick the status before the type, redact only gateway failures * fix(proxy): pass shared-handler error messages through as on main Redacting timeout and 5xx messages changed what every shared-handler route returns, which relabeling the error type does not need, so messages go back to the verbatim text main sends. * fix(proxy): redact only the status-less failure main never sent * fix(proxy): answer a generic 500 for a status-less failure and expect auth_error on the /user/new 403 A failure with no status now answers "Internal server error" instead of its own text, since that text can name a guardrail's private host and pattern redaction misses cluster DNS names and IPv6 addresses. The /user/new integration test now expects the auth_error type main's route catch-alls give a 403 the proxy raises * fix(proxy): answer only a 4xx or 5xx carried status and fall back to the route default otherwise A numeric-string status outside 400 to 599 made /v1/messages, images, rerank and the OpenAI pass-through answer HTTP 200 with an error body or drop the connection, and non-ASCII digits or a 5000-digit string slipped through isdecimal. error_status_code now accepts an int or a three-digit ASCII string in 400 to 599, the range the chat route already enforces * refactor(proxy): parse a three-digit string status without recursing --- .../common_utils/openai_error_payload.py | 48 +++-- .../proxy/guardrails/guardrail_endpoints.py | 3 +- litellm/proxy/utils.py | 35 ++-- .../test_guardrail_submissions.py | 4 +- .../test_user_self_update_guard.py | 2 +- .../test_user_delete_revokes_keys.py | 4 +- .../test_apply_guardrail_endpoint.py | 2 +- .../common_utils/test_openai_error_payload.py | 31 ++- .../credential_endpoints/test_endpoints.py | 17 ++ .../guardrails/test_guardrail_endpoints.py | 35 +++- .../test_key_management_endpoints.py | 2 +- .../proxy/utils/helpers/test_error_helpers.py | 181 +++++++++++++++++- 12 files changed, 311 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index 53b2a05d40f..a948c390375 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -28,6 +28,8 @@ _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( _AUTH_FAILURE_STATUS_CODES: Final = frozenset({status.HTTP_401_UNAUTHORIZED, status.HTTP_403_FORBIDDEN}) +_HTTP_ERROR_STATUSES: Final = range(status.HTTP_400_BAD_REQUEST, 600) + _ERROR_BODY: Final = TypeAdapter(dict[str, object]) _FAILOVER_SQLSTATES: Final = frozenset({"25006", "57P01", "57P02", "57P03"}) @@ -47,13 +49,26 @@ def attribute_of(value: object, name: str, default: object = None) -> object: def error_status_code(exc: object, default: int) -> int: - """The HTTP status an exception carries as ``status_code`` or, the way ``ProxyException`` - stores it, as a stringified ``code``; ``default`` when it carries neither.""" - carried: Final = attribute_of(exc, "status_code") - if isinstance(carried, int) and not isinstance(carried, bool): + """The 4xx or 5xx HTTP status an exception carries as ``status_code`` (an int or a decimal string) or, + the way ``ProxyException`` stores it, as a stringified ``code``; ``default`` when it carries neither.""" + carried: Final = _http_error_status(attribute_of(exc, "status_code")) + if carried is not None: return carried stringified: Final = attribute_of(exc, "code") - return int(stringified) if isinstance(stringified, str) and stringified.isdecimal() else default + from_code: Final = _http_error_status(stringified) if isinstance(stringified, str) else None + return default if from_code is None else from_code + + +def _http_error_status(value: object) -> int | None: + match value: + case bool(): + return None + case int() if value in _HTTP_ERROR_STATUSES: + return value + case str() if len(value) == 3 and value.isascii() and value.isdigit() and int(value) in _HTTP_ERROR_STATUSES: + return int(value) + case _: + return None def openai_error_type(exc: object, status_code: int) -> str: @@ -105,9 +120,9 @@ def proxy_exception_for(exc: Exception, default_status_code: int) -> ProxyExcept code=status.HTTP_503_SERVICE_UNAVAILABLE, ) case HTTPException(status_code=status_code, detail=detail): - return _openai_shaped_proxy_exception(exc, str(detail), status_code) + return openai_shaped_proxy_exception(exc, str(detail), status_code) case _: - return _openai_shaped_proxy_exception( + return openai_shaped_proxy_exception( exc, str(exc) or type(exc).__name__, error_status_code(exc, default_status_code) ) @@ -176,14 +191,19 @@ def _route_failure_error_type(exc: Exception, raised_status: int, answered_statu return openai_error_type(exc, answered_status) -def _openai_shaped_proxy_exception(exc: Exception, message: str, status_code: int) -> ProxyException: +def openai_shaped_proxy_exception(exc: Exception, message: str, status_code: int) -> ProxyException: answered_status: Final = answered_status_code(message, status_code) - return ProxyException( - message=message, - type=_catch_all_error_type(exc, answered_status), - param=openai_error_param(exc), - code=answered_status, - ) + return _proxy_exception(exc, message, answered_status, _catch_all_error_type(exc, answered_status)) + + +def status_typed_proxy_exception(exc: HTTPException) -> ProxyException: + message: Final = str(exc.detail) + answered_status: Final = answered_status_code(message, exc.status_code) + return _proxy_exception(exc, message, answered_status, openai_error_type(exc, answered_status)) + + +def _proxy_exception(exc: Exception, message: str, answered_status: int, error_type: str) -> ProxyException: + return ProxyException(message=message, type=error_type, param=openai_error_param(exc), code=answered_status) def _catch_all_error_type(exc: Exception, status_code: int) -> str: diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index a5822bfb0e3..363b9162e00 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import CALLBACK_VAR_ENCRYPTED_PREFIX +from litellm.proxy.common_utils.openai_error_payload import status_typed_proxy_exception from litellm.proxy.common_utils.path_utils import is_within, safe_join from litellm.proxy.guardrails.content_filter_data import CATEGORIES_DIR, DATA_ROOTS, category_dirs, find_category_file from litellm.proxy.guardrails.guardrail_hooks.custom_code.bounded_execution import ( @@ -2533,7 +2534,7 @@ async def apply_guardrail( e = transformed_exception except Exception: verbose_proxy_logger.exception("apply_guardrail: post_call_failure_hook failed") - raise handle_exception_on_proxy(e) + raise handle_exception_on_proxy(status_typed_proxy_exception(e) if isinstance(e, HTTPException) else e) # Success logging outside except so a hook error never triggers failure handlers. response = await _emit_guardrail_success_logs( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b1055bb9124..74e16628bc8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -85,8 +85,8 @@ from litellm.proxy._types import ( from litellm.proxy.bug_report_config import build_proxy_bug_report from litellm.proxy.common_utils.openai_error_payload import ( attribute_of, - litellm_call_id_headers, - openai_error_param, + error_status_code, + openai_shaped_proxy_exception, with_litellm_call_id, ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error @@ -9186,39 +9186,30 @@ def _recreate_writer_on_read_only_transaction(prisma_client: "PrismaClient | Non asyncio.create_task(prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction")) +def _client_message(e: Exception) -> str: + status_is_none: Final = hasattr(e, "status_code") and attribute_of(e, "status_code") is None + return "Internal server error" if status_is_none else strip_bug_report_notice(str(e)) + + def handle_exception_on_proxy(e: Exception, litellm_call_id: str | None = None) -> ProxyException: """ Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible """ - from fastapi import status - verbose_proxy_logger.exception("Exception: %s", e) if PrismaDBExceptionHandler.is_read_only_transaction_error(e): from litellm.proxy.proxy_server import prisma_client _recreate_writer_on_read_only_transaction(prisma_client) - headers: Final = litellm_call_id_headers(litellm_call_id) - if isinstance(e, HTTPException): - return ProxyException( - message=getattr(e, "detail", f"error({e})"), - type=ProxyErrorTypes.internal_server_error, - param=openai_error_param(e), - headers=headers, - code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): + if isinstance(e, ProxyException): return with_litellm_call_id(e, litellm_call_id) - _status_code: Final = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) + if isinstance(e, HTTPException): + return with_litellm_call_id(openai_shaped_proxy_exception(e, str(e.detail), e.status_code), litellm_call_id) if should_report_bug(e): verbose_proxy_logger.error(bug_report_notice(build_proxy_bug_report(e))) - return ProxyException( - message=strip_bug_report_notice(str(e)), - type=ProxyErrorTypes.internal_server_error, - param=openai_error_param(e), - headers=headers, - code=_status_code, - ) + message: Final = _client_message(e) + status_code: Final = error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR) + return with_litellm_call_id(openai_shaped_proxy_exception(e, message, status_code), litellm_call_id) def require_enterprise_license(feature: str | None = None) -> None: diff --git a/tests/integration/authorization/test_guardrail_submissions.py b/tests/integration/authorization/test_guardrail_submissions.py index 16b9cbf31dc..fb095010eb2 100644 --- a/tests/integration/authorization/test_guardrail_submissions.py +++ b/tests/integration/authorization/test_guardrail_submissions.py @@ -432,7 +432,7 @@ def test_team_guardrail_submissions_require_admin_approval( f"Guardrail '{first_name}' not found. " "Please ensure the guardrail is configured in your LiteLLM proxy." ), - "type": "internal_server_error", + "type": "invalid_request_error", "param": None, "code": "404", } @@ -481,7 +481,7 @@ def test_team_guardrail_submissions_require_admin_approval( f"Guardrail '{second_name}' not found. " "Please ensure the guardrail is configured in your LiteLLM proxy." ), - "type": "internal_server_error", + "type": "invalid_request_error", "param": None, "code": "404", } diff --git a/tests/integration/authorization/test_user_self_update_guard.py b/tests/integration/authorization/test_user_self_update_guard.py index 2c495beb77e..4294e8a088d 100644 --- a/tests/integration/authorization/test_user_self_update_guard.py +++ b/tests/integration/authorization/test_user_self_update_guard.py @@ -279,7 +279,7 @@ def test_new_user_with_admin_role_requires_proxy_admin(gateway: Gateway, role: s assert refused.json() == { "error": { "message": f"Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). Attempted to create user with role: {role}. Your role: internal_user", - "type": "internal_server_error", + "type": "auth_error", "param": None, "code": "403", } diff --git a/tests/integration/management/test_user_delete_revokes_keys.py b/tests/integration/management/test_user_delete_revokes_keys.py index 513a80b7319..2a1bd210b49 100644 --- a/tests/integration/management/test_user_delete_revokes_keys.py +++ b/tests/integration/management/test_user_delete_revokes_keys.py @@ -170,7 +170,7 @@ def test_user_delete_revokes_keys_and_memberships(gateway: Gateway, peer: Gatewa assert info.json() == { "error": { "message": f"User {user_id} not found", - "type": "internal_server_error", + "type": "invalid_request_error", "param": None, "code": "404", } @@ -337,7 +337,7 @@ def test_scim_user_delete_revokes_keys(gateway: Gateway, peer: Gateway, monkeypa expected_info: Final = { "error": { "message": f"User {user_id} not found", - "type": "internal_server_error", + "type": "invalid_request_error", "param": None, "code": "404", } diff --git a/tests/integration/observability/test_apply_guardrail_endpoint.py b/tests/integration/observability/test_apply_guardrail_endpoint.py index 9fa7d1a74db..f50679bd54e 100644 --- a/tests/integration/observability/test_apply_guardrail_endpoint.py +++ b/tests/integration/observability/test_apply_guardrail_endpoint.py @@ -156,7 +156,7 @@ def test_apply_guardrail_returns_a_block_error_for_generic_guardrails(gateway: G assert error.model_dump() == { "error": { "message": _BLOCK_REASON, - "type": "internal_server_error", + "type": "invalid_request_error", "param": None, "code": 400, } diff --git a/tests/unit/proxy/common_utils/test_openai_error_payload.py b/tests/unit/proxy/common_utils/test_openai_error_payload.py index acefc4c1e2c..d4ab811d505 100644 --- a/tests/unit/proxy/common_utils/test_openai_error_payload.py +++ b/tests/unit/proxy/common_utils/test_openai_error_payload.py @@ -111,10 +111,13 @@ def test_the_default_status_stands_when_the_exception_carries_none(default: int) assert error_status_code(ValueError("boom"), default) == default -@pytest.mark.parametrize("carried_status", [True, False, "429", None, 429.0]) -def test_a_non_int_carried_status_falls_back_to_the_default(carried_status: object): - """True is an int in Python but not an HTTP status, and a stringified one would break - every caller that compares the code numerically.""" +@pytest.mark.parametrize( + "carried_status", + [True, False, "slow down", "", None, 429.0, 200, 1301, 0, "200", "1301", "0", "0429", " 429", "\u0664\u0662\u0669", "4" * 5000], +) +def test_a_carried_status_that_is_not_an_error_status_falls_back_to_the_default(carried_status: object): + """True is an int in Python but not an HTTP status, and answering a 200 or a 1301 either reads + as success or makes uvicorn drop the connection.""" class _Carrier(Exception): status_code = carried_status @@ -122,6 +125,22 @@ def test_a_non_int_carried_status_falls_back_to_the_default(carried_status: obje assert error_status_code(_Carrier("boom"), 500) == 500 +def test_a_status_carried_as_a_decimal_string_is_answered_as_that_status(): + """Provider stream-error chunks pass the vendor's code through as a string ``status_code``, + and the shared handler used to answer it as is, so dropping it turned a rate limit into a 500.""" + + class _Carrier(Exception): + status_code = "429" + + assert error_status_code(_Carrier("slow down"), 500) == 429 + assert proxy_exception_for(_Carrier("slow down"), default_status_code=500).to_dict() == { + "message": "slow down", + "type": "rate_limit_error", + "param": None, + "code": "429", + } + + def test_a_proxy_exception_keeps_the_status_it_was_raised_with(): """ProxyException stores its status as the string ``code`` rather than ``status_code``, so a route tail that rewraps one used to answer a 4xx rejection as a 500.""" @@ -130,8 +149,8 @@ def test_a_proxy_exception_keeps_the_status_it_was_raised_with(): assert error_status_code(rejection, 500) == 400 -@pytest.mark.parametrize("carried_code", [None, "None", "", "rate_limited", "4xx", 404]) -def test_a_code_that_is_not_a_decimal_string_falls_back_to_the_default(carried_code: object): +@pytest.mark.parametrize("carried_code", [None, "None", "", "rate_limited", "4xx", 404, "200", "1301"]) +def test_a_code_that_is_not_a_stringified_error_status_falls_back_to_the_default(carried_code: object): """Only ProxyException's stringified status is a status; ``code`` on anything else (OpenAI's ``invalid_api_key``, a stray int) says nothing about the HTTP answer.""" diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index 47dd6d561b3..3d8e92988d9 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -3277,6 +3277,23 @@ class TestCredentialDisplayName: assert listed["openai-prod"]["credential_values"]["api_key"] != "sk-old" +def test_looking_up_a_missing_credential_is_labelled_a_client_error_not_a_gateway_failure( + restore_credential_list: None, +): + """Naming a credential that does not exist is the caller's own mistake, and the shared + handler used to label every one of these internal_server_error, so an SDK branching on + the type read a 404 it can never retry past as a gateway outage and backed off forever.""" + response: Final = _call_as("GET", "/credentials/by_name/definitely-not-there") + + assert response.status_code == 404, f"lookup answered {response.status_code}: {response.text}" + assert response.json()["error"] == { + "message": "Credential not found. Got credential name: definitely-not-there", + "type": "invalid_request_error", + "param": None, + "code": "404", + } + + def test_get_credential_by_model_masks_the_deployment_credentials(monkeypatch): import asyncio diff --git a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py index 6807a4fb569..7d912478877 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py @@ -13,7 +13,7 @@ import yaml from fastapi import HTTPException from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_endpoints import ( CreateGuardrailRequest, ListGuardrailsResponse, @@ -2198,6 +2198,39 @@ def _patch_apply_guardrail_env(mocker, guardrail_result): return mock_guardrail +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("vendor_status", "expected_type"), + [(401, "authentication_error"), (403, "permission_error")], +) +async def test_a_guardrail_vendors_rejection_is_labelled_the_way_chat_labels_it( + mocker, vendor_status: int, expected_type: str +): + """A vendor refusing the guardrail's own credentials is not the caller's key being rejected, + so it must not read as the proxy's auth_error the way a 401 the proxy raises does.""" + mock_guardrail: Final = _patch_apply_guardrail_env(mocker, None) + mock_guardrail.apply_guardrail.side_effect = HTTPException( + status_code=vendor_status, detail="vendor refused the request" + ) + mocker.patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", AsyncMock(return_value=None) + ) + + with pytest.raises(ProxyException) as raised: + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello"), + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert raised.value.to_dict() == { + "message": "vendor refused the request", + "type": expected_type, + "param": None, + "code": str(vendor_status), + } + + @pytest.mark.asyncio async def test_apply_guardrail_forwards_metadata_to_guardrail(mocker): """Client-supplied metadata must reach apply_guardrail via request_data so diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index a0c32f2235f..62dd367b8e5 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -18750,7 +18750,7 @@ async def test_list_keys_rejects_invalid_expires(): expires="expred", ) - assert exc_info.value.code == "400" + assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error") assert "Invalid expires value" in str(exc_info.value.message) diff --git a/tests/unit/proxy/utils/helpers/test_error_helpers.py b/tests/unit/proxy/utils/helpers/test_error_helpers.py index df399c8b1d2..b20857dbdaf 100644 --- a/tests/unit/proxy/utils/helpers/test_error_helpers.py +++ b/tests/unit/proxy/utils/helpers/test_error_helpers.py @@ -1,13 +1,20 @@ import asyncio import json +from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx +import openai import pytest from fastapi import HTTPException from prisma.errors import DataError +import litellm from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.common_request_processing import proxy_exception_from_route_error +from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.utils import get_error_message_str, handle_exception_on_proxy +from litellm.types.router import RouterErrors def normalize(value): @@ -96,7 +103,7 @@ def test_handle_exception_on_proxy_happy_path_http_exception(): assert snapshot == { "is_proxy_exception": True, "message": "forbidden", - "type": ProxyErrorTypes.internal_server_error.value, + "type": "auth_error", "code": "403", } @@ -156,7 +163,7 @@ def test_handle_exception_on_proxy_uses_attached_status_code_when_present(): assert snapshot == { "code": "418", "message": "teapot", - "type": ProxyErrorTypes.internal_server_error.value, + "type": "invalid_request_error", } @@ -176,6 +183,176 @@ def test_handle_exception_on_proxy_error_path_none_input_wraps_as_500(): } +@pytest.mark.parametrize( + "status_code, expected_type", + [ + (400, "invalid_request_error"), + (401, "auth_error"), + (403, "auth_error"), + (404, "invalid_request_error"), + (408, "invalid_request_error"), + (409, "invalid_request_error"), + (422, "invalid_request_error"), + (429, "rate_limit_error"), + (500, ProxyErrorTypes.internal_server_error.value), + (503, ProxyErrorTypes.internal_server_error.value), + ], +) +def test_handle_exception_on_proxy_labels_an_http_exception_by_the_status_it_answers( + status_code: int, expected_type: str +): + """A management route that rejects a caller's own argument answers 400, and calling that + internal_server_error tells the caller's retry loop the gateway is down when nothing but + their request is wrong.""" + result: Final = handle_exception_on_proxy(HTTPException(status_code=status_code, detail="boom")) + + assert (result.code, result.type) == (str(status_code), expected_type) + + +def _exception_carrying_status(carried: int | None, message: str = "boom") -> Exception: + class _CarriedStatus(Exception): + status_code = carried + + return _CarriedStatus(message) + + +@pytest.mark.parametrize( + "status_code, expected_type", + [ + (400, "invalid_request_error"), + (404, "invalid_request_error"), + (408, "invalid_request_error"), + (429, "rate_limit_error"), + (500, ProxyErrorTypes.internal_server_error.value), + ], +) +def test_handle_exception_on_proxy_labels_a_carried_status_the_same_way(status_code: int, expected_type: str): + """Routes also reach this handler with plain exceptions carrying their own status, and + those must be named by that status rather than all landing on the 500 label.""" + result: Final = handle_exception_on_proxy(_exception_carrying_status(status_code)) + + assert (result.code, result.type) == (str(status_code), expected_type) + + +def test_handle_exception_on_proxy_keeps_the_type_the_exception_already_names(): + """A rate limit litellm raised already names itself, and overwriting that with the status + map's guess would drop what the client branches on.""" + result: Final = handle_exception_on_proxy( + litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-4o") + ) + + assert (result.code, result.type) == ("429", "throttling_error") + + +@pytest.mark.parametrize( + "exc", + [ + openai.InternalServerError( + message="upstream down", + response=httpx.Response(500, request=httpx.Request("GET", "https://api.openai.com/v1/batches/b")), + body={"type": "server_error"}, + ), + litellm.Timeout(message="took too long", model="gpt-6.1-sol", llm_provider="openai"), + litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-6.1-sol"), + litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-6.1-sol"), + litellm.BadRequestError(message="bad input", model="gpt-6.1-sol", llm_provider="openai"), + HTTPException(status_code=404, detail="no such batch"), + ], + ids=["provider_5xx_naming_its_own_type", "timeout", "rate_limit", "bad_key", "bad_request", "http_404"], +) +def test_handle_exception_on_proxy_labels_a_failure_the_way_the_chat_route_does(exc: Exception): + """Batch, file, and management routes answer through this handler while chat answers through the route + error path, and one client calling both must read the same type for the same failure.""" + shared: Final = handle_exception_on_proxy(exc) + chat: Final = proxy_exception_from_route_error(exc) + + assert (shared.code, shared.type) == (chat.code, chat.type) + + +@pytest.mark.parametrize( + "exc", + [ + HTTPException(status_code=400, detail="bad budget_duration"), + HTTPException(status_code=401, detail="not allowed"), + HTTPException(status_code=403, detail="admin only"), + HTTPException(status_code=404, detail="no such team"), + litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-6.1-sol"), + ValueError(f"{RouterErrors.no_deployments_with_tag_routing.value}=['paid']"), + _exception_carrying_status(409, "duplicate credential"), + ], + ids=["http_400", "http_401", "http_403", "http_404", "provider_401", "tag_denial", "carried_409"], +) +def test_handle_exception_on_proxy_labels_a_failure_the_way_route_catch_alls_do(exc: Exception): + """Management routes answer some failures through their own catch-all and the rest through this handler, + so a client reading /key/list and /key/generate must see the same type and status for the same failure.""" + shared: Final = handle_exception_on_proxy(exc) + catch_all: Final = proxy_exception_for(exc, default_status_code=500) + + assert (shared.code, shared.type, shared.message) == (catch_all.code, catch_all.type, catch_all.message) + + +@pytest.mark.parametrize( + "message, expected", + [ + (f"{RouterErrors.no_deployments_with_provider_budget_routing.value}, try later", ("429", "rate_limit_error")), + (f"{RouterErrors.no_deployments_with_tag_routing.value}=['paid']", ("401", "auth_error")), + ], + ids=["no_deployment_left", "no_deployment_for_tags"], +) +def test_handle_exception_on_proxy_names_a_router_failure_by_the_status_it_answers( + message: str, expected: tuple[str, str] +): + """A router failure that carries no status answers the 429 or 401 its message maps to, and a 429 labelled + internal_server_error reads as an outage to a client that branches on the type.""" + result: Final = handle_exception_on_proxy(ValueError(message)) + + assert (result.code, result.type) == expected + + +def test_handle_exception_on_proxy_keeps_the_path_a_rejected_request_names(): + """A provider rejecting a batch input names the caller's own bucket path, and redacting the /tmp part of it + leaves the caller nothing to fix their request with.""" + result: Final = handle_exception_on_proxy( + litellm.BadRequestError( + message="Invalid input_file_id gs://my-bucket/tmp/batch_input.jsonl", + model="gemini-3.1-pro-preview", + llm_provider="vertex_ai", + ) + ) + + assert result.code == "400" + assert "gs://my-bucket/tmp/batch_input.jsonl" in result.message + + +def test_handle_exception_on_proxy_keeps_the_message_an_exception_with_a_status_carries(): + """Main sends a status-bearing failure's message verbatim, and relabeling its type must not change that text.""" + result: Final = handle_exception_on_proxy(_exception_carrying_status(503, "Cannot connect to host 10.12.4.7:8443")) + + assert (result.code, result.message) == ("503", "Cannot connect to host 10.12.4.7:8443") + + +@pytest.mark.parametrize( + "message", + [ + "boom", + "Cannot connect to host 10.12.4.7:8443", + "Cannot connect to host guardrail.internal.svc.cluster.local:443", + "Cannot connect to host [fd00::12]:8443", + ], + ids=["plain", "ipv4", "cluster_dns", "ipv6"], +) +def test_handle_exception_on_proxy_answers_a_generic_500_for_an_exception_whose_status_is_none(message: str): + """A guardrail whose upstream refuses the connection raises with status_code=None, which crashed the handler + into a 503 blaming the auth database, and its text names private hosts no virtual key was ever sent.""" + result: Final = handle_exception_on_proxy(_exception_carrying_status(None, message)) + + assert (result.code, result.type, result.message) == ( + "500", + ProxyErrorTypes.internal_server_error.value, + "Internal server error", + ) + + @pytest.mark.parametrize( "exc", [ From 160a0b72cff4dbbc41bd5de1130966a207f0c17d Mon Sep 17 00:00:00 2001 From: simonp9 <349053457+simonp9@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:34:10 -0700 Subject: [PATCH 26/64] fix(mcp): bill both calls of a streamed chat completions MCP auto-execute (#48899) * fix(mcp): bill both rounds of a streamed chat completions MCP auto-execute A streamed /v1/chat/completions request with an MCP tool set to require_approval "never" runs two model calls on one logging object, and the final-stream dedupe kept only the first round's spend row. The final usage chunk also showed only the second round. The MCP stream iterator now holds the first round's assembled response, prices it, and logs one success row that carries the usage and cost of both rounds, with the fixed margin charged once. The final usage chunk shows the same summed usage. A single-round stream still logs its one row when the answer finishes. When the follow-up round fails, the first round is logged from a copy of the logging object taken before the follow-up opened, and the failure row carries only the follow-up's partial usage with no second fixed margin. With proxy post-call guardrails armed, only (assembled, cache_hit) tuples are parked for the proxy, and the deferred log is armed only while the iterator reads, so the proxy's orphaned-callback cleanup never sees it after the Router's first-chunk read * fix(mcp): send the usage chunk of a streamed chat MCP request that ends in one round When a streamed chat completions MCP request ends in its first round (the model answers without a tool call, or no tool results come back), the iterator drained the inner stream to fire its spend log and dropped the usage chunk that stream produced, so a client that asked for include_usage got no usage chunk at all. The drained usage chunk is now held and sent to the client after the final content chunk and before [DONE], the same place a plain stream sends it, so the last usage chunk matches the billed row. A two-round request still drops the first round's usage chunk and sends the merged one from the follow-up, and a stream without include_usage still sends none * fix(mcp): hold round one's log for the guardrail, bill the chunk cost once, and keep round one's id When the follow-up call of a streamed chat MCP auto-execute fails, round one's success log now goes through the proxy's deferred guardrail release instead of straight to the handlers, so a post-call guardrail runs on it and the row carries it. Without a proxy release it logs directly as before The last usage chunk now takes its cost from the combined row's billed total, so a fixed cost margin is counted once in the chunk, as it is on the row. The follow-up's usage chunk is held until the follow-up ends to make that possible The combined row now keeps round one's id, the stream's first chunk id, which matches how a single-round stream is logged * fix(mcp): bill every chat MCP stream round once on failure, hang-up, and sync drop When the follow-up call fails, round 1's success row is now logged on a copy of the logging object that does not share the request's guardrail results, so a post-call guardrail that reports on the failed request counts once instead of on both rows A client hang-up during the follow-up now bills round 1 plus the follow-up's partial answer on one success row with round 1's id and the fixed margin counted once, the same way a plain stream's hang-up bills its partial answer The last usage chunk now carries round 1's usage and cost when only round 1 is billed, and a synchronous reader that drops the stream mid-answer bills round 1 on the loop it read with instead of discarding it The deferred release's logged_on argument is typed and every release closure accepts it, and the parked dispatch constants and protocols both MCP billing paths use now live in one module. The unit tests wait on a logging signal instead of a fixed delay and check the billed row through the proxy stream hook instead of private callback state * test(mcp): drive the unassembled follow-up through an upstream audio delta that is not base64 The test replaced litellm.stream_chunk_builder with a stub that raises. Round 2 now streams an audio delta whose data is not valid base64, so the real builder raises, the stream parks only best-effort usage, and the client's last usage chunk must carry round 1's tokens and cost * fix(mcp): count the fixed margin once in the combined stream row's usage cost * fix(mcp): bill round 2's partial answer when a client hangs up on a routed stream On the proxy the Router wraps the MCP stream, and when a client hangs up during round 2 the router's own cleanup closes that stream before the proxy's disconnect cleanup runs. The close logged round 1 alone, so round 2's partial answer was never billed. The MCP stream now bills round 1 plus round 2's partial answer as soon as the hang-up's cancellation reaches its read, before anything closes it * fix(mcp): count the follow-up's own messages in round 2's partial prompt The follow-up stream shares the request's logger, so it took the client's original messages as its prompt. When round 2 ended early, from a hang-up or a broken upstream, its partial usage counted only those messages and left out the assistant tool call and the tool results the follow-up actually sent. The follow-up stream now counts the messages the follow-up call sent * fix(mcp): bill round 2's partial when a routed stream's client hangs up between reads On a proxy the stream the disconnect cleanup holds is the router's fallback wrapper, so the cleanup never found the MCP stream's in-flight billing and released round 1 alone when the client hung up while the proxy sat at a yield. The round billing now registers its in-flight biller on the request's logging object, next to the parked release, and the cleanup looks it up there, so any wrapper chain bills round 1 plus round 2's partial once * fix(mcp): send the MCP stream usage chunk under the billed row id and count the partial off the loop A streamed chat completion that runs a LiteLLM-hosted MCP tool bills both model calls on one row under round 1's chatcmpl id, but the usage chunk it sent the client was round 2's held chunk and kept round 2's id, so a caller could not join the usage chunk to its row. The usage chunk now carries round 1's id in every path that sends it, with or without a deferred post-call guardrail, while its usage and cost and every content chunk's id stay as they were When the client hangs up during round 2, assembling round 2's partial answer counts the follow-up messages, tool result included, with the local tokenizer. That ran on the event loop and stalled every other request on the worker for about a second on a large tool result. The assembly now runs on a worker thread through asyncify from both the reader's cancellation and the proxy's disconnect cleanup, shielded from cancellation, and round 1 is checked again after the await so a release that happened meanwhile is never billed twice * test(mcp): wait on the loop with a hang guard instead of a 300 ms deadline * fix(mcp): give a proxy-rejected follow-up's failure row round 1's own cost breakdown * refactor(streaming): fold the partial response's empty guards into one check --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 22 + .../litellm_core_utils/streaming_handler.py | 35 +- litellm/proxy/common_request_processing.py | 45 +- litellm/proxy/utils.py | 3 + .../responses/mcp/chat_completions_handler.py | 101 +- .../mcp/chat_stream_round_billing.py | 281 +++++ .../responses/mcp/parked_stream_dispatch.py | 72 ++ litellm/responses/mcp/stream_round_billing.py | 58 +- .../integration/_support/hook_count_probe.py | 18 + .../integration/mcp/test_mcp_llm_endpoints.py | 261 +++- .../proxy/test_common_request_processing.py | 28 + .../mcp/test_chat_completions_handler.py | 1050 ++++++++++++++++- .../mcp/test_parked_stream_dispatch.py | 17 + 13 files changed, 1894 insertions(+), 97 deletions(-) create mode 100644 litellm/responses/mcp/chat_stream_round_billing.py create mode 100644 litellm/responses/mcp/parked_stream_dispatch.py create mode 100644 tests/unit/responses/mcp/test_parked_stream_dispatch.py diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4dd33fa5565..8bfb0b60693 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -568,6 +568,11 @@ def _resolve_mantle_region_for_cost( _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) _REDACTED_RESULT: Final = TypeAdapter(object) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +def _without_guardrail_information(metadata: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in metadata.items() if key != _GUARDRAIL_INFORMATION_KEY} def _provider_response_id(source: object) -> str | None: @@ -2566,6 +2571,23 @@ class Logging(LiteLLMLoggingBaseClass): attempt.model_call_details = dict(self.model_call_details) return attempt + def guardrail_scoped_copy(self) -> "Logging": + request_params: Final = _LITELLM_PARAMS.validate_python( + self.litellm_params or {} # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped + ) + litellm_params: Final = { + **request_params, + **{ + key: _without_guardrail_information(_METADATA_MAPPING.validate_python(value)) + for key, value in request_params.items() + if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict) + }, + } + scoped: Final = copy.copy(self) + scoped.litellm_params = litellm_params + scoped.model_call_details = {**self.model_call_details, "litellm_params": litellm_params} + return scoped + def dispatch_attempt_failure_handlers( self, exception: Exception, traceback_exception: str ) -> Coroutine[None, None, None]: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 682b2a2f538..098a404ea14 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2486,6 +2486,25 @@ class CustomStreamWrapper: def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: self._earlier_rounds_spend = earlier + def partial_response(self) -> "ModelResponse | TextCompletionResponse | None": + partial: Final = ( + litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages if isinstance(self.messages, list) else None, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + if self.chunks + else None + ) + usage: Final[object] = getattr(partial, "usage", None) + if partial is None or not isinstance(usage, Usage): + return None + if self.model: + partial.model = self.model + backfill_missing_cache_usage_fields(usage) + return partial + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks @@ -2495,23 +2514,15 @@ class CustomStreamWrapper: moves it into terminal failure usage before opening the next attempt, so a recovered attempt logs its own usage. """ - if self.logging_obj is None or not self.chunks: + if self.logging_obj is None: return try: - partial_response: Final = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages if isinstance(self.messages, list) else None, - logging_obj=self.logging_obj, - count_prompt_tokens=self.count_prompt_tokens, - ) + partial_response: Final = self.partial_response() if partial_response is None: return - usage: Final = cast(Usage | None, getattr(partial_response, "usage", None)) - if usage is None: + usage: Final[object] = getattr(partial_response, "usage", None) + if not isinstance(usage, Usage): return - if self.model: - partial_response.model = self.model - backfill_missing_cache_usage_fields(usage) partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 if self._earlier_rounds_spend is None: self.logging_obj.model_call_details["combined_usage_object"] = usage diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a79d8fe29db..48eb58adaf2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -140,6 +140,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, ) +from litellm.responses.mcp.parked_stream_dispatch import in_flight_round_biller, parked_success from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task from litellm.router import Router from litellm.router_api import RoutingContext @@ -535,6 +536,16 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return await _dispatch_disconnect_billing_event(logging_obj, partial_response) +async def _bill_in_flight_round_on_disconnect(request_data: Mapping[str, object]) -> None: + bill_in_flight_round: Final = in_flight_round_biller(request_data.get("litellm_logging_obj")) + if bill_in_flight_round is None: + return + try: + await bill_in_flight_round() + except Exception as e: # noqa: BLE001 # disconnect billing is best-effort; never break stream teardown + verbose_proxy_logger.debug("Failed to bill the in-flight round after client disconnect: %s", e) + + async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool: try: await logging_obj.dispatch_success_handlers( @@ -2083,6 +2094,13 @@ class _ClientStreamFailureReport: ) +def _metadata_logged_on(logged_on: LiteLLMLoggingObj) -> Mapping[str, object]: + logged_params: Final = object_mapping(logged_on.model_call_details.get("litellm_params")) or MappingProxyType({}) + return MappingProxyType( + {key: logged_params[key] for key in ("metadata", "litellm_metadata") if key in logged_params} + ) + + class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict, served_model: str | None = None): self.data = data @@ -4151,11 +4169,21 @@ class ProxyBaseLLMRequestProcessing: _captured_user_api_key_dict: Final = user_api_key_dict _captured_logging_obj: Final = logging_obj - async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None: + async def _on_deferred_stream_complete( + assembled_response: object, cache_hit: object, logged_on: LiteLLMLoggingObj | None = None + ) -> None: await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( - captured_data=_captured_data, + captured_data=_captured_data + if logged_on is None + else { + **_captured_data, + **{ + key: value for key, value in _metadata_logged_on(logged_on).items() if key in _captured_data + }, + "litellm_logging_obj": logged_on, + }, captured_user_api_key_dict=_captured_user_api_key_dict, - captured_logging_obj=_captured_logging_obj, + captured_logging_obj=_captured_logging_obj if logged_on is None else logged_on, assembled_response=assembled_response, cache_hit=cache_hit, ) @@ -4180,14 +4208,8 @@ class ProxyBaseLLMRequestProcessing: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) - case (assembled_response, cache_hit): - await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers( - assembled_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) + case _ if (success := parked_success(args)) is not None: + await success.dispatch(_captured_native_logging_obj) case _: verbose_proxy_logger.error( "Deferred stream logging dropped: unexpected stored args shape %s", @@ -4524,6 +4546,7 @@ class ProxyBaseLLMRequestProcessing: client_disconnected, ) if recorded_client_disconnect: + await _bill_in_flight_round_on_disconnect(_REQUEST_FIELDS.validate_python(request_data)) deferred_stream_logging_armed: Final = _deferred_stream_logging_is_armed(request_data) ProxyLogging.fire_deferred_stream_logging(request_data) # A disconnect-time success event (the deferred-guardrail flush diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 74e16628bc8..ddaa8d2e69d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,6 +91,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction +from litellm.responses.mcp.parked_stream_dispatch import held_round_spend from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage @@ -4943,6 +4944,8 @@ class ProxyLogging: if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS): return False match _args: + case (ModelResponse(), *_) if (held := held_round_spend(logging_obj)) is not None: + logging_obj.record_billed_rounds_for_failure(held) case (ModelResponse() as assembled, *_): logging_obj.record_assembled_response_for_failure(assembled) case (BilledRoundsDispatch() as billed,): diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index cd072825ef5..1e07b768f62 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -1,15 +1,18 @@ """Helpers for handling MCP-aware `/chat/completions` requests.""" +import asyncio import logging from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +import anyio from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack from litellm._internal_context import internal_sub_call from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.chat_stream_round_billing import ChatStreamRoundBilling, engaged_round_billing from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) @@ -19,7 +22,7 @@ from litellm.responses.mcp.round_billing import ( first_round_billed_on_failure, summed_cost_breakdown, ) -from litellm.types.utils import Message, ModelResponse +from litellm.types.utils import Message, ModelResponse, Usage from litellm.utils import CustomStreamWrapper if TYPE_CHECKING: @@ -258,10 +261,20 @@ async def acompletion_with_mcp( self.follow_up_stream = None self.follow_up_iterator = None self.follow_up_exhausted = False + self._round_billing: ChatStreamRoundBilling | None = engaged_round_billing(initial_stream.logging_obj) + self._unsent_usage_chunk: ModelResponseStream | None = None def __aiter__(self): return self + def __del__(self) -> None: + self.release_abandoned() + + def release_abandoned(self) -> None: + round_billing: Final[object] = getattr(self, "_round_billing", None) + if isinstance(round_billing, ChatStreamRoundBilling): + round_billing.release_abandoned() + def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( @@ -321,18 +334,54 @@ async def acompletion_with_mcp( return chunk - async def _drain_inner_stream(self): + async def _drain_inner_stream(self) -> tuple[ModelResponseStream, ...]: try: - while True: - await self._stream_iterator.__anext__() - except StopAsyncIteration: - pass + return tuple([chunk async for chunk in initial_stream]) except Exception: logging.getLogger("LiteLLM").exception( "Error draining inner MCP stream after final chunk; spend logging may be incomplete" ) + return () - async def __anext__(self): + async def _complete_first_round(self) -> ModelResponseStream | None: + drained: Final = await self._drain_inner_stream() + if self._round_billing is not None: + self._round_billing.complete_first_round() + usage_chunks: Final = tuple( + chunk for chunk in drained if isinstance(getattr(chunk, "usage", None), Usage) + ) + return usage_chunks[-1] if usage_chunks else None + + async def _follow_up_or_release_first_round(self, usage_chunk: ModelResponseStream | None) -> None: + if self.tool_results and self.complete_response: + await self._prepare_follow_up_call() + if self.follow_up_stream is not None: + return + self._unsent_usage_chunk = usage_chunk + if self._round_billing is not None: + self._round_billing.finish() + + async def __anext__(self) -> ModelResponseStream: + round_billing: Final = self._round_billing + if round_billing is None: + return await self._next_chunk() + with round_billing.armed_while_reading(): + try: + return await self._next_chunk() + except StopAsyncIteration: + usage_chunk: Final = round_billing.finish() + if usage_chunk is None: + raise + return usage_chunk + except asyncio.CancelledError: + with anyio.CancelScope(shield=True): + await round_billing.bill_in_flight_round() + raise + except Exception: + round_billing.fail() + raise + + async def _next_chunk(self) -> ModelResponseStream: # Phase 1: Collect and yield initial stream chunks if not self.stream_exhausted: # Get the iterator from the stream wrapper @@ -362,29 +411,22 @@ async def acompletion_with_mcp( if is_final: self.stream_exhausted = True + usage_chunk: Final = await self._complete_first_round() await self._process_tool_calls() chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - # Drain inner stream so CustomStreamWrapper fires its - # end-of-stream handler (dispatch_success_handlers → - # _ProxyDBLogger → LiteLLM_SpendLogs). The CSW may - # yield one usage chunk before raising StopAsyncIteration. - await self._drain_inner_stream() + await self._follow_up_or_release_first_round(usage_chunk) return chunk except StopAsyncIteration: self.stream_exhausted = True + drained_usage_chunk: Final = await self._complete_first_round() # Process tool calls after stream is exhausted await self._process_tool_calls() # If we have chunks, yield the final one with metadata if self.collected_chunks: final_chunk = self.collected_chunks[-1] final_chunk = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) - # If we have tool results, prepare follow-up call - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - await self._drain_inner_stream() + await self._follow_up_or_release_first_round(drained_usage_chunk) return final_chunk # Phase 2: Yield follow-up stream chunks if available @@ -400,7 +442,6 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.debug("Follow-up chunk yielded: %s", chunk) - return chunk except StopAsyncIteration: self.follow_up_exhausted = True from litellm._logging import verbose_logger @@ -408,6 +449,10 @@ async def acompletion_with_mcp( verbose_logger.debug("Follow-up stream exhausted") # After follow-up stream is exhausted, check if we need to raise StopAsyncIteration raise StopAsyncIteration + client_chunk: Final = ( + chunk if self._round_billing is None else self._round_billing.client_chunk(chunk) + ) + return await self._next_chunk() if client_chunk is None else client_chunk # If we're here and follow_up_stream is None but we expected it, log a warning if ( @@ -420,7 +465,11 @@ async def acompletion_with_mcp( verbose_logger.warning("Follow-up stream was not created despite having tool results") - raise StopAsyncIteration + unsent_usage_chunk: Final = self._unsent_usage_chunk + if unsent_usage_chunk is None: + raise StopAsyncIteration + self._unsent_usage_chunk = None + return unsent_usage_chunk async def _process_tool_calls(self): """Process tool calls after streaming completes.""" @@ -488,10 +537,16 @@ async def acompletion_with_mcp( # This ensures the patch works correctly in tests import litellm + earlier_rounds: Final = None if self._round_billing is None else self._round_billing.bank_first_round() follow_up_response: Final = await litellm.acompletion(**follow_up_call_args) # Ensure follow-up response is a CustomStreamWrapper if isinstance(follow_up_response, CustomStreamWrapper): + if earlier_rounds is not None: + follow_up_response.bill_failure_after_earlier_rounds(earlier_rounds) + follow_up_response.messages = follow_up_messages + if self._round_billing is not None: + self._round_billing.follow_up_started(follow_up_response.partial_response) self.follow_up_stream = follow_up_response from litellm._logging import verbose_logger @@ -548,11 +603,13 @@ async def acompletion_with_mcp( def __aiter__(self): return self._custom_iterator + async def aclose(self) -> None: + self._custom_iterator.release_abandoned() + await super().aclose() + def __iter__(self): # For synchronous iteration, create a sync wrapper if self._sync_iterator is None: - import asyncio - try: self._sync_loop = asyncio.get_event_loop() except RuntimeError: diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py new file mode 100644 index 00000000000..370de663606 --- /dev/null +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -0,0 +1,281 @@ +import asyncio +from collections.abc import Callable, Generator +from contextlib import contextmanager +from dataclasses import dataclass, replace +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.litellm_core_utils.asyncify import asyncify, is_event_loop_running +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.parked_stream_dispatch import ( + DISPATCH_RELEASE, + HELD_ROUND_SPEND, + IN_FLIGHT_ROUND, + PARKED_DISPATCH, + SERVED_FROM_CACHE, + DeferredStreamRelease, + parked_success, +) +from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown +from litellm.responses.streaming_iterator import spawn_logging_task +from litellm.types.utils import ModelResponse, ModelResponseStream, Usage + +Parked: TypeAlias = tuple[object, ...] +PartialAnswer: TypeAlias = Callable[[], object] + + +@dataclass(frozen=True, slots=True) +class _AssembledRound: + response: ModelResponse + usage: Usage + cache_hit: bool + parked: Parked + + +@dataclass(frozen=True, slots=True) +class _PricedRound: + response: ModelResponse + usage: Usage + price: UncachedPrice + cache_hit: bool + + @property + def spend(self) -> BilledRoundsSpend: + return BilledRoundsSpend(self.usage, self.price.cost or 0.0, self.price.cost_breakdown) + + +@dataclass(frozen=True, slots=True) +class _HeldUsageChunk: + chunk: ModelResponseStream + usage: Usage + + +@dataclass(frozen=True, slots=True) +class _FirstRound: + priced: _PricedRound + parked: Parked + logged_alone: Parked + banked: LiteLLMLoggingObj | None = None + usage_chunk: _HeldUsageChunk | None = None + + def client_usage_chunk(self, billed: ModelResponse | None) -> ModelResponseStream | None: + if self.usage_chunk is None: + return None + usage: Final = ( + self.priced.usage.model_copy(update=MappingProxyType({"cost": self.priced.price.cost})) + if billed is None + else merged_chat_usage(self.priced.usage, self.usage_chunk.usage).model_copy( + update=MappingProxyType({"cost": _cost(billed)}) + ) + ) + return self.usage_chunk.chunk.model_copy( + update=MappingProxyType({"id": self.priced.response.id, "usage": usage}) + ) + + +def _cost(response: ModelResponse) -> object: + hidden_params: Final = get_hidden_params(response) + return None if hidden_params is None else hidden_params.get("response_cost") + + +def _assembled_round(parked: Parked | None) -> _AssembledRound | None: + match parked: + case (ModelResponse() as response, cache_hit) if isinstance(usage := getattr(response, "usage", None), Usage): + return _AssembledRound(response, usage, cache_hit is True, parked) + case _: + return None + + +def engaged_round_billing(logging_obj: object) -> "ChatStreamRoundBilling | None": + return ChatStreamRoundBilling(logging_obj) if isinstance(logging_obj, LiteLLMLoggingObj) else None + + +class ChatStreamRoundBilling: + def __init__(self, logging_obj: LiteLLMLoggingObj) -> None: + self._logging_obj: Final = logging_obj + self._first: _FirstRound | None = None + self._partial_follow_up: PartialAnswer | None = None + self._reading_loop: asyncio.AbstractEventLoop | None = None + + @contextmanager + def armed_while_reading(self) -> Generator[None]: + armed: Final = getattr(self._logging_obj, DISPATCH_RELEASE, None) is None + if armed: + setattr(self._logging_obj, DISPATCH_RELEASE, self._log_success) + self._reading_loop = asyncio.get_running_loop() + try: + yield + finally: + if armed and self._peek_parked() is None: + self._disarm() + + def complete_first_round(self) -> None: + assembled: Final = _assembled_round(self._peek_parked()) + if assembled is not None: + self._first = _FirstRound(self._priced(assembled), assembled.parked, assembled.parked) + setattr(self._logging_obj, HELD_ROUND_SPEND, self._held_round_one_spend) + + def bank_first_round(self) -> BilledRoundsSpend | None: + self._bank_apart_from_the_follow_up() + return None if self._first is None else self._first.priced.spend + + def client_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: + usage: Final[object] = getattr(chunk, "usage", None) + if self._first is None or not isinstance(usage, Usage): + return chunk + self._first = replace(self._first, usage_chunk=_HeldUsageChunk(chunk, usage)) + return None + + def finish(self) -> ModelResponseStream | None: + first: Final = self._first + billed: Final = self._release_rounds() + return None if first is None else first.client_usage_chunk(billed) + + def fail(self) -> None: + self._bank_apart_from_the_follow_up() + first: Final = self._first + if first is None or first.banked is None or not self._holds_round_one_alone(first): + self.finish() + return + self._first = None + setattr(self._logging_obj, PARKED_DISPATCH, None) + self._log_through_the_release((*first.parked, first.banked.guardrail_scoped_copy())) + + def follow_up_started(self, partial_follow_up: PartialAnswer) -> None: + self._partial_follow_up = partial_follow_up + setattr(self._logging_obj, IN_FLIGHT_ROUND, self.bill_in_flight_round) + + async def bill_in_flight_round(self) -> None: + partial_follow_up: Final = self._partial_follow_up + if partial_follow_up is None or self._first_round_held_alone() is None: + return + final_round: Final = _assembled_round((await asyncify(partial_follow_up)(), False)) + first: Final = self._first_round_held_alone() + if final_round is None or first is None: + return + self._first = None + setattr( + self._logging_obj, PARKED_DISPATCH, self._billed_for_both_rounds(first.priced, self._priced(final_round)) + ) + + def release_abandoned(self) -> None: + if not self._owns_release(): + return + if is_event_loop_running() or self._idle_reading_loop() is not None: + self.finish() + return + setattr(self._logging_obj, PARKED_DISPATCH, None) + self._first = None + self._disarm() + + def _bank_apart_from_the_follow_up(self) -> None: + first: Final = self._first + if first is None or first.banked is not None or self._peek_parked() is not first.parked: + return + banked: Final = self._logging_obj.attempt_scoped_copy() + if not self._owns_release(): + self._first = replace(first, banked=banked) + return + logged_alone: Final = (*first.parked, banked) + self._first = replace(first, logged_alone=logged_alone, banked=banked) + setattr(self._logging_obj, PARKED_DISPATCH, logged_alone) + + def _release_rounds(self) -> ModelResponse | None: + parked: Final = self._take_parked() + first: Final = self._first + self._first = None + if parked is None: + self._disarm() + return None + if first is None or parked is first.logged_alone: + self._release(parked) + return None + final_round: Final = _assembled_round(parked) + if final_round is None: + self._release(first.logged_alone) + return None + billed, cache_hit = self._billed_for_both_rounds(first.priced, self._priced(final_round)) + self._release((billed, cache_hit)) + return billed + + def _billed_for_both_rounds(self, first: _PricedRound, final: _PricedRound) -> tuple[ModelResponse, bool]: + first_breakdown: Final = first.price.cost_breakdown + final_breakdown: Final = final.price.cost_breakdown + rounds: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_breakdown) + usage: Final[object] = getattr(rounds, "usage", None) + charged_usage: Final = ( + usage.model_copy(update=MappingProxyType({"cost": _cost(rounds)})) if isinstance(usage, Usage) else usage + ) + billed: Final = rounds.model_copy(update=MappingProxyType({"id": first.response.id, "usage": charged_usage})) + cache_hit: Final = first.cache_hit and final.cache_hit + set_hidden_params(billed, {**(get_hidden_params(rounds) or {}), "cache_hit": cache_hit}) + self._logging_obj.cost_breakdown = summed_cost_breakdown( + first.response, first_breakdown, final.response, final_breakdown + ) + return billed, cache_hit + + def _priced(self, assembled: _AssembledRound) -> _PricedRound: + served: Final = assembled.cache_hit or served_from_cache(assembled.response) + price: Final = SERVED_FROM_CACHE if served else self._logging_obj.uncached_price(assembled.response) + priced: Final = assembled.response.model_copy() + set_hidden_params( + priced, + {**(get_hidden_params(assembled.response) or {}), "response_cost": price.cost, "cache_hit": served}, + ) + return _PricedRound(priced, assembled.usage, price, served) + + def _peek_parked(self) -> Parked | None: + parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None) + return parked if isinstance(parked, tuple) else None + + def _take_parked(self) -> Parked | None: + parked: Final = self._peek_parked() + setattr(self._logging_obj, PARKED_DISPATCH, None) + return parked + + def _first_round_held_alone(self) -> _FirstRound | None: + first: Final = self._first + return first if first is not None and self._holds_round_one_alone(first) else None + + def _held_round_one_spend(self) -> BilledRoundsSpend | None: + first: Final = self._first_round_held_alone() + return None if first is None else first.priced.spend + + def _holds_round_one_alone(self, first: _FirstRound) -> bool: + parked: Final = self._peek_parked() + return parked is first.parked or parked is first.logged_alone + + def _owns_release(self) -> bool: + release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None) + return release == self._log_success + + def _release(self, parked: Parked) -> None: + if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, PARKED_DISPATCH, parked) + return + self._log_through_the_release(parked) + + def _log_through_the_release(self, parked: Parked) -> None: + armed: Final[DeferredStreamRelease | None] = getattr(self._logging_obj, DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) + logged: Final = (self._log_success if armed is None else armed)(*parked) + idle_reading_loop: Final = None if is_event_loop_running() else self._idle_reading_loop() + if idle_reading_loop is not None: + idle_reading_loop.run_until_complete(logged) + return + spawn_logging_task(asyncio.get_running_loop(), logged, task_name="Chat completions MCP stream success log") + + def _idle_reading_loop(self) -> asyncio.AbstractEventLoop | None: + loop: Final = self._reading_loop + return loop if loop is not None and not loop.is_closed() and not loop.is_running() else None + + def _disarm(self) -> None: + if self._owns_release(): + setattr(self._logging_obj, DISPATCH_RELEASE, None) + + async def _log_success(self, *parked: object) -> None: + success: Final = parked_success(parked) + if success is not None: + await success.dispatch(self._logging_obj) diff --git a/litellm/responses/mcp/parked_stream_dispatch.py b/litellm/responses/mcp/parked_stream_dispatch.py new file mode 100644 index 00000000000..792363f0bab --- /dev/null +++ b/litellm/responses/mcp/parked_stream_dispatch.py @@ -0,0 +1,72 @@ +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Final, Protocol + +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +PARKED_DISPATCH: Final = "_deferred_stream_complete_args" +DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +IN_FLIGHT_ROUND: Final = "_bill_in_flight_round_on_disconnect" +HELD_ROUND_SPEND: Final = "_held_round_spend_for_failure" +SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) + + +class DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +class DeferredStreamRelease(Protocol): + def __call__(self, *parked: object) -> Coroutine[object, object, None]: ... + + +class InFlightRoundBiller(Protocol): + def __call__(self) -> Coroutine[object, object, None]: ... + + +class HeldRoundSpend(Protocol): + def __call__(self) -> BilledRoundsSpend | None: ... + + +def success_dispatcher(logger: LiteLLMLoggingObj) -> DispatchesSuccessHandlers: + return logger + + +@dataclass(frozen=True, slots=True) +class ParkedSuccess: + assembled: object + cache_hit: object + logged_on: LiteLLMLoggingObj | None = None + + async def dispatch(self, request_logger: LiteLLMLoggingObj) -> None: + logger: Final = request_logger if self.logged_on is None else self.logged_on + await success_dispatcher(logger).dispatch_success_handlers( + self.assembled, cache_hit=self.cache_hit, start_time=None, end_time=None, prefer_async_handlers=True + ) + + +def in_flight_round_biller(logging_obj: object) -> InFlightRoundBiller | None: + biller: Final[InFlightRoundBiller | None] = getattr(logging_obj, IN_FLIGHT_ROUND, None) + return biller if isinstance(logging_obj, LiteLLMLoggingObj) and callable(biller) else None + + +def held_round_spend(logging_obj: object) -> BilledRoundsSpend | None: + held: Final[HeldRoundSpend | None] = getattr(logging_obj, HELD_ROUND_SPEND, None) + return held() if isinstance(logging_obj, LiteLLMLoggingObj) and callable(held) else None + + +def parked_success(parked: tuple[object, ...]) -> ParkedSuccess | None: + match parked: + case (assembled, cache_hit): + return ParkedSuccess(assembled, cache_hit) + case (assembled, cache_hit, LiteLLMLoggingObj() as logged_on): + return ParkedSuccess(assembled, cache_hit, logged_on) + case _: + return None diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py index 443147230bf..5cb256a8ac0 100644 --- a/litellm/responses/mcp/stream_round_billing.py +++ b/litellm/responses/mcp/stream_round_billing.py @@ -2,11 +2,18 @@ import asyncio from collections.abc import Callable, Coroutine, Sequence from dataclasses import dataclass from datetime import datetime -from typing import Final, Protocol, TypeAlias +from typing import Final, TypeAlias from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.parked_stream_dispatch import ( + DISPATCH_RELEASE, + PARKED_DISPATCH, + SERVED_FROM_CACHE, + parked_success, + success_dispatcher, +) from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, spawn_logging_task from litellm.responses.utils import ResponseAPILoggingUtils @@ -16,25 +23,6 @@ from litellm.types.utils import CostBreakdown, ModelResponse, Usage Parked: TypeAlias = tuple[object, ...] RestoresHeaders: TypeAlias = Callable[[ResponsesAPIStreamingResponse], None] -_PARKED_DISPATCH: Final = "_deferred_stream_complete_args" -_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" -_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) - - -class _DispatchesSuccessHandlers(Protocol): - async def dispatch_success_handlers( - self, - result: object = None, - start_time: object = None, - end_time: object = None, - cache_hit: object = None, - prefer_async_handlers: bool = False, - ) -> None: ... - - -def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers: - return logger - def _event_loop_running() -> bool: try: @@ -198,8 +186,8 @@ class StreamedRoundBilling: self._billed: _BilledRounds | None = None self._banked: Parked | None = None self._banked_row: _RoundRow | None = None - if getattr(logging_obj, _DISPATCH_RELEASE, None) is None: - setattr(logging_obj, _DISPATCH_RELEASE, self._run_parked) + if getattr(logging_obj, DISPATCH_RELEASE, None) is None: + setattr(logging_obj, DISPATCH_RELEASE, self._run_parked) def complete_round( self, @@ -286,7 +274,7 @@ class StreamedRoundBilling: def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice: if served_from_cache(response): - return _SERVED_FROM_CACHE + return SERVED_FROM_CACHE match round_parked: case (ModelResponse() as assembled, bool()): return self._logging_obj.uncached_price(assembled) @@ -297,7 +285,7 @@ class StreamedRoundBilling: self._banked = parked self._banked_row = row if parked is not None: - setattr(self._logging_obj, _PARKED_DISPATCH, parked) + setattr(self._logging_obj, PARKED_DISPATCH, parked) def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None: row: Final = self._banked_row @@ -320,7 +308,7 @@ class StreamedRoundBilling: @staticmethod async def _log_success(success: _SuccessLog) -> None: success.carry_bridged_usage() - await _success_dispatcher(success.logger).dispatch_success_handlers( + await success_dispatcher(success.logger).dispatch_success_handlers( success.row.logged, start_time=success.logger.start_time, end_time=success.row.end_time, @@ -329,10 +317,10 @@ class StreamedRoundBilling: ) def _take_parked(self) -> Parked | None: - parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None) + parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None) if not isinstance(parked, tuple): return None - setattr(self._logging_obj, _PARKED_DISPATCH, None) + setattr(self._logging_obj, PARKED_DISPATCH, None) return parked async def _run_parked(self, *parked: object) -> None: @@ -341,29 +329,27 @@ class StreamedRoundBilling: await dispatch case (dispatch,) if asyncio.iscoroutine(dispatch): await dispatch - case (assembled, cache_hit): - await _success_dispatcher(self._logging_obj).dispatch_success_handlers( - assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True - ) + case _ if (success := parked_success(parked)) is not None: + await success.dispatch(self._logging_obj) case _: return def _owns_release(self) -> bool: - release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None) return release == self._run_parked def _release(self, parked: Parked) -> None: - if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release(): - setattr(self._logging_obj, _PARKED_DISPATCH, parked) + if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, PARKED_DISPATCH, parked) return - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) spawn_logging_task( asyncio.get_running_loop(), self._run_parked(*parked), task_name="Responses MCP stream success logging" ) def _disarm(self) -> None: if self._owns_release(): - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) @staticmethod def _close(parked: Parked | None) -> None: diff --git a/tests/integration/_support/hook_count_probe.py b/tests/integration/_support/hook_count_probe.py index 0d4056e8fb9..3cd72f1a6ba 100644 --- a/tests/integration/_support/hook_count_probe.py +++ b/tests/integration/_support/hook_count_probe.py @@ -75,6 +75,24 @@ class BlockingPostCallGuardrail(CustomGuardrail): raise HTTPException(status_code=400, detail="blocked by the hook-count guardrail") +class UsageRecordingPostCallGuardrail(CustomGuardrail): + """Records the token usage of the response it is handed in the request's guardrail information, so the spend row + shows which response the guardrail saw.""" + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: object, response: object + ) -> None: + usage: Final = getattr(response, "usage", None) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={ + "prompt_tokens": getattr(usage, "prompt_tokens", None), + "completion_tokens": getattr(usage, "completion_tokens", None), + }, + request_data=data, + guardrail_status="success", + ) + + probe: Final = HookCountProbe() sync_only_probe: Final = SyncOnlyHookCountProbe() diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 4b967802893..f6743a8751a 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -46,9 +46,10 @@ from pydantic import JsonValue, TypeAdapter pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest.mark.shares_worker_state] -@pytest.fixture(scope="session") -def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: - directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway") +@contextmanager +def _private_gateway( + directory: Path, guardrails: Sequence[Mapping[str, JsonValue]] = () +) -> Iterator[tuple[Gateway, Mapping[str, str]]]: stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = directory / "stored-responses.yaml" (directory / f"{STREAM_BLOCK_MODULE}.py").write_text(STREAM_BLOCK_SOURCE) @@ -68,7 +69,8 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat { "guardrail_name": STREAM_BLOCK, "litellm_params": {"guardrail": f"{STREAM_BLOCK_MODULE}.BlockGuard", "mode": "post_call"}, - } + }, + *guardrails, ], } ) @@ -81,15 +83,61 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat yield private, services -@pytest.fixture -def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: - private, services = _llm_gateway +@contextmanager +def _serving(served: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + private, services = served with pytest.MonkeyPatch.context() as environment: for name, value in services.items(): environment.setenv(name, value) yield private +@pytest.fixture(scope="session") +def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-gateway")) as served: + yield served + + +@pytest.fixture +def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + with _serving(_llm_gateway) as private: + yield private + + +ROUND_USAGE_GUARDRAIL: Final = "mcp-llm-round-usage" +CONTENT_FILTER: Final = f"mcp-llm-content-filter-{uuid.uuid4().hex[:8]}" + + +@pytest.fixture(scope="module") +def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + round_usage: Final[Mapping[str, JsonValue]] = { + "guardrail_name": ROUND_USAGE_GUARDRAIL, + "litellm_params": { + "guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail", + "mode": "post_call", + "default_on": True, + }, + } + content_filter: Final[Mapping[str, JsonValue]] = { + "guardrail_name": CONTENT_FILTER, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "default_on": False, + "blocked_words": [{"keyword": "forbidden-keeper", "action": "BLOCK"}], + }, + } + guardrails: Final = (round_usage, content_filter) + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), guardrails) as served: + yield served + + +@pytest.fixture +def guarded_gateway(_guarded_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + with _serving(_guarded_llm_gateway) as private: + yield private + + Surface = Literal["chat", "responses", "messages", "messages_bridge", "responses_bridge"] SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge") ADD: Final = {"a": 2, "b": 3} @@ -1032,6 +1080,184 @@ def test_nonstream_failed_follow_up_bills_round_one_on_the_failure_row( ) +def _streamed_chunks(response: httpx.Response) -> tuple[dict[str, object], ...]: + return tuple( + json.loads(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def test_streaming_chat_auto_execute_logs_the_answer_and_bills_both_rounds(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert response.status_code == 200, response.text + chunks: Final = _streamed_chunks(response) + text: Final = "".join( + str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices") + ) + assert text == ANSWER, response.text + usages: Final = tuple(chunk["usage"] for chunk in chunks if chunk.get("usage")) + assert usages, response.text + rows: Final = _spend_rows(response.headers["x-litellm-call-id"], "acompletion", "success") + assert len(rows) == 1, rows + row: Final = rows[0] + assert row["request_id"] == chunks[0]["id"], (row, chunks[0]) + assert (row["prompt_tokens"], row["completion_tokens"]) == ( + FIRST_ROUND[0] + FOLLOW_UP_ROUND[0], + FIRST_ROUND[1] + FOLLOW_UP_ROUND[1], + ), row + assert float(str(row["spend"])) == pytest.approx(_price(FIRST_ROUND, FOLLOW_UP_ROUND)), row + final_usage: Final = usages[-1] + assert (final_usage["prompt_tokens"], final_usage["completion_tokens"], final_usage["total_tokens"]) == ( + row["prompt_tokens"], + row["completion_tokens"], + sum((*FIRST_ROUND, *FOLLOW_UP_ROUND)), + ), usages + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_bills_round_one_once(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + succeeded: Final = _spend_rows(call_id, "acompletion", "success") + assert len(succeeded) == 1, succeeded + assert (succeeded[0]["prompt_tokens"], succeeded[0]["completion_tokens"]) == FIRST_ROUND, succeeded + assert float(str(succeeded[0]["spend"])) == pytest.approx(_price(FIRST_ROUND)), succeeded + failed: Final = _spend_rows(call_id, "acompletion", "failure") + assert [float(str(row["spend"])) for row in failed] == [0.0], failed + + +def _guardrail_information(metadata: JsonValue) -> tuple[JsonValue, ...]: + entries: Final = object_value(metadata).get("guardrail_information") + return tuple(entries) if isinstance(entries, list) else () + + +def _round_usage_seen_by_the_guardrail(call_id: str) -> list[JsonValue]: + rows: Final = read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s AND status=%s', + (call_id, "acompletion", "success"), + ) + entries: Final = tuple(chain.from_iterable(_guardrail_information(row["metadata"]) for row in rows)) + return [ + entry.get("guardrail_response") + for entry in entries + if isinstance(entry, dict) and entry.get("guardrail_name") == ROUND_USAGE_GUARDRAIL + ] + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_runs_the_post_call_guardrail_on_round_one( + guarded_gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + succeeded: Final = _spend_rows(call_id, "acompletion", "success") + assert len(succeeded) == 1, succeeded + assert (succeeded[0]["prompt_tokens"], succeeded[0]["completion_tokens"]) == FIRST_ROUND, succeeded + assert _round_usage_seen_by_the_guardrail(call_id) == [ + {"prompt_tokens": FIRST_ROUND[0], "completion_tokens": FIRST_ROUND[1]} + ] + + +def _chat_rows(call_id: str, statuses: Sequence[str]) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s', + (call_id, "acompletion"), + ), + lambda found: _statuses(found) == sorted(statuses), + seconds=70, + ) + + +def _entries_of(guardrail_name: str, row: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + entry + for entry in _guardrail_information(row["metadata"]) + if isinstance(entry, dict) and entry.get("guardrail_name") == guardrail_name + ) + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_counts_its_post_call_content_filter_once( + guarded_gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send( + _granted_key(rig), + [AUTO], + stream=True, + stream_options={"include_usage": True}, + guardrails=[CONTENT_FILTER], + ) + assert FOLLOW_UP_OUTAGE in response.text, response.text + rows: Final = _chat_rows(response.headers["x-litellm-call-id"], ("success", "failure")) + filtered: Final = tuple( + (row["status"], entry.get("guardrail_status")) for row in rows for entry in _entries_of(CONTENT_FILTER, row) + ) + assert filtered == (("failure", "guardrail_failed_to_respond"),), rows + (guardrail_id,) = frozenset( + str(entry.get("guardrail_id") or CONTENT_FILTER) + for row in rows + for entry in _entries_of(CONTENT_FILTER, row) + ) + counted: Final = eventually( + lambda: read_rows( + 'SELECT requests_evaluated, flagged_count FROM "LiteLLM_DailyGuardrailMetrics" WHERE guardrail_id=%s', + (guardrail_id,), + ), + lambda found: bool(found), + seconds=70, + ) + assert counted == [{"requests_evaluated": 1, "flagged_count": 1}], counted + + +@pytest.mark.parametrize("include_usage", (True, False), ids=("include-usage", "no-usage")) +def test_streaming_chat_auto_execute_answered_in_one_round_ends_with_the_usage_it_billed( + gateway: Gateway, include_usage: bool +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send( + _granted_key(rig), + [AUTO], + prompt=f"{DIRECT} {rig.alias}", + stream=True, + stream_options={"include_usage": include_usage}, + ) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + chunks: Final = _streamed_chunks(response) + text: Final = "".join( + str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices") + ) + assert text == ANSWER, response.text + rows: Final = _spend_rows(response.headers["x-litellm-call-id"], "acompletion", "success") + assert len(rows) == 1 and (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == FIRST_ROUND, rows + usage_at: Final = tuple(index for index, chunk in enumerate(chunks) if chunk.get("usage")) + if not include_usage: + assert usage_at == (), response.text + return + assert usage_at == (len(chunks) - 1,), response.text + finished_at: Final = tuple( + index + for index, chunk in enumerate(chunks) + if chunk.get("choices") and chunk["choices"][0].get("finish_reason") is not None + ) + assert finished_at and finished_at[-1] < usage_at[0], response.text + usage: Final = chunks[-1]["usage"] + assert (usage["prompt_tokens"], usage["completion_tokens"], usage["total_tokens"]) == ( + *FIRST_ROUND, + sum(FIRST_ROUND), + ), usage + + NO_CACHE: Final = {"no-cache": True} @@ -1269,6 +1495,27 @@ def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_an assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows +def test_streaming_chat_follow_up_broken_mid_answer_bills_round_one_and_its_partial_answer_with_one_fixed_margin( + gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True) + with _rig(gateway, "chat", upstream="messages_bridge", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert "the sum" in response.text, response.text + assert ANSWER not in response.text, response.text + assert len(rig.upstream_tools()) == 2 + round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one)) + assert _statuses(rows) == ["failure", "success"], rows + (success,) = (row for row in rows if row["status"] == "success") + (failure,) = (row for row in rows if row["status"] == "failure") + assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows + assert float(str(success["spend"])) == pytest.approx(round_one), rows + partial_answer: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"]))) + assert min(partial_answer) > 0, rows + assert float(str(failure["spend"])) == pytest.approx(_price(partial_answer)), rows + + @pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once( gateway: Gateway, upstream: StreamedUpstream diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 2d08609c93b..2db8941d304 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -83,6 +83,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( team_member_budget_entity_id, ) from litellm.proxy.utils import ProxyLogging +from litellm.responses.mcp.parked_stream_dispatch import IN_FLIGHT_ROUND from litellm.router import Router from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.utils import ModelResponseListIterator @@ -9051,6 +9052,33 @@ class TestStreamingClientDisconnectBilling: proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + @pytest.mark.asyncio + async def test_disconnect_whose_in_flight_round_billing_raises_still_closes_the_stream_and_frees_the_slot(self): + import types + + response = await self._start_partial_stream() + closed: Final = AsyncMock() + + async def cannot_price_the_in_flight_round() -> None: + raise RuntimeError("the in-flight round could not be priced") + + setattr(response.logging_obj, IN_FLIGHT_ROUND, cannot_price_the_in_flight_round) + stream: Final = types.SimpleNamespace(chunks=[], messages=None, aclose=closed) + proxy_logging_obj: Final = types.SimpleNamespace(arelease_max_parallel_requests_on_disconnect=AsyncMock()) + + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=stream, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=proxy_logging_obj, + ) + + closed.assert_awaited_once() + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + async def _bill_and_collect_success_event(self, prepare=None, request_data=None): recorder = _RecordingSuccessLogger() original_callbacks = litellm.callbacks diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ec5968c2a1a..c0d32742e8a 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1,17 +1,25 @@ import asyncio +import gc import json +import queue import sys +import threading import types -from collections.abc import Mapping, Sequence +import warnings +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager from datetime import datetime from types import MappingProxyType from typing import Final, cast -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from httpx import Response from mcp.types import Tool as MCPTool +from starlette.responses import StreamingResponse +from starlette.types import Message import litellm from litellm.integrations.custom_logger import CustomLogger @@ -25,7 +33,8 @@ from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.utils import ModelResponse, StandardLoggingPayload +from litellm.types.utils import ModelResponse, ModelResponseStream, StandardLoggingPayload, Usage +from litellm.utils import CustomStreamWrapper @pytest.mark.asyncio @@ -1556,6 +1565,7 @@ _CHAT_ROUND_TWO: Final = MappingProxyType( ), } ) +_NEVER_LOGGED_SECONDS: Final = 60 _CONTEXT_WINDOW_ERROR: Final = MappingProxyType( { "error": MappingProxyType( @@ -1575,16 +1585,31 @@ class _ChatRoundRecorder(CustomLogger): super().__init__() self.logged: tuple[StandardLoggingPayload, ...] = () self.failed: tuple[StandardLoggingPayload, ...] = () + self._recorded: Final = asyncio.Condition() async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + await self._announce() async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + await self._announce() + + async def _announce(self) -> None: + async with self._recorded: + self._recorded.notify_all() + + async def settled(self, logged: int = 0, failed: int = 0) -> None: + async with self._recorded: + await asyncio.wait_for( + self._recorded.wait_for(lambda: len(self.logged) >= logged and len(self.failed) >= failed), + timeout=_NEVER_LOGGED_SECONDS, + ) + await _settled() def _chat_round_cost(body: Mapping[str, object]) -> float: @@ -1600,18 +1625,20 @@ async def _settled() -> None: await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) -async def _auto_execute_chat_echo( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], call_id: str -) -> tuple[_ChatRoundRecorder, ModelResponse | Exception]: +def _serve_chat_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response] +) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") litellm.in_memory_llm_clients_cache.flush_cache() respx_mock.post("https://mcp-chat-rounds.example.com/v1/chat/completions").mock(side_effect=list(rounds)) - recorder: Final = _ChatRoundRecorder() - logging_obj: Final = Logging( + + +def _chat_rounds_logging_obj(recorder: _ChatRoundRecorder, call_id: str, stream: bool) -> Logging: + return Logging( model="openai/gpt-5.6", messages=[{"role": "user", "content": "Call echo with PING"}], - stream=False, + stream=stream, call_type="acompletion", start_time=datetime(2026, 1, 1), litellm_call_id=call_id, @@ -1619,6 +1646,10 @@ async def _auto_execute_chat_echo( dynamic_async_success_callbacks=[recorder], dynamic_async_failure_callbacks=[recorder], ) + + +@contextmanager +def _echo_tool_served() -> Iterator[None]: echo_tool: Final = MCPTool.model_validate( {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, by_name=False, @@ -1635,6 +1666,16 @@ async def _auto_execute_chat_echo( AsyncMock(return_value=[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}]), ), ): + yield + + +async def _auto_execute_chat_echo( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], call_id: str +) -> tuple[_ChatRoundRecorder, ModelResponse | Exception]: + _serve_chat_rounds(respx_mock, monkeypatch, rounds) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, call_id, stream=False) + with _echo_tool_served(): try: response: Final = await litellm.acompletion( model="openai/gpt-5.6", @@ -1712,3 +1753,994 @@ async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_follo assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71, 19) assert recorder.failed[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) assert recorder.failed[0]["response_cost"] > 0 + + +_FIXED_MARGIN: Final = 0.001 + + +def _chat_stream_events(body: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + completed: Final = json.loads(json.dumps(body, default=dict)) + choice: Final = completed["choices"][0] + message: Final = choice["message"] + calls: Final = message.get("tool_calls") or [] + delta: Final = ( + {**message, "tool_calls": [{**call, "index": index} for index, call in enumerate(calls)]} if calls else message + ) + head: Final = {"id": completed["id"], "object": "chat.completion.chunk", "created": 1, "model": completed["model"]} + return ( + {**head, "choices": [{"index": 0, "delta": delta, "finish_reason": None}]}, + {**head, "choices": [{"index": 0, "delta": {}, "finish_reason": choice["finish_reason"]}]}, + {**head, "choices": [], "usage": completed["usage"]}, + ) + + +def _sse(events: Sequence[Mapping[str, object]]) -> bytes: + return b"".join(f"data: {json.dumps(event)}\n\n".encode() for event in events) + + +def _streamed(body: Mapping[str, object]) -> Response: + return Response( + 200, + content=_sse(_chat_stream_events(body)) + b"data: [DONE]\n\n", + headers={"content-type": "text/event-stream"}, + ) + + +class _CutOnceTheAnswerStarts(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + raise httpx.ReadError("connection reset mid-stream") + + +def _cut_once_the_answer_starts() -> Response: + return Response(200, stream=_CutOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"}) + + +async def _open_chat_echo_stream( + logging_obj: Logging, metadata: dict[str, object] | None = None +) -> CustomStreamWrapper: + stream: Final = await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + api_key="sk-mcp-chat-rounds", + api_base="https://mcp-chat-rounds.example.com/v1", + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + metadata=metadata, + ) + assert isinstance(stream, CustomStreamWrapper) + return stream + + +async def _read_all(stream: CustomStreamWrapper) -> tuple[ModelResponseStream, ...]: + return tuple([chunk async for chunk in stream]) + + +async def _read_until(stream: CustomStreamWrapper, content: str) -> None: + async for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content == content: + return + pytest.fail(f"the stream never sent {content!r}") + + +def _streamed_usage(chunks: Sequence[ModelResponseStream]) -> Usage: + usages: Final = tuple(usage for chunk in chunks if isinstance(usage := getattr(chunk, "usage", None), Usage)) + assert len(usages) == 1, usages + return usages[0] + + +def _assert_one_stream_row_bills_both_rounds_once( + logged: Sequence[StandardLoggingPayload], chunks: Sequence[ModelResponseStream] +) -> None: + charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN + assert len(logged) == 1 + row: Final = logged[0] + assert row["id"] == chunks[0].id + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert row["metadata"]["usage_object"]["cost"] == pytest.approx(row["response_cost"]) + usage: Final = _streamed_usage(chunks) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71 + 149, 19 + 5, 90 + 154) + assert usage.cost == pytest.approx(row["response_cost"]) + + +def _prompt_tokens_the_follow_up_sent(respx_mock: respx.MockRouter) -> int: + follow_up_request: Final = json.loads(respx_mock.calls[1].request.content) + return litellm.token_counter(model="gpt-5.6", messages=follow_up_request["messages"]) + + +def _assert_one_row_bills_round_one(logged: Sequence[StandardLoggingPayload]) -> None: + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (71, 19) + assert logged[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_read_to_the_end_logs_one_row_whose_usage_cost_counts_the_fixed_margin_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-read-to-the-end", stream=True) + + with _echo_tool_served(): + chunks: Final = await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_when_the_proxy_releases( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-rounds-deferred", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_all(stream) + await _settled() + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await recorder.settled(logged=1) + + assert logged_before_release == () + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_rounds_on_the_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.exceptions import GuardrailRaisedException + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-rounds-blocked", stream=True) + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_all(stream) + assert ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked) + await logging_obj.async_failure_handler(blocked, "blocked by output-policy") + await recorder.settled(failed=1) + + charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71 + 149, 19 + 5) + assert recorder.failed[0]["response_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_rejected_by_the_proxy_during_the_follow_up_bills_round_one_with_its_own_breakdown( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from fastapi import HTTPException + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-rejected-mid-follow-up", stream=True) + rejected: Final = HTTPException(status_code=400, detail="rejected by a stream hook") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_until(stream, "PING") + assert ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, rejected) + await logging_obj.async_failure_handler(rejected, "rejected by a stream hook") + await recorder.settled(failed=1) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + failed: Final = recorder.failed[0] + breakdown: Final = failed["cost_breakdown"] + assert (failed["prompt_tokens"], failed["completion_tokens"]) == (71, 19) + assert failed["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) + assert breakdown is not None + assert breakdown["total_cost"] == pytest.approx(failed["response_cost"]) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round_once_the_answer_finishes( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_TWO),)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-one-round", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + async for chunk in stream: + if chunk.choices and chunk.choices[0].finish_reason is not None: + break + await recorder.settled(logged=1) + logged_once_the_answer_finished: Final = recorder.logged + after_the_answer: Final = await _read_all(stream) + await _settled() + + assert len(logged_once_the_answer_finished) == 1 + assert recorder.logged == logged_once_the_answer_finished + assert recorder.logged[0]["id"] == _CHAT_ROUND_TWO["id"] + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (149, 5) + assert after_the_answer[-1:] and isinstance(after_the_answer[-1].usage, Usage) + usage: Final = _streamed_usage(after_the_answer) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + recorder.logged[0]["prompt_tokens"], + recorder.logged[0]["completion_tokens"], + recorder.logged[0]["total_tokens"], + ) + assert recorder.logged[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_TWO)) + assert recorder.failed == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("closes", (True, False), ids=("aclose", "dropped")) +async def test_stream_auto_execute_bills_the_finished_first_round_when_the_caller_stops_reading_the_follow_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, closes: bool +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-caller-stops", stream=True) + + async def stop_reading_once_the_follow_up_answers() -> None: + stream: Final = await _open_chat_echo_stream(logging_obj) + await _read_until(stream, "PING") + if closes: + await stream.aclose() + + with _echo_tool_served(): + await stop_reading_once_the_follow_up_answers() + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_first_round_breaks_mid_answer_bills_only_its_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_cut_once_the_answer_starts(),)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-first-round-breaks", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + with pytest.raises(litellm.ServiceUnavailableError): + await _read_all(stream) + await recorder.settled(failed=1) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert recorder.failed[0]["completion_tokens"] > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging, _failure_fields_to_lift + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-breaks", stream=True) + request_data: Final = MappingProxyType({"litellm_logging_obj": logging_obj}) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError) as broken: + await _read_all(stream) + if not ProxyLogging._discard_deferred_stream_logging_for_failure(dict(request_data), broken.value): + ProxyLogging.fire_deferred_stream_logging(dict(request_data)) + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert len(recorder.failed) == 1 + failed: Final = recorder.failed[0] + unmargined: Final = ( + litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=failed["prompt_tokens"], + completion_tokens=failed["completion_tokens"], + total_tokens=failed["total_tokens"], + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + - _FIXED_MARGIN + ) + assert failed["prompt_tokens"] == _prompt_tokens_the_follow_up_sent(respx_mock) + assert failed["completion_tokens"] > 0 + assert failed["response_cost"] == pytest.approx(unmargined) + request_failure: Final = _failure_fields_to_lift(request_data) + assert cast(StandardLoggingPayload, request_failure["standard_logging_object"])["id"] != recorder.logged[0]["id"] + assert request_failure["response_cost"] == pytest.approx(unmargined) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_whose_follow_up_breaks_reports_the_request_guardrail_result_on_the_failure_row_only( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-breaks-guarded", stream=True) + metadata: Final[dict[str, object]] = {} + request_data: Final = {"metadata": metadata, "litellm_logging_obj": logging_obj} + content_filter: Final = CustomGuardrail(guardrail_name="request-content-filter", event_hook="post_call") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj, metadata) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError) as broken: + await _read_all(stream) + content_filter.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=str(broken.value), + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + ) + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, broken.value): + ProxyLogging.fire_deferred_stream_logging(request_data) + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert recorder.logged[0]["guardrail_information"] is None + assert [entry["guardrail_name"] for entry in recorder.failed[0]["guardrail_information"] or ()] == [ + "request-content-filter" + ] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_served_through_the_proxy_stream_hook_logs_both_rounds_once_when_it_ends( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-proxy-stream-hook", stream=True) + + with _echo_tool_served(): + served: Final = ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=await _open_chat_echo_stream(logging_obj), + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj}, + ) + chunks: Final = tuple([chunk async for chunk in served]) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_follow_up_streams_nothing_still_bills_the_first_round_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + nothing_streamed: Final = Response(200, content=b"data: [DONE]\n\n", headers={"content-type": "text/event-stream"}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), nothing_streamed)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-empty-follow-up", stream=True) + + with _echo_tool_served(): + await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +_ANSWER_WITH_UNREADABLE_AUDIO: Final = MappingProxyType( + { + **_CHAT_ROUND_TWO, + "choices": ( + MappingProxyType( + { + "index": 0, + "finish_reason": "stop", + "message": MappingProxyType( + { + "role": "assistant", + "content": "PING", + "audio": MappingProxyType({"id": "audio_echo", "data": "not base64", "transcript": "PING"}), + } + ), + } + ), + ), + } +) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with_the_usage_of_the_round_it_billed( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_ANSWER_WITH_UNREADABLE_AUDIO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-unassembled", stream=True) + + with _echo_tool_served(): + chunks: Final = await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert any(chunk.choices and chunk.choices[0].delta.content == "PING" for chunk in chunks) + usage: Final = _streamed_usage(chunks) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71, 19, 90) + assert usage.cost == pytest.approx(recorder.logged[0]["response_cost"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +@pytest.mark.parametrize("answer", [_CHAT_ROUND_TWO, _ANSWER_WITH_UNREADABLE_AUDIO], ids=("assembled", "unassembled")) +async def test_stream_auto_execute_sends_its_usage_chunk_under_the_id_of_the_row_it_billed( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + deferred_for_a_post_call_guardrail: bool, + answer: Mapping[str, object], +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(answer))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-usage-chunk-id", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_all(stream) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await recorder.settled(logged=1) + + usage_chunks: Final = tuple(chunk for chunk in chunks if isinstance(getattr(chunk, "usage", None), Usage)) + answer_chunks: Final = tuple(chunk for chunk in chunks if chunk.choices and chunk.choices[0].delta.content == "PING") + assert len(recorder.logged) == 1 + assert [chunk.id for chunk in answer_chunks] == [_CHAT_ROUND_TWO["id"]] + assert len(usage_chunks) == 1 + assert usage_chunks[0].id == chunks[0].id == recorder.logged[0]["id"] == _CHAT_ROUND_ONE["id"] + + +def _assert_one_row_bills_round_one_and_the_partial_answer( + recorder: _ChatRoundRecorder, respx_mock: respx.MockRouter +) -> None: + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert row["id"] == _CHAT_ROUND_ONE["id"] + partial_prompt, partial_completion = (row["prompt_tokens"] - 71, row["completion_tokens"] - 19) + assert partial_prompt == _prompt_tokens_the_follow_up_sent(respx_mock) + assert partial_completion > 0 + partial_answer: Final = litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=partial_prompt, + completion_tokens=partial_completion, + total_tokens=partial_prompt + partial_completion, + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + assert row["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE) + partial_answer - _FIXED_MARGIN) + assert recorder.failed == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_during_the_follow_up_bills_round_one_and_the_partial_answer_on_one_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-mid-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_until(stream, "PING") + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +def _chat_rounds_router() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "mcp-chat-rounds", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-mcp-chat-rounds", + "api_base": "https://mcp-chat-rounds.example.com/v1", + }, + } + ] + ) + + +class _StallsOnceTheAnswerStarts(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + await asyncio.Event().wait() + + +def _streams_content(body: object, content: str) -> bool: + lines: Final = body.decode().splitlines() if isinstance(body, bytes) else () + events: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: {")) + return any(event["choices"] and event["choices"][0]["delta"].get("content") == content for event in events) + + +async def _serve_until_the_client_hangs_up_after(response: StreamingResponse, content: str) -> None: + read: Final = asyncio.Event() + + async def receive() -> Message: + await read.wait() + return {"type": "http.disconnect"} + + async def send(message: Message) -> None: + if message["type"] == "http.response.body" and _streams_content(message.get("body"), content): + read.set() + + uvicorn_scope: Final = {"type": "http", "asgi": {"version": "3.0", "spec_version": "2.3"}, "headers": []} + await response(uvicorn_scope, receive, send) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_answer_when_the_client_hangs_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, create_response + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds( + respx_mock, + monkeypatch, + ( + _streamed(_CHAT_ROUND_ONE), + Response(200, stream=_StallsOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"}), + ), + ) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-hang-up", stream=True) + + with _echo_tool_served(): + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + ) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + served: Final = await create_response( + proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ), + "text/event-stream", + {}, + ) + assert isinstance(served, StreamingResponse) + await _serve_until_the_client_hangs_up_after(served, "PING") + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +async def _stop_reading_once_the_proxy_sent(served: AsyncIterator[object], content: str) -> None: + async for frame in served: + if _streams_content(frame, content): + return + pytest.fail(f"the proxy never sent {content!r}") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_answer_when_hung_up_between_reads( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-hang-up-between-reads", stream=True) + + with _echo_tool_served(): + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + ) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + served: Final = proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ) + await _stop_reading_once_the_proxy_sent(served, "PING") + await served.aclose() + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +_LONGEST_WAIT_FOR_THE_LOOP_TO_MOVE_SECONDS: Final = 10 + + +class _CountsTheToolResultWhileTheLoopRuns: + def __init__(self, loop: asyncio.AbstractEventLoop, count: Callable[..., int]) -> None: + self._loop: Final = loop + self._count: Final = count + self._counting: Final = asyncio.Event() + self._loop_moved: Final = threading.Event() + self.loop_moved_while_counting: tuple[bool, ...] = () + + def __call__(self, model: str | None = None, messages: list[object] | None = None, **kwargs: object) -> int: + if any(isinstance(message, Mapping) and message.get("role") == "tool" for message in messages or ()): + self._loop.call_soon_threadsafe(self._counting.set) + moved: Final = self._loop_moved.wait(timeout=_LONGEST_WAIT_FOR_THE_LOOP_TO_MOVE_SECONDS) + self.loop_moved_while_counting = (*self.loop_moved_while_counting, moved) + return self._count(model=model, messages=messages, **kwargs) + + async def keep_the_loop_moving(self) -> None: + await self._counting.wait() + self._loop_moved.set() + + +@contextmanager +def _partial_count_watched(monkeypatch: pytest.MonkeyPatch) -> Iterator[_CountsTheToolResultWhileTheLoopRuns]: + from litellm.litellm_core_utils import streaming_chunk_builder_utils + + counter: Final = _CountsTheToolResultWhileTheLoopRuns( + asyncio.get_running_loop(), streaming_chunk_builder_utils.token_counter + ) + monkeypatch.setattr(streaming_chunk_builder_utils, "token_counter", counter) + other_task: Final = asyncio.create_task(counter.keep_the_loop_moving()) + try: + yield counter + finally: + other_task.cancel() + + +class _StallsAfterTheAnswerStarts(httpx.AsyncByteStream): + def __init__(self) -> None: + self.stalled: Final = asyncio.Event() + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + self.stalled.set() + await asyncio.Event().wait() + + +@pytest.mark.asyncio +async def test_stream_auto_execute_cancelled_mid_answer_counts_the_partial_off_the_loop_and_bills_it_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + answer: Final = _StallsAfterTheAnswerStarts() + _serve_chat_rounds( + respx_mock, + monkeypatch, + (_streamed(_CHAT_ROUND_ONE), Response(200, stream=answer, headers={"content-type": "text/event-stream"})), + ) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-cancelled-off-the-loop", stream=True) + + with _echo_tool_served(), _partial_count_watched(monkeypatch) as counter: + stream: Final = await _open_chat_echo_stream(logging_obj) + reading: Final = asyncio.create_task(_read_all(stream)) + await asyncio.wait_for(answer.stalled.wait(), timeout=10) + reading.cancel() + with pytest.raises(asyncio.CancelledError): + await reading + await stream.aclose() + await recorder.settled(logged=1) + + assert counter.loop_moved_while_counting == (True,) + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_routed_through_the_proxy_and_closed_between_reads_counts_the_partial_off_the_loop( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-closed-off-the-loop", stream=True) + + with _echo_tool_served(), _partial_count_watched(monkeypatch) as counter: + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + ) + served: Final = proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ) + await _stop_reading_once_the_proxy_sent(served, "PING") + await served.aclose() + await recorder.settled(logged=1) + + assert counter.loop_moved_while_counting == (True,) + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +class _StallsBeforeTheAnswerStarts(httpx.AsyncByteStream): + def __init__(self) -> None: + self.opened: Final = asyncio.Event() + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b": the answer has not started\n\n" + self.opened.set() + await asyncio.Event().wait() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_before_the_answer_starts_bills_round_one_alone_on_one_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + answer: Final = _StallsBeforeTheAnswerStarts() + _serve_chat_rounds( + respx_mock, + monkeypatch, + (_streamed(_CHAT_ROUND_ONE), Response(200, stream=answer, headers={"content-type": "text/event-stream"})), + ) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-before-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + reading: Final = asyncio.create_task(_read_all(stream)) + await asyncio.wait_for(answer.opened.wait(), timeout=10) + reading.cancel() + with pytest.raises(asyncio.CancelledError): + await reading + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert recorder.failed == () + + +async def _read_through_the_usage_chunk( + chunks: AsyncIterator[ModelResponseStream], read: tuple[ModelResponseStream, ...] = () +) -> tuple[ModelResponseStream, ...]: + chunk: Final = await anext(chunks) + if isinstance(getattr(chunk, "usage", None), Usage): + return (*read, chunk) + return await _read_through_the_usage_chunk(chunks, (*read, chunk)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_after_its_whole_answer_bills_both_rounds_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-after-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_through_the_usage_chunk(aiter(stream)) + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + assert recorder.failed == () + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bridged_to_messages_bills_round_one_when_its_follow_up_breaks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-bridged-to-messages", stream=True) + + async def bridged(stream: CustomStreamWrapper) -> AsyncIterator[ModelResponseStream]: + async for chunk in stream: + yield chunk + + with _echo_tool_served(): + relayed: Final = bridged(await _open_chat_echo_stream(logging_obj)) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=relayed, route_type="anthropic_messages", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError): + _ = [chunk async for chunk in relayed] + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +def _read_synchronously_until_the_follow_up_answers( + logging_obj: Logging, loop: asyncio.AbstractEventLoop, before_dropping: Callable[[], None] = lambda: None +) -> None: + stream: Final = loop.run_until_complete(_open_chat_echo_stream(logging_obj)) + for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content == "PING": + before_dropping() + return + pytest.fail("the stream never sent 'PING'") + + +_CHAT_STREAM_BILLING_NAMES: Final = ("MCPStreamingIterator", "ChatStreamRoundBilling", "_log_success") + + +def _names_chat_stream_billing(text: str) -> bool: + return any(name in text for name in _CHAT_STREAM_BILLING_NAMES) + + +class _UnraisableRecorder: + def __init__(self) -> None: + self.raised: tuple[str, ...] = () + + def __call__(self, unraisable: "sys.UnraisableHookArgs") -> None: + self.raised = (*self.raised, f"{unraisable.object!r}: {unraisable.exc_value!r}") + + +def test_stream_auto_execute_read_synchronously_and_dropped_mid_answer_bills_the_finished_first_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + unraisable: Final = _UnraisableRecorder() + loop: Final = asyncio.new_event_loop() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + asyncio.set_event_loop(loop) + + try: + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _read_synchronously_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped", stream=True), loop + ) + gc.collect() + loop.run_until_complete(recorder.settled(logged=1)) + finally: + asyncio.set_event_loop(None) + loop.close() + + _assert_one_row_bills_round_one(recorder.logged) + assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] + + +def test_stream_auto_execute_read_synchronously_and_dropped_after_its_loop_closed_drops_without_errors( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + unraisable: Final = _UnraisableRecorder() + loop: Final = asyncio.new_event_loop() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + asyncio.set_event_loop(loop) + + try: + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _read_synchronously_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped-after-close", stream=True), + loop, + before_dropping=loop.close, + ) + gc.collect() + finally: + asyncio.set_event_loop(None) + loop.close() + + assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] diff --git a/tests/unit/responses/mcp/test_parked_stream_dispatch.py b/tests/unit/responses/mcp/test_parked_stream_dispatch.py new file mode 100644 index 00000000000..dd380c582c8 --- /dev/null +++ b/tests/unit/responses/mcp/test_parked_stream_dispatch.py @@ -0,0 +1,17 @@ +from typing import Final + +import pytest + +from litellm.responses.mcp.parked_stream_dispatch import parked_success +from litellm.types.utils import ModelResponse + +_ASSEMBLED: Final = ModelResponse(id="chatcmpl-parked", model="gpt-5.6") + + +@pytest.mark.parametrize( + "parked", + [(), (_ASSEMBLED,), (_ASSEMBLED, False, "not a logger"), (_ASSEMBLED, False, None, None)], + ids=("empty", "assembled-alone", "third-part-is-not-a-logger", "four-parts"), +) +def test_a_parked_dispatch_of_an_unknown_shape_is_never_dispatched(parked: tuple[object, ...]): + assert parked_success(parked) is None From 4786dd861b1b8c001117e0403cc2c606c193896c Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:36:01 -0700 Subject: [PATCH 27/64] fix(proxy)!: scope /health by wildcard keys and team model allowlists the way auth does (#49155) * test(proxy): cover /health scoping for wildcard keys and team model allowlists * fix(proxy): scope /health with the model access check auth runs on a request * fix(proxy): scope /health to the model group an allowed name routes to A key alias, team alias, or router model_group_alias lets auth admit a requested name by checking its target, so /health also kept the deployment under the requested name, one the key's requests never reach. Keep only the served model group each admitted name dispatches to, through the same alias layers /v1/models mirrors, and let alias_map read its own read-only output. * test(proxy): cover /health skipping a deployment that has no model name for a restricted key * fix(health): show deployments reached through wildcard and alias allowlists, and by id when an alias shadows the name * chore(router): drop should_include_deployment from the outside-use allowlist now that /health no longer calls it * fix(health): never list a default-fallback deployment for an unserved grant, and keep access-group keys to their group's deployments * fix(health): resolve ?model= through the caller's aliases the way chat does * fix(proxy): apply user and agent model limits to /health and yield between checks /health now runs the same user and agent model checks chat runs, so a personal key whose owner may call only some models, or an agent key limited by its access group, no longer sees or probes other deployments. Each per-name check yields to the event loop, and ?model= reaches the wildcard deployment serving an allowed name * fix(proxy): check owner and agent grants for wildcard routes in /health and fail on non-access auth errors * fix(proxy): answer /health auth lookup failures the way chat does and record ?model= checks under the requested name * test(proxy): pass recorded_model_name in the direct health save tests * fix(proxy): count alias names as /health candidates so aliased names a key may call get probed * fix(health): narrow wildcard-served deployments to the key's access groups * fix(health): narrow every deployment serving a name to the key's access groups, team-owned ones included * fix(health): narrow a team deployment's internal name to the key's access groups like its public name * fix(health): keep a team deployment's internal name narrow when a limit holds only that name An agent or key limit that holds one team deployment's internal name no longer reaches every deployment behind the team's public name. The internal name maps to its public name only when the public name is also admitted. * test(health): cover all-proxy-models next to * for a key that also holds an access group * refactor(auth): keep the _as_proxy_exception name so main's tests still find it * fix(health): keep an internal team name a wildcard grant admits as its own target --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 47 +- .../proxy/common_utils/model_listing_utils.py | 11 +- .../health_endpoints/_health_endpoints.py | 427 ++++++++++- ...st_auth_checks_object_access_and_lookup.py | 10 + .../health_endpoints/test_health_endpoints.py | 720 +++++++++++++++++- .../unit/proxy/test_health_check_functions.py | 2 + 6 files changed, 1164 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b8d9530d2eb..f80296f4bd7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -84,6 +84,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AgentAccessGroupCeiling, CeilingResolver, resolve_agent_access_group_ceiling, ) @@ -1684,6 +1685,7 @@ def get_actual_routes(allowed_routes: list) -> list: KEY_END_USER_BUDGET_ID_METADATA_FIELD: Final = "end_user_budget_id" _KEY_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) +_OBJECTS_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: @@ -5085,6 +5087,36 @@ async def can_agent_call_model( ) +async def _agent_access_group_ceilings( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver +) -> tuple[AgentAccessGroupCeiling, ...]: + if not valid_token.agent_id: + return () + + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings + + managed: Final = managed_agent_policy(valid_token) + if managed is not None: + return await resolve_managed_agent_ceilings(managed) + unmanaged: Final = await resolve_ceiling(valid_token.agent_id) + return (unmanaged,) if unmanaged is not None else () + + +def _managed_agent_model_grants(valid_token: UserAPIKeyAuth) -> frozenset[str]: + managed: Final = managed_agent_policy(valid_token) + granted: Final = (managed.object_permission or {}).get("models") if managed is not None else None + if not isinstance(granted, (list, tuple)): + return frozenset() + return frozenset(name for name in _OBJECTS_ADAPTER.validate_python(granted) if isinstance(name, str)) + + +async def agent_model_grants( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling +) -> frozenset[str]: + ceilings: Final = await _agent_access_group_ceilings(valid_token, resolve_ceiling) + return _managed_agent_model_grants(valid_token).union(*(ceiling.models for ceiling in ceilings)) + + async def _check_agent_access_group_model_access( model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str] valid_token: UserAPIKeyAuth | None, @@ -5092,20 +5124,9 @@ async def _check_agent_access_group_model_access( resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> Literal[True]: """Attached groups naming no model deny every model; the empty allowlist in ``_can_object_call_model`` allows.""" - if not model or valid_token is None or not valid_token.agent_id: + if not model or valid_token is None: return True - - from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings - - managed: Final = managed_agent_policy(valid_token) - unmanaged: Final = await resolve_ceiling(valid_token.agent_id) if managed is None else None - ceilings: Final = ( - await resolve_managed_agent_ceilings(managed) - if managed is not None - else (unmanaged,) - if unmanaged is not None - else () - ) + ceilings: Final = await _agent_access_group_ceilings(valid_token, resolve_ceiling) dispatched: Final = _resolve_team_alias( model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router ) diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index ad4c3bd8770..a9424018d9a 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -219,7 +219,7 @@ def caller_alias_maps( def alias_map(aliases: object) -> Mapping[str, str]: try: - entries: Final = _ALIAS_ENTRIES.validate_python(aliases, strict=True) + entries: Final = _ALIAS_ENTRIES.validate_python(aliases) except ValidationError: return _NO_ALIASES return MappingProxyType( @@ -287,6 +287,15 @@ def _dispatch(model_id: str, layers: Sequence[_AliasLayer]) -> tuple[str, bool]: return reduce(_through_layer, layers, (model_id, False)) +def dispatched_model(model_id: str, aliases: CallerAliases) -> str: + target, _ = _dispatch(model_id, _alias_layers(aliases)) + return target + + +def requestable_alias_names(aliases: CallerAliases) -> frozenset[str]: + return frozenset(chain.from_iterable(layer.aliases for layer in _alias_layers(aliases))) + + def _lookup_id(name: str, lookup_by_response: Mapping[str, str], lookup_ids: Container[str]) -> str | None: if name in lookup_by_response: return lookup_by_response[name] diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 59c63a18f0f..8545e0609db 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -7,13 +7,15 @@ import secrets import time import traceback from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from itertools import chain, groupby from types import MappingProxyType from typing import Final, Literal, TypedDict import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response, status -from pydantic import TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter from typing_extensions import ReadOnly import litellm @@ -31,7 +33,9 @@ from litellm.proxy._types import ( CallInfo, EnterpriseLicenseData, Litellm_EntityType, + LiteLLM_UserTable, LitellmUserRoles, + ModelAccessDeniedProxyException, ProxyException, SpecialModelNames, UserAPIKeyAuth, @@ -39,13 +43,28 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import ( _resolve_key_models_for_auth_check, # pyright: ignore[reportPrivateUsage] # the auth layer's sentinel resolution, reused so /health scopes exactly like a request + agent_model_grants, + can_agent_call_model, + can_key_call_resolved_model, + can_user_call_model, + get_user_object, + key_model_aliases_for_auth_check, + model_matches_any_wildcard_pattern_in_list, + team_model_aliases_for_auth_check, +) +from litellm.proxy.auth.auth_exception_handler import ( + _as_proxy_exception, # pyright: ignore[reportPrivateUsage] # one auth-failure mapping, shared with user_api_key_auth ) from litellm.proxy.auth.auth_utils import ( _BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check ) -from litellm.proxy.auth.model_checks import get_key_models from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils import http_parsing_utils +from litellm.proxy.common_utils.model_listing_utils import ( + CallerAliases, + dispatched_model, + requestable_alias_names, +) from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -79,6 +98,7 @@ from litellm.router_utils.clientside_credential_handler import ( _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the router path clientside_credential_keys, ) +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.secret_managers.main import get_secret_bool from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields from litellm.types.router import reject_server_owned_wif_fields @@ -253,6 +273,28 @@ def get_callback_identifier(callback): router: Final = APIRouter() _OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECTS: Final = TypeAdapter(tuple[object, ...]) +_DEPLOYMENTS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +class _ModelGrants(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + models: tuple[object, ...] = () + + +class _ModelInfoAccessGroups(BaseModel): + model_config = ConfigDict(frozen=True) + + access_groups: tuple[object, ...] + + +class _DeploymentAccessGroups(BaseModel): + model_config = ConfigDict(frozen=True) + + model_info: _ModelInfoAccessGroups + + services = ( Literal[ "slack_budget_alerts", @@ -1045,42 +1087,339 @@ def _strip_admin_only_fields_from_health_result(result: dict) -> dict: return out -def _health_accessible_model_names( - user_api_key_dict: UserAPIKeyAuth, llm_router: Router | None -) -> frozenset[str] | None: - """Model names the caller may health-check, or None when the key is unrestricted.""" +def _key_models_unrestricted(user_api_key_dict: UserAPIKeyAuth) -> bool: granted_models: Final = _resolve_key_models_for_auth_check(user_api_key_dict) - if not granted_models or SpecialModelNames.all_proxy_models.value in granted_models: + return not granted_models or SpecialModelNames.all_proxy_models.value in granted_models + + +def _own_team_names(deployment: object, team_id: str | None) -> tuple[str, str] | None: + match deployment: + case { + "model_name": str(internal_name), + "model_info": {"team_id": str(owner_team_id), "team_public_model_name": str(public_name)}, + } if owner_team_id == team_id: + return (internal_name, public_name) + case _: + return None + + +def _names_caller_may_request(deployment: object, team_id: str | None) -> frozenset[str]: + team_names: Final = _own_team_names(deployment, team_id) + if team_names is not None: + return frozenset(team_names) + match deployment: + case {"model_name": str(name)}: + return frozenset((name,)) + case _: + return frozenset() + + +def _public_names_by_internal_name(llm_model_list: Sequence[object], team_id: str | None) -> Mapping[str, str]: + team_names: Final = (_own_team_names(deployment, team_id) for deployment in llm_model_list) + return MappingProxyType(dict(names for names in team_names if names is not None)) + + +async def _key_owner(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_UserTable | None: + if user_api_key_dict.team_id is not None or user_api_key_dict.user_id is None: return None + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + try: + return await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: # noqa: BLE001 # auth reads an unloadable user as no user limit, so /health does too + return None + + +async def _key_may_call_model( + model_name: str, + user_api_key_dict: UserAPIKeyAuth, + owner: LiteLLM_UserTable | None, + llm_router: Router | None, +) -> bool: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + await asyncio.sleep(0) + try: + await can_key_call_resolved_model( + model=model_name, llm_model_list=None, valid_token=user_api_key_dict, llm_router=llm_router + ) + await can_agent_call_model( + model=model_name, + valid_token=user_api_key_dict, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await can_user_call_model( + model=model_name, + llm_router=llm_router, + user_object=owner, + key_model_aliases=key_model_aliases_for_auth_check(user_api_key_dict), + ) + except ModelAccessDeniedProxyException: + return False + except ProxyException: + raise + except Exception as failure: + if isinstance(failure, HTTPException) and failure.status_code == status.HTTP_403_FORBIDDEN: + return False + raise _as_proxy_exception(failure) from failure + return True + + +async def _caller_aliases(user_api_key_dict: UserAPIKeyAuth) -> CallerAliases: + from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj + + router_settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict, prisma_client, proxy_logging_obj + ) + return CallerAliases( + team_model_aliases_for_auth_check(user_api_key_dict), + key_model_aliases_for_auth_check(user_api_key_dict), + router_settings.get("model_group_alias") if router_settings is not None else None, + ) + + +def _model_auth_checks(model_name: str, aliases: CallerAliases) -> str: + return resolve_model_group_alias(aliases.router, model_name) or model_name + + +def _deployment_id(deployment: object) -> str | None: + match deployment: + case {"model_info": {"id": str(model_id)}}: + return model_id + case _: + return None + + +def _granted_names(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: + grants: Final = ( + *_ModelGrants.model_validate(user_api_key_dict).models, + *_OBJECTS.validate_python(user_api_key_dict.team_models), + ) + return frozenset(name for name in grants if isinstance(name, str)) + + +async def _names_any_limit_grants(user_api_key_dict: UserAPIKeyAuth, owner: LiteLLM_UserTable | None) -> frozenset[str]: + owner_grants: Final = _ModelGrants.model_validate(owner).models if owner is not None else () + return ( + _granted_names(user_api_key_dict) + | frozenset(name for name in owner_grants if isinstance(name, str)) + | await agent_model_grants(user_api_key_dict) + ) + + +def _alias_names_caller_may_request(aliases: CallerAliases, llm_router: Router | None) -> frozenset[str]: + router_aliases: Final = frozenset(llm_router.registry().aliases()) if llm_router is not None else frozenset[str]() + return requestable_alias_names(aliases) | router_aliases + + +def _served_group(model_name: str, aliases: CallerAliases, llm_router: Router | None) -> str: + dispatched: Final = dispatched_model(model_name, aliases) + return llm_router.registry().resolve(dispatched) if llm_router is not None else dispatched + + +def _name_id_pairs(deployment: object, team_id: str | None) -> tuple[tuple[str, str], ...]: + model_id: Final = _deployment_id(deployment) + if model_id is None: + return () + return tuple((name, model_id) for name in _names_caller_may_request(deployment, team_id) | {model_id}) + + +def _ids_by_requested_name(llm_model_list: Sequence[object], team_id: str | None) -> Mapping[str, frozenset[str]]: + pairs: Final = sorted(chain.from_iterable(_name_id_pairs(deployment, team_id) for deployment in llm_model_list)) + return MappingProxyType( + {name: frozenset(model_id for _, model_id in group) for name, group in groupby(pairs, key=lambda pair: pair[0])} + ) + + +def _ids_of(deployments: Iterable[object]) -> frozenset[str]: + return frozenset(model_id for model_id in map(_deployment_id, deployments) if model_id is not None) + + +def _access_groups_of(deployment: object) -> frozenset[str]: + match deployment: + case {"model_info": {"access_groups": list()}}: + access_groups: Final = _DeploymentAccessGroups.model_validate(deployment).model_info.access_groups + return frozenset(access_group for access_group in access_groups if isinstance(access_group, str)) + case _: + return frozenset() + + +@dataclass(frozen=True, slots=True) +class _AccessGroupGrants: + granted_names: frozenset[str] + ids_in_granted_groups: frozenset[str] + + def requested_name(self, name: str, public_names: Mapping[str, str], admitted_names: frozenset[str]) -> str: + public_name: Final = public_names.get(name) + if public_name is None or public_name not in admitted_names or self._grants(name): + return name + return public_name + + def _grants(self, name: str) -> bool: + return name in self.granted_names or model_matches_any_wildcard_pattern_in_list(name, tuple(self.granted_names)) + + def narrow(self, serving_ids: frozenset[str], served_names: frozenset[str]) -> frozenset[str]: + in_granted_groups: Final = serving_ids & self.ids_in_granted_groups + granted_directly: Final = not self.granted_names.isdisjoint(served_names) + return serving_ids if granted_directly or not in_granted_groups else in_granted_groups + + +def _access_group_grants(llm_model_list: Sequence[object], user_api_key_dict: UserAPIKeyAuth) -> _AccessGroupGrants: + granted: Final = _granted_names(user_api_key_dict) + in_granted_groups: Final = ( + deployment for deployment in llm_model_list if not granted.isdisjoint(_access_groups_of(deployment)) + ) + return _AccessGroupGrants(granted, _ids_of(in_granted_groups)) + + +def _ids_serving_group( + group: str, + ids_by_name: Mapping[str, frozenset[str]], + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router | None, + access_group_grants: _AccessGroupGrants, +) -> frozenset[str]: + named: Final = ids_by_name.get(group) + if named is not None: + return access_group_grants.narrow(named, frozenset((group,))) if llm_router is None: - return frozenset(granted_models) - registry: Final = llm_router.registry() + return frozenset() + team_id: Final = user_api_key_dict.team_id + wildcard_served: Final = llm_router.registry().deployments(group, source="wildcard", team_id=team_id) + served_names: Final = frozenset((group,)).union( + *(_names_caller_may_request(deployment, team_id) for deployment in wildcard_served) + ) + return access_group_grants.narrow(_ids_of(wildcard_served), served_names) + + +async def _names_key_may_request( + names: Iterable[str], + aliases: CallerAliases, + user_api_key_dict: UserAPIKeyAuth, + owner: LiteLLM_UserTable | None, + llm_router: Router | None, +) -> frozenset[str]: return frozenset( - get_key_models( - user_api_key_dict=user_api_key_dict, - proxy_model_list=registry.names(team_id=user_api_key_dict.team_id), - model_access_groups=registry.access_groups(), + [ + name + for name in names + if await _key_may_call_model(_model_auth_checks(name, aliases), user_api_key_dict, owner, llm_router) + ] + ) + + +def _ids_serving( + groups: Iterable[str], + ids_by_name: Mapping[str, frozenset[str]], + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router | None, + access_group_grants: _AccessGroupGrants, +) -> frozenset[str]: + return frozenset[str]().union( + *( + _ids_serving_group(group, ids_by_name, user_api_key_dict, llm_router, access_group_grants) + for group in groups ) ) -def _caller_may_probe_deployment( - deployment: Mapping[str, object], - allowed_models: frozenset[str] | None, +def _shadowed_ids( + llm_model_list: Sequence[object], admitted_names: frozenset[str], reached_ids: frozenset[str], team_id: str | None +) -> frozenset[str]: + ids_of_admitted_names: Final = ( + _deployment_id(deployment) + for deployment in llm_model_list + if not admitted_names.isdisjoint(_names_caller_may_request(deployment, team_id)) + ) + return frozenset(model_id for model_id in ids_of_admitted_names if model_id is not None) - reached_ids + + +@dataclass(frozen=True, slots=True) +class _ReachableDeployments: + model_ids: frozenset[str] + groups: frozenset[str] + + def include(self, deployment: object, team_id: str | None) -> bool: + model_id: Final = _deployment_id(deployment) + if model_id is None: + return not self.groups.isdisjoint(_names_caller_may_request(deployment, team_id)) + return model_id in self.model_ids + + +_NOTHING_REACHABLE: Final = _ReachableDeployments(frozenset(), frozenset()) + + +async def _deployments_key_may_reach( + llm_model_list: Sequence[object], + aliases: CallerAliases, + user_api_key_dict: UserAPIKeyAuth, + owner: LiteLLM_UserTable | None, llm_router: Router | None, +) -> _ReachableDeployments: + candidate_names: Final = (await _names_any_limit_grants(user_api_key_dict, owner)).union( + _alias_names_caller_may_request(aliases, llm_router), + *(_names_caller_may_request(deployment, user_api_key_dict.team_id) for deployment in llm_model_list), + ) + admitted_names: Final = await _names_key_may_request(candidate_names, aliases, user_api_key_dict, owner, llm_router) + access_group_grants: Final = _access_group_grants(llm_model_list, user_api_key_dict) + public_names: Final = _public_names_by_internal_name(llm_model_list, user_api_key_dict.team_id) + groups: Final = frozenset( + _served_group(access_group_grants.requested_name(name, public_names, admitted_names), aliases, llm_router) + for name in admitted_names + ) + ids_by_name: Final = _ids_by_requested_name(llm_model_list, user_api_key_dict.team_id) + reached_ids: Final = _ids_serving(groups, ids_by_name, user_api_key_dict, llm_router, access_group_grants) + admitted_ids: Final = await _names_key_may_request( + _shadowed_ids(llm_model_list, admitted_names, reached_ids, user_api_key_dict.team_id), + aliases, + user_api_key_dict, + owner, + llm_router, + ) + id_groups: Final = frozenset(_served_group(model_id, aliases, llm_router) for model_id in admitted_ids) + return _ReachableDeployments( + reached_ids | _ids_serving(id_groups, ids_by_name, user_api_key_dict, llm_router, access_group_grants), groups + ) + + +def _model_param_target( + model: str, + aliases: CallerAliases, + llm_model_list: Sequence[Mapping[str, object]], team_id: str | None, - caller_is_admin: bool, + llm_router: Router | None, +) -> str: + group: Final = _served_group(model, aliases, llm_router) + if llm_router is None or deployments_targeted_by_name(llm_model_list, group, team_id): + return group + patterns: Final = frozenset( + deployment.get("model_name") + for deployment in llm_router.registry().deployments(group, source="wildcard", team_id=team_id) + ) + match tuple(patterns): + case (str(pattern),): + return pattern + case _: + return group + + +def _caller_may_probe_deployment( + deployment: Mapping[str, object], reachable: _ReachableDeployments, team_id: str | None, caller_is_admin: bool ) -> bool: """Same deployment visibility rule as routing: another team's deployment is never in scope, team-less callers included.""" if not caller_is_admin and not deployment_usable_by_team(deployment, team_id): return False - if allowed_models is None: - return True - if llm_router is None: - return deployment.get("model_name") in allowed_models - model: Final = dict(deployment) - registry: Final = llm_router.registry() - return any(registry.includes(model, model_name=name, team_id=team_id) for name in allowed_models) + return reachable.include(deployment, team_id) def _resolve_targeted_model_ids( @@ -1137,6 +1476,7 @@ def _filter_health_check_results_by_model_ids(results: dict, allowed_model_ids: async def _perform_health_check_and_save( model_list, target_model, + recorded_model_name: str | None, cli_model, details, prisma_client, @@ -1159,8 +1499,7 @@ async def _perform_health_check_and_save( # Optionally save health check result to database (non-blocking) if prisma_client is not None: - # For CLI model, use cli_model name; for router models, use target_model - model_name_for_db: Final = cli_model if cli_model is not None else target_model + model_name_for_db: Final = cli_model if cli_model is not None else recorded_model_name if model_name_for_db is not None: asyncio.create_task( _save_health_check_to_db( @@ -1290,6 +1629,7 @@ async def health_endpoint( cli_result: Final = await _perform_health_check_and_save( model_list=[], target_model=None, + recorded_model_name=None, cli_model=user_model, details=health_check_details, prisma_client=prisma_client, @@ -1304,15 +1644,39 @@ async def health_endpoint( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": "Model list not initialized"}, ) - allowed_models: Final = _health_accessible_model_names(user_api_key_dict, llm_router) - restrict_to_allowed_models: Final = not is_admin or allowed_models is not None + restrict_to_allowed_models: Final = not is_admin or not _key_models_unrestricted(user_api_key_dict) + aliases: Final = await _caller_aliases(user_api_key_dict) if restrict_to_allowed_models else None + reachable: Final = ( + await _deployments_key_may_reach( + _OBJECTS.validate_python(llm_model_list), + aliases, + user_api_key_dict, + await _key_owner(user_api_key_dict), + llm_router, + ) + if aliases is not None + else _NOTHING_REACHABLE + ) + requested_model: Final = ( + _model_param_target( + model, + aliases, + _DEPLOYMENTS.validate_python(llm_model_list), + user_api_key_dict.team_id, + llm_router, + ) + if model is not None and aliases is not None + else model + ) _llm_model_list: Final = [ m for m in copy.deepcopy(llm_model_list) if not restrict_to_allowed_models - or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id, is_admin) + or _caller_may_probe_deployment(m, reachable, user_api_key_dict.team_id, is_admin) ] - targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id, user_api_key_dict.team_id) + targeted_ids: Final = _resolve_targeted_model_ids( + _llm_model_list, requested_model, model_id, user_api_key_dict.team_id + ) if restrict_to_allowed_models and targeted_ids is not None and not targeted_ids: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -1370,7 +1734,8 @@ async def health_endpoint( else: router_result: Final = await _perform_health_check_and_save( model_list=_llm_model_list, - target_model=target_model, + target_model=target_model if model_id else requested_model, + recorded_model_name=target_model, cli_model=None, details=health_check_details, prisma_client=prisma_client, diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 0bdc5f1910a..6b59f6d40dc 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -47,6 +47,7 @@ from litellm.proxy.auth.auth_checks import ( can_object_call_model, _can_object_call_vector_stores, _check_agent_access_group_model_access, + agent_model_grants, check_end_user_budget, _check_team_member_budget, personal_key_owner_teams, @@ -11294,6 +11295,15 @@ async def test_agent_access_groups_naming_no_model_deny_every_model(): assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied +@pytest.mark.asyncio +async def test_agent_model_grants_names_every_model_the_agents_access_groups_allow(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=[]) + resolve, asked = _agent_model_ceiling_resolver(frozenset({"gpt-5", "openai/gpt-5.6"})) + + assert await agent_model_grants(agent_key, resolve) == frozenset({"gpt-5", "openai/gpt-5.6"}) + assert asked == ["agent-1"] + + @pytest.mark.asyncio async def test_agent_without_access_groups_adds_no_model_ceiling(): agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 9847787efce..2258eb95fb8 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -1,8 +1,9 @@ import asyncio import copy +import itertools import json import time -from collections.abc import Awaitable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from datetime import datetime, timedelta from importlib import import_module @@ -22,7 +23,7 @@ import litellm import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_module from litellm.litellm_core_utils.health_check_helpers import TEST_IMAGE_BASE64 from litellm.models.credentials import CredentialItem -from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.health_endpoints._health_endpoints import ( _db_health_readiness_check, @@ -35,6 +36,7 @@ from litellm.proxy.health_endpoints._health_endpoints import ( test_model_connection as health_test_model_connection, ) from litellm.router import Router +from litellm.types.router import RouterModelGroupAliasItem from litellm.types.utils import oauth_token_exchange_litellm_params from litellm.types.workload_identity import ( ANTHROPIC_WIF_KWARGS_KEYS, @@ -925,7 +927,11 @@ async def test_test_model_connection_stored_audio_base_model_follows_the_stored_ deployment: Final = MappingProxyType( { "model_name": "custom-name-audio", - "litellm_params": {**deployment_params, "api_key": "fake-azure-key", "api_base": "https://x.openai.azure.com"}, + "litellm_params": { + **deployment_params, + "api_key": "fake-azure-key", + "api_base": "https://x.openai.azure.com", + }, "model_info": deployment_model_info, } ) @@ -1948,7 +1954,6 @@ def test_health_liveness_endpoint(proxy_client): assert duration_ms < 100, f"Health check took {duration_ms:.2f}ms, expected < 100ms for a simple endpoint" - def test_proxy_client_serves_full_proxy_routes_after_gateway_import(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy.proxy_server import app @@ -2472,6 +2477,7 @@ def _proxy_health_globals( llm_router: object, use_background_health_checks: bool = False, health_check_results: Mapping[str, object] | None = None, + prisma_client: object = None, ) -> Iterator[None]: with ( patch( # test-quality-ok: proxy module global, no injection seam @@ -2481,7 +2487,7 @@ def _proxy_health_globals( "litellm.proxy.proxy_server.llm_router", llm_router ), patch( # test-quality-ok: proxy module global, no injection seam - "litellm.proxy.proxy_server.prisma_client", None + "litellm.proxy.proxy_server.prisma_client", prisma_client ), patch( # test-quality-ok: proxy module global, no injection seam "litellm.proxy.proxy_server.use_background_health_checks", use_background_health_checks @@ -3812,7 +3818,11 @@ def test_clean_endpoint_data_never_displays_credential_fields(credential_field, async def _live_probed_model_ids( - model_list: Sequence[Mapping[str, object]], user_api_key_dict: UserAPIKeyAuth, model: str | None = None + model_list: Sequence[Mapping[str, object]], + user_api_key_dict: UserAPIKeyAuth, + model: str | None = None, + llm_router: Router | None = None, + prisma_client: object = None, ) -> set[str]: from fastapi import Response @@ -3825,7 +3835,7 @@ async def _live_probed_model_ids( return {"healthy_endpoints": [], "unhealthy_endpoints": [], "healthy_count": 0, "unhealthy_count": 0} with ( - _proxy_health_globals(model_list, _router_for(model_list)), + _proxy_health_globals(model_list, llm_router or _router_for(model_list), prisma_client=prisma_client), patch( # test-quality-ok: the model list handed to the probe is the assertion; no injection seam "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", side_effect=fake_perform, @@ -4120,6 +4130,7 @@ async def _live_narrowed_model_ids( user_api_key_dict: UserAPIKeyAuth, model: str | None = None, model_id: str | None = None, + llm_router: Router | None = None, ) -> set[str]: from fastapi import Response @@ -4130,7 +4141,7 @@ async def _live_narrowed_model_ids( return probed, [], {} with ( - _proxy_health_globals(model_list, _router_for(model_list)), + _proxy_health_globals(model_list, llm_router or _router_for(model_list)), patch( # test-quality-ok: the probe is the provider edge; which deployments reach it is the assertion "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe ), @@ -4289,6 +4300,699 @@ async def test_health_endpoint_keeps_a_team_only_public_name_off_a_team_less_key probe.assert_not_awaited() +_PROVIDER_PREFIXED_MODEL_LIST: Final = [ + { + "model_name": "openai/gpt-5.4-nano", + "litellm_params": {"model": "openai/gpt-5.4-nano"}, + "model_info": {"id": "id-prefixed-nano"}, + }, + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-mini"}, + }, + { + "model_name": "claude-haiku-5-5", + "litellm_params": {"model": "anthropic/claude-haiku-5-5"}, + "model_info": {"id": "id-haiku"}, + }, +] +_PROVIDER_PREFIXED_CACHED_RESULTS: Final = { + "healthy_endpoints": [ + {"model": "openai/gpt-5.4-nano", "model_id": "id-prefixed-nano"}, + {"model": "openai/gpt-5.4-mini", "model_id": "id-mini"}, + {"model": "anthropic/claude-haiku-5-5", "model_id": "id-haiku"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 3, + "unhealthy_count": 0, +} +_PROVIDER_WILDCARD_KEY: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["openai/*"]) +_RESTRICTED_TEAM_KEY: Final = UserAPIKeyAuth( + api_key="hashed-test-key", models=[], team_id="team-a", team_models=["gpt-5.4-mini"] +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller", "expected_ids"), + [(_PROVIDER_WILDCARD_KEY, {"id-prefixed-nano", "id-mini"}), (_RESTRICTED_TEAM_KEY, {"id-mini"})], +) +async def test_health_endpoint_probes_the_deployments_auth_lets_the_key_request( + caller: UserAPIKeyAuth, expected_ids: set[str] +): + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, caller) + + assert probed == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller", "expected_ids"), + [(_PROVIDER_WILDCARD_KEY, ["id-prefixed-nano", "id-mini"]), (_RESTRICTED_TEAM_KEY, ["id-mini"])], +) +async def test_health_endpoint_returns_the_deployments_auth_lets_the_key_request_on_background_cache_path( + caller: UserAPIKeyAuth, expected_ids: list[str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_for(_PROVIDER_PREFIXED_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ): + result: Final = await health_endpoint(response=Response(), user_api_key_dict=caller, model=None, model_id=None) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +async def test_health_endpoint_refuses_a_deployment_the_team_allowlist_forbids(use_background_health_checks: bool): + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_for(_PROVIDER_PREFIXED_MODEL_LIST), + use_background_health_checks=use_background_health_checks, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), user_api_key_dict=_RESTRICTED_TEAM_KEY, model=None, model_id="id-prefixed-nano" + ) + + assert refused.value.status_code == 403 + probe.assert_not_awaited() + + +_KEY_ALIASED_AWAY_FROM_HAIKU: Final = UserAPIKeyAuth( + api_key="hashed-test-key", models=["gpt-5.4-mini"], aliases={"claude-haiku-5-5": "gpt-5.4-mini"} +) +_TEAM_ALIASED_AWAY_FROM_HAIKU: Final = UserAPIKeyAuth( + api_key="hashed-test-key", + models=["claude-haiku-5-5"], + team_id="team-a", + team_model_aliases={"claude-haiku-5-5": "gpt-5.4-mini"}, +) +_MINI_ONLY_KEY: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-mini"]) +_ALIAS_SHADOW_CASES: Final = [ + pytest.param(_KEY_ALIASED_AWAY_FROM_HAIKU, {}, id="key-alias"), + pytest.param(_TEAM_ALIASED_AWAY_FROM_HAIKU, {}, id="team-alias"), + pytest.param(_MINI_ONLY_KEY, {"claude-haiku-5-5": "gpt-5.4-mini"}, id="router-alias"), +] + + +def _router_with_aliases(model_group_alias: Mapping[str, str]) -> Router: + return Router(model_list=copy.deepcopy(_PROVIDER_PREFIXED_MODEL_LIST), model_group_alias=dict(model_group_alias)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_probes_the_deployment_an_alias_sends_the_key_to_not_the_one_it_shadows( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + probed: Final = await _live_probed_model_ids( + _PROVIDER_PREFIXED_MODEL_LIST, caller, llm_router=_router_with_aliases(model_group_alias) + ) + + assert probed == {"id-mini"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_returns_the_deployment_an_alias_sends_the_key_to_on_background_cache_path( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_with_aliases(model_group_alias), + use_background_health_checks=True, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ): + result: Final = await health_endpoint(response=Response(), user_api_key_dict=caller, model=None, model_id=None) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-mini"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +async def test_health_endpoint_refuses_the_deployment_a_key_alias_shadows(use_background_health_checks: bool): + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_for(_PROVIDER_PREFIXED_MODEL_LIST), + use_background_health_checks=use_background_health_checks, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), user_api_key_dict=_KEY_ALIASED_AWAY_FROM_HAIKU, model=None, model_id="id-haiku" + ) + + assert refused.value.status_code == 403 + probe.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_targets_the_deployment_an_alias_name_sends_the_key_to( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + probed: Final = await _live_narrowed_model_ids( + _PROVIDER_PREFIXED_MODEL_LIST, + caller, + model="claude-haiku-5-5", + llm_router=_router_with_aliases(model_group_alias), + ) + + assert probed == {"id-mini"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_records_an_alias_check_under_the_name_the_caller_asked_for( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + async def fake_probe(model_list, details=True, max_concurrency=None, instrumentation_context=None): + return [{"model": m["litellm_params"]["model"], "model_id": m["model_info"]["id"]} for m in model_list], [], {} + + with ( + _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, _router_with_aliases(model_group_alias), prisma_client=MagicMock() + ), + patch( # test-quality-ok: the probe is the provider edge; the endpoint has no injection seam for it + "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe + ), + patch( # test-quality-ok: the history row is the assertion; the endpoint has no injection seam for its writer + "litellm.proxy.health_endpoints._health_endpoints._save_health_check_to_db", new_callable=AsyncMock + ) as save, + ): + result: Final = await health_endpoint( + response=Response(), user_api_key_dict=caller, model="claude-haiku-5-5", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-mini"] + assert save.call_args.args[1] == "claude-haiku-5-5" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_targets_the_deployment_an_alias_name_sends_the_key_to_on_background_cache_path( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_with_aliases(model_group_alias), + use_background_health_checks=True, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ): + result: Final = await health_endpoint( + response=Response(), user_api_key_dict=caller, model="claude-haiku-5-5", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-mini"] + + +@pytest.mark.asyncio +async def test_health_endpoint_never_probes_a_deployment_without_a_model_name_for_a_restricted_key(): + nameless_deployment: Final = { + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-nameless"}, + } + + probed: Final = await _live_probed_model_ids( + [*_PROVIDER_PREFIXED_MODEL_LIST, nameless_deployment], + _MINI_ONLY_KEY, + llm_router=_router_for(_PROVIDER_PREFIXED_MODEL_LIST), + ) + + assert probed == {"id-mini"} + + +@pytest.mark.asyncio +async def test_health_endpoint_never_probes_the_deployment_a_keys_router_alias_sends_it_to_but_forbids(): + caller: Final = UserAPIKeyAuth( + api_key="hashed-test-key", + models=["gpt-5.4-mini", "openai/gpt-5.4-nano"], + router_settings={"model_group_alias": {"gpt-5.4-mini": "claude-haiku-5-5"}}, + ) + + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, caller) + + assert probed == {"id-prefixed-nano"} + + +_WILDCARD_DEPLOYMENT: Final = { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-wild"}, +} +_CATCH_ALL_DEPLOYMENT: Final = {"model_name": "*", "litellm_params": {"model": "*"}, "model_info": {"id": "id-star"}} +_TEAM_WILDCARD_DEPLOYMENT: Final = { + "model_name": "openai-wild_team-a_7c1e", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-team-wild", "team_id": "team-a", "team_public_model_name": "openai/*"}, +} + + +def _team_key(team_models: list[str], team_model_aliases: Mapping[str, str] | None = None) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="hashed-test-key", + models=[], + team_id="team-a", + team_models=team_models, + team_model_aliases=dict(team_model_aliases) if team_model_aliases is not None else None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_list", "model_group_alias", "caller", "expected_ids"), + [ + pytest.param( + [_WILDCARD_DEPLOYMENT, _PROVIDER_PREFIXED_MODEL_LIST[2]], + {}, + _team_key(["openai/gpt-5.4-nano"]), + {"id-wild"}, + id="team-name-served-by-provider-wildcard", + ), + pytest.param( + [_TEAM_WILDCARD_DEPLOYMENT, _PROVIDER_PREFIXED_MODEL_LIST[2]], + {}, + _team_key(["openai/gpt-5.4-nano"]), + {"id-team-wild"}, + id="team-name-served-by-teams-own-wildcard", + ), + pytest.param( + [_CATCH_ALL_DEPLOYMENT], + {}, + _team_key(["gpt-5.4-mini"]), + {"id-star"}, + id="team-name-served-by-catch-all", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {"gpt-4": "gpt-5.4-mini"}, + _team_key(["gpt-4"]), + {"id-mini"}, + id="team-allows-router-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {}, + _team_key(["my-alias"], {"my-alias": "gpt-5.4-mini"}), + {"id-mini"}, + id="team-allows-team-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {"friendly-mini": RouterModelGroupAliasItem(model="gpt-5.4-mini", hidden=True)}, + UserAPIKeyAuth(api_key="hashed-test-key", models=["friendly-*"]), + {"id-mini"}, + id="wildcard-key-allows-hidden-router-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {}, + _team_key(["friendly-*"], {"friendly-mini": "gpt-5.4-mini"}), + {"id-mini"}, + id="wildcard-team-allows-team-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {"claude-haiku-5-5": "gpt-5.4-mini"}, + UserAPIKeyAuth(api_key="hashed-test-key", models=[]), + {"id-prefixed-nano", "id-mini", "id-haiku"}, + id="unrestricted-key-calls-shadowed-deployment-by-id", + ), + ], +) +async def test_health_endpoint_probes_every_deployment_the_callers_allowed_names_route_to( + model_list: Sequence[Mapping[str, object]], + model_group_alias: Mapping[str, str | RouterModelGroupAliasItem], + caller: UserAPIKeyAuth, + expected_ids: set[str], +): + llm_router: Final = Router( + model_list=copy.deepcopy([dict(deployment) for deployment in model_list]), + model_group_alias=dict(model_group_alias), + ) + + probed: Final = await _live_probed_model_ids(llm_router.get_model_list() or [], caller, llm_router=llm_router) + + assert probed == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "caller", + [ + pytest.param( + UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-mini", "gpt-4-retired"]), id="key-grant" + ), + pytest.param(_team_key(["gpt-5.4-mini", "gpt-4-retired"]), id="team-grant"), + ], +) +async def test_health_endpoint_never_probes_the_default_fallback_for_a_granted_name_no_deployment_serves( + caller: UserAPIKeyAuth, +) -> None: + llm_router: Final = Router( + model_list=copy.deepcopy(_PROVIDER_PREFIXED_MODEL_LIST), default_fallbacks=["claude-haiku-5-5"] + ) + + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, caller, llm_router=llm_router) + + assert probed == {"id-mini"} + + +def _personal_key_whose_owner_may_call(models: list[str]) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import user_api_key_cache + + owner_id: Final = "health-scope-owner-" + "-".join(models) + user_api_key_cache.set_cache(key=owner_id, value=LiteLLM_UserTable(user_id=owner_id, models=models)) + return UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_id=owner_id) + + +def _managed_agent_key_granted(models: list[str]) -> UserAPIKeyAuth: + from litellm.types.agents import AgentResponse + + policy: Final = AgentResponse( + agent_id="managed", agent_name="Managed", agent_card_params={}, object_permission={"models": models} + ) + return UserAPIKeyAuth(api_key="hashed-test-key", models=[]).model_copy(update={"managed_agent_policy": policy}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("make_caller", "models", "expected_ids"), + [ + pytest.param(_personal_key_whose_owner_may_call, ["gpt-5.4-mini"], {"id-mini"}, id="key-owners-models"), + pytest.param(_managed_agent_key_granted, ["gpt-5.4-mini"], {"id-mini"}, id="managed-agents-models"), + pytest.param(_managed_agent_key_granted, [], set(), id="managed-agent-without-model-grants"), + ], +) +async def test_health_endpoint_probes_only_the_models_the_keys_user_or_agent_may_call( + make_caller: Callable[[list[str]], UserAPIKeyAuth], models: list[str], expected_ids: set[str] +): + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, make_caller(models)) + + assert probed == expected_ids + + +@pytest.mark.asyncio +async def test_health_endpoint_refuses_a_model_id_the_keys_owner_may_not_call(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals(_PROVIDER_PREFIXED_MODEL_LIST, _router_for(_PROVIDER_PREFIXED_MODEL_LIST)), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=_personal_key_whose_owner_may_call(["gpt-5.4-mini"]), + model=None, + model_id="id-prefixed-nano", + ) + + assert refused.value.status_code == 403 + probe.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "make_caller", + [ + pytest.param(_personal_key_whose_owner_may_call, id="key-owner"), + pytest.param(_managed_agent_key_granted, id="managed-agent"), + ], +) +async def test_health_endpoint_keeps_the_wildcard_serving_a_name_the_keys_user_or_agent_may_call( + make_caller: Callable[[list[str]], UserAPIKeyAuth], +): + model_list: Final = [_WILDCARD_DEPLOYMENT, *_PROVIDER_PREFIXED_MODEL_LIST] + + probed: Final = await _live_probed_model_ids(model_list, make_caller(["openai/gpt-5.6"])) + + assert probed == {"id-wild"} + + +@pytest.mark.asyncio +async def test_health_endpoint_answers_an_agent_key_acting_for_an_unknown_user_the_way_chat_does(): + from litellm.types.agents import AgentCaller + + users_table: Final = MagicMock() + users_table.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + agent_key: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-mini"]).model_copy( + update={"agent_caller": AgentCaller(user_id="user-deleted-before-this-request")} + ) + + with pytest.raises(ProxyException) as failed: + await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, agent_key, prisma_client=users_table) + + assert (failed.value.code, failed.value.type) == ("401", ProxyErrorTypes.auth_error) + assert "user-deleted-before-this-request" in failed.value.message + + +@pytest.mark.asyncio +async def test_health_endpoint_lets_other_requests_run_while_it_scopes_a_large_model_list(): + model_list: Final = [ + {"model_name": f"model-{n}", "litellm_params": {"model": f"openai/model-{n}"}, "model_info": {"id": f"id-{n}"}} + for n in range(50) + ] + other_request_turns: Final = itertools.count() + scoping_done: Final = asyncio.Event() + + async def scope() -> set[str]: + try: + return await _live_probed_model_ids( + model_list, UserAPIKeyAuth(api_key="hashed-test-key", models=["model-0"]) + ) + finally: + scoping_done.set() + + async def other_request() -> None: + while not scoping_done.is_set(): + next(other_request_turns) + await asyncio.sleep(0) + + probed, _ = await asyncio.gather(scope(), other_request()) + + assert probed == {"id-0"} + assert next(other_request_turns) >= len(model_list) + + +_WILDCARD_CACHED_RESULTS: Final = { + "healthy_endpoints": [ + {"model": "openai/gpt-5.6", "model_id": "id-wild"}, + {"model": "openai/gpt-5.6", "model_id": "id-team-wild"}, + {"model": "anthropic/claude-haiku-5-5", "model_id": "id-haiku"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 3, + "unhealthy_count": 0, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +@pytest.mark.parametrize( + ("wildcard_deployment", "expected_id"), + [ + pytest.param(_WILDCARD_DEPLOYMENT, "id-wild", id="proxy-wildcard"), + pytest.param(_TEAM_WILDCARD_DEPLOYMENT, "id-team-wild", id="teams-own-wildcard"), + ], +) +async def test_health_endpoint_model_param_reaches_the_wildcard_serving_a_team_allowed_name( + wildcard_deployment: Mapping[str, object], expected_id: str, use_background_health_checks: bool +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + model_list: Final = [wildcard_deployment, _PROVIDER_PREFIXED_MODEL_LIST[2]] + + async def fake_probe(model_list, details=True, max_concurrency=None, instrumentation_context=None): + return [{"model": m["litellm_params"]["model"], "model_id": m["model_info"]["id"]} for m in model_list], [], {} + + with ( + _proxy_health_globals( + model_list, + _router_for(model_list), + use_background_health_checks=use_background_health_checks, + health_check_results=_WILDCARD_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe is the provider edge; which deployments reach it is the assertion + "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe + ), + ): + result: Final = await health_endpoint( + response=Response(), user_api_key_dict=_team_key(["openai/gpt-5.6"]), model="openai/gpt-5.6", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == [expected_id] + + +_ACCESS_GROUP_SPLIT_MODEL_LIST: Final = [ + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-mini-premium", "access_groups": ["premium"]}, + }, + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-mini-standard"}, + }, +] + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_access_groups_deployments_of_a_shared_name() -> None: + caller: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["premium"]) + + probed: Final = await _live_probed_model_ids(_ACCESS_GROUP_SPLIT_MODEL_LIST, caller) + + assert probed == {"id-mini-premium"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct_grant", ["gpt-5.4-mini", "*", "all-proxy-models"]) +async def test_health_endpoint_probes_every_deployment_of_a_name_the_key_also_holds_directly(direct_grant: str) -> None: + caller: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["premium", direct_grant]) + + probed: Final = await _live_probed_model_ids(_ACCESS_GROUP_SPLIT_MODEL_LIST, caller) + + assert probed == {"id-mini-premium", "id-mini-standard"} + + +_ACCESS_GROUP_SPLIT_WILDCARD_MODEL_LIST: Final = [ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-wild-premium", "access_groups": ["premium"]}, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-wild-standard", "access_groups": ["standard"]}, + }, +] + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_access_groups_wildcard_deployments_serving_a_name() -> None: + caller: Final = _personal_key_whose_owner_may_call(["openai/gpt-5.4-nano"]).model_copy( + update={"models": ["premium"]} + ) + + probed: Final = await _live_probed_model_ids(_ACCESS_GROUP_SPLIT_WILDCARD_MODEL_LIST, caller) + + assert probed == {"id-wild-premium"} + + +def _team_deployment( + internal_name: str, public_name: str, model: str, model_id: str, access_groups: list[str] +) -> Mapping[str, object]: + return { + "model_name": internal_name, + "litellm_params": {"model": model}, + "model_info": { + "id": model_id, + "team_id": "team-a", + "team_public_model_name": public_name, + "access_groups": access_groups, + }, + } + + +def _team_split_model_list(public_name: str, model: str) -> list[Mapping[str, object]]: + return [ + _team_deployment("model_name_team-a_7c1e4b2a", public_name, model, "id-team-premium", ["premium"]), + _team_deployment("model_name_team-a_9d2f6e01", public_name, model, "id-team-standard", ["standard"]), + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("public_name", "model"), + [ + pytest.param("openai/*", "openai/*", id="team-wildcard"), + pytest.param("gpt-5.4-mini", "openai/gpt-5.4-mini", id="team-model"), + ], +) +async def test_health_endpoint_probes_only_the_access_groups_team_deployments_serving_a_name( + public_name: str, model: str +) -> None: + probed: Final = await _live_probed_model_ids(_team_split_model_list(public_name, model), _team_key(["premium"])) + + assert probed == {"id-team-premium"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("internal_name_grant", ["model_name_team-a_9d2f6e01", "model_name_team-a_*"]) +async def test_health_endpoint_probes_a_team_deployment_whose_internal_name_the_key_holds( + internal_name_grant: str, +) -> None: + caller: Final = _team_key(["premium", internal_name_grant]) + + probed: Final = await _live_probed_model_ids(_team_split_model_list("gpt-5.4-mini", "openai/gpt-5.4-mini"), caller) + + assert probed == {"id-team-premium", "id-team-standard"} + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_team_deployment_whose_internal_name_an_agent_holds() -> None: + from litellm.types.agents import AgentResponse + + policy: Final = AgentResponse( + agent_id="managed", + agent_name="Managed", + agent_card_params={}, + object_permission={"models": ["model_name_team-a_9d2f6e01"]}, + ) + caller: Final = _team_key(["gpt-5.4-mini"]).model_copy(update={"managed_agent_policy": policy}) + + probed: Final = await _live_probed_model_ids(_team_split_model_list("gpt-5.4-mini", "openai/gpt-5.4-mini"), caller) + + assert probed == {"id-team-standard"} + + def test_resolve_targeted_model_ids_lets_model_id_win_over_model(): resolve = _health_endpoints_module._resolve_targeted_model_ids diff --git a/tests/unit/proxy/test_health_check_functions.py b/tests/unit/proxy/test_health_check_functions.py index 2c0198cd538..79eec44958f 100644 --- a/tests/unit/proxy/test_health_check_functions.py +++ b/tests/unit/proxy/test_health_check_functions.py @@ -741,6 +741,7 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c result = await _perform_health_check_and_save( model_list=model_list, target_model=None, + recorded_model_name=None, cli_model=None, details=True, prisma_client=None, @@ -777,6 +778,7 @@ async def test_perform_health_check_and_save_forwards_skip_disabled_background_f await _perform_health_check_and_save( model_list=model_list, target_model=None, + recorded_model_name=None, cli_model=None, details=True, prisma_client=None, From 43c07a0f1246c95da710408df4d7846bc5f848e2 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:48:32 -0700 Subject: [PATCH 28/64] fix(router)!: check and count deployment itpm and otpm on every Responses WebSocket turn (#49069) * fix(router): check and count deployment itpm and otpm on every Responses WebSocket turn The success hook only reconciled an ITPM/OTPM reservation found in top-level metadata or the standard logging object, but on the proxy the reservation lives in litellm_params.metadata, so no reservation was ever settled there. The hook now always calls the reconcile, which already returns early without a reservation and clears what it settled, so a second success event is a no-op. A Responses WebSocket turn after the first now has its deployment's current-minute ITPM and OTPM counters checked and is refused with a 429 frame once a counter is at its limit. An admitted turn carries a zero reservation on the current minute's keys, so its actual tokens land on those counters on both the native and the managed relay. * fix(responses): settle the socket open's itpm/otpm reservation on the first turn only After the first turn takes its logging copy, both websocket relays drop the open's ITPM/OTPM reservation keys from the connection metadata, the same way the budget reservation is detached. Later turns, unsolicited backend events, turns after limits are removed mid-socket, and sockets with no turn admission no longer settle or refund the open's reservation a second time. * refactor(responses): read the itpm/otpm reservation keys from core metadata helpers instead of router_utils * fix(responses): keep a socket open's hold for its first real turn and check tokens before request limits A turn that will settle the socket open's itpm/otpm reservation no longer counts that reservation against itself or swaps it for a zero one, so Codex's warmup frame on a managed provider, or a refused first frame on the native socket, no longer blocks the first real prompt or strands the hold for the minute. The token check now runs before the rpm/rpd checks, so a turn refused for tokens never uses up a request slot. Token-check tests pin the clock, and the managed relay settle test drives a real streamed call through a custom provider. * test(router): pin the router clock in the socket turn admission tests Both tests seeded this minute and the next to dodge a minute boundary. They now use the file's pinned-minute fixture and seed only that minute, so neither depends on when the test runs. * fix(router): check a responses socket's open without reserving and read shared token counters first A Responses WebSocket open now runs the same itpm/otpm check a turn runs and reserves nothing. It used to reserve an estimate that falls back to the model's max output, so a socket could not open under any otpm below that. The open still records a zero reservation for its own minute, so the first frame, which the relay forwards without a turn check, counts its actual tokens like every later turn. This drops the plumbing that only carried the open's hold into the first turn. The turn admission callable takes no arguments again, and the native relay detaches only the budget reservation, as before this change. The itpm/otpm turn check reads Redis first and falls back to this instance's in-memory count only without Redis or when the circuit breaker is open. These counters go down when a call settles, so an in-memory copy at the limit could refuse turns after the shared count dropped. The tpm check keeps its in-memory shortcut, since its counter only grows within a minute. Responses usage that reports input_tokens_details.cached_tokens now leaves cached input out of itpm, as chat usage already did. * fix(router): leave a socket turn's cached input out of itpm The native Responses WebSocket relay logs a turn with its list of events, so the deployment rate limit hook finds no usage on the response and falls back to the standard logging payload. That fallback read cached input from metadata.cache_read_input_tokens, a key the payload's cleaned metadata never carries, so every cached token of a socket turn was counted against itpm. The fallback now reads the cached count from the payload's usage_object, the same usage the logging path already built for cost, so a native socket turn counts only its uncached input. The managed relay and HTTP /v1/responses already counted correctly through the response's own usage * fix(router): mark only responses socket opens as check-only and count a turn on the minute it finishes The check-only branch for a socket open keyed on a websocket value in the request kwargs, which any HTTP body can carry and which the realtime route also sets. The router now passes an opens_responses_socket flag into set_io_token_rate_limit_request_kwargs only for its _aresponses_websocket call, and the flag lives in the same context value as the request kwargs, so no request body can set it and realtime opens reserve and get refused as before A zero reservation holds nothing to refund, so a socket turn now counts its real tokens on the minute the call finishes, read from the clock at settle time. A non-zero reservation still settles on the keys it reserved on so its refund hits the same counter When the shared Redis read returns nothing, the check now falls back to this instance's in-memory count instead of admitting the turn with no count at all * fix(router): keep io token reservation keys out of the metadata sent to providers The itpm/otpm pre-call check wrote its reservation keys into both metadata and litellm_metadata. On the routes the router serves through its generic helper (Responses, Anthropic Messages, and the rest), metadata is the client's own field and goes to the provider as is. So _litellm_itpm_cache_key and _litellm_otpm_cache_key, which carry the deployment id and the model name, reached the provider, came back in the echoed response metadata, and counted toward the provider's 16-key cap. The managed Responses socket relay merged each turn's reservation into the frame's metadata the same way, and a native socket opened with metadata sent the keys on every turn through its request defaults The router now passes the limiter the metadata field it uses for its own fields on that call (metadata for chat, litellm_metadata for the generic routes), and the reservation goes only there. metadata_channel is a required keyword, so a new caller has to name the field it writes. The managed relay puts each turn's reservation only on litellm_metadata With the reservation only on litellm_metadata, a retried call on the generic routes refunded its first attempt twice and pushed the deployment's counters below zero. The failure hook read the reservation from the logging copy in litellm_params.metadata and cleared only that copy, so the router's litellm_metadata still held it and the retry refunded it again. The hooks now also read and clear litellm_params.litellm_metadata, which is the router's own dict The proxy does not strip client keys that start with _litellm_ on the way in, so dropping that prefix at the provider boundary would also drop keys clients send today The io token tests now run Redis in process with fakeredis instead of a TCP fake server * test(responses): count socket turns through real logging callbacks instead of a Logging subclass The native relay test now hands the rate limit check to a real Logging object as its success and failure callback, so the turn's tokens go through the same dispatch path production uses. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../internal_call_metadata.py | 6 + litellm/llms/custom_httpx/llm_http_handler.py | 4 +- litellm/responses/main.py | 6 +- litellm/responses/streaming_iterator.py | 103 ++-- litellm/router.py | 26 +- .../io_token_rate_limit_check.py | 209 ++++++--- ...test_responses_websocket_turn_admission.py | 113 ++++- .../test_responses_websocket_all_providers.py | 441 +++++++++++++++++- .../test_router/test_io_token_rate_limits.py | 356 ++++++++++++-- tests/unit/test_router/test_router.py | 6 +- .../test_router/test_router_helper_utils.py | 374 ++++++++++++++- 11 files changed, 1492 insertions(+), 152 deletions(-) diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index f3dbd612bee..453d4358b35 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -35,6 +35,12 @@ from litellm.types.utils import ( BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) +ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" +OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" +ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" +OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key" +IO_TOKEN_RESERVATION_KEYS: Final = frozenset({ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY}) + MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups" """Where auth records the model access groups that authorized the request, for the spend writer. diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 61fc23589ac..9a2b352dd86 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -249,7 +249,7 @@ if TYPE_CHECKING: FakeAnthropicMessagesStreamIterator, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig - from litellm.responses.streaming_iterator import ResponsesWsFrameGuard + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard, ResponsesWsTurnAdmission from litellm.types.google_genai.main import GenerateContentResponse from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -7269,7 +7269,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, - turn_admission: Callable[[], Awaitable[None]] | None = None, + turn_admission: "ResponsesWsTurnAdmission | None" = None, frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs: object, ) -> Exception | None: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 5c8350589b8..9259103cb5c 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1,7 +1,7 @@ import asyncio import contextvars import json -from collections.abc import Awaitable, Callable, Coroutine, Generator, Iterable, Mapping, Sequence +from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass from functools import partial @@ -81,7 +81,7 @@ from litellm.types.llms.openai import ( # Handle ResponseText import with fallback if TYPE_CHECKING: - from litellm.responses.streaming_iterator import ResponsesWsFrameGuard + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard, ResponsesWsTurnAdmission from litellm.types.llms.openai import ResponseText else: ResponseText = str # Fallback for ResponseText import @@ -2867,7 +2867,7 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, - turn_admission: Callable[[], Awaitable[None]] | None = None, + turn_admission: "ResponsesWsTurnAdmission | None" = None, responses_ws_frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs, ) -> Exception | None: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index f0bb2201162..fb91ba6a618 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2145,18 +2145,32 @@ async def _enforce_frame_project_quota( ) -async def _admit_frame_turn(turn_admission: Callable[[], Awaitable[None]] | None, raw_message: str) -> Exception | None: +@dataclass(frozen=True, slots=True) +class _AdmittedTurn: + reservation: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) + + +@dataclass(frozen=True, slots=True) +class _RefusedTurn: + rejection: Exception + + +ResponsesWsTurnAdmission: TypeAlias = Callable[[], Awaitable[Mapping[str, object]]] + + +async def _admit_frame_turn( + turn_admission: ResponsesWsTurnAdmission | None, raw_message: str +) -> _AdmittedTurn | _RefusedTurn: frame: Final = _json_object_or_none(raw_message) if turn_admission is None or frame is None or frame.get("type") != "response.create": - return None + return _AdmittedTurn() try: - await turn_admission() + return _AdmittedTurn(await turn_admission()) except RateLimitError as refusal: - return refusal + return _RefusedTurn(refusal) except Exception as e: verbose_logger.exception("Responses WS turn admission failed: %s", e) - return e - return None + return _RefusedTurn(e) def _turn_refusal_frame(refusal: RateLimitError) -> str: @@ -2191,6 +2205,7 @@ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) _PER_TURN_METADATA_KEYS: Final = frozenset({"metadata", "litellm_metadata"}) +_NO_IO_TOKEN_RESERVATION: Final[Mapping[str, object]] = MappingProxyType({}) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" @@ -2318,6 +2333,17 @@ def _set_turn_guardrail_information( _replace_guardrail_information(litellm_params.get(key), guardrail_information) +def _set_turn_io_token_reservation(turn_logging_obj: LiteLLMLoggingObj, reservation: Mapping[str, object]) -> None: + litellm_params: Final = _json_object_or_none( + turn_logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped + ) + if litellm_params is None or not reservation: + return + for bucket in (litellm_params.get(key) for key in _PER_TURN_METADATA_KEYS): + if _is_json_object(bucket): + bucket.update(reservation) + + def _replace_guardrail_information(bucket: object, guardrail_information: tuple[Mapping[str, object], ...]) -> None: if not _is_json_object(bucket): return @@ -2458,7 +2484,7 @@ class ResponsesWebSocketStreaming: authorized_model: str | None = None, custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, - turn_admission: Callable[[], Awaitable[None]] | None = None, + turn_admission: ResponsesWsTurnAdmission | None = None, frame_guard: ResponsesWsFrameGuard | None = None, ): self.websocket = websocket @@ -2480,10 +2506,13 @@ class ResponsesWebSocketStreaming: # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults - self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission + self.turn_admission: ResponsesWsTurnAdmission | None = turn_admission def _start_turn( - self, client_frame: str | None, guardrail_information: tuple[Mapping[str, object], ...] = () + self, + client_frame: str | None, + guardrail_information: tuple[Mapping[str, object], ...] = (), + io_token_reservation: Mapping[str, object] = _NO_IO_TOKEN_RESERVATION, ) -> None: start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time turn_logging_obj: Final = self.logging_obj.copy_for_new_call( @@ -2491,6 +2520,7 @@ class ResponsesWebSocketStreaming: ) self._detach_budget_reservation() _set_turn_guardrail_information(turn_logging_obj, guardrail_information) + _set_turn_io_token_reservation(turn_logging_obj, io_token_reservation) turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame)) turn_logging_obj.pre_call(input=client_frame, api_key="") self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj)) @@ -2510,11 +2540,14 @@ class ResponsesWebSocketStreaming: } def _record_client_frame( - self, client_frame: str, guardrail_information: tuple[Mapping[str, object], ...] = () + self, + client_frame: str, + guardrail_information: tuple[Mapping[str, object], ...] = (), + io_token_reservation: Mapping[str, object] = _NO_IO_TOKEN_RESERVATION, ) -> None: frame: Final = _json_object_or_none(client_frame) if frame is not None and frame.get("type") == "response.create": - self._start_turn(client_frame, guardrail_information) + self._start_turn(client_frame, guardrail_information, io_token_reservation) async def _guard_frame(self, message: str) -> GuardedResponsesWsFrame | None: """Run the proxy's pre_call guardrails on one ``response.create`` frame. @@ -2533,15 +2566,18 @@ class ResponsesWebSocketStreaming: return None async def _forward_guarded_frame( - self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + self, message: str, turn_admission: ResponsesWsTurnAdmission | None = None ) -> None: guarded: Final = await self._guard_frame(message) - if guarded is None or not await self._admit_or_reject_turn(guarded, message, turn_admission): + if guarded is None: + return + admitted: Final = await self._admit_or_reject_turn(guarded, message, turn_admission) + if admitted is None: return if guarded.pii_tokens: self._remember_pii_tokens(guarded.pii_tokens) masked: Final = await self._mask_response_create(guarded.frame) - self._record_client_frame(masked, guarded.guardrail_information) + self._record_client_frame(masked, guarded.guardrail_information, admitted.reservation) await self.backend_ws.send(masked) def _remember_pii_tokens(self, pii_tokens: Mapping[str, str]) -> None: @@ -2941,15 +2977,17 @@ class ResponsesWebSocketStreaming: return True async def _admit_or_reject_turn( - self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: Callable[[], Awaitable[None]] | None - ) -> bool: - rejection: Final = await _admit_frame_turn(turn_admission, message) - if rejection is None: - return True - if self._failure is None: - self._failure = rejection - await _reject_guarded_turn(self.websocket, guarded, rejection) - return False + self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: ResponsesWsTurnAdmission | None + ) -> _AdmittedTurn | None: + verdict: Final = await _admit_frame_turn(turn_admission, message) + match verdict: + case _AdmittedTurn(): + return verdict + case _RefusedTurn(rejection=rejection): + if self._failure is None: + self._failure = rejection + await _reject_guarded_turn(self.websocket, guarded, rejection) + return None async def _send_turn_refusal(self, refusal: RateLimitError) -> None: try: @@ -3052,7 +3090,7 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, - turn_admission: Callable[[], Awaitable[None]] | None = None, + turn_admission: ResponsesWsTurnAdmission | None = None, frame_guard: ResponsesWsFrameGuard | None = None, guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, **kwargs: object, @@ -3076,7 +3114,7 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () - self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission + self.turn_admission: ResponsesWsTurnAdmission | None = turn_admission self._failure: Exception | None = None # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} @@ -3473,7 +3511,7 @@ class ManagedResponsesWebSocketHandler: # ------------------------------------------------------------------ async def _process_response_create( - self, raw_message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + self, raw_message: str, turn_admission: ResponsesWsTurnAdmission | None = None ) -> None: """ Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``, @@ -3520,11 +3558,11 @@ class ManagedResponsesWebSocketHandler: guarded: Final = await self._guard_frame(raw_message) if guarded is None: return - rejection: Final = await _admit_frame_turn(turn_admission, raw_message) - if rejection is not None: + verdict: Final = await _admit_frame_turn(turn_admission, raw_message) + if isinstance(verdict, _RefusedTurn): if self._failure is None: - self._failure = rejection - await _reject_guarded_turn(self.websocket, guarded, rejection) + self._failure = verdict.rejection + await _reject_guarded_turn(self.websocket, guarded, verdict.rejection) return if guarded.pii_tokens: self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens} @@ -3555,6 +3593,11 @@ class ManagedResponsesWebSocketHandler: ) self._update_proxy_request(call_kwargs, requested_model or self.model_group or model) call_kwargs.update(self.extra_kwargs) + turn_litellm_metadata: Final = call_kwargs.get("litellm_metadata") + call_kwargs["litellm_metadata"] = { + **(turn_litellm_metadata if _is_json_object(turn_litellm_metadata) else {}), + **verdict.reservation, + } try: terminal_event: Final = await self._stream_and_forward(model, call_kwargs) diff --git a/litellm/router.py b/litellm/router.py index af2181cedb5..233e7581a40 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -378,6 +378,7 @@ from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check impor EncryptedContentAffinityCheck, ) from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + async_admit_io_token_turn, async_io_token_refund_failure, build_io_token_rate_limit_headers, deployment_has_io_token_limits, @@ -5113,6 +5114,7 @@ class Router(DeprecatedRouterMembers): kwargs: dict, function_name: str | None = None, splits_stream_timeout: bool = False, + opens_responses_socket: bool = False, ) -> None: """ 5 jobs: @@ -5202,7 +5204,12 @@ class Router(DeprecatedRouterMembers): # here before it's wiped below, instead of relying on that attempt's # (possibly still-pending) failure event to do it. refund_stale_reservation_before_retry(self._core_cache, kwargs) - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=deployment_has_io_token_limits(deployment)) + set_io_token_rate_limit_request_kwargs( + kwargs, + store_in_context=deployment_has_io_token_limits(deployment), + opens_responses_socket=opens_responses_socket, + metadata_channel=metadata_variable_name, + ) kwargs[metadata_variable_name].setdefault( ROUTING_REQUEST_TAGS_METADATA_KEY, @@ -6615,6 +6622,7 @@ class Router(DeprecatedRouterMembers): kwargs=kwargs, function_name=function_name, splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + opens_responses_socket=handler_name == "_aresponses_websocket", ) data: Final = deployment["litellm_params"].copy() @@ -10465,9 +10473,20 @@ class Router(DeprecatedRouterMembers): return deployment if current is None else current.model_dump(exclude_none=True) async def async_admit_deployment_turn( - self, model_group: str, opened_on: Mapping[str, object], request_kwargs: Mapping[str, object] - ) -> None: + self, + model_group: str, + opened_on: Mapping[str, object], + request_kwargs: Mapping[str, object], + ) -> Mapping[str, object]: deployment: Final = self._deployment_as_configured_now(opened_on) + rate_limit_check: Final = next( + (callback for callback in self._routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)), None + ) + io_token_admission: Final[Mapping[str, object]] = ( + MappingProxyType({}) + if rate_limit_check is None + else await async_admit_io_token_turn(rate_limit_check.dual_cache, deployment, get_utc_datetime()) + ) budget_limiter: Final = self._core_get_router_deployment_budget_limiter() if budget_limiter is not None: try: @@ -10488,6 +10507,7 @@ class Router(DeprecatedRouterMembers): await self.async_routing_strategy_pre_call_checks( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict deployment=dict(without_io_token_limits(deployment)), parent_otel_span=None ) + return io_token_admission async def _admit_turn_through_routing_strategy( self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object] diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 7580c7dadd3..80e1f3fe29c 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -14,6 +14,9 @@ from __future__ import annotations import contextlib import contextvars from collections.abc import Mapping, MutableMapping +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx @@ -23,7 +26,15 @@ from litellm import token_counter from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION +from litellm.litellm_core_utils.internal_call_metadata import ( + IO_TOKEN_RESERVATION_KEYS, + ITPM_CACHE_KEY, + ITPM_RESERVED_KEY, + OTPM_CACHE_KEY, + OTPM_RESERVED_KEY, +) from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET @@ -39,18 +50,27 @@ else: RoutingArgsTTL: Final = 60 -_io_token_rate_limit_request_kwargs: Final[contextvars.ContextVar[dict[str, object] | None]] = contextvars.ContextVar( - "io_token_rate_limit_request_kwargs", + +@dataclass(frozen=True, slots=True) +class _IoTokenRateLimitRequest: + kwargs: dict[str, object] | None + opens_responses_socket: bool + metadata_channel: str + + +_io_token_rate_limit_request: Final[contextvars.ContextVar[_IoTokenRateLimitRequest | None]] = contextvars.ContextVar( + "io_token_rate_limit_request", default=None, ) -ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" -OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" -ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" -OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key" - -def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, object] | None, store_in_context: bool = True) -> None: +def set_io_token_rate_limit_request_kwargs( + kwargs: dict[str, object] | None, + store_in_context: bool = True, + *, + opens_responses_socket: bool = False, + metadata_channel: str, +) -> None: # The reservation sentinels are server-only, but `metadata` is caller # controlled on proxy requests. Strip any client-supplied copies here (this # runs before the router stashes its own reservation) so a forged @@ -64,11 +84,14 @@ def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, object] | None, sto # every other deployment overwrite the slot with None instead of the # kwargs; overwriting (rather than skipping) also releases a previous # request's kwargs when a context is reused. - _io_token_rate_limit_request_kwargs.set(kwargs if store_in_context else None) + _io_token_rate_limit_request.set( + _IoTokenRateLimitRequest(kwargs, opens_responses_socket, metadata_channel) if store_in_context else None + ) def get_io_token_rate_limit_request_kwargs() -> dict[str, object] | None: - return _io_token_rate_limit_request_kwargs.get() + request: Final = _io_token_rate_limit_request.get() + return None if request is None else request.kwargs def seconds_until_minute_reset() -> int: @@ -224,27 +247,32 @@ def _otpm_reservation( return default if from_deployment is None else min(from_deployment, default) +def _usage_field(container: object, name: str) -> object: + if is_str_keyed_mapping(container): + return container.get(name) + return getattr(container, name, None) + + +def _cached_prompt_tokens(usage: object) -> int: + reported: Final = ( + _usage_field(_usage_field(usage, "prompt_tokens_details"), "cached_tokens"), + _usage_field(_usage_field(usage, "input_tokens_details"), "cached_tokens"), + _usage_field(usage, "cache_read_input_tokens"), + ) + return next((int(tokens) for tokens in reported if isinstance(tokens, int | float) and tokens), 0) + + def _get_usage_tokens(usage: object) -> tuple[int, int, int]: if usage is None: return 0, 0, 0 if hasattr(usage, "prompt_tokens") or hasattr(usage, "input_tokens"): prompt = int(getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", 0) or 0) completion = int(getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", 0) or 0) - cached = 0 - details: object = getattr(usage, "prompt_tokens_details", None) - if details is not None: - cached = int(getattr(details, "cached_tokens", 0) or 0) - if not cached: - cached = int(getattr(usage, "cache_read_input_tokens", 0) or 0) - return prompt, completion, cached + return prompt, completion, _cached_prompt_tokens(usage) if isinstance(usage, dict): prompt = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0) completion = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0) - details = usage.get("prompt_tokens_details") or {} - cached = int(details.get("cached_tokens", 0) or 0) if isinstance(details, dict) else 0 - if not cached: - cached = int(usage.get("cache_read_input_tokens") or 0) - return prompt, completion, cached + return prompt, completion, _cached_prompt_tokens(usage) return 0, 0, 0 @@ -298,19 +326,36 @@ def _resolve_reconcile_usage_tokens( if isinstance(standard_logging_object, dict): prompt_tokens = int(standard_logging_object.get("prompt_tokens") or 0) completion_tokens = int(standard_logging_object.get("completion_tokens") or 0) - cached_tokens = 0 - metadata: Final = standard_logging_object.get("metadata") - if isinstance(metadata, dict): - cached_tokens = int(metadata.get("cache_read_input_tokens") or 0) + logged_usage: Final = _usage_field(standard_logging_object.get("metadata"), "usage_object") + logged_cached_tokens: Final = _cached_prompt_tokens(logged_usage) # Same rationale as _usage_is_present: a bare total_tokens with no # prompt/completion breakdown can't be split, so it isn't treated # as resolved usage - the reservation is kept instead of refunded. if prompt_tokens or completion_tokens: - return max(0, prompt_tokens - cached_tokens), completion_tokens, True + return max(0, prompt_tokens - logged_cached_tokens), completion_tokens, True return 0, 0, False +def _reservation( + itpm_reserved: int, otpm_reserved: int, itpm_cache_key: str | None, otpm_cache_key: str | None +) -> Mapping[str, object]: + return MappingProxyType( + { + ITPM_RESERVED_KEY: itpm_reserved, + OTPM_RESERVED_KEY: otpm_reserved, + ITPM_CACHE_KEY: itpm_cache_key, + OTPM_CACHE_KEY: otpm_cache_key, + } + ) + + +def _stash_reservation_in_metadata(request: _IoTokenRateLimitRequest | None, reservation: Mapping[str, object]) -> None: + metadata: Final = None if request is None or not request.kwargs else request.kwargs.get(request.metadata_channel) + if isinstance(metadata, dict): + metadata.update(reservation) + + def _key_for_this_minute(reserved_key: str) -> str: counter, _, _ = reserved_key.rpartition(":") return f"{counter}:{get_utc_datetime().strftime('%H-%M')}" @@ -326,30 +371,6 @@ def _reconcile_increments( ) -def _stash_reservation_in_metadata( - request_kwargs: dict[str, object] | None, - *, - itpm_reserved: int, - otpm_reserved: int, - itpm_cache_key: str | None, - otpm_cache_key: str | None, -) -> None: - if not request_kwargs: - return - reservation: Final = { - ITPM_RESERVED_KEY: itpm_reserved, - OTPM_RESERVED_KEY: otpm_reserved, - ITPM_CACHE_KEY: itpm_cache_key, - OTPM_CACHE_KEY: otpm_cache_key, - } - for channel in ("metadata", "litellm_metadata"): - existing = request_kwargs.get(channel) - if isinstance(existing, dict): - existing.update(reservation) - elif channel == "metadata": - request_kwargs[channel] = dict(reservation) - - def _extract_reservation(reservation: Mapping[str, int | str | None]) -> tuple[int, int, str | None, str | None]: itpm_cache_key: Final = reservation.get(ITPM_CACHE_KEY) otpm_cache_key: Final = reservation.get(OTPM_CACHE_KEY) @@ -399,7 +420,7 @@ def _clear_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> None: for channel in _reservation_channels(kwargs): channel_dict = _as_mutable_mapping(channel) if channel_dict is not None: - for key in (ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY): + for key in IO_TOKEN_RESERVATION_KEYS: channel_dict.pop(key, None) @@ -499,7 +520,8 @@ def io_token_pre_call_check( if itpm_limit is None and otpm_limit is None: return deployment - request_kwargs: Final = get_io_token_rate_limit_request_kwargs() + request: Final = _io_token_rate_limit_request.get() + request_kwargs: Final = None if request is None else request.kwargs _model, served_model = _tokenizer_and_served_models(deployment) estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model) otpm_to_reserve: Final = ( @@ -546,11 +568,13 @@ def io_token_pre_call_check( raise _stash_reservation_in_metadata( - request_kwargs, - itpm_reserved=itpm_reserved, - otpm_reserved=otpm_reserved, - itpm_cache_key=itpm_key if itpm_limit is not None else None, - otpm_cache_key=otpm_key if otpm_limit is not None else None, + request, + _reservation( + itpm_reserved, + otpm_reserved, + itpm_key if itpm_limit is not None else None, + otpm_key if otpm_limit is not None else None, + ), ) return deployment @@ -565,7 +589,14 @@ async def async_io_token_pre_call_check( if itpm_limit is None and otpm_limit is None: return deployment - request_kwargs: Final = get_io_token_rate_limit_request_kwargs() + request: Final = _io_token_rate_limit_request.get() + request_kwargs: Final = None if request is None else request.kwargs + if request is not None and request.opens_responses_socket: + _stash_reservation_in_metadata( + request, + await async_admit_io_token_turn(dual_cache, deployment, get_utc_datetime(), parent_otel_span), + ) + return deployment _model, served_model = _tokenizer_and_served_models(deployment) estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model) otpm_to_reserve: Final = ( @@ -618,15 +649,65 @@ async def async_io_token_pre_call_check( raise _stash_reservation_in_metadata( - request_kwargs, - itpm_reserved=itpm_reserved, - otpm_reserved=otpm_reserved, - itpm_cache_key=itpm_key if itpm_limit is not None else None, - otpm_cache_key=otpm_key if otpm_limit is not None else None, + request, + _reservation( + itpm_reserved, + otpm_reserved, + itpm_key if itpm_limit is not None else None, + otpm_key if otpm_limit is not None else None, + ), ) return deployment +async def _async_redis_usage(redis_cache: RedisCache | None, key: str, parent_otel_span: Span | None) -> int | None: + if redis_cache is None: + return None + with contextlib.suppress(RedisCircuitBreakerOpenError): + return await redis_cache.async_get_cache(key=key, parent_otel_span=parent_otel_span) + return None + + +@with_service_target(ROUTER_USAGE_TARGET) +async def _async_shared_usage(dual_cache: DualCache, key: str, parent_otel_span: Span | None) -> int | None: + shared: Final = await _async_redis_usage(dual_cache.redis_cache, key, parent_otel_span) + if shared is not None: + return shared + return await dual_cache.async_get_cache(key=key, local_only=True) + + +async def _raise_if_at_limit( + dual_cache: DualCache, key: str, limit: int | None, limit_label: str, parent_otel_span: Span | None +) -> None: + if limit is None: + return + usage: Final = await _async_shared_usage(dual_cache, key, parent_otel_span) + if usage is not None and usage >= limit: + raise _rate_limit_error(limit_label, limit, usage) + + +@with_service_target(ROUTER_USAGE_TARGET) +async def async_admit_io_token_turn( + dual_cache: DualCache, + deployment: Mapping[str, object], + now: datetime, + parent_otel_span: Span | None = None, +) -> Mapping[str, object]: + configured: Final = dict(deployment) + itpm_limit, otpm_limit = get_deployment_io_token_limits(configured) + if itpm_limit is None and otpm_limit is None: + return MappingProxyType({}) + cache_keys: Final = _get_cache_keys(configured, now.strftime("%H-%M")) + if cache_keys is None: + return MappingProxyType({}) + itpm_key, otpm_key = cache_keys + await _raise_if_at_limit(dual_cache, itpm_key, itpm_limit, "ITPM", parent_otel_span) + await _raise_if_at_limit(dual_cache, otpm_key, otpm_limit, "OTPM", parent_otel_span) + return _reservation( + 0, 0, itpm_key if itpm_limit is not None else None, otpm_key if otpm_limit is not None else None + ) + + @with_service_target(ROUTER_USAGE_TARGET) def io_token_reconcile_success( dual_cache: DualCache, diff --git a/tests/integration/routing/test_responses_websocket_turn_admission.py b/tests/integration/routing/test_responses_websocket_turn_admission.py index 87b56a92979..7b0e14d044f 100644 --- a/tests/integration/routing/test_responses_websocket_turn_admission.py +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -518,9 +518,7 @@ def test_a_tpm_limit_lowered_while_the_socket_is_open_applies_to_its_next_turn( assert deployment.forwarded_turns() == 1 -def test_a_socket_on_an_otpm_limited_deployment_reserves_once_and_serves_every_turn( - limited: Gateway, cert: tuple[Path, Path] -) -> None: +def test_a_socket_on_an_otpm_limited_deployment_serves_every_turn(limited: Gateway, cert: tuple[Path, Path]) -> None: with responses_peer(cert) as peer, limited.scenario() as scenario: deployment: Final = _native(scenario, peer, otpm=10_000) key: Final = scenario.key(models=[deployment.name]) @@ -536,6 +534,115 @@ def test_a_socket_on_an_otpm_limited_deployment_reserves_once_and_serves_every_t assert deployment.forwarded_turns() == 5 +IO_TOKEN_LIMIT: Final = 10_000 +TOKENS_PER_TURN: Final = MappingProxyType({"itpm": USAGE["input_tokens"], "otpm": USAGE["output_tokens"]}) + + +def _io_tokens_counted(identity: str, kind: str, minute: str) -> int: + with _redis() as redis: + return sum(int(redis.get(key) or 0) for key in redis.scan_iter(f"global_router:{identity}:*:{kind}:{minute}")) + + +def _add_io_tokens(deployment: Deployment, kind: str, minute: str, tokens: int) -> None: + with _redis() as redis: + redis.incrby(f"global_router:{deployment.identity}:openai/{PROVIDER_MODEL}:{kind}:{minute}", tokens) + + +async def _say_until_counted( + connection: websockets.ClientConnection, deployment: Deployment, counted: Callable[[], float], expected: int +) -> tuple[dict[str, JsonValue], ...]: + received: Final = await _say(connection, deployment.name) + await _until(counted, lambda tokens: tokens == expected) + return received + + +@pytest.mark.parametrize("kind", ("otpm", "itpm")) +def test_a_turn_is_refused_once_earlier_turns_fill_the_deployment_io_token_limit( + limited: Gateway, cert: tuple[Path, Path], kind: str +) -> None: + per_turn: Final = TOKENS_PER_TURN[kind] + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, **{kind: IO_TOKEN_LIMIT}) + key: Final = scenario.key(models=[deployment.name]) + minute: Final = _fresh_minute() + counted: Final = partial(_io_tokens_counted, deployment.identity, kind, minute) + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + first: Final = await _say_until_counted(connection, deployment, counted, per_turn) + await asyncio.to_thread(_add_io_tokens, deployment, kind, minute, IO_TOKEN_LIMIT - 5 * per_turn) + answered: Final = tuple( + [ + await _say_until_counted( + connection, deployment, counted, IO_TOKEN_LIMIT - (5 - number) * per_turn + ) + for number in range(2, 6) + ] + ) + return (first, *answered, await _say(connection, deployment.name)) + + turns: Final = asyncio.run(session()) + + assert [completed(received)["status"] for received in turns[:5]] == ["completed"] * 5 + _assert_refused(turns[5]) + assert deployment.forwarded_turns() == 5 + assert held(counted, lambda tokens: tokens == IO_TOKEN_LIMIT, holding=3, seconds=20) == IO_TOKEN_LIMIT + + +def test_a_socket_at_half_its_deployment_otpm_serves_ten_turns(limited: Gateway, cert: tuple[Path, Path]) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, otpm=IO_TOKEN_LIMIT) + key: Final = scenario.key(models=[deployment.name]) + minute: Final = _fresh_minute() + _add_io_tokens(deployment, "otpm", minute, IO_TOKEN_LIMIT // 2) + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + return tuple([await _say(connection, deployment.name) for _ in range(10)]) + + turns: Final = asyncio.run(session()) + + assert [completed(received)["status"] for received in turns] == ["completed"] * 10 + assert deployment.forwarded_turns() == 10 + expected: Final = IO_TOKEN_LIMIT // 2 + 10 * USAGE["output_tokens"] + counted: Final = partial(_io_tokens_counted, deployment.identity, "otpm", minute) + assert held(counted, lambda tokens: tokens == expected, holding=3, seconds=20) == expected + + +@pytest.mark.parametrize("relay", ("native", "managed")) +def test_every_turn_counts_its_actual_tokens_toward_the_deployment_itpm_and_otpm( + limited: Gateway, cert: tuple[Path, Path], relay: str +) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + limits: Final = {"itpm": 1_000_000, "otpm": 1_000_000} + deployment: Final = _native(scenario, peer, **limits) if relay == "native" else _managed(scenario, **limits) + key: Final = scenario.key(models=[deployment.name]) + minute: Final = _fresh_minute() + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + return tuple([await _say(connection, deployment.name) for _ in range(3)]) + + turns: Final = asyncio.run(session()) + + usages: Final = tuple(object_value(completed(received)["usage"]) for received in turns) + expected: Final = { + "itpm": sum(int(str(usage["input_tokens"])) for usage in usages), + "otpm": sum(int(str(usage["output_tokens"])) for usage in usages), + } + assert any(tokens > 0 for tokens in expected.values()), usages + counted: Final = { + kind: held( + partial(_io_tokens_counted, deployment.identity, kind, minute), + lambda tokens, kind=kind: tokens == expected[kind], + holding=3, + seconds=20, + ) + for kind in expected + } + assert counted == expected + + def test_usage_based_routing_refuses_a_turn_only_once_earlier_turns_used_up_the_deployment_tpm( usage_routed: Gateway, cert: tuple[Path, Path] ) -> None: diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index ab52e6d29c3..d0e26798090 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -8,15 +8,21 @@ Tests that: """ import asyncio +import itertools import json -from collections.abc import Mapping -from typing import Final -from unittest.mock import MagicMock +from collections.abc import AsyncIterator, Mapping +from datetime import datetime, timezone +from functools import partial +from typing import TYPE_CHECKING, Final +from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +import respx from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig +from litellm.llms.custom_llm import CustomLLM from litellm.llms.databricks.responses.transformation import ( DatabricksResponsesAPIConfig, ) @@ -42,6 +48,13 @@ from litellm.llms.volcengine.responses.transformation import ( VolcEngineResponsesAPIConfig, ) from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.types.utils import GenericStreamingChunk + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ResponsesWsTurnAdmission + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck class TestResponsesAPIWebSocketSupport: @@ -1591,7 +1604,7 @@ def _admission_refusing_the_next_turn(): from litellm.exceptions import RateLimitError - return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), None]) + return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), {}]) class _ClientSocket: @@ -1630,7 +1643,7 @@ class _StalledClientSocket(_ClientSocket): def _admission_crashing_the_next_turn(): from unittest.mock import AsyncMock - return AsyncMock(side_effect=[ValueError("custom check failed"), None]) + return AsyncMock(side_effect=[ValueError("custom check failed"), {}]) def _guard_blocking(word: str): @@ -1655,6 +1668,229 @@ def _guard_recording(bill_refusal: MagicMock): return guard +_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-34" +_TURNED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-35" +_OPENED_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc) +_TURNED_AT: Final = datetime(2026, 1, 1, 12, 35, 5, tzinfo=timezone.utc) +_WS_DEPLOYMENT: Final = { + "litellm_params": {"model": "openai/gpt-5-mini", "otpm": 4096}, + "model_info": {"id": "ws-dep"}, +} +_OTHER_USAGE: Final = 10 +_OUTPUT_TOKENS_PER_TURN: Final = 2 +_UPSTREAM_ERROR: Final = json.dumps({"type": "error", "error": {"type": "server_error", "message": "upstream failed"}}) + + +def _turn_completed() -> str: + return json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_turn", + "status": "completed", + "output": [], + "usage": {"input_tokens": 5, "output_tokens": _OUTPUT_TOKENS_PER_TURN, "total_tokens": 7}, + }, + } + ) + + +async def _socket_opened_on(dual_cache: "DualCache") -> Mapping[str, object]: + from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn + + await dual_cache.async_increment_cache(key=_OPENED_ON, value=_OTHER_USAGE, ttl=60) + await dual_cache.async_increment_cache(key=_TURNED_ON, value=_OTHER_USAGE, ttl=60) + return await async_admit_io_token_turn(dual_cache, _WS_DEPLOYMENT, _OPENED_AT) + + +def _turns_finish_at(monkeypatch: pytest.MonkeyPatch, finished_at: datetime) -> None: + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: finished_at + ) + + +def _later_turns_admitted_through(dual_cache: "DualCache") -> "ResponsesWsTurnAdmission": + from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn + + return partial(async_admit_io_token_turn, dual_cache, _WS_DEPLOYMENT, _TURNED_AT) + + +_ITPM_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:itpm:12-34" +_ITPM_DEPLOYMENT: Final = { + "litellm_params": {"model": "openai/gpt-5-mini", "itpm": 40_000}, + "model_info": {"id": "ws-dep"}, +} +_TURN_INPUT_TOKENS: Final = 1_000 + + +def _turn_completed_reading_cache(cached_tokens: int) -> str: + return json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_cached_turn", + "status": "completed", + "output": [], + "usage": { + "input_tokens": _TURN_INPUT_TOKENS, + "input_tokens_details": {"cached_tokens": cached_tokens}, + "output_tokens": _OUTPUT_TOKENS_PER_TURN, + "total_tokens": _TURN_INPUT_TOKENS + _OUTPUT_TOKENS_PER_TURN, + }, + }, + } + ) + + +def _connection_logged_through(rate_limits: "ModelRateLimitingCheck", opened_with: Mapping[str, object]) -> "Logging": + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CallTypes + + connection: Final = Logging( + model="gpt-5-mini", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id="connection-call", + function_id="connection-call", + dynamic_async_success_callbacks=[rate_limits], + dynamic_async_failure_callbacks=[rate_limits], + ) + connection.update_environment_variables( + litellm_params={"metadata": dict(opened_with), CallTypes.aresponses_websocket.value: True}, + optional_params={}, + ) + return connection + + +def _provider_turn_stream(_request: httpx.Request) -> httpx.Response: + completed: Final = { + "id": "resp_turn", + "object": "response", + "created_at": 1767271505, + "status": "completed", + "model": "gpt-5-mini", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 5, "output_tokens": _OUTPUT_TOKENS_PER_TURN, "total_tokens": 7}, + } + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text="".join( + f"event: {kind}\ndata: {json.dumps({'type': kind, 'response': completed})}\n\n" + for kind in ("response.created", "response.completed") + ), + ) + + +class _ProviderResponsesSocket: + """A provider's native Responses socket that answers each ``response.create`` with a completed turn.""" + + def __init__(self) -> None: + self.sent: Final[list[dict[str, object]]] = [] + self._replies: Final[asyncio.Queue[str]] = asyncio.Queue() + + def connect(self, _url: str, **_kwargs: object) -> "_ProviderResponsesSocket": + return self + + async def __aenter__(self) -> "_ProviderResponsesSocket": + return self + + async def __aexit__(self, *_exc_info: object) -> None: + return None + + async def send(self, message: str) -> None: + self.sent.append(json.loads(message)) + await self._replies.put(_turn_completed()) + + async def recv(self, decode: bool = False) -> str: + return await self._replies.get() + + async def close(self) -> None: + return None + + +class _TwoOutputTokenTurns(CustomLLM): + def __init__(self, failing_turns: frozenset[int] = frozenset()) -> None: + super().__init__() + self._failing_turns: Final = failing_turns + self._turns: Final = itertools.count(1) + + async def astreaming(self, *args: object, **kwargs: object) -> AsyncIterator[GenericStreamingChunk]: + if next(self._turns) in self._failing_turns: + raise RuntimeError("upstream failed") + yield GenericStreamingChunk( + text="hi", + is_finished=True, + finish_reason="stop", + index=0, + tool_use=None, + usage={"prompt_tokens": 5, "completion_tokens": _OUTPUT_TOKENS_PER_TURN, "total_tokens": 7}, + ) + + +async def _run_managed_relay( + monkeypatch: pytest.MonkeyPatch, + provider: CustomLLM, + rate_limits: "ModelRateLimitingCheck", + frames: tuple[str, ...], + turn_admission: "ResponsesWsTurnAdmission | None", + opened_with: Mapping[str, object], +) -> _ClientSocket: + import litellm + + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "ws-turn-llm", "custom_handler": provider}]) + return await _relay_turns_to(monkeypatch, "ws-turn-llm/turns", rate_limits, frames, turn_admission, opened_with) + + +async def _relay_turns_to( + monkeypatch: pytest.MonkeyPatch, + model: str, + rate_limits: "ModelRateLimitingCheck", + frames: tuple[str, ...], + turn_admission: "ResponsesWsTurnAdmission | None", + opened_with: Mapping[str, object], + **connection: object, +) -> _ClientSocket: + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler, drain_logging_tasks + + monkeypatch.setattr(litellm, "callbacks", [rate_limits]) + first, *later = frames + websocket: Final = _ClientSocket(*later) + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model=model, + logging_obj=Logging( + model=model, + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + litellm_metadata={"user_api_key_hash": "k1", **opened_with}, + first_message=first, + turn_admission=turn_admission, + **connection, + ) + + await handler.run() + await drain_logging_tasks() + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + return websocket + + def _errors(websocket: _ClientSocket) -> list[tuple[object, object]]: return [(frame.get("status"), frame["error"]) for frame in websocket.sent if frame.get("type") == "error"] @@ -1722,6 +1958,195 @@ class TestWebSocketTurnAdmission: assert isinstance(failure, RateLimitError) assert "crossed budget" in str(failure) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("first_turn_end", "first_turn_tokens"), + [ + pytest.param(_turn_completed(), _OUTPUT_TOKENS_PER_TURN, id="first_turn_completed"), + pytest.param(_UPSTREAM_ERROR, 0, id="first_turn_failed"), + ], + ) + async def test_native_relay_counts_every_turn_on_the_minute_it_finishes( + self, first_turn_end: str, first_turn_tokens: int, monkeypatch: pytest.MonkeyPatch + ) -> None: + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.caching.dual_cache import DualCache + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + _turns_finish_at(monkeypatch, _TURNED_AT) + dual_cache: Final = DualCache() + opened_with: Final = await _socket_opened_on(dual_cache) + first, second, third = _turn_frames() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock( + side_effect=[first_turn_end, _turn_completed(), _turn_completed(), Exception("stop")] + ) + handler: Final = _make_streaming( + websocket=_ClientSocket(second, third), + backend_ws=backend_ws, + request_data={}, + logging_obj=_connection_logged_through(ModelRateLimitingCheck(dual_cache), opened_with), + first_message=first, + turn_admission=_later_turns_admitted_through(dual_cache), + ) + + await handler.client_to_backend() + await handler.backend_to_client() + await drain_logging_tasks() + + assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE + assert ( + await dual_cache.async_get_cache(key=_TURNED_ON) + == _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("cached_tokens", "counted_input_tokens"), + [ + pytest.param(800, _TURN_INPUT_TOKENS - 800, id="cached_input"), + pytest.param(0, _TURN_INPUT_TOKENS, id="uncached_input"), + ], + ) + async def test_native_relay_leaves_a_turn_s_cached_input_out_of_the_deployment_s_itpm( + self, cached_tokens: int, counted_input_tokens: int, monkeypatch: pytest.MonkeyPatch + ) -> None: + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.caching.dual_cache import DualCache + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + _turns_finish_at(monkeypatch, _OPENED_AT) + dual_cache: Final = DualCache() + opened_with: Final = await async_admit_io_token_turn(dual_cache, _ITPM_DEPLOYMENT, _OPENED_AT) + first, _, _ = _turn_frames() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=[_turn_completed_reading_cache(cached_tokens), Exception("stop")]) + handler: Final = _make_streaming( + websocket=_ClientSocket(), + backend_ws=backend_ws, + request_data={}, + logging_obj=_connection_logged_through(ModelRateLimitingCheck(dual_cache), opened_with), + first_message=first, + ) + + await handler.client_to_backend() + await handler.backend_to_client() + await drain_logging_tasks() + + assert await dual_cache.async_get_cache(key=_ITPM_OPENED_ON) == counted_input_tokens + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("failing_turns", "first_turn_tokens"), + [ + pytest.param(frozenset(), _OUTPUT_TOKENS_PER_TURN, id="first_turn_completed"), + pytest.param(frozenset({1}), 0, id="first_turn_failed"), + ], + ) + async def test_managed_relay_counts_every_turn_on_the_minute_it_finishes( + self, monkeypatch: pytest.MonkeyPatch, failing_turns: frozenset[int], first_turn_tokens: int + ) -> None: + from litellm.caching.dual_cache import DualCache + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + _turns_finish_at(monkeypatch, _TURNED_AT) + dual_cache: Final = DualCache() + opened_with: Final = await _socket_opened_on(dual_cache) + + await _run_managed_relay( + monkeypatch, + _TwoOutputTokenTurns(failing_turns), + ModelRateLimitingCheck(dual_cache), + _turn_frames(), + _later_turns_admitted_through(dual_cache), + opened_with, + ) + + assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE + assert ( + await dual_cache.async_get_cache(key=_TURNED_ON) + == _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN + ) + + @pytest.mark.asyncio + async def test_managed_relay_sends_the_provider_only_the_metadata_each_turn_carries( + self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter + ) -> None: + import litellm + from litellm.caching.dual_cache import DualCache + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + _turns_finish_at(monkeypatch, _TURNED_AT) + provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=_provider_turn_stream) + client_metadata: Final = {"session": "abc", "team": "search"} + dual_cache: Final = DualCache() + + await _relay_turns_to( + monkeypatch, + "openai/gpt-5-mini", + ModelRateLimitingCheck(dual_cache), + tuple( + json.dumps({"type": "response.create", "input": word, "metadata": client_metadata}) + for word in ("one", "two") + ), + _later_turns_admitted_through(dual_cache), + await _socket_opened_on(dual_cache), + api_key="sk-fake", + ) + + assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [ + client_metadata, + client_metadata, + ] + assert await dual_cache.async_get_cache(key=_TURNED_ON) == _OTHER_USAGE + 2 * _OUTPUT_TOKENS_PER_TURN + + @pytest.mark.asyncio + async def test_native_relay_sends_the_provider_only_the_metadata_the_socket_opened_with( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from unittest.mock import patch + + from litellm import Router + + monkeypatch.setenv("LITELLM_RUST", "false") + provider: Final = _ProviderResponsesSocket() + client_metadata: Final = {"session": "abc", "team": "search"} + first, second, _ = _turn_frames() + router: Final = Router( + model_list=[ + { + "model_name": "ws-io", + "litellm_params": { + "model": "openai/gpt-5-mini", + "api_key": "sk-fake", + "itpm": 40_000, + "otpm": 4096, + }, + "model_info": {"id": "ws-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + with patch("websockets.connect", provider.connect): + await router._aresponses_websocket( + model="ws-io", + websocket=_ClientSocket(second), + first_message=first, + metadata=dict(client_metadata), + ) + + assert [frame.get("metadata") for frame in provider.sent] == [client_metadata, client_metadata] + @pytest.mark.asyncio async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch: pytest.MonkeyPatch): from unittest.mock import AsyncMock @@ -1875,7 +2300,7 @@ class TestWebSocketTurnAdmission: request_data={}, first_message=first, frame_guard=_guard_recording(bill_refusal), - turn_admission=AsyncMock(side_effect=[refusal, None]), + turn_admission=AsyncMock(side_effect=[refusal, {}]), ) close_time_failure: Final = await handler.bidirectional_forward() @@ -1904,7 +2329,7 @@ class TestWebSocketTurnAdmission: request_data={}, first_message=first, frame_guard=_guard_recording(bill_refusal), - turn_admission=AsyncMock(side_effect=[refusal, None]), + turn_admission=AsyncMock(side_effect=[refusal, {}]), ) relay: Final = asyncio.ensure_future(handler.bidirectional_forward()) await asyncio.wait_for(client.send_started.wait(), timeout=1) @@ -1955,7 +2380,7 @@ class TestWebSocketTurnAdmission: ), first_message=first, frame_guard=_guard_recording(bill_refusal), - turn_admission=AsyncMock(side_effect=[refusal, None]), + turn_admission=AsyncMock(side_effect=[refusal, {}]), ) close_time_failure: Final = await handler.run() diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index e5ff4531f84..819c6b9d4f5 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -4,14 +4,19 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). import asyncio from collections.abc import Awaitable, Callable, Mapping -from datetime import datetime +from datetime import datetime, timezone from typing import Final +from unittest.mock import AsyncMock, MagicMock import pytest +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeRedis as FakeAsyncRedis import litellm from litellm import Router from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -22,6 +27,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( OTPM_RESERVED_KEY, _explicit_max_tokens, _reservation_value, + async_admit_io_token_turn, async_io_token_pre_call_check, async_io_token_reconcile_success, build_io_token_rate_limit_headers, @@ -35,6 +41,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) +from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterCacheEnum from litellm.types.utils import ModelResponse, Usage from litellm.utils import get_utc_datetime @@ -115,7 +122,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 10, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") @@ -173,7 +180,8 @@ class TestModelRateLimitingCheckIOTokens: ], "max_tokens": 10, "metadata": {}, - } + }, + metadata_channel="metadata", ) with pytest.raises(litellm.RateLimitError) as exc_info: @@ -203,7 +211,8 @@ class TestModelRateLimitingCheckIOTokens: "messages": [{"role": "user", "content": "hi"}], "max_tokens": max_tokens, "metadata": {}, - } + }, + metadata_channel="metadata", ) async def _attempt(): @@ -255,7 +264,8 @@ class TestModelRateLimitingCheckIOTokens: { "max_tokens": 5, "metadata": {}, - } + }, + metadata_channel="metadata", ) async def _attempt(): @@ -320,7 +330,7 @@ class TestModelRateLimitingCheckIOTokens: } request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") @@ -363,15 +373,15 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "opus", } request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) # Second zero-estimate call still fits within the itpm=2 limit. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}, metadata_channel="metadata") await check.async_pre_call_check(deployment) # A third exceeds the limit and is rejected. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}, metadata_channel="metadata") with pytest.raises(litellm.RateLimitError): await check.async_pre_call_check(deployment) @@ -392,7 +402,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 0, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(zero_output_kwargs) + set_io_token_rate_limit_request_kwargs(zero_output_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] @@ -404,7 +414,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 5, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(normal_output_kwargs) + set_io_token_rate_limit_request_kwargs(normal_output_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] @@ -432,7 +442,7 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "no-max-tokens", } request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") if use_sync_check: check.pre_call_check(deployment) @@ -462,7 +472,7 @@ class TestModelRateLimitingCheckIOTokens: async def _attempt() -> dict[str, object] | None: request_kwargs: Final = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") try: await check.async_pre_call_check(deployment) except litellm.RateLimitError: @@ -493,7 +503,7 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "roomy", } request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) @@ -507,7 +517,7 @@ class TestModelRateLimitingCheckIOTokens: async def _reserved(model_info: dict[str, str]) -> int: request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") deployment: Final = { "litellm_params": {"model": "azure/litellm-io-custom-deployment", "otpm": otpm_limit}, "model_info": model_info, @@ -536,7 +546,7 @@ class TestModelRateLimitingCheckIOTokens: } async def _pre_call(request_kwargs: dict[str, object]) -> None: - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") if use_sync_check: check.pre_call_check(deployment) else: @@ -563,12 +573,14 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "opus", } request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 100, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == 100 - set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": "hi"}], "metadata": {}}) + set_io_token_rate_limit_request_kwargs( + {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}, metadata_channel="metadata" + ) with pytest.raises(litellm.RateLimitError): await check.async_pre_call_check(deployment) @@ -585,7 +597,7 @@ class TestModelRateLimitingCheckIOTokens: with_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 300, "metadata": {}} for request_kwargs in (without_cap, with_cap): - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") if use_sync_check: check.pre_call_check(deployment) else: @@ -602,7 +614,7 @@ class TestModelRateLimitingCheckIOTokens: async def reserved(deployment_cap: Mapping[str, int], request_shape: Mapping[str, object]) -> int: request_kwargs: Final = {**request_shape, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check( { "litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, **deployment_cap}, @@ -631,7 +643,7 @@ class TestModelRateLimitingCheckIOTokens: async def reserved(deployment_cap: Mapping[str, int]) -> int: request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") deployment: Final = { "litellm_params": {"model": UNMAPPED_MODEL, "otpm": otpm_limit, **deployment_cap}, "model_info": {"id": f"io-deployment-cap-above-default-{deployment_cap}"}, @@ -667,7 +679,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 10, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") check.pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") @@ -803,7 +815,7 @@ class TestModelRateLimitingCheckIOTokens: # The retry proceeds to stash deployment B's own reservation on the # same dict; it starts clean, unaffected by A's cleared sentinels. - set_io_token_rate_limit_request_kwargs(shared_kwargs) + set_io_token_rate_limit_request_kwargs(shared_kwargs, metadata_channel="metadata") itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b @@ -819,7 +831,7 @@ class TestModelRateLimitingCheckIOTokens: "litellm_metadata": {OTPM_RESERVED_KEY: 7}, "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, } - set_io_token_rate_limit_request_kwargs(forged) + set_io_token_rate_limit_request_kwargs(forged, metadata_channel="metadata") stored = get_io_token_rate_limit_request_kwargs() assert ITPM_RESERVED_KEY not in stored["metadata"] @@ -845,7 +857,7 @@ class TestModelRateLimitingCheckIOTokens: }, } # The router sanitizes the request kwargs before the call runs. - set_io_token_rate_limit_request_kwargs(kwargs) + set_io_token_rate_limit_request_kwargs(kwargs, metadata_channel="metadata") response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), @@ -880,7 +892,8 @@ class TestModelRateLimitingCheckIOTokens: "messages": [{"role": "user", "content": "hello world"}], "max_tokens": 5, "metadata": {}, - } + }, + metadata_channel="metadata", ) with pytest.raises(RuntimeError): @@ -945,7 +958,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 5, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") result = await async_io_token_pre_call_check(dual_cache, deployment) @@ -1147,7 +1160,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 5, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): with pytest.raises(litellm.RateLimitError): @@ -1495,15 +1508,15 @@ class TestContextSlotRetention: "messages": [{"role": "user", "content": "x" * 1000}], "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, } - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is None assert ITPM_RESERVED_KEY not in kwargs["metadata"] assert ITPM_CACHE_KEY not in kwargs["metadata"] - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is kwargs - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is None @pytest.mark.asyncio @@ -1516,7 +1529,7 @@ class TestContextSlotRetention: } ] ) - set_io_token_rate_limit_request_kwargs(None) + set_io_token_rate_limit_request_kwargs(None, metadata_channel="metadata") kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} deployment = router.get_deployment_by_model_group_name("plain") assert deployment is not None @@ -1534,7 +1547,7 @@ class TestContextSlotRetention: ], optional_pre_call_checks=["enforce_model_rate_limits"], ) - set_io_token_rate_limit_request_kwargs(None) + set_io_token_rate_limit_request_kwargs(None, metadata_channel="metadata") kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} deployment = router.get_deployment_by_model_group_name("limited") assert deployment is not None @@ -1551,7 +1564,7 @@ async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_tok "max_tokens": 5, "metadata": metadata, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await async_io_token_pre_call_check(DualCache(), deployment) return metadata[ITPM_RESERVED_KEY] @@ -1571,3 +1584,278 @@ async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_tok ) assert named == served + + +def _reply_using(prompt_tokens: int, completion_tokens: int) -> ModelResponse: + return ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ) + + +def _proxy_success_kwargs(reservation: Mapping[str, object]) -> dict[str, object]: + return { + "litellm_params": {"metadata": {**reservation, "user_api_key_hash": "hashed-key"}}, + "standard_logging_object": { + "hidden_params": {"litellm_model_name": "openai/gpt-5-mini"}, + "metadata": {"user_api_key_hash": "hashed-key"}, + }, + } + + +async def _deliver_success( + check: ModelRateLimitingCheck, delivery: str, kwargs: dict[str, object], reply: object +) -> None: + if delivery == "async": + await check.async_log_success_event(kwargs, reply, None, None) + return + check.log_success_event(kwargs, reply, None, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deliveries", + [("async",), ("sync",), ("async", "async"), ("async", "sync"), ("sync", "async"), ("sync", "sync")], +) +async def test_success_settles_a_reservation_held_only_in_litellm_params_metadata_exactly_once( + deliveries: tuple[str, ...], +) -> None: + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:itpm:12-34" + otpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:otpm:12-34" + await dual_cache.async_increment_cache(key=itpm_key, value=40, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=4096, ttl=60) + kwargs: Final = _proxy_success_kwargs( + {ITPM_RESERVED_KEY: 40, OTPM_RESERVED_KEY: 4096, ITPM_CACHE_KEY: itpm_key, OTPM_CACHE_KEY: otpm_key} + ) + + for delivery in deliveries: + await _deliver_success(check, delivery, kwargs, _reply_using(prompt_tokens=7, completion_tokens=3)) + + assert await dual_cache.async_get_cache(key=itpm_key) == 7 + assert await dual_cache.async_get_cache(key=otpm_key) == 3 + + +_TURN_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc) +_ACTUAL_TOKENS: Final = {"itpm": 5, "otpm": 2} + + +def _turn_deployment(kind: str) -> dict[str, object]: + return {"litellm_params": {"model": "openai/gpt-5-mini", kind: 100}, "model_info": {"id": f"ws-turn-{kind}"}} + + +def _turn_key(kind: str, minute: str = "12-34") -> str: + return f"global_router:ws-turn-{kind}:openai/gpt-5-mini:{kind}:{minute}" + + +@pytest.fixture +def settled_on_the_turn_s_minute(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: _TURN_AT + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("used", [100, 101, 250]) +async def test_turn_admission_refuses_once_the_minutes_tokens_reach_the_limit(kind: str, used: int) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key(kind), value=used, ttl=60) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage={used}"): + await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("used", [0, 1, 99]) +async def test_an_admitted_turn_reserves_nothing_and_counts_its_actual_tokens(kind: str, used: int) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key(kind), value=used, ttl=60) + actual: Final = _ACTUAL_TOKENS + + reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + counter_key: Final = reservation[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] + assert isinstance(counter_key, str) + assert await dual_cache.async_get_cache(key=counter_key) == used + await async_io_token_reconcile_success( + dual_cache, _proxy_success_kwargs(reservation), _reply_using(actual["itpm"], actual["otpm"]) + ) + + assert await dual_cache.async_get_cache(key=counter_key) == used + actual[kind] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment", + [ + {"litellm_params": {"model": "openai/gpt-5-mini", "tpm": 1}, "model_info": {"id": "ws-turn-unlimited"}}, + {"litellm_params": {"model": "openai/gpt-5-mini", "otpm": 1}}, + ], +) +async def test_a_turn_without_io_token_counters_is_admitted_with_nothing_to_settle( + deployment: dict[str, object], +) -> None: + assert await async_admit_io_token_turn(DualCache(), deployment, _TURN_AT) == {} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("used_by_others", [0, 1, 99]) +async def test_a_socket_open_reserves_nothing_and_its_first_turn_counts_its_actual_tokens_on_its_own_minute( + kind: str, used_by_others: int +) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key(kind), value=used_by_others, ttl=60) + socket_open: Final[dict[str, object]] = {"metadata": {}} + set_io_token_rate_limit_request_kwargs(socket_open, opens_responses_socket=True, metadata_channel="metadata") + + assert await async_io_token_pre_call_check(dual_cache, _turn_deployment(kind)) == _turn_deployment(kind) + assert await dual_cache.async_get_cache(key=_turn_key(kind)) == used_by_others + + opened_with: Final = socket_open["metadata"] + assert isinstance(opened_with, dict) + await async_io_token_reconcile_success( + dual_cache, _proxy_success_kwargs(opened_with), _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]) + ) + + assert await dual_cache.async_get_cache(key=_turn_key(kind)) == used_by_others + _ACTUAL_TOKENS[kind] + + +class _RedisCacheOnFakeServer(RedisCache): + def __init__(self, server: FakeServer) -> None: + super().__init__(host="io-token-turn-redis", connection_pool=FakeRedis(server=server).connection_pool) + self._fake_async_client: Final = FakeAsyncRedis(server=server) + + def init_async_client(self) -> FakeAsyncRedis: + return self._fake_async_client + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_turn_is_refused_on_this_instance_s_count_when_reading_the_shared_counter_fails( + kind: str, +) -> None: + server: Final = FakeServer() + dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnFakeServer(server)) + server.connected = False + await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=100) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage=100"): + await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("delivery", ["async", "sync"]) +async def test_a_turn_finishing_in_the_next_minute_counts_its_actual_tokens_on_that_minute( + kind: str, delivery: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", + lambda: datetime(2026, 1, 1, 12, 35, 10, tzinfo=timezone.utc), + ) + dual_cache: Final = DualCache() + reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + await _deliver_success( + ModelRateLimitingCheck(dual_cache=dual_cache), + delivery, + _proxy_success_kwargs(reservation), + _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]), + ) + + assert [await dual_cache.async_get_cache(key=_turn_key(kind, minute)) for minute in ("12-34", "12-35")] == [ + None, + _ACTUAL_TOKENS[kind], + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_turn_is_admitted_once_the_shared_counter_drops_below_the_limit_this_instance_last_read( + kind: str, +) -> None: + redis_cache: Final = _RedisCacheOnFakeServer(FakeServer()) + dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=100) + await redis_cache.async_increment(key=_turn_key(kind), value=99) + + admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "responses_reply", + [ + pytest.param( + ResponsesAPIResponse( + id="resp_cached", + created_at=0, + output=[], + usage=ResponseAPIUsage( + input_tokens=50, + output_tokens=5, + total_tokens=55, + input_tokens_details=InputTokensDetails(cached_tokens=30), + ), + ), + id="usage_object", + ), + pytest.param( + { + "usage": { + "input_tokens": 50, + "output_tokens": 5, + "total_tokens": 55, + "input_tokens_details": {"cached_tokens": 30}, + } + }, + id="usage_dict", + ), + ], +) +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +async def test_a_turn_counts_only_its_uncached_responses_input_toward_itpm(responses_reply: object) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key("itpm"), value=10, ttl=60) + + reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment("itpm"), _TURN_AT) + await async_io_token_reconcile_success(dual_cache, _proxy_success_kwargs(reservation), responses_reply) + + assert await dual_cache.async_get_cache(key=_turn_key("itpm")) == 10 + 50 - 30 + + +def _shared_counter_answering(answer: int | Exception) -> DualCache: + redis_cache: Final = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(side_effect=[answer]) + return DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_later_turn_is_refused_when_another_instance_filled_the_shared_minute(kind: str) -> None: + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage=100"): + await async_admit_io_token_turn(_shared_counter_answering(100), _turn_deployment(kind), _TURN_AT) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_later_turn_is_admitted_on_this_instance_s_count_when_the_shared_counter_is_unreachable( + kind: str, +) -> None: + dual_cache: Final = _shared_counter_answering(RedisCircuitBreakerOpenError("circuit open")) + await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=99) + + admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 1ac44b5976d..a6436e47507 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2200,7 +2200,11 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): ) def inject_alias_into_kwargs( - deployment: dict, kwargs: dict, function_name: str | None = None, splits_stream_timeout: bool = False + deployment: dict, + kwargs: dict, + function_name: str | None = None, + splits_stream_timeout: bool = False, + opens_responses_socket: bool = False, ) -> None: # Simulate the alias leaking into kwargs (as happens when # _ageneric_api_call_with_fallbacks sets kwargs["model"] = alias before diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 5f1b6032ca8..5f31ded626d 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -4,13 +4,15 @@ import gc import json import os import weakref -from collections.abc import Iterator, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from datetime import datetime, timezone from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm from litellm import Router @@ -538,14 +540,14 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( assert await router.cache.async_get_cache(key=tpm_key) == 17 -async def _kwargs_routed_to(router: Router, model: str, function_name: str) -> dict[str, object]: +async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> dict[str, object]: routed: Final[dict[str, object]] = {} async def call(**kwargs: object) -> None: routed.update(kwargs) call.__name__ = function_name - await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call) + await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call, **request) return routed @@ -715,6 +717,363 @@ async def test_responses_websocket_turns_do_not_reserve_deployment_otpm(): await admit_turn() +def _io_limited_router(kind: str, **litellm_params: object) -> Router: + return Router( + model_list=[ + { + "model_name": "ws-io", + "litellm_params": { + "model": "openai/ws-io-unpriced", + "api_key": "sk-fake", + kind: 10_000, + **litellm_params, + }, + "model_info": {"id": "ws-io-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + +def _minutes_counter_key(kind: str, minute: datetime, model: str = "openai/ws-io-unpriced") -> str: + return f"global_router:ws-io-dep:{model}:{kind}:{minute.strftime('%H-%M')}" + + +async def _count_minutes_tokens( + router: Router, kind: str, tokens: int, minute: datetime, model: str = "openai/ws-io-unpriced" +) -> None: + await router.cache.async_increment_cache(key=_minutes_counter_key(kind, minute, model), value=tokens, ttl=120) + + +async def _settle_turn( + router: Router, reservation: Mapping[str, object], input_tokens: int, output_tokens: int +) -> None: + from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + rate_limits: Final = next( + check for check in router.optional_callbacks or () if isinstance(check, ModelRateLimitingCheck) + ) + await rate_limits.async_log_success_event( + {"litellm_params": get_litellm_params(litellm_metadata=dict(reservation))}, + ResponsesWebSocketTokenUsageProcessor.create_logging_object( + [ + { + "type": "response.completed", + "response": {"usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}}, + } + ] + ), + None, + None, + ) + + +_LONG_OUTPUT_MODEL: Final = "openai/gpt-5.4-mini" + + +def _long_output_router() -> tuple[Router, Mapping[str, int]]: + max_output: Final = litellm.get_model_info(_LONG_OUTPUT_MODEL)["max_output_tokens"] + assert isinstance(max_output, int) + limits: Final = {"itpm": 10_000, "otpm": max_output // 2} + return _io_limited_router("itpm", model=_LONG_OUTPUT_MODEL, otpm=limits["otpm"]), limits + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_responses_websocket_opens_without_counting_tokens_until_the_minute_reaches_deployment_io_limit( + kind: str, router_minute_pinned: datetime +): + router, limits = _long_output_router() + + await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) + + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(counted, router_minute_pinned, _LONG_OUTPUT_MODEL)) + for counted in ("itpm", "otpm") + ] == [None, None] + await _count_minutes_tokens(router, kind, limits[kind], router_minute_pinned, _LONG_OUTPUT_MODEL) + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit={limits[kind]}, current usage="): + await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) + + +@pytest.mark.asyncio +async def test_responses_websocket_first_turn_counts_its_actual_tokens_toward_deployment_itpm_and_otpm( + router_minute_pinned: datetime, +): + router, _ = _long_output_router() + socket_metadata: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()))[ + "litellm_metadata" + ] + assert isinstance(socket_metadata, dict) + + await _settle_turn(router, socket_metadata, input_tokens=20, output_tokens=15) + + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(counted, router_minute_pinned, _LONG_OUTPUT_MODEL)) + for counted in ("itpm", "otpm") + ] == [20, 15] + + +_ADMITTED_LATE_IN_THE_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 50, tzinfo=timezone.utc) +_FINISHED_IN_THE_NEXT_MINUTE: Final = datetime(2026, 1, 1, 12, 1, 15, tzinfo=timezone.utc) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("turn", ["first", "later"]) +async def test_responses_websocket_turn_finishing_in_the_next_minute_counts_its_tokens_toward_that_minute( + kind: str, turn: str, monkeypatch: pytest.MonkeyPatch +): + _pin_router_clock(monkeypatch, _ADMITTED_LATE_IN_THE_MINUTE) + router: Final = _io_limited_router(kind) + socket: Final = await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) + admit_turn: Final = socket["turn_admission"] + assert callable(admit_turn) + reservation: Final = socket["litellm_metadata"] if turn == "first" else await admit_turn() + assert isinstance(reservation, Mapping) + _pin_router_clock(monkeypatch, _FINISHED_IN_THE_NEXT_MINUTE) + + await _settle_turn(router, reservation, input_tokens=10_000, output_tokens=10_000) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000, current usage=10000"): + await admit_turn() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [{}, {"websocket": True}], ids=["plain", "websocket_key"]) +async def test_messages_call_reserves_max_tokens_against_deployment_otpm_whatever_websocket_key_its_body_carries( + body: Mapping[str, object], router_minute_pinned: datetime +): + router: Final = _io_limited_router("otpm") + + for _ in range(3): + await _kwargs_routed_to(router, "ws-io", "anthropic_messages", max_tokens=3_000, **body) + + assert await router.cache.async_get_cache(key=_minutes_counter_key("otpm", router_minute_pinned)) == 9_000 + with pytest.raises(litellm.RateLimitError, match="OTPM limit=10000, current usage=12000"): + await _kwargs_routed_to(router, "ws-io", "anthropic_messages", max_tokens=3_000, **body) + + +async def _realtime_opens_until_refused(monkeypatch: pytest.MonkeyPatch, **socket: object) -> tuple[int, str]: + monkeypatch.setattr(litellm, "callbacks", []) + router: Final = _io_limited_router("otpm", otpm=50) + for opened in range(51): + try: + await _kwargs_routed_to(router, "ws-io", "_arealtime", **socket) + except litellm.RateLimitError as refused: + return opened, str(refused) + pytest.fail("realtime opens never reached the deployment's otpm") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") +async def test_realtime_open_reserves_and_is_refused_at_deployment_otpm_like_an_open_without_a_socket( + monkeypatch: pytest.MonkeyPatch, +): + with_socket: Final = await _realtime_opens_until_refused(monkeypatch, websocket=object()) + + assert "OTPM limit=50, current usage=" in with_socket[1] + assert with_socket == await _realtime_opens_until_refused(monkeypatch) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_responses_websocket_turn_is_refused_once_the_minutes_tokens_reach_deployment_io_limit( + kind: str, router_minute_pinned: datetime +): + router: Final = _io_limited_router(kind) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + await _count_minutes_tokens(router, kind, 10_000, router_minute_pinned) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000, current usage="): + await admit_turn() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_responses_websocket_turn_refused_for_tokens_leaves_the_daily_request_limit_unused( + kind: str, router_minute_pinned: datetime +): + router: Final = _io_limited_router(kind, rpd=2) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + await _count_minutes_tokens(router, kind, 10_000, router_minute_pinned) + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000"): + await admit_turn() + await _count_minutes_tokens(router, kind, -10_000, router_minute_pinned) + + await admit_turn() + + with pytest.raises(litellm.RateLimitError, match="RPD"): + await admit_turn() + + +_PROVIDER_RESPONSE: Final = { + "id": "resp_io", + "object": "response", + "created_at": 1767268830, + "status": "completed", + "model": "ws-io-unpriced", + "output": [ + { + "type": "message", + "id": "msg_io", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_metadata", [None, {"session": "abc", "team": "search"}], ids=["without_metadata", "with_metadata"] +) +async def test_responses_call_sends_the_provider_only_the_client_metadata_and_settles_deployment_io_tokens( + client_metadata: Mapping[str, str] | None, + router_minute_pinned: datetime, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_PROVIDER_RESPONSE) + ) + router: Final = _io_limited_router("itpm", otpm=10_000) + + await router.aresponses( + model="ws-io", + input="hi", + max_output_tokens=50, + **({} if client_metadata is None else {"metadata": dict(client_metadata)}), + ) + await drain_logging_tasks() + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [client_metadata] + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned)) + for kind in ("itpm", "otpm") + ] == [10, 5] + + +@pytest.mark.parametrize( + "client_metadata", [None, {"session": "abc", "team": "search"}], ids=["without_metadata", "with_metadata"] +) +def test_sync_responses_call_sends_the_provider_only_the_client_metadata( + client_metadata: Mapping[str, str] | None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +): + litellm.in_memory_llm_clients_cache.flush_cache() + provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_PROVIDER_RESPONSE) + ) + + _io_limited_router("itpm", otpm=10_000).responses( + model="ws-io", + input="hi", + max_output_tokens=50, + **({} if client_metadata is None else {"metadata": dict(client_metadata)}), + ) + + assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [client_metadata] + + +_ANTHROPIC_PROVIDER_RESPONSE: Final = { + "id": "msg_io", + "type": "message", + "role": "assistant", + "model": "ws-io-unpriced", + "content": [{"type": "text", "text": "Done."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, +} + + +async def _responses_call(router: Router) -> object: + return await router.aresponses(model="ws-io", input="hi", max_output_tokens=50) + + +async def _messages_call(router: Router) -> object: + return await router.aanthropic_messages(model="ws-io", messages=[{"role": "user", "content": "hi"}], max_tokens=50) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "provider_url", "provider_response", "call"), + [ + pytest.param( + "openai/ws-io-unpriced", + "https://api.openai.com/v1/responses", + _PROVIDER_RESPONSE, + _responses_call, + id="responses", + ), + pytest.param( + "anthropic/ws-io-unpriced", + "https://api.anthropic.com/v1/messages", + _ANTHROPIC_PROVIDER_RESPONSE, + _messages_call, + id="messages", + ), + ], +) +async def test_retried_call_settles_deployment_io_tokens_to_the_successful_attempt_s_tokens( + model: str, + provider_url: str, + provider_response: Mapping[str, object], + call: Callable[[Router], Awaitable[object]], + router_minute_pinned: datetime, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + provider: Final = respx_mock.post(provider_url).mock( + side_effect=[ + httpx.Response(500, json={"error": {"message": "upstream failed", "type": "server_error"}}), + httpx.Response(200, json=provider_response), + ] + ) + router: Final = Router( + model_list=[ + { + "model_name": "ws-io", + "litellm_params": {"model": model, "api_key": "sk-fake", "itpm": 10_000, "otpm": 10_000}, + "model_info": {"id": "ws-io-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + num_retries=1, + retry_after=0, + ) + + await call(router) + await drain_logging_tasks() + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert provider.call_count == 2 + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned, model)) + for kind in ("itpm", "otpm") + ] == [10, 5] + + @pytest.mark.asyncio async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment(): """ @@ -1194,10 +1553,17 @@ def _rpm_tpm_router(model_id: str) -> Router: ) +def _pin_router_clock(monkeypatch: pytest.MonkeyPatch, pinned: datetime) -> None: + monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned) + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: pinned + ) + + @pytest.fixture def router_minute_pinned(monkeypatch: pytest.MonkeyPatch) -> datetime: pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) - monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned) + _pin_router_clock(monkeypatch, pinned) return pinned From df47ec01717cef4e5614d34081a568be4e5de4c6 Mon Sep 17 00:00:00 2001 From: quietjules <349048734+quietjules@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:27:35 -0700 Subject: [PATCH 29/64] fix(spend): keep a NaN or infinite cost from wedging daily spend writes (#48829) * fix(spend): keep a NaN or infinite cost from wedging daily spend writes A non-finite provider-reported cost now counts as not reported, so the token price bills the request. Any other non-finite cost is logged and recorded as 0, and spend log JSON writes NaN and Infinity as null so Postgres accepts the row. * fix(spend): price non-finite provider costs on every route and drain parked NaN batches Perplexity, OpenRouter pass-through, and Responses streaming now treat a NaN or infinite reported cost as unreported, so the token price bills the request. A daily spend batch parked in the Redis buffer by an older version with a NaN spend now drains as 0 instead of failing every flush. The cost header never carries NaN or Infinity, and a set holding a non-finite float no longer breaks spend log JSON * fix(spend): record non-finite savings as 0 and warn when a parked NaN batch drains An overflowing price on a compressed request made compression savings infinite, which still wedged the daily spend batch. Every savings column now records a non-finite value as 0 with a warning. The Redis drain warns when it reads a NaN an older version parked, since the spend summed into it is lost from daily spend. The writer test now runs the real payload builder and spend log insert * fix(spend): keep non-finite costs out of headers, responses, and admin JSON Writing NaN as null is now opt-in for spend log columns, the Redis buffer, and streamed frames, so admin writes and logging integrations keep their output. The cost calculator records an overflowing total as 0 with a warning, cost headers leave out non-finite values, a non-streaming NaN usage.cost is dropped instead of failing the response, and a parked NaN spend log row reads back as 0. * fix(spend): keep caller-set call ids from splitting the non-finite cost warning The overflow warning strips newlines from litellm_call_id, and neither non-finite cost warning prints the raw value, which said nothing beyond "non-finite" * fix(spend): drop non-finite Responses and OpenRouter costs and the breakdown of an unpriceable request * fix(types): drop a usage cost breakdown that holds a non-finite number * fix(streaming): keep a dropped non-finite final usage cost from falling back to an earlier chunk's cost * fix(spend): keep the dropped-cost flag out of reach of provider usage keys A provider usage key named dropped_non_finite_cost shadowed the Usage method on the instance, so the stream builder raised and dropped the request's spend. The builder now reads the flag through the class, and Usage sets the flag after copying provider keys, so a _dropped_non_finite_cost key cannot set it either. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 21 ++- litellm/litellm_core_utils/litellm_logging.py | 25 ++- litellm/litellm_core_utils/safe_json_dumps.py | 23 ++- .../streaming_chunk_builder_utils.py | 9 +- .../llms/openrouter/chat/transformation.py | 8 +- .../pass_through/typesafe_logging_handler.py | 6 +- litellm/llms/perplexity/cost_calculator.py | 8 +- litellm/proxy/common_request_processing.py | 15 +- litellm/proxy/db/db_spend_update_writer.py | 12 +- .../redis_update_buffer.py | 24 ++- .../proxy/hooks/proxy_track_cost_callback.py | 3 +- litellm/proxy/spend_tracking/savings.py | 6 +- .../spend_tracking/spend_tracking_utils.py | 11 +- litellm/types/llms/openai.py | 10 +- litellm/types/utils.py | 30 +++- .../test_non_finite_cost_spend_writes.py | 166 ++++++++++++++++++ .../test_convert_dict_to_response.py | 44 +++++ .../litellm_core_utils/test_core_helpers.py | 28 ++- .../test_litellm_logging.py | 54 ++++++ .../test_safe_json_dumps.py | 25 +++ .../test_streaming_chunk_builder_utils.py | 41 +++++ .../test_openrouter_chat_transformation.py | 29 +++ ...est_openrouter_responses_transformation.py | 25 +++ .../test_perplexity_cost_calculator.py | 22 +++ .../test_redis_update_buffer.py | 130 ++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 54 ++++++ ...st_typesafe_passthrough_logging_handler.py | 29 +++ .../unit/proxy/spend_tracking/test_savings.py | 28 +++ .../proxy/test_common_request_processing.py | 49 ++++++ tests/unit/test_main.py | 22 +++ 30 files changed, 906 insertions(+), 51 deletions(-) create mode 100644 tests/integration/spend/test_non_finite_cost_spend_writes.py diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 28f45d2bb2c..2f81154ab83 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -982,25 +982,36 @@ def _usage_field(usage: object, name: str) -> object: return getattr(usage, name, None) -def _cost_number(value: object) -> float | None: - if isinstance(value, bool) or not isinstance(value, int | float): +def finite_cost(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float) or not math.isfinite(value): return None return float(value) +def single_log_line(value: object) -> str: + return str(value).replace("\r", "").replace("\n", "") + + +def finite_reported_cost(value: object) -> float | None: + cost: Final = finite_cost(value) + if cost is None and isinstance(value, float): + verbose_logger.warning("Ignoring a non-finite provider-reported cost %s", value) + return cost + + def _reported_cost_value(value: object) -> float | None: """A reported cost is a number, or a breakdown whose total lives under ``total_cost``.""" raw: Final[object] = value if isinstance(value, Mapping): - return _cost_number(_usage_field(raw, "total_cost")) - return _cost_number(value) + return finite_reported_cost(_usage_field(raw, "total_cost")) + return finite_reported_cost(value) def response_cost_in_hidden_params(response: _CarriesHiddenParams) -> float | None: additional_headers: Final[object] = response.hidden_params.get("additional_headers") if not isinstance(additional_headers, Mapping): return None - return _cost_number(_USAGE_BLOCK_ADAPTER.validate_python(additional_headers).get(RESPONSE_COST_HEADER)) + return finite_reported_cost(_USAGE_BLOCK_ADAPTER.validate_python(additional_headers).get(RESPONSE_COST_HEADER)) def override_stamped_response_cost(response: _CarriesHiddenParams, cost: float) -> None: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8bfb0b60693..cfe8c1739c6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -76,6 +76,7 @@ from litellm.litellm_core_utils.classifier_logging import ( is_classifier_call, ) from litellm.litellm_core_utils.core_helpers import ( + finite_cost, get_provider_response_headers_from_hidden_params, is_expected_client_error, native_messages_reported_cost, @@ -84,6 +85,7 @@ from litellm.litellm_core_utils.core_helpers import ( proxy_stamped_used_client_oauth_token, reconstruct_model_name, set_response_cost_in_hidden_params, + single_log_line, ) from litellm.litellm_core_utils.error_normalization import normalize_error from litellm.litellm_core_utils.get_litellm_params import get_litellm_params @@ -2093,7 +2095,7 @@ class Logging(LiteLLMLoggingBaseClass): litellm_model_name=litellm_model_name, router_model_id=router_model_id or hidden_params.get("model_id"), ) - return hidden_params["response_cost"] + return self._recorded_response_cost(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"] @@ -2193,7 +2195,7 @@ class Logging(LiteLLMLoggingBaseClass): litellm_model_name=litellm_model_name, router_model_id=router_model_id, ) - return total_response_cost + return self._recorded_response_cost(total_response_cost) except Exception as e: # error calculating cost debug_info = StandardLoggingModelCostFailureDebugInformation( error_str=str(e), @@ -2219,6 +2221,15 @@ class Logging(LiteLLMLoggingBaseClass): _response_cost_calculator = response_cost_calculator + def _recorded_response_cost(self, cost: float | None) -> float | None: + if cost is None or finite_cost(cost) is not None: + return cost + verbose_logger.warning( + "Recording a non-finite response_cost as 0 for litellm_call_id=%s", single_log_line(self.litellm_call_id) + ) + self.cost_breakdown = None + return 0.0 + def uncached_price(self, result: object) -> UncachedPrice: previous_cost_breakdown: Final = self.cost_breakdown previous_billed_token_rates: Final = self.billed_token_rates @@ -7582,7 +7593,15 @@ def get_standard_logging_object_payload( ) custom_pricing: Final = False if spilled_over else logging_obj.custom_pricing_for(init_response_obj) raw_response_cost: Final = kwargs.get("response_cost") - llm_response_cost: Final[float] = raw_response_cost or 0.0 + numeric_response_cost: Final = raw_response_cost if isinstance(raw_response_cost, int | float) else None + priced_response_cost: Final = finite_cost(numeric_response_cost) + if numeric_response_cost is not None and priced_response_cost is None: + verbose_logger.warning( + "Logging a non-finite response_cost=%s as 0 for litellm_call_id=%s", + numeric_response_cost, + single_log_line(logging_obj.litellm_call_id), + ) + llm_response_cost: Final[float] = priced_response_cost or 0.0 guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information")) finalized_cost: Final = _finalized_cost_of(logging_obj) response_cost: Final[float] = ( diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 24c401281ef..909e1222ae6 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -1,4 +1,5 @@ import json +import math from collections.abc import Callable from typing import Any, Final @@ -14,16 +15,23 @@ def strip_null_bytes(value: str) -> str: return value.replace("\x00", "") +def _finite_or_none(value: float | None) -> float | None: + return None if isinstance(value, float) and not math.isfinite(value) else value + + def safe_json_structure( data: object, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, value_transform: Callable[[str | None, str], str] | None = None, key: str | None = None, + *, + non_finite_as_null: bool = False, ) -> object: """ Rebuild data out of JSON-native pieces while detecting circular references. If a circular reference is detected then a marker string is returned. - NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors. + NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors. With + non_finite_as_null, NaN and infinite floats become null, since jsonb rejects them. value_transform, when given, is applied to every string leaf (and to the str() fallback for non-serializable objects) with the mapping key the leaf @@ -43,7 +51,7 @@ def safe_json_structure( cleaned = obj.replace("\x00", "") if "\x00" in obj else obj return _transform(key, cleaned) if isinstance(obj, (int, float, bool, type(None))): - return obj + return _finite_or_none(obj) if non_finite_as_null else obj # Check for circular reference. if id(obj) in seen: return "CircularReference Detected" @@ -65,8 +73,9 @@ def safe_json_structure( seen.remove(id(obj)) return tuple_result elif isinstance(obj, set): - sortable_values: Final[list[Any]] = [_serialize(item, seen, depth + 1, key) for item in obj] - set_result: Final = sorted(sortable_values) + members: Final[list[Any]] = [_serialize(item, seen, depth + 1, key) for item in obj] + present: Final = sorted(member for member in members if member is not None) + set_result: Final = present + [None] * members.count(None) seen.remove(id(obj)) return set_result elif isinstance(obj, BaseModel): @@ -88,6 +97,10 @@ def safe_dumps( data: object, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, value_transform: Callable[[str | None, str], str] | None = None, + *, + non_finite_as_null: bool = False, ) -> str: """Serialize data to JSON text through safe_json_structure.""" - return json.dumps(safe_json_structure(data, max_depth, value_transform), default=str) + return json.dumps( + safe_json_structure(data, max_depth, value_transform, non_finite_as_null=non_finite_as_null), default=str + ) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 94351e72ab1..574f1103c1a 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -15,6 +15,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import ( PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS, chunk_usage, + finite_cost, provider_reported_usage_cost, stamp_litellm_bill_on_usage, ) @@ -1014,7 +1015,9 @@ class ChunkProcessor: prompt_tokens_details, cache_creation_token_details ) - if usage_chunk_dict["cost"] is not None: + if Usage.dropped_non_finite_cost(usage_chunk): + cost = None + elif usage_chunk_dict["cost"] is not None: cost = usage_chunk_dict["cost"] prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details) @@ -1444,8 +1447,8 @@ def _set_stream_builder_response_cost( def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None: if logging_obj is None: return - if _logged_provider(logging_obj) in PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS and isinstance( - getattr(usage, "cost", None), (int, float) + if _logged_provider(logging_obj) in PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS and ( + finite_cost(getattr(usage, "cost", None)) is not None ): return computed_cost: Final = logging_obj.response_cost_calculator(result=response) diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index 9aaf552e79a..cca6815399f 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Final, cast import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.core_helpers import finite_cost from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -225,8 +226,11 @@ class OpenrouterConfig(OpenAIGPTConfig): additional_headers: Final = cast( # cast-ok: preserve mapping operations on response metadata dict[str, object], hidden_params["additional_headers"] ) - if isinstance(response_cost, str | int | float): - additional_headers["llm_provider-x-litellm-response-cost"] = float(response_cost) + reported_cost: Final = ( + finite_cost(float(response_cost)) if isinstance(response_cost, str | int | float) else None + ) + if reported_cost is not None: + additional_headers["llm_provider-x-litellm-response-cost"] = reported_cost except Exception: # If we can't extract cost, continue without it - don't fail the response pass diff --git a/litellm/llms/pass_through/typesafe_logging_handler.py b/litellm/llms/pass_through/typesafe_logging_handler.py index 9eae57774c2..7a41d2a5d14 100644 --- a/litellm/llms/pass_through/typesafe_logging_handler.py +++ b/litellm/llms/pass_through/typesafe_logging_handler.py @@ -6,6 +6,7 @@ import httpx from pydantic import TypeAdapter, ValidationError import litellm +from litellm.litellm_core_utils.core_helpers import finite_reported_cost from litellm.litellm_core_utils.decisions_pricing import apply_decisions_pricing from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import ( @@ -100,9 +101,10 @@ class TypeSafePassthroughLoggingHandler: f"{custom_llm_provider}/{model}" for model in candidate_models if model is not None ) pricing: Final = _pricing_for(candidate_model_keys) + reported_cost: Final = finite_reported_cost(usage.cost) if custom_llm_provider == "openrouter" else None response_cost: Final = ( - usage.cost - if custom_llm_provider == "openrouter" and usage.cost is not None + reported_cost + if reported_cost is not None else input_tokens * pricing.input_cost_per_token + output_tokens * pricing.output_cost_per_token ) usage_object: Final = Usage( diff --git a/litellm/llms/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index 2d762adf958..c3d178771a5 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -6,6 +6,7 @@ Helper util for handling perplexity-specific cost calculation from datetime import datetime from typing import Final +from litellm.litellm_core_utils.core_helpers import finite_reported_cost from litellm.litellm_core_utils.llm_cost_calc.utils import TokenRates, apply_off_peak_pricing from litellm.types.utils import Usage from litellm.utils import get_model_info @@ -34,10 +35,9 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non total_cost = cost_info.get("total_cost") elif isinstance(cost_info, (int, float)) and not isinstance(cost_info, bool): total_cost = float(cost_info) - if total_cost is not None: - # Return total cost as completion_cost (prompt_cost=0) since Perplexity - # doesn't break down by input/output in their cost object - return (0.0, float(total_cost)) + reported_cost: Final = None if total_cost is None else finite_reported_cost(float(total_cost)) + if reported_cost is not None: + return (0.0, reported_cost) ## FALLBACK: Calculate cost manually if Perplexity doesn't provide it ## GET MODEL INFO diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 48eb58adaf2..762f552dadc 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -60,6 +60,7 @@ from litellm.litellm_core_utils.bug_report import ( strip_bug_report_notice, ) from litellm.litellm_core_utils.core_helpers import ( + finite_cost, get_or_create_metadata_bucket, independent_snapshot, is_expected_client_error, @@ -1812,14 +1813,16 @@ def _get_cost_breakdown_from_logging_obj( def _header_value(value: float | None) -> str | None: - return str(value) if value is not None else None + return str(value) if finite_cost(value) is not None else None def _response_cost_header_value(hidden_params: Mapping[str, object]) -> float | str: if hidden_params.get("cache_hit") is True: return 0.0 response_cost: Final = hidden_params.get("response_cost") - return response_cost if isinstance(response_cost, (int, float, str)) and response_cost else "" + if isinstance(response_cost, str): + return response_cost + return response_cost if isinstance(response_cost, int | float) and finite_cost(response_cost) else "" def cost_headers( @@ -4509,7 +4512,7 @@ class ProxyBaseLLMRequestProcessing: """ if isinstance(chunk, dict): # Use safe_dumps for proper JSON serialization with circular reference detection - chunk_str: Final = safe_dumps(chunk) + chunk_str: Final = safe_dumps(chunk, non_finite_as_null=True) return f"{STREAM_SSE_DATA_PREFIX}{chunk_str}\n\n" else: return chunk @@ -4885,7 +4888,11 @@ class ProxyBaseLLMRequestProcessing: obj, model_name, litellm_logging_obj ) if maybe_modified is not None: - lines[idx] = "data: " + safe_dumps(maybe_modified) + ("\r" if ln.endswith("\r") else "") + lines[idx] = ( + "data: " + + safe_dumps(maybe_modified, non_finite_as_null=True) + + ("\r" if ln.endswith("\r") else "") + ) return "\n".join(lines) return None except Exception: diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index fe804a82e96..d9fb236632b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -34,6 +34,7 @@ from litellm.constants import ( DB_SPEND_UPDATE_JOB_NAME, INTERNAL_CALL_ORIGIN_METADATA_KEY, ) +from litellm.litellm_core_utils.core_helpers import finite_cost, single_log_line from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( @@ -652,7 +653,14 @@ class DBSpendUpdateWriter: end_time=end_time, llm_router=llm_router, ) - payload["spend"] = response_cost or 0.0 + billable_cost: Final = finite_cost(response_cost) + if response_cost is not None and billable_cost is None: + verbose_proxy_logger.warning( + "Spend tracking - recording a non-finite response_cost=%s as 0 for request_id=%s", + response_cost, + single_log_line(payload["request_id"]), + ) + payload["spend"] = billable_cost or 0.0 if isinstance(payload["startTime"], datetime): payload["startTime"] = payload["startTime"].isoformat() if isinstance(payload["endTime"], datetime): @@ -702,7 +710,7 @@ class DBSpendUpdateWriter: spawn_logging_task( asyncio.get_running_loop(), self._batch_database_updates( - response_cost=response_cost, + response_cost=billable_cost, user_id=user_id, hashed_token=hashed_token, team_id=team_id, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 3a4375ac54a..459f95972ce 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -8,7 +8,7 @@ import asyncio import json from collections.abc import Mapping, Sequence from datetime import datetime -from functools import reduce +from functools import partial, reduce from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast @@ -109,6 +109,16 @@ _ValueT = TypeVar("_ValueT") _DECODED_JSON: Final = TypeAdapter(object) +def _zero_for_non_finite_constant(constant: str) -> float: + verbose_proxy_logger.warning( + "Redis spend buffer - reading a parked %s spend as 0, the spend it held is lost", constant + ) + return 0.0 + + +_load_buffered_transaction: Final = partial(json.loads, parse_constant=_zero_for_non_finite_constant) + + def _spend_log_json_default(value: object) -> str: return value.isoformat() if isinstance(value, datetime) else str(value) @@ -118,7 +128,7 @@ def _encode_spend_log_row(row: SpendLogRow) -> str: def _decode_spend_log_row(encoded: str) -> dict[str, object] | None: - decoded: Final = _DECODED_JSON.validate_python(json.loads(encoded)) + decoded: Final = _DECODED_JSON.validate_python(_load_buffered_transaction(encoded)) return decoded if isinstance(decoded, dict) else None @@ -825,13 +835,13 @@ class RedisUpdateBuffer: if slot_result is None: daily_results.append(None) else: - list_of_daily = [json.loads(t) for t in slot_result] + list_of_daily = [_load_buffered_transaction(t) for t in slot_result] aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(list_of_daily) daily_results.append(aggregated) window_spend: Final = ( WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( - tuple(json.loads(transaction) for transaction in raw_results[6]) + tuple(_load_buffered_transaction(transaction) for transaction in raw_results[6]) ) if raw_results[6] is not None else None @@ -880,7 +890,7 @@ class RedisUpdateBuffer: ) if list_of_transactions is None: return None - return [json.loads(transaction) for transaction in list_of_transactions] + return [_load_buffered_transaction(transaction) for transaction in list_of_transactions] async def get_all_daily_spend_update_transactions_from_redis_buffer( self, @@ -1010,9 +1020,9 @@ class RedisUpdateBuffer: Parses the list of transactions from Redis """ if isinstance(list_of_transactions, list): - return [json.loads(transaction) for transaction in list_of_transactions] + return [_load_buffered_transaction(transaction) for transaction in list_of_transactions] else: - return [json.loads(list_of_transactions)] + return [_load_buffered_transaction(list_of_transactions)] @staticmethod def _combine_list_of_transactions( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 9c4a4bc775b..0d333d4fd08 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -16,6 +16,7 @@ from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLE from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, + finite_cost, get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, get_parent_otel_span_from_kwargs, @@ -703,7 +704,7 @@ class _ProxyDBLogger(CustomLogger): raw_cost: Final = _str_keyed_fields( sl_object if sl_object is not None else kwargs # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload ).get("response_cost") - priced_cost: Final = raw_cost if isinstance(raw_cost, (int, float)) else None + priced_cost: Final = finite_cost(raw_cost) response_cost: Final = ( await apply_search_free_tier( kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index 28602fbbefc..8e79ea5255f 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -689,9 +689,13 @@ def compute_savings_spend( cost_breakdown=cost_breakdown, ) ) - return SavingsSpend( + savings: Final = SavingsSpend( compression=compression, prompt_caching=prompt_caching, autorouter=0.0 if autorouter is None else autorouter, gateway_injected_caching=gateway_injected_caching, ) + if all(isfinite(value) for value in savings): + return savings + verbose_proxy_logger.warning("Spend tracking - recording non-finite savings %s as 0 for model=%s", savings, model) + return SavingsSpend(*(value if isfinite(value) else 0.0 for value in savings)) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 09b7ce7ee9a..f8d5cba4229 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -912,7 +912,7 @@ def get_logging_payload( user=metadata.get("user_api_key_user_id", "") or "", team_id=metadata.get("user_api_key_team_id", "") or "", organization_id=metadata.get("user_api_key_org_id") or "", - metadata=safe_dumps(clean_metadata), + metadata=safe_dumps(clean_metadata, non_finite_as_null=True), cache_key=cache_key, spend=kwargs.get("response_cost", 0), total_tokens=usage.get("total_tokens", standard_logging_total_tokens), @@ -1224,7 +1224,7 @@ def _get_messages_for_spend_logs_payload( messages: Final = standard_logging_payload.get("messages") if messages is not None: try: - return safe_dumps(messages) + return safe_dumps(messages, non_finite_as_null=True) except Exception: return "{}" return "{}" @@ -1919,7 +1919,8 @@ def _get_proxy_server_request_for_spend_logs_payload( _request_body_json_str: Final = safe_dumps( _stored_request_without_requester_ip_address(stored_request_body) if drop_requester_ip_address - else stored_request_body + else stored_request_body, + non_finite_as_null=True, ) if LITELLM_TRUNCATED_PAYLOAD_FIELD in _request_body_json_str: verbose_proxy_logger.info( @@ -2000,7 +2001,7 @@ def _get_response_for_spend_logs_payload( if isinstance(sanitized_response, str): result_str = strip_null_bytes(sanitized_response) else: - result_str = safe_dumps(sanitized_response) + result_str = safe_dumps(sanitized_response, non_finite_as_null=True) if LITELLM_TRUNCATED_PAYLOAD_FIELD in result_str: verbose_proxy_logger.info( "Spend Log: response was truncated before storing in DB. %s", @@ -2071,7 +2072,7 @@ def spend_log_row_with_retained_metadata( return row metadata: Final = _SPEND_LOGS_METADATA_ADAPTER.validate_json(metadata_json) retained: Final = {name: value for name, value in metadata.items() if fields.keeps(name)} - return {**row, "metadata": safe_dumps(retained)} + return {**row, "metadata": safe_dumps(retained, non_finite_as_null=True)} def _get_status_for_spend_log( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index d9d86cf54b3..a54d66c4beb 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1,4 +1,5 @@ import builtins +import math from collections.abc import Iterable, Mapping, Sequence from enum import Enum from os import PathLike @@ -83,6 +84,7 @@ from typing_extensions import ( override, ) +from litellm._logging import verbose_logger from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject, LiteLLMBaseModel from litellm.types.responses.main import ( CustomToolCallOutputItem, @@ -1446,9 +1448,11 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject): @classmethod def parse_cost(cls, v: object) -> object: """Normalise cost: accept either a float or a dict with a ``total_cost`` key.""" - if isinstance(v, dict): - return v.get("total_cost") - return v + cost: Final = v.get("total_cost") if isinstance(v, dict) else v + if isinstance(cost, float) and not math.isfinite(cost): + verbose_logger.warning("Dropping a non-finite provider usage cost") + return None + return cost model_config = ConfigDict(extra="allow") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 780c9ef6824..55aaf606175 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1,4 +1,5 @@ import json +import math import re import time from collections.abc import Collection, Mapping, Sequence @@ -1971,6 +1972,19 @@ class ServerToolUse(LiteLLMBaseModel): return getattr(self, key) +_COST_BREAKDOWN_ADAPTER: Final = TypeAdapter(Mapping[object, object]) + + +def _is_non_finite_float(value: object) -> bool: + return isinstance(value, float) and not math.isfinite(value) + + +def _carries_non_finite_number(cost: object) -> bool: + if not isinstance(cost, Mapping): + return _is_non_finite_float(cost) + return any(_is_non_finite_float(member) for member in _COST_BREAKDOWN_ADAPTER.validate_python(cost).values()) + + class Usage(SafeAttributeModel, CompletionUsage): _cache_creation_input_tokens: int = PrivateAttr( 0 @@ -1978,6 +1992,7 @@ class Usage(SafeAttributeModel, CompletionUsage): _cache_read_input_tokens: int = PrivateAttr( 0 ) # hidden param for prompt caching. Might change, once openai introduces their equivalent. + _dropped_non_finite_cost: bool = PrivateAttr(False) server_tool_use: ServerToolUse | None = None cost: float | None = None @@ -2084,10 +2099,14 @@ class Usage(SafeAttributeModel, CompletionUsage): else: # maintain openai compatibility in usage object if possible del self.server_tool_use - if cost is not None: - self.cost = cost - else: + dropped_non_finite_cost: Final = cost is not None and _carries_non_finite_number(cost) + if cost is None: del self.cost + elif dropped_non_finite_cost: + verbose_logger.warning("Dropping a non-finite provider usage cost") + del self.cost + else: + self.cost = cost ## ANTHROPIC MAPPING ## if "cache_creation_input_tokens" in params and isinstance(params["cache_creation_input_tokens"], int): @@ -2102,6 +2121,11 @@ class Usage(SafeAttributeModel, CompletionUsage): for k, v in params.items(): setattr(self, k, v) + self._dropped_non_finite_cost = dropped_non_finite_cost + + @staticmethod + def dropped_non_finite_cost(usage: "Usage") -> bool: + return usage._dropped_non_finite_cost def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator diff --git a/tests/integration/spend/test_non_finite_cost_spend_writes.py b/tests/integration/spend/test_non_finite_cost_spend_writes.py new file mode 100644 index 00000000000..9a4d2dd45a4 --- /dev/null +++ b/tests/integration/spend/test_non_finite_cost_spend_writes.py @@ -0,0 +1,166 @@ +import json +from hashlib import sha256 +from pathlib import Path +from typing import Final, Literal +from uuid import uuid4 + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.shares_worker_state + +PROMPT_TOKENS: Final = 5 +COMPLETION_TOKENS: Final = 3 +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +TOKEN_PRICE: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE +OVERFLOWING_INPUT_RATE: Final = 1e308 +POISONED_UPSTREAM: Final = "poisoned-upstream" +CONTROL_UPSTREAM: Final = "control-upstream" +DAILY_TOTALS: Final = "SELECT model_group, SUM(spend)::float8 AS spend, SUM(api_requests)::int8 AS api_requests" +DAILY_USER_QUERY: Final = f'{DAILY_TOTALS} FROM "LiteLLM_DailyUserSpend" WHERE api_key=%s GROUP BY model_group' +DAILY_TEAM_QUERY: Final = f'{DAILY_TOTALS} FROM "LiteLLM_DailyTeamSpend" WHERE api_key=%s GROUP BY model_group' +DAILY_TAG_QUERY: Final = f'{DAILY_TOTALS} FROM "LiteLLM_DailyTagSpend" WHERE api_key=%s AND tag=%s GROUP BY model_group' + +Poison = Literal["reported_nan_cost", "overflowing_price"] + + +def _frame(payload: dict[str, JsonValue]) -> bytes: + return b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def _stream(upstream_model: str, reported_cost: float | None) -> tuple[bytes, ...]: + chunk: Final = {"id": f"chatcmpl-{uuid4().hex}", "object": "chat.completion.chunk", "created": 1} + usage: Final[dict[str, JsonValue]] = { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + } + reported_usage: Final = usage if reported_cost is None else {**usage, "cost": reported_cost} + return ( + _frame({**chunk, "model": upstream_model, "choices": [{"index": 0, "delta": {"content": "ok"}}]}), + _frame( + { + **chunk, + "model": upstream_model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": reported_usage, + } + ), + b"data: [DONE]\n\n", + ) + + +def _deployment(alias: str, upstream_model: str, wire_url: str, input_rate: float) -> dict[str, JsonValue]: + return { + "model_name": alias, + "litellm_params": { + "model": f"openai/{upstream_model}", + "api_base": f"{wire_url}/v1", + "api_key": "integration-provider-key", + "input_cost_per_token": input_rate, + "output_cost_per_token": OUTPUT_RATE, + }, + } + + +def _proxy_config(directory: Path, deployments: tuple[dict[str, JsonValue], ...]) -> Path: + config: Final = directory / "non_finite_cost_config.yaml" + config.write_text( + json.dumps( + { + "model_list": list(deployments), + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_spend_logs": False, + "proxy_batch_write_at": 1, + "proxy_batch_polling_interval": 1, + }, + "router_settings": {"disable_cooldowns": True}, + } + ) + ) + return config + + +def _by_model_group(rows: list[dict[str, JsonValue]]) -> dict[str, float]: + return {str(row["model_group"]): float(str(row["spend"])) for row in rows} + + +def _both_requests_rows(query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]: + return eventually(lambda: read_rows(query, parameters), lambda rows: len(rows) == 2, seconds=70) + + +@pytest.mark.parametrize("poison", ["reported_nan_cost", "overflowing_price"]) +@pytest.mark.timeout(240) +def test_a_non_finite_cost_never_blocks_the_next_requests_spend_rows( + gateway: Gateway, tmp_path: Path, poison: Poison +) -> None: + poisoned_alias: Final = f"non-finite-cost-{uuid4().hex}" + control_alias: Final = f"finite-cost-{uuid4().hex}" + tag: Final = f"non-finite-cost-{uuid4().hex}" + poisoned_spend: Final = TOKEN_PRICE if poison == "reported_nan_cost" else 0.0 + poisoned_input_rate: Final = INPUT_RATE if poison == "reported_nan_cost" else OVERFLOWING_INPUT_RATE + + def respond(request: Request) -> Reply: + assert request.target.endswith("/chat/completions"), request.target + upstream_model: Final = json.loads(request.body)["model"] + reported_cost: Final = ( + float("nan") if upstream_model == POISONED_UPSTREAM and poison == "reported_nan_cost" else None + ) + return Reply(content_type="text/event-stream", chunks=_stream(upstream_model, reported_cost)) + + with wire_server(answering_model_discovery(respond)) as wire: + config: Final = _proxy_config( + tmp_path, + ( + _deployment(poisoned_alias, POISONED_UPSTREAM, wire.url, poisoned_input_rate), + _deployment(control_alias, CONTROL_UPSTREAM, wire.url, INPUT_RATE), + ), + ) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + team: Final = scenario.team() + key: Final = scenario.key(user_id=f"non-finite-cost-{uuid4().hex}", team_id=team) + digest: Final = sha256(key.encode()).hexdigest() + responses: Final = tuple( + candidate.request( + "POST", + "/v1/chat/completions", + { + "model": alias, + "messages": [{"role": "user", "content": f"non-finite cost {uuid4().hex}"}], + "stream": True, + "metadata": {"tags": [tag]}, + }, + key=key, + ) + for alias in (poisoned_alias, control_alias) + ) + assert all(response.status_code == 200 for response in responses), [response.text for response in responses] + + expected: Final = {poisoned_alias: pytest.approx(poisoned_spend), control_alias: pytest.approx(TOKEN_PRICE)} + spend_logs: Final = _both_requests_rows( + 'SELECT model_group, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,) + ) + assert _by_model_group(spend_logs) == expected, spend_logs + daily_rows: Final = ( + _both_requests_rows(DAILY_USER_QUERY, (digest,)), + _both_requests_rows(DAILY_TEAM_QUERY, (digest,)), + _both_requests_rows(DAILY_TAG_QUERY, (digest, tag)), + ) + assert [_by_model_group(rows) for rows in daily_rows] == [expected] * 3, daily_rows + assert [{row["api_requests"] for row in rows} for rows in daily_rows] == [{1}] * 3, daily_rows + key_total: Final = poisoned_spend + TOKEN_PRICE + key_spend: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) == pytest.approx(key_total), + seconds=70, + return_last_on_timeout=True, + ) + assert float(str(key_spend[0]["spend"])) == pytest.approx(key_total), key_spend diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index de8d6a7a8c5..d0f5bfc6b9f 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -3089,6 +3089,50 @@ def test_convert_to_model_response_object_function_output(): assert result._hidden_params is not None +@pytest.mark.parametrize( + "reported_cost", + [ + float("nan"), + float("inf"), + {"input_tokens_cost": 0.0001, "request_cost": 0.005, "total_cost": float("nan")}, + {"input_tokens_cost": float("inf"), "request_cost": 0.005, "total_cost": 0.0051}, + ], + ids=["nan", "inf", "breakdown-with-nan-total", "breakdown-with-inf-part"], +) +def test_a_non_finite_provider_usage_cost_is_dropped_so_the_response_is_strict_json(reported_cost: object): + response: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-non-finite-cost", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4, "cost": reported_cost}, + }, + model_response_object=ModelResponse(), + ) + + usage: Final = json.loads(json.dumps(response.model_dump(), allow_nan=False))["usage"] + assert "cost" not in usage + assert usage["total_tokens"] == 4 + + +def test_a_finite_provider_cost_breakdown_with_non_string_keys_is_kept(): + response: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-int-key-cost", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4, "cost": {1: 0.5}}, + }, + model_response_object=ModelResponse(), + ) + + assert response.usage.cost == {1: 0.5} + + _URL_CITATION: Final = { "type": "url_citation", "url_citation": { diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 1c86c4e30d8..93defc0f04f 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -26,15 +26,24 @@ from litellm.litellm_core_utils.core_helpers import ( normalize_drop_params, preserve_upstream_non_openai_attributes, process_response_headers, + provider_reported_usage_cost, qualify_agentic_followup_model, qualify_provider_stripped_model, reconstruct_model_name, redact_nested_match_and_regex_keys, + response_cost_in_hidden_params, set_provider_response_headers_in_hidden_params, unbind_budget_reservation_from_callbacks, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.types.utils import Delta, ImageResponse, ModelResponseStream, StreamingChoices, TranscriptionResponse +from litellm.types.utils import ( + Delta, + ImageResponse, + ModelResponse, + ModelResponseStream, + StreamingChoices, + TranscriptionResponse, +) def test_parent_otel_span_reads_dict_like_metadata_by_index(): @@ -790,6 +799,23 @@ class TestReceiveStampSurvivesASecondMetadataBucket: assert get_litellm_metadata_from_kwargs(kwargs)["litellm_received_at"] is own +@pytest.mark.parametrize("reported", [float("nan"), float("inf"), {"total_cost": float("nan")}]) +def test_a_non_finite_provider_reported_cost_counts_as_unreported(reported: object) -> None: + assert provider_reported_usage_cost({"cost": reported}, "openrouter") is None + assert provider_reported_usage_cost({"cost": {"total_cost": 0.25}}, "openrouter") == 0.25 + + +@pytest.mark.parametrize("header_cost", [float("nan"), float("inf")]) +def test_a_non_finite_response_cost_header_counts_as_unreported(header_cost: float) -> None: + response: Final = ModelResponse() + response._hidden_params = {"additional_headers": {RESPONSE_COST_HEADER: header_cost}} + finite: Final = ModelResponse() + finite._hidden_params = {"additional_headers": {RESPONSE_COST_HEADER: 0.25}} + + assert response_cost_in_hidden_params(response) is None + assert response_cost_in_hidden_params(finite) == 0.25 + + def test_billed_usage_cost_bills_any_chat_upstream_without_its_own_reader() -> None: from types import SimpleNamespace diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index b3ae3bf6240..336e9f7d24a 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -6308,6 +6308,60 @@ def test_standard_logging_payload_resolves_kubernetes_pod_routing_from_metadata_ assert payload["metadata"][KUBERNETES_POD_ROUTING_KEY] == expected +@pytest.mark.parametrize("response_cost", [float("nan"), float("inf"), float("-inf")]) +def test_get_standard_logging_object_payload_logs_a_non_finite_response_cost_as_zero( + logging_obj: LitellmLogging, response_cost: float +) -> None: + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload + + now: Final = datetime(2026, 1, 1) + payload: Final = get_standard_logging_object_payload( + kwargs={"model": "gpt-4o", "messages": [], "litellm_params": {}, "response_cost": response_cost}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["response_cost"] == 0 + + +def test_an_overflowing_custom_price_records_a_zero_response_cost(caplog: pytest.LogCaptureFixture) -> None: + messages: Final = [{"role": "user", "content": "hello world"}] + start_time: Final = datetime.datetime.now() + logging_obj: Final = Logging( + model="gpt-4o", + messages=messages, + stream=False, + call_type="completion", + start_time=start_time, + litellm_call_id="overflow\nforged log line", + function_id="overflow", + ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + response: Final = litellm.completion( + model="gpt-4o", + mock_response="hi there", + messages=messages, + input_cost_per_token=1e308, + output_cost_per_token=1e308, + litellm_logging_obj=logging_obj, + ) + logging_obj._process_hidden_params_and_response_cost( + response, start_time=start_time, end_time=datetime.datetime.now() + ) + + assert response._hidden_params["response_cost"] == 0.0 + payload: Final = logging_obj.model_call_details["standard_logging_object"] + assert json.loads(json.dumps(payload, allow_nan=False, default=str))["response_cost"] == 0.0 + warnings: Final = [r.getMessage() for r in caplog.records if "non-finite response_cost" in r.getMessage()] + assert warnings == ["Recording a non-finite response_cost as 0 for litellm_call_id=overflowforged log line"] + + def test_get_standard_logging_object_payload_carries_matched_access_groups(logging_obj): """Access groups stamped at auth time reach the logging payload, so integrations see what a request billed.""" from datetime import datetime diff --git a/tests/unit/litellm_core_utils/test_safe_json_dumps.py b/tests/unit/litellm_core_utils/test_safe_json_dumps.py index 1f2664f33cb..1b8bdd21326 100644 --- a/tests/unit/litellm_core_utils/test_safe_json_dumps.py +++ b/tests/unit/litellm_core_utils/test_safe_json_dumps.py @@ -169,6 +169,31 @@ def test_null_byte_stripped_in_fallback_str(): assert json.loads(out)["obj"] == "objrepr" +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_non_finite_floats_are_written_as_json_null(value: float) -> None: + out = safe_dumps( + {"cost_breakdown": {"total_cost": value, "output_cost": 0.006}, "costs": [value]}, non_finite_as_null=True + ) + + assert json.loads(out, parse_constant=lambda token: pytest.fail(f"non-JSON token {token}")) == { + "cost_breakdown": {"total_cost": None, "output_cost": 0.006}, + "costs": [None], + } + + +def test_a_set_holding_a_non_finite_float_sorts_its_finite_members_and_puts_null_last() -> None: + assert json.loads(safe_dumps({"costs": {2.0, float("inf"), 1.0}}, non_finite_as_null=True)) == { + "costs": [1.0, 2.0, None] + } + + +def test_non_finite_floats_keep_their_value_unless_the_caller_asks_for_null() -> None: + assert json.loads(safe_dumps({"timeout": float("inf"), "costs": [float("-inf")]})) == { + "timeout": float("inf"), + "costs": [float("-inf")], + } + + def test_clean_strings_are_not_run_through_replace(): """Regression for LIT-3910. diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 28587bf4749..c8fa9c94c17 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1077,6 +1077,47 @@ def test_cost_field_in_usage_chunks(): assert usage.completion_tokens == 5 + +def _perplexity_usage_chunk(cost: Mapping[str, float], finish_reason: str | None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-pplx", + created=1745513206, + model="sonar", + choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=Delta(content="ok"))], + usage={"prompt_tokens": 26, "completion_tokens": 1, "total_tokens": 27, "cost": dict(cost)}, + ) + + +@pytest.mark.parametrize("poisoned_field", ["request_cost", "total_cost"]) +def test_a_non_finite_final_stream_cost_does_not_fall_back_to_an_earlier_chunk_cost(poisoned_field: str): + final_cost: Final = {"input_tokens_cost": 0.0012, "output_tokens_cost": 1e-05, "request_cost": 0.0025} + chunks: Final = [ + _perplexity_usage_chunk({"input_tokens_cost": 0, "output_tokens_cost": 0, "total_cost": 0}, None), + _perplexity_usage_chunk({**final_cost, "total_cost": 0.00371, poisoned_field: float("nan")}, "stop"), + ] + + response: Final = stream_chunk_builder(chunks=chunks) + + assert response is not None + assert getattr(response.usage, "cost", None) is None + assert response.usage.total_tokens == 27 + + +@pytest.mark.parametrize("provider_key", ["dropped_non_finite_cost", "_dropped_non_finite_cost"]) +def test_a_provider_usage_key_named_like_the_dropped_cost_flag_keeps_the_stream_cost(provider_key: str): + chunk: Final = ModelResponseStream( + id="chatcmpl-flag-key", + created=1745513206, + model="sonar", + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="ok"))], + usage={"prompt_tokens": 26, "completion_tokens": 1, "total_tokens": 27, "cost": 0.00371, provider_key: True}, + ) + + response: Final = stream_chunk_builder(chunks=[chunk]) + + assert response is not None + assert getattr(response.usage, "cost", None) == 0.00371 + def test_stream_chunk_builder_tolerates_trailing_chunk_without_choices(): """Regression for https://github.com/BerriAI/litellm/issues/32051 diff --git a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py index 59355dfabcb..888c47505a4 100644 --- a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py +++ b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py @@ -1,6 +1,7 @@ import json from typing import Final +from unittest.mock import MagicMock import httpx import pytest @@ -464,6 +465,34 @@ def test_openrouter_cost_tracking_non_streaming(): ) + +@pytest.mark.parametrize("cost", [float("nan"), float("inf")]) +def test_openrouter_chat_leaves_a_non_finite_provider_cost_out_of_the_response_headers(cost: float) -> None: + body: Final = { + "id": "gen-123", + "object": "chat.completion", + "created": 1234567890, + "model": "openai/gpt-5.6-luna", + "choices": [{"message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop", "index": 0}], + "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11, "cost": cost}, + } + + result: Final = OpenrouterConfig().transform_response( + model="openrouter/openai/gpt-5.6-luna", + raw_response=httpx.Response( + 200, content=json.dumps(body).encode(), headers={"content-type": "application/json"} + ), + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "Reply with ok"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert "llm_provider-x-litellm-response-cost" not in result._hidden_params.get("additional_headers", {}) + def test_openrouter_cost_tracking_streaming(): """ Test OpenRouter cost tracking for streaming completions. diff --git a/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py b/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py index 19517e58b31..78d7161c388 100644 --- a/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py +++ b/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py @@ -9,6 +9,8 @@ reasoning.encrypted_content for multi-turn stateless workflows. Related issue: https://github.com/BerriAI/litellm/issues/22189 """ +import json +from typing import Final from unittest.mock import MagicMock import httpx @@ -157,6 +159,29 @@ def test_openrouter_responses_returns_provider_reported_cost(cost): assert get_response_cost_from_hidden_params(response.hidden_params) == cost +@pytest.mark.parametrize("cost", [float("nan"), float("inf")]) +def test_openrouter_responses_drops_a_non_finite_provider_cost(cost: float) -> None: + body: Final = { + "id": "resp_test", + "created_at": 1, + "model": "unmapped-model", + "object": "response", + "output": [], + "status": "completed", + "usage": {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5, "cost": cost}, + } + response: Final = OpenRouterResponsesAPIConfig().transform_response_api_response( + model="openrouter/unmapped-model", + raw_response=httpx.Response( + 200, content=json.dumps(body).encode(), headers={"content-type": "application/json"} + ), + logging_obj=MagicMock(), + ) + + assert get_response_cost_from_hidden_params(response.hidden_params) is None + assert json.loads(json.dumps(response.model_dump(), allow_nan=False, default=str))["usage"]["cost"] is None + + @pytest.mark.parametrize( "event_type", ["response.completed", "response.incomplete", "response.failed"], diff --git a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py index 83c71479311..380f85445e5 100644 --- a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py @@ -6,9 +6,11 @@ search queries, and reasoning tokens. """ import json +import logging import math import os from datetime import datetime, timezone +from typing import Final from unittest.mock import patch import pytest @@ -140,6 +142,26 @@ class TestPerplexityCostCalculator: assert prompt_cost == 0.0 assert completion_cost == 0.008 + @pytest.mark.parametrize( + "reported_cost", + [float("nan"), float("inf"), {"total_cost": float("nan")}], + ids=["nan", "inf", "nan-total-cost-dict"], + ) + def test_a_non_finite_perplexity_cost_is_priced_like_no_reported_cost_with_a_warning( + self, reported_cost: float | dict, caplog: pytest.LogCaptureFixture + ): + unreported: Final = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + reported: Final = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + reported.cost = reported_cost + + token_priced: Final = perplexity_cost_per_token(model="sonar-pro", usage=unreported) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + reported_priced: Final = perplexity_cost_per_token(model="sonar-pro", usage=reported) + + assert sum(token_priced) > 0 + assert reported_priced == token_priced + assert [record for record in caplog.records if "non-finite provider-reported cost" in record.getMessage()] + OFF_PEAK_MODEL = "sonar-off-peak-test" OFF_PEAK_WINDOW = "14:00-00:00" INSIDE_WINDOW = datetime(2026, 9, 3, 17, 25, tzinfo=timezone.utc) diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 24c12e06c9c..65da5630f33 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -1,10 +1,12 @@ import json +import logging from datetime import datetime, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm._logging import verbose_proxy_logger from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DailyTagSpendTransaction @@ -379,6 +381,104 @@ async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_fail assert restored_spend["org_member_list_transactions"] == {member_key: 1.5} +@pytest.mark.asyncio +async def test_a_single_parked_batch_with_a_nan_spend_drains_as_zero( + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock +): + mock_redis_cache.async_lpop = AsyncMock( + return_value=json.dumps({"key_list_transactions": {"nan-key": float("nan"), "ok-key": 0.5}}, allow_nan=True) + ) + + drained: Final = await redis_update_buffer.get_all_update_transactions_from_redis_buffer() + + assert drained is not None + assert drained["key_list_transactions"] == {"nan-key": 0.0, "ok-key": 0.5} + + +PARKED_DAILY_KEY: Final = "user-1_2025-01-15_api-key-123_gpt-4o_openai_/chat/completions" + + +def _parked_daily_batch(spend: float) -> str: + return json.dumps( + { + PARKED_DAILY_KEY: { + "spend": spend, + "prompt_tokens": 12, + "completion_tokens": 7, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + }, + allow_nan=True, + ) + + +@pytest.mark.asyncio +async def test_a_nan_spend_parked_by_an_older_version_drains_from_the_pipeline_as_zero( + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock +): + entity_batch: Final = json.dumps({"key_list_transactions": {"hashed-key": float("nan")}}, allow_nan=True) + window_batches: Final = tuple( + json.dumps( + [ + { + "entity_type": "key", + "entity_id": "hashed-key", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": spend, + "started_at": None, + } + ], + allow_nan=True, + ) + for spend in (float("nan"), 3.0) + ) + mock_redis_cache.async_lpop_pipeline = AsyncMock( + return_value=[ + [entity_batch], + [_parked_daily_batch(float("nan")), _parked_daily_batch(0.0042)], + None, + None, + None, + None, + list(window_batches), + ] + ) + + ( + db_spend, + daily_user, + *_rest, + window_spend, + ) = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + + assert db_spend is not None + assert db_spend["key_list_transactions"] == {"hashed-key": 0} + assert daily_user is not None + assert daily_user[PARKED_DAILY_KEY]["spend"] == 0.0042 + assert daily_user[PARKED_DAILY_KEY]["api_requests"] == 2 + assert window_spend is not None + assert [transaction["spend"] for transaction in window_spend] == [3.0] + + +@pytest.mark.asyncio +async def test_a_nan_spend_parked_by_an_older_version_drains_from_a_daily_key_as_zero( + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, caplog: pytest.LogCaptureFixture +): + mock_redis_cache.async_lpop = AsyncMock( + return_value=[_parked_daily_batch(float("inf")), _parked_daily_batch(0.0042)] + ) + + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + daily_tag: Final = await redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + + assert daily_tag is not None + assert daily_tag[PARKED_DAILY_KEY]["spend"] == 0.0042 + assert [record.getMessage() for record in caplog.records if "Infinity" in record.getMessage()] != [] + + @pytest.mark.asyncio async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): """When redis_cache is None, should return all Nones""" @@ -724,6 +824,10 @@ class _ListRedis: def __init__(self) -> None: self.rows: list[str] = [] + async def async_rpush(self, key: str, values: list[str]) -> int: + self.rows.extend(values) + return len(self.rows) + async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int: self.rows.extend(values) pushed_len = len(self.rows) @@ -752,6 +856,32 @@ async def test_store_spend_logs_in_redis_drops_oldest_rows_past_the_cap(): assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == () +@pytest.mark.asyncio +@pytest.mark.parametrize("spend", [float("nan"), float("inf")]) +async def test_a_non_finite_daily_spend_buffered_in_redis_drains_as_zero(spend: float): + buffer: Final = RedisUpdateBuffer(redis_cache=_ListRedis()) + queue: Final = DailySpendUpdateQueue() + await queue.add_update(json.loads(_parked_daily_batch(spend))) + await queue.add_update(json.loads(_parked_daily_batch(0.0042))) + + await buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue) + drained: Final = await buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + + assert drained is not None + assert drained[PARKED_DAILY_KEY]["spend"] == 0.0 + assert drained[PARKED_DAILY_KEY]["api_requests"] == 2 + + +@pytest.mark.asyncio +async def test_a_spend_log_row_parked_with_a_nan_spend_is_read_back_with_zero_spend(): + buffer: Final = RedisUpdateBuffer(redis_cache=_ListRedis()) + buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True) + + assert await buffer.store_spend_logs_in_redis([{"request_id": "nan-row", "spend": float("nan")}]) is True + + assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == ({"request_id": "nan-row", "spend": 0.0},) + + @pytest.mark.asyncio async def test_store_spend_logs_in_redis_reports_failure_without_redis(): buffer = RedisUpdateBuffer(redis_cache=None) diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index c3a93fd89be..298b153f621 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -2226,6 +2226,60 @@ async def test_shutdown_drain_finishes_the_rollups_update_database_started(): assert (team_row["api_requests"], team_row["spend"]) == (1, 0.1) +@pytest.mark.asyncio +@pytest.mark.parametrize("response_cost", [float("nan"), float("inf"), float("-inf")]) +async def test_update_database_records_a_non_finite_response_cost_as_zero(response_cost: float) -> None: + db_writer: Final = DBSpendUpdateWriter() + prisma: Final = _tool_usage_prisma() + prisma.get_request_status = MagicMock(return_value="success") + response: Final = litellm.ModelResponse( + model="gpt-5.6-luna", + usage=litellm.Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={ + "model": "gpt-5.6-luna", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "req-non-finite-cost", + "response_cost": response_cost, + "litellm_params": {"metadata": {"user_api_key": "hashed-key", "user_api_key_team_id": "test-team"}}, + }, + completion_response=response, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=response_cost, + ) + await drain_logging_tasks() + + assert [row["spend"] for row in prisma.spend_log_transactions] == [0] + spend: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert spend["key_list_transactions"] == {"test-token": 0} + assert spend["team_list_transactions"] == {"test-team": 0} + daily_user: Final = ( + await db_writer.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + daily_team: Final = ( + await db_writer.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + assert daily_user is not None and daily_team is not None + (user_row,) = daily_user.values() + (team_row,) = daily_team.values() + assert (user_row["api_requests"], user_row["spend"]) == (1, 0) + assert (team_row["api_requests"], team_row["spend"]) == (1, 0) + + @pytest.mark.asyncio async def test_shutdown_drain_waits_for_an_in_flight_rollup_enqueue(): """ diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py index 4011bc0294a..7f8aac983f5 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py @@ -323,6 +323,35 @@ def test_openrouter_decisions_response_uses_provider_cost_without_cost_map_row( assert response["kwargs"]["combined_usage_object"].total_tokens == 302 +@pytest.mark.parametrize("reported_cost", [float("nan"), float("inf")]) +def test_openrouter_non_finite_cost_logs_the_token_price(monkeypatch: pytest.MonkeyPatch, reported_cost: float) -> None: + monkeypatch.setitem( + litellm.model_cost, + "openrouter/typesafe/jev-1.13", + {"input_cost_per_token": 1.0, "output_cost_per_token": 3.0, "litellm_provider": "openrouter", "mode": "chat"}, + ) + logging_obj: Final = _logging_obj() + + result: Final = TypeSafePassthroughLoggingHandler.typesafe_passthrough_handler( + httpx_response=_response(), + response_body={ + "model": "typesafe/jev-1.13-20260917", + "usage": {"input_tokens": 2, "output_tokens": 1, "cost": reported_cost}, + }, + logging_obj=logging_obj, + url_route="https://openrouter.ai/api/alpha/decisions", + result='{"answers": {}}', + start_time=datetime(2000, 1, 1), + end_time=datetime(2000, 1, 1), + cache_hit=False, + request_body={"model": "typesafe/jev-1.13"}, + custom_llm_provider="openrouter", + ) + + assert result["kwargs"]["response_cost"] == pytest.approx(5.0) + assert logging_obj.model_call_details["response_cost"] == pytest.approx(5.0) + + def test_success_handler_dispatches_openrouter_to_the_shared_handler(): logging_obj = _logging_obj() normalized = PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( diff --git a/tests/unit/proxy/spend_tracking/test_savings.py b/tests/unit/proxy/spend_tracking/test_savings.py index f9566fd1ba0..8503a124b01 100644 --- a/tests/unit/proxy/spend_tracking/test_savings.py +++ b/tests/unit/proxy/spend_tracking/test_savings.py @@ -137,6 +137,34 @@ def test_compression_savings_priced_at_input_rate(): assert result.prompt_caching == 0.0 +def test_an_overflowing_deployment_price_records_finite_savings() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "overflowing-sonnet", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "input_cost_per_token": 1e308, + "cache_read_input_token_cost": 1e-07, + }, + }, + ] + ) + deployment_id: Final = router.get_model_list(model_name="overflowing-sonnet")[0]["model_info"]["id"] + + result: Final = compute_savings_spend( + model="claude-sonnet-5-5", + custom_llm_provider="anthropic", + compression_saved_tokens=4389, + gateway_injected_cache=True, + usage_object=_caching_usage(read=1000, written=0), + model_id=deployment_id, + llm_router=lambda: router, + ) + + assert result == (0.0, 0.0, 0.0, 0.0) + + def test_prompt_caching_savings_priced_at_input_minus_cache_read(): input_cost, cache_read_cost = _anthropic_costs("claude-sonnet-5") # A model that supports prompt caching must charge less to read from cache; diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 2db8941d304..a7596b51e3e 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -12496,6 +12496,55 @@ class TestResponseCostFromLoggingObjForCacheHits: ) assert {name: headers.get(name) for name in cost_headers} == dict.fromkeys(cost_headers, expected_header) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("response_cost", "expected_header"), + ((float("nan"), None), (float("inf"), None), (0.0042, "0.0042"), (2, "2")), + ids=("nan", "inf", "finite", "int"), + ) + async def test_a_non_finite_response_cost_sends_no_cost_header( + self, response_cost: float, expected_header: str | None + ): + response: Final = litellm.ModelResponse() + response._hidden_params = {"response_cost": response_cost} + + headers: Final = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=response, + request_data={"model": "gpt-5.4-mini"}, + request=Request({"type": "http", "method": "POST", "headers": []}), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + logging_obj=self._logging_obj(), + version=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + ) + + assert headers.get("x-litellm-response-cost") == expected_header + + @pytest.mark.asyncio + async def test_an_overflowing_cost_breakdown_sends_no_infinite_cost_header(self): + response: Final = litellm.ModelResponse() + response._hidden_params = {"response_cost": 0.0} + logging_obj: Final = self._logging_obj() + logging_obj.cost_breakdown = { + "original_cost": float("inf"), + "input_cost": float("inf"), + "output_cost": 0.0012, + "total_cost": float("inf"), + } + + headers: Final = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=response, + request_data={"model": "gpt-5.4-mini"}, + request=Request({"type": "http", "method": "POST", "headers": []}), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + logging_obj=logging_obj, + version=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + ) + + cost_headers: Final = {name: value for name, value in headers.items() if name.startswith("x-litellm-response-cost")} + assert cost_headers == {"x-litellm-response-cost-output": "0.0012"} + def test_a_fresh_zero_cost_v1_messages_call_still_sends_no_cost_header(self): logging_obj: Final = self._logging_obj() logging_obj.model_call_details["response_cost"] = 0.0 diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index b0bf1cb879e..aaec486b60f 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -5900,6 +5900,28 @@ def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): assert response._hidden_params["response_cost"] == pytest.approx(0.5) +@pytest.mark.parametrize("reported_cost", [float("nan"), float("inf")]) +def test_stream_chunk_builder_bills_the_token_price_when_the_provider_cost_is_not_finite(reported_cost: float): + usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") + usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=reported_cost) + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello "), + _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() + ) + + assert response is not None + prompt_cost, completion_cost = litellm.cost_per_token(model="gpt-4o", usage_object=response.usage) + token_price: Final = prompt_cost + completion_cost + assert token_price > 0 + assert getattr(response.usage, "cost", None) == pytest.approx(token_price) + assert response._hidden_params["response_cost"] == pytest.approx(token_price) + + def test_completion_cost_on_a_stamped_stream_still_bills_the_reported_charge(): usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) From 5d5e3685c1d52841082c199a88f8190f35e7be9b Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:48:43 -0700 Subject: [PATCH 30/64] fix(logging): warn when drop_params drops a requested reasoning level on the remaining paths (#48933) * fix(logging): warn when drop_params drops a requested reasoning level on the remaining paths * fix(logging): keep guardrail and token-count reads of a request from logging drop warnings Only the real Anthropic-messages adapter request reports a dropped thinking budget. Guardrail, shadow-eval and token-count translations of the same request now drop it silently, so one request logs one warning. The 422 retry also no longer hashes each loc entry, so a list inside loc cannot raise TypeError. * fix(bedrock): keep the DeepSeek R1 Converse reasoning_effort drop at debug level Only allowed_openai_params reaches that branch, and it puts reasoning_effort back after the provider mapping, so a warning there claimed a drop that never happened * fix(logging): warn on Gemini non-integer budget drops and keep the auto-router baseline silent Gemini 2.x now logs the drop of a non-integer thinking budget. The auto-router savings baseline builds a body it never sends, so it no longer warns. Logged values are capped with reprlib, Azure AI marks only its drop_params retry as gated, and Bedrock Chat Completions names the model without its route prefix. * fix(logging): warn on the Gemini 3, Magistral, Ollama, Fireworks and DeepSeek ARN reasoning drops Gemini 3 and later now log a dropped thinking budget, Magistral logs the level it swaps for a reasoning prompt, Ollama logs a level it maps to think off, Fireworks logs chat_template_kwargs effort keys dropped for a model the map does not list as reasoning, and DeepSeek on Bedrock logs a drop when the model is addressed by an ARN that advertises the param. * fix(logging): warn on the remaining Ollama, DeepSeek and Tencent reasoning drops The ollama/ completion route now shares the chat route's think mapping, so a level it turns off logs the same warning. DeepSeek warns when it drops thinking.budget_tokens, and Tencent warns when an adaptive-only model gets a non-integer budget. Magistral no longer warns on low, medium, high or enabled thinking, the values its reasoning prompt stands in for. The Gemini 3 budget warning no longer claims the model uses its default level. * fix(reasoning): warn when Claude drops a non-string reasoning_effort Claude chat and native /v1/messages dropped a reasoning_effort like 5 or {"effort": 5} with no log line. Both now log one warn-only WARNING. A Magistral thinking.type sent as a list no longer raises TypeError. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../reasoning_effort_utils.py | 77 ++++++++++- litellm/llms/anthropic/chat/transformation.py | 2 + litellm/llms/anthropic/common_utils.py | 33 +++-- .../messages/adapters/transformation.py | 18 ++- .../llms/anthropic/messages/native/handler.py | 8 +- .../messages/native/transformation.py | 11 +- .../responses_adapters/transformation.py | 4 +- .../llms/anthropic/prompt_cache_prediction.py | 8 +- litellm/llms/azure_ai/chat/transformation.py | 25 +++- .../chat/chat_completions/transformation.py | 17 +++ .../bedrock/chat/converse_transformation.py | 17 ++- litellm/llms/deepseek/chat/transformation.py | 13 ++ litellm/llms/fal_ai/chat/transformation.py | 13 +- .../llms/fireworks_ai/chat/transformation.py | 14 +- .../fireworks_ai/completion/transformation.py | 7 +- .../llms/hosted_vllm/chat/transformation.py | 2 +- litellm/llms/mistral/chat/transformation.py | 25 +++- litellm/llms/oci/chat/transformation.py | 16 +++ litellm/llms/ollama/chat/transformation.py | 8 +- litellm/llms/ollama/common_utils.py | 12 ++ .../llms/ollama/completion/transformation.py | 8 +- litellm/llms/openai/common_utils.py | 22 +-- litellm/llms/openai/openai.py | 8 +- litellm/llms/tencent/chat/transformation.py | 7 + .../vertex_and_google_ai_studio_gemini.py | 26 +++- .../reasoning_effort_capability.py | 8 +- .../test_anthropic_guardrail_handler.py | 16 +++ .../test_anthropic_chat_transformation.py | 42 ++++++ ...hropic_messages_adapters_transformation.py | 16 ++- ..._handler_reasoning_effort_normalization.py | 26 ++-- .../native/test_anthropic_messages_handler.py | 10 +- .../native/test_anthropic_messages_router.py | 17 ++- .../test_reasoning_effort_translation.py | 46 +++++++ .../test_responses_adapters_transformation.py | 33 +++-- .../messages/test_reasoning_effort_fields.py | 5 +- .../anthropic/test_anthropic_common_utils.py | 18 ++- .../test_anthropic_prompt_cache_prediction.py | 40 ++++++ .../chat/test_azure_ai_transformation.py | 61 +++++++++ ...bedrock_chat_completions_transformation.py | 52 ++++--- .../chat/test_converse_transformation.py | 47 +++++++ .../chat/test_deepseek_chat_transformation.py | 41 ++++++ .../chat/test_fal_ai_chat_transformation.py | 33 ++++- .../test_fireworks_ai_chat_transformation.py | 35 +++++ ...works_ai_text_completion_transformation.py | 37 +++++ .../test_mistral_chat_transformation.py | 105 +++++++++++++++ .../oci/chat/test_oci_chat_transformation.py | 41 +++--- .../ollama/test_ollama_chat_transformation.py | 35 +++++ .../test_ollama_completion_transformation.py | 39 ++++++ .../llms/openai/test_openai_common_utils.py | 127 +++++++++++++++++- .../chat/test_tencent_chat_transformation.py | 26 ++++ ...test_vertex_and_google_ai_studio_gemini.py | 72 ++++++++++ tests/unit/test_main.py | 36 ++++- 52 files changed, 1318 insertions(+), 147 deletions(-) diff --git a/litellm/litellm_core_utils/reasoning_effort_utils.py b/litellm/litellm_core_utils/reasoning_effort_utils.py index acd20aeb37d..e9c94e32f8f 100644 --- a/litellm/litellm_core_utils/reasoning_effort_utils.py +++ b/litellm/litellm_core_utils/reasoning_effort_utils.py @@ -1,6 +1,12 @@ +import reprlib +from collections.abc import Collection, Mapping +from dataclasses import dataclass from typing import Final, Literal +from pydantic import TypeAdapter + import litellm +from litellm._logging import verbose_logger from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, @@ -11,6 +17,54 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping OpenAIStyleReasoningEffort = Literal["minimal", "low", "medium", "high"] +_DROP_WARNING: Final = "Dropping %s=%s for model=%r because %s" +_DROP_PARAMS_DROP_WARNING: Final = "Dropping %s=%s for model=%r (drop_params=True) because %s" +_NON_NUMERIC_BUDGET_REASON: Final = "it is not a number. Send budget_tokens as an integer to set a reasoning level." +_NON_STRING_EFFORT_REASON: Final = ( + "it is not a string. Send reasoning_effort as a level name such as 'high' to set a reasoning level." +) +_REASONING_PARAMS: Final = ("reasoning", "reasoning_effort", "thinking") +OPAQUE_VALUE: Final = TypeAdapter(object) + + +def warn_dropped_reasoning( + param: str, value: object, model: str | None, reason: str, *, drop_params_gated: bool +) -> None: + verbose_logger.warning( + _DROP_PARAMS_DROP_WARNING if drop_params_gated else _DROP_WARNING, param, reprlib.repr(value), model, reason + ) + + +def warn_non_string_reasoning_effort(value: object, model: str | None) -> None: + effort: Final = as_str_mapping(value) + requested_level: Final = effort.get("effort") if effort is not None else value + if requested_level is None or isinstance(value, str): + return + warn_dropped_reasoning("reasoning_effort", value, model, _NON_STRING_EFFORT_REASON, drop_params_gated=False) + + +def warn_dropped_reasoning_params( + dropped: Collection[object], + request: Mapping[str, object], + model: str | None, + reason: str, + *, + drop_params_gated: bool, +) -> None: + for param in _REASONING_PARAMS: + if param in dropped and param in request: + warn_dropped_reasoning(param, request[param], model, reason, drop_params_gated=drop_params_gated) + + +@dataclass(frozen=True, slots=True) +class DroppedThinkingBudget: + budget_tokens: object + + def warn(self, model: str) -> None: + warn_dropped_reasoning( + "thinking.budget_tokens", self.budget_tokens, model, _NON_NUMERIC_BUDGET_REASON, drop_params_gated=True + ) + def reasoning_effort_from_thinking_budget( budget_tokens: float, @@ -29,7 +83,9 @@ def reasoning_effort_from_thinking_budget( return "minimal" -def reasoning_effort_from_thinking(thinking: object, *, drop_params: bool) -> OpenAIStyleReasoningEffort | None: +def _effort_or_dropped_budget( + thinking: object, *, drop_params: bool +) -> OpenAIStyleReasoningEffort | DroppedThinkingBudget | None: fields: Final = as_str_mapping(thinking) if fields is None or fields.get("type") != "enabled": return None @@ -37,10 +93,27 @@ def reasoning_effort_from_thinking(thinking: object, *, drop_params: bool) -> Op if isinstance(budget, (int, float)) and not isinstance(budget, bool): return reasoning_effort_from_thinking_budget(budget) if drop_params or litellm.drop_params is True: - return None + return DroppedThinkingBudget(budget) raise non_numeric_thinking_budget_error(budget) +def reasoning_effort_from_thinking( + thinking: object, *, model: str, drop_params: bool +) -> OpenAIStyleReasoningEffort | None: + effort: Final = _effort_or_dropped_budget(thinking, drop_params=drop_params) + if not isinstance(effort, DroppedThinkingBudget): + return effort + effort.warn(model) + return None + + +def reasoning_effort_from_thinking_silently( + thinking: object, *, drop_params: bool +) -> OpenAIStyleReasoningEffort | None: + effort: Final = _effort_or_dropped_budget(thinking, drop_params=drop_params) + return None if isinstance(effort, DroppedThinkingBudget) else effort + + def non_numeric_thinking_budget_error(budget: object) -> UnsupportedParamsError: return UnsupportedParamsError( message=( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 865247afc63..d773e31c733 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -48,6 +48,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import place_mid_conversation_system, split_leading_system_run, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_non_string_reasoning_effort from litellm.llms.base_llm.base_utils import ( ResponseFormatSchema, as_json_object, @@ -1951,6 +1952,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if isinstance(effort_value, dict): effort_value = effort_value.get("effort") if not isinstance(effort_value, str): + warn_non_string_reasoning_effort(OPAQUE_VALUE.validate_python(value), model) continue mapped_thinking = AnthropicConfig._map_reasoning_effort( reasoning_effort=effort_value, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 82977f97854..6d1ac24af55 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -37,7 +37,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of, ) -from litellm.litellm_core_utils.reasoning_effort_utils import non_numeric_thinking_budget_error +from litellm.litellm_core_utils.reasoning_effort_utils import DroppedThinkingBudget, non_numeric_thinking_budget_error from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.wif import ( aget_anthropic_wif_token, @@ -610,11 +610,11 @@ def _numeric_budget_tokens(budget_tokens: object) -> int | None: return None -def _legacy_budget_tokens(budget_tokens: object, *, drop_params: bool) -> int | None: +def _legacy_budget_tokens(budget_tokens: object, *, drop_params: bool) -> int | DroppedThinkingBudget: if (numeric := _numeric_budget_tokens(budget_tokens)) is not None: return numeric if drop_params or litellm.drop_params is True: - return None + return DroppedThinkingBudget(budget_tokens) raise non_numeric_thinking_budget_error(budget_tokens) @@ -1097,6 +1097,22 @@ class AnthropicModelInfo(BaseLLMModelInfo): custom_llm_provider: str, drop_params: bool, ) -> None: + dropped: Final = AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model_silently( + model=model, + optional_params=optional_params, + custom_llm_provider=custom_llm_provider, + drop_params=drop_params, + ) + if dropped is not None: + dropped.warn(model) + + @staticmethod + def translate_legacy_thinking_for_adaptive_model_silently( + model: str, + optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param like the sibling helpers + custom_llm_provider: str, + drop_params: bool, + ) -> DroppedThinkingBudget | None: """Translate legacy ``thinking.type=enabled`` to adaptive for the adaptive-thinking models that reject it (4.7+ and the 5 families). Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the @@ -1105,12 +1121,12 @@ class AnthropicModelInfo(BaseLLMModelInfo): ``output_config.effort`` is never overridden. """ if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider): - return + return None if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider): - return + return None thinking: Final = validated_object_mapping(optional_params.get("thinking")) if thinking is None or thinking.get("type") != "enabled": - return + return None budget_tokens: Final = _legacy_budget_tokens(thinking.get("budget_tokens", 0), drop_params=drop_params) display: Final = thinking.get("display") @@ -1119,8 +1135,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): thinking, rebuilt_keys=frozenset(("type", "budget_tokens", "display")), drop_params=drop_params ) optional_params["thinking"] = {**carried_keys, "type": "adaptive", **supported_display} - if budget_tokens is None: - return + if isinstance(budget_tokens, DroppedThinkingBudget): + return budget_tokens effort: Final = AnthropicModelInfo._legacy_budget_to_effort( model=model, budget_tokens=budget_tokens, custom_llm_provider=custom_llm_provider ) @@ -1129,6 +1145,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): "effort": effort, **(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})), } + return None @staticmethod def _legacy_budget_to_effort(model: str, budget_tokens: int, custom_llm_provider: str) -> str: diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index 8ea6293bbe6..bc9649ed867 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -134,6 +134,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( ) from litellm.litellm_core_utils.reasoning_effort_utils import ( reasoning_effort_from_thinking, + reasoning_effort_from_thinking_silently, ) from litellm.llms.anthropic.common_utils import ( eager_input_streaming_flag, @@ -297,6 +298,7 @@ class AnthropicAdapter: custom_llm_provider=custom_llm_provider, api_base=api_base, drop_params=drop_params, + report_reasoning_drops=True, ) return translated_body, tool_name_mapping @@ -666,7 +668,9 @@ class LiteLLMAnthropicMessagesAdapter: def translate_anthropic_thinking_to_reasoning_effort( thinking: AnthropicThinkingParam, *, + model: str, drop_params: bool = True, + report_reasoning_drops: bool = False, ) -> str | None: """ Translate Anthropic's thinking parameter to OpenAI's reasoning_effort. @@ -685,7 +689,9 @@ class LiteLLMAnthropicMessagesAdapter: if thinking_type == "disabled": return "none" elif thinking_type == "enabled": - return reasoning_effort_from_thinking(thinking, drop_params=drop_params) + if report_reasoning_drops: + return reasoning_effort_from_thinking(thinking, model=model, drop_params=drop_params) + return reasoning_effort_from_thinking_silently(thinking, drop_params=drop_params) elif thinking_type == "adaptive": # Adaptive thinking: effort is controlled by output_config.effort, # not budget_tokens. Return a default; caller should override with @@ -752,7 +758,7 @@ class LiteLLMAnthropicMessagesAdapter: return {"thinking": thinking} else: reasoning_effort: Final = LiteLLMAnthropicMessagesAdapter.translate_anthropic_thinking_to_reasoning_effort( - thinking + thinking, model=model ) if reasoning_effort: return { @@ -1147,6 +1153,7 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, drop_params: bool = True, + report_reasoning_drops: bool = False, ) -> None: """Translate Anthropic thinking to either thinking or reasoning_effort. @@ -1207,7 +1214,10 @@ class LiteLLMAnthropicMessagesAdapter: return reasoning_effort: Final = declared_effort or self.translate_anthropic_thinking_to_reasoning_effort( - cast(AnthropicThinkingParam, thinking), drop_params=drop_params + cast(AnthropicThinkingParam, thinking), + model=model, + drop_params=drop_params, + report_reasoning_drops=report_reasoning_drops, ) if not reasoning_effort: return @@ -1267,6 +1277,7 @@ class LiteLLMAnthropicMessagesAdapter: api_base: str | None = None, preserve_midturn_system: bool = False, drop_params: bool = True, + report_reasoning_drops: bool = False, ) -> tuple[ChatCompletionRequest, dict[str, str]]: """ This is used by the beta Anthropic Adapter, for translating anthropic `/v1/messages` requests to the openai format. @@ -1321,6 +1332,7 @@ class LiteLLMAnthropicMessagesAdapter: new_kwargs=new_kwargs, custom_llm_provider=custom_llm_provider, drop_params=drop_params, + report_reasoning_drops=report_reasoning_drops, ) ## CONVERT STOP_SEQUENCES self._translate_stop_sequences_to_openai( diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 6f97a019a96..1da7c2efdda 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.common_utils import ( + AnthropicModelInfo, strip_provider_specific_fields_from_anthropic_messages, tools_marked_for_a_tool_config_point, ) @@ -36,7 +37,6 @@ from litellm.router_utils.reasoning_effort_capability import ( keeps_effort, model_group_from_metadata, resolve_deployment_reasoning_effort, - with_adaptive_thinking_as_effort, ) from litellm.types.llms.anthropic import AllAnthropicToolsValues from litellm.types.llms.anthropic_messages.anthropic_request import AnthropicMetadata @@ -185,9 +185,11 @@ def _with_deployment_output_effort( ).value def translate(source: Mapping[str, object]) -> Mapping[str, object]: - return with_adaptive_thinking_as_effort( - source, model=model, custom_llm_provider=custom_llm_provider, drop_params=drop_params + translated: Final = dict(source) + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=model, optional_params=translated, custom_llm_provider=custom_llm_provider, drop_params=drop_params ) + return translated requested_output: Final = as_str_mapping(params.get("output_config")) caller_output: Final = None if requested_output is None else resolve(requested_output) diff --git a/litellm/llms/anthropic/messages/native/transformation.py b/litellm/llms/anthropic/messages/native/transformation.py index ef855eff180..c148a0defe8 100644 --- a/litellm/llms/anthropic/messages/native/transformation.py +++ b/litellm/llms/anthropic/messages/native/transformation.py @@ -21,6 +21,7 @@ from litellm.constants import ANTHROPIC_BILLING_METADATA_PREFIX from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_non_string_reasoning_effort from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.anthropic.messages.convertible_blocks import drop_non_string_tool_result_texts from litellm.llms.anthropic.messages.safeguards import proxy_serves_safeguards @@ -103,6 +104,7 @@ _MESSAGES_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): _workload_identity_eligible: ClassVar[bool] = True + _report_reasoning_drops: ClassVar[bool] = True @property def custom_llm_provider(self) -> str | None: @@ -695,12 +697,17 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ): anthropic_messages_optional_request_params.pop("output_format") # pyright: ignore[reportUnknownMemberType] # the request params arrive as a bare dict + requested_reasoning_effort: Final = OPAQUE_VALUE.validate_python( + anthropic_messages_optional_request_params.get("reasoning_effort") # pyright: ignore[reportUnknownMemberType] # the request params arrive as a bare dict + ) self._translate_reasoning_effort_to_anthropic( model=model, optional_params=anthropic_messages_optional_request_params, max_tokens=max_tokens, custom_llm_provider=self._resolved_provider, ) + if self._report_reasoning_drops: + warn_non_string_reasoning_effort(requested_reasoning_effort, model) AnthropicModelInfo.maybe_drop_disabled_thinking( model=model, @@ -708,12 +715,14 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): custom_llm_provider=self._resolved_provider, ) - AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + dropped_budget: Final = AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model_silently( model=model, optional_params=anthropic_messages_optional_request_params, custom_llm_provider=self._resolved_provider, drop_params=getattr(litellm_params, "drop_params", None) is True, ) + if dropped_budget is not None and self._report_reasoning_drops: + dropped_budget.warn(model) self._translate_adaptive_effort_for_non_adaptive_model( model=model, diff --git a/litellm/llms/anthropic/messages/responses_adapters/transformation.py b/litellm/llms/anthropic/messages/responses_adapters/transformation.py index 94c608f90e6..0630bef0311 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/messages/responses_adapters/transformation.py @@ -534,6 +534,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: output_config: dict[str, object] | None = None, auto_summary_level: str = "detailed", *, + model: str, drop_params: bool = True, ) -> dict[str, object] | None: """ @@ -555,7 +556,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(output_config, dict) and output_config.get("effort"): effort = output_config["effort"] elif thinking_type == "enabled": - budget_effort: Final = reasoning_effort_from_thinking(thinking, drop_params=drop_params) + budget_effort: Final = reasoning_effort_from_thinking(thinking, model=model, drop_params=drop_params) if budget_effort is None: return None effort = budget_effort @@ -667,6 +668,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: thinking, output_config=cast(dict[str, object] | None, output_config), auto_summary_level=auto_summary_level, + model=model, drop_params=drop_params, ) if reasoning: diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py index 27e003ce8b6..2da071ebdd2 100644 --- a/litellm/llms/anthropic/prompt_cache_prediction.py +++ b/litellm/llms/anthropic/prompt_cache_prediction.py @@ -9,7 +9,7 @@ from dataclasses import dataclass, field from functools import reduce from itertools import accumulate, groupby from types import MappingProxyType -from typing import Annotated, Final, Literal, Protocol, TypeAlias +from typing import Annotated, ClassVar, Final, Literal, Protocol, TypeAlias import httpx from pydantic import ConfigDict, Field, JsonValue, StrictInt, TypeAdapter, ValidationError @@ -81,6 +81,10 @@ _DEPLOYMENT_OPTIONS: Final = frozenset( ) +class _BaselineAnthropicMessagesConfig(AnthropicMessagesConfig): + _report_reasoning_drops: ClassVar[bool] = False + + class _StrictModel(LiteLLMBaseModel): model_config = ConfigDict(extra="forbid", frozen=True, strict=True) @@ -661,7 +665,7 @@ def prepare_native_baseline_body(request: Mapping[str, object], model: str) -> M TypeAdapter(tuple[str, ...]).validate_python(owned.get("additional_drop_params") or ()), dict(options), ) - config: Final = AnthropicMessagesConfig() + config: Final = _BaselineAnthropicMessagesConfig() body: Final = config.transform_anthropic_messages_request( model=resolved_model, messages=strip_provider_specific_fields_from_anthropic_messages(messages), diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 5381447f28f..5214d31fb2e 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( prompt_cache_breakpoint_for_wire, with_prompt_cache_breakpoint, ) +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning_params from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.azure_ai.common_utils import ( api_key_header_for_base, @@ -40,6 +41,14 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +_REQUEST_BODY: Final = TypeAdapter(dict[str, object]) + + +def _request_model(request: Mapping[str, object]) -> str | None: + model: Final = request.get("model") + return model if isinstance(model, str) else None + + class AzureFoundryErrorStrings(str, enum.Enum): SET_EXTRA_PARAMETERS_TO_PASS_THROUGH = "Set extra-parameters to 'pass-through'" @@ -510,7 +519,13 @@ class AzureAIStudioConfig(OpenAIConfig): request_data = self._drop_extra_params_from_request_data(request_data, error_text) if "Extra inputs are not permitted" in error_text and self._error_has_tool_level_extra_fields(error_text): request_data = self._drop_tool_level_extra_fields(request_data, error_text) - data: Final = drop_params_from_unprocessable_entity_error(e=e, data=request_data) + data: Final = drop_params_from_unprocessable_entity_error( + e=e, + data=request_data, + model=_request_model(_REQUEST_BODY.validate_python(request_data)), + drop_params_gated="Extra inputs are not permitted" in error_text + and not self._error_has_tool_level_extra_fields(error_text), + ) return data def _drop_tool_level_extra_fields(self, request_data: dict, error_text: str) -> dict: @@ -526,6 +541,14 @@ class AzureAIStudioConfig(OpenAIConfig): def _drop_extra_params_from_request_data(self, request_data: dict, error_text: str) -> dict: params_to_drop: Final = self._extract_params_to_drop_from_error_text(error_text) if params_to_drop: + request: Final = _REQUEST_BODY.validate_python(request_data) + warn_dropped_reasoning_params( + params_to_drop, + request, + _request_model(request), + "Azure AI rejected it as an extra parameter its deployment does not pass through.", + drop_params_gated=False, + ) for param in params_to_drop: if param in request_data: request_data.pop(param, None) diff --git a/litellm/llms/bedrock/chat/chat_completions/transformation.py b/litellm/llms/bedrock/chat/chat_completions/transformation.py index 123e4a9e6b4..dd2277a6c35 100644 --- a/litellm/llms/bedrock/chat/chat_completions/transformation.py +++ b/litellm/llms/bedrock/chat/chat_completions/transformation.py @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( inline_remote_image_urls, inline_remote_media, ) +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token @@ -42,6 +43,7 @@ from litellm.llms.bedrock.common_utils import ( bedrock_rejects_stop_sequences, bedrock_runtime_chat_completions_serves_reasoning_inline, split_bedrock_region_path, + strip_bedrock_routing_prefix, ) from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig @@ -121,6 +123,13 @@ def without_refused_reasoning_effort(model: str, params: Mapping[str, object]) - effort: Final = params.get("reasoning_effort") if not isinstance(effort, str) or effort not in chat_completions_reasoning_efforts_refused_for(model): return params + warn_dropped_reasoning( + "reasoning_effort", + effort, + strip_bedrock_routing_prefix(model), + "AWS's Chat Completions endpoint rejects that level for this model family, so AWS applies its default effort.", + drop_params_gated=False, + ) return _without_params(params, frozenset(("reasoning_effort",))) @@ -445,6 +454,14 @@ class AmazonBedrockRuntimeChatCompletionsConfig(OpenAILikeChatConfig): ), status_code=400, ) + if malformed_effort: + warn_dropped_reasoning( + "reasoning_effort", + raw_params["reasoning_effort"], + strip_bedrock_routing_prefix(model), + "Bedrock's Chat Completions endpoint takes it only as a string, so AWS applies its default effort.", + drop_params_gated=True, + ) rejected_stop: Final = frozenset(("stop",)) if bedrock_rejects_stop_sequences(model) else frozenset[str]() return dict( without_refused_reasoning_effort( diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 5e621a5f322..3800489ec05 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -50,6 +50,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, @@ -187,6 +188,8 @@ def _tools_the_model_accepts( _OUTPUT_CONFIG_BLOCK: Final = TypeAdapter(OutputConfigBlock) +_DEEPSEEK_THINKING_REASON: Final = "DeepSeek on Bedrock reasons natively and takes no thinking field." +_DEEPSEEK_R1_EFFORT_REASON: Final = "DeepSeek R1 on Bedrock always reasons and rejects reasoning_effort." def _has_native_text_format(output_config: object) -> bool: @@ -640,6 +643,10 @@ class AmazonConverseConfig(BaseConfig): or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) ) + def _warn_native_reasoning_drop(self, param: str, value: object, model: str, reason: str) -> None: + if param in self.get_supported_openai_params(model): + warn_dropped_reasoning(param, value, model, reason, drop_params_gated=False) + def _model_rejects_reasoning_effort_param(self, model: str, base_model: str) -> bool: """Whether the model returns a 400 for every ``reasoning_effort`` shape on Converse. @@ -1143,9 +1150,8 @@ class AmazonConverseConfig(BaseConfig): "tool_choice": {"type": "auto", "disable_parallel_tool_use": not value} } if param == "thinking" and drop_thinking_param: - verbose_logger.debug( - "Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively.", - model, + self._warn_native_reasoning_drop( + param, OPAQUE_VALUE.validate_python(value), model, _DEEPSEEK_THINKING_REASON ) elif param == "thinking" and not self._is_openai_gpt_reasoning_model(model): if ( @@ -1179,10 +1185,7 @@ class AmazonConverseConfig(BaseConfig): drop_params=drop_params, ) elif param == "reasoning_effort" and isinstance(value, str) and drop_reasoning_effort_param: - verbose_logger.debug( - "Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it always reasons and rejects it.", - model, - ) + self._warn_native_reasoning_drop(param, value, model, _DEEPSEEK_R1_EFFORT_REASON) elif ( param == "reasoning_effort" and isinstance(value, str) diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index ecaa0a2ef16..5b7c2a2c410 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -6,10 +6,12 @@ from collections.abc import Coroutine, Mapping, Sequence from typing import Final, Literal, cast, overload import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, extract_search_results_text, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -18,6 +20,16 @@ from litellm.utils import supports_reasoning, supports_vision from ...openai.chat.gpt_transformation import OpenAIGPTConfig SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "xhigh", "max", "ultra") +_NO_BUDGET_REASON: Final = "DeepSeek takes thinking only as enabled or disabled, with no budget." + + +def _warn_on_dropped_budget(thinking: object, model: str) -> None: + fields: Final = as_str_mapping(thinking) + if fields is None or fields.get("type") != "enabled" or "budget_tokens" not in fields: + return + warn_dropped_reasoning( + "thinking.budget_tokens", fields["budget_tokens"], model, _NO_BUDGET_REASON, drop_params_gated=False + ) class DeepSeekChatConfig(OpenAIGPTConfig): @@ -53,6 +65,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): # Then re-add only if valid for DeepSeek thinking_value: Final = optional_params.pop("thinking", None) reasoning_effort: Final = optional_params.pop("reasoning_effort", None) + _warn_on_dropped_budget(OPAQUE_VALUE.validate_python(thinking_value), model) # Handle thinking parameter - accept both enabled and disabled, ignore budget_tokens if isinstance(thinking_value, dict) and thinking_value.get("type") in ("enabled", "disabled"): diff --git a/litellm/llms/fal_ai/chat/transformation.py b/litellm/llms/fal_ai/chat/transformation.py index 115b1f840de..848ae498cdd 100644 --- a/litellm/llms/fal_ai/chat/transformation.py +++ b/litellm/llms/fal_ai/chat/transformation.py @@ -12,7 +12,9 @@ from typing import TYPE_CHECKING, Final import httpx from pydantic import ConfigDict, TypeAdapter +import litellm from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.base import LiteLLMBaseModel @@ -29,6 +31,10 @@ PROVIDER_PREFIX: Final[str] = "fal_ai/" PASSTHROUGH_PARAMS: Final[frozenset[str]] = frozenset(("reasoning", "temperature", "top_p")) REASONING_DISABLED_EFFORTS: Final[frozenset[str]] = frozenset(("none", "minimal")) REASONING_ENABLED_EFFORTS: Final[frozenset[str]] = frozenset(("low", "medium", "high")) +_UNSUPPORTED_EFFORT_REASON: Final = ( + f"fal only turns reasoning off for {sorted(REASONING_DISABLED_EFFORTS)} " + f"and on for {sorted(REASONING_ENABLED_EFFORTS)}." +) class _FalUsage(LiteLLMBaseModel): @@ -120,9 +126,10 @@ class FalAIChatConfig(BaseConfig): return False if isinstance(value, str) and value in REASONING_ENABLED_EFFORTS: return True - if drop_params: - return None - raise FalAIError(status_code=400, message=f"Unsupported reasoning_effort {value!r} for {model}") + if not (drop_params or litellm.drop_params is True): + raise FalAIError(status_code=400, message=f"Unsupported reasoning_effort {value!r} for {model}") + warn_dropped_reasoning("reasoning_effort", value, model, _UNSUPPORTED_EFFORT_REASON, drop_params_gated=True) + return None def _translate_param(self, param: str, value: object, model: str, drop_params: bool) -> tuple[str, object] | None: if param in ("temperature", "top_p"): diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index fa5794d71b3..92c6c8439b3 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_response_only_fields, unpack_legacy_defs, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -97,6 +98,7 @@ def _json_schema_response_format(schema: object, name: str) -> Mapping[str, obje EFFORT_KWARG_KEYS: Final = frozenset({"enable_thinking", "thinking", "reasoning_budget", "low_effort"}) +_NOT_A_REASONING_MODEL_REASON: Final = "litellm's model map does not list it as a reasoning model." def _bool_from_kwargs(kwargs: Mapping[str, object], keys: tuple[str, ...]) -> bool | None: @@ -107,6 +109,13 @@ def _bool_from_kwargs(kwargs: Mapping[str, object], keys: tuple[str, ...]) -> bo return None +def warn_dropped_effort_kwargs(chat_template_kwargs: object, effort: object, model: str) -> None: + if effort != "none": + warn_dropped_reasoning( + "chat_template_kwargs", chat_template_kwargs, model, _NOT_A_REASONING_MODEL_REASON, drop_params_gated=False + ) + + def effort_from_chat_template_kwargs(kwargs: Mapping[str, object]) -> object: enable_thinking: Final = _bool_from_kwargs(kwargs, ("enable_thinking", "thinking")) if enable_thinking is False: @@ -438,10 +447,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): if effort is None: return () if not supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): - verbose_logger.debug( - "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs effort keys.", - model, - ) + warn_dropped_effort_kwargs(OPAQUE_VALUE.validate_python(chat_template_kwargs), effort, model) return () return (("reasoning_effort", effort),) diff --git a/litellm/llms/fireworks_ai/completion/transformation.py b/litellm/llms/fireworks_ai/completion/transformation.py index 1f5d4f306a9..611e60cedd4 100644 --- a/litellm/llms/fireworks_ai/completion/transformation.py +++ b/litellm/llms/fireworks_ai/completion/transformation.py @@ -2,6 +2,7 @@ from collections.abc import Mapping from typing import Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUserMessage from litellm.utils import supports_reasoning @@ -12,6 +13,7 @@ from ..chat.transformation import ( NIM_VLLM_STRIP_PARAMS, FireworksAIConfig, effort_from_chat_template_kwargs, + warn_dropped_effort_kwargs, ) from ..common_utils import FireworksAIMixin, resolve_fireworks_resource_name @@ -129,10 +131,7 @@ class FireworksAITextCompletionConfig(FireworksAIMixin, BaseTextCompletionConfig ) return result if not supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): - verbose_logger.debug( - "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs effort keys.", - model, - ) + warn_dropped_effort_kwargs(OPAQUE_VALUE.validate_python(chat_template_kwargs), effort, model) return result return {**result, "reasoning_effort": effort} diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 63e6cb2b5da..9f9503f93b6 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -99,7 +99,7 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): thinking_effort: Final = ( None if "reasoning_effort" in non_default_params - else reasoning_effort_from_thinking(thinking, drop_params=drop_params) + else reasoning_effort_from_thinking(thinking, model=model, drop_params=drop_params) ) if thinking_effort is not None: non_default_params["reasoning_effort"] = thinking_effort diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 15e348453e0..a402e2c4048 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_internal_message_fields, strip_none_values_from_message, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, @@ -74,6 +75,12 @@ _MISTRAL_REQUEST_ROUTING: Final = TypeAdapter(MistralRequestRouting) _MISTRAL_CACHE_KEY_PARAMS: Final = TypeAdapter(MistralCacheKeyParams) _MISTRAL_REASONING_EFFORT_PARAMS: Final = TypeAdapter(MistralReasoningEffortParams) _MISTRAL_FILE_PART: Final = TypeAdapter(MistralFilePart) +_REASONING_PROMPT_REASON: Final = "litellm sends Magistral a reasoning system prompt instead of the requested level." +_NON_LEVEL_EFFORT_REASON: Final = ( + "the Mistral API takes it only as a level such as 'high', so Mistral applies its default effort." +) +_DOCUMENTED_MAGISTRAL_EFFORTS: Final = frozenset({"low", "medium", "high"}) +_REASONING_ON_THINKING_TYPES: Final = frozenset({"enabled", "adaptive"}) def _reasoning_level_and_summary(reasoning_effort: object) -> tuple[object, object]: @@ -83,6 +90,16 @@ def _reasoning_level_and_summary(reasoning_effort: object) -> tuple[object, obje return requested.get("effort"), requested.get("summary") +def _warn_unless_reasoning_prompt_honors(param: str, value: object, model: str) -> None: + if param == "reasoning_effort" and isinstance(value, str) and value in _DOCUMENTED_MAGISTRAL_EFFORTS: + return + thinking: Final = as_str_mapping(value) + thinking_type: Final = thinking.get("type") if thinking is not None else None + if param == "thinking" and isinstance(thinking_type, str) and thinking_type in _REASONING_ON_THINKING_TYPES: + return + warn_dropped_reasoning(param, value, model, _REASONING_PROMPT_REASON, drop_params_gated=False) + + def _is_json_array(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a decoded JSON array is a list return isinstance(value, list) @@ -253,6 +270,7 @@ class MistralConfig(OpenAIGPTConfig): if param in ("reasoning_effort", "thinking") and "magistral" in model.lower(): # Flag that we need to add reasoning system prompt optional_params["_add_reasoning_prompt"] = True + _warn_unless_reasoning_prompt_honors(param, OPAQUE_VALUE.validate_python(value), model) if param in ("parallel_tool_calls", "prompt_cache_key"): optional_params[param] = value return optional_params @@ -260,13 +278,18 @@ class MistralConfig(OpenAIGPTConfig): def _reasoning_level_to_send( self, model: str, requested: MistralReasoningEffortParams, drop_params: bool ) -> str | None: - level, summary = _reasoning_level_and_summary(requested.get("reasoning_effort")) + requested_effort: Final = requested.get("reasoning_effort") + level, summary = _reasoning_level_and_summary(requested_effort) sendable: Final = ( _accepted_reasoning_effort(model, level, self.custom_llm_provider) if isinstance(level, str) else None ) if sendable is not None and summary is None: return sendable if drop_params or litellm.drop_params is True: + if sendable is None and level is not None: + warn_dropped_reasoning( + "reasoning_effort", requested_effort, model, _NON_LEVEL_EFFORT_REASON, drop_params_gated=True + ) return sendable raise UnsupportedParamsError( message=( diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 8257a75b4a5..9c827822dc1 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -21,6 +21,7 @@ from litellm._logging import verbose_logger from litellm.constants import DEFAULT_OCI_CHAT_MAX_TOKENS from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.prompt_templates.common_utils import messages_with_file_parts_as_image_urls +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -428,6 +429,13 @@ class OCIChatConfig(BaseConfig): alias = param_map.get(key) if key == "reasoning_effort" and alias and not _model_supports_reasoning_effort(model): if drop_params or litellm.drop_params: + warn_dropped_reasoning( + key, + value, + model, + "the model map lists no reasoning_effort levels for it, so OCI would reject reasoningEffort.", + drop_params_gated=True, + ) continue raise OCIError( status_code=400, @@ -451,6 +459,14 @@ class OCIChatConfig(BaseConfig): if key == "n" and (value is None or value == 1): continue if drop_params or litellm.drop_params: + if key == "reasoning_effort": + warn_dropped_reasoning( + key, + value, + model, + "OCI's Cohere chat API has no reasoning effort field.", + drop_params_gated=True, + ) continue raise OCIError( status_code=400, diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index f9e7071db41..755d647da1d 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_images_from_message, extract_reasoning_content, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -29,7 +30,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import ModelResponse, ModelResponseStream -from ..common_utils import OllamaError +from ..common_utils import OllamaError, ollama_think if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -172,10 +173,7 @@ class OllamaChatConfig(BaseConfig): if value.get("json_schema") and value["json_schema"].get("schema"): optional_params["format"] = value["json_schema"]["schema"] if param == "reasoning_effort" and value is not None: - if model.startswith("gpt-oss"): - optional_params["think"] = value - else: - optional_params["think"] = value in {"low", "medium", "high"} + optional_params["think"] = ollama_think(model, OPAQUE_VALUE.validate_python(value)) ### FUNCTION CALLING LOGIC ### # Ollama 0.4+ supports native tool calling - pass tools directly # and let Ollama handle model capability detection diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index f313c62f026..fcf6a396a13 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -8,11 +8,14 @@ from pydantic import BaseModel, TypeAdapter from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import ModelResponse, ModelResponseStream _PROVIDER_FIELDS: Final = TypeAdapter(dict[str, object]) +_THINK_ON_EFFORTS: Final = frozenset({"low", "medium", "high"}) +_THINK_OFF_REASON: Final = "Ollama takes think only as on or off, and litellm maps this level to off." _DECODED_JSON: Final = TypeAdapter(object) @@ -296,3 +299,12 @@ class OllamaModelInfo(BaseLLMModelInfo): Return the base model name for Ollama (no-op). """ return model + + +def ollama_think(model: str, reasoning_effort: object) -> object: + if model.startswith("gpt-oss"): + return reasoning_effort + think: Final = reasoning_effort in _THINK_ON_EFFORTS + if isinstance(reasoning_effort, str) and not think and reasoning_effort != "none": + warn_dropped_reasoning("reasoning_effort", reasoning_effort, model, _THINK_OFF_REASON, drop_params_gated=False) + return think diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 239b4e32645..6dbcb050e11 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, inline_remote_image_urls, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -43,7 +44,7 @@ from litellm.types.utils import ( generate_id, ) -from ..common_utils import OllamaError, OllamaModelInfo, convert_image +from ..common_utils import OllamaError, OllamaModelInfo, convert_image, ollama_think if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -249,10 +250,7 @@ class OllamaConfig(BaseConfig): elif param == "stop": optional_params["stop"] = value elif param == "reasoning_effort" and value is not None: - if model.startswith("gpt-oss"): - optional_params["think"] = value - else: - optional_params["think"] = value in {"low", "medium", "high"} + optional_params["think"] = ollama_think(model, OPAQUE_VALUE.validate_python(value)) elif param == "response_format" and isinstance(value, dict): if value["type"] == "json_object": optional_params["format"] = "json" diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index aa14a985afe..c113cb351d8 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -32,6 +32,7 @@ import litellm from litellm.exceptions import error_response_without_credentials from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning_params from litellm.litellm_core_utils.token_counter import token_counter, usage_tokenizer_model from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( @@ -165,17 +166,11 @@ _DECODED_JSON: Final = TypeAdapter(object) def drop_params_from_unprocessable_entity_error( e: openai.UnprocessableEntityError | httpx.HTTPStatusError, data: Mapping[str, object], + *, + model: str | None, + drop_params_gated: bool, ) -> dict[str, object]: - """ - Helper function to read OpenAI UnprocessableEntityError and drop the params that raised an error from the error message. - - Args: - e (UnprocessableEntityError): The UnprocessableEntityError exception - data (Dict[str, Any]): The original data dictionary containing all parameters - - Returns: - Dict[str, Any]: A new dictionary with invalid parameters removed - """ + """Return a copy of ``data`` without the params a 422 response names as invalid.""" invalid_params: Final[list[str]] = [] if isinstance(e, httpx.HTTPStatusError): error_json: Final = decoded_json_object(_DECODED_JSON.validate_python(e.response.json())) @@ -203,6 +198,13 @@ def drop_params_from_unprocessable_entity_error( invalid_params.append(error_dict["loc"][1]) new_data: Final = {k: v for k, v in data.items() if k not in invalid_params} + warn_dropped_reasoning_params( + invalid_params, + data, + model, + "the provider answered 422 naming it an invalid param.", + drop_params_gated=drop_params_gated, + ) return new_data diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index bef99fd0719..bdd9f285823 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -940,7 +940,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): except openai.UnprocessableEntityError as e: ## check if body contains unprocessable params - related issue https://github.com/BerriAI/litellm/issues/4800 if litellm.drop_params is True or drop_params is True: - inference_params = drop_params_from_unprocessable_entity_error(e, inference_params) + inference_params = drop_params_from_unprocessable_entity_error( + e, inference_params, model=model, drop_params_gated=True + ) else: raise e # e.message @@ -1097,7 +1099,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): except openai.UnprocessableEntityError as e: ## check if body contains unprocessable params - related issue https://github.com/BerriAI/litellm/issues/4800 if litellm.drop_params is True or drop_params is True: - data = drop_params_from_unprocessable_entity_error(e, data) + data = drop_params_from_unprocessable_entity_error(e, data, model=model, drop_params_gated=True) else: raise e # e.message @@ -1252,7 +1254,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): except openai.UnprocessableEntityError as e: ## check if body contains unprocessable params - related issue https://github.com/BerriAI/litellm/issues/4800 if litellm.drop_params is True or drop_params is True: - data = drop_params_from_unprocessable_entity_error(e, data) + data = drop_params_from_unprocessable_entity_error(e, data, model=model, drop_params_gated=True) else: raise e except ( diff --git a/litellm/llms/tencent/chat/transformation.py b/litellm/llms/tencent/chat/transformation.py index 1f0bcd218d6..23a6ccd52c6 100644 --- a/litellm/llms/tencent/chat/transformation.py +++ b/litellm/llms/tencent/chat/transformation.py @@ -9,12 +9,15 @@ from typing import Final, TypedDict from typing_extensions import ReadOnly import litellm +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.secret_managers.main import get_secret_str from litellm.utils import supports_reasoning from ...openai.chat.gpt_transformation import OpenAIGPTConfig +_ADAPTIVE_BUDGET_REASON: Final = "it is not an integer, so litellm sends adaptive thinking with no budget." + class ThinkingPayload(TypedDict, total=False): """Tencent TokenHub `thinking` object. @@ -102,6 +105,10 @@ class TencentChatConfig(OpenAIGPTConfig): if isinstance(budget, int): coerced_with_budget: Final[ThinkingPayload] = {"type": "adaptive", "budget_tokens": budget} return coerced_with_budget + if budget is not None: + warn_dropped_reasoning( + "thinking.budget_tokens", budget, model, _ADAPTIVE_BUDGET_REASON, drop_params_gated=False + ) coerced: Final[ThinkingPayload] = {"type": "adaptive"} return coerced diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index c506192f253..1ef5267df14 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc from litellm.litellm_core_utils.prompt_templates.factory import ( encode_tool_call_id_with_signature, ) +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -201,6 +202,13 @@ def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsErr _SAMPLING_DEPRECATION_UNMAPPED_MODELS: Final = "models outside litellm's model map" +_NON_INTEGER_BUDGET_REASON: Final = ( + "Gemini takes budget_tokens only as an integer, so it uses its default thinking budget." +) +_LEVEL_ONLY_BUDGET_REASON: Final = ( + "Gemini 3 and later set reasoning by thinking level, not budget, so litellm does not send the budget. " + "Send reasoning_effort to pick a level." +) _sampling_deprecation_lock: Final = threading.Lock() @@ -1466,6 +1474,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): params["includeThoughts"] = False else: params["includeThoughts"] = sends_thoughts + if thinking_budget is not None: + warn_dropped_reasoning( + "thinking.budget_tokens", + thinking_budget, + model, + _LEVEL_ONLY_BUDGET_REASON, + drop_params_gated=False, + ) # Follow provider defaults unless explicitly opted into legacy behavior. if litellm.enable_gemini_default_thinking_level_low is True: params["thinkingLevel"] = ( @@ -1478,8 +1494,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # For older Gemini models, use thinkingBudget if thinking_enabled and not VertexGeminiConfig._is_thinking_budget_zero(thinking_budget): params["includeThoughts"] = sends_thoughts - if thinking_budget is not None and isinstance(thinking_budget, int): + if isinstance(thinking_budget, int): params["thinkingBudget"] = thinking_budget + elif thinking_budget is not None: + warn_dropped_reasoning( + "thinking.budget_tokens", + thinking_budget, + model, + _NON_INTEGER_BUDGET_REASON, + drop_params_gated=False, + ) return params diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py index 3563fcd96ae..c315a31b65d 100644 --- a/litellm/router_utils/reasoning_effort_capability.py +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -38,7 +38,7 @@ from litellm.litellm_core_utils.reasoning_effort_strength import ( ReasoningEffortRounding, nearest_declared_reasoning_effort, ) -from litellm.litellm_core_utils.reasoning_effort_utils import reasoning_effort_from_thinking +from litellm.litellm_core_utils.reasoning_effort_utils import reasoning_effort_from_thinking_silently from litellm.llms.anthropic.common_utils import AnthropicModelInfo, validated_object_mapping from litellm.types.llms.anthropic import AnthropicThinkingParam from litellm.types.llms.openai import REASONING_EFFORT @@ -143,7 +143,7 @@ def with_adaptive_thinking_as_effort( params: Mapping[str, object], *, model: str, custom_llm_provider: str, drop_params: bool ) -> Mapping[str, object]: translated: Final = dict(params) - AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model_silently( model=model, optional_params=translated, custom_llm_provider=custom_llm_provider, drop_params=drop_params ) return translated @@ -151,7 +151,7 @@ def with_adaptive_thinking_as_effort( def thinking_budget_effort(thinking: object, *, model: str, custom_llm_provider: str, drop_params: bool) -> str | None: if custom_llm_provider in _PROVIDERS_SENDING_THINKING_AS_EFFORT: - return reasoning_effort_from_thinking(thinking, drop_params=drop_params) + return reasoning_effort_from_thinking_silently(thinking, drop_params=drop_params) adaptive_config: Final = as_str_mapping( with_adaptive_thinking_as_effort( {"thinking": thinking}, model=model, custom_llm_provider=custom_llm_provider, drop_params=drop_params @@ -162,7 +162,7 @@ def thinking_budget_effort(thinking: object, *, model: str, custom_llm_provider: return adaptive_effort if isinstance(adaptive_effort, str) else None if "thinking" in (get_supported_openai_params(model=model, custom_llm_provider=custom_llm_provider) or ()): return None - return reasoning_effort_from_thinking(thinking, drop_params=drop_params) + return reasoning_effort_from_thinking_silently(thinking, drop_params=drop_params) def provider_reads_thinking(thinking: object, *, model: str, custom_llm_provider: str, drop_params: bool) -> bool: diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index d815e1eedd3..c2ee58b67ed 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -6,6 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic """ import json +import logging import re from typing import Any, Literal, Optional, Final from unittest.mock import MagicMock, patch @@ -3391,3 +3392,18 @@ class TestAnthropicMessagesHandlerPostCallHookResponse: native = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]} assert AnthropicMessagesHandler().post_call_hook_response(native) is native + + +def test_reading_a_dropped_budget_for_a_guardrail_logs_no_drop_warning(caplog: pytest.LogCaptureFixture) -> None: + data: Final = { + "model": "gemini-2.5-flash", + "max_tokens": 64, + "thinking": {"type": "enabled", "budget_tokens": "high"}, + "messages": [{"role": "user", "content": "hi"}], + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + messages: Final = AnthropicMessagesHandler().get_structured_messages(data) + + assert messages == [{"role": "user", "content": "hi"}] + assert not [record for record in caplog.records if "Dropping" in record.getMessage()] diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 4d5d6cc17a9..cbb60caa57b 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1,5 +1,6 @@ import copy import json +import logging from itertools import chain from typing import Final from unittest.mock import MagicMock, patch @@ -3189,6 +3190,47 @@ def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model( ) +def _reasoning_effort_drop_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple( + record.getMessage() for record in caplog.records if "Dropping reasoning_effort=" in record.getMessage() + ) + + +@pytest.mark.parametrize("reasoning_effort_value", [5, {"effort": 5, "summary": "concise"}]) +def test_reasoning_effort_that_is_not_a_level_name_is_dropped_with_one_warning( + reasoning_effort_value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = AnthropicConfig().map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort_value}, + optional_params={}, + model="claude-sonnet-4-6-20260219", + drop_params=False, + ) + + warnings: Final = _reasoning_effort_drop_warnings(caplog) + assert "thinking" not in result + assert "output_config" not in result + assert len(warnings) == 1 + assert f"reasoning_effort={reasoning_effort_value!r}" in warnings[0] + assert "claude-sonnet-4-6-20260219" in warnings[0] + + +@pytest.mark.parametrize("reasoning_effort_value", [{"summary": "concise"}, {"effort": "low", "summary": "concise"}]) +def test_reasoning_effort_with_a_level_name_or_no_level_logs_no_drop( + reasoning_effort_value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + AnthropicConfig().map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort_value}, + optional_params={}, + model="claude-sonnet-4-6-20260219", + drop_params=False, + ) + + assert _reasoning_effort_drop_warnings(caplog) == () + + @pytest.mark.parametrize( "model,budget_tokens,expected", [ diff --git a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py index c30644e7fdc..e0882653324 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py @@ -1,6 +1,7 @@ import base64 import functools import json +import logging from collections.abc import Mapping from types import MappingProxyType from typing import Any, Final, cast @@ -2871,14 +2872,21 @@ def test_translate_completion_input_params_rejects_a_non_numeric_budget_without_ @pytest.mark.parametrize("budget", ["2048", True]) def test_translate_completion_input_params_drops_a_non_numeric_budget_under_drop_params( - budget: object, monkeypatch: pytest.MonkeyPatch + budget: object, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ): monkeypatch.setattr(litellm, "drop_params", False) - translated: Final = AnthropicAdapter().translate_completion_input_params( - _gemini_request_with_budget(budget, drop_params=True) - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + translated: Final = AnthropicAdapter().translate_completion_input_params( + _gemini_request_with_budget(budget, drop_params=True) + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert translated is not None assert not {"reasoning_effort", "thinking"} & translated.keys() + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in (f"thinking.budget_tokens={budget!r}", "gemini-2.5-flash", "(drop_params=True)") + ) def test_translate_openai_content_to_anthropic_reasoning_content_without_thinking_blocks(): diff --git a/tests/unit/llms/anthropic/messages/adapters/test_handler_reasoning_effort_normalization.py b/tests/unit/llms/anthropic/messages/adapters/test_handler_reasoning_effort_normalization.py index e504513dd36..27ed1f5ebd6 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_handler_reasoning_effort_normalization.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_handler_reasoning_effort_normalization.py @@ -8,6 +8,7 @@ leaves the adapter, in the shape the target expects. import functools import json +import logging from collections.abc import Mapping from typing import Final @@ -510,26 +511,33 @@ class TestADeclaredDeploymentRoundsTheMessagesTierOntoItsLevels: @pytest.mark.asyncio @pytest.mark.parametrize("budget", ["2048", True]) @pytest.mark.parametrize("litellm_params, model_info, upstream_url, upstream_body", _MALFORMED_BUDGET_TARGETS) - async def test_drop_params_drops_a_non_numeric_budget( + async def test_drop_params_drops_a_non_numeric_budget_with_one_warning( self, local_model_cost_map: None, respx_mock: respx.MockRouter, + caplog: pytest.LogCaptureFixture, litellm_params: Mapping[str, object], model_info: Mapping[str, object], upstream_url: str, upstream_body: Mapping[str, object], budget: object, ) -> None: - sent: Final = await _body_sent_upstream( - respx_mock, - {**litellm_params, "drop_params": True}, - model_info, - upstream_url, - upstream_body, - thinking={"type": "enabled", "budget_tokens": budget}, - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + sent: Final = await _body_sent_upstream( + respx_mock, + {**litellm_params, "drop_params": True}, + model_info, + upstream_url, + upstream_body, + thinking={"type": "enabled", "budget_tokens": budget}, + ) assert not {"reasoning_effort", "reasoning", "thinking"} & sent.keys() + dropped: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + assert len(dropped) == 1 + assert f"thinking.budget_tokens={budget!r}" in dropped[0] + assert "(drop_params=True)" in dropped[0] + assert str(litellm_params["model"]).split("/", 1)[1] in dropped[0] @pytest.mark.asyncio @pytest.mark.parametrize( diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index 5ce6d03f477..af3108e5072 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -159,7 +159,7 @@ async def test_openai_model_does_not_forward_stream_options_to_responses_api(): await litellm.anthropic.messages.acreate( max_tokens=100, messages=[{"role": "user", "content": "Hello, how are you?"}], - model="openai/gpt-5.5", + model="openai/gpt-6.1-sol", api_key="test-api-key", stream_options={"include_usage": True}, ) @@ -669,7 +669,9 @@ class TestThinkingSummaryPreservation: ) thinking = {"type": "enabled", "budget_tokens": 5000, "summary": "concise"} - result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( + thinking, model="openai/gpt-6.1-sol" + ) assert result == {"effort": "high", "summary": "concise"} def test_responses_adapter_no_summary_by_default(self): @@ -683,7 +685,9 @@ class TestThinkingSummaryPreservation: try: litellm.reasoning_auto_summary = False thinking = {"type": "enabled", "budget_tokens": 5000} - result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( + thinking, model="openai/gpt-6.1-sol" + ) assert result == {"effort": "high"} assert result is not None and "summary" not in result finally: diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py index 596cb968a0e..1d0d70d8e60 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import base64 import json +import logging import struct import time import uuid @@ -891,24 +892,32 @@ async def test_native_messages_send_no_effort_for_a_default_effort_beside_a_thin @pytest.mark.asyncio @pytest.mark.parametrize("model_info", [None, _DECLARING_LOW_AND_HIGH]) @pytest.mark.parametrize("drop_globally", [False, True]) -async def test_native_messages_drop_params_drops_a_non_numeric_thinking_budget_on_adaptive_claude( +async def test_native_messages_drop_params_drops_a_non_numeric_thinking_budget_on_adaptive_claude_with_one_warning( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, model_info: Mapping[str, object] | None, drop_globally: bool, ) -> None: monkeypatch.setattr(litellm, "drop_params", drop_globally) request_drop: Final = {} if drop_globally else {"drop_params": True} - dropped: Final = await _native_opus_body_sent( - respx_mock, model_info, thinking={"type": "enabled", "budget_tokens": "lots"}, **request_drop - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + dropped: Final = await _native_opus_body_sent( + respx_mock, model_info, thinking={"type": "enabled", "budget_tokens": "lots"}, **request_drop + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] adaptive_alone: Final = await _native_opus_body_sent( respx_mock, model_info, thinking={"type": "adaptive"}, **request_drop ) assert "output_config" not in adaptive_alone assert dropped == adaptive_alone + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in ("thinking.budget_tokens='lots'", "claude-opus-5-5", "(drop_params=True)") + ) _MIXED_GROUP: Final = "gpt-5.6-mixed" diff --git a/tests/unit/llms/anthropic/messages/native/test_reasoning_effort_translation.py b/tests/unit/llms/anthropic/messages/native/test_reasoning_effort_translation.py index 50dc0deabfd..28abb348371 100644 --- a/tests/unit/llms/anthropic/messages/native/test_reasoning_effort_translation.py +++ b/tests/unit/llms/anthropic/messages/native/test_reasoning_effort_translation.py @@ -1,5 +1,8 @@ """Tests for ``reasoning_effort`` translation on the Anthropic /v1/messages route.""" +import logging +from typing import Final + import pytest from litellm.constants import ( @@ -46,6 +49,49 @@ def test_reasoning_effort_maps_to_output_config_for_adaptive_model( assert result.get("output_config") == {"effort": expected_effort} +def _reasoning_effort_drop_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple( + record.getMessage() for record in caplog.records if "Dropping reasoning_effort=" in record.getMessage() + ) + + +@pytest.mark.parametrize("reasoning_effort", [5, {"effort": "high"}]) +def test_reasoning_effort_that_is_not_a_string_is_dropped_with_one_warning( + reasoning_effort: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-opus-4-7", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params={"max_tokens": 1024, "reasoning_effort": reasoning_effort}, + litellm_params={}, + headers={}, + ) + + warnings: Final = _reasoning_effort_drop_warnings(caplog) + assert "reasoning_effort" not in result + assert "output_config" not in result + assert len(warnings) == 1 + assert f"reasoning_effort={reasoning_effort!r}" in warnings[0] + assert "claude-opus-4-7" in warnings[0] + + +@pytest.mark.parametrize("reasoning_effort", ["high", {"summary": "concise"}]) +def test_reasoning_effort_that_is_a_string_or_names_no_level_logs_no_drop( + reasoning_effort: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-opus-4-7", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params={"max_tokens": 1024, "reasoning_effort": reasoning_effort}, + litellm_params={}, + headers={}, + ) + + assert _reasoning_effort_drop_warnings(caplog) == () + + def test_reasoning_effort_none_clears_thinking_and_output_config(): config = AnthropicMessagesConfig() optional_params = { diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_transformation.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_transformation.py index f6fcb9e6f1e..ece0cf5d423 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_transformation.py @@ -44,6 +44,7 @@ def _make_request(**overrides) -> AnthropicMessagesRequest: _ADAPTER = LiteLLMAnthropicToResponsesAPIAdapter() +_MODEL: Final = "openai/gpt-6.1-sol" # --------------------------------------------------------------------------- @@ -957,13 +958,13 @@ class TestTranslateThinkingToReasoning: """Anthropic thinking param -> Responses API reasoning param.""" def test_budget_high_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}, model=_MODEL) # Default (reasoning_auto_summary=False): only effort, no summary assert result == {"effort": "high"} assert result is not None and "summary" not in result def test_budget_above_threshold_high_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 50000}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 50000}, model=_MODEL) assert result is not None assert result["effort"] == "high" assert "summary" not in result @@ -973,7 +974,8 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result == {"effort": "medium"} assert result is not None and "summary" not in result @@ -983,13 +985,14 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result == {"effort": "low"} assert result is not None and "summary" not in result def test_budget_minimal_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 500}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 500}, model=_MODEL) assert result == {"effort": "minimal"} assert result is not None and "summary" not in result @@ -998,7 +1001,8 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result_high is not None assert result_high["effort"] == "high" @@ -1006,7 +1010,8 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result_medium is not None assert result_medium["effort"] == "medium" @@ -1015,23 +1020,24 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result_low is not None assert result_low["effort"] == "low" assert "summary" not in result_low def test_disabled_type_returns_none(self): - result = _ADAPTER.translate_thinking_to_reasoning({"type": "disabled"}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "disabled"}, model=_MODEL) assert result is None def test_non_dict_returns_none(self): - result = _ADAPTER.translate_thinking_to_reasoning("enabled") # type: ignore + result = _ADAPTER.translate_thinking_to_reasoning("enabled", model=_MODEL) # type: ignore assert result is None def test_missing_budget_defaults_to_minimal(self): """Missing budget_tokens defaults to 0, below the low threshold -> minimal.""" - result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled"}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled"}, model=_MODEL) assert result == {"effort": "minimal"} assert result is not None and "summary" not in result @@ -1042,7 +1048,7 @@ class TestTranslateThinkingToReasoning: original = litellm.reasoning_auto_summary try: litellm.reasoning_auto_summary = True - result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}, model=_MODEL) assert result == {"effort": "high", "summary": "detailed"} finally: litellm.reasoning_auto_summary = original @@ -1059,7 +1065,8 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result == {"effort": "medium", "summary": "detailed"} finally: diff --git a/tests/unit/llms/anthropic/messages/test_reasoning_effort_fields.py b/tests/unit/llms/anthropic/messages/test_reasoning_effort_fields.py index 3432c9503de..7128c04a056 100644 --- a/tests/unit/llms/anthropic/messages/test_reasoning_effort_fields.py +++ b/tests/unit/llms/anthropic/messages/test_reasoning_effort_fields.py @@ -162,7 +162,8 @@ class TestAdapterAdaptiveThinking: adapter = LiteLLMAnthropicMessagesAdapter() result = adapter.translate_anthropic_thinking_to_reasoning_effort( - {"type": "adaptive"} + {"type": "adaptive"}, + model="openai/gpt-6.1-sol", ) assert result == "medium" @@ -198,6 +199,7 @@ class TestAdapterAdaptiveThinking: result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( thinking={"type": "adaptive"}, output_config={"effort": "xhigh"}, + model="openai/gpt-6.1-sol", ) assert result is not None assert result["effort"] == "xhigh" @@ -210,6 +212,7 @@ class TestAdapterAdaptiveThinking: result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( thinking={"type": "adaptive"}, + model="openai/gpt-6.1-sol", ) assert result is not None assert result["effort"] == "medium" diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index 3ec72f29d34..c1f28e33969 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -11,6 +11,7 @@ Verifies that: """ import json +import logging import os import sys import threading @@ -4530,16 +4531,25 @@ def test_shared_legacy_thinking_translation_rejects_a_non_numeric_budget(budget_ @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize("budget_tokens", _NON_NUMERIC_BUDGETS) -def test_shared_legacy_thinking_translation_drops_a_non_numeric_budget_under_drop_params(budget_tokens: object) -> None: +def test_shared_legacy_thinking_translation_drops_a_non_numeric_budget_under_drop_params( + budget_tokens: object, caplog: pytest.LogCaptureFixture +) -> None: from litellm.llms.anthropic.common_utils import AnthropicModelInfo optional_params: Final = {"thinking": {"type": "enabled", "budget_tokens": budget_tokens}} - AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( - model="claude-opus-5", optional_params=optional_params, custom_llm_provider="azure_ai", drop_params=True - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model="claude-opus-5", optional_params=optional_params, custom_llm_provider="azure_ai", drop_params=True + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert optional_params == {"thinking": {"type": "adaptive"}} + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in (f"thinking.budget_tokens={budget_tokens!r}", "claude-opus-5", "(drop_params=True)") + ) @pytest.mark.parametrize( diff --git a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py index 4d0ebef4f2d..d95f831df3e 100644 --- a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py +++ b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Mapping from datetime import datetime from types import SimpleNamespace @@ -597,6 +598,45 @@ async def test_native_count_preserves_settings_and_requires_every_prefix( assert unexpected.call_count == 0 +@pytest.mark.usefixtures("local_model_cost_map") +def test_native_baseline_body_drops_a_non_numeric_budget_without_logging_a_drop( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + request: Final = { + "model": "claude-opus-5", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + "thinking": {"type": "enabled", "budget_tokens": "high"}, + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + body: Final = prepare_native_baseline_body(request, "anthropic/claude-opus-5") + + assert body is not None + assert body["thinking"] == {"type": "adaptive"} + assert [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] == [] + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_native_baseline_body_drops_a_non_string_reasoning_effort_without_logging_a_drop( + caplog: pytest.LogCaptureFixture, +) -> None: + request: Final = { + "model": "claude-opus-5", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + "reasoning_effort": 5, + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + body: Final = prepare_native_baseline_body(request, "anthropic/claude-opus-5") + + assert body is not None + assert "reasoning_effort" not in body + assert [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] == [] + + def test_native_baseline_body_matches_the_sent_body_when_drop_params_meets_an_auto_added_beta( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index 145d45317c3..9f68a4fb73f 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -1,4 +1,5 @@ import json +import logging import traceback from collections.abc import Callable from types import MappingProxyType @@ -758,6 +759,66 @@ def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): assert result["tools"][1]["function"]["name"] == "read_file" +def test_extra_parameters_error_drops_reasoning_effort_with_one_warning(caplog: pytest.LogCaptureFixture) -> None: + import httpx + + error_text: Final = json.dumps( + { + "error": { + "message": "Extra parameters ['reasoning_effort'] are not allowed when extra-parameters is not set " + "or set to be 'error'. Set extra-parameters to 'pass-through' to pass to the model." + } + } + ) + error: Final = httpx.HTTPStatusError( + message="422", + request=httpx.Request("POST", "https://example.services.ai.azure.com/models/chat/completions"), + response=httpx.Response(422, text=error_text), + ) + request_data: Final = { + "model": "Kimi-K2.6", + "messages": [{"role": "user", "content": "Say hi."}], + "reasoning_effort": "high", + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + retried: Final = AzureAIStudioConfig().transform_request_on_unprocessable_entity_error(error, request_data) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + + assert retried == {"model": "Kimi-K2.6", "messages": [{"role": "user", "content": "Say hi."}]} + assert len(drop_warnings) == 1 + assert all(fragment in drop_warnings[0] for fragment in ("reasoning_effort='high'", "Kimi-K2.6")) + assert "drop_params" not in drop_warnings[0] + + +def test_extra_inputs_error_names_drop_params_in_the_reasoning_drop_warning(caplog: pytest.LogCaptureFixture) -> None: + import httpx + + detail: Final = [ + {"type": "extra_forbidden", "loc": ["body", "reasoning_effort"], "msg": "Extra inputs are not permitted"} + ] + error: Final = httpx.HTTPStatusError( + message="422", + request=httpx.Request("POST", "https://example.services.ai.azure.com/models/chat/completions"), + response=httpx.Response(422, text=json.dumps({"error": {"message": json.dumps({"detail": detail})}})), + ) + request_data: Final = { + "model": "Mistral-Large-3", + "messages": [{"role": "user", "content": "Say hi."}], + "reasoning_effort": "high", + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + retried: Final = AzureAIStudioConfig().transform_request_on_unprocessable_entity_error(error, request_data) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + + assert retried == {"model": "Mistral-Large-3", "messages": [{"role": "user", "content": "Say hi."}]} + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] for fragment in ("reasoning_effort='high'", "Mistral-Large-3", "(drop_params=True)") + ) + + def _find_key_anywhere(obj, key: str) -> bool: if isinstance(obj, dict): if key in obj: diff --git a/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py b/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py index 3faae2e3289..f4dcff1969e 100644 --- a/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py +++ b/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py @@ -1,6 +1,7 @@ """Bedrock Runtime Chat Completions: the default for GPT 5.6 and newer, ``bedrock/chat_completions/`` for the rest.""" import json +import logging import re import httpx @@ -706,18 +707,30 @@ def test_with_max_completion_tokens_leaves_other_params_alone(): @pytest.mark.parametrize( - "model", - ["us.xai.grok-4.6", "bedrock/us-gov-west-1/us.xai.grok-4.6"], + ("model", "logged_model"), + [ + ("us.xai.grok-4.6", "us.xai.grok-4.6"), + ("chat_completions/us.xai.grok-4.6", "us.xai.grok-4.6"), + ("bedrock/us-gov-west-1/us.xai.grok-4.6", "us-gov-west-1/us.xai.grok-4.6"), + ], ) -def test_map_openai_params_drops_reasoning_effort_none_for_grok(model): +def test_map_openai_params_drops_reasoning_effort_none_for_grok_with_one_warning( + model: str, logged_model: str, caplog: pytest.LogCaptureFixture +) -> None: cfg = AmazonBedrockRuntimeChatCompletionsConfig() - mapped = cfg.map_openai_params( - non_default_params={"reasoning_effort": "none", "max_tokens": 64}, - optional_params={}, - model=model, - drop_params=False, - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + mapped = cfg.map_openai_params( + non_default_params={"reasoning_effort": "none", "max_tokens": 64}, + optional_params={}, + model=model, + drop_params=False, + ) + drop_warnings = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert "reasoning_effort" not in mapped + assert len(drop_warnings) == 1 + assert "reasoning_effort='none'" in drop_warnings[0] + assert f"model={logged_model!r}" in drop_warnings[0] + assert "drop_params" not in drop_warnings[0] def test_map_openai_params_keeps_reasoning_effort_low_for_grok(): @@ -750,17 +763,24 @@ def test_map_openai_params_refuses_a_non_string_reasoning_effort_without_drop_pa @pytest.mark.parametrize("reasoning_effort", [["low"], {"effort": "low"}, 5], ids=["list", "object", "int"]) @pytest.mark.parametrize("drop_params_via", ["request", "litellm.drop_params"]) def test_map_openai_params_drops_a_non_string_reasoning_effort_under_drop_params( - monkeypatch, model, reasoning_effort, drop_params_via + monkeypatch, caplog, model, reasoning_effort, drop_params_via ): monkeypatch.setattr(litellm, "drop_params", drop_params_via == "litellm.drop_params") - mapped = AmazonBedrockRuntimeChatCompletionsConfig().map_openai_params( - non_default_params={"reasoning_effort": reasoning_effort, "max_tokens": 64}, - optional_params={}, - model=model, - drop_params=drop_params_via == "request", - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + mapped = AmazonBedrockRuntimeChatCompletionsConfig().map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort, "max_tokens": 64}, + optional_params={}, + model=model, + drop_params=drop_params_via == "request", + ) + drop_warnings = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert "reasoning_effort" not in mapped assert mapped["max_completion_tokens"] == 64 + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in (f"reasoning_effort={reasoning_effort!r}", repr(model), "(drop_params=True)") + ) def test_map_openai_params_keeps_reasoning_effort_none_for_gpt56(): diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 743e7e74431..524caaba7a1 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -13,6 +13,7 @@ import pytest import respx import litellm +from litellm._logging import verbose_logger from litellm import ModelResponse, Router, completion from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, @@ -11096,6 +11097,52 @@ def test_a_deployment_with_a_batch_control_plane_endpoint_keeps_it_out_of_chat_r assert "aws_bedrock_control_plane_endpoint" not in route.calls[0].request.content.decode() +def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: + return tuple( + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + ) + + +_DEEPSEEK_R1_PROFILE_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.deepseek.r1-v1:0" + + +@pytest.mark.parametrize( + "param, value", (("thinking", {"type": "enabled", "budget_tokens": 1024}), ("reasoning_effort", "high")) +) +def test_bedrock_deepseek_arn_reasoning_drop_warns_once( + param: str, value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + optional_params: Final = AmazonConverseConfig().map_openai_params( + non_default_params={param: value, "max_tokens": 100}, + optional_params={}, + model=_DEEPSEEK_R1_PROFILE_ARN, + drop_params=False, + ) + + warnings: Final = _warnings_naming(caplog, param) + assert param not in optional_params + assert len(warnings) == 1 + assert _DEEPSEEK_R1_PROFILE_ARN in warnings[0] + + +@pytest.mark.parametrize( + "param, value", (("thinking", {"type": "enabled", "budget_tokens": 1024}), ("reasoning_effort", "high")) +) +def test_bedrock_deepseek_r1_drop_of_an_unadvertised_param_does_not_warn( + param: str, value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + AmazonConverseConfig().map_openai_params( + non_default_params={param: value, "max_tokens": 100}, + optional_params={}, + model="us.deepseek.r1-v1:0", + drop_params=False, + ) + + assert _warnings_naming(caplog, param) == () + + def _converse_bound_thinking_request( headers: dict[str, str] | None, body_betas: tuple[str, ...] = (), drop_params: bool = True ) -> dict[str, object]: diff --git a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py index 3fbb12881da..56af47aeebe 100644 --- a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py +++ b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Iterator from typing import Final @@ -6,6 +7,7 @@ import httpx import litellm import pytest import respx +from litellm._logging import verbose_logger from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.deepseek.chat.transformation import DeepSeekChatConfig @@ -951,3 +953,42 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) == "https://gateway.example.com/deepseek/chat/completions" ) + + +def _budget_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple( + r.getMessage() + for r in caplog.records + if r.levelno == logging.WARNING and "Dropping thinking.budget_tokens=" in r.getMessage() + ) + + +def test_enabled_thinking_budget_dropped_for_deepseek_warns_once(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = DeepSeekChatConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 4000}}, + optional_params={}, + model="deepseek-reasoner", + drop_params=False, + ) + + warnings: Final = _budget_warnings(caplog) + assert result["thinking"] == {"type": "enabled"} + assert len(warnings) == 1 + assert "budget_tokens=4000" in warnings[0] + assert "deepseek-reasoner" in warnings[0] + + +@pytest.mark.parametrize( + "thinking", ({"type": "enabled"}, {"type": "disabled", "budget_tokens": 4000}), ids=("no-budget", "disabled") +) +def test_deepseek_thinking_without_a_requested_budget_does_not_warn( + thinking: dict[str, object], caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = DeepSeekChatConfig().map_openai_params( + non_default_params={"thinking": thinking}, optional_params={}, model="deepseek-reasoner", drop_params=False + ) + + assert result["thinking"] == {"type": thinking["type"]} + assert _budget_warnings(caplog) == () diff --git a/tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py b/tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py index b7a79f03f8d..f4a7455ff27 100644 --- a/tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py +++ b/tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py @@ -1,3 +1,6 @@ +import logging +from typing import Final + import httpx import pytest @@ -242,11 +245,33 @@ def test_map_openai_params_maps_reasoning_effort(effort, expected): assert mapped["reasoning"] is expected -def test_map_openai_params_drops_unknown_reasoning_effort_when_dropping(): - mapped = FalAIChatConfig().map_openai_params( - non_default_params={"reasoning_effort": "extreme"}, optional_params={}, model=MODEL, drop_params=True - ) +@pytest.mark.parametrize("drop_globally", [False, True]) +def test_map_openai_params_drops_unknown_reasoning_effort_with_one_warning( + drop_globally: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(litellm, "drop_params", drop_globally) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + mapped: Final = FalAIChatConfig().map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, + optional_params={}, + model=MODEL, + drop_params=not drop_globally, + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert "reasoning" not in mapped + assert len(drop_warnings) == 1 + assert all(fragment in drop_warnings[0] for fragment in ("reasoning_effort='xhigh'", MODEL, "(drop_params=True)")) + + +def test_map_openai_params_rejects_unknown_reasoning_effort_without_drop_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(FalAIError) as exc_info: + FalAIChatConfig().map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, optional_params={}, model=MODEL, drop_params=False + ) + assert exc_info.value.status_code == 400 @pytest.mark.parametrize("effort", [{"level": "low"}, ["low"], 1]) diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 28f4fdead9e..ed166d206c8 100644 --- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Callable from typing import Final from unittest.mock import MagicMock, patch @@ -9,6 +10,7 @@ import respx from pydantic import ValidationError import litellm +from litellm._logging import verbose_logger from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -2361,6 +2363,39 @@ def test_global_disable_flag_with_transform_messages_helper( assert json.loads(route.calls.last.request.content)["messages"] == [{"role": "user", "content": content}] +def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: + return tuple( + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + ) + + +def test_chat_template_kwargs_effort_dropped_for_non_reasoning_model_warns_once( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + FireworksAIConfig().map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"reasoning_budget": 512}}}, + _NON_REASONING_MODEL, + ) + + warnings: Final = _warnings_naming(caplog, "chat_template_kwargs") + assert len(warnings) == 1 + assert "'reasoning_budget': 512" in warnings[0] + assert "llama-v3-70b-instruct" in warnings[0] + + +def test_chat_template_kwargs_thinking_off_for_non_reasoning_model_does_not_warn( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + FireworksAIConfig().map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": False}}}, + _NON_REASONING_MODEL, + ) + + assert _warnings_naming(caplog, "chat_template_kwargs") == () + + _CODEX_DEVELOPER_MESSAGE_AFTER_A_USER_TURN: Final = ( {"role": "system", "content": "You are Codex"}, {"role": "user", "content": "Hi there"}, diff --git a/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py b/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py index e5a77aa8d41..d35442d4095 100644 --- a/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py +++ b/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py @@ -1,7 +1,11 @@ +import logging +from typing import Final + import pytest import litellm +from litellm._logging import verbose_logger from litellm.llms.fireworks_ai.completion.transformation import ( @@ -204,3 +208,36 @@ def test_transform_text_completion_request_keeps_sdk_rejected_keys_in_extra_body assert "chat_template_kwargs" not in data["extra_body"] assert "best_of" not in data["extra_body"] assert "response_format" not in data + + +def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: + return tuple( + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + ) + + +def test_chat_template_kwargs_effort_dropped_for_non_reasoning_model_warns_once( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + FireworksAITextCompletionConfig().map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"reasoning_budget": 512}}}, + _NON_REASONING_MODEL, + ) + + warnings: Final = _warnings_naming(caplog, "chat_template_kwargs") + assert len(warnings) == 1 + assert "'reasoning_budget': 512" in warnings[0] + assert "llama-v3-70b-instruct" in warnings[0] + + +def test_chat_template_kwargs_thinking_off_for_non_reasoning_model_does_not_warn( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + FireworksAITextCompletionConfig().map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": False}}}, + _NON_REASONING_MODEL, + ) + + assert _warnings_naming(caplog, "chat_template_kwargs") == () diff --git a/tests/unit/llms/mistral/test_mistral_chat_transformation.py b/tests/unit/llms/mistral/test_mistral_chat_transformation.py index c340f77c317..720a1b4f458 100644 --- a/tests/unit/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/unit/llms/mistral/test_mistral_chat_transformation.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Mapping from typing import Final, List, cast from unittest.mock import MagicMock, patch @@ -9,6 +10,7 @@ import pytest import respx from openai import AsyncOpenAI +from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY from litellm.llms.mistral.chat.transformation import ( MistralChatResponseIterator, @@ -1670,6 +1672,109 @@ def test_mistral_transform_request_drops_an_empty_tool_message_whose_call_id_is_ assert [m.get("role") for m in request["messages"]] == ["user", "assistant"] +def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: + return tuple( + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + ) + + +@pytest.mark.parametrize( + "param, value", + ( + ("reasoning_effort", "none"), + ("reasoning_effort", "xhigh"), + ("thinking", {"type": "disabled"}), + ("thinking", {"type": ["enabled"]}), + ), +) +def test_magistral_reasoning_request_replaced_by_prompt_warns_once( + param: str, value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = MistralConfig().map_openai_params( + non_default_params={param: value}, + optional_params={}, + model="magistral-medium-latest", + drop_params=False, + ) + + warnings: Final = _warnings_naming(caplog, param) + assert result["_add_reasoning_prompt"] is True + assert len(warnings) == 1 + assert "magistral-medium-latest" in warnings[0] + + +@pytest.mark.parametrize( + "param, value", + ( + ("reasoning_effort", "low"), + ("reasoning_effort", "medium"), + ("reasoning_effort", "high"), + ("thinking", {"type": "enabled", "budget_tokens": 1024}), + ), +) +def test_magistral_reasoning_request_the_prompt_honors_does_not_warn( + param: str, value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = MistralConfig().map_openai_params( + non_default_params={param: value}, + optional_params={}, + model="magistral-medium-latest", + drop_params=False, + ) + + assert result["_add_reasoning_prompt"] is True + assert _warnings_naming(caplog, param) == () + + +def test_non_magistral_reasoning_effort_is_sent_without_warning(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = MistralConfig().map_openai_params( + non_default_params={"reasoning_effort": "high"}, + optional_params={}, + model="mistral-medium-latest", + drop_params=False, + ) + + assert result["reasoning_effort"] == "high" + assert _warnings_naming(caplog, "reasoning_effort") == () + + +@pytest.mark.parametrize("reasoning_effort", [5, {"effort": 5}, {"effort": 5, "summary": "auto"}]) +def test_mistral_warns_once_when_drop_params_drops_a_level_it_cannot_send( + caplog: pytest.LogCaptureFixture, reasoning_effort: object +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = MistralConfig().map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort}, + optional_params={}, + model="mistral-medium-latest", + drop_params=True, + ) + + assert "reasoning_effort" not in result + warnings: Final = _warnings_naming(caplog, "reasoning_effort") + assert len(warnings) == 1 + assert "model='mistral-medium-latest'" in warnings[0] + assert "(drop_params=True)" in warnings[0] + + +@pytest.mark.parametrize("reasoning_effort", [{"summary": "auto"}, {"effort": "high", "summary": "auto"}]) +def test_mistral_stays_quiet_under_drop_params_when_no_level_is_lost( + caplog: pytest.LogCaptureFixture, reasoning_effort: object +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + MistralConfig().map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort}, + optional_params={}, + model="mistral-medium-latest", + drop_params=True, + ) + + assert _warnings_naming(caplog, "reasoning_effort") == () + + @pytest.mark.parametrize( ("model", "custom_llm_provider"), [ diff --git a/tests/unit/llms/oci/chat/test_oci_chat_transformation.py b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py index 0535279d177..f22b5e32593 100644 --- a/tests/unit/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py @@ -1,4 +1,5 @@ import datetime +import logging from typing import Final import httpx import pytest @@ -1103,19 +1104,25 @@ class TestOCIReasoningEffort: ) assert chat_request.get("reasoningEffort") == "HIGH" - def test_reasoning_effort_unsupported_on_cohere_dropped(self): - """drop_params=True → silently drop reasoning_effort for Cohere.""" + def test_reasoning_effort_unsupported_on_cohere_dropped_with_one_warning(self, caplog): from litellm.llms.oci.chat.transformation import OCIChatConfig config = OCIChatConfig() - result = config.map_openai_params( - non_default_params={"reasoning_effort": "low"}, - optional_params={}, - model="cohere.command-latest", - drop_params=True, - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result = config.map_openai_params( + non_default_params={"reasoning_effort": "low"}, + optional_params={}, + model="cohere.command-latest", + drop_params=True, + ) + drop_warnings = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert "reasoning_effort" not in result assert "reasoningEffort" not in result + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in ("reasoning_effort='low'", "cohere.command-latest", "(drop_params=True)") + ) def test_reasoning_effort_unsupported_on_cohere_raises(self): """drop_params=False → raise rather than ship a payload Cohere will reject.""" @@ -2110,16 +2117,20 @@ class TestReasoningEffortGating: assert "reasoning_effort" not in OCIChatConfig().get_supported_openai_params("meta.llama-3.3-70b-instruct") @pytest.mark.parametrize("model", SERVED_GROK_4_20_MODELS) - def test_reasoning_effort_dropped_under_drop_params(self, local_model_cost_map, model): - result = OCIChatConfig().map_openai_params( - non_default_params={"reasoning_effort": "low", "temperature": 0.1}, - optional_params={}, - model=model, - drop_params=True, - ) + def test_reasoning_effort_dropped_under_drop_params_with_one_warning(self, local_model_cost_map, model, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result = OCIChatConfig().map_openai_params( + non_default_params={"reasoning_effort": "low", "temperature": 0.1}, + optional_params={}, + model=model, + drop_params=True, + ) + drop_warnings = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert "reasoningEffort" not in result assert "reasoning_effort" not in result assert result["temperature"] == 0.1 + assert len(drop_warnings) == 1 + assert all(fragment in drop_warnings[0] for fragment in ("reasoning_effort='low'", model, "(drop_params=True)")) def test_reasoning_effort_raises_without_drop_params_for_gated_grok(self, local_model_cost_map, monkeypatch): monkeypatch.setattr(litellm, "drop_params", False) diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index df9f5fd9d14..af1b85cf325 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -1,5 +1,6 @@ import asyncio import inspect +import logging import os import sys from typing import Final, cast @@ -23,6 +24,7 @@ import json from unittest.mock import ANY, MagicMock, patch import litellm +from litellm._logging import verbose_logger from litellm.types.utils import Choices, EmbeddingResponse, Message, ModelResponse, ModelResponseStream from unittest import mock @@ -1114,3 +1116,36 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) == "http://gateway.example.com/ollama/api/chat" ) + + +def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: + return tuple( + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + ) + + +@pytest.mark.parametrize("effort", ("minimal", "xhigh")) +def test_reasoning_effort_mapped_to_think_off_warns_once(effort: str, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = OllamaChatConfig().map_openai_params( + non_default_params={"reasoning_effort": effort}, optional_params={}, model="qwen3", drop_params=False + ) + + warnings: Final = _warnings_naming(caplog, "reasoning_effort") + assert result["think"] is False + assert len(warnings) == 1 + assert f"reasoning_effort={effort!r}" in warnings[0] + assert "qwen3" in warnings[0] + + +@pytest.mark.parametrize("effort, think", (("none", False), ("low", True), ("high", True))) +def test_reasoning_effort_with_a_matching_think_value_does_not_warn( + effort: str, think: bool, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = OllamaChatConfig().map_openai_params( + non_default_params={"reasoning_effort": effort}, optional_params={}, model="qwen3", drop_params=False + ) + + assert result["think"] is think + assert _warnings_naming(caplog, "reasoning_effort") == () diff --git a/tests/unit/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py index 248e745ebff..edea39bb10c 100644 --- a/tests/unit/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_completion_transformation.py @@ -1,6 +1,7 @@ import base64 import io import json +import logging import sys from typing import Final from unittest.mock import MagicMock, patch @@ -10,6 +11,7 @@ import pytest import respx import litellm +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.ollama.completion.transformation import ( @@ -990,3 +992,40 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) == "http://gateway.example.com/ollama/api/generate" ) + +def _reasoning_effort_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple( + r.getMessage() + for r in caplog.records + if r.levelno == logging.WARNING and "Dropping reasoning_effort=" in r.getMessage() + ) + + +@pytest.mark.parametrize("effort", ("minimal", "xhigh")) +def test_reasoning_effort_mapped_to_think_off_warns_once(effort: str, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = OllamaConfig().map_openai_params( + non_default_params={"reasoning_effort": effort}, optional_params={}, model="qwen3:0.6b", drop_params=False + ) + + warnings: Final = _reasoning_effort_warnings(caplog) + assert result["think"] is False + assert len(warnings) == 1 + assert f"reasoning_effort={effort!r}" in warnings[0] + assert "qwen3:0.6b" in warnings[0] + + +@pytest.mark.parametrize( + "model, effort, think", + (("qwen3:0.6b", "none", False), ("qwen3:0.6b", "high", True), ("gpt-oss:20b", "xhigh", "xhigh")), +) +def test_reasoning_effort_ollama_honors_does_not_warn( + model: str, effort: str, think: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = OllamaConfig().map_openai_params( + non_default_params={"reasoning_effort": effort}, optional_params={}, model=model, drop_params=False + ) + + assert result["think"] == think + assert _reasoning_effort_warnings(caplog) == () diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 46c14fcfd8f..f0865cca7f9 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -1,4 +1,6 @@ import asyncio +import json +import logging from typing import Final from unittest.mock import MagicMock, patch @@ -9,7 +11,11 @@ import respx import litellm from litellm.litellm_core_utils.token_counter import token_counter -from litellm.llms.openai.common_utils import BaseOpenAILLM, is_openai_backed_api_base +from litellm.llms.openai.common_utils import ( + BaseOpenAILLM, + drop_params_from_unprocessable_entity_error, + is_openai_backed_api_base, +) # Test parameters for different API functions API_FUNCTION_PARAMS = [ @@ -319,6 +325,104 @@ def _completion_kwargs(provider: str, client, **overrides) -> dict: } +_REASONING_EFFORT_422: Final = { + "error": { + "message": json.dumps( + {"detail": [{"loc": ["body", "reasoning_effort"], "msg": "Extra inputs are not permitted"}]} + ) + } +} +_CHAT_COMPLETION: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "gpt-6.1-sol", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +_CHAT_COMPLETION_STREAM: Final = ( + 'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,"model":"gpt-6.1-sol",' + '"choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}\n\n' + "data: [DONE]\n\n" +) + + +def _answer_422_while_reasoning_effort_is_sent(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + if "reasoning_effort" in body: + return httpx.Response(422, json=_REASONING_EFFORT_422) + if body.get("stream"): + return httpx.Response(200, text=_CHAT_COMPLETION_STREAM, headers={"content-type": "text/event-stream"}) + return httpx.Response(200, json=_CHAT_COMPLETION) + + +async def _complete_through_the_422_retry(call_shape: str) -> str: + transport: Final = httpx.MockTransport(_answer_422_while_reasoning_effort_is_sent) + kwargs: Final = { + "model": "gpt-6.1-sol", + "messages": [{"role": "user", "content": "hi"}], + "reasoning_effort": "high", + } + if call_shape == "sync": + sync_client: Final = openai.OpenAI(api_key="sk-not-a-real-key", http_client=httpx.Client(transport=transport)) + return litellm.completion(**kwargs, client=sync_client).choices[0].message.content + async_client: Final = openai.AsyncOpenAI( + api_key="sk-not-a-real-key", http_client=httpx.AsyncClient(transport=transport) + ) + if call_shape == "async": + return (await litellm.acompletion(**kwargs, client=async_client)).choices[0].message.content + stream: Final = await litellm.acompletion(**kwargs, client=async_client, stream=True) + return "".join([chunk.choices[0].delta.content or "" async for chunk in stream]) + + +@pytest.mark.parametrize("call_shape", ["sync", "async", "async_stream"]) +def test_a_422_naming_reasoning_effort_drops_it_with_one_warning( + call_shape: str, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + content: Final = asyncio.run(_complete_through_the_422_retry(call_shape)) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + + assert content == "ok" + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in ("reasoning_effort='high'", "gpt-6.1-sol", "(drop_params=True)", "422") + ) + + +def test_a_422_with_an_unhashable_loc_still_drops_reasoning_effort_with_one_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + detail: Final = [ + {"loc": ["body", ["messages", 0]], "msg": "Invalid"}, + {"loc": ["body", "reasoning_effort"], "msg": "Extra inputs are not permitted"}, + ] + + def answer(request: httpx.Request) -> httpx.Response: + if "reasoning_effort" in json.loads(request.content): + return httpx.Response(422, json={"error": {"message": json.dumps({"detail": detail})}}) + return httpx.Response(200, json=_CHAT_COMPLETION) + + monkeypatch.setattr(litellm, "drop_params", True) + client: Final = openai.OpenAI( + api_key="sk-not-a-real-key", http_client=httpx.Client(transport=httpx.MockTransport(answer)) + ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + response: Final = litellm.completion( + model="gpt-6.1-sol", + messages=[{"role": "user", "content": "hi"}], + reasoning_effort="high", + client=client, + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + + assert response.choices[0].message.content == "ok" + assert len(drop_warnings) == 1 + assert "reasoning_effort='high'" in drop_warnings[0] + + @pytest.mark.parametrize("provider", ["openai", "azure"]) def test_sync_output_limit_400_maps_to_length_truncated_response(provider): response = litellm.completion( @@ -588,3 +692,24 @@ async def test_azure_gateway_and_image_variations_use_the_callers_async_session( ) assert response.data[0].url == "https://example.com/image.png" assert not session.is_closed + + + +def test_reasoning_drop_warning_caps_a_large_caller_value(caplog: pytest.LogCaptureFixture) -> None: + error: Final = httpx.HTTPStatusError( + message="422", + request=httpx.Request("POST", "https://example.invalid/v1/chat/completions"), + response=httpx.Response(422, json=_REASONING_EFFORT_422), + ) + data: Final[dict[str, object]] = {"model": "gpt-6.1-sol", "messages": [], "reasoning_effort": "x" * 5000} + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + retried: Final = drop_params_from_unprocessable_entity_error( + error, data, model="gpt-6.1-sol", drop_params_gated=True + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + + assert "reasoning_effort" not in retried + assert len(drop_warnings) == 1 + assert len(drop_warnings[0]) < 300 + assert all(fragment in drop_warnings[0] for fragment in ("reasoning_effort='xxx", "'gpt-6.1-sol'", "422")) diff --git a/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py b/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py index 4d6d252ae6e..1bf646fc2ad 100644 --- a/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py +++ b/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py @@ -1,5 +1,10 @@ +import logging +from typing import Final from unittest.mock import patch +import pytest + +from litellm._logging import verbose_logger from litellm.llms.tencent.chat.transformation import TencentChatConfig @@ -208,6 +213,27 @@ class TestAdaptiveThinkingCoercion: assert result["extra_body"]["thinking"] == {"type": "adaptive", "budget_tokens": 4096} + @pytest.mark.parametrize("budget, warned", (("high", 1), (4096, 0)), ids=("non-integer", "integer")) + def test_enabled_thinking_on_adaptive_only_model_warns_only_when_the_budget_is_dropped( + self, local_model_cost_map, budget: object, warned: int, caplog: pytest.LogCaptureFixture + ): + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = TencentChatConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget}}, + optional_params={}, + model="tencent/minimax-m3", + drop_params=False, + ) + + warnings: Final = tuple( + r.getMessage() + for r in caplog.records + if r.levelno == logging.WARNING and f"Dropping thinking.budget_tokens={budget!r}" in r.getMessage() + ) + assert result["extra_body"]["thinking"]["type"] == "adaptive" + assert result["extra_body"]["thinking"].get("budget_tokens") == (None if warned else budget) + assert len(warnings) == warned + def test_disabled_thinking_kept_for_adaptive_only_model(self, local_model_cost_map): config = TencentChatConfig() result = config.map_openai_params( diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 67c78b62151..98d1b8f1300 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1693,6 +1693,78 @@ def test_vertex_ai_map_thinking_param_without_budget_tokens_for_gemini_3(): assert result["thinkingConfig"] == {"includeThoughts": True} +def _budget_drop_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple( + r.getMessage() + for r in caplog.records + if r.levelno == logging.WARNING and "thinking.budget_tokens" in r.getMessage() + ) + + +@pytest.mark.parametrize("budget", ("high", 1024.5, {"tokens": 1024})) +def test_gemini_non_integer_thinking_budget_drop_warns_once(budget: object, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget}}, + optional_params={}, + model="gemini-2.5-flash", + drop_params=False, + ) + + warnings: Final = _budget_drop_warnings(caplog) + assert "thinkingBudget" not in result["thinkingConfig"] + assert len(warnings) == 1 + assert f"thinking.budget_tokens={budget!r}" in warnings[0] + assert "gemini-2.5-flash" in warnings[0] + + +def test_gemini_integer_thinking_budget_is_sent_without_warning(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 2048}}, + optional_params={}, + model="gemini-2.5-flash", + drop_params=False, + ) + + assert result["thinkingConfig"]["thinkingBudget"] == 2048 + assert _budget_drop_warnings(caplog) == () + + +def test_gemini_3_thinking_budget_drop_warns_once(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 2048}}, + optional_params={}, + model="gemini-3-pro-preview", + drop_params=False, + ) + + warnings: Final = _budget_drop_warnings(caplog) + assert "thinkingBudget" not in result["thinkingConfig"] + assert len(warnings) == 1 + assert "thinking.budget_tokens=2048" in warnings[0] + assert "gemini-3-pro-preview" in warnings[0] + + +@pytest.mark.parametrize( + "thinking", + ({"type": "enabled"}, {"type": "enabled", "budget_tokens": 0}, {"type": "disabled", "budget_tokens": 2048}), +) +def test_gemini_3_thinking_without_a_dropped_budget_does_not_warn( + thinking: dict[str, object], caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + VertexGeminiConfig().map_openai_params( + non_default_params={"thinking": thinking}, + optional_params={}, + model="gemini-3-pro-preview", + drop_params=False, + ) + + assert _budget_drop_warnings(caplog) == () + + def test_vertex_ai_map_tools(): v = VertexGeminiConfig() optional_params = {} diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index aaec486b60f..9f2d2a9bbc5 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -10262,9 +10262,10 @@ def test_a_non_numeric_thinking_budget_on_adaptive_claude_is_rejected_with_a_400 @pytest.mark.parametrize("drop_globally", [False, True]) @pytest.mark.parametrize("model_info", [None, _OPUS_DECLARING_LOW_AND_HIGH]) @pytest.mark.parametrize("body_sent", [_anthropic_chat_body_sent, _bedrock_converse_request_fields_sent]) -def test_drop_params_drops_a_non_numeric_thinking_budget_on_adaptive_claude( +def test_drop_params_drops_a_non_numeric_thinking_budget_on_adaptive_claude_with_one_warning( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, body_sent: Callable[..., Mapping[str, object]], model_info: Mapping[str, object] | None, drop_globally: bool, @@ -10272,13 +10273,40 @@ def test_drop_params_drops_a_non_numeric_thinking_budget_on_adaptive_claude( monkeypatch.setattr(litellm, "drop_params", drop_globally) request_drop: Final = {} if drop_globally else {"drop_params": True} - dropped: Final = body_sent( - respx_mock, model_info, thinking={"type": "enabled", "budget_tokens": "lots"}, **request_drop - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + dropped: Final = body_sent( + respx_mock, model_info, thinking={"type": "enabled", "budget_tokens": "lots"}, **request_drop + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] adaptive_alone: Final = body_sent(respx_mock, model_info, thinking={"type": "adaptive"}, **request_drop) assert "output_config" not in adaptive_alone assert dropped == adaptive_alone + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in ("thinking.budget_tokens='lots'", "claude-opus-5-5", "(drop_params=True)") + ) + + +@pytest.mark.parametrize("model_info", [None, _KIMI_DEPLOYMENT]) +def test_drop_params_drops_a_non_numeric_thinking_budget_on_hosted_vllm_with_one_warning( + respx_mock: respx.MockRouter, caplog: pytest.LogCaptureFixture, model_info: Mapping[str, object] | None +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + sent: Final = _chat_body_sent( + respx_mock, + "hosted_vllm/my-kimi", + None, + model_info, + thinking={"type": "enabled", "budget_tokens": "lots"}, + drop_params=True, + ) + + assert not {"reasoning_effort", "thinking"} & sent.keys() + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + assert len(drop_warnings) == 1 + assert all(fragment in drop_warnings[0] for fragment in ("thinking.budget_tokens='lots'", "my-kimi")) @pytest.mark.parametrize("model", ["azure_ai/my-kimi", "hosted_vllm/my-kimi"]) From d18efdfe1362194edfb98d814c85307af35d2038 Mon Sep 17 00:00:00 2001 From: alvcol Date: Fri, 30 Oct 2026 16:57:48 +0100 Subject: [PATCH 31/64] fix(vllm): send the pass-through key as a bearer on both /vllm routes (#46861) * fix(vllm): send the deployment api_key on /vllm passthrough calls VLLMModelInfo.get_api_key always returned None, so /vllm/{endpoint} calls to a router model reached the vLLM server with no auth header even when the deployment had an api_key, and a server started with --api-key answered 401 (#23443). Return the deployment key (or VLLM_API_KEY) and send it as Authorization: Bearer, the header vllm serve --api-key checks. x-api-key was never actually sent before, since the key was always None. * fix(vllm): send the pass-through key as a bearer on both /vllm routes vllm serve --api-key only reads Authorization: Bearer, so a keyed vLLM answered 401 on both pass-through routes. Send the key as a bearer and keep x-api-key for gateways that read it. Router models send their own deployment key and never fall back to VLLM_API_KEY, which belongs to the /vllm route's own server. * fix(vllm): send only the bearer on router models and stop doubling /v1 on /vllm Router models now send the same Authorization: Bearer header /v1/chat/completions sends for that deployment, without the x-api-key copy that route never sent. The VLLM_API_BASE route keeps x-api-key next to the bearer, as it sent it before. A router model whose api_base already ends in /v1 no longer gets /v1/v1 on /vllm/v1 paths. * fix(vllm): send no placeholder key from keyless router models and stop doubling /v1 on the VLLM_API_BASE route * fix(vllm): send VLLM_API_KEY when a router model with no api_base is forwarded to VLLM_API_BASE * fix(vllm): keep a router model's own key on the VLLM_API_BASE server and resolve an oidc VLLM_API_KEY * fix(vllm): treat an empty key as no key so pass-through sends no empty bearer * fix(vllm): send no bearer for a Google oidc key and match VLLM_API_BASE across a trailing slash --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/vllm/common_utils.py | 20 +- .../llms/vllm/passthrough/transformation.py | 35 +++- .../llm_passthrough_endpoints.py | 3 +- tests/unit/passthrough/test_main.py | 175 ++++++++++++++++++ .../test_llm_pass_through_endpoints.py | 62 +++++++ 5 files changed, 289 insertions(+), 6 deletions(-) diff --git a/litellm/llms/vllm/common_utils.py b/litellm/llms/vllm/common_utils.py index 6484cdc88e3..014fabd5187 100644 --- a/litellm/llms/vllm/common_utils.py +++ b/litellm/llms/vllm/common_utils.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final import httpx @@ -10,6 +11,14 @@ from litellm.types.llms.openai import AllMessageValues from litellm.utils import add_path_to_api_base +def endpoint_under_api_base(endpoint: str, api_base: str | None) -> str: + return ( + endpoint.lstrip("/").removeprefix("v1/") + if api_base is not None and api_base.rstrip("/").endswith("/v1") + else endpoint + ) + + class VLLMError(BaseLLMException): def __init__( self, @@ -39,9 +48,14 @@ class VLLMModelInfo(BaseLLMModelInfo): api_key: str | None = None, api_base: str | None = None, ) -> dict: - if api_key is not None: - headers["x-api-key"] = api_key - return headers + upstream_key: Final = self.upstream_api_key(api_key, api_base) + return {**headers, **self.upstream_auth_headers(upstream_key)} if upstream_key else headers + + def upstream_api_key(self, api_key: str | None, api_base: str | None) -> str | None: + return api_key + + def upstream_auth_headers(self, api_key: str) -> Mapping[str, str]: + return {"Authorization": f"Bearer {api_key}", "x-api-key": api_key} @staticmethod def get_api_base(api_base: str | None = None) -> str | None: diff --git a/litellm/llms/vllm/passthrough/transformation.py b/litellm/llms/vllm/passthrough/transformation.py index e7da6e54410..2b253e7fa38 100644 --- a/litellm/llms/vllm/passthrough/transformation.py +++ b/litellm/llms/vllm/passthrough/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Final from httpx import Response +from litellm.constants import PLACEHOLDER_API_KEY from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.llms.openai_like.passthrough.transformation import ( @@ -12,8 +13,14 @@ from litellm.llms.openai_like.passthrough.transformation import ( logged_text_completion_stream, relays_text_completion, ) +from litellm.secret_managers.main import ( + GOOGLE_OIDC_SECRET_PREFIX, + UnresolvableSecretReference, + get_secret_str, + resolve_local_oidc_reference, +) -from ..common_utils import VLLMModelInfo +from ..common_utils import VLLMModelInfo, endpoint_under_api_base if TYPE_CHECKING: from httpx import URL @@ -22,6 +29,30 @@ if TYPE_CHECKING: class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig): + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return None if api_key == PLACEHOLDER_API_KEY else api_key + + def upstream_api_key(self, api_key: str | None, api_base: str | None) -> str | None: + if api_key is not None or not self._is_vllm_api_base(api_base): + return api_key + vllm_api_key: Final = get_secret_str("VLLM_API_KEY") + if vllm_api_key is None or vllm_api_key.startswith(GOOGLE_OIDC_SECRET_PREFIX): + return None + match resolve_local_oidc_reference(vllm_api_key): + case UnresolvableSecretReference(): + return None + case str() as resolved: + return resolved + + @staticmethod + def _is_vllm_api_base(api_base: str | None) -> bool: + vllm_api_base: Final = get_secret_str("VLLM_API_BASE") + return api_base is not None and vllm_api_base is not None and api_base.rstrip("/") == vllm_api_base.rstrip("/") + + def upstream_auth_headers(self, api_key: str) -> Mapping[str, str]: + return {"Authorization": f"Bearer {api_key}"} + def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return "stream" in request_data @@ -40,7 +71,7 @@ class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig): raise Exception("VLLM api base not found") return ( - self.format_url(endpoint, base_target_url, request_query_params), + self.format_url(endpoint_under_api_base(endpoint, base_target_url), base_target_url, request_query_params), base_target_url, ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 21078dfc848..8de7441121f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -72,6 +72,7 @@ from litellm.llms.openai.common_utils import is_openai_agents_api_path, openai_a from litellm.llms.openai.workload_identity import get_workload_identity_bearer_token_for_api_base from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, validate_oss_request from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.llms.vllm.common_utils import endpoint_under_api_base from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import enforced_model_allowlists @@ -600,7 +601,7 @@ async def vllm_proxy_route( ) return await llm_passthrough_factory_proxy_route( - endpoint=endpoint, + endpoint=endpoint_under_api_base(endpoint, get_secret_str("VLLM_API_BASE")), request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, diff --git a/tests/unit/passthrough/test_main.py b/tests/unit/passthrough/test_main.py index 78fa83d126b..4997d736d0a 100644 --- a/tests/unit/passthrough/test_main.py +++ b/tests/unit/passthrough/test_main.py @@ -51,3 +51,178 @@ async def test_allm_passthrough_route_hosted_vllm_sends_normalized_model(respx_m outbound: Final = json.loads(route.calls[0].request.content) assert outbound["model"] == "my-deployment" assert outbound["messages"] == [{"role": "user", "content": "Hello"}] + + +@pytest.mark.asyncio +async def test_allm_passthrough_route_hosted_vllm_sends_deployment_api_key(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + api_base=_API_BASE, + api_key="sk-upstream", + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert route.calls[0].request.headers["authorization"] == "Bearer sk-upstream" + assert "x-api-key" not in route.calls[0].request.headers + + +@pytest.mark.parametrize("endpoint", ["v1/chat/completions", "chat/completions"]) +@pytest.mark.asyncio +async def test_allm_passthrough_route_hosted_vllm_reaches_v1_paths_under_an_api_base_ending_in_v1( + respx_mock: respx.MockRouter, endpoint: str +) -> None: + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint=endpoint, + model="hosted_vllm/my-deployment", + api_base=f"{_API_BASE}/v1", + api_key="sk-upstream", + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert route.call_count == 1 + + +@pytest.mark.asyncio +async def test_allm_passthrough_route_hosted_vllm_sends_no_auth_for_a_keyless_deployment( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("VLLM_API_KEY", "sk-vllm-route-only") + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + api_base=_API_BASE, + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert "authorization" not in route.calls[0].request.headers + assert "x-api-key" not in route.calls[0].request.headers + + +@pytest.mark.parametrize( + ("vllm_api_key", "expected_authorization"), + [("sk-vllm-env", "Bearer sk-vllm-env"), ("oidc/env/VLLM_OIDC_TOKEN", "Bearer sk-vllm-oidc")], +) +@pytest.mark.asyncio +async def test_allm_passthrough_route_keyless_hosted_vllm_without_api_base_sends_the_vllm_api_base_key( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, vllm_api_key: str, expected_authorization: str +) -> None: + monkeypatch.setenv("VLLM_API_BASE", _API_BASE) + monkeypatch.setenv("VLLM_API_KEY", vllm_api_key) + monkeypatch.setenv("VLLM_OIDC_TOKEN", "sk-vllm-oidc") + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert route.calls[0].request.headers["authorization"] == expected_authorization + + +@pytest.mark.parametrize( + ("model_api_base", "vllm_api_base"), [(f"{_API_BASE}/", _API_BASE), (_API_BASE, f"{_API_BASE}/")] +) +@pytest.mark.asyncio +async def test_allm_passthrough_route_keyless_hosted_vllm_matches_the_vllm_api_base_across_a_trailing_slash( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, model_api_base: str, vllm_api_base: str +) -> None: + monkeypatch.setenv("VLLM_API_BASE", vllm_api_base) + monkeypatch.setenv("VLLM_API_KEY", "sk-vllm-env") + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post(url__regex=rf"^{_API_BASE}/+v1/chat/completions$").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + api_base=model_api_base, + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert route.calls[0].request.headers["authorization"] == "Bearer sk-vllm-env" + + +@pytest.mark.parametrize("vllm_api_key", ["", "oidc/google/audience"]) +@pytest.mark.asyncio +async def test_allm_passthrough_route_unusable_vllm_api_key_sends_no_authorization_header( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, vllm_api_key: str +) -> None: + monkeypatch.setenv("VLLM_API_BASE", _API_BASE) + monkeypatch.setenv("VLLM_API_KEY", vllm_api_key) + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert "authorization" not in route.calls[0].request.headers + + +@pytest.mark.parametrize( + ("api_base", "api_key", "expected_authorization"), + [ + (None, "sk-upstream", "Bearer sk-upstream"), + (None, "sk-hosted-env", "Bearer sk-hosted-env"), + (_API_BASE, "sk-hosted-env", "Bearer sk-hosted-env"), + (None, None, "Bearer sk-hosted-env"), + ], +) +@pytest.mark.asyncio +async def test_allm_passthrough_route_hosted_vllm_on_the_vllm_api_base_sends_its_own_resolved_key( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + api_base: str | None, + api_key: str | None, + expected_authorization: str, +) -> None: + monkeypatch.setenv("VLLM_API_BASE", _API_BASE) + monkeypatch.setenv("VLLM_API_KEY", "sk-vllm-env") + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.setenv("HOSTED_VLLM_API_KEY", "sk-hosted-env") + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + api_base=api_base, + api_key=api_key, + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert route.calls[0].request.headers["authorization"] == expected_authorization diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 806d88b2467..aab227e9494 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2759,6 +2759,68 @@ class TestVLLMProxyRoute: mock_factory_route.assert_awaited_once() +@pytest.mark.parametrize( + ("vllm_api_base", "request_path"), + [ + ("http://vllm-upstream.test:8090", "/vllm/v1/chat/completions"), + ("http://vllm-upstream.test:8090/v1", "/vllm/v1/chat/completions"), + ("http://vllm-upstream.test:8090/v1", "/vllm/chat/completions"), + ], +) +def test_vllm_provider_route_reaches_v1_paths_with_its_key_as_a_bearer( + monkeypatch: pytest.MonkeyPatch, vllm_api_base: str, request_path: str +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setenv("VLLM_API_BASE", vllm_api_base) + monkeypatch.setenv("VLLM_API_KEY", "sk-vllm-upstream") + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setitem( + proxy_server.app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual") + ) + + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post("http://vllm-upstream.test:8090/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"id": "chatcmpl-1"}) + ) + response: Final = TestClient(proxy_server.app).post( + request_path, + json={"model": "Qwen/Qwen3-0.6B", "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert (response.status_code, response.json()) == (200, {"id": "chatcmpl-1"}) + assert route.calls.last.request.headers["authorization"] == "Bearer sk-vllm-upstream" + assert route.calls.last.request.headers["x-api-key"] == "sk-vllm-upstream" + + +def test_vllm_provider_route_with_an_empty_key_sends_no_credential_headers(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setenv("VLLM_API_BASE", "http://vllm-upstream.test:8090") + monkeypatch.setenv("VLLM_API_KEY", "") + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setitem( + proxy_server.app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual") + ) + + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post("http://vllm-upstream.test:8090/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"id": "chatcmpl-1"}) + ) + response: Final = TestClient(proxy_server.app).post( + "/vllm/v1/chat/completions", + json={"model": "Qwen/Qwen3-0.6B", "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert (response.status_code, response.json()) == (200, {"id": "chatcmpl-1"}) + assert "authorization" not in route.calls.last.request.headers + assert "x-api-key" not in route.calls.last.request.headers + + class TestGigachatProxyRoute: @pytest.mark.asyncio @patch( # test-quality-ok: patching litellm internal for unit test isolation From b195d15b937ab5d09a368ed3972ba3e2cd48d19a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:28:19 -0700 Subject: [PATCH 32/64] fix(vertex_ai,azure_ai)!: auto-add the compaction, thinking display, and tool change betas on Claude routes (#44747) * fix(vertex_ai): forward and auto-add the compact-2026-09-04 beta on Claude routes * test(vertex_ai): prove the compact-2026-09-04 beta reaches the wire through both Claude handlers * fix(vertex_ai): move a client-sent body compact-2026-09-04 beta to the request header * refactor(vertex_ai): build the Claude beta set without mutation and inject the client in the messages wire test * fix(vertex_ai): keep client body betas in the body and answer 400 on a malformed anthropic_beta * fix(rust): allow the compact-2026-09-04 beta on vertex_ai in the Rust beta table The beta header config now maps compact-2026-09-04 for vertex_ai, so the Rust table that mirrors it must stop rejecting that beta there, and the Rust Vertex messages route adds it for a compaction request like the Python route does. * fix(vertex_ai): keep client body betas out of the Vertex body next to litellm betas, honor drop_params, and send the compaction beta on count_tokens Litellm's own betas replace a client's body anthropic_beta again, as on main, so a body beta Vertex rejects no longer turns a 200 into a 400. A malformed beta drops under drop_params and names its field otherwise. inline-tools-2026-09-15 joins compact-2026-09-04 as header-only, since Vertex rejects both in the body. Vertex count_tokens sends compact-2026-09-04 when the request carries compaction. * fix(logging): log and bill Anthropic messages compaction responses on the Rust route The Rust route hands the served Anthropic body to the logger, which validates it against AnthropicResponse. That model had no compaction content block and no compaction stop reason, so a compaction reply raised and wrote no spend row. * fix(vertex_ai,azure_ai): forward compact-2026-09-04 to Azure AI Foundry, merge Vertex body betas, and redact compaction summaries Azure AI Foundry needs the compact-2026-09-04 header for compaction, a replayed compaction block, and count tokens, so the beta map, the native compaction providers, and the Rust beta table now allow it there. Vertex keeps a client's body betas next to the ones litellm adds, as the Anthropic route does, and turn_off_message_logging now redacts compaction summaries too * fix(vertex_ai,azure_ai): auto-add thinking display and tool change betas, redact streamed compaction summaries Folds the header changes from #44103 into this PR. Vertex AI and Azure AI Claude routes now add thinking-display-updates-2026-08-18 when thinking.display is "updates" and mid-conversation-tool-changes-2026-07-01 when a tool_addition block is present, matching the Anthropic and Bedrock routes. With turn_off_message_logging on, compaction summaries from streamed messages and streamed chat completions are now redacted in spend logs and callback payloads, like the non-streamed response already was. * fix(vertex_ai,azure_ai,rust): add tool change betas on count tokens and Rust messages, bill Rust server tool replies Vertex and Azure count tokens now add mid-conversation-tool-changes-2026-07-01 when the history carries a tool_addition or tool_removal block, through one shared helper with the compaction beta. The Rust /v1/messages route adds thinking-display-updates-2026-08-18 and mid-conversation-tool-changes-2026-07-01 when a request uses them, on the providers that accept each beta. Rust-route logging parses the served body with the Anthropic response parser, so web_search, pause_turn, mcp_tool_use and container_upload replies get a spend row. Legacy enabled thinking keeps display updates on anthropic too, and the Vertex transform reads drop_params through drops_unsupported_params. * style(rust): rewrap beta map test line with cargo fmt * test(rust): expect the display beta on Databricks, which main's beta list now accepts --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../src/providers/anthropic/beta.rs | 34 +- .../src/anthropic/messages/transformation.rs | 99 +++++- .../src/azure_ai/messages/transformation.rs | 5 + .../src/vertex_ai/messages/transformation.rs | 12 + litellm/anthropic_beta_headers_config.json | 9 +- litellm/anthropic_beta_headers_manager.py | 39 ++- litellm/litellm_core_utils/litellm_logging.py | 15 +- litellm/litellm_core_utils/redact_messages.py | 1 + litellm/llms/anthropic/common_utils.py | 41 ++- .../messages/native/transformation.py | 3 +- .../anthropic/count_tokens/handler.py | 1 + .../anthropic/count_tokens/transformation.py | 8 +- .../anthropic/messages/transformation.py | 30 +- .../anthropic/transformation.py | 82 ++++- .../count_tokens/handler.py | 25 +- litellm/types/llms/anthropic.py | 15 +- .../test_litellm_logging.py | 89 ++++++ .../test_redact_messages.py | 5 +- .../anthropic/test_anthropic_common_utils.py | 61 +++- .../count_tokens/test_token_counter.py | 53 +++ ...azure_anthropic_messages_transformation.py | 88 +++++ .../test_azure_anthropic_transformation.py | 84 +++++ ...artner_models_anthropic_messages_config.py | 195 ++++++++++- ...partner_models_anthropic_transformation.py | 302 +++++++++++++++++- .../count_tokens/test_handler.py | 60 ++++ .../test_anthropic_beta_headers_filtering.py | 51 ++- 26 files changed, 1323 insertions(+), 84 deletions(-) create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs index 80095cac36c..01ed023999d 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs @@ -199,7 +199,7 @@ impl AnthropicBeta { &[BetaProvider::BedrockConverse, BetaProvider::Databricks] } Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse], - Self::AdvisorTool20260301 | Self::Compact20260904 => &[ + Self::AdvisorTool20260301 => &[ BetaProvider::AzureAi, BetaProvider::BedrockConverse, BetaProvider::Bedrock, @@ -207,6 +207,12 @@ impl AnthropicBeta { BetaProvider::VertexAi, BetaProvider::Databricks, ], + Self::Compact20260904 => &[ + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::BedrockMantle, + BetaProvider::Databricks, + ], Self::Bash20241022 | Self::Bash20250124 | Self::McpServers20251204 @@ -256,14 +262,18 @@ impl AnthropicBeta { Self::InterleavedThinking20250514 | Self::WebSearch20250305 => { &[BetaProvider::BedrockConverse, BetaProvider::Bedrock] } - Self::MidConversationOutputConfig20260701 - | Self::MidConversationToolChanges20260701 => &[ + Self::MidConversationOutputConfig20260701 => &[ BetaProvider::AzureAi, BetaProvider::BedrockConverse, BetaProvider::BedrockMantle, BetaProvider::VertexAi, BetaProvider::Databricks, ], + Self::MidConversationToolChanges20260701 => &[ + BetaProvider::BedrockConverse, + BetaProvider::BedrockMantle, + BetaProvider::Databricks, + ], Self::Output128k20250219 | Self::TokenEfficientTools20250219 => &[ BetaProvider::AzureAi, BetaProvider::BedrockConverse, @@ -483,7 +493,7 @@ mod tests { let config = beta_headers_config(); let config_columns: Vec = config .keys() - .filter(|column| column.as_str() != "description") + .filter(|column| column.as_str() != "description" && !column.ends_with("_header_only")) .cloned() .collect(); let beta_provider_columns: Vec = BetaProvider::VARIANTS @@ -502,12 +512,16 @@ mod tests { let config = beta_headers_config(); for provider in BetaProvider::VARIANTS { for beta in &AnthropicBeta::KNOWN { - let expected = config - .get(provider.as_ref()) - .and_then(serde_json::Value::as_object) - .and_then(|column| column.get(beta.as_str())) - .and_then(serde_json::Value::as_str) - .map(str::to_string); + let cell = |column: &str| { + config + .get(column) + .and_then(serde_json::Value::as_object) + .and_then(|values| values.get(beta.as_str())) + .and_then(serde_json::Value::as_str) + .map(str::to_string) + }; + let expected = + cell(&format!("{provider}_header_only")).or_else(|| cell(provider.as_ref())); let actual = beta.on(*provider).map(|name| name.to_string()); assert_eq!(actual, expected, "provider {provider}, beta {beta}"); if matches!(beta, AnthropicBeta::AdvancedToolUse20251120) diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 6657cbf262f..fb6d7259bf9 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -1,7 +1,8 @@ use litellm_auth::CredentialPlacement; use litellm_llms_types::{ formats::messages::{ - ContextEdit, ContextManagement, Message, MessagesOptionalParams, MessagesRequest, Speed, + ContentBlock, ContentBlockType, ContextEdit, ContextManagement, Message, + MessagesOptionalParams, MessagesRequest, Speed, ThinkingConfig, ThinkingDisplay, }, providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, DEFAULT_HEADERS}, recognized::Recognized, @@ -203,6 +204,10 @@ fn feature_betas(request: &MessagesRequest) -> BetaSet { .then_some(AnthropicBeta::FastMode20260201), messages_carry_output_config(&request.messages) .then_some(AnthropicBeta::PerTurnControl20260701), + uses_thinking_display_updates(params) + .then_some(AnthropicBeta::ThinkingDisplayUpdates20260818), + messages_carry_tool_changes(&request.messages) + .then_some(AnthropicBeta::MidConversationToolChanges20260701), has_advisor_tool(tools).then_some(AnthropicBeta::AdvisorTool20260301), is_tool_search_used(tools).then_some(AnthropicBeta::AdvancedToolUse20251120), ] @@ -253,6 +258,36 @@ fn messages_carry_output_config(messages: &[Message]) -> bool { .any(|message| message.extra.contains_key("output_config")) } +fn uses_thinking_display_updates(params: &MessagesOptionalParams) -> bool { + let display = match params.thinking.as_ref().and_then(Recognized::known) { + Some(ThinkingConfig::Enabled(thinking)) => thinking.display.as_ref(), + Some(ThinkingConfig::Adaptive(thinking)) => thinking.display.as_ref(), + _ => None, + }; + display.and_then(Recognized::known) == Some(&ThinkingDisplay::Updates) +} + +fn is_tool_reference_change(block: &ContentBlock) -> bool { + matches!( + &block.block_type, + Some(ContentBlockType::Other(block_type)) + if block_type == "tool_addition" || block_type == "tool_removal" + ) && block + .extra + .get("tool") + .and_then(|tool| tool.get("type")) + .and_then(Value::as_str) + == Some("tool_reference") +} + +fn messages_carry_tool_changes(messages: &[Message]) -> bool { + messages + .iter() + .filter(|message| message.role == "system") + .flat_map(Message::blocks) + .any(is_tool_reference_change) +} + fn unsupported_param(model: &str, param: &str, value: &str, hint: &str) -> Error { Error::InvalidRequest(crate::ErrorDetail::UnsupportedParameter { model: model.into(), @@ -1071,6 +1106,41 @@ mod tests { json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}]}), &[] )] + #[case::adaptive_display_updates( + json!({"thinking": {"type": "adaptive", "display": "updates"}}), + &["thinking-display-updates-2026-08-18"] + )] + #[case::enabled_display_updates( + json!({"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "updates"}}), + &["thinking-display-updates-2026-08-18"] + )] + #[case::summarized_display(json!({"thinking": {"type": "adaptive", "display": "summarized"}}), &[])] + #[case::disabled_thinking(json!({"thinking": {"type": "disabled"}}), &[])] + #[case::system_tool_addition( + json!({"messages": [ + {"role": "user", "content": "Use the ping tool"}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ]}), + &["mid-conversation-tool-changes-2026-07-01"] + )] + #[case::system_tool_removal( + json!({"messages": [ + {"role": "system", "content": [{"type": "tool_removal", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ]}), + &["mid-conversation-tool-changes-2026-07-01"] + )] + #[case::user_tool_addition( + json!({"messages": [ + {"role": "user", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ]}), + &[] + )] + #[case::tool_addition_without_a_tool_reference( + json!({"messages": [ + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_definition"}}]}, + ]}), + &[] + )] fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { assert_eq!(feature_betas(&request(fields)), betas(expected)); } @@ -1093,6 +1163,33 @@ mod tests { ); } + #[rstest] + #[case::vertex_ai( + BetaProvider::VertexAi, + &["mid-conversation-tool-changes-2026-07-01", "thinking-display-updates-2026-08-18"] + )] + #[case::azure_ai( + BetaProvider::AzureAi, + &["mid-conversation-tool-changes-2026-07-01", "thinking-display-updates-2026-08-18"] + )] + #[case::bedrock_converse(BetaProvider::BedrockConverse, &["thinking-display-updates-2026-08-18"])] + #[case::databricks(BetaProvider::Databricks, &["thinking-display-updates-2026-08-18"])] + fn tool_change_and_display_betas_reach_the_providers_that_accept_them( + #[case] provider: BetaProvider, + #[case] expected: &[&str], + ) { + let fields = json!({ + "thinking": {"type": "adaptive", "display": "updates"}, + "messages": [ + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ], + }); + assert_eq!( + provider_feature_betas(&request(fields), provider), + betas(expected) + ); + } + #[rstest] fn hosts_without_a_beta_map_column_get_safeguards_without_a_derived_beta() { let client = headers(&[("anthropic-beta", "claude-code-20250219")]); diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index bf934fb4d4c..d578acea68c 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -545,6 +545,11 @@ mod tests { &[("anthropic-beta", "web-search-2025-03-05")], &[("x-api-key", "k"), ("anthropic-beta", "structured-outputs-2025-11-13,web-search-2025-03-05")] )] + #[case::compaction_needs_the_compact_2026_09_04_beta( + json!({"compaction": {"type": "summarize"}}), + &[], + &[("x-api-key", "k"), ("anthropic-beta", "compact-2026-09-04")] + )] #[case::no_feature_needs_a_beta(json!({}), &[], &[("x-api-key", "k")])] #[case::safeguards_need_the_dangerous_tool_use_beta( json!({"safeguards": [{"type": "dangerous_tool_use"}]}), diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs index 46e5a68dcb8..e5e3a056fab 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -798,6 +798,18 @@ mod tests { ]}}), &["compact-2026-01-12", "context-management-2025-06-27"] )] + #[case::compaction(json!({"compaction": {"type": "summarize"}}), &["compact-2026-09-04"])] + #[case::thinking_display_updates( + json!({"thinking": {"type": "adaptive", "display": "updates"}}), + &["thinking-display-updates-2026-08-18"] + )] + #[case::tool_addition( + json!({"messages": [ + {"role": "user", "content": "Use the ping tool"}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]} + ]}), + &["mid-conversation-tool-changes-2026-07-01"] + )] #[case::advisor_tool_is_not_a_vertex_beta( json!({"tools": [{"type": "advisor_20260301", "name": "advisor"}]}), &[] diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index cee5ce15469..253074b0e1f 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -1,5 +1,5 @@ { - "description": "Mapping of Anthropic beta headers for each provider. Keys are input header names, values are provider-specific header names (or null if unsupported). Only headers present in mapping keys with non-null values can be forwarded.", + "description": "Mapping of Anthropic beta headers for each provider. Keys are input header names, values are provider-specific header names (or null if unsupported). Only headers present in mapping keys with non-null values can be forwarded. A _header_only column lists betas that provider accepts only as the anthropic-beta HTTP header: they are forwarded there and dropped from the anthropic_beta body field. Proxies that predate these columns ignore them and read the provider column alone.", "anthropic": { "advisor-tool-2026-03-01": "advisor-tool-2026-03-01", "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20", @@ -46,6 +46,7 @@ "bash_20250124": null, "code-execution-2025-08-25": "code-execution-2025-08-25", "compact-2026-01-12": null, + "compact-2026-09-04": "compact-2026-09-04", "computer-use-2025-01-24": "computer-use-2025-01-24", "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", @@ -59,6 +60,7 @@ "mcp-client-2025-11-20": "mcp-client-2025-11-20", "mcp-client-2025-04-04": "mcp-client-2025-04-04", "mcp-servers-2025-12-04": null, + "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", "output-128k-2025-02-19": null, "structured-output-2024-03-01": null, "per-turn-control-2026-07-01": "per-turn-control-2026-07-01", @@ -202,6 +204,7 @@ "mcp-client-2025-11-20": null, "mcp-client-2025-04-04": null, "mcp-servers-2025-12-04": null, + "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", "output-128k-2025-02-19": null, "structured-output-2024-03-01": null, "per-turn-control-2026-07-01": "per-turn-control-2026-07-01", @@ -218,6 +221,10 @@ "web-search-2025-03-05": "web-search-2025-03-05", "inline-tools-2026-09-15": "inline-tools-2026-09-15" }, + "vertex_ai_header_only": { + "compact-2026-09-04": "compact-2026-09-04", + "inline-tools-2026-09-15": "inline-tools-2026-09-15" + }, "databricks": { "advisor-tool-2026-03-01": null, "advanced-tool-use-2025-11-20": null, diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index 07730a0fbf8..be2ff07e088 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -30,7 +30,7 @@ from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from functools import cache from importlib.resources import files -from typing import Final, TypeAlias +from typing import Final, Literal, TypeAlias import httpx from pydantic import TypeAdapter @@ -47,6 +47,9 @@ BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]] _BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig) _DECODED_JSON: Final = TypeAdapter(object) +BetaPlacement: TypeAlias = Literal["header", "body"] +_HEADER_ONLY_COLUMN_SUFFIX: Final = "_header_only" + class GetAnthropicBetaHeadersConfig: """ @@ -241,6 +244,31 @@ def get_provider_name(provider: str) -> str: return aliases.get(provider, provider) +def _beta_column(name: str) -> Mapping[str, str | None]: + column: Final[Mapping[str, str | None]] = _load_beta_headers_config().get(name, {}) + return column + + +def _header_only_column(canonical_provider: str) -> Mapping[str, str | None]: + return _beta_column(f"{canonical_provider}{_HEADER_ONLY_COLUMN_SUFFIX}") + + +def _provider_beta_mapping(provider: str, placement: BetaPlacement) -> Mapping[str, str | None]: + canonical_provider: Final = get_provider_name(provider) + column: Final = _beta_column(canonical_provider) + header_only: Final = _header_only_column(canonical_provider) + match placement: + case "header": + return {**column, **header_only} + case "body": + return {beta: name for beta, name in column.items() if beta not in header_only} + + +def get_header_only_betas(provider: str) -> frozenset[str]: + header_only: Final = _header_only_column(get_provider_name(provider)) + return frozenset(beta for beta, name in header_only.items() if name is not None) + + def has_beta_headers_config(provider: str) -> bool: config: Final = _bundled_beta_headers_config() aliases: Final = config.get("provider_aliases", {}) @@ -251,6 +279,7 @@ def has_beta_headers_config(provider: str) -> bool: def filter_and_transform_beta_headers( beta_headers: list[str], provider: str, + placement: BetaPlacement = "header", ) -> list[str]: """ Filter and transform beta headers based on provider's mapping configuration. @@ -263,6 +292,7 @@ def filter_and_transform_beta_headers( Args: beta_headers: List of Anthropic beta header values provider: Provider name (e.g., "anthropic", "bedrock", "vertex_ai") + placement: Whether the betas travel as the anthropic-beta header or the anthropic_beta body field Returns: List of filtered and transformed beta headers for the provider @@ -270,11 +300,7 @@ def filter_and_transform_beta_headers( if not beta_headers: return [] - config: Final = _load_beta_headers_config() - provider = get_provider_name(provider) - - # Get the header mapping for this provider - provider_mapping: Final = config.get(provider, {}) + provider_mapping: Final = _provider_beta_mapping(provider, placement) filtered_headers: Final[set[str]] = set() @@ -426,6 +452,7 @@ def update_request_with_filtered_beta( filtered_body_betas: Final = filter_and_transform_beta_headers( beta_headers=existing_body_betas, provider=provider, + placement="body", ) if filtered_body_betas: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index cfe8c1739c6..26949d9bbfb 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -569,6 +569,7 @@ def _resolve_mantle_region_for_cost( _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) +_SERVED_ANTHROPIC_BODY: Final = TypeAdapter(dict[str, object]) _REDACTED_RESULT: Final = TypeAdapter(object) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" @@ -5067,13 +5068,15 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.warning("LiteLLM: the anthropic_messages stream assembled no response, logging an empty one") return litellm.ModelResponse(model=self.model) else: - from litellm.types.llms.anthropic import AnthropicResponse - - pydantic_result: Final = AnthropicResponse.model_validate({"stop_sequence": None, **(served_body or {})}) - import httpx - + anthropic_body: Final = ( + served_body + if served_body is not None + else _SERVED_ANTHROPIC_BODY.validate_python( + result.model_dump() if isinstance(result, BaseModel) else result + ) + ) result = litellm.AnthropicConfig().transform_parsed_response( - completion_response=pydantic_result.model_dump(), + completion_response={"stop_sequence": None, **anthropic_body}, raw_response=httpx.Response( status_code=200, headers={}, diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 548d088d833..981859c620b 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -249,6 +249,7 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ "reasoningContentBlocks", "citations", "compaction_blocks", + "compaction_start", "compaction_delta", ) if key in fields diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 6d1ac24af55..e1f9a7d540f 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -58,6 +58,7 @@ from litellm.types.llms.anthropic import ( ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER, ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, ANTHROPIC_TOOL_SEARCH_TOOL_TYPES, + LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS, AllAnthropicToolsValues, AnthropicMcpServerTool, AnthropicMessagesToolChoice, @@ -126,6 +127,8 @@ _JSON_VALUE: Final = TypeAdapter(JsonValue) _CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/") +NATIVE_COMPACTION_PROVIDERS: Final = frozenset({"anthropic", "azure_ai", "vertex_ai"}) + def is_claude_tool_target(custom_llm_provider: str | None, model: str) -> bool: if custom_llm_provider == "anthropic": @@ -204,7 +207,7 @@ def requires_native_compaction_beta( optional_params: Mapping[str, object], messages: Sequence[object], ) -> bool: - return custom_llm_provider == "anthropic" and ( + return custom_llm_provider in NATIVE_COMPACTION_PROVIDERS and ( optional_params.get("compaction") is not None or carries_signed_compaction_block(messages) ) @@ -246,6 +249,27 @@ def context_management_betas(context_management: JsonValue) -> tuple[str, ...]: ) +def count_tokens_feature_betas( + custom_llm_provider: str, + optional_params: Mapping[str, object], + messages: Sequence[object], +) -> tuple[str, ...]: + return tuple( + beta + for beta, used in ( + ( + ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value, + requires_native_compaction_beta(custom_llm_provider, optional_params, messages), + ), + ( + ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, + AnthropicModelInfo().is_mid_conversation_tool_change_used(messages), + ), + ) + if used + ) + + def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None) -> bool: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.utils import supports_prompt_caching @@ -645,10 +669,12 @@ class AnthropicModelInfo(BaseLLMModelInfo): file_ids: Final = get_file_ids_from_messages(messages) return len(file_ids) > 0 - def is_thinking_display_updates_used(self, thinking: AnthropicThinkingParam | None) -> bool: - if not isinstance(thinking, dict): - return False - return thinking.get("type") in ("adaptive", "enabled") and thinking.get("display") == "updates" + def is_thinking_display_updates_used(self, thinking: object) -> bool: + match thinking: + case {"type": "adaptive" | "enabled", "display": "updates"}: + return True + case _: + return False def thinking_betas(self, thinking: AnthropicThinkingParam | None) -> tuple[str, ...]: return ( @@ -1130,7 +1156,10 @@ class AnthropicModelInfo(BaseLLMModelInfo): budget_tokens: Final = _legacy_budget_tokens(thinking.get("budget_tokens", 0), drop_params=drop_params) display: Final = thinking.get("display") - supported_display: Final = {"display": display} if display in ("summarized", "omitted") else {} + keeps_display: Final = display in ("summarized", "omitted") or ( + display == "updates" and custom_llm_provider in LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS + ) + supported_display: Final = {"display": display} if keeps_display else {} carried_keys: Final = carried_thinking_keys( thinking, rebuilt_keys=frozenset(("type", "budget_tokens", "display")), drop_params=drop_params ) diff --git a/litellm/llms/anthropic/messages/native/transformation.py b/litellm/llms/anthropic/messages/native/transformation.py index c148a0defe8..5281ad99620 100644 --- a/litellm/llms/anthropic/messages/native/transformation.py +++ b/litellm/llms/anthropic/messages/native/transformation.py @@ -41,6 +41,7 @@ from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header from litellm.types.router import GenericLiteLLMParams from ...common_utils import ( + NATIVE_COMPACTION_PROVIDERS, OBJECT_LIST_ADAPTER, AnthropicError, AnthropicModelInfo, @@ -128,7 +129,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): "tool_choice", "thinking", "context_management", - *(("compaction",) if self._resolved_provider == "anthropic" else ()), + *(("compaction",) if self._resolved_provider in NATIVE_COMPACTION_PROVIDERS else ()), "output_format", "inference_geo", "speed", diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 4b3f263201a..c324f0f0980 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -84,6 +84,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): headers: Final = self.get_required_headers( api_key=api_key, litellm_params=litellm_params, + messages=messages, ) if not {"x-api-key", "Authorization"} & headers.keys(): raise AnthropicError( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index 6e3cfce5eb4..c02c92a95d6 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -4,11 +4,13 @@ Azure AI Anthropic CountTokens API transformation logic. Extends the base Anthropic CountTokens transformation with Azure authentication. """ +from collections.abc import Sequence from typing import Final from urllib.parse import urlsplit from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.anthropic.common_utils import count_tokens_feature_betas, merge_anthropic_beta_headers from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) @@ -28,6 +30,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): self, api_key: str | None, litellm_params: dict[str, object] | None = None, + messages: Sequence[object] = (), ) -> dict[str, str]: """ Get the required headers for the Azure AI Anthropic CountTokens API. @@ -38,6 +41,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): Args: api_key: The Azure AI API key litellm_params: Optional LiteLLM parameters for additional auth config + messages: The request messages, read for betas the counted features need Returns: Dictionary of required headers with both x-api-key and Azure authentication @@ -46,7 +50,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): headers: Final = { "Content-Type": "application/json", "anthropic-version": "2023-06-01", - "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, + "anthropic-beta": merge_anthropic_beta_headers( + ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, count_tokens_feature_betas("azure_ai", {}, messages) + ), **({"x-api-key": api_key} if api_key else {}), } diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py index da82766c2b9..98fa8ed2b77 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py @@ -1,9 +1,15 @@ from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter + from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.core_helpers import drops_unsupported_params -from litellm.llms.anthropic.common_utils import AnthropicModelInfo, context_management_betas +from litellm.llms.anthropic.common_utils import ( + AnthropicModelInfo, + context_management_betas, + requires_native_compaction_beta, +) from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, messages_carry_output_config, @@ -11,6 +17,7 @@ from litellm.llms.anthropic.messages.native.transformation import ( from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_HOSTED_TOOLS, + ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, ) from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header from litellm.types.llms.vertex_ai import VertexPartnerProvider @@ -19,6 +26,8 @@ from litellm.types.router import GenericLiteLLMParams from ....vertex_llm_base import VertexBase from ..output_params_utils import sanitize_vertex_anthropic_output_params +_OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object]) + _VERSION_HEADER_VERTEX_READS_FROM_BODY: Final = "anthropic-version" _VERTEX_DEFAULT_VERSION: Final = "vertex-2023-10-16" _VERSIONS_SERVED_AS_VERTEX_DEFAULT: Final = frozenset({"", "2023-01-01", "2023-06-01"}) @@ -190,14 +199,23 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert if messages_carry_output_config(messages): beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value) - beta_values.update( + typed_params: Final = _OPTIONAL_PARAMS.validate_python(optional_params) + if requires_native_compaction_beta("vertex_ai", typed_params, messages): + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value) + + tool_change_betas: Final = ( + (ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,) + if anthropic_model_info.is_mid_conversation_tool_change_used(messages) + else () + ) + all_beta_values: Final = beta_values.union( anthropic_model_info.thinking_betas( optional_params.get("thinking"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # optional_params is the base config's bare dict - ) + ), + tool_change_betas, ) - - if beta_values: - request_headers["anthropic-beta"] = ",".join(beta_values) + if all_beta_values: + request_headers["anthropic-beta"] = ",".join(sorted(all_beta_values)) return request_headers, api_base diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 79fd9f9e8b1..3d1690fca33 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -2,19 +2,23 @@ ## Handler file for calling claude-3 on vertex ai import base64 import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final import httpx -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm +from litellm.anthropic_beta_headers_manager import get_header_only_betas +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, inline_remote_image_urls from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse from ....anthropic.chat.transformation import AnthropicConfig +from ....anthropic.common_utils import AnthropicError, requires_native_compaction_beta from .output_params_utils import sanitize_vertex_anthropic_output_params if TYPE_CHECKING: @@ -36,6 +40,32 @@ def _vertex_labels_header(labels: object, header_names: Mapping[str, object]) -> return base64.b64encode(json.dumps(string_labels).encode()).decode() if string_labels else None +_BETA_VALUE: Final[TypeAdapter[str | tuple[str, ...] | None]] = TypeAdapter(str | tuple[str, ...] | None) + +_REQUEST_MESSAGES: Final = TypeAdapter(Sequence[object]) + + +def _beta_names(value: object) -> frozenset[str] | None: + try: + parsed: Final = _BETA_VALUE.validate_python(value) + except ValidationError: + return None + names: Final = parsed.split(",") if isinstance(parsed, str) else parsed or () + return frozenset(name.strip() for name in names if name.strip()) + + +def _client_betas(value: object, field: str, drop_params: bool) -> frozenset[str]: + names: Final = _beta_names(value) + if names is not None: + return names + if drop_params: + return frozenset() + raise AnthropicError( + status_code=400, + message=f"{field} must be a string or a list of strings, got {type(value).__name__}", + ) + + class VertexAIError(Exception): def __init__(self, status_code, message): self.status_code = status_code @@ -157,21 +187,43 @@ class VertexAIAnthropicConfig(AnthropicConfig): if context_management: self._add_context_management_beta_headers(beta_set, context_management) - extra_headers: Final = optional_params.get("extra_headers") or {} - anthropic_beta_value: Final = extra_headers.get("anthropic-beta", "") - if isinstance(anthropic_beta_value, str) and anthropic_beta_value: - for beta in anthropic_beta_value.split(","): - beta = beta.strip() - if beta: - beta_set.add(beta) - elif isinstance(anthropic_beta_value, list): - beta_set.update(anthropic_beta_value) - + request_params: Final = _NON_DEFAULT_PARAMS.validate_python(optional_params) + compaction_betas: Final = ( + frozenset({ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value}) + if requires_native_compaction_beta( + "vertex_ai", request_params, _REQUEST_MESSAGES.validate_python(data["messages"]) + ) + else frozenset[str]() + ) + drop_params: Final = drops_unsupported_params(litellm_params.get("drop_params")) + extra_headers: Final = _NON_DEFAULT_PARAMS.validate_python(request_params.get("extra_headers") or {}) data.pop("extra_headers", None) + injected_betas: Final = frozenset(beta_set) | _client_betas( + extra_headers.get("anthropic-beta"), "The anthropic-beta header", drop_params + ) + client_body_betas: Final = _client_betas( + _NON_DEFAULT_PARAMS.validate_python(data).get("anthropic_beta"), "anthropic_beta", drop_params + ) - if beta_set: - data["anthropic_beta"] = list(beta_set) - headers["anthropic-beta"] = ",".join(beta_set) + requested_betas: Final = injected_betas | client_body_betas + header_only_betas: Final = get_header_only_betas("vertex_ai") if requested_betas else frozenset[str]() + body_betas: Final = requested_betas - header_only_betas + if body_betas: + data["anthropic_beta"] = sorted(body_betas) + elif "anthropic_beta" in data: + del data["anthropic_beta"] + header_betas: Final = ( + injected_betas + | compaction_betas + | (client_body_betas & header_only_betas) + | _client_betas( + _NON_DEFAULT_PARAMS.validate_python(headers).get("anthropic-beta"), + "The anthropic-beta header", + drop_params, + ) + ) + if header_betas: + headers["anthropic-beta"] = ",".join(sorted(header_betas)) return data diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 180761f15d9..e341c8f5f48 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -6,12 +6,19 @@ Unlike Gemini models which use Google's token counting API, partner models use their respective publisher-specific count-tokens endpoints. """ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Any, Final +from pydantic import TypeAdapter + from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.llms.anthropic.common_utils import OBJECT_LIST_ADAPTER, is_claude_model_target, tools_without_custom_field +from litellm.llms.anthropic.common_utils import ( + OBJECT_LIST_ADAPTER, + count_tokens_feature_betas, + is_claude_model_target, + tools_without_custom_field, +) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -23,6 +30,9 @@ _COUNT_TOKENS_PUBLISHERS: Final = ( ) +_REQUEST_MESSAGES: Final = TypeAdapter(Sequence[object]) + + def count_tokens_request_body( request_data: Mapping[str, object], model: str, drop_params: object ) -> Mapping[str, object]: @@ -43,6 +53,14 @@ def count_tokens_publisher(model: str) -> str | None: return next((publisher for marker, publisher in _COUNT_TOKENS_PUBLISHERS if marker in model), None) +def count_tokens_headers(access_token: str, request_data: Mapping[str, object]) -> Mapping[str, str]: + authorization: Final = {"Authorization": f"Bearer {access_token}"} + betas: Final = count_tokens_feature_betas( + "vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"]) + ) + return {**authorization, "anthropic-beta": ",".join(betas)} if betas else authorization + + class VertexAIPartnerModelsTokenCounter(VertexBase): """ Token counter for Vertex AI Partner Models. @@ -175,8 +193,7 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): api_base=litellm_params.get("api_base"), ) - # Prepare headers - headers: Final = {"Authorization": f"Bearer {access_token}"} + headers: Final = dict(count_tokens_headers(access_token, request_data)) # Get async HTTP client from litellm import LlmProviders diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index a648376c542..152614e4385 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -713,6 +713,14 @@ class AnthropicResponseContentBlockRedactedThinking(LiteLLMBaseModel): data: str +class AnthropicResponseContentBlockCompaction(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["compaction"] + content: str | None = None + signature: str | None = None + + class AnthropicResponseContentBlockServerToolUse(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) @@ -743,7 +751,7 @@ class AnthropicOutputTokensDetails(LiteLLMBaseModel): thinking_tokens: int | None = None -AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal"] +AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal", "compaction"] class AnthropicResponse(LiteLLMBaseModel): @@ -761,6 +769,7 @@ class AnthropicResponse(LiteLLMBaseModel): | AnthropicResponseContentBlockToolUse | AnthropicResponseContentBlockThinking | AnthropicResponseContentBlockRedactedThinking + | AnthropicResponseContentBlockCompaction | AnthropicResponseContentBlockServerToolUse | AnthropicResponseContentBlockWebSearchToolResult ] @@ -808,6 +817,9 @@ class AnthropicThinkingParam(TypedDict, total=False): block_binding: ReadOnly[AnthropicThinkingBlockBinding] +LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS: Final = frozenset(("anthropic", "azure_ai", "vertex_ai")) + + class ANTHROPIC_HOSTED_TOOLS(str, Enum): WEB_SEARCH = "web_search" BASH = "bash" @@ -828,6 +840,7 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): CONTEXT_MANAGEMENT_2025_06_27 = "context-management-2025-06-27" COMPACT_2026_01_12 = "compact-2026-01-12" COMPACT_2026_09_04 = "compact-2026-09-04" + INLINE_TOOLS_2026_09_15 = "inline-tools-2026-09-15" STRUCTURED_OUTPUT_2025_09_25 = "structured-outputs-2025-11-13" ADVANCED_TOOL_USE_2025_11_20 = "advanced-tool-use-2025-11-20" FAST_MODE_2026_02_01 = "fast-mode-2026-02-01" diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 336e9f7d24a..d0e31cdd57a 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -6956,6 +6956,95 @@ def test_handle_anthropic_messages_response_logging_reparses_the_served_body_wit assert logged.usage.completion_tokens == 7 # type: ignore[attr-defined] +@pytest.mark.parametrize( + "compaction_block", + [ + {"type": "compaction", "content": "Summary of the chat so far.", "signature": "sig-abc"}, + {"type": "compaction", "signature": "sig-abc"}, + ], + ids=["with_summary", "without_content"], +) +def test_handle_anthropic_messages_response_logging_reparses_a_served_compaction_response( + compaction_block: dict[str, str], +): + logging_obj = _anthropic_messages_logging_obj() + logging_obj.optional_params = {} + + logged = logging_obj._handle_anthropic_messages_response_logging( + result={ + "id": "msg_compacted", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [compaction_block], + "stop_reason": "compaction", + "stop_sequence": None, + "usage": {"input_tokens": 410, "output_tokens": 95}, + } + ) + + assert isinstance(logged, ModelResponse) + assert logged.id == "msg_compacted" + assert logged.usage.prompt_tokens == 410 + assert logged.usage.completion_tokens == 95 + assert logged.choices[0].message.provider_specific_fields["compaction_blocks"] == [compaction_block] + + +_WEB_SEARCH_REPLY: Final = ( + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "weather"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [{"type": "web_search_result", "url": "https://example.com", "title": "Weather"}], + }, + {"type": "text", "text": "It is sunny."}, +) + + +@pytest.mark.parametrize( + ("content", "stop_reason"), + [ + (list(_WEB_SEARCH_REPLY), "end_turn"), + (list(_WEB_SEARCH_REPLY), "pause_turn"), + ([{"type": "text", "text": "The context window is full."}], "model_context_window_exceeded"), + ( + [ + {"type": "mcp_tool_use", "id": "mcptoolu_1", "name": "lookup", "server_name": "docs", "input": {}}, + {"type": "mcp_tool_result", "tool_use_id": "mcptoolu_1", "is_error": False, "content": [{"type": "text", "text": "ok"}]}, + {"type": "text", "text": "Found it."}, + ], + "end_turn", + ), + ([{"type": "container_upload", "file_id": "file_1"}, {"type": "text", "text": "Uploaded."}], "end_turn"), + ], + ids=["web_search", "pause_turn", "model_context_window_exceeded", "mcp_tool_use", "container_upload"], +) +def test_handle_anthropic_messages_response_logging_bills_a_served_server_tool_response( + content: list[dict[str, object]], stop_reason: str +): + logging_obj = _anthropic_messages_logging_obj() + logging_obj.optional_params = {} + + logged = logging_obj._handle_anthropic_messages_response_logging( + result={ + "id": "msg_server_tools", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": content, + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": {"input_tokens": 1200, "output_tokens": 80}, + } + ) + + assert isinstance(logged, ModelResponse) + assert logged.id == "msg_server_tools" + assert logged.usage.prompt_tokens == 1200 + assert logged.usage.completion_tokens == 80 + assert logged.choices[0].message.content == content[-1]["text"] + + def test_handle_anthropic_messages_response_logging_passes_model_response_through(): """Anthropic-native path already yields a ModelResponse; it must be returned unchanged.""" logging_obj = _anthropic_messages_logging_obj() diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 5b64e42d8a8..9e385c588ae 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -57,9 +57,10 @@ def _provider_fields() -> dict[str, object]: "thinking_blocks": [{"type": "thinking", "thinking": SECRET}], "reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}], "citations": [{"cited_text": SECRET}], - "thought_signatures": [SIGNATURE], - "compaction_blocks": [{"type": "compaction", "content": SECRET}], + "compaction_blocks": [{"type": "compaction", "content": SECRET, "signature": "compaction-signature"}], + "compaction_start": {"type": "compaction", "content": SECRET}, "compaction_delta": {"type": "compaction_delta", "content": SECRET}, + "thought_signatures": [SIGNATURE], } diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index c1f28e33969..0808ce0ec3c 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -3544,12 +3544,11 @@ class TestWifProviderAllowlist: monkeypatch.setenv("WIF_TEST_JWT", "jwt-assertion-value") def test_vertex_anthropic_never_mints_or_sends_the_assertion(self, monkeypatch, wif_engine): + import litellm from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) - import litellm - poster, calls = wif_engine self._env_only_wif(monkeypatch) @@ -3775,7 +3774,7 @@ class TestModelDiscovery: ) monkeypatch.setattr("litellm.module_level_client", client) - with pytest.raises(Exception, match="invalid x-api-key") as exc_info: # noqa: B017, PT011 # the callee raises a bare Exception; match pins the sanitized text + with pytest.raises(Exception, match="invalid x-api-key") as exc_info: # noqa: B017 # the callee raises a bare Exception; match pins the sanitized text AnthropicModelInfo().get_models(api_base="https://api.anthropic.com") assert "invalid x-api-key" in str(exc_info.value) @@ -4238,15 +4237,18 @@ def test_thinking_display_beta_requires_active_thinking(thinking: object, expect @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize( - ("display", "expected_thinking"), + ("provider", "display", "expected_thinking"), ( - ("summarized", {"type": "adaptive", "display": "summarized"}), - ("omitted", {"type": "adaptive", "display": "omitted"}), - ("updates", {"type": "adaptive"}), + ("azure_ai", "summarized", {"type": "adaptive", "display": "summarized"}), + ("azure_ai", "omitted", {"type": "adaptive", "display": "omitted"}), + ("azure_ai", "updates", {"type": "adaptive", "display": "updates"}), + ("vertex_ai", "updates", {"type": "adaptive", "display": "updates"}), + ("anthropic", "updates", {"type": "adaptive", "display": "updates"}), + ("databricks", "updates", {"type": "adaptive"}), ), ) def test_shared_legacy_thinking_translation_preserves_supported_display( - display: str, expected_thinking: dict[str, str] + provider: str, display: str, expected_thinking: dict[str, str] ) -> None: from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -4257,7 +4259,7 @@ def test_shared_legacy_thinking_translation_preserves_supported_display( AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( model="claude-opus-5", optional_params=optional_params, - custom_llm_provider="azure_ai", + custom_llm_provider=provider, drop_params=False, ) @@ -4552,6 +4554,47 @@ def test_shared_legacy_thinking_translation_drops_a_non_numeric_budget_under_dro ) +_SIGNED_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "summary so far", "signature": "sig"} + + +@pytest.mark.parametrize( + "provider,expected", + [("anthropic", True), ("vertex_ai", True), ("azure_ai", True), ("bedrock", False), ("bedrock_converse", False)], +) +@pytest.mark.parametrize( + "optional_params,messages", + [ + ({"compaction": {"type": "summarize"}}, [{"role": "user", "content": "hi"}]), + ({}, [{"role": "assistant", "content": [_SIGNED_COMPACTION_BLOCK]}, {"role": "user", "content": "go on"}]), + ], + ids=["compaction_param", "signed_compaction_block"], +) +def test_requires_native_compaction_beta_per_provider( + provider: str, expected: bool, optional_params: dict[str, object], messages: list[object] +) -> None: + from litellm.llms.anthropic.common_utils import requires_native_compaction_beta + + assert requires_native_compaction_beta(provider, optional_params, messages) is expected + + +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "vertex_ai"]) +@pytest.mark.parametrize( + "optional_params,messages", + [ + ({"max_tokens": 64}, [{"role": "user", "content": "hi"}]), + ({}, [{"role": "assistant", "content": [{"type": "compaction", "content": "unsigned"}]}]), + ({"compaction": None}, [{"role": "user", "content": "hi"}]), + ], + ids=["no_signal", "unsigned_block", "compaction_none"], +) +def test_requires_native_compaction_beta_is_false_without_a_signal( + provider: str, optional_params: dict[str, object], messages: list[object] +) -> None: + from litellm.llms.anthropic.common_utils import requires_native_compaction_beta + + assert requires_native_compaction_beta(provider, optional_params, messages) is False + + @pytest.mark.parametrize( ("custom_llm_provider", "model", "expected"), [ diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index 4913d2b0d1d..7807975be24 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -215,3 +215,56 @@ async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest assert result is not None assert result.total_tokens == 6 assert route.called + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("assistant_content", "expected_beta"), + [ + ( + [{"type": "compaction", "content": "Summary of the chat so far.", "signature": "sig-abc"}], + "compact-2026-09-04,token-counting-2024-11-01", + ), + ("Use Recipe and Step tables.", "token-counting-2024-11-01"), + ], + ids=("signed_compaction_block", "plain_turn"), +) +async def test_count_sends_the_compact_beta_only_when_the_history_replays_a_compaction_block( + monkeypatch: pytest.MonkeyPatch, assistant_content: object, expected_beta: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 376})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-5-5", + messages=[ + {"role": "assistant", "content": assistant_content}, + {"role": "user", "content": "What tables did we pick?"}, + ], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-5-5", + ) + assert result is not None + assert result.total_tokens == 376 + assert route.calls.last.request.headers["anthropic-beta"] == expected_beta + + +@pytest.mark.asyncio +async def test_count_sends_the_tool_change_beta_when_the_history_adds_a_tool(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 541})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-5-5", + messages=[ + {"role": "user", "content": "Use the ping tool once."}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-5-5", + ) + assert result is not None + assert result.total_tokens == 541 + assert route.calls.last.request.headers["anthropic-beta"] == "mid-conversation-tool-changes-2026-07-01,token-counting-2024-11-01" diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 5e0395f84dc..1c4097ad7e0 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -607,3 +607,91 @@ async def test_outgoing_azure_request_keeps_an_object_output_format_under_drop_p body: Final = await _azure_request_body_for_output_format(output_format, drop_params=True) assert body["output_format"] == output_format + + +_SIGNED_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "Summary of the chat so far.", "signature": "sig"} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("messages", "params"), + [ + ([{"role": "user", "content": "Summarize the chat so far."}], {"compaction": {"type": "summarize"}}), + ( + [{"role": "assistant", "content": [_SIGNED_COMPACTION_BLOCK]}, {"role": "user", "content": "Go on"}], + {}, + ), + ], + ids=["compaction_param", "signed_compaction_block"], +) +async def test_outgoing_azure_compaction_request_carries_the_compact_2026_09_04_beta( + messages: list[object], params: dict[str, object] +) -> None: + """Foundry answers 400 "compaction: this parameter requires anthropic-beta: compact-2026-09-04" for the param + and 400 "Input tag 'compaction' ... does not match" for a replayed block without the header, and 200 with it + (live, claude-sonnet-5-5 on Azure AI Foundry, 2026-10-20).""" + received: Final[list[httpx.Request]] = [] + + def foundry_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _foundry_answer(request) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(foundry_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=256, + messages=messages, + model="azure_ai/claude-sonnet-5-5", + api_key="test-api-key", + api_base="https://test-resource.services.ai.azure.com", + client=upstream, + **params, + ) + finally: + await upstream.close() + + assert len(received) == 1, received + assert received[0].headers["anthropic-beta"].split(",").count("compact-2026-09-04") == 1, dict(received[0].headers) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("feature", ("thinking", "output_config", "tool_addition", "tool_removal", "none")) +@pytest.mark.parametrize("explicit_beta", (False, True)) +def test_messages_feature_beta_headers(feature: str, explicit_beta: bool) -> None: + from typing import Final + + from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig + from litellm.types.llms.anthropic import ( + ANTHROPIC_BETA_HEADER_VALUES, + ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, + ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, + ) + + beta: Final = ( + ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + if feature == "thinking" + else ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value + if feature == "output_config" + else ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER + ) + message: Final = ( + {"role": "system", "content": [], "output_config": {"effort": "high"}} + if feature == "output_config" + else {"role": "system", "content": [{"type": feature, "tool": {"type": "tool_reference", "name": "ping"}}]} + if feature in ("tool_addition", "tool_removal") + else {"role": "user", "content": "Reply with OK"} + ) + headers, _ = AzureAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={"anthropic-beta": f"existing-beta,{beta}" if explicit_beta else "existing-beta"}, + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Hello"}, message], + optional_params={"thinking": {"type": "adaptive", "display": "updates"}} if feature == "thinking" else {}, + litellm_params={"vertex_ai_project": "test-project", "vertex_ai_location": "us-east5"}, + api_key="test-key", + api_base="https://example.com/anthropic", + ) + + betas: Final = headers["anthropic-beta"].split(",") + assert betas.count(beta) == int(feature != "none" or explicit_beta) + assert betas.count("existing-beta") == 1 diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py index 1179b0dd5da..6ff2a6973a3 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py @@ -5,10 +5,12 @@ sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) ) +from typing import Final from unittest.mock import MagicMock, patch import pytest +from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig from litellm.types.router import GenericLiteLLMParams @@ -484,3 +486,85 @@ def test_chat_flagged_model_keeps_mid_conversation_system_role_in_place(local_mo "role": "system", "content": [{"type": "text", "text": "Answer with exactly one word."}], } + + +@pytest.mark.parametrize( + ("optional_params", "messages", "expected"), + [ + ({"compaction": {"type": "summarize"}}, [{"role": "user", "content": "Hello"}], True), + ( + {}, + [ + {"role": "assistant", "content": [{"type": "compaction", "content": "Summary.", "signature": "sig"}]}, + {"role": "user", "content": "Go on"}, + ], + True, + ), + ({}, [{"role": "user", "content": "Hello"}], False), + ], + ids=["compaction_param", "signed_compaction_block", "no_compaction"], +) +def test_azure_chat_compaction_sends_the_compact_2026_09_04_beta( + local_beta_headers_config: None, optional_params: dict[str, object], messages: list[object], expected: bool +) -> None: + headers: Final = AzureAnthropicConfig().update_headers_with_optional_anthropic_beta( + headers={}, optional_params=optional_params, messages=messages + ) + + filtered: Final = update_headers_with_filtered_beta(headers=headers, provider="azure_ai") + + assert ("compact-2026-09-04" in filtered.get("anthropic-beta", "").split(",")) is expected + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("thinking_type", ("adaptive", "enabled")) +@pytest.mark.parametrize("display", ("summarized", "omitted", "updates")) +def test_chat_preserves_thinking_display_and_beta(thinking_type: str, display: str) -> None: + from typing import Final + + from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig + from litellm.types.llms.anthropic import ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + + config: Final = AzureAnthropicConfig() + optional_params: Final = config.map_openai_params( + non_default_params={"thinking": {"type": thinking_type, "display": display, "budget_tokens": 2048}}, + optional_params={}, + model="claude-fable-5-1", + drop_params=False, + ) + headers: Final = config.validate_environment( + headers={"anthropic-beta": "existing-beta"}, + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Reply with OK"}], + optional_params=optional_params, + litellm_params={}, + api_key="test-key", + ) + result: Final = config.transform_request( + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Reply with OK"}], + optional_params=optional_params, + litellm_params={}, + headers=headers, + ) + + assert result["thinking"]["type"] == "adaptive" + assert result["thinking"]["display"] == display + assert (ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER in headers["anthropic-beta"].split(",")) == (display == "updates") + assert "existing-beta" in headers["anthropic-beta"].split(",") + + +@pytest.mark.parametrize("thinking", ("enabled", {"type": "bogus"}, {"type": "adaptive", "display": "future-mode"})) +def test_beta_detection_preserves_unrecognized_thinking(thinking: object) -> None: + from typing import Final + + from litellm.types.llms.anthropic import ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + + params: Final = {"thinking": thinking} + headers: Final = AzureAnthropicConfig().validate_environment( + headers={}, model="claude-fable-5-1", messages=[], optional_params=params, + litellm_params={}, api_key="test-key", + ) + + assert params["thinking"] == thinking + assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER not in headers.get("anthropic-beta", "").split(",") diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 7f16957ddc4..fc57dd13302 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -8,11 +8,13 @@ from threading import get_ident from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import requests import responses import litellm +from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.messages.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, @@ -115,7 +117,11 @@ def _vertex_client() -> MagicMock: async def _send_to_vertex(client: MagicMock, caller_headers: Mapping[str, str], drop_params: bool) -> None: with responses.RequestsMock(assert_all_requests_are_fired=False) as oauth: - oauth.add(responses.POST, "https://oauth2.googleapis.com/token", json={"access_token": "vertex-token", "expires_in": 3600}) + oauth.add( + responses.POST, + "https://oauth2.googleapis.com/token", + json={"access_token": "vertex-token", "expires_in": 3600}, + ) await litellm.anthropic.messages.acreate( model="vertex_ai/claude-sonnet-5", messages=[{"role": "user", "content": "hi"}], @@ -289,7 +295,12 @@ def test_no_safeguards_leaves_dangerous_tool_use_beta_header_out(): assert "dangerous-tool-use-2026-09-03" not in updated_headers.get("anthropic-beta", "") -def _validate_vertex_headers(client_headers, messages, thinking: Mapping[str, object] | None = None): +def _validate_vertex_headers( + client_headers: dict[str, str], + messages: list[object], + optional_params: dict[str, object] | None = None, + thinking: Mapping[str, object] | None = None, +) -> dict[str, str]: config = VertexAIPartnerModelsAnthropicMessagesConfig() litellm_params = { "vertex_ai_project": "test-project", @@ -305,7 +316,10 @@ def _validate_vertex_headers(client_headers, messages, thinking: Mapping[str, ob headers=client_headers, model="claude-opus-5-5", messages=messages, - optional_params={"max_tokens": 64, **({"thinking": thinking} if thinking is not None else {})}, + optional_params={ + **(optional_params or {"max_tokens": 64}), + **({"thinking": thinking} if thinking is not None else {}), + }, litellm_params=litellm_params, api_base=None, ) @@ -1040,3 +1054,178 @@ def test_vertex_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_f and info.get("supports_mid_conversation_system") is not True ] assert missing == [] + + +_SIGNED_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "summary so far", "signature": "sig"} + + +@pytest.mark.parametrize( + "client_headers", + [{"anthropic-beta": "compact-2026-09-04"}, {}], + ids=["client_sends_beta", "client_omits_beta"], +) +@pytest.mark.parametrize( + "optional_params,messages", + [ + ({"max_tokens": 64, "compaction": {"type": "summarize"}}, [{"role": "user", "content": "Hello"}]), + ( + {"max_tokens": 64}, + [{"role": "assistant", "content": [_SIGNED_COMPACTION_BLOCK]}, {"role": "user", "content": "Go on"}], + ), + ], + ids=["compaction_param", "signed_compaction_block"], +) +def test_native_compaction_reaches_vertex_with_compact_2026_09_04_beta( + local_beta_headers_config: None, + client_headers: dict[str, str], + optional_params: dict[str, object], + messages: list[object], +) -> None: + """Vertex serves on-demand compaction only behind compact-2026-09-04 (Google Cloud: beta on + https://platform.claude.com/docs/en/build-with-claude/compaction-on-demand, read 2026-10-05) + and answers 400 "compaction: Extra inputs are not permitted" without it, so the beta rides + along with the `compaction` param or a signed compaction block exactly once, whether or not + the client sent it, and survives the Vertex beta filter.""" + headers: Final = _validate_vertex_headers(client_headers, messages, optional_params) + + filtered: Final = update_headers_with_filtered_beta(headers=headers, provider="vertex_ai") + + assert filtered["anthropic-beta"].split(",").count("compact-2026-09-04") == 1 + + +def test_no_compaction_signal_leaves_compact_2026_09_04_beta_out() -> None: + headers: Final = _validate_vertex_headers({}, [{"role": "user", "content": "Hello"}]) + + assert "compact-2026-09-04" not in headers.get("anthropic-beta", "") + + +_VERTEX_MESSAGE_RESPONSE: Final = { + "id": "msg_vrtx_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "pong"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, +} + + +@pytest.mark.asyncio +async def test_vertex_messages_request_sends_compact_2026_09_04_beta_on_the_wire( + local_beta_headers_config: None, +) -> None: + """A /v1/messages compaction request to a Vertex Claude deployment reaches Vertex with the beta in the + anthropic-beta request header and never in the body, the shape Vertex accepts (live, 2026-10-05).""" + mock_response: Final = MagicMock() + mock_response.status_code = 200 + mock_response.text = json.dumps(_VERTEX_MESSAGE_RESPONSE) + mock_response.headers = httpx.Headers({}) + mock_response.json.return_value = _VERTEX_MESSAGE_RESPONSE + + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.post = AsyncMock(return_value=mock_response) + + with patch( + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token_async", + AsyncMock(return_value=("token", "test-project")), + ): + await litellm.anthropic.messages.acreate( + model="vertex_ai/claude-sonnet-4-6", + max_tokens=64, + messages=[{"role": "user", "content": "Reply with the single word: pong"}], + compaction={"type": "summarize"}, + vertex_project="test-project", + vertex_location="us-east5", + client=client, + ) + + sent: Final = client.post.call_args.kwargs + sent_body: Final = sent["json"] if "json" in sent else json.loads(sent["data"]) + assert sent["headers"]["anthropic-beta"].split(",").count("compact-2026-09-04") == 1 + assert sent_body["compaction"] == {"type": "summarize"} + assert "anthropic_beta" not in sent_body + + +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +@pytest.mark.parametrize("feature", ("thinking", "output_config", "tool_addition", "tool_removal", "none")) +@pytest.mark.parametrize("explicit_beta", (False, True)) +def test_messages_feature_beta_headers(feature: str, explicit_beta: bool) -> None: + from typing import Final + from unittest.mock import patch + + from google.oauth2.credentials import Credentials + + from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.messages.transformation import ( + VertexAIPartnerModelsAnthropicMessagesConfig, + ) + from litellm.types.llms.anthropic import ( + ANTHROPIC_BETA_HEADER_VALUES, + ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, + ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, + ) + + beta: Final = ( + ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + if feature == "thinking" + else ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value + if feature == "output_config" + else ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER + ) + message: Final = ( + {"role": "system", "content": [], "output_config": {"effort": "high"}} + if feature == "output_config" + else {"role": "system", "content": [{"type": feature, "tool": {"type": "tool_reference", "name": "ping"}}]} + if feature in ("tool_addition", "tool_removal") + else {"role": "user", "content": "Reply with OK"} + ) + with patch( + "google.auth.default", + return_value=(MagicMock(spec=Credentials, token="test-token", expired=False), "test-project"), + ): + headers, _ = VertexAIPartnerModelsAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={"anthropic-beta": f"existing-beta,{beta}" if explicit_beta else "existing-beta"}, + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Hello"}, message], + optional_params={"thinking": {"type": "adaptive", "display": "updates"}} if feature == "thinking" else {}, + litellm_params={"vertex_ai_project": "test-project", "vertex_ai_location": "us-east5"}, + api_key="test-key", + api_base="https://example.com/anthropic", + ) + + betas: Final = headers["anthropic-beta"].split(",") + assert betas.count(beta) == int(feature != "none" or explicit_beta) + assert betas.count("existing-beta") == 1 + + if feature in ("tool_addition", "tool_removal"): + filtered_headers: Final = update_headers_with_filtered_beta(headers=headers, provider="vertex_ai") + assert filtered_headers.get("anthropic-beta", "").split(",").count(beta) == 1 + + +@pytest.mark.parametrize("thinking", ("enabled", {"type": "bogus"}, {"type": "adaptive", "display": "future-mode"})) +def test_beta_detection_preserves_unrecognized_thinking(thinking: object) -> None: + from typing import Final + from unittest.mock import patch + + from google.oauth2.credentials import Credentials + + from litellm.types.llms.anthropic import ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + + params: Final = {"thinking": thinking} + with patch( + "google.auth.default", + return_value=(MagicMock(spec=Credentials, token="test-token", expired=False), "test-project"), + ): + headers, _ = VertexAIPartnerModelsAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="claude-fable-5-1", + messages=[], + optional_params=params, + litellm_params={"vertex_ai_project": "test-project", "vertex_ai_location": "us-east5"}, + api_key="test-key", + api_base="https://example.com/anthropic", + ) + + assert params["thinking"] == thinking + assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER not in headers.get("anthropic-beta", "").split(",") diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 3ae208bb3a8..0a582c95749 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -4,7 +4,7 @@ import copy import json from collections.abc import Mapping from typing import Final -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import httpx import pytest @@ -12,10 +12,12 @@ import respx from google.oauth2.credentials import Credentials import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.anthropic_beta_headers_manager import ( update_headers_with_filtered_beta, update_request_with_filtered_beta, ) +from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) @@ -475,8 +477,10 @@ def test_vertex_ai_anthropic_extra_headers_beta_propagation(): def test_vertex_ai_anthropic_inline_tools_beta_survives_the_chat_beta_filter(local_beta_headers_config: None) -> None: """The chat path runs the Vertex beta filter over both the header and the `anthropic_beta` body field right - before the request goes out, so inline-tools-2026-09-15 must survive both (source and date in - tests/unit/test_anthropic_beta_headers_filtering.py).""" + before the request goes out, so inline-tools-2026-09-15 must survive it as the header (source and date in + tests/unit/test_anthropic_beta_headers_filtering.py) and stay out of the body field, which Vertex rejects + with 400 "Unexpected value(s) `inline-tools-2026-09-15` for the `anthropic_beta` parameter" (live, global, + claude-sonnet-5-5, 2026-10-20).""" config: Final = VertexAIAnthropicConfig() headers: Final[dict[str, str]] = {} optional_params: Final[dict[str, object]] = { @@ -497,7 +501,7 @@ def test_vertex_ai_anthropic_inline_tools_beta_survives_the_chat_beta_filter(loc ) assert filtered_headers["anthropic-beta"].split(",").count("inline-tools-2026-09-15") == 1 - assert "inline-tools-2026-09-15" in filtered_request["anthropic_beta"] + assert "inline-tools-2026-09-15" not in filtered_request.get("anthropic_beta", []) def test_vertex_ai_anthropic_extra_headers_beta_merged_with_auto_betas(): @@ -1232,3 +1236,293 @@ def test_refusal_stop_details_kept_through_vertex_anthropic_transform(): provider_fields: Final = result.choices[0].message.provider_specific_fields assert provider_fields["stop_reason"] == "refusal" assert provider_fields["stop_details"] == ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + + +_SIGNED_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "summary so far", "signature": "sig"} +_CLEAR_TOOL_USES: Final = {"edits": [{"type": "clear_tool_uses_20250919"}]} + + +def _vertex_request(optional_params: dict[str, object], headers: dict[str, str], litellm_params: dict[str, object] | None = None) -> dict: + return VertexAIAnthropicConfig().transform_request( + model="claude-opus-4-6", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"max_tokens": 100, "is_vertex_request": True, **optional_params}, + litellm_params=litellm_params or {}, + headers=headers, + ) + + +@pytest.mark.parametrize( + "optional_params,messages", + [ + ( + {"max_tokens": 100, "is_vertex_request": True, "compaction": {"type": "summarize"}}, + [{"role": "user", "content": "Hello"}], + ), + ( + {"max_tokens": 100, "is_vertex_request": True}, + [ + { + "role": "assistant", + "content": "", + "provider_specific_fields": {"compaction_blocks": [_SIGNED_COMPACTION_BLOCK]}, + }, + {"role": "user", "content": "Go on"}, + ], + ), + ], + ids=["compaction_param", "signed_compaction_block_sent_back"], +) +def test_vertex_ai_anthropic_native_compaction_adds_compact_2026_09_04_beta( + local_beta_headers_config: None, optional_params: dict[str, object], messages: list[dict[str, object]] +) -> None: + """The chat path skips the Anthropic beta injection for Vertex requests, so the Vertex config + adds compact-2026-09-04 itself (Google Cloud: beta on + https://platform.claude.com/docs/en/build-with-claude/compaction-on-demand, read 2026-10-05) + as the HTTP header only, since Vertex answered 400 "Unexpected value(s) `compact-2026-09-04` + for the `anthropic_beta` parameter" with the name in the body field (us-east5, + claude-sonnet-4-6, 2026-10-05), and the header must survive the Vertex beta filter.""" + headers: Final[dict[str, str]] = {} + + result: Final = VertexAIAnthropicConfig().transform_request( + model="claude-opus-4-6", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers=headers, + ) + filtered_headers, filtered_body = update_request_with_filtered_beta( + headers=headers, request_data=result, provider="vertex_ai" + ) + + assert filtered_headers["anthropic-beta"].split(",").count("compact-2026-09-04") == 1 + assert "compact-2026-09-04" not in filtered_body.get("anthropic_beta", []) + + +def test_vertex_ai_anthropic_no_compaction_signal_leaves_compact_2026_09_04_beta_out() -> None: + headers: Final[dict[str, str]] = {} + + result: Final = _vertex_request({}, headers) + + assert "compact-2026-09-04" not in result.get("anthropic_beta", []) + assert "compact-2026-09-04" not in headers.get("anthropic-beta", "") + + +def test_vertex_ai_anthropic_client_sent_compact_2026_09_04_stays_out_of_the_body_field( + local_beta_headers_config: None, +) -> None: + headers: Final[dict[str, str]] = {} + + result: Final = _vertex_request( + { + "compaction": {"type": "summarize"}, + "extra_headers": {"anthropic-beta": "compact-2026-09-04,interleaved-thinking-2025-05-14"}, + }, + headers, + ) + + assert headers["anthropic-beta"].split(",").count("compact-2026-09-04") == 1 + assert result["anthropic_beta"] == ["interleaved-thinking-2025-05-14"] + + +def test_vertex_ai_anthropic_compaction_keeps_the_betas_the_client_sent_as_headers( + local_beta_headers_config: None, +) -> None: + headers: Final = {"anthropic-beta": "context-1m-2025-08-07"} + + result: Final = _vertex_request({"compaction": {"type": "summarize"}}, headers) + + assert sorted(headers["anthropic-beta"].split(",")) == ["compact-2026-09-04", "context-1m-2025-08-07"] + assert "anthropic_beta" not in result + + +@pytest.mark.parametrize("header_only_beta", ["compact-2026-09-04", "inline-tools-2026-09-15"]) +def test_vertex_ai_anthropic_client_sent_body_header_only_beta_moves_to_the_header( + local_beta_headers_config: None, header_only_beta: str +) -> None: + """Vertex rejects these betas inside the body `anthropic_beta` field with 400 "Unexpected value(s)" and + accepts them as the header (live, global, claude-sonnet-5-5, 2026-10-20), so a client that opts in through that field still + gets the beta as a header and keeps its other betas in the body.""" + headers: Final[dict[str, str]] = {} + + result: Final = _vertex_request( + {"anthropic_beta": [header_only_beta, "interleaved-thinking-2025-05-14"]}, headers + ) + + assert headers["anthropic-beta"] == header_only_beta + assert result["anthropic_beta"] == ["interleaved-thinking-2025-05-14"] + + +@pytest.mark.parametrize( + "client_beta,body,header", + [ + ( + "context-1m-2025-08-07", + ["context-1m-2025-08-07", "context-management-2025-06-27"], + "context-management-2025-06-27", + ), + ( + "inline-tools-2026-09-15", + ["context-management-2025-06-27"], + "context-management-2025-06-27,inline-tools-2026-09-15", + ), + ], + ids=["body_beta", "header_only_beta"], +) +def test_vertex_ai_anthropic_litellm_betas_merge_with_the_client_body_betas( + local_beta_headers_config: None, client_beta: str, body: list[str], header: str +) -> None: + headers: Final[dict[str, str]] = {} + + result: Final = _vertex_request({"context_management": _CLEAR_TOOL_USES, "anthropic_beta": [client_beta]}, headers) + + assert result["anthropic_beta"] == body + assert headers["anthropic-beta"] == header + + +@pytest.mark.parametrize( + "anthropic_beta", + [["interleaved-thinking-2025-05-14"], ("interleaved-thinking-2025-05-14",)], + ids=["list", "sdk_tuple"], +) +def test_vertex_ai_anthropic_client_sent_body_betas_stay_out_of_the_header(anthropic_beta: object) -> None: + headers: Final[dict[str, str]] = {} + + result: Final = _vertex_request({"anthropic_beta": anthropic_beta}, headers) + + assert result["anthropic_beta"] == ["interleaved-thinking-2025-05-14"] + assert "anthropic-beta" not in headers + + +_MALFORMED_BETAS: Final = pytest.mark.parametrize( + "malformed", [5, [1, 2], {"beta": "x"}], ids=["int", "list_of_ints", "dict"] +) +_MALFORMED_BETA_SOURCES: Final = pytest.mark.parametrize( + "source,field", + [("anthropic_beta", "anthropic_beta"), ("extra_headers", "The anthropic-beta header")], + ids=["body_field", "extra_header"], +) + + +def _malformed_beta_params(source: str, malformed: object) -> dict[str, object]: + if source == "extra_headers": + return {"extra_headers": {"anthropic-beta": malformed}} + return {"anthropic_beta": malformed} + + +@_MALFORMED_BETAS +@_MALFORMED_BETA_SOURCES +def test_vertex_ai_anthropic_malformed_anthropic_beta_answers_400_naming_its_field( + malformed: object, source: str, field: str +) -> None: + with pytest.raises(AnthropicError) as excinfo: + _vertex_request(_malformed_beta_params(source, malformed), {}) + + assert excinfo.value.status_code == 400 + assert excinfo.value.message.startswith(f"{field} must be") + + +@_MALFORMED_BETAS +@_MALFORMED_BETA_SOURCES +@pytest.mark.parametrize("auto_beta", [False, True], ids=["alone", "next_to_litellm_beta"]) +def test_vertex_ai_anthropic_malformed_anthropic_beta_is_dropped_under_drop_params( + local_beta_headers_config: None, malformed: object, source: str, field: str, auto_beta: bool +) -> None: + headers: Final[dict[str, str]] = {} + params: Final = {**_malformed_beta_params(source, malformed), **({"context_management": _CLEAR_TOOL_USES} if auto_beta else {})} + + result: Final = _vertex_request(params, headers, litellm_params={"drop_params": True}) + + assert result.get("anthropic_beta") == (["context-management-2025-06-27"] if auto_beta else None) + assert headers.get("anthropic-beta") == ("context-management-2025-06-27" if auto_beta else None) + + +_VERTEX_MESSAGE_RESPONSE: Final = { + "id": "msg_vrtx_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "pong"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, +} + + +def test_vertex_ai_anthropic_completion_sends_compact_2026_09_04_beta_on_the_wire(local_beta_headers_config: None) -> None: + """A chat completion with `compaction` on a Vertex Claude deployment reaches Vertex with the beta in the + anthropic-beta request header and never in the body field, the shape Vertex accepts (live, 2026-10-05).""" + mock_response: Final = MagicMock() + mock_response.status_code = 200 + mock_response.text = json.dumps(_VERTEX_MESSAGE_RESPONSE) + mock_response.headers = {} + mock_response.json.return_value = _VERTEX_MESSAGE_RESPONSE + client: Final = MagicMock(spec=HTTPHandler) + client.post.return_value = mock_response + + with patch( + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token", + return_value=("token", "test-project"), + ): + litellm.completion( + model="vertex_ai/claude-sonnet-4-6", + messages=[{"role": "user", "content": "Reply with the single word: pong"}], + compaction={"type": "summarize"}, + vertex_project="test-project", + vertex_location="us-east5", + client=client, + ) + + sent: Final = client.post.call_args.kwargs + sent_body: Final = sent["json"] if "json" in sent else json.loads(sent["data"]) + assert sent["headers"]["anthropic-beta"].split(",").count("compact-2026-09-04") == 1 + assert sent_body["compaction"] == {"type": "summarize"} + assert "anthropic_beta" not in sent_body + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("thinking_type", ("adaptive", "enabled")) +@pytest.mark.parametrize("display", ("summarized", "omitted", "updates")) +def test_chat_preserves_thinking_display_and_beta(thinking_type: str, display: str) -> None: + from typing import Final + + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig + from litellm.types.llms.anthropic import ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + + config: Final = VertexAIAnthropicConfig() + optional_params: Final = config.map_openai_params( + non_default_params={"thinking": {"type": thinking_type, "display": display, "budget_tokens": 2048}}, + optional_params={}, + model="claude-fable-5-1", + drop_params=False, + ) + optional_params["extra_headers"] = {"anthropic-beta": "existing-beta"} + headers: Final = {} + result: Final = config.transform_request( + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Reply with OK"}], + optional_params=optional_params, + litellm_params={}, + headers=headers, + ) + + assert result["thinking"]["type"] == "adaptive" + assert result["thinking"]["display"] == display + assert (ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER in headers["anthropic-beta"].split(",")) == (display == "updates") + assert "existing-beta" in headers["anthropic-beta"].split(",") + + +@pytest.mark.parametrize("thinking", ("enabled", {"type": "bogus"}, {"type": "adaptive", "display": "future-mode"})) +def test_beta_detection_preserves_unrecognized_thinking(thinking: object) -> None: + from typing import Final + + from litellm.types.llms.anthropic import ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + + headers: Final = {} + result: Final = VertexAIAnthropicConfig().transform_request( + model="claude-fable-5-1", messages=[{"role": "user", "content": "Hello"}], + optional_params={"thinking": thinking}, litellm_params={}, headers=headers, + ) + + assert result["thinking"] == thinking + assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER not in headers.get("anthropic-beta", "").split(",") diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py new file mode 100644 index 00000000000..c03f08ab610 --- /dev/null +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py @@ -0,0 +1,60 @@ +from typing import Final + +import pytest + +from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_headers + +_HISTORY: Final = ({"role": "user", "content": "Hello"},) + + +@pytest.mark.parametrize( + "request_data", + [ + {"messages": [{"role": "assistant", "content": [{"type": "compaction", "content": "so far", "signature": "sig"}]}]}, + {"messages": list(_HISTORY), "compaction": {"type": "summarize"}}, + ], + ids=["signed_compaction_block", "compaction_param"], +) +def test_count_tokens_on_a_compaction_request_sends_the_compact_beta(request_data: dict[str, object]) -> None: + assert count_tokens_headers("token", request_data) == { + "Authorization": "Bearer token", + "anthropic-beta": "compact-2026-09-04", + } + + +@pytest.mark.parametrize( + "request_data", + [ + {"messages": list(_HISTORY)}, + {"messages": [{"role": "assistant", "content": [{"type": "compaction", "content": "so far", "signature": ""}]}]}, + ], + ids=["plain", "unsigned_compaction_block"], +) +def test_count_tokens_without_compaction_sends_no_beta(request_data: dict[str, object]) -> None: + assert count_tokens_headers("token", request_data) == {"Authorization": "Bearer token"} + + +_TOOL_ADDITION: Final = { + "role": "system", + "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}], +} + + +@pytest.mark.parametrize( + ("request_data", "expected_beta"), + [ + ({"messages": [*_HISTORY, _TOOL_ADDITION]}, "mid-conversation-tool-changes-2026-07-01"), + ( + {"messages": [*_HISTORY, _TOOL_ADDITION], "compaction": {"type": "summarize"}}, + "compact-2026-09-04,mid-conversation-tool-changes-2026-07-01", + ), + ], + ids=["tool_addition", "tool_addition_with_compaction"], +) +def test_count_tokens_on_a_tool_change_request_sends_the_tool_change_beta( + request_data: dict[str, object], expected_beta: str +) -> None: + assert count_tokens_headers("token", request_data) == { + "Authorization": "Bearer token", + "anthropic-beta": expected_beta, + } diff --git a/tests/unit/test_anthropic_beta_headers_filtering.py b/tests/unit/test_anthropic_beta_headers_filtering.py index deba7fe9587..c8fef95cf98 100644 --- a/tests/unit/test_anthropic_beta_headers_filtering.py +++ b/tests/unit/test_anthropic_beta_headers_filtering.py @@ -85,7 +85,8 @@ class TestAnthropicBetaHeadersFiltering: filtered = filter_and_transform_beta_headers( beta_headers=all_headers, provider=provider ) - assert ("compact-2026-09-04" in filtered) is (provider == "anthropic") + assert ("compact-2026-09-04" in filtered) is (provider in ("anthropic", "azure_ai", "vertex_ai")) + assert ("mid-conversation-tool-changes-2026-07-01" in filtered) is (provider != "bedrock_converse") for header in unsupported_headers: assert ( @@ -406,15 +407,49 @@ class TestAnthropicBetaHeadersFiltering: test_case["expected"] in filtered ), f"Header '{test_case['input']}' should be mapped to '{test_case['expected']}' for {test_case['provider']}, but got: {filtered}" - def test_filter_and_transform_beta_headers_vertex_ai_keeps_compact(self): - """Vertex AI supports compact context edits, so the compact beta header - must be forwarded instead of stripped (it was previously mapped to null, - which broke compact_20260112 context edits over /v1/messages).""" - filtered = filter_and_transform_beta_headers( - beta_headers=["compact-2026-01-12"], provider="vertex_ai" + @pytest.mark.parametrize("beta", ["compact-2026-09-04", "inline-tools-2026-09-15"]) + def test_vertex_ai_header_only_beta_is_forwarded_as_the_header_and_dropped_from_the_body( + self, beta: str + ) -> None: + headers, body = update_request_with_filtered_beta( + headers={"anthropic-beta": f"{beta},context-1m-2025-08-07"}, + request_data={"anthropic_beta": [beta, "context-1m-2025-08-07"]}, + provider="vertex_ai", ) - assert filtered == ["compact-2026-01-12"] + assert set(headers["anthropic-beta"].split(",")) == {beta, "context-1m-2025-08-07"} + assert body["anthropic_beta"] == ["context-1m-2025-08-07"] + + def test_header_only_betas_stay_out_of_the_provider_columns_older_proxies_read(self) -> None: + """Proxies released before header-only columns existed fetch this map from main, read only the + provider column, and copy its betas into the body field too, where Vertex answers 400. So a + header-only beta may sit in its provider column only if older proxies already forwarded it.""" + header_only_columns: Final = { + column.removesuffix("_header_only"): betas + for column, betas in self.config.items() + if column.endswith("_header_only") + } + forwarded_by_older_proxies: Final = { + provider: {beta for beta in betas if self.config[provider].get(beta) is not None} + for provider, betas in header_only_columns.items() + } + + assert "vertex_ai" in header_only_columns + assert {provider: betas for provider, betas in forwarded_by_older_proxies.items() if betas} == { + "vertex_ai": {"inline-tools-2026-09-15"} + } + + def test_filter_and_transform_beta_headers_vertex_ai_keeps_compact(self): + """Vertex AI serves both compaction betas (threshold compaction behind compact-2026-01-12 + and on-demand compaction behind compact-2026-09-04, Google Cloud: beta on + https://platform.claude.com/docs/en/build-with-claude/compaction-on-demand, read 2026-10-05), + so the Vertex filter forwards both instead of stripping them: a dropped header makes + Vertex answer 400 "compaction: Extra inputs are not permitted".""" + filtered = filter_and_transform_beta_headers( + beta_headers=["compact-2026-01-12", "compact-2026-09-04"], provider="vertex_ai" + ) + + assert filtered == ["compact-2026-01-12", "compact-2026-09-04"] @pytest.mark.parametrize("provider", ["bedrock_converse", "bedrock"]) def test_fine_grained_tool_streaming_forwarded_for_bedrock(self, provider): From a084d136358a2375c0999b56afadf64ebaa4d107 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:40:58 -0700 Subject: [PATCH 33/64] fix(passthrough): log the URL model and declared provider on Gemini and Vertex failures (#49075) * fix(passthrough): log the URL model and declared provider on Gemini and Vertex failures Pass-through failures on Gemini and Vertex routes carry the model only in the URL, so the failure spend row and the Prometheus failure counter logged an empty model. The failure payload now falls back to the model parsed from the URL on those routes, and the Prometheus provider label honors the provider the request declares before inferring one from the model name * fix(passthrough): keep model and provider on mid-stream and Vertex failure logs Build the failure payload from the URL at both streamed relay sites so a 200 stream that dies mid-body still logs the model, derive vertex_ai for Vertex routes that never declare a provider, and have the Vertex logging handler reuse get_vertex_model_id_from_url instead of its own regex * fix(passthrough): bound the Prometheus provider label and name Vertex failures like the success handler Declared providers only reach the api_provider label when they name a registered provider, otherwise the label falls back to inference. Vertex failure rows now use the success handler's model naming (vertex_ai/ prefix on rawPredict and streamRawPredict, vertex_ai/search_api on :search) and its host-based provider detection, and the detected-stream failure sites are covered by tests. * fix(passthrough): name Google failures from the URL on untracked methods too A Gemini or Vertex pass-through failure on a method the success handler does not track (countTokens, for one) still logged model "" and no provider, although the model sits in the URL. The failure payload now falls back to the Google AI host for the provider and reads /models/ from the URL, so countTokens 400s name their model and provider like generateContent ones do. Non-Google hosts keep the tracked-route rule and still log nothing extra. * fix(passthrough): trust the declared Google provider and Vertex-shaped paths on custom bases The failure path only named the Google provider from the URL when the method was one the success handler tracks or the host was a Google API host, so a Gemini countTokens failure behind a custom GEMINI_API_BASE (or a Vertex one behind a custom base) still logged an empty model. A route that declares gemini or vertex_ai now wins outright, and a URL whose path carries /projects/

/locations// counts as Vertex on any host, while a tracked method on a non-Google host no longer passes as Vertex. * fix(pass_through): stop the URL model parser at the query string so a configured key never lands in a spend row's model * test(pass_through): run the URL-model failure tests through the real ProxyLogging with a recording callback * fix(passthrough): scope declared provider to pass-through failures and match Google hosts only Prometheus trusts the resolved provider as before and reads the request-declared provider only on pass-through failures. Vertex attribution matches Google hosts, strips operation poll suffixes, prefixes vertex_ai/ only for rawPredict, reads the URL path only, covers Vertex Live WebSocket failures, and labels a router alias failure with the alias the client requested * fix(passthrough): attribute failures from the URL only on built-in Google routes and keep other routes' metric labels A user-defined pass-through route to a Google host keeps its empty model and provider, since its URL names an upstream model rather than a LiteLLM-managed one. Prometheus failure labels read the route's provider and model_group only for gemini and vertex_ai, so Azure, OpenRouter, Bedrock and other pass-through routes keep their earlier labels. The Gemini URL parser drops the query string, Google host matching needs a dot or dash boundary, and the Vertex Claude prefix follows alt=sse like the success handler. * fix(passthrough): log the URL model on guardrail blocks and the named model on Vertex Live failures A pre-call guardrail block on a pass-through route now builds its failure payload the same way upstream failures do, so the model and provider read from the URL reach spend logs and metrics. Vertex Live websocket failures (connect errors, upstream rejections, and the credential failure) now carry the model from the ?model= query parameter under the vertex_ai provider * fix(passthrough): mark guardrail-block and Vertex Live credential failures as pass-through calls Prometheus reads a failure's declared provider only on pass-through calls, so without call_type it fell back to guessing from the model name. That labeled a guardrail block on the Gemini route as vertex_ai and the Live credential failure as None * fix(passthrough): let the route name a failure's provider and give guardrail blocks their own spend row * fix(pass-through): keep URL model attribution off routes that name another provider A built-in route that names a non-Google provider, like /openai pointed at Gemini's OpenAI-compatible host, no longer takes its failure model from the Google URL. A guardrail block response now returns the x-litellm-call-id of its spend row. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/prometheus.py | 53 +- litellm/llms/vertex_ai/common_utils.py | 2 +- .../llm_passthrough_endpoints.py | 12 +- .../gemini_passthrough_logging_handler.py | 8 +- .../vertex_passthrough_logging_handler.py | 63 +- .../pass_through_endpoints.py | 136 ++- .../integrations/test_prometheus_labels.py | 97 +++ .../vertex_ai/test_vertex_ai_common_utils.py | 3 + .../test_llm_pass_through_endpoints.py | 53 +- .../test_pass_through_endpoints.py | 772 +++++++++++++++--- 10 files changed, 1025 insertions(+), 174 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index e6318525b88..5c2f0f553a8 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -82,6 +82,7 @@ from litellm.types.proxy.carried_budget_state import ( UserBudgetSnapshot, ) from litellm.types.utils import ( + CallTypes, StandardLoggingGuardrailInformation, StandardLoggingPayload, StandardLoggingZeroCostDiagnostic, @@ -279,6 +280,11 @@ def _logged_fields(value: object) -> Mapping[str, object]: return _NO_LOGGED_FIELDS +_GOOGLE_PASS_THROUGH_PROVIDERS: Final = frozenset( + (litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value) +) + + def _is_llm_provider(name: str) -> bool: routable: Final = ( name in litellm.provider_list or JSONProviderRegistry.exists(name) or is_registered_custom_provider(name) @@ -2942,6 +2948,22 @@ class PrometheusLogger(CustomLogger): return False + @staticmethod + def _google_pass_through_provider(request_data: Mapping[str, object]) -> str | None: + declared_provider: Final = request_data.get("custom_llm_provider") + if request_data.get("call_type") != CallTypes.pass_through.value or not isinstance(declared_provider, str): + return None + return declared_provider if declared_provider in _GOOGLE_PASS_THROUGH_PROVIDERS else None + + @staticmethod + def _failed_request_requested_model(request_data: Mapping[str, object]) -> object: + requested_model: Final = request_data.get("model", "") + if PrometheusLogger._google_pass_through_provider(request_data) is None: + return requested_model + request_metadata: Final = _logged_fields(_logged_fields(request_data.get("litellm_params")).get("metadata")) + model_group: Final = request_metadata.get("model_group") + return model_group if isinstance(model_group, str) and model_group else requested_model + @staticmethod def _extract_api_provider_from_request_data(request_data: dict) -> str | None: """ @@ -2950,22 +2972,27 @@ class PrometheusLogger(CustomLogger): A request can fail before a deployment is resolved, so the provider is not always known. Prefer the resolved ``custom_llm_provider`` on ``litellm_params``, then any provider recovered onto a partial - ``standard_logging_object`` (e.g. a stream that broke mid-flight), and - finally infer it from the requested model name (e.g. ``gpt-4o-mini`` -> - ``openai``) since the proxy's failure ``request_data`` usually carries - only the client-supplied model. Return ``None`` when it cannot be - determined so the label emits empty rather than a guess. + ``standard_logging_object`` (e.g. a stream that broke mid-flight). On a + Gemini or Vertex pass-through failure only, next take the provider the + failure payload names (``gemini`` or ``vertex_ai``), so every other + pass-through route keeps its earlier label. Finally infer it from the + requested model name (e.g. ``gpt-4o-mini`` -> ``openai``) since the + proxy's failure ``request_data`` usually carries only the client-supplied + model. Return ``None`` when it cannot be determined so the label emits + empty rather than a guess. """ litellm_params: Final = request_data.get("litellm_params") or {} - provider = litellm_params.get("custom_llm_provider") - if provider: - return provider standard_logging_object: Final = request_data.get("standard_logging_object") or {} - provider = standard_logging_object.get("custom_llm_provider") - if provider: - return provider + resolved_provider: Final = litellm_params.get("custom_llm_provider") or standard_logging_object.get( + "custom_llm_provider" + ) + if resolved_provider: + return resolved_provider + google_provider: Final = PrometheusLogger._google_pass_through_provider(request_data) + if google_provider is not None: + return google_provider model: Final = litellm_params.get("model") or request_data.get("model") - if not model: + if not isinstance(model, str) or not model: return None try: return litellm.get_llm_provider(model=model)[1] or None @@ -3052,7 +3079,7 @@ class PrometheusLogger(CustomLogger): team_alias=user_api_key_dict.team_alias, org_id=user_api_key_dict.org_id, org_alias=user_api_key_dict.organization_alias, - requested_model=_bounded_requested_model_label(request_data.get("model", "")), + requested_model=_bounded_requested_model_label(self._failed_request_requested_model(request_data)), status_code=str(status_code), exception_status=str(status_code), exception_class=self._get_exception_class_name(original_exception), diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 0fa80cb2cad..ad64d852ef9 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1257,7 +1257,7 @@ def get_vertex_model_id_from_url(url: str) -> str | None: `https://${LOCATION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/publishers/google/models/${MODEL_ID}:streamGenerateContent` """ - match: Final = re.search(r"/models/([^:]+)", url) + match: Final = re.search(r"/models/([^:?#]+)", url) return match.group(1) if match else None diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 8de7441121f..33b18c3d939 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -32,6 +32,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm import get_llm_provider from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, AWS_ERROR_TYPE_HEADER, @@ -149,7 +150,7 @@ from litellm.types.passthrough_endpoints.tinyfish import ( ) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials from litellm.types.router import LiteLLMParamsTypedDict -from litellm.types.utils import LlmProviders +from litellm.types.utils import CallTypes, LlmProviders from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager @@ -4105,7 +4106,12 @@ async def vertex_ai_live_websocket_passthrough( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, - request_data={}, + request_data={ + "model": model or "", + "custom_llm_provider": litellm.LlmProviders.VERTEX_AI.value, + "call_type": CallTypes.pass_through.value, + "litellm_call_id": str(uuid.uuid4()), + }, ) if websocket.client_state != WebSocketState.DISCONNECTED: await websocket.close(code=1011, reason=VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON) @@ -4140,6 +4146,8 @@ async def vertex_ai_live_websocket_passthrough( forward_headers=False, endpoint="/vertex_ai/live", accept_websocket=False, + builtin_provider_route=True, + model=model, setup_model_rewriter=_build_vertex_live_setup_model_rewriter( vertex_project=resolved_project, vertex_location=resolved_location, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py index 7a5d5a1e259..d449f62af2c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -1,4 +1,3 @@ -import re from collections.abc import Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final @@ -9,6 +8,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.gemini.videos.transformation import GeminiVideoConfig +from litellm.llms.vertex_ai.common_utils import get_vertex_model_id_from_url from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( GeminiSseChunkParser, @@ -207,11 +207,7 @@ class GeminiPassthroughLoggingHandler: @staticmethod def extract_model_from_url(url: str) -> str: - pattern: Final = r"/models/([^:]+)" - match: Final = re.search(pattern, url) - if match: - return match.group(1) - return "unknown" + return get_vertex_model_id_from_url(url) or "unknown" @staticmethod def _create_gemini_response_logging_payload_for_generate_content( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index b3cbad91624..d5c162296d3 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -3,7 +3,7 @@ import re from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast -from urllib.parse import urlparse +from urllib.parse import parse_qs, urlparse import httpx from pydantic import ConfigDict, TypeAdapter @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import from litellm.llms.vertex_ai.common_utils import ( get_vertex_ai_lyria_generation_cost, get_vertex_location_from_url, + get_vertex_model_id_from_url, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as VertexModelResponseIterator, @@ -47,6 +48,19 @@ from litellm.types.utils import ( ) vertex_search_api_config: Final = VertexSearchAPIVectorStoreConfig() +GEMINI_API_HOST_SUFFIX: Final = "generativelanguage.googleapis.com" +GOOGLE_AI_HOST_SUFFIXES: Final = ( + GEMINI_API_HOST_SUFFIX, + "aiplatform.googleapis.com", + "discoveryengine.googleapis.com", +) +_VERTEX_MULTI_REGION_HOST: Final = re.compile(r"aiplatform\.[a-z0-9-]+\.rep\.googleapis\.com") + + +def _is_host_or_subdomain(hostname: str, suffix: str) -> bool: + return hostname == suffix or hostname.endswith((f".{suffix}", f"-{suffix}")) + + if TYPE_CHECKING: from litellm.types.utils import LiteLLMBatch @@ -684,11 +698,18 @@ class VertexPassthroughLoggingHandler: @staticmethod def extract_model_from_url(url: str) -> str: - pattern: Final = r"/models/([^:]+)" - match: Final = re.search(pattern, url) - if match: - return match.group(1) - return "unknown" + return get_vertex_model_id_from_url(url) or "unknown" + + @staticmethod + def model_from_url_route(url_route: str) -> str | None: + parsed_url: Final = urlparse(url_route) + model: Final = get_vertex_model_id_from_url(parsed_url.path) + if model is None: + return "vertex_ai/search_api" if ":search" in parsed_url.path else None + streamed: Final = "sse" in parse_qs(parsed_url.query).get("alt", ()) + if parsed_url.path.endswith((":rawPredict", ":streamRawPredict")) and not streamed: + return f"vertex_ai/{model}" + return model @staticmethod def extract_model_name_from_vertex_path(vertex_model_path: str) -> str: @@ -734,18 +755,32 @@ class VertexPassthroughLoggingHandler: return None @staticmethod - def _get_custom_llm_provider(url: str, logging_obj: LiteLLMLoggingObj) -> str: + def is_google_ai_url(url: str) -> bool: hostname: Final = urlparse(url).hostname - if logging_obj.model_call_details.get("custom_llm_provider") == litellm.LlmProviders.GEMINI.value or ( - hostname is not None - and ( - hostname == "generativelanguage.googleapis.com" - or hostname.endswith(".generativelanguage.googleapis.com") - ) - ): + return hostname is not None and ( + any(_is_host_or_subdomain(hostname, suffix) for suffix in GOOGLE_AI_HOST_SUFFIXES) + or _VERTEX_MULTI_REGION_HOST.fullmatch(hostname) is not None + ) + + @staticmethod + def is_gemini_api_host(url: str) -> bool: + hostname: Final = urlparse(url).hostname + return hostname is not None and ( + hostname == GEMINI_API_HOST_SUFFIX or hostname.endswith(f".{GEMINI_API_HOST_SUFFIX}") + ) + + @staticmethod + def custom_llm_provider_from_url(url: str) -> str: + if VertexPassthroughLoggingHandler.is_gemini_api_host(url): return litellm.LlmProviders.GEMINI.value return litellm.LlmProviders.VERTEX_AI.value + @staticmethod + def _get_custom_llm_provider(url: str, logging_obj: LiteLLMLoggingObj) -> str: + if logging_obj.model_call_details.get("custom_llm_provider") == litellm.LlmProviders.GEMINI.value: + return litellm.LlmProviders.GEMINI.value + return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url) + @staticmethod def _is_multimodal_embedding_response(json_response: dict) -> bool: """ diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index bda2bfb9e04..cfac28db1a9 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3,6 +3,7 @@ import asyncio import copy import json import posixpath +import re import traceback from base64 import b64encode from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence @@ -77,6 +78,7 @@ from litellm.llms.base_llm.managed_resources.utils import ( from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.eu_jev.common_utils import eu_jev_scoped_model, eu_jev_upstream_body from litellm.llms.oss_decision import validate_oss_request +from litellm.llms.vertex_ai.common_utils import get_vertex_model_id_from_url from litellm.passthrough import BasePassthroughUtils from litellm.proxy._lazy_features import lazy_owned_routes from litellm.proxy._types import ( @@ -147,11 +149,12 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( EndpointType, PassthroughStandardLoggingPayload, ) -from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage +from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( is_tinyfish_agent_url, ) +from .llm_provider_handlers.vertex_passthrough_logging_handler import VertexPassthroughLoggingHandler from .streaming_handler import PassThroughStreamingHandler from .success_handler import PassThroughEndpointLogging from .upstream_usage_headers import ( @@ -874,13 +877,45 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | Non metadata.setdefault("standard_logging_guardrail_information", list(entries)) +_GOOGLE_PASSTHROUGH_PROVIDERS: Final = frozenset( + (litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value) +) + + +def _passthrough_google_provider(url_route: str, custom_llm_provider: str | None) -> str | None: + if custom_llm_provider is not None: + return custom_llm_provider if custom_llm_provider in _GOOGLE_PASSTHROUGH_PROVIDERS else None + if VertexPassthroughLoggingHandler.is_google_ai_url(url_route): + return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route) + return None + + +_OPERATION_POLL_SUFFIX: Final = re.compile(r"/operations/[^/]+$") + + +def _without_operation_poll_suffix(model: str | None) -> str | None: + return _OPERATION_POLL_SUFFIX.sub("", model) if model else model + + +def _passthrough_google_model(url_route: str, google_provider: str | None) -> str | None: + match google_provider: + case "gemini": + return _without_operation_poll_suffix(get_vertex_model_id_from_url(urlparse(url_route).path)) + case "vertex_ai": + return _without_operation_poll_suffix(VertexPassthroughLoggingHandler.model_from_url_route(url_route)) + case _: + return None + + def _build_passthrough_failure_request_payload( parsed_body: dict | None, kwargs: dict | None, logging_obj: LiteLLMLoggingObj | None, custom_llm_provider: str | None, upstream_usage: UpstreamReportedUsage | None = None, + url: httpx.URL | None = None, upstream_model: str | None = None, + builtin_provider_route: bool = False, ) -> dict: """Build the ``request_data`` dict passed to ``post_call_failure_hook``. @@ -893,16 +928,25 @@ def _build_passthrough_failure_request_payload( comes paired with a usage object, so both keys are written together. ``upstream_model`` is the model name the upstream reported, which replaces the body's ``model`` on the failure row just as it does on a success row. + ``builtin_provider_route`` lets a Gemini or Vertex URL name the failure's + model and provider. A user-defined route never gets that, since its URL + names an upstream model rather than a LiteLLM-managed one. """ request_payload: Final[dict] = dict(cast(Mapping[str, object], parsed_body or {})) # cast-ok: json body if kwargs: request_payload.update(kwargs) if logging_obj is not None: request_payload["litellm_logging_obj"] = logging_obj - if "model" not in request_payload and parsed_body and isinstance(parsed_body, dict): - request_payload["model"] = parsed_body.get("model", "") - if "custom_llm_provider" not in request_payload and custom_llm_provider: - request_payload["custom_llm_provider"] = custom_llm_provider + url_route: Final = str(url) if url is not None else "" + google_provider: Final = ( + _passthrough_google_provider(url_route, custom_llm_provider) if builtin_provider_route else None + ) + if "model" not in request_payload: + request_payload["model"] = _passthrough_google_model(url_route, google_provider) or "" + resolved_provider: Final = custom_llm_provider or google_provider + request_payload.pop("custom_llm_provider", None) + if resolved_provider: + request_payload["custom_llm_provider"] = resolved_provider if upstream_usage is not None: request_payload["response_cost"] = upstream_usage.response_cost or 0.0 request_payload["combined_usage_object"] = Usage(total_tokens=upstream_usage.total_tokens or 0) @@ -1203,6 +1247,7 @@ async def pass_through_request( # Initialize variables ######################################################### litellm_call_id: Final = str(uuid.uuid4()) + builtin_provider_route: Final = not request_dispatched_to_pass_through_endpoint(request) url: httpx.URL | None = None # parsed request body @@ -1612,6 +1657,8 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, + builtin_provider_route=builtin_provider_route, upstream_model=upstream_model, ), logging_obj=logging_obj, @@ -1651,6 +1698,8 @@ async def pass_through_request( kwargs=kwargs, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, + url=url, + builtin_provider_route=builtin_provider_route, ), ), managed_id_provider=_managed_id_provider, @@ -1710,6 +1759,8 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, + builtin_provider_route=builtin_provider_route, upstream_model=upstream_model, ), logging_obj=logging_obj, @@ -1749,6 +1800,8 @@ async def pass_through_request( kwargs=kwargs, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, + url=url, + builtin_provider_route=builtin_provider_route, ), ), managed_id_provider=_managed_id_provider, @@ -1811,6 +1864,8 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, + builtin_provider_route=builtin_provider_route, upstream_model=upstream_model, ) failure_request_payload["response_body"] = response_body @@ -1991,7 +2046,14 @@ async def pass_through_request( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, - request_data=e.request_data, + request_data=_build_passthrough_failure_request_payload( + parsed_body=e.request_data, + kwargs={"call_type": CallTypes.pass_through.value, "litellm_call_id": litellm_call_id}, + logging_obj=None, + custom_llm_provider=custom_llm_provider, + url=url, + builtin_provider_route=builtin_provider_route, + ), ) except Exception: verbose_proxy_logger.warning( @@ -2010,6 +2072,7 @@ async def pass_through_request( content=json.dumps(error_body), status_code=200, media_type="application/json", + headers={"x-litellm-call-id": litellm_call_id}, ) except Exception as e: custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -2033,19 +2096,14 @@ async def pass_through_request( # Monitoring: Trigger post_call_failure_hook # for pass through endpoint failure ######################################################### - request_payload: Final[dict] = _parsed_body or {} - # add user_api_key_dict, litellm_call_id, passthrough_logging_payloa for logging - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - if logging_obj is not None: - request_payload["litellm_logging_obj"] = logging_obj - - if "model" not in request_payload and _parsed_body and isinstance(_parsed_body, dict): - request_payload["model"] = _parsed_body.get("model", "") - if "custom_llm_provider" not in request_payload and custom_llm_provider: - request_payload["custom_llm_provider"] = custom_llm_provider - + request_payload: Final = _build_passthrough_failure_request_payload( + parsed_body=_parsed_body, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + url=url, + builtin_provider_route=builtin_provider_route, + ) _carry_guardrail_logging_info(request_payload, post_call_guardrail_data) await proxy_logging_obj.post_call_failure_hook( @@ -2547,6 +2605,8 @@ async def websocket_passthrough_request( cost_per_request: float | None = None, accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, + builtin_provider_route: bool = False, + model: str | None = None, ): """ WebSocket passthrough request handler. @@ -2560,6 +2620,8 @@ async def websocket_passthrough_request( endpoint: The endpoint path (for logging purposes) cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream + builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider + model: The model the caller named before connecting, logged on a failure before the setup frame arrives """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj @@ -2912,19 +2974,17 @@ async def websocket_passthrough_request( except InvalidStatus as exc: verbose_proxy_logger.exception("WebSocket passthrough (%s): upstream rejected WebSocket connection", endpoint) - # Prepare request payload for logging - request_payload = {} - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - if logging_obj is not None: - request_payload["litellm_logging_obj"] = logging_obj - - # Log the connection failure using the same pattern as HTTP await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=exc, - request_data=request_payload, + request_data=_build_passthrough_failure_request_payload( + parsed_body=None if model is None else {"model": model}, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=None, + url=httpx.URL(target), + builtin_provider_route=builtin_provider_route, + ), traceback_str=traceback.format_exc( limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, ), @@ -2940,19 +3000,17 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): unexpected error while proxying WebSocket", endpoint ) - # Prepare request payload for logging - request_payload = {} - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - if logging_obj is not None: - request_payload["litellm_logging_obj"] = logging_obj - - # Log the unexpected error using the same pattern as HTTP await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, - request_data=request_payload, + request_data=_build_passthrough_failure_request_payload( + parsed_body=None if model is None else {"model": model}, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=None, + url=httpx.URL(target), + builtin_provider_route=builtin_provider_route, + ), traceback_str=traceback.format_exc( limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, ), diff --git a/tests/unit/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py index bff89b6cbf1..878f19c66f8 100644 --- a/tests/unit/integrations/test_prometheus_labels.py +++ b/tests/unit/integrations/test_prometheus_labels.py @@ -238,6 +238,53 @@ def test_extract_api_provider_from_request_data_failure_path(): assert extract({"model": "some-unknown-model-xyz"}) is None +_PASS_THROUGH_CALL_TYPE: Final = "pass_through_endpoint" + + +@pytest.mark.parametrize( + "request_data, expected_provider", + [ + ({"model": "gpt-4o-mini", "call_type": "acompletion", "custom_llm_provider": "anthropic"}, "openai"), + ({"model": "gpt-4o-mini", "litellm_params": {"custom_llm_provider": "humanloop"}}, "humanloop"), + ({"model": "gemini-3.8-flash", "call_type": _PASS_THROUGH_CALL_TYPE}, "vertex_ai"), + ( + {"model": "gemini-3.8-flash", "call_type": _PASS_THROUGH_CALL_TYPE, "custom_llm_provider": "gemini"}, + "gemini", + ), + ( + { + "call_type": _PASS_THROUGH_CALL_TYPE, + "custom_llm_provider": "gemini", + "standard_logging_object": {"custom_llm_provider": "vertex_ai"}, + }, + "vertex_ai", + ), + ], +) +def test_extract_api_provider_reads_the_request_declared_provider_only_on_pass_through_failures( + request_data: dict[str, object], expected_provider: str +): + from litellm.integrations.prometheus import PrometheusLogger + + assert PrometheusLogger._extract_api_provider_from_request_data(request_data) == expected_provider + + +@pytest.mark.parametrize( + "declared_provider", + ["azure", "bedrock", "openrouter", "client-supplied-value", "", {"nested": "object"}, 42], +) +def test_extract_api_provider_ignores_pass_through_declared_values_outside_gemini_and_vertex( + declared_provider: object, +): + from litellm.integrations.prometheus import PrometheusLogger + + extract = PrometheusLogger._extract_api_provider_from_request_data + declared = {"call_type": _PASS_THROUGH_CALL_TYPE, "custom_llm_provider": declared_provider} + + assert extract({"model": "gpt-4o-mini", **declared}) == "openai" + assert extract({"model": "no-such-model", **declared}) is None + + def test_extract_api_provider_swallows_unknown_model_but_logs_unexpected_errors(): """ get_llm_provider raises BadRequestError for a model that maps to no @@ -849,6 +896,56 @@ async def test_failure_hook_prefers_request_data_provider_over_exception_provide ) == ["azure"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type, custom_llm_provider, expected_requested_model", + [ + (_PASS_THROUGH_CALL_TYPE, "vertex_ai", "gcp/google/gemini-2.5-flash"), + ("acompletion", "vertex_ai", "other"), + (_PASS_THROUGH_CALL_TYPE, "bedrock", "other"), + ], +) +async def test_failure_hook_labels_a_pass_through_failure_with_the_router_model_the_client_requested( + call_type: str, custom_llm_provider: str, expected_requested_model: str +): + from unittest.mock import patch + + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy._types import UserAPIKeyAuth + + router = litellm.Router( + model_list=[ + { + "model_name": "gcp/google/gemini-2.5-flash", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "p", + "vertex_location": "us-central1", + }, + } + ] + ) + _clear_prometheus_registry() + try: + with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam + await PrometheusLogger().async_post_call_failure_hook( + request_data={ + "model": "gemini-2.5-flash", + "call_type": call_type, + "custom_llm_provider": custom_llm_provider, + "metadata": {}, + "litellm_params": {"metadata": {"model_group": "gcp/google/gemini-2.5-flash"}}, + }, + original_exception=Exception("upstream 400"), + user_api_key_dict=UserAPIKeyAuth(token="tok"), + ) + assert { + s.labels.get("requested_model") for s in _collected_samples("litellm_proxy_failed_requests_metric_total") + } == {expected_requested_model} + finally: + _clear_prometheus_registry() + + def test_model_group_in_deployment_metrics(): """ Test that model_group label is present on the deployment-scoped metrics diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 152ab2ffc13..3b222be6e5e 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1293,6 +1293,9 @@ def test_get_vertex_model_id_from_url(): model_id = get_vertex_model_id_from_url(url) assert model_id == "gemini-pro" + url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro?key=provider-secret#frag" + assert get_vertex_model_id_from_url(url) == "gemini-pro" + # Test with invalid URL url = "https://invalid-url.com" model_id = get_vertex_model_id_from_url(url) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index aab227e9494..ff27ac54a5b 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -25,12 +25,14 @@ from starlette.datastructures import FormData import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.prometheus import PrometheusLogger from tests._master_key import MASTER_KEY as SHARED_MASTER_KEY from litellm.caching.caching import DualCache -from litellm.types.utils import CallTypesLiteral +from litellm.types.utils import CallTypes, CallTypesLiteral from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException +from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id from litellm.proxy.utils import ProxyLogging from litellm.passthrough.main import raise_for_provider_status from litellm.llms.bedrock.common_utils import BedrockError @@ -6490,6 +6492,55 @@ class TestVertexAILiveWebsocketPassthrough: assert "default_vertex_config" in close_kwargs["reason"] assert len(close_kwargs["reason"].encode("utf-8")) <= 123 + @pytest.mark.asyncio + @pytest.mark.parametrize("model, expected_model", [("gemini-3.8-live", "gemini-3.8-live"), (None, "")]) + async def test_credential_failure_logs_the_named_model_under_vertex_ai( + self, monkeypatch: pytest.MonkeyPatch, model: str | None, expected_model: str + ): + from litellm.proxy.pass_through_endpoints import ( + llm_passthrough_endpoints as passthrough_module, + ) + + class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.request_data: tuple[dict[str, object], ...] = () + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + self.request_data = (*self.request_data, request_data) + + recorder: Final = _FailureRecorder() + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache()) + ) + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None) + self._clear_vertex_env(monkeypatch) + ensure_token: Final = AsyncMock(side_effect=Exception("Unable to find your credentials")) + + with patch.object(passthrough_module.vertex_llm_base, "ensure_access_token_async", ensure_token): + await passthrough_module.vertex_ai_live_websocket_passthrough( + websocket=self._websocket(), + model=model, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert [ + ( + d["model"], + PrometheusLogger._extract_api_provider_from_request_data(d), + get_spend_logs_id(CallTypes.pass_through.value, {}, d) is not None, + ) + for d in recorder.request_data + ] == [(expected_model, "vertex_ai", True)] + class TestAnthropicProxyRoute: """The /anthropic passthrough route: custom auth headers must not clobber the diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 7b62853b0e9..a9098ac3d59 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -30,20 +30,26 @@ from starlette.datastructures import UploadFile as StarletteUploadFile import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.prometheus import PrometheusLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( + VertexPassthroughLoggingHandler, +) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, SafeRouteAdder, + _build_passthrough_failure_request_payload, _registered_pass_through_routes, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, @@ -61,8 +67,11 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id +from litellm.proxy.utils import ProxyLogging from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypes from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, @@ -140,9 +149,7 @@ def test_is_multipart(): ), ], ) -def test_get_endpoint_type_classifies_anthropic_messages_routes( - url: str, expected_endpoint_type: EndpointType -) -> None: +def test_get_endpoint_type_classifies_anthropic_messages_routes(url: str, expected_endpoint_type: EndpointType) -> None: assert HttpPassThroughEndpointHelpers.get_endpoint_type(url) == expected_endpoint_type @@ -430,53 +437,18 @@ async def test_pass_through_request_failure_handler(): Critical Test: When a users pass through endpoint request fails, we must log the failure code, exception in litellm spend logs. """ - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client") as mock_get_client: - with patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processing: - # Setup mock for post_call_failure_hook and pre_call_hook - mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.pre_call_hook = AsyncMock() + request_failure = httpx.HTTPError("Request failed") + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") - # Setup mock for httpx client - mock_client = MagicMock() - mock_client.client = MagicMock() - mock_client.client.request = AsyncMock(side_effect=httpx.HTTPError("Request failed")) - mock_get_client.return_value = mock_client + failure: Final = await _run_passthrough_failure( + "http://test.com", None, request_failure, user_api_key_dict=user_api_key_dict + ) - # Mock headers for custom headers - mock_processing.get_custom_headers.return_value = {} - - # Create mock request - mock_request = MagicMock(spec=Request) - mock_request.method = "POST" - mock_request.body = AsyncMock(return_value=b'{"test": "data"}') - mock_request.headers = Headers({}) - - # Create a simple empty QueryParams - mock_request.query_params = QueryParams({}) - - # Create mock user API key dict - mock_user_api_key_dict = MagicMock() - - # Call the function with a target that will trigger an HTTPError - with pytest.raises(ProxyException): - await pass_through_request( - request=mock_request, - target="http://test.com", - custom_headers={}, - user_api_key_dict=mock_user_api_key_dict, - ) - - # Assert post_call_failure_hook was called - mock_proxy_logging.post_call_failure_hook.assert_called_once() - - # Verify the arguments to post_call_failure_hook - call_args = mock_proxy_logging.post_call_failure_hook.call_args[1] - assert call_args["user_api_key_dict"] == mock_user_api_key_dict - assert isinstance(call_args["original_exception"], TypeError) # Now expecting TypeError - assert "traceback_str" in call_args + assert failure.user_api_key_dict is user_api_key_dict + assert failure.original_exception is request_failure + assert "Request failed" in (failure.traceback_str or "") + assert failure.request_data["contents"] == "not a list" + assert failure.request_data["model"] == "" @pytest.mark.asyncio @@ -3656,9 +3628,7 @@ async def test_pass_through_request_preserves_target_query_without_client_query( @pytest.mark.asyncio -async def test_config_pass_through_without_configured_headers_preserves_no_authorization( - tmp_path, monkeypatch -): +async def test_config_pass_through_without_configured_headers_preserves_no_authorization(tmp_path, monkeypatch): proxy = await _boot_db_backed_proxy( tmp_path, monkeypatch, @@ -3725,9 +3695,7 @@ async def test_pass_through_rpm_limit_is_sequential_and_scoped_per_key(tmp_path, proxy = await _boot_db_backed_proxy( tmp_path, monkeypatch, - config_pass_through_endpoints=[ - {"path": "/rpm-rerank", "target": "http://upstream.test/rerank", "auth": True} - ], + config_pass_through_endpoints=[{"path": "/rpm-rerank", "target": "http://upstream.test/rerank", "auth": True}], db_pass_through_endpoints=[], master_key=MASTER_KEY, ) @@ -4902,8 +4870,10 @@ from fastapi import HTTPException as _FastAPIHTTPException from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, + ModifyResponseException, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + _PT_MODULE = "litellm.proxy.pass_through_endpoints.pass_through_endpoints" @@ -8182,8 +8152,10 @@ async def test_user_defined_passthrough_is_neither_tracked_nor_enforced(metadata } ) body: Final = { - "model": "upstream-only-model", metadata_slot: { - "model_group": "managed-model", "customer_label": "retained", + "model": "upstream-only-model", + metadata_slot: { + "model_group": "managed-model", + "customer_label": "retained", "user_api_key_team_model_max_budget": budget, }, } @@ -8191,48 +8163,84 @@ async def test_user_defined_passthrough_is_neither_tracked_nor_enforced(metadata assert await limiter.is_team_within_model_budget("shared-team", budget, None, "managed-model") start: Final = datetime.now() logging_obj: Final = LiteLLMLoggingObj( - model="upstream-only-model", messages=[], stream=False, call_type="pass_through_endpoint", - start_time=start, litellm_call_id="custom-budget", function_id="custom-budget", kwargs={}, + model="upstream-only-model", + messages=[], + stream=False, + call_type="pass_through_endpoint", + start_time=start, + litellm_call_id="custom-budget", + function_id="custom-budget", + kwargs={}, dynamic_async_success_callbacks=[limiter], ) payload: Final = { - "url": "https://upstream.test/echo", "request_body": body, "request_method": "POST", "cost_per_request": 0.25, + "url": "https://upstream.test/echo", + "request_body": body, + "request_method": "POST", + "cost_per_request": 0.25, } kwargs: Final = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( - request=request, user_api_key_dict=auth, passthrough_logging_payload=payload, logging_obj=logging_obj, - _parsed_body=body, litellm_call_id="custom-budget", + request=request, + user_api_key_dict=auth, + passthrough_logging_payload=payload, + logging_obj=logging_obj, + _parsed_body=body, + litellm_call_id="custom-budget", ) logging_obj.update_environment_variables( - model="upstream-only-model", user="unknown", optional_params={}, - litellm_params=kwargs["litellm_params"], call_type="pass_through_endpoint", + model="upstream-only-model", + user="unknown", + optional_params={}, + litellm_params=kwargs["litellm_params"], + call_type="pass_through_endpoint", ) response: Final = httpx.Response( - 200, request=httpx.Request("POST", "https://upstream.test/echo"), json={"ok": True}, + 200, + request=httpx.Request("POST", "https://upstream.test/echo"), + json={"ok": True}, ) await PassThroughEndpointLogging().pass_through_async_success_handler( - httpx_response=response, response_body={"ok": True}, request_body=body, logging_obj=logging_obj, - url_route="https://upstream.test/echo", result=response.text, start_time=start, end_time=datetime.now(), - cache_hit=False, **kwargs, + httpx_response=response, + response_body={"ok": True}, + request_body=body, + logging_obj=logging_obj, + url_route="https://upstream.test/echo", + result=response.text, + start_time=start, + end_time=datetime.now(), + cache_hit=False, + **kwargs, ) assert logging_obj.model_call_details["response_cost"] == 0.25 assert await limiter.is_team_within_model_budget("shared-team", budget, None, "managed-model") metadata: Final = kwargs["litellm_params"]["metadata"] assert (metadata["model_group"], metadata["customer_label"]) == ("managed-model", "retained") - assert metadata.keys().isdisjoint({ - "user_api_key_model_max_budget", "user_api_key_team_model_max_budget", - "user_api_key_user_model_max_budget", "user_api_key_end_user_model_max_budget", - }) + assert metadata.keys().isdisjoint( + { + "user_api_key_model_max_budget", + "user_api_key_team_model_max_budget", + "user_api_key_user_model_max_budget", + "user_api_key_end_user_model_max_budget", + } + ) @pytest.mark.parametrize("metadata_slot", ["metadata", "litellm_metadata"]) def test_builtin_passthrough_pins_model_group_to_the_resolved_model(metadata_slot: str) -> None: - request: Final = Request({ - "type": "http", "method": "POST", "path": "/gemini/v1beta/models/gemini-2.5-flash:generateContent", - "headers": [], "query_string": b"", - }) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/gemini/v1beta/models/gemini-2.5-flash:generateContent", + "headers": [], + "query_string": b"", + } + ) kwargs: Final = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( - request=request, user_api_key_dict=UserAPIKeyAuth(token="hash", user_id="u-1"), - passthrough_logging_payload=MagicMock(), logging_obj=MagicMock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(token="hash", user_id="u-1"), + passthrough_logging_payload=MagicMock(), + logging_obj=MagicMock(), _parsed_body={"contents": [], metadata_slot: {"model_group": "unbounded-client-choice"}}, ) assert kwargs["litellm_params"]["metadata"]["model_group"] == "gemini-2.5-flash" @@ -8591,6 +8599,581 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) +_GEMINI_FLASH_TARGET: Final = "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.8-flash:generateContent" +_VERTEX_FLASH_TARGET: Final = ( + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/google/models/gemini-3.8-flash:generateContent" +) +_VERTEX_CLAUDE_RAW_PREDICT_TARGET: Final = ( + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/anthropic/models/claude-opus-5:rawPredict" +) +_VERTEX_SEARCH_TARGET: Final = ( + "https://discoveryengine.googleapis.com/v1/projects/p/locations/global" + "/dataStores/default/servingConfigs/default:search" +) +_RELAYED_GEMINI_COUNT_TOKENS_TARGET: Final = "http://127.0.0.1:30732/v1beta/models/gemini-3.8-flash:countTokens" +_TRANSLATE_TARGET: Final = "https://translate.googleapis.com/v3/projects/p/locations/global:translateText" +_NON_GOOGLE_VERTEX_SHAPED_TARGET: Final = "https://llm.example.com/v1/projects/a/locations/b/models/my-model:predict" +_VERTEX_LIVE_TARGET: Final = ( + "wss://us-central1-aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" +) +_VEO_OPERATION_SUFFIX: Final = "/models/veo-3.0-generate-001/operations/op-8f2c1" +_RELAYED_VERTEX_COUNT_TOKENS_TARGET: Final = ( + "http://127.0.0.1:30732/v1/projects/p/locations/us-east5/publishers/google/models/gemini-3.8-flash:countTokens" +) + + +@dataclass(frozen=True, slots=True) +class _RecordedFailure: + request_data: dict[str, object] + original_exception: Exception + user_api_key_dict: UserAPIKeyAuth + traceback_str: str | None + + +class _FailureHookRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: tuple[_RecordedFailure, ...] = () + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + self.failures = ( + *self.failures, + _RecordedFailure(request_data, original_exception, user_api_key_dict, traceback_str), + ) + + +@dataclass(frozen=True, slots=True) +class _PassthroughRun: + failure: _RecordedFailure + response: Response | None + + +async def _drained_pass_through_response( + request: Request, + target: str, + custom_llm_provider: str | None, + user_api_key_dict: UserAPIKeyAuth, + stream: bool, +) -> Response | None: + try: + received: Final = await pass_through_request( + request=request, + target=target, + custom_headers={}, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + stream=stream, + ) + if isinstance(received, StreamingResponse): + async for _ in received.body_iterator: + pass + except (ProxyException, httpx.HTTPError): + return None + return received if isinstance(received, Response) else None + + +async def _run_passthrough_failure( + target: str, + custom_llm_provider: str | None, + upstream: httpx.Response | Exception, + stream: bool = False, + user_api_key_dict: UserAPIKeyAuth | None = None, + dispatched_endpoint: Callable[..., object] | None = None, + guardrail: CustomGuardrail | None = None, + body: bytes = b'{"contents": "not a list"}', +) -> _RecordedFailure: + run: Final = await _run_passthrough( + target, custom_llm_provider, upstream, stream, user_api_key_dict, dispatched_endpoint, guardrail, body + ) + return run.failure + + +async def _run_passthrough( + target: str, + custom_llm_provider: str | None, + upstream: httpx.Response | Exception, + stream: bool = False, + user_api_key_dict: UserAPIKeyAuth | None = None, + dispatched_endpoint: Callable[..., object] | None = None, + guardrail: CustomGuardrail | None = None, + body: bytes = b'{"contents": "not a list"}', +) -> _PassthroughRun: + recorder: Final = _FailureHookRecorder() + with ( + pytest.MonkeyPatch.context() as litellm_patch, + patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside pass_through_request; there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache()) + ), + respx.mock(assert_all_called=False) as upstream_mock, + ): + litellm_patch.setattr(litellm, "disable_aiohttp_transport", True) + litellm_patch.setattr(litellm, "callbacks", [recorder] if guardrail is None else [guardrail, recorder]) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream_route: Final = upstream_mock.route(method="POST") + if isinstance(upstream, Exception): + upstream_route.mock(side_effect=upstream) + else: + upstream_route.mock(return_value=upstream) + mock_request: Final = MagicMock(spec=Request) + mock_request.method = "POST" + proxy_url: Final = httpx.URL(target.replace("https://", "http://test-proxy.com/")) + mock_request.url = proxy_url + mock_request.scope = { + "type": "http", + "method": "POST", + "path": proxy_url.path, + **({"endpoint": dispatched_endpoint} if dispatched_endpoint is not None else {}), + } + mock_request.body = AsyncMock(return_value=body) + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + response: Final = await _drained_pass_through_response( + mock_request, + target, + custom_llm_provider, + user_api_key_dict or UserAPIKeyAuth(api_key="sk-test"), + stream, + ) + await asyncio.sleep(0) + litellm.in_memory_llm_clients_cache.flush_cache() + assert len(recorder.failures) == 1 + return _PassthroughRun(recorder.failures[0], response) + + +def _upstream_400(target: str) -> httpx.Response: + return httpx.Response( + status_code=400, + headers={"content-type": "application/json"}, + content=json.dumps({"error": {"code": 400, "status": "INVALID_ARGUMENT"}}).encode(), + request=httpx.Request("POST", target), + ) + + +def _upstream_400_as_event_stream(target: str) -> httpx.Response: + return httpx.Response( + status_code=400, + headers={"content-type": "text/event-stream"}, + content=b'data: {"error": {"code": 400, "status": "INVALID_ARGUMENT"}}\n\n', + request=httpx.Request("POST", target), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target, custom_llm_provider, stream, expected_provider", + [ + (_GEMINI_FLASH_TARGET, "gemini", False, "gemini"), + (_GEMINI_FLASH_TARGET, "gemini", True, "gemini"), + (_VERTEX_FLASH_TARGET, None, False, "vertex_ai"), + (_VERTEX_FLASH_TARGET, None, True, "vertex_ai"), + (_GEMINI_FLASH_TARGET.replace(":generateContent", ":countTokens"), "gemini", False, "gemini"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":countTokens"), None, False, "vertex_ai"), + (_RELAYED_GEMINI_COUNT_TOKENS_TARGET, "gemini", False, "gemini"), + ], +) +async def test_passthrough_upstream_error_logs_model_from_url( + target: str, custom_llm_provider: str | None, stream: bool, expected_provider: str +): + failure: Final = await _run_passthrough_failure(target, custom_llm_provider, _upstream_400(target), stream=stream) + assert failure.request_data["model"] == "gemini-3.8-flash" + assert failure.request_data["custom_llm_provider"] == expected_provider + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target", + [ + _RELAYED_VERTEX_COUNT_TOKENS_TARGET, + _TRANSLATE_TARGET, + _NON_GOOGLE_VERTEX_SHAPED_TARGET, + ], +) +async def test_passthrough_upstream_error_to_a_host_outside_gemini_and_vertex_gets_no_url_attribution(target: str): + failure: Final = await _run_passthrough_failure(target, None, _upstream_400(target)) + assert failure.request_data["model"] == "" + assert failure.request_data.get("custom_llm_provider") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("target", [_GEMINI_FLASH_TARGET, _VERTEX_FLASH_TARGET]) +async def test_passthrough_upstream_error_on_a_user_defined_route_to_google_gets_no_url_attribution(target: str): + failure: Final = await _run_passthrough_failure( + target, None, _upstream_400(target), dispatched_endpoint=_marked_pass_through_endpoint() + ) + assert failure.request_data["model"] == "" + assert failure.request_data.get("custom_llm_provider") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target, custom_llm_provider, expected_provider", + [ + (_GEMINI_FLASH_TARGET, "gemini", "gemini"), + (_VERTEX_FLASH_TARGET, None, "vertex_ai"), + ], +) +async def test_passthrough_detected_stream_upstream_error_logs_model_from_url( + target: str, custom_llm_provider: str | None, expected_provider: str +): + failure: Final = await _run_passthrough_failure( + target, custom_llm_provider, _upstream_400_as_event_stream(target), stream=False + ) + assert failure.request_data["model"] == "gemini-3.8-flash" + assert failure.request_data["custom_llm_provider"] == expected_provider + + +@pytest.mark.asyncio +async def test_passthrough_internal_failure_logs_model_from_url(): + failure: Final = await _run_passthrough_failure( + _GEMINI_FLASH_TARGET, "gemini", httpx.ReadTimeout("upstream timed out") + ) + assert failure.request_data["model"] == "gemini-3.8-flash" + assert failure.request_data["custom_llm_provider"] == "gemini" + + +def _upstream_dropping_stream(target: str) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "text/event-stream"}, + stream=_UpstreamDroppingMidStream(), + request=httpx.Request("POST", target), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize( + "target, custom_llm_provider, expected_provider", + [ + (_GEMINI_FLASH_TARGET.replace(":generateContent", ":streamGenerateContent?alt=sse"), "gemini", "gemini"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":streamGenerateContent?alt=sse"), None, "vertex_ai"), + ], +) +async def test_passthrough_mid_stream_drop_logs_model_from_url( + target: str, custom_llm_provider: str | None, expected_provider: str, stream: bool +): + failure: Final = await _run_passthrough_failure( + target, custom_llm_provider, _upstream_dropping_stream(target), stream=stream + ) + assert isinstance(failure.original_exception, httpx.ReadError) + assert failure.request_data["model"] == "gemini-3.8-flash" + assert failure.request_data["custom_llm_provider"] == expected_provider + + +class _ModifyResponseOnPreCallGuardrail(CustomGuardrail): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: str, + ) -> None: + self.raise_passthrough_exception(violation_message="blocked by policy", request_data=data, detection_info=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target, custom_llm_provider, body, expected_model, expected_provider", + [ + (_GEMINI_FLASH_TARGET, "gemini", b'{"contents": "not a list"}', "gemini-3.8-flash", "gemini"), + (_VERTEX_FLASH_TARGET, None, b'{"contents": "not a list"}', "gemini-3.8-flash", "vertex_ai"), + (_VERTEX_FLASH_TARGET, None, b'{"custom_llm_provider": "gemini"}', "gemini-3.8-flash", "vertex_ai"), + ], +) +async def test_passthrough_guardrail_block_logs_model_from_url( + target: str, custom_llm_provider: str | None, body: bytes, expected_model: str, expected_provider: str +): + guardrail: Final = _ModifyResponseOnPreCallGuardrail( + guardrail_name="block-policy", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + failure: Final = await _run_passthrough_failure( + target, custom_llm_provider, _upstream_400(target), guardrail=guardrail, body=body + ) + assert isinstance(failure.original_exception, ModifyResponseException) + assert failure.request_data["model"] == expected_model + assert failure.request_data["custom_llm_provider"] == expected_provider + assert PrometheusLogger._extract_api_provider_from_request_data(failure.request_data) == expected_provider + assert get_spend_logs_id(CallTypes.pass_through.value, {}, failure.request_data) is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [b'{"contents": "not a list"}', b'{"litellm_call_id": "client-chosen-id"}']) +async def test_passthrough_guardrail_block_response_names_the_call_id_of_its_spend_row(body: bytes): + guardrail: Final = _ModifyResponseOnPreCallGuardrail( + guardrail_name="block-policy", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + run: Final = await _run_passthrough( + _GEMINI_FLASH_TARGET, "gemini", _upstream_400(_GEMINI_FLASH_TARGET), guardrail=guardrail, body=body + ) + assert isinstance(run.failure.original_exception, ModifyResponseException) + assert run.response is not None + assert run.response.headers["x-litellm-call-id"] == get_spend_logs_id( + CallTypes.pass_through.value, {}, run.failure.request_data + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [b'{"contents": "not a list"}', b'{"custom_llm_provider": "vertex_ai"}']) +async def test_passthrough_guardrail_block_on_a_user_defined_route_gets_no_url_attribution(body: bytes): + guardrail: Final = _ModifyResponseOnPreCallGuardrail( + guardrail_name="block-policy", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + failure: Final = await _run_passthrough_failure( + _GEMINI_FLASH_TARGET, + None, + _upstream_400(_GEMINI_FLASH_TARGET), + dispatched_endpoint=_marked_pass_through_endpoint(), + guardrail=guardrail, + body=body, + ) + assert isinstance(failure.original_exception, ModifyResponseException) + assert failure.request_data["model"] == "" + assert "custom_llm_provider" not in failure.request_data + + +def _websocket_connect_failure(kind: str) -> Exception: + from websockets.datastructures import Headers as WebSocketHeaders + from websockets.exceptions import InvalidStatus + from websockets.http11 import Response as WebSocketHandshakeResponse + + if kind == "rejected_handshake": + return InvalidStatus(WebSocketHandshakeResponse(403, "Forbidden", WebSocketHeaders())) + return OSError("connection refused") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_kind", ["rejected_handshake", "unreachable_host"]) +@pytest.mark.parametrize( + "target, builtin_provider_route, expected_provider", + [ + (_VERTEX_LIVE_TARGET, True, "vertex_ai"), + (_VERTEX_LIVE_TARGET, False, None), + ("wss://api.openai.com/v1/realtime?model=gpt-realtime", True, None), + ], +) +async def test_websocket_passthrough_connect_failure_logs_the_provider_of_the_upstream_url( + failure_kind: str, target: str, builtin_provider_route: bool, expected_provider: str | None +): + failure: Final = await _run_websocket_connect_failure(failure_kind, target, builtin_provider_route) + assert failure.request_data.get("custom_llm_provider") == expected_provider + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_kind", ["rejected_handshake", "unreachable_host"]) +@pytest.mark.parametrize("model, expected_model", [("gemini-3.8-live", "gemini-3.8-live"), (None, "")]) +async def test_websocket_passthrough_connect_failure_logs_the_model_the_caller_named( + failure_kind: str, model: str | None, expected_model: str +): + failure: Final = await _run_websocket_connect_failure(failure_kind, _VERTEX_LIVE_TARGET, True, model=model) + assert failure.request_data["model"] == expected_model + + +async def _run_websocket_connect_failure( + failure_kind: str, target: str, builtin_provider_route: bool, model: str | None = None +) -> _RecordedFailure: + recorder: Final = _FailureHookRecorder() + with ( + pytest.MonkeyPatch.context() as litellm_patch, + patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside websocket_passthrough_request; there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache()) + ), + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", + side_effect=_websocket_connect_failure(failure_kind), + ), + ): + litellm_patch.setattr(litellm, "callbacks", [recorder]) + await websocket_passthrough_request( + websocket=_client_websocket(_pending_receive), + target=target, + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + endpoint="/vertex_ai/live", + accept_websocket=False, + builtin_provider_route=builtin_provider_route, + model=model, + ) + assert len(recorder.failures) == 1 + return recorder.failures[0] + + +@pytest.mark.parametrize( + "parsed_body, url, custom_llm_provider, expected_model, expected_provider", + [ + ({"model": "from-body"}, _GEMINI_FLASH_TARGET, "gemini", "from-body", "gemini"), + ({"contents": []}, _GEMINI_FLASH_TARGET, "gemini", "gemini-3.8-flash", "gemini"), + ({"contents": []}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash", "vertex_ai"), + ({"contents": []}, _VERTEX_FLASH_TARGET, "anthropic", "", "anthropic"), + ({}, "https://generativelanguage.googleapis.com/v1beta/openai/models/gemini-3.8-flash", "openai", "", "openai"), + ({"custom_llm_provider": "from-body"}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash", "vertex_ai"), + ({"custom_llm_provider": "gemini"}, _VERTEX_FLASH_TARGET, "anthropic", "", "anthropic"), + ({"contents": []}, _GEMINI_FLASH_TARGET, None, "gemini-3.8-flash", "gemini"), + ( + {}, + "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.8-flash?key=provider-secret", + "gemini", + "gemini-3.8-flash", + "gemini", + ), + ( + {}, + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/google/models/gemini-3.8-flash?key=provider-secret", + None, + "gemini-3.8-flash", + "vertex_ai", + ), + ({}, _VERTEX_CLAUDE_RAW_PREDICT_TARGET, None, "vertex_ai/claude-opus-5", "vertex_ai"), + ( + {}, + _VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":streamRawPredict?alt=sse"), + None, + "claude-opus-5", + "vertex_ai", + ), + ({}, "https://generativelanguage.googleapis.com/v1beta/models?pageToken=/models/x", "gemini", "", "gemini"), + ( + {"contents": []}, + "https://evilaiplatform.googleapis.com/v1/projects/p/locations/l" + "/publishers/google/models/gemini-3.8-flash:generateContent", + None, + "", + None, + ), + ({}, _VERTEX_SEARCH_TARGET, None, "vertex_ai/search_api", "vertex_ai"), + ({}, f"{_VERTEX_FLASH_TARGET}?tag=rawPredict", None, "gemini-3.8-flash", "vertex_ai"), + ( + {}, + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs" + "?filter=:search", + None, + "", + "vertex_ai", + ), + ( + {}, + f"https://generativelanguage.googleapis.com/v1beta{_VEO_OPERATION_SUFFIX}", + "gemini", + "veo-3.0-generate-001", + "gemini", + ), + ( + {}, + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1" + f"/publishers/google{_VEO_OPERATION_SUFFIX}", + None, + "veo-3.0-generate-001", + "vertex_ai", + ), + ( + {"contents": []}, + "https://aiplatform.us.rep.googleapis.com/v1/projects/p/locations/us" + "/publishers/google/models/gemini-3.8-flash:generateContent", + None, + "gemini-3.8-flash", + "vertex_ai", + ), + ( + {}, + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs", + None, + "", + "vertex_ai", + ), + ( + {"contents": []}, + _GEMINI_FLASH_TARGET.replace(":generateContent", ":countTokens"), + "gemini", + "gemini-3.8-flash", + "gemini", + ), + ( + {"contents": []}, + _VERTEX_FLASH_TARGET.replace(":generateContent", ":countTokens"), + None, + "gemini-3.8-flash", + "vertex_ai", + ), + ( + {"contents": []}, + "https://aiplatform.googleapis.com/v1/projects/p/locations/global" + "/publishers/google/models/gemini-3.8-flash:countTokens", + None, + "gemini-3.8-flash", + "vertex_ai", + ), + ( + {}, + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/cachedContents", + None, + "", + "vertex_ai", + ), + ({"contents": []}, _RELAYED_GEMINI_COUNT_TOKENS_TARGET, "gemini", "gemini-3.8-flash", "gemini"), + ({"contents": []}, _RELAYED_VERTEX_COUNT_TOKENS_TARGET, None, "", None), + ({"contents": []}, _TRANSLATE_TARGET, None, "", None), + ({"contents": []}, _NON_GOOGLE_VERTEX_SHAPED_TARGET, None, "", None), + ({"contents": []}, "https://api.example.com/v1/models/gpt-x:countTokens", None, "", None), + ({"contents": []}, "https://api.example.com/v1/models/gpt-x:predict", None, "", None), + ({"contents": []}, "https://api.example.com/v1/models/gpt-x:thing", None, "", None), + (None, None, None, "", None), + ], +) +def test_build_passthrough_failure_request_payload_model_and_provider_precedence( + parsed_body: dict | None, + url: str | None, + custom_llm_provider: str | None, + expected_model: str, + expected_provider: str | None, +): + request_payload: Final = _build_passthrough_failure_request_payload( + parsed_body=parsed_body, + kwargs=None, + logging_obj=None, + custom_llm_provider=custom_llm_provider, + url=httpx.URL(url) if url else None, + builtin_provider_route=True, + ) + assert request_payload["model"] == expected_model + assert request_payload.get("custom_llm_provider") == expected_provider + + +@pytest.mark.parametrize( + "url_route, expected_model", + [ + (_VERTEX_FLASH_TARGET, "gemini-3.8-flash"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":streamGenerateContent?alt=sse"), "gemini-3.8-flash"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":predict"), "gemini-3.8-flash"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":embedContent"), "gemini-3.8-flash"), + (_VERTEX_CLAUDE_RAW_PREDICT_TARGET, "vertex_ai/claude-opus-5"), + (_VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":streamRawPredict"), "vertex_ai/claude-opus-5"), + ( + _VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":streamRawPredict?alt=sse"), + "claude-opus-5", + ), + (_VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":rawPredict?alt=sse"), "claude-opus-5"), + (f"{_VERTEX_FLASH_TARGET}?tag=rawPredict", "gemini-3.8-flash"), + (_VERTEX_SEARCH_TARGET, "vertex_ai/search_api"), + ("https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs", None), + ( + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs" + "?filter=:search", + None, + ), + ], +) +def test_vertex_failure_model_name_matches_success_handler_naming(url_route: str, expected_model: str | None): + assert VertexPassthroughLoggingHandler.model_from_url_route(url_route) == expected_model + + @pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"]) def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_body(client_metadata_key: str): """A provider route that resolved a router deployment stashes its model_info on request.state. That @@ -8598,7 +9181,9 @@ def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_bo the call to (LIT-1761: passthrough successes carried model_id="").""" mock_request = MagicMock(spec=Request) mock_request.method = "POST" - mock_request.url = httpx.URL("http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent") + mock_request.url = httpx.URL( + "http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent" + ) mock_request.headers = Headers({}) mock_request.scope = {} mock_request.state = SimpleNamespace( @@ -9668,12 +10253,14 @@ async def _drain_logging_worker(): await GLOBAL_LOGGING_WORKER.stop() yield + @pytest.fixture() def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.mark.usefixtures("_drain_logging_worker", "_vcr_outcome_gate") def test_update_pass_through_route_updates_registry(): """ @@ -9763,9 +10350,7 @@ def mock_request(): return self._dict.values() class MockRequest: - def __init__( - self, headers=None, method="POST", request_body: Optional[dict] = None - ): + def __init__(self, headers=None, method="POST", request_body: Optional[dict] = None): self.headers = headers or {} self.query_params = QueryParams() self.method = method @@ -9873,9 +10458,9 @@ def test_pass_through_routes_support_all_methods(): path = route.path methods = set(route.methods) allowed = PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES.get(path, expected_methods) - assert ( - methods == allowed - ), f"Route {path} does not support all methods. Supported: {methods}, Expected: {allowed}" + assert methods == allowed, ( + f"Route {path} does not support all methods. Supported: {methods}, Expected: {allowed}" + ) check_router_methods(llm_router) @@ -9905,19 +10490,11 @@ def test_is_bedrock_agent_runtime_route(): # Test agent runtime endpoints (should return True) assert _is_bedrock_agent_runtime_route("/knowledgebases/kb-123/retrieve") is True - assert ( - _is_bedrock_agent_runtime_route("/agents/knowledgebases/kb-123/retrieve") - is True - ) + assert _is_bedrock_agent_runtime_route("/agents/knowledgebases/kb-123/retrieve") is True # Test regular bedrock runtime endpoints (should return False) - assert ( - _is_bedrock_agent_runtime_route("/guardrail/test-id/version/1/apply") is False - ) - assert ( - _is_bedrock_agent_runtime_route("/model/cohere.command-r-v1:0/converse") - is False - ) + assert _is_bedrock_agent_runtime_route("/guardrail/test-id/version/1/apply") is False + assert _is_bedrock_agent_runtime_route("/model/cohere.command-r-v1:0/converse") is False assert _is_bedrock_agent_runtime_route("/some/random/endpoint") is False @@ -9984,9 +10561,7 @@ def test_custom_pricing_used_in_cost_calculation(): # Verify custom cost is different from standard cost (unless prices happen to match) # This confirms custom pricing is actually being applied - assert ( - custom_cost != standard_cost - ), "Custom pricing should produce different cost than standard pricing" + assert custom_cost != standard_cost, "Custom pricing should produce different cost than standard pricing" # Test 3: Custom pricing with cache_read_input_token_cost and input_cost_per_token_batches # This specifically tests the parameters that were causing the original issue @@ -10562,6 +11137,7 @@ def test_init_kwargs_client_metadata_cannot_mark_a_pass_through_call_as_a_silent assert "is_silent_experiment" not in json.loads(row["metadata"]) assert result["litellm_params"]["metadata"]["client_key"] == "kept" + @pytest.mark.asyncio async def test_direct_passthrough_keeps_mapped_baggage_tags_outside_upstream_body(): from litellm.proxy.common_utils.http_parsing_utils import read_request_body From 8ac9b757370a27f51de80a8dcc84b26f743fe583 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:43:01 -0700 Subject: [PATCH 34/64] fix(model_prices): rolling registry audit, add Gemini batch and Azure Astra flex tiers (#49114) * fix(model_prices): add EmpirioLabs model rows from the provider catalog Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(model_prices): align EmpirioLabs rows with the registry endpoint vocabulary and schema Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(model_prices): absorb Gemini long-prompt batch rates and Azure GPT-6 Astra long-context flex rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(model_prices): drop EmpirioLabs rows from the registry audit --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 36 ++++++++++++++----- model_prices_and_context_window.json | 36 ++++++++++++++----- 2 files changed, 54 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0799a298942..28a8fe83fa2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -8715,7 +8715,7 @@ "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -8742,7 +8742,11 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 }, "azure/gpt-6-astra-2026-09-03": { "cache_creation_input_token_cost": 1.25e-05, @@ -8763,7 +8767,7 @@ "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -8790,7 +8794,11 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 }, "azure/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, @@ -30856,8 +30864,9 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/embeddings#model-versions", - "tpm": 10000000 + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "tpm": 10000000, + "input_cost_per_token_batches": 7.5e-08 }, "gemini/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, @@ -31713,7 +31722,10 @@ "input_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, - "supports_url_context": true + "supports_url_context": true, + "cache_read_input_token_cost_above_200k_tokens_batches": 2.5e-07, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06 }, "gemini/gemini-2.5-computer-use-preview-10-2025": { "deprecation_date": "2026-07-28", @@ -32280,7 +32292,10 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "output_cost_per_token_flex": 6e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06 }, "gemini/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -32343,7 +32358,10 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "output_cost_per_token_flex": 6e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06 }, "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0799a298942..28a8fe83fa2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -8715,7 +8715,7 @@ "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -8742,7 +8742,11 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 }, "azure/gpt-6-astra-2026-09-03": { "cache_creation_input_token_cost": 1.25e-05, @@ -8763,7 +8767,7 @@ "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -8790,7 +8794,11 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 }, "azure/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, @@ -30856,8 +30864,9 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/embeddings#model-versions", - "tpm": 10000000 + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "tpm": 10000000, + "input_cost_per_token_batches": 7.5e-08 }, "gemini/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, @@ -31713,7 +31722,10 @@ "input_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, - "supports_url_context": true + "supports_url_context": true, + "cache_read_input_token_cost_above_200k_tokens_batches": 2.5e-07, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06 }, "gemini/gemini-2.5-computer-use-preview-10-2025": { "deprecation_date": "2026-07-28", @@ -32280,7 +32292,10 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "output_cost_per_token_flex": 6e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06 }, "gemini/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -32343,7 +32358,10 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "output_cost_per_token_flex": 6e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06 }, "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, From 9432bf6b8ee686d0fbd06863e5f75889d383c8a2 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:01:44 -0700 Subject: [PATCH 35/64] perf(streaming): stop raising caught AttributeErrors inside pydantic on every streamed chunk (#49166) * perf(streaming): stop raising AttributeErrors inside pydantic on every streamed chunk Every later plain chunk on openai and azure streams raised and caught 12 AttributeErrors inside pydantic. Assigning an undeclared extra goes through a class lookup that raises in the model metaclass and is never memoized, and the hidden_params getter went through the private attribute fallback set_extra_attribute writes the extra and fields_set directly when a cached per-class check says pydantic would do the same, and falls back to setattr otherwise. The hidden_params getter reads the private storage directly. Emitted chunk JSON is unchanged * fix(streaming): match pydantic's own extra-field bookkeeping on every supported version Take the direct extra write only when the installed pydantic itself adds an assigned extra to model_fields_set (2.13 and later) and the class defines no attribute of that name, so pydantic 2.11, 2.12 and 2.14 see the same model state as a plain setattr. A deleted _hidden_params raises AttributeError again, and the mutable test models carry frozen-ok reasons. * fix(streaming): keep the normal write for models that do not allow extra fields openai's construct gives every model an extras dict, so a forbid model built that way still needs setattr's ValueError instead of a silent extra write. * test(streaming): cover a ClassVar declared without a value in the extra-field parity test * fix(streaming): type the streamed choice's logprobs as the parsed value main now stores --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 7 +- litellm/litellm_core_utils/hidden_params.py | 35 ++- .../litellm_core_utils/streaming_handler.py | 17 +- litellm/types/utils.py | 32 ++- .../litellm_core_utils/test_hidden_params.py | 139 +++++++++++- .../test_streaming_handler.py | 211 +++++++++++++++--- 6 files changed, 388 insertions(+), 53 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 2f81154ab83..20e10f63f83 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -9,12 +9,13 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import TypeIs # noqa: TID251 # narrows a duck-typed value to the Protocol its caller reads import litellm from litellm._logging import verbose_logger from litellm.constants import KUBERNETES_POD_ROUTING_KEY +from litellm.litellm_core_utils.hidden_params import set_extra_attribute from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionFinishReason @@ -710,9 +711,9 @@ def preserve_upstream_non_openai_attributes( _assign_attributes(model_response, original_chunk.model_dump(include=frozenset(candidate_keys)).items()) -def _assign_attributes(target: object, values: Iterable[tuple[str, object]]) -> None: +def _assign_attributes(target: BaseModel, values: Iterable[tuple[str, object]]) -> None: for key, value in values: - setattr(target, key, value) + set_extra_attribute(target, key, value) def safe_deep_copy(data): diff --git a/litellm/litellm_core_utils/hidden_params.py b/litellm/litellm_core_utils/hidden_params.py index a74337d2549..c04d7614b90 100644 --- a/litellm/litellm_core_utils/hidden_params.py +++ b/litellm/litellm_core_utils/hidden_params.py @@ -1,8 +1,9 @@ from collections.abc import Iterator, Mapping, MutableMapping +from functools import cache, lru_cache from types import MappingProxyType from typing import Final, Protocol, cast -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, create_model from litellm.types.llms.base import HiddenParams @@ -60,6 +61,38 @@ def stored_attribute(obj: object, name: str) -> object: return extras.get(name) +@cache +def _assigned_extras_join_fields_set() -> bool: + probe: Final = create_model("_ExtraAssignmentProbe", __config__=ConfigDict(extra="allow"))() + setattr(probe, "assigned", None) + return "assigned" in probe.model_fields_set + + +@lru_cache(maxsize=1024) +def _writes_plain_extra(model_class: type[BaseModel], name: str) -> bool: + config: Final = model_class.model_config + return ( + model_class.__setattr__ is BaseModel.__setattr__ + and config.get("extra") == "allow" + and not config.get("frozen") + and not config.get("validate_assignment") + and not name.startswith("_") + and name not in model_class.__pydantic_fields__ + and name not in model_class.__class_vars__ + and not hasattr(model_class, name) + and _assigned_extras_join_fields_set() + ) + + +def set_extra_attribute(model: BaseModel, name: str, value: object) -> None: + extras: Final = model.__pydantic_extra__ + if extras is None or not _writes_plain_extra(type(model), name): + setattr(model, name, value) + return + extras[name] = value + model.__pydantic_fields_set__.add(name) + + def _get_hidden_params_storage(obj: object) -> object | None: if isinstance(obj, dict): return obj.get(_HIDDEN_PARAMS_ATTR) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 098a404ea14..b1b3e20f218 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_report from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, get_hidden_params, + set_extra_attribute, set_hidden_param, stored_attribute, ) @@ -1035,7 +1036,7 @@ class CustomStreamWrapper: _initial_delta: Final = model_response.choices[0].delta.model_dump() _initial_delta.pop("role", None) - model_response.choices[0].delta = Delta(**_initial_delta) + set_extra_attribute(model_response.choices[0], "delta", Delta(**_initial_delta)) return model_response def _prebuilt_plain_delta(self, chunk: object) -> Delta | None: @@ -1061,13 +1062,13 @@ class CustomStreamWrapper: def _set_role_on_fresh_delta(self, model_response: ModelResponseStream) -> None: delta: Final = model_response.choices[0].delta if self.sent_first_chunk is False: - delta.role = "assistant" + set_extra_attribute(delta, "role", "assistant") self.sent_first_chunk = True return if not _holds_only_scalars(delta): self.strip_role_from_delta(model_response) return - delta.role = None + set_extra_attribute(delta, "role", None) if "provider_specific_fields" not in delta.model_fields_set: delta.provider_specific_fields = None @@ -1186,8 +1187,8 @@ class CustomStreamWrapper: and _holds_only_scalars(upstream_delta) ) if can_reuse_choice and upstream_choice is not None: - model_response.choices[0].finish_reason = None # rebind-ok: keep reused choice - model_response.choices[0].index = upstream_choice.index # rebind-ok: keep reused choice + set_extra_attribute(model_response.choices[0], "finish_reason", None) + set_extra_attribute(model_response.choices[0], "index", upstream_choice.index) _prepend_upstream_choice_extras( model_response.choices[0], upstream_choice, @@ -1212,11 +1213,7 @@ class CustomStreamWrapper: else: return model_response.system_fingerprint = upstream_chunk.system_fingerprint - setattr( - model_response, - "citations", - stored_attribute(upstream_chunk, "citations"), - ) + set_extra_attribute(model_response, "citations", stored_attribute(upstream_chunk, "citations")) preserve_upstream_non_openai_attributes( model_response=model_response, original_chunk=original_chunk, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 55aaf606175..f8115571dfd 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -75,6 +75,7 @@ from litellm.types.workload_identity import ( ) from ..litellm_core_utils.core_helpers import map_finish_reason, process_response_headers +from ..litellm_core_utils.hidden_params import set_extra_attribute from . import litellm_params as _litellm_params from .agents import LiteLLMSendMessageResponse from .guardrails import GuardrailEventHooks @@ -2149,6 +2150,13 @@ class UnreportedUsage(Usage): class StreamingChoices(OpenAIObject): + if TYPE_CHECKING: + finish_reason: str | None + index: int + delta: Delta + enhancements: object + logprobs: ChoiceLogprobs | JsonValue + def __init__( self, finish_reason=None, @@ -2162,25 +2170,22 @@ class StreamingChoices(OpenAIObject): # https://github.com/BerriAI/litellm/issues/8455 params.pop("message", None) super().__init__(**params) - if finish_reason: - self.finish_reason = map_finish_reason(finish_reason) - else: - self.finish_reason = None - self.index = index + set_extra_attribute(self, "finish_reason", map_finish_reason(finish_reason) if finish_reason else None) + set_extra_attribute(self, "index", index) if delta is not None: if isinstance(delta, Delta): - self.delta = delta + set_extra_attribute(self, "delta", delta) elif isinstance(delta, dict): - self.delta = Delta(**delta) + set_extra_attribute(self, "delta", Delta(**delta)) else: - self.delta = Delta() + set_extra_attribute(self, "delta", Delta()) if enhancements is not None: - self.enhancements = enhancements + set_extra_attribute(self, "enhancements", enhancements) if isinstance(logprobs, (dict, BaseModel)): - self.logprobs = CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs) + set_extra_attribute(self, "logprobs", CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs)) else: - self.logprobs = logprobs + set_extra_attribute(self, "logprobs", logprobs) def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator @@ -2234,7 +2239,10 @@ class ModelResponseBase(OpenAIObject): @property def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation - return self._hidden_params + private: Final = self.__pydantic_private__ + if private is None or "_hidden_params" not in private: + return self._hidden_params + return private["_hidden_params"] @hidden_params.setter def hidden_params(self, hidden_params: dict[str, object]) -> None: # mutable-ok: API requires mutation diff --git a/tests/unit/litellm_core_utils/test_hidden_params.py b/tests/unit/litellm_core_utils/test_hidden_params.py index c0485a07696..b32692c67ba 100644 --- a/tests/unit/litellm_core_utils/test_hidden_params.py +++ b/tests/unit/litellm_core_utils/test_hidden_params.py @@ -1,7 +1,11 @@ -from typing import Final +import pickle +from collections.abc import Callable +from typing import ClassVar, Final import pytest +from openai import BaseModel as OpenAIBaseModel from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel, ConfigDict, PrivateAttr from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, @@ -9,13 +13,14 @@ from litellm.litellm_core_utils.hidden_params import ( get_hidden_params_storage, get_or_create_hidden_params, served_from_cache, + set_extra_attribute, set_hidden_param, set_hidden_params, ) from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig from litellm.types.decisions import DecisionsResponse from litellm.types.llms.base import HiddenParams -from litellm.types.utils import ModelResponse, TextChoices, TextCompletionResponse, Usage +from litellm.types.utils import ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse, Usage def test_get_and_set_hidden_params_on_plain_object() -> None: @@ -309,3 +314,133 @@ def test_served_from_cache_is_true_only_for_a_response_the_cache_stamped( assert served_from_cache(model_response) is expected assert served_from_cache({HIDDEN_PARAMS_ATTR: dict(hidden_params)}) is expected assert served_from_cache(object()) is False + + +class _ExtraModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="allow") + + declared: str | None = None + marker: ClassVar[str | None] = None + unset_marker: ClassVar[int] + _secret: str = PrivateAttr(default="hidden") + + @property + def doubled(self) -> str | None: + return self.declared + + @doubled.setter + def doubled(self, value: str) -> None: + self.declared = value * 2 + + +class _FrozenExtraModel(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + +class _ValidatedExtraModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="allow", validate_assignment=True) + + __pydantic_extra__: dict[str, int] + + +class _IgnoredExtraModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="ignore") + + +class _MirroringModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="allow") + + def __setattr__(self, name: str, value: object) -> None: + super().__setattr__(name, value) + super().__setattr__(f"{name}_copy", value) + + +class _NoneClassAttribute: + inherited = None + + +class _InheritingExtraModel(_NoneClassAttribute, BaseModel): # frozen-ok: assignment tests need attribute writes + model_config = ConfigDict(extra="allow") + + +class _ForbiddingOpenAIModel(OpenAIBaseModel): # frozen-ok: assignment tests need attribute writes + model_config = ConfigDict(extra="forbid") + + +def _assignment_outcome(model: BaseModel, assign: Callable[[], None], name: str) -> tuple[object, ...]: + try: + assign() + except (AttributeError, ValueError) as error: + return type(error), str(error) + return model.model_dump(), getattr(model, name), model.model_fields_set, model.model_dump_json() + + +@pytest.mark.parametrize( + "build, name, value", + [ + pytest.param(lambda: _ExtraModel(existing=1, later=3), "fresh", 1, id="new-extra"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "existing", 2, id="existing-extra-keeps-order"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "declared", "x", id="declared-field"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "doubled", "ab", id="property"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "_free", 1, id="undeclared-underscore-name"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "_secret", "x", id="private-attribute"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "marker", 1, id="class-var"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "unset_marker", 1, id="class-var-without-value"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "model_config", 1, id="class-attribute"), + pytest.param(_FrozenExtraModel, "fresh", 1, id="frozen"), + pytest.param(_ValidatedExtraModel, "fresh", "5", id="validate-assignment"), + pytest.param(_IgnoredExtraModel, "fresh", 1, id="extra-ignore"), + pytest.param(_MirroringModel, "fresh", 1, id="setattr-override"), + pytest.param(_InheritingExtraModel, "inherited", 1, id="class-attribute-set-to-none"), + pytest.param( + lambda: _ForbiddingOpenAIModel.construct(existing=1), "fresh", 1, id="extra-forbid-built-by-construct" + ), + ], +) +def test_set_extra_attribute_matches_setattr(build: Callable[[], BaseModel], name: str, value: object) -> None: + through_setattr: Final = build() + through_helper: Final = build() + + expected: Final = _assignment_outcome(through_setattr, lambda: setattr(through_setattr, name, value), name) + actual: Final = _assignment_outcome(through_helper, lambda: set_extra_attribute(through_helper, name, value), name) + + assert actual == expected + + +@pytest.mark.parametrize("response_type", [ModelResponse, ModelResponseStream]) +def test_hidden_params_raises_attribute_error_once_deleted( + response_type: type[ModelResponse | ModelResponseStream], +) -> None: + response: Final = response_type() + del response._hidden_params + + assert getattr(response, "hidden_params", "fallback") == "fallback" + + +def _round_trip(response: ModelResponse | ModelResponseStream) -> ModelResponse | ModelResponseStream: + return pickle.loads(pickle.dumps(response)) + + +@pytest.mark.parametrize("response_type", [ModelResponse, ModelResponseStream]) +@pytest.mark.parametrize( + "rebuild", + [ + pytest.param(lambda response: response, id="init"), + pytest.param(lambda response: response.model_copy(), id="model-copy"), + pytest.param(lambda response: response.model_copy(deep=True), id="deep-copy"), + pytest.param(_round_trip, id="pickle"), + ], +) +def test_hidden_params_is_the_stored_dict_on_every_construction_path( + response_type: type[ModelResponse | ModelResponseStream], + rebuild: Callable[[ModelResponse | ModelResponseStream], ModelResponse | ModelResponseStream], +) -> None: + original: Final = response_type() + original.hidden_params["model_id"] = "deployment-1" + response: Final = rebuild(original) + + response.hidden_params["api_base"] = "https://example.com" + + assert response.hidden_params is response._hidden_params + assert get_hidden_params(response) == response._hidden_params + assert response._hidden_params["api_base"] == "https://example.com" diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index b9c8feb223d..148b7193ef6 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -11,7 +11,7 @@ import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path -from types import CodeType, FrameType, MappingProxyType +from types import CodeType, FrameType, MappingProxyType, TracebackType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -6292,53 +6292,76 @@ def test_openai_chunk_creator_skips_full_chunk_dump_and_delta_rebuild( assert finish_output.choices[0].delta.content is None -def _openai_chunk(chunk_id: str, delta: Mapping[str, object]) -> ChatCompletionChunk: +def _openai_chunk( + chunk_id: str, + delta: Mapping[str, object], + choice_extras: Mapping[str, object] = MappingProxyType({}), + finish_reason: str | None = None, + chunk_extras: Mapping[str, object] = MappingProxyType({}), +) -> ChatCompletionChunk: return ChatCompletionChunk( id=chunk_id, created=1, model="gpt-4o", object="chat.completion.chunk", - choices=[{"index": 0, "delta": delta, "finish_reason": None}], + choices=[{"index": 0, "delta": delta, "finish_reason": finish_reason, **choice_extras}], + **chunk_extras, ) -def _missed_attributes(wrapper: CustomStreamWrapper, chunks: Sequence[ChatCompletionChunk]) -> tuple[str, ...]: - attribute_lookup: Final = BaseModel.__getattr__.__code__ - misses: Final[list[str]] = [] +def _raised_attribute_errors(wrapper: CustomStreamWrapper, chunks: Sequence[ChatCompletionChunk]) -> tuple[str, ...]: + raised: Final[list[str]] = [] - def record_miss(frame: FrameType, event: str, arg: object) -> Callable[[FrameType, str, object], object]: - if event == "exception" and isinstance(arg, tuple) and isinstance(arg[1], AttributeError): - misses.append(str(frame.f_locals["item"])) - return record_miss - - def trace_lookups(frame: FrameType, event: str, _arg: object) -> Callable[[FrameType, str, object], object] | None: - return record_miss if event == "call" and frame.f_code is attribute_lookup else None + def record_raise(frame: FrameType, event: str, arg: object) -> Callable[[FrameType, str, object], object]: + if event == "exception" and isinstance(arg, tuple): + _, error, traceback = arg + if isinstance(error, AttributeError) and isinstance(traceback, TracebackType) and traceback.tb_next is None: + raised.append(f"{frame.f_code.co_qualname}: {error}") + return record_raise previous_tracer: Final = sys.gettrace() - sys.settrace(trace_lookups) + sys.settrace(record_raise) try: tuple(wrapper.chunk_creator(chunk=chunk) for chunk in chunks) finally: sys.settrace(previous_tracer) - return tuple(misses) + return tuple(raised) -def test_openai_chunk_creator_builds_one_delta_and_misses_no_attribute_per_later_chunk( - logging_obj: Logging, +_AZURE_CHOICE_EXTRAS: Final = MappingProxyType( + {"content_filter_results": {"hate": {"filtered": False, "severity": "safe"}}} +) +_OPENAI_AND_AZURE_CHOICE_EXTRAS: Final = ( + pytest.param("openai", MappingProxyType({}), id="openai"), + pytest.param("azure", _AZURE_CHOICE_EXTRAS, id="azure"), +) + + +@pytest.mark.parametrize(("custom_llm_provider", "choice_extras"), _OPENAI_AND_AZURE_CHOICE_EXTRAS) +def test_openai_chunk_creator_builds_one_delta_and_raises_no_attribute_error_per_later_chunk( + logging_obj: Logging, custom_llm_provider: str, choice_extras: Mapping[str, object] ) -> None: wrapper: Final = CustomStreamWrapper( completion_stream=iter(()), model="gpt-4o", - custom_llm_provider="openai", + custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, ) - wrapper.chunk_creator(chunk=_openai_chunk("chunk-1", {"role": "assistant", "content": ""})) - plain_chunks: Final = ( - _openai_chunk("chunk-2", {"content": "hello"}), - _openai_chunk("chunk-3", {"content": " world"}), + + def later_chunk(chunk_id: str, delta: Mapping[str, object]) -> ChatCompletionChunk: + return _openai_chunk( + chunk_id, delta, choice_extras, chunk_extras={"service_tier": "default", "obfuscation": chunk_id} + ) + + nested_delta: Final = {"content": "!", "reasoning_details": [{"type": "reasoning.text", "text": "thought"}]} + first_chunk_of_each_shape: Final = ( + later_chunk("chunk-1", {"role": "assistant", "content": ""}), + later_chunk("chunk-2", nested_delta), ) - nested_chunk: Final = _openai_chunk( - "chunk-4", {"content": "!", "reasoning_details": [{"type": "reasoning.text", "text": "thought"}]} + tuple(wrapper.chunk_creator(chunk=chunk) for chunk in first_chunk_of_each_shape) + plain_chunks: Final = ( + later_chunk("chunk-3", {"content": "hello"}), + later_chunk("chunk-4", {"content": " world"}), ) delta_init: Final = Delta.__init__.__code__ streaming_choices_init: Final = StreamingChoices.__init__.__code__ @@ -6348,7 +6371,145 @@ def test_openai_chunk_creator_builds_one_delta_and_misses_no_attribute_per_later assert sum(call.code is delta_init for call in calls) == len(plain_chunks) assert sum(call.code is streaming_choices_init for call in calls) == len(plain_chunks) assert [output.choices[0].delta.content if output else None for output in outputs] == ["hello", " world"] - assert _missed_attributes(wrapper, (*plain_chunks, nested_chunk)) == () + assert _raised_attribute_errors(wrapper, (*plain_chunks, later_chunk("chunk-5", nested_delta))) == () + + +_OPENAI_STREAM_JSON: Final = ( + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":null,"index":0,' + '"delta":{"provider_specific_fields":null,"refusal":null,"content":"","role":"assistant",' + '"function_call":null,"tool_calls":null,"audio":null},"logprobs":null}],' + '"provider_specific_fields":null,"citations":null,"service_tier":"default","obfuscation":"o0"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":null,"index":0,' + '"delta":{"provider_specific_fields":null,"refusal":null,"content":"one","role":null,' + '"function_call":null,"tool_calls":null,"audio":null},"logprobs":null}],' + '"provider_specific_fields":null,"citations":null,"service_tier":"default","obfuscation":"o1"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":null,"index":0,' + '"delta":{"provider_specific_fields":null,"refusal":null,"content":" two","role":null,' + '"function_call":null,"tool_calls":null,"audio":null},"logprobs":null}],' + '"provider_specific_fields":null,"citations":null,"service_tier":"default","obfuscation":"o2"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":"stop","index":0,' + '"delta":{"provider_specific_fields":null,"content":null,"role":null,"function_call":null,' + '"tool_calls":null,"audio":null},"logprobs":null}],"provider_specific_fields":null,' + '"service_tier":"default","obfuscation":"o3"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[],"provider_specific_fields":null,' + '"usage":{"completion_tokens":3,"prompt_tokens":9,"total_tokens":12,"completion_tokens_details":null,' + '"prompt_tokens_details":null}}' + ), +) +_AZURE_STREAM_JSON: Final = ( + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"content_filter_results":{"hate":{"filtered":false,' + '"severity":"safe"}},"finish_reason":null,"index":0,"delta":{"provider_specific_fields":null,' + '"refusal":null,"content":"","role":"assistant","function_call":null,"tool_calls":null,"audio":null},' + '"logprobs":null}],"provider_specific_fields":null,"citations":null,"service_tier":"default",' + '"obfuscation":"o0"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"content_filter_results":{"hate":{"filtered":false,' + '"severity":"safe"}},"finish_reason":null,"index":0,"delta":{"provider_specific_fields":null,' + '"refusal":null,"content":"one","role":null,"function_call":null,"tool_calls":null,"audio":null},' + '"logprobs":null}],"provider_specific_fields":null,"citations":null,"service_tier":"default",' + '"obfuscation":"o1"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"content_filter_results":{"hate":{"filtered":false,' + '"severity":"safe"}},"finish_reason":null,"index":0,"delta":{"provider_specific_fields":null,' + '"refusal":null,"content":" two","role":null,"function_call":null,"tool_calls":null,"audio":null},' + '"logprobs":null}],"provider_specific_fields":null,"citations":null,"service_tier":"default",' + '"obfuscation":"o2"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":"stop","index":0,' + '"delta":{"provider_specific_fields":null,"content":null,"role":null,"function_call":null,' + '"tool_calls":null,"audio":null},"logprobs":null}],"provider_specific_fields":null,' + '"service_tier":"default","obfuscation":"o3"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[],"provider_specific_fields":null,' + '"usage":{"completion_tokens":3,"prompt_tokens":9,"total_tokens":12,"completion_tokens_details":null,' + '"prompt_tokens_details":null}}' + ), +) + + +@pytest.mark.parametrize( + ("custom_llm_provider", "choice_extras", "expected_json"), + [ + pytest.param("openai", MappingProxyType({}), _OPENAI_STREAM_JSON, id="openai"), + pytest.param("azure", _AZURE_CHOICE_EXTRAS, _AZURE_STREAM_JSON, id="azure"), + ], +) +def test_openai_chunk_creator_emits_each_chunk_as_the_same_json( + logging_obj: Logging, + custom_llm_provider: str, + choice_extras: Mapping[str, object], + expected_json: tuple[str | None, ...], +) -> None: + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="gpt-4o", + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + stream_options={"include_usage": True}, + ) + chunks: Final = ( + _openai_chunk( + "chatcmpl-1", + {"role": "assistant", "content": ""}, + choice_extras, + chunk_extras={"service_tier": "default", "obfuscation": "o0"}, + ), + _openai_chunk( + "chatcmpl-1", + {"content": "one"}, + choice_extras, + chunk_extras={"service_tier": "default", "obfuscation": "o1"}, + ), + _openai_chunk( + "chatcmpl-1", + {"content": " two"}, + choice_extras, + chunk_extras={"service_tier": "default", "obfuscation": "o2"}, + ), + _openai_chunk( + "chatcmpl-1", {}, choice_extras, "stop", chunk_extras={"service_tier": "default", "obfuscation": "o3"} + ), + ChatCompletionChunk( + id="chatcmpl-1", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[], + service_tier="default", + obfuscation="o4", + usage={"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12}, + ), + ) + + outputs: Final = tuple(wrapper.chunk_creator(chunk=chunk) for chunk in chunks) + + assert tuple( + output.model_copy(update={"created": 1}).model_dump_json() if output else None for output in outputs + ) == expected_json def test_chunk_creator_skips_full_dump_of_model_response_stream_chunks( From ab3f6caddcd096b820359a7ae2829e9440b77c4f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:18:29 -0700 Subject: [PATCH 36/64] fix(gemini): flag computer use support on the four models the Gemini computer use docs list (#49252) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 10 +++++++--- model_prices_and_context_window.json | 10 +++++++--- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 28a8fe83fa2..e1131117fcf 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -31883,7 +31883,8 @@ }, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "supports_computer_use": true }, "gemini/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, @@ -31949,7 +31950,8 @@ "input_cost_per_token_flex": 2.5e-07, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, - "supports_audio_input": true + "supports_audio_input": true, + "supports_computer_use": true }, "gemini/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -32196,7 +32198,8 @@ }, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "supports_computer_use": true }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -76558,6 +76561,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 28a8fe83fa2..e1131117fcf 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -31883,7 +31883,8 @@ }, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "supports_computer_use": true }, "gemini/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, @@ -31949,7 +31950,8 @@ "input_cost_per_token_flex": 2.5e-07, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, - "supports_audio_input": true + "supports_audio_input": true, + "supports_computer_use": true }, "gemini/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -32196,7 +32198,8 @@ }, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "supports_computer_use": true }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -76558,6 +76561,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, From 372c096388f3102419e0177f501f42126c7170c7 Mon Sep 17 00:00:00 2001 From: Etienne Chabert Date: Fri, 30 Oct 2026 18:20:00 +0100 Subject: [PATCH 37/64] fix(proxy): evict a deleted key from every worker's auth cache (#47609) * fix(proxy): evict a deleted key from every worker's auth cache /key/delete dropped the key object only from the in-memory cache of the worker that served the call. Auth serves key objects cache-first, so every other worker, and the Redis copy when enable_redis_auth_cache is on, kept authenticating the deleted key until user_api_key_cache_ttl expired. /key/update already evicts and broadcasts through evict_and_broadcast, delete now does the same for the tokens it deletes * fix(proxy): evict only the deleted rows and broadcast them in one batch Scope the cross-worker eviction on /key/delete to the rows the caller was allowed to delete and send it through delete_cache_key_objects, so a caller-supplied name can no longer drop another team's cached entries on every worker * refactor(proxy): drop the comment above the delete eviction call * fix(proxy): run the usage cache eviction beside the auth eviction so a hung Redis costs one wait * test(proxy): wait on the broadcast instead of a sleep and type the peer eviction tests * test(proxy): wait for the peer to apply the broadcast instead of polling its cache * test(proxy): cover hashed keys and aliases when deleting a key warmed on a peer * fix(key_management): evict every requested key on delete, including keys whose row is already gone A batch delete that names a live key and a key whose row was removed outside the API evicted only the live key, so the gone key kept authenticating from every worker's cache. Every requested key hash is evicted now, while raw non-key cache keys are still left alone * test(key_management): split the gone-row eviction assertion into one comprehension per worker --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 37 +-- .../key_management_endpoints.py | 15 +- .../management/test_key_delete_revocation.py | 15 +- tests/unit/proxy/auth/test_auth_checks.py | 50 +++- .../test_key_management_endpoints.py | 231 +++++++++++++++++- 5 files changed, 317 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f80296f4bd7..79a46ef355e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3457,28 +3457,35 @@ async def delete_cache_key_objects( proxy_logging_obj: ProxyLogging | None, ) -> None: """ - Evict a batch of key objects, for callers that delete keys in bulk rather than through - `/key/delete`. Auth resolves a cached key object without re-reading its team, so a key left - cached after its row is gone keeps buying access until its TTL expires. - - Evicting locally only reaches this worker; `_delete_cache_key_object` itself broadcasts each - token, so a deleted key left in a peer worker's in-memory cache still authenticates there until - its TTL expires. + Evict a batch of deleted key objects here and on every other worker. Auth resolves a cached key + object without re-reading its team, so a key left cached after its row is gone keeps buying + access until its TTL expires. Best-effort per key: the rows are already deleted by the time this runs, so an unreachable cache backend must not abort the caller partway through its own cascade. - Every token is dropped from this worker's memory first, then from the auth cache and the - internal usage cache in one batch each, so a deleted key stops authenticating here the instant - its rows are gone instead of refilling from Redis on the next miss. Every other worker then - gets the eviction in one pipelined broadcast: a publish per key would pass the publisher's - pending cap on a bulk delete, and the dropped keys would keep authenticating there until their - TTL. + Every token is dropped from this worker's memory first, then from the auth cache's Redis layer, + so a deleted key stops authenticating here the instant its rows are gone instead of refilling + from Redis on the next miss. Every other worker then gets the eviction in one pipelined + broadcast: a publish per key would pass the publisher's pending cap on a bulk delete, and the + dropped keys would keep authenticating there until their TTL. The internal usage cache is not + an auth refill source, so its batch delete runs beside that chain and a hung Redis costs one + bounded wait there instead of another in line. """ evict_local(cache_keys=hashed_tokens, user_api_key_cache=user_api_key_cache) + usage_eviction: Final = ( + () + if proxy_logging_obj is None + else (evict_shared(cache_keys=hashed_tokens, cache=proxy_logging_obj.internal_usage_cache.dual_cache),) + ) + _ = await asyncio.gather( + _evict_auth_cache_then_broadcast(hashed_tokens=hashed_tokens, user_api_key_cache=user_api_key_cache), + *usage_eviction, + ) + + +async def _evict_auth_cache_then_broadcast(hashed_tokens: Sequence[str], user_api_key_cache: UserApiKeyCache) -> None: await evict_shared(cache_keys=hashed_tokens, cache=user_api_key_cache) - if proxy_logging_obj is not None: - await evict_shared(cache_keys=hashed_tokens, cache=proxy_logging_obj.internal_usage_cache.dual_cache) await broadcast_evictions(cache_keys=hashed_tokens) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a6814b593ca..3021b21de51 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -50,6 +50,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo _delete_cache_key_object, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export can_team_access_model, delete_cache_key_object, + delete_cache_key_objects, get_jwt_key_mapping_cache_keys_for_token, get_key_end_user_budget_id, get_org_object, @@ -78,7 +79,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( from litellm.proxy.common_utils.openai_error_payload import openai_error_param, proxy_exception_for from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time -from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache, is_user_key_cache_key from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.management.teams.authz import TEAM_ADMIN_ONLY, TEAM_OR_ORG_ADMIN, is_team_admin @@ -5483,7 +5484,7 @@ async def delete_verification_tokens( - List of keys being deleted, this contains information about the key_alias, token, and user_id being deleted, this is passed down to the KeyManagementEventHooks to delete the keys from the secret manager and handle audit logs """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj failed_tokens: list = [] try: @@ -5567,11 +5568,11 @@ async def delete_verification_tokens( verbose_proxy_logger.debug(traceback.format_exc()) raise e - for key in tokens: - user_api_key_cache.delete_cache(key) - # remove hash token from cache - hashed_token = hash_token(cast(str, key)) - user_api_key_cache.delete_cache(hashed_token) + await delete_cache_key_objects( + hashed_tokens=tuple(token for token in hashed_tokens if is_user_key_cache_key(token)), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) # After credential invalidation, so a failure here can never keep a deleted key alive. for deleted_key in authorized_keys: diff --git a/tests/integration/management/test_key_delete_revocation.py b/tests/integration/management/test_key_delete_revocation.py index 5b3527a52ae..bdc91dfe5f6 100644 --- a/tests/integration/management/test_key_delete_revocation.py +++ b/tests/integration/management/test_key_delete_revocation.py @@ -171,11 +171,8 @@ def test_deleting_a_key_revokes_it(gateway: Gateway, spelling: str) -> None: assert deleted_rows == (_DeletedKeyRow(token=key_hash),) -def test_deleting_a_key_revokes_it_on_a_warmed_peer(gateway: Gateway, peer: Gateway) -> None: - pytest.skip( - "BUG: POST /key/delete does not broadcast the key eviction, so a warmed peer keeps answering 200 for about " - "60 seconds until its in-memory cache TTL expires" - ) +@pytest.mark.parametrize("spelling", ("keys", "hashed_keys", "key_aliases")) +def test_deleting_a_key_revokes_it_on_a_warmed_peer(gateway: Gateway, peer: Gateway, spelling: str) -> None: backend: Final = "integration-keys-auth-delete" prompts: Final = ( "delete warm primary " + uuid.uuid4().hex, @@ -208,9 +205,13 @@ def test_deleting_a_key_revokes_it_on_a_warmed_peer(gateway: Gateway, peer: Gate ) assert len(warm_responses) == 2 assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,) * 2 - deleted: Final = gateway.request("POST", "/key/delete", {"keys": [key]}) + identifier: Final = {"keys": key, "hashed_keys": sha256(key.encode()).hexdigest(), "key_aliases": alias}[ + spelling + ] + delete_body: Final = {"key_aliases" if spelling == "key_aliases" else "keys": [identifier]} + deleted: Final = gateway.request("POST", "/key/delete", delete_body) assert deleted.status_code == 200, deleted.text - assert _DeleteResponse.model_validate_json(deleted.content).deleted_keys == (key,), deleted.text + assert _DeleteResponse.model_validate_json(deleted.content).deleted_keys == (identifier,), deleted.text peer_refusal: Final = eventually( lambda: _chat(peer, model, key, prompts[2]), lambda response: response.status_code == 401, diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e783f01769e..9d4fe67614e 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -46,7 +46,7 @@ from litellm.proxy.auth.auth_checks import ( _team_soft_budget_check, common_checks, ) -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import InternalUsageCache, ProxyLogging from litellm.proxy.utils import CallInfo from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER @@ -2425,3 +2425,51 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ assert refilled == [] assert redis.entries == {} + + +class _SignallingPublishRedisClient(_SlowPublishRedisClient): + def __init__(self) -> None: + super().__init__() + self.sent = asyncio.Event() + + async def _send(self, messages: Sequence[str]) -> None: + await super()._send(messages) + self.sent.set() + + +class _DeleteAfterBroadcastRedisCache(_AuthRedisCache): + def __init__(self, broadcast_sent: asyncio.Event) -> None: + super().__init__({}) + self._broadcast_sent = broadcast_sent + self.finished = False + + async def delete_cache_keys(self, keys: Sequence[str]) -> None: + await self._broadcast_sent.wait() + self.finished = True + + +@pytest.mark.asyncio +async def test_bulk_eviction_broadcasts_while_the_usage_cache_redis_is_wedged( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16)) + monkeypatch.setattr(pubsub_module, "_pending_publishes", set()) + client: Final = _SignallingPublishRedisClient() + usage_redis: Final = _DeleteAfterBroadcastRedisCache(broadcast_sent=client.sent) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=None) + proxy_logging_obj.internal_usage_cache = InternalUsageCache( + dual_cache=DualCache(redis_cache=usage_redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer + ) + hashed_tokens: Final = ("token-a", "token-b") + + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_PubSubRedisCache(client), + ): + await delete_cache_key_objects( + hashed_tokens=hashed_tokens, user_api_key_cache=UserApiKeyCache(), proxy_logging_obj=proxy_logging_obj + ) + await _drain_publishes() + + assert usage_redis.finished + assert {json.loads(message)["cache_key"] for message in client.published} == set(hashed_tokens) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 62dd367b8e5..c26a4bf19b8 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -7,7 +7,7 @@ from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch -from uuid import UUID +from uuid import UUID, uuid4 import pytest import yaml @@ -16,6 +16,7 @@ from fastapi.testclient import TestClient from pydantic import JsonValue, ValidationError import litellm +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.models.object_permission import LiteLLM_ObjectPermissionTable @@ -41,12 +42,15 @@ from litellm.proxy._types import ( RegenerateKeyRequest, ResetSpendRequest, UpdateKeyRequest, + hash_token, ) from litellm.proxy.auth.auth_checks import ( delete_cache_key_object, + get_key_object, jwt_key_mapping_cache_key, ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -5755,6 +5759,231 @@ async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypat ) +class _ChannelSubscription: + """Marks a message done when the subscriber comes back for the next one, so ``messages.join()`` + returns once every delivered message has been applied""" + + def __init__(self) -> None: + self.messages: asyncio.Queue[dict[str, object]] = asyncio.Queue() + self._holding_message = False + + async def subscribe(self, *channels: str) -> None: + return None + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> dict[str, object] | None: + if self._holding_message: + self._holding_message = False + self.messages.task_done() + try: + message: Final = await asyncio.wait_for(self.messages.get(), timeout) + except asyncio.TimeoutError: + return None + self._holding_message = True + return message + + async def aclose(self) -> None: + return None + + +class _ChannelPipeline: + def __init__(self, redis: "_SharedCoordinationRedis") -> None: + self._redis = redis + self._messages: tuple[str, ...] = () + + def publish(self, channel: str, message: str) -> "_ChannelPipeline": + self._messages = (*self._messages, message) + return self + + async def execute(self) -> None: + for message in self._messages: + self._redis.deliver(message) + + +class _SharedCoordinationRedis: + """The coordination Redis two workers share: whatever one worker publishes reaches the other's subscriber""" + + namespace = None + + def __init__(self) -> None: + self.published: tuple[str, ...] = () + self._subscriptions: tuple[_ChannelSubscription, ...] = () + + def init_pubsub_client(self) -> "_SharedCoordinationRedis": + return self + + def pubsub(self) -> _ChannelSubscription: + subscription: Final = _ChannelSubscription() + self._subscriptions = (*self._subscriptions, subscription) + return subscription + + def pipeline(self, transaction: bool = True) -> _ChannelPipeline: + return _ChannelPipeline(self) + + async def publish(self, channel: str, message: str) -> int: + self.deliver(message) + return 1 + + def deliver(self, message: str) -> None: + self.published = (*self.published, message) + for subscription in self._subscriptions: + subscription.messages.put_nowait({"type": "message", "data": message.encode()}) + + async def until_every_subscriber_applied_its_messages(self) -> None: + _ = await asyncio.gather(*(subscription.messages.join() for subscription in self._subscriptions)) + + +def _owned_key_row(hashed_token: str) -> LiteLLM_VerificationToken: + return LiteLLM_VerificationToken( + token=hashed_token, + user_id="user-123", + team_id=None, + key_alias="doomed-key", + spend=0.0, + max_budget=None, + models=[], + aliases={}, + config={}, + permissions={}, + metadata={}, + model_max_budget={}, + model_spend={}, + soft_budget_cooldown=False, + allowed_routes=[], + ) + + +def _prisma_with_key_row(key_row: LiteLLM_VerificationToken) -> AsyncMock: + prisma_client: Final = AsyncMock() + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[key_row]) + prisma_client.db.litellm_jwtkeymapping = _CascadingJWTMappingTable([]) + prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + prisma_client.delete_data = AsyncMock(return_value=[key_row.token]) + prisma_client.get_data = AsyncMock(return_value=None) + return prisma_client + + +@pytest.mark.asyncio +async def test_delete_verification_tokens_stops_the_key_on_a_peer_worker(monkeypatch: pytest.MonkeyPatch) -> None: + """Auth serves key objects cache-first, so a peer worker that is not told about the delete keeps + authenticating the deleted key from its own copy until user_api_key_cache_ttl expires""" + raw_key: Final = "sk-" + uuid4().hex + hashed_token: Final = hash_token(raw_key) + deleting_worker_cache: Final = UserApiKeyCache() + peer_worker_cache: Final = UserApiKeyCache() + for cache in (deleting_worker_cache, peer_worker_cache): + cache.set_cache(key=hashed_token, value=UserAPIKeyAuth(token=hashed_token, user_id="user-123")) + coordination_redis: Final = _SharedCoordinationRedis() + peer_subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=coordination_redis, user_api_key_cache=peer_worker_cache + ) + peer_subscriber.start() + prisma_client: Final = _prisma_with_key_row(_owned_key_row(hashed_token)) + monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", coordination_redis) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + + try: + await delete_verification_tokens( + tokens=[raw_key], + user_api_key_cache=deleting_worker_cache, + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value + ), + ) + await coordination_redis.until_every_subscriber_applied_its_messages() + finally: + await peer_subscriber.stop() + + assert deleting_worker_cache.get_cache(key=hashed_token) is None + with pytest.raises(ProxyException) as denied: + await get_key_object( + hashed_token=hashed_token, prisma_client=prisma_client, user_api_key_cache=peer_worker_cache + ) + assert denied.value.code == "401" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deleter", + [ + UserAPIKeyAuth(user_id="admin-user", api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value), + UserAPIKeyAuth(user_id="user-123", api_key="sk-owner", user_role=LitellmUserRoles.INTERNAL_USER.value), + ], + ids=["proxy_admin", "key_owner"], +) +async def test_delete_verification_tokens_evicts_a_requested_key_whose_row_is_already_gone( + monkeypatch: pytest.MonkeyPatch, deleter: UserAPIKeyAuth +) -> None: + """A key whose row was removed outside the API stays cached on every worker, so deleting it in a + batch beside a live key must still stop it everywhere instead of only evicting the rows found""" + live_key: Final = "sk-" + uuid4().hex + gone_key: Final = "sk-" + uuid4().hex + live_token: Final = hash_token(live_key) + gone_token: Final = hash_token(gone_key) + deleting_worker_cache: Final = UserApiKeyCache() + peer_worker_cache: Final = UserApiKeyCache() + for cache in (deleting_worker_cache, peer_worker_cache): + for token in (live_token, gone_token): + cache.set_cache(key=token, value=UserAPIKeyAuth(token=token, user_id="user-123")) + coordination_redis: Final = _SharedCoordinationRedis() + peer_subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=coordination_redis, user_api_key_cache=peer_worker_cache + ) + peer_subscriber.start() + monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", coordination_redis) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _prisma_with_key_row(_owned_key_row(live_token))) + + try: + await delete_verification_tokens( + tokens=[live_key, gone_key], user_api_key_cache=deleting_worker_cache, user_api_key_dict=deleter + ) + await coordination_redis.until_every_subscriber_applied_its_messages() + finally: + await peer_subscriber.stop() + + assert [deleting_worker_cache.get_cache(key=token) for token in (live_token, gone_token)] == [None, None] + assert [peer_worker_cache.get_cache(key=token) for token in (live_token, gone_token)] == [None, None] + + +@pytest.mark.asyncio +async def test_delete_verification_tokens_leaves_unowned_cache_entries_alone_on_every_worker(monkeypatch: pytest.MonkeyPatch) -> None: + """/key/delete accepts any string, and a name that is not a key row passes through unhashed. A key + owner padding the request with another team's spend counter name must not wipe that counter from + every worker's memory: only the rows the caller was allowed to delete are evicted and broadcast""" + raw_key: Final = "sk-" + uuid4().hex + hashed_token: Final = hash_token(raw_key) + other_team_counter: Final = "spend:team:other-team" + deleting_worker_cache: Final = UserApiKeyCache() + peer_worker_cache: Final = UserApiKeyCache() + peer_spend_counters: Final = InMemoryCache() + peer_spend_counters.set_cache(other_team_counter, 42.0) + peer_worker_cache.set_cache(key=hashed_token, value=UserAPIKeyAuth(token=hashed_token, user_id="user-123")) + coordination_redis: Final = _SharedCoordinationRedis() + peer_subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=coordination_redis, + user_api_key_cache=peer_worker_cache, + additional_in_memory_caches=(peer_spend_counters,), + ) + peer_subscriber.start() + monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", coordination_redis) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _prisma_with_key_row(_owned_key_row(hashed_token))) + + try: + await delete_verification_tokens( + tokens=[other_team_counter, raw_key], + user_api_key_cache=deleting_worker_cache, + user_api_key_dict=UserAPIKeyAuth( + user_id="user-123", api_key="sk-owner", user_role=LitellmUserRoles.INTERNAL_USER.value + ), + ) + await coordination_redis.until_every_subscriber_applied_its_messages() + finally: + await peer_subscriber.stop() + + assert peer_worker_cache.get_cache(key=hashed_token) is None + assert peer_spend_counters.get_cache(other_team_counter) == 42.0 + assert coordination_redis.published == (json.dumps({"cache_key": hashed_token}),) + + @pytest.mark.asyncio async def test_delete_key_fn_persists_deleted_keys(monkeypatch): from litellm.proxy._types import KeyRequest From 17d1f5833689739003abf379950137d7e8eeb2f5 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 10:31:57 -0700 Subject: [PATCH 38/64] fix(ui): preserve team model grants on unrelated settings saves (#49258) --- .../ModelSelect/ModelSelect.test.tsx | 33 +++++++++++ .../components/ModelSelect/ModelSelect.tsx | 5 +- .../team/TeamInfo.integration.test.tsx | 59 ++++++++++++++++++- .../src/components/team/TeamInfo.test.tsx | 1 - .../src/components/team/TeamInfo.tsx | 5 +- 5 files changed, 97 insertions(+), 6 deletions(-) diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx index 7ab757cfed9..d7394b1f949 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx @@ -775,6 +775,39 @@ describe("ModelSelect", () => { expect(mockOnChange).toHaveBeenCalledWith(["gpt-4"]); }); + it.each([ + { value: "all-proxy-models", label: "All Proxy Models", includeSpecialOptions: true }, + { value: "no-default-models", label: "No Default Models", includeSpecialOptions: false }, + ])( + "should label retained $value without offering it again after removal", + async ({ value, label, includeSpecialOptions }) => { + const user = userEvent.setup(); + mockUseOrganization.mockReturnValue({ + data: createMockOrganization(["gpt-4"]), + isLoading: false, + } as unknown as ReturnType); + const props = { + onChange: mockOnChange, + context: "team" as const, + organizationID: "org-1", + options: { includeSpecialOptions }, + }; + const { rerender } = renderWithProviders(); + + expect(screen.getByLabelText(label)).toBeInTheDocument(); + await openModelList(user); + const retainedGroup = within(screen.getByRole("listbox")).getAllByRole("group")[0]; + expect(within(retainedGroup).getByText("Unavailable")).toBeInTheDocument(); + await user.click(within(retainedGroup).getByRole("option", { name: label })); + expect(mockOnChange).toHaveBeenCalledWith([]); + + rerender(); + expectNotOffered(label); + expectNotOffered("claude-3"); + expectOffered("gpt-4"); + }, + ); + it("should keep the other selections when removing one unavailable model alongside a special option", async () => { const user = userEvent.setup(); renderWithProviders( diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx index f169382b4ed..c7c428f2ccf 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx @@ -240,7 +240,10 @@ export const ModelSelect = (props: ModelSelectProps) => { ]; const offeredByValue = new Map(offeredGroups.flatMap((group) => group.items).map((option) => [option.value, option])); - const selectedOptions = value.map((v) => offeredByValue.get(v) ?? { label: v, value: v }); + const selectedOptions = value.map( + (v) => + offeredByValue.get(v) ?? MODEL_SENTINEL_OPTIONS.find((option) => option.value === v) ?? { label: v, value: v }, + ); const groups: ModelOptionGroup[] = [ ...unavailableGroups( selectedOptions, diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx index fc054e51e5b..545c504281b 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx @@ -1,6 +1,6 @@ import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import * as networking from "@/components/networking"; -import { fireEvent, screen, waitFor } from "@testing-library/react"; +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; @@ -817,7 +817,7 @@ describe("TeamInfoView - member budget reset prompt", () => { }); }); -describe("TeamInfoView - model rate limit picker", () => { +describe("TeamInfoView - model settings", () => { const props = { teamId: "123", onUpdate: vi.fn(), @@ -836,6 +836,61 @@ describe("TeamInfoView - model rate limit picker", () => { vi.clearAllMocks(); }); + it.each([ + { name: "unrestricted", models: [], organizationId: null }, + { name: "explicit all models", models: ["all-proxy-models"], organizationId: null }, + { name: "access groups only", models: ["no-default-models"], organizationId: null }, + { name: "retained explicit models", models: ["retired-model", "openai/*"], organizationId: null }, + { name: "unrestricted within an organization", models: [], organizationId: "org-1" }, + ])("preserves $name model grants when saving only a budget", async ({ models, organizationId }) => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ models, organization_id: organizationId }), + ); + mockUseTeam.mockReturnValue({ + data: { team_id: "123", organization_models: organizationId ? ["gpt-4"] : undefined }, + isLoading: false, + } as unknown as ReturnType); + + renderWithProviders(); + await user.click(await screen.findByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + if (models.length === 0) { + expect(await screen.findByLabelText("All Proxy Models")).toBeInTheDocument(); + } + fireEvent.change(await screen.findByLabelText("Max Budget (USD)"), { target: { value: "100" } }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalledTimes(1)); + const payload = vi.mocked(networking.teamUpdateCall).mock.calls[0][1]; + expect(payload).toMatchObject({ team_id: "123", max_budget: 100 }); + expect(payload).not.toHaveProperty("models"); + }); + + it("discards a cancelled model clear but saves an intentional clear of unrestricted access", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: [] })); + + renderWithProviders(); + await user.click(await screen.findByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await user.click(within(await screen.findByLabelText("All Proxy Models")).getByRole("button")); + await user.click(screen.getByRole("button", { name: "Cancel" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + expect(await screen.findByLabelText("All Proxy Models")).toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "100" } }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalledTimes(1)); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1]).not.toHaveProperty("models"); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await user.click(within(await screen.findByLabelText("All Proxy Models")).getByRole("button")); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalledTimes(2)); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[1][1].models).toEqual(["no-default-models"]); + }); + it("offers access-group models, not no-default-models, in the model rate limit picker", async () => { const user = userEvent.setup({ delay: null }); vi.mocked(networking.teamInfoCall).mockResolvedValue( diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index e4bf67ab72f..46081ae96f5 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -2421,7 +2421,6 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { const alwaysSent = { team_id: "123", team_alias: "Test Team", - models: ["gpt-4"], tpm_limit: 1000, rpm_limit: 1000, tpd_limit: null, diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 0d20d160283..329552d735a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -72,6 +72,7 @@ import { } from "../key_team_helpers/ModelMaxBudgetEditor"; import { modelMaxBudgetUpdate, StoredModelMaxBudget } from "../key_team_helpers/modelMaxBudgetPayload"; import { + ALL_PROXY_MODELS, computeTeamModelBadges, normalizeTeamModelSelection, TeamAccessGroupModelGrant, @@ -515,7 +516,7 @@ const computeEffectiveGuardrails = (info: TeamInfoRecord, globalGuardrailNames: const toTeamFormValues = (info: TeamInfoRecord, effectiveGuardrails: string[]): TeamUpdateFormValues => ({ team_alias: info.team_alias, - models: info.models, + models: info.models.length === 0 ? [ALL_PROXY_MODELS] : info.models, max_budget: info.max_budget, soft_budget: info.soft_budget, soft_budget_alerting_emails: Array.isArray(info.metadata?.soft_budget_alerting_emails) @@ -1064,7 +1065,7 @@ const TeamInfoView: React.FC = ({ const updateData: any = { team_id: teamId, team_alias: values.team_alias, - models: normalizeTeamModelSelection(values.models), + ...(form.getFieldState("models").isDirty ? { models: normalizeTeamModelSelection(values.models) } : {}), tpm_limit: numberOrNull(values.tpm_limit), rpm_limit: numberOrNull(values.rpm_limit), tpd_limit: numberOrNull(values.tpd_limit), From 707ac1ffee620e2906ef0fb5f9ad699cfd65a16d Mon Sep 17 00:00:00 2001 From: islomjon-aiand Date: Sat, 31 Oct 2026 02:55:26 +0900 Subject: [PATCH 39/64] fix(aiand)!: show ai&'s name and logo, and sync ai&'s model list (#49158) * feat(ui): show ai&'s name and logo in the dashboard The Add Model form already lists ai& from its provider metadata, but the dashboard had no entry for it, so the provider showed no logo and model tables printed the raw aiand slug. Adds ai& to the dashboard provider list with its official mark, plated in dark mode so the red and black mark keep their colours, and the aiand/openai/gpt-oss-120b model placeholder * fix(aiand): retire deepseek-v4-flash, glm-5.2 and kimi-k2.7-code ai& retired these three models on 2026-10-27 and https://api.aiand.com/v1/models now lists ten. The rows keep their prices with a deprecation_date, so earlier spend still prices, and aiand/* stops listing them, the same way baseten handles retired rows. The other ten rows already match the live prices, context windows and output limits * feat(aiand): add deepseek-v4.1-flash-1017 ai& launches deepseek-ai/deepseek-v4.1-flash-1017 on 2026-11-02. It is served on https://staging-api.aiand.com/v1/models with the same price, context window, output limit and capabilities as deepseek-v4.1-flash, so its row copies that one * fix(aiand): retire deepseek-v4.1-flash ai& retires deepseek-ai/deepseek-v4.1-flash on 2026-10-30, now that deepseek-v4.1-flash-1017 is live in production on https://api.aiand.com/v1/models with the same price, limits and capabilities. Its row keeps its prices with a deprecation_date, like the other three retired ai& rows * fix(aiand): retire deepseek-v4.1-flash on 2026-11-06 https://api.aiand.com/v1/models now gives deepseek-v4.1-flash a deprecation_date of 2026-11-05T15:00Z, which is 2026-11-06 00:00 in Japan, so its row uses 2026-11-06 instead of 2026-10-30 and aiand/* keeps listing it until then. The wildcard test now leaves a model out only once its retirement date has passed * test(aiand): derive retired models from the rows instead of pinning dates The wildcard test now reads each row's deprecation_date, so it holds whichever dates ai& sets, and the test that only restated the JSON dates is gone * fix(aiand): retire models on their UTC date while the proxy runs, and add muse-spark-open The wildcard listing now checks each row's deprecation_date when it expands, against the UTC date the field documents, so a model drops on its day without a restart. deepseek-v4.1-flash now carries 2026-11-05, the UTC day ai& retires it, and meta/muse-spark-open gets its price row from ai&'s live /v1/models * fix(aiand): hide retired models only from bundled listings, including the wildcard health check * fix(aiand): share the UTC retirement check with Together AI and the health check sort --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/__init__.py | 21 ++-- .../litellm_core_utils/llm_request_utils.py | 21 ++-- .../litellm_core_utils/model_retirement.py | 19 ++++ ...odel_prices_and_context_window_backup.json | 34 ++++++ litellm/utils.py | 10 +- model_prices_and_context_window.json | 34 ++++++ .../test_llm_request_utils.py | 38 +++++++ .../baseten/chat/test_baseten_completions.py | 33 +++++- .../llms/openai_like/test_aiand_provider.py | 100 +++++++++++++++++- .../public/assets/logos/aiand.svg | 4 + .../components/provider_info_helpers.test.tsx | 12 +++ .../src/components/provider_info_helpers.tsx | 5 + .../src/lib/logoTreatments.ts | 1 + 13 files changed, 296 insertions(+), 36 deletions(-) create mode 100644 litellm/litellm_core_utils/model_retirement.py create mode 100644 ui/litellm-dashboard/public/assets/logos/aiand.svg diff --git a/litellm/__init__.py b/litellm/__init__.py index 5f14c605cd0..f64531f86fb 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -7,7 +7,6 @@ _apply_process_defaults() ### INIT VARIABLES ######################### import threading import os -import datetime from collections.abc import Mapping, Sequence from typing import ( @@ -550,6 +549,7 @@ _key_management_system: Optional["KeyManagementSystem"] = None output_parse_pii: bool = False ############################################# from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete +from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired model_cost = get_model_cost_map(url=model_cost_map_url) cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount @@ -756,16 +756,6 @@ def is_openai_finetune_model(key: str) -> bool: return key.startswith("ft:") and not key.count(":") > 1 -def _is_past_deprecation_date(value: Mapping[str, object]) -> bool: - deprecation_date: Final = value.get("deprecation_date") - if not isinstance(deprecation_date, str): - return False - try: - return datetime.date.fromisoformat(deprecation_date) <= datetime.date.today() - except ValueError: - return False - - def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None: for key, value in tuple(model_cost_map.items()): if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key): @@ -873,12 +863,12 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje if "-to-" not in key and "fireworks-ai-default" not in key: fireworks_ai_models.add(key) elif value.get("litellm_provider") == "together_ai": - if _is_past_deprecation_date(value): + if is_hidden_as_retired(value): together_ai_models.discard(key) elif not key.startswith("together-ai-"): together_ai_models.add(key) elif value.get("litellm_provider") == "baseten": - if _is_past_deprecation_date(value): + if is_hidden_as_retired(value): baseten_models.discard(key) else: baseten_models.add(key) @@ -1051,7 +1041,10 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje elif value.get("litellm_provider") == "coralbricks": coralbricks_models.add(key) elif value.get("litellm_provider") == "aiand": - aiand_models.add(key) + if is_hidden_as_retired(value): + aiand_models.discard(key) + else: + aiand_models.add(key) elif value.get("litellm_provider") == "reducto": reducto_models.add(key) elif value.get("litellm_provider") == "bedrock_mantle": diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index 9c2ec0b8984..84f161e0c62 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -1,4 +1,3 @@ -import datetime from collections.abc import Iterable, Iterator, Mapping from dataclasses import dataclass from typing import Final, cast @@ -9,6 +8,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider +from litellm.litellm_core_utils.model_retirement import is_past_deprecation_date def _form_field_value(value: object) -> str: @@ -168,7 +168,9 @@ class RankedChatModel: def rank_chat_models_from_llm_provider(custom_llm_provider: str) -> tuple[RankedChatModel, ...]: - known_models: Final[Iterable[str]] = litellm.models_by_provider.get(custom_llm_provider, ()) + from litellm.utils import bundled_provider_models + + known_models: Final = bundled_provider_models(custom_llm_provider) ranked_models: Final = ( RankedChatModel( model=model, route_rank=route_rank, cost=cost, is_past_deprecation_date=is_past_deprecation_date @@ -187,18 +189,11 @@ _COST_MAP_DEPRECATION: Final = TypeAdapter(_CostMapDeprecation) def _is_past_deprecation_date(cost_map_key: str) -> bool: - deprecation_date: Final = _COST_MAP_DEPRECATION.validate_python( - litellm.model_cost.get(cost_map_key) or {} # pyright: ignore[reportUnknownMemberType] # untyped cost map JSON - ).get("deprecation_date") - if not isinstance(deprecation_date, str): - return False - try: - return ( - datetime.date.fromisoformat(deprecation_date) - <= datetime.datetime.now(datetime.timezone.utc).astimezone().date() + return is_past_deprecation_date( + _COST_MAP_DEPRECATION.validate_python( + litellm.model_cost.get(cost_map_key) or {} # pyright: ignore[reportUnknownMemberType] # untyped cost map JSON ) - except ValueError: - return False + ) def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float, bool]]: diff --git a/litellm/litellm_core_utils/model_retirement.py b/litellm/litellm_core_utils/model_retirement.py new file mode 100644 index 00000000000..3bfb834733d --- /dev/null +++ b/litellm/litellm_core_utils/model_retirement.py @@ -0,0 +1,19 @@ +import datetime +from collections.abc import Mapping +from typing import Final + +PROVIDERS_HIDING_RETIRED_MODELS: Final = frozenset({"aiand", "baseten", "together_ai"}) + + +def is_past_deprecation_date(value: Mapping[str, object]) -> bool: + deprecation_date: Final = value.get("deprecation_date") + if not isinstance(deprecation_date, str): + return False + try: + return datetime.date.fromisoformat(deprecation_date) <= datetime.datetime.now(datetime.timezone.utc).date() + except ValueError: + return False + + +def is_hidden_as_retired(value: Mapping[str, object]) -> bool: + return value.get("litellm_provider") in PROVIDERS_HIDING_RETIRED_MODELS and is_past_deprecation_date(value) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e1131117fcf..ea9cacbdabe 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -72249,6 +72249,7 @@ }, "aiand/deepseek-ai/deepseek-v4-flash": { "cache_read_input_token_cost": 8e-08, + "deprecation_date": "2026-10-27", "input_cost_per_token": 1.5e-07, "litellm_provider": "aiand", "max_input_tokens": 1048576, @@ -72278,6 +72279,22 @@ "supports_vision": false }, "aiand/deepseek-ai/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 2e-08, + "deprecation_date": "2026-11-05", + "input_cost_per_token": 3e-07, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/deepseek-ai/deepseek-v4.1-flash-1017": { "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 3e-07, "litellm_provider": "aiand", @@ -72307,8 +72324,24 @@ "supports_tool_choice": true, "supports_vision": true }, + "aiand/meta/muse-spark-open": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 943718, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "aiand/moonshotai/kimi-k2.7-code": { "cache_read_input_token_cost": 2e-07, + "deprecation_date": "2026-10-27", "input_cost_per_token": 7.5e-07, "litellm_provider": "aiand", "max_input_tokens": 262144, @@ -72399,6 +72432,7 @@ }, "aiand/zai-org/glm-5.2": { "cache_read_input_token_cost": 3e-07, + "deprecation_date": "2026-10-27", "input_cost_per_token": 1e-06, "litellm_provider": "aiand", "max_input_tokens": 1048576, diff --git a/litellm/utils.py b/litellm/utils.py index 6f56cc65f45..2ca8a49f968 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -94,6 +94,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( ) from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, GetModelCostMap from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases +from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens @@ -8806,6 +8807,12 @@ def _get_valid_models_from_provider_api( return [] +def bundled_provider_models(provider: str) -> tuple[str, ...]: + bundled: Final[tuple[str, ...]] = tuple(litellm.models_by_provider.get(provider, ())) + cost_map: Final = _builtin_cost_map() + return tuple(model for model in bundled if not is_hidden_as_retired(cost_map.get(model, {}))) + + def get_valid_models( check_provider_endpoint: bool | None = None, custom_llm_provider: str | None = None, @@ -8871,8 +8878,7 @@ def get_valid_models( ) ) else: - models_for_provider = copy.deepcopy(litellm.models_by_provider.get(provider, [])) - valid_models.extend(models_for_provider) + valid_models.extend(bundled_provider_models(provider)) return valid_models except Exception as e: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e1131117fcf..ea9cacbdabe 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -72249,6 +72249,7 @@ }, "aiand/deepseek-ai/deepseek-v4-flash": { "cache_read_input_token_cost": 8e-08, + "deprecation_date": "2026-10-27", "input_cost_per_token": 1.5e-07, "litellm_provider": "aiand", "max_input_tokens": 1048576, @@ -72278,6 +72279,22 @@ "supports_vision": false }, "aiand/deepseek-ai/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 2e-08, + "deprecation_date": "2026-11-05", + "input_cost_per_token": 3e-07, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/deepseek-ai/deepseek-v4.1-flash-1017": { "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 3e-07, "litellm_provider": "aiand", @@ -72307,8 +72324,24 @@ "supports_tool_choice": true, "supports_vision": true }, + "aiand/meta/muse-spark-open": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 943718, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "aiand/moonshotai/kimi-k2.7-code": { "cache_read_input_token_cost": 2e-07, + "deprecation_date": "2026-10-27", "input_cost_per_token": 7.5e-07, "litellm_provider": "aiand", "max_input_tokens": 262144, @@ -72399,6 +72432,7 @@ }, "aiand/zai-org/glm-5.2": { "cache_read_input_token_cost": 3e-07, + "deprecation_date": "2026-10-27", "input_cost_per_token": 1e-06, "litellm_provider": "aiand", "max_input_tokens": 1048576, diff --git a/tests/unit/litellm_core_utils/test_llm_request_utils.py b/tests/unit/litellm_core_utils/test_llm_request_utils.py index e7350fae395..5243d308543 100644 --- a/tests/unit/litellm_core_utils/test_llm_request_utils.py +++ b/tests/unit/litellm_core_utils/test_llm_request_utils.py @@ -1,3 +1,5 @@ +import datetime +import time from typing import Final import httpx @@ -245,3 +247,39 @@ def test_rank_chat_models_flags_a_retired_model_the_provider_lists_without_its_c (ranked_model.model, ranked_model.is_past_deprecation_date) for ranked_model in rank_chat_models_from_llm_provider("vertex_ai") ] == [("fake-retired-vertex-chat", True), ("fake-live-vertex-chat", False)] + + +@pytest.mark.parametrize("host_zone", ["Etc/GMT+12", "Etc/GMT-14"]) +def test_rank_chat_models_flags_retirement_by_the_utc_date_whatever_the_host_zone( + monkeypatch: pytest.MonkeyPatch, host_zone: str +) -> None: + utc_today: Final = datetime.datetime.now(datetime.timezone.utc).date() + retirement_dates: Final = { + "deepseek/fake-retires-today-chat": utc_today, + "deepseek/fake-retires-tomorrow-chat": utc_today + datetime.timedelta(days=1), + } + for name, retirement_date in retirement_dates.items(): + monkeypatch.setitem( + litellm.model_cost, + name, + { + "input_cost_per_token": 1e-12, + "output_cost_per_token": 1e-12, + "mode": "chat", + "litellm_provider": PICKER_PROVIDER, + "deprecation_date": retirement_date.isoformat(), + }, + ) + monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, tuple(retirement_dates)) + monkeypatch.setenv("TZ", host_zone) + time.tzset() + try: + flags: Final = { + ranked_model.model: ranked_model.is_past_deprecation_date + for ranked_model in rank_chat_models_from_llm_provider(PICKER_PROVIDER) + } + finally: + monkeypatch.undo() + time.tzset() + + assert flags == {"deepseek/fake-retires-today-chat": True, "deepseek/fake-retires-tomorrow-chat": False} diff --git a/tests/unit/llms/baseten/chat/test_baseten_completions.py b/tests/unit/llms/baseten/chat/test_baseten_completions.py index 4f135a1a106..407b39319e9 100644 --- a/tests/unit/llms/baseten/chat/test_baseten_completions.py +++ b/tests/unit/llms/baseten/chat/test_baseten_completions.py @@ -7,8 +7,9 @@ import pytest import respx import litellm +from litellm.litellm_core_utils.model_retirement import is_past_deprecation_date from litellm.llms.baseten.chat import BasetenConfig -from litellm.utils import get_valid_models +from litellm.utils import _model_cache, get_valid_models class TestBasetenRouting: @@ -65,7 +66,7 @@ class TestBasetenModelDiscovery: servable_registry_models: Final = { key for key, value in litellm.model_cost.items() - if value.get("litellm_provider") == "baseten" and not litellm._is_past_deprecation_date(value) + if value.get("litellm_provider") == "baseten" and not is_past_deprecation_date(value) } assert servable_registry_models @@ -92,6 +93,20 @@ class TestBasetenModelDiscovery: litellm.baseten_models.discard(model) litellm.add_known_models({}) + def test_wildcard_drops_a_model_whose_retirement_date_passes_while_running(self, monkeypatch: pytest.MonkeyPatch): + model: Final = "baseten/acme/retires-while-running" + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "baseten", "deprecation_date": "2999-01-01"}) + try: + litellm.add_known_models({model: litellm.model_cost[model]}) + assert model in get_valid_models(custom_llm_provider="baseten") + + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "baseten", "deprecation_date": "2000-01-01"}) + + assert model not in get_valid_models(custom_llm_provider="baseten") + finally: + litellm.baseten_models.discard(model) + litellm.add_known_models({}) + @pytest.mark.parametrize( "value, expected", [ @@ -102,7 +117,7 @@ class TestBasetenModelDiscovery: ], ) def test_is_past_deprecation_date(self, value, expected): - assert litellm._is_past_deprecation_date(value) is expected + assert is_past_deprecation_date(value) is expected def test_get_models_calls_model_api(self, respx_mock: respx.MockRouter) -> None: route: Final = respx_mock.get("https://inference.baseten.co/v1/models").mock( @@ -156,6 +171,18 @@ class TestBasetenModelDiscovery: assert models == ["baseten/brand-new/model"] + @patch.dict(os.environ, {"BASETEN_API_KEY": "test-key"}) + def test_live_catalog_keeps_a_model_the_cost_map_dates_as_retired(self, respx_mock: respx.MockRouter) -> None: + assert is_past_deprecation_date(litellm.model_cost["baseten/zai-org/GLM-5"]) + _model_cache.flush_cache() + respx_mock.get("https://inference.baseten.co/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "zai-org/GLM-5"}]}) + ) + + models: Final = get_valid_models(check_provider_endpoint=True, custom_llm_provider="baseten") + + assert models == ["baseten/zai-org/GLM-5"] + class TestBasetenReasoningEffort: def test_reasoning_effort_passthrough_with_drop_params(self): diff --git a/tests/unit/llms/openai_like/test_aiand_provider.py b/tests/unit/llms/openai_like/test_aiand_provider.py index 804f3a3b6eb..0274b58f19a 100644 --- a/tests/unit/llms/openai_like/test_aiand_provider.py +++ b/tests/unit/llms/openai_like/test_aiand_provider.py @@ -1,4 +1,6 @@ +import datetime import json +import time from typing import Final import httpx @@ -8,6 +10,8 @@ import respx import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.health_check_helpers import wildcard_health_check_models +from litellm.litellm_core_utils.model_retirement import is_past_deprecation_date from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.proxy.auth.model_checks import get_known_models_from_wildcard from litellm.utils import ProviderConfigManager @@ -89,11 +93,99 @@ def _aiand_catalog() -> frozenset[str]: return frozenset(name for name, info in litellm.model_cost.items() if info.get("litellm_provider") == "aiand") -def test_aiand_wildcard_lists_its_priced_models(): - catalog: Final = _aiand_catalog() +def _utc_today() -> datetime.date: + return datetime.datetime.now(datetime.timezone.utc).date() - assert catalog - assert frozenset(get_known_models_from_wildcard("aiand/*")) == catalog + +def test_aiand_wildcard_lists_only_models_not_yet_retired(): + catalog: Final = _aiand_catalog() + retired: Final = frozenset( + name + for name in catalog + if (retired_on := litellm.model_cost[name].get("deprecation_date")) is not None + and datetime.date.fromisoformat(retired_on) <= _utc_today() + ) + + assert catalog - retired + assert frozenset(get_known_models_from_wildcard("aiand/*")) == catalog - retired + + +def test_aiand_wildcard_drops_a_model_whose_retirement_date_passes_while_running(monkeypatch: pytest.MonkeyPatch): + model: Final = "aiand/acme/retires-while-running" + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "aiand", "deprecation_date": "2999-01-01"}) + try: + litellm.add_known_models({model: litellm.model_cost[model]}) + assert model in get_known_models_from_wildcard("aiand/*") + + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "aiand", "deprecation_date": "2000-01-01"}) + + assert model not in get_known_models_from_wildcard("aiand/*") + finally: + litellm.aiand_models.discard(model) + litellm.add_known_models({}) + + +def test_aiand_models_by_provider_leaves_out_models_retired_at_load(): + retired: Final = frozenset( + name + for name, row in litellm.model_cost.items() + if row.get("litellm_provider") == "aiand" and is_past_deprecation_date(row) + ) + + assert retired + assert retired.isdisjoint(litellm.models_by_provider["aiand"]) + + +def test_aiand_wildcard_health_check_skips_a_model_whose_retirement_date_passes_while_running( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "aiand/acme/free-and-retiring" + live_row: Final = { + "litellm_provider": "aiand", + "mode": "chat", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "deprecation_date": "2999-01-01", + } + monkeypatch.setitem(litellm.model_cost, model, live_row) + try: + litellm.add_known_models({model: live_row}) + assert model in wildcard_health_check_models(wildcard="aiand/*", custom_llm_provider="aiand") + + monkeypatch.setitem(litellm.model_cost, model, {**live_row, "deprecation_date": "2000-01-01"}) + + assert model not in wildcard_health_check_models(wildcard="aiand/*", custom_llm_provider="aiand") + finally: + litellm.aiand_models.discard(model) + litellm.add_known_models({}) + + +@pytest.mark.parametrize("host_zone", ["Etc/GMT+12", "Etc/GMT-14"]) +def test_aiand_models_retire_on_the_utc_date_whatever_the_host_zone(monkeypatch: pytest.MonkeyPatch, host_zone: str): + today: Final = _utc_today() + rows: Final = { + "aiand/acme/retires-today": {"litellm_provider": "aiand", "deprecation_date": today.isoformat()}, + "aiand/acme/retires-tomorrow": { + "litellm_provider": "aiand", + "deprecation_date": (today + datetime.timedelta(days=1)).isoformat(), + }, + } + for name, row in rows.items(): + monkeypatch.setitem(litellm.model_cost, name, row) + monkeypatch.setenv("TZ", host_zone) + time.tzset() + try: + litellm.add_known_models(rows) + listed: Final = frozenset(get_known_models_from_wildcard("aiand/*")) + + assert "aiand/acme/retires-today" not in listed + assert "aiand/acme/retires-tomorrow" in listed + finally: + monkeypatch.undo() + time.tzset() + for name in rows: + litellm.aiand_models.discard(name) + litellm.add_known_models({}) def test_aiand_calls_are_billed_from_their_cost_rows(): diff --git a/ui/litellm-dashboard/public/assets/logos/aiand.svg b/ui/litellm-dashboard/public/assets/logos/aiand.svg new file mode 100644 index 00000000000..486c73b8dbc --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/aiand.svg @@ -0,0 +1,4 @@ + + + + diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 810b010f25e..02c31ecc082 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -32,6 +32,18 @@ describe("provider_info_helpers", () => { expect(getPlaceholder(Providers.BIGMODEL)).toMatch(/^bigmodel\//); }); + it("resolves ai& metadata, name, logo and model placeholder", () => { + const metadata = PROVIDER_CREATE_FIELDS.find((provider) => provider.provider === "AIAND"); + expect(provider_map.AIAND).toBe(metadata?.litellm_provider); + expect(getProviderLogoAndName("aiand")).toEqual({ + displayName: "ai&", + logo: providerLogoMap[Providers.AIAND], + }); + expect(getProviderLogoAndName("AIAND")).toEqual(getProviderLogoAndName("aiand")); + expect(providerLogoMap[Providers.AIAND]).toContain("aiand"); + expect(getPlaceholder(Providers.AIAND)).toBe("aiand/openai/gpt-oss-120b"); + }); + it("maps Microsoft 365 Copilot to its chat model provider and placeholder", () => { expect(provider_map.MICROSOFT_365_COPILOT).toBe("microsoft_365_copilot"); expect(getPlaceholder(Providers.MICROSOFT_365_COPILOT)).toBe("microsoft_365_copilot/chat"); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 1ed32c160e6..6448bd936cc 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -1,6 +1,7 @@ import { resolveLogoSrc } from "@/lib/assetPaths"; import a2aAgentLogo from "../../public/assets/logos/a2a_agent.png"; import ai21Logo from "../../public/assets/logos/ai21.svg"; +import aiandLogo from "../../public/assets/logos/aiand.svg"; import aimlApiLogo from "../../public/assets/logos/aiml_api.svg"; import anthropicLogo from "../../public/assets/logos/anthropic.svg"; import assemblyaiSmallLogo from "../../public/assets/logos/assemblyai_small.png"; @@ -76,6 +77,7 @@ export enum Providers { A2A_Agent = "A2A Agent", AI21 = "Ai21", AI21_CHAT = "Ai21 Chat", + AIAND = "ai&", AIML = "AI/ML API", AIOHTTP_OPENAI = "Aiohttp Openai", Anthropic = "Anthropic", @@ -207,6 +209,7 @@ export const provider_map: Record = { A2A_Agent: "a2a_agent", AI21: "ai21", AI21_CHAT: "ai21_chat", + AIAND: "aiand", AIML: "aiml", AIOHTTP_OPENAI: "aiohttp_openai", Anthropic: "anthropic", @@ -344,6 +347,7 @@ export const providerLogoMap: Partial> = { [Providers.A2A_Agent]: a2aAgentLogo.src, [Providers.AI21]: ai21Logo.src, [Providers.AI21_CHAT]: ai21Logo.src, + [Providers.AIAND]: aiandLogo.src, [Providers.AIML]: aimlApiLogo.src, [Providers.AIOHTTP_OPENAI]: openaiSmallLogo.src, [Providers.Anthropic]: anthropicLogo.src, @@ -472,6 +476,7 @@ export const getProviderLogoAndName = (providerValue: string): { logo: string; d const providerPlaceholderMap: Partial> = { [Providers.AIML]: "aiml/flux-pro/v1.1", + [Providers.AIAND]: "aiand/openai/gpt-oss-120b", [Providers.Anthropic]: "claude-3-opus", [Providers.Azure]: "my-deployment", [Providers.Azure_AI_Studio]: "azure_ai/command-r-plus", diff --git a/ui/litellm-dashboard/src/lib/logoTreatments.ts b/ui/litellm-dashboard/src/lib/logoTreatments.ts index fc19c5c364d..4cd6f65a613 100644 --- a/ui/litellm-dashboard/src/lib/logoTreatments.ts +++ b/ui/litellm-dashboard/src/lib/logoTreatments.ts @@ -32,6 +32,7 @@ const TREATMENT_BY_ASSET: Readonly> = { "v0.svg": "invert", "vercel.svg": "invert", "watsonx.svg": "invert", + "aiand.svg": "plate", "aiml_api.svg": "plate", "akto.svg": "plate", "aws.svg": "plate", From ccdc8c7d10f8b042f576b1adfb670a02011cddcd Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Fri, 30 Oct 2026 10:59:28 -0700 Subject: [PATCH 40/64] docs(daily-activity): describe guardrail spend fields and reporting key IDs in OpenAPI (#49236) * docs(daily-activity): describe guardrail spend and reporting key IDs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(daily-activity): describe reporting key IDs on /user/daily/activity Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 17 +++++++-- .../daily_activity_scopes.py | 6 ++- .../internal_user_endpoints.py | 3 +- .../common_daily_activity.py | 27 +++++++++----- ui/litellm-dashboard/src/lib/http/schema.d.ts | 37 +++++++++++++------ 5 files changed, 64 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index c76b9be7c33..dc8b9415020 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3491,6 +3491,7 @@ "additionalProperties": { "$ref": "#/components/schemas/GuardrailMetric" }, + "description": "Spend by guardrail ID, already included in key and model spend. Populated on user, team and organization aggregated activity only, without a model filter", "title": "Guardrails", "type": "object" }, @@ -3812,10 +3813,12 @@ "GuardrailKeyMetric": { "properties": { "metadata": { - "$ref": "#/components/schemas/KeyMetadata" + "$ref": "#/components/schemas/KeyMetadata", + "description": "API key alias and existence information" }, "metrics": { - "$ref": "#/components/schemas/GuardrailSpendMetrics" + "$ref": "#/components/schemas/GuardrailSpendMetrics", + "description": "Guardrail spend and unpriced units for this API key" } }, "required": [ @@ -3828,6 +3831,7 @@ "GuardrailMetadata": { "properties": { "guardrail_name": { + "description": "Guardrail display name, falling back to the guardrail ID", "title": "Guardrail Name", "type": "string" } @@ -3844,14 +3848,17 @@ "additionalProperties": { "$ref": "#/components/schemas/GuardrailKeyMetric" }, + "description": "Per-key guardrail spend and unpriced units, limited to keys in the day's API key breakdown", "title": "Api Key Breakdown", "type": "object" }, "metadata": { - "$ref": "#/components/schemas/GuardrailMetadata" + "$ref": "#/components/schemas/GuardrailMetadata", + "description": "Guardrail name information" }, "metrics": { - "$ref": "#/components/schemas/GuardrailSpendMetrics" + "$ref": "#/components/schemas/GuardrailSpendMetrics", + "description": "Guardrail spend and unpriced units across all matching API keys" } }, "required": [ @@ -3865,11 +3872,13 @@ "properties": { "spend": { "default": 0.0, + "description": "Guardrail spend in USD, already included in key and model spend", "title": "Spend", "type": "number" }, "untracked_units": { "default": 0, + "description": "Guardrail units with no known price, excluded from spend", "title": "Untracked Units", "type": "integer" } diff --git a/litellm/proxy/management_endpoints/daily_activity_scopes.py b/litellm/proxy/management_endpoints/daily_activity_scopes.py index ca1a96f9d84..7ff73ac6e2e 100644 --- a/litellm/proxy/management_endpoints/daily_activity_scopes.py +++ b/litellm/proxy/management_endpoints/daily_activity_scopes.py @@ -216,7 +216,11 @@ def _user_query( start_date: str | None = Query(default=None, description="Start date in YYYY-MM-DD format"), end_date: str | None = Query(default=None, description="End date in YYYY-MM-DD format"), model: str | None = Query(default=None, description="Filter by specific model"), - api_key: str | None = Query(default=None, description="Filter by specific API key"), + api_key: str | None = Query( + default=None, + description="Filter by reporting API key ID. Missing keys are reported as unattributed. " + "For a literal key unattributed or a key starting with raw:, add an extra raw: prefix", + ), user_id: str | None = Query( default=None, description="Filter by specific user ID. Admins can filter by any user or omit for global view. " diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 6950cf08a25..54a2d4062d2 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -2922,7 +2922,8 @@ async def get_user_daily_activity( ), api_key: str | None = fastapi.Query( default=None, - description="Filter by specific API key", + description="Filter by reporting API key ID. Missing keys are reported as unattributed. " + "For a literal key unattributed or a key starting with raw:, add an extra raw: prefix", ), user_id: str | None = fastapi.Query( default=None, diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 3f1a7980b27..37f59711050 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -68,29 +68,34 @@ class MetricWithMetadata(MetricBase): class GuardrailSpendMetrics(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - spend: float = 0.0 - untracked_units: int = 0 + spend: float = Field(default=0.0, description="Guardrail spend in USD, already included in key and model spend") + untracked_units: int = Field(default=0, description="Guardrail units with no known price, excluded from spend") class GuardrailKeyMetric(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - metrics: GuardrailSpendMetrics - metadata: KeyMetadata + metrics: GuardrailSpendMetrics = Field(description="Guardrail spend and unpriced units for this API key") + metadata: KeyMetadata = Field(description="API key alias and existence information") class GuardrailMetadata(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - guardrail_name: str + guardrail_name: str = Field(description="Guardrail display name, falling back to the guardrail ID") class GuardrailMetric(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - metrics: GuardrailSpendMetrics - metadata: GuardrailMetadata - api_key_breakdown: Mapping[str, GuardrailKeyMetric] = Field(default_factory=dict) + metrics: GuardrailSpendMetrics = Field( + description="Guardrail spend and unpriced units across all matching API keys" + ) + metadata: GuardrailMetadata = Field(description="Guardrail name information") + api_key_breakdown: Mapping[str, GuardrailKeyMetric] = Field( + default_factory=dict, + description="Per-key guardrail spend and unpriced units, limited to keys in the day's API key breakdown", + ) class BreakdownMetrics(LiteLLMBaseModel): @@ -103,7 +108,11 @@ class BreakdownMetrics(LiteLLMBaseModel): endpoints: dict[str, MetricWithMetadata] = Field(default_factory=dict) # endpoint -> {metrics, metadata} api_keys: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata} entities: dict[str, MetricWithMetadata] = Field(default_factory=dict) # entity -> {metrics, metadata} - guardrails: Mapping[str, GuardrailMetric] = Field(default_factory=dict) # guardrail_id -> spend + guardrails: Mapping[str, GuardrailMetric] = Field( + default_factory=dict, + description="Spend by guardrail ID, already included in key and model spend. " + "Populated on user, team and organization aggregated activity only, without a model filter", + ) class DailySpendData(LiteLLMBaseModel): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 960616e2631..aafa732276b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27870,7 +27870,10 @@ export interface components { entities?: { [key: string]: components["schemas"]["MetricWithMetadata"]; }; - /** Guardrails */ + /** + * Guardrails + * @description Spend by guardrail ID, already included in key and model spend. Populated on user, team and organization aggregated activity only, without a model filter + */ guardrails?: { [key: string]: components["schemas"]["GuardrailMetric"]; }; @@ -33448,32 +33451,44 @@ export interface components { }; /** GuardrailKeyMetric */ GuardrailKeyMetric: { + /** @description API key alias and existence information */ metadata: components["schemas"]["KeyMetadata"]; + /** @description Guardrail spend and unpriced units for this API key */ metrics: components["schemas"]["GuardrailSpendMetrics"]; }; /** GuardrailMetadata */ GuardrailMetadata: { - /** Guardrail Name */ + /** + * Guardrail Name + * @description Guardrail display name, falling back to the guardrail ID + */ guardrail_name: string; }; /** GuardrailMetric */ GuardrailMetric: { - /** Api Key Breakdown */ + /** + * Api Key Breakdown + * @description Per-key guardrail spend and unpriced units, limited to keys in the day's API key breakdown + */ api_key_breakdown?: { [key: string]: components["schemas"]["GuardrailKeyMetric"]; }; + /** @description Guardrail name information */ metadata: components["schemas"]["GuardrailMetadata"]; + /** @description Guardrail spend and unpriced units across all matching API keys */ metrics: components["schemas"]["GuardrailSpendMetrics"]; }; /** GuardrailSpendMetrics */ GuardrailSpendMetrics: { /** * Spend + * @description Guardrail spend in USD, already included in key and model spend * @default 0 */ spend: number; /** * Untracked Units + * @description Guardrail units with no known price, excluded from spend * @default 0 */ untracked_units: number; @@ -76633,7 +76648,7 @@ export interface operations { end_date?: string | null; /** @description Filter by specific model */ model?: string | null; - /** @description Filter by specific API key */ + /** @description Filter by reporting API key ID. Missing keys are reported as unattributed. For a literal key unattributed or a key starting with raw:, add an extra raw: prefix */ api_key?: string | null; /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ user_id?: string | null; @@ -76682,7 +76697,7 @@ export interface operations { end_date?: string | null; /** @description Filter by specific model */ model?: string | null; - /** @description Filter by specific API key */ + /** @description Filter by reporting API key ID. Missing keys are reported as unattributed. For a literal key unattributed or a key starting with raw:, add an extra raw: prefix */ api_key?: string | null; /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ user_id?: string | null; @@ -76727,7 +76742,7 @@ export interface operations { end_date?: string | null; /** @description Filter by specific model */ model?: string | null; - /** @description Filter by specific API key */ + /** @description Filter by reporting API key ID. Missing keys are reported as unattributed. For a literal key unattributed or a key starting with raw:, add an extra raw: prefix */ api_key?: string | null; /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ user_id?: string | null; @@ -76773,7 +76788,7 @@ export interface operations { end_date?: string | null; /** @description Filter by specific model */ model?: string | null; - /** @description Filter by specific API key */ + /** @description Filter by reporting API key ID. Missing keys are reported as unattributed. For a literal key unattributed or a key starting with raw:, add an extra raw: prefix */ api_key?: string | null; /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ user_id?: string | null; @@ -76820,7 +76835,7 @@ export interface operations { end_date?: string | null; /** @description Filter by specific model */ model?: string | null; - /** @description Filter by specific API key */ + /** @description Filter by reporting API key ID. Missing keys are reported as unattributed. For a literal key unattributed or a key starting with raw:, add an extra raw: prefix */ api_key?: string | null; /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ user_id?: string | null; @@ -76866,7 +76881,7 @@ export interface operations { end_date?: string | null; /** @description Filter by specific model */ model?: string | null; - /** @description Filter by specific API key */ + /** @description Filter by reporting API key ID. Missing keys are reported as unattributed. For a literal key unattributed or a key starting with raw:, add an extra raw: prefix */ api_key?: string | null; /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ user_id?: string | null; @@ -76912,7 +76927,7 @@ export interface operations { end_date?: string | null; /** @description Filter by specific model */ model?: string | null; - /** @description Filter by specific API key */ + /** @description Filter by reporting API key ID. Missing keys are reported as unattributed. For a literal key unattributed or a key starting with raw:, add an extra raw: prefix */ api_key?: string | null; /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ user_id?: string | null; @@ -76958,7 +76973,7 @@ export interface operations { end_date?: string | null; /** @description Filter by specific model */ model?: string | null; - /** @description Filter by specific API key */ + /** @description Filter by reporting API key ID. Missing keys are reported as unattributed. For a literal key unattributed or a key starting with raw:, add an extra raw: prefix */ api_key?: string | null; /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ user_id?: string | null; From f7f87e6193e151db09ce3db40f26dcff5333a59b Mon Sep 17 00:00:00 2001 From: Ahmet Selim Fedakar <122402435+selimfedakar@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:01:00 +0300 Subject: [PATCH 41/64] fix(responses): normalize empty function_call_output to an empty tool result (#48845) * fix(responses): normalize empty function_call_output to an empty tool result * test(responses): preserve empty tool outputs through guardrails --------- Co-authored-by: Tin --- .../transformation.py | 10 ++++--- ...test_openai_responses_guardrail_handler.py | 14 +++++++--- .../test_anthropic_responses_bridge.py | 26 +++++++++++++++++ ...test_function_call_output_normalization.py | 17 ++++++++++- .../test_google_ai_studio_responses_wire.py | 28 +++++++++++++++++++ 5 files changed, 86 insertions(+), 9 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index ce8841c9190..8f0a85d67de 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -398,23 +398,25 @@ def tool_output_to_chat_content( if single_media is not None: return [single_media] if isinstance(output, list): + parts: Final = cast(Sequence[object], output) # cast-ok: checked list shape, elements remain unvalidated normalized_blocks: Final[list[ToolMessageContentPart]] = [] text_acc: Final[list[str]] = [] - for part in output: + for part in parts: if not isinstance(part, dict): continue part_type = part.get("type") if part_type in ("input_text", "output_text", "text"): txt = part.get("text") - if isinstance(txt, str) and txt: + if isinstance(txt, str): text_acc.append(txt) - normalized_blocks.append({"type": "text", "text": txt}) + if txt: + normalized_blocks.append({"type": "text", "text": txt}) elif (media := decode_responses_media(part)) is not None: normalized_blocks.append(media) if any(b.get("type") != "text" for b in normalized_blocks): return normalized_blocks - if text_acc: + if any(text_acc) or len(text_acc) == len(parts): return "".join(text_acc) try: return json.dumps(output) diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index 673e880a1c0..d732d922375 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -724,16 +724,23 @@ class TestToolResultItemsReachTextsGuardrails: assert blocked.value.status_code == 400 @pytest.mark.asyncio - @pytest.mark.parametrize("empty_output", ["", None], ids=["empty_string", "null"]) - async def test_empty_tool_output_sends_no_text_and_stays_as_sent(self, empty_output: str | None) -> None: + @pytest.mark.parametrize(("call_item", "output_type"), _TOOL_TURNS, ids=["function", "custom"]) + @pytest.mark.parametrize( + "empty_output", + ["", None, [], *[[{"type": kind, "text": ""}] for kind in ("input_text", "output_text", "text")]], + ids=["empty_string", "null", "empty_list", "input_text", "output_text", "text"], + ) + async def test_empty_tool_output_sends_no_text_and_stays_as_sent( + self, call_item: dict[str, str], output_type: str, empty_output: object + ) -> None: guardrail: Final = RecordingMaskingGuardrail(guardrail_name="test") - call_item, output_type = _TOOL_TURNS[0] result: Final = await OpenAIResponsesHandler().process_input_messages( _tool_result_request(call_item, output_type, empty_output), guardrail ) assert guardrail.seen_texts == [["Weather?"]] + assert guardrail.seen_message_contents == [["Weather?", None]] assert result["input"][2] == {"type": output_type, "call_id": "call_1", "output": empty_output} @pytest.mark.asyncio @@ -1103,7 +1110,6 @@ class TestToolResultItemsReachTextsGuardrails: ["visible"], id="non_dict_part", ), - pytest.param([], ["[]"], id="empty"), pytest.param({"result": "SECRET forecast"}, ['{"result": "SECRET forecast"}'], id="dict"), pytest.param( [{"type": "input_image", "file_id": "file-map"}], diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index ef198b487b9..0341511496c 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -23,6 +23,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( ) from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import Message, ModelResponse +from tests.unit.proxy.conftest import httpx_transport def test_response_api_handler_merges_metadata_and_service_tier_without_error(): @@ -593,6 +594,31 @@ def _anthropic_message(*content: Mapping[str, object]) -> httpx.Response: ) +@pytest.mark.usefixtures(httpx_transport.__name__) +@pytest.mark.parametrize("output", ["", [], [{"type": "input_text", "text": ""}]]) +def test_empty_responses_tool_output_reaches_anthropic_without_json_text( + respx_mock: respx.MockRouter, output: str | list[dict[str, str]] +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "ok"}) + ) + response: Final = litellm.responses( + model=_MODEL, + api_key="mock_api_key", + input=[ + {"role": "user", "content": "Check the weather"}, + {"type": "function_call", "call_id": "call_empty", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_empty", "output": output}, + ], + ) + sent: Final = json.loads(route.calls.last.request.content) + assert response.output_text == "ok" + assert sent["messages"][-1] == { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "call_empty", "content": ""}], + } + + def _replay(output: Sequence[BaseModel]) -> list[object]: return [*_FIRST_INPUT, *(item.model_dump(exclude_none=True) for item in output), *_NEXT_INPUT] diff --git a/tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py b/tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py index 65800e42124..fce06598d29 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py +++ b/tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py @@ -50,9 +50,10 @@ def test_image_detail_mime_and_adjacent_text_survive_normalization(kind: str, de else "data:application/octet-stream;base64,AAAA", **({} if nested_url else metadata), } - parts: Final = [{"type": "input_text", "text": "Screenshot follows"}, image] + parts: Final = [{"type": "input_text", "text": ""}, {"type": "input_text", "text": "Screenshot follows"}, image] original: Final = deepcopy(parts) assert _normalize_media_content(kind, parts) == [ + *([{"type": "text", "text": ""}] if kind == "direct" else []), {"type": "text", "text": "Screenshot follows"}, {"type": "image_url", "image_url": {"url": "data:application/octet-stream;base64,AAAA", **metadata}}, ] @@ -97,6 +98,12 @@ def test_media_normalization_preserves_formats_references_and_metadata(kind: str ("output", "expected"), [ (None, ""), + ("", ""), + ([], ""), + ([{"type": "input_text", "text": ""}], ""), + ([{"type": kind, "text": ""} for kind in ("input_text", "output_text", "text")], ""), + ("[]", "[]"), + ([{"type": "input_text", "text": " \n"}], " \n"), ('{"ok":true}', '{"ok":true}'), ([{"type": "input_text", "text": "hello"}, {"type": "input_text", "text": " world"}], "hello world"), ({"type": "file", "status": "ready"}, '{"type": "file", "status": "ready"}'), @@ -106,6 +113,14 @@ def test_media_normalization_preserves_formats_references_and_metadata(kind: str ), ({"type": "input_image", "format": 17}, '{"type": "input_image", "format": 17}'), ([{"type": "unknown", "value": 17}], '[{"type": "unknown", "value": 17}]'), + ([{"type": "input_text", "text": 17}], '[{"type": "input_text", "text": 17}]'), + ([{"type": "input_text"}], '[{"type": "input_text"}]'), + ([{"type": "input_text", "text": None}], '[{"type": "input_text", "text": null}]'), + ( + [{"type": "input_text", "text": ""}, {"type": "unknown", "value": 17}], + '[{"type": "input_text", "text": ""}, {"type": "unknown", "value": 17}]', + ), + ([{"type": "input_text", "text": ""}, 17], '[{"type": "input_text", "text": ""}, 17]'), ( [{"type": "input_text", "text": "keep"}, {"type": "input_image", "file_id": "file-only"}], "keep", diff --git a/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_wire.py b/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_wire.py index f63088a606f..c7acc70ba95 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_wire.py +++ b/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_wire.py @@ -84,6 +84,34 @@ async def test_web_search_preview_is_sent_as_google_search() -> None: assert request_body["tools"] == [{"googleSearch": {}}] +@pytest.mark.asyncio +@pytest.mark.parametrize("output", ["", [], [{"type": "input_text", "text": ""}]]) +async def test_empty_responses_tool_output_reaches_gemini_without_json_text( + output: str | list[dict[str, str]], +) -> None: + with respx.mock() as mock_router: + route: Final = mock_router.post(url__regex=_GEMINI_URL).mock( + return_value=httpx.Response(status_code=200, json=_function_call_response("signature")) + ) + response: Final = await litellm.aresponses( + model="gemini/gemini-2.5-flash", + api_key="test-key", + input=[ + {"role": "user", "content": "Check the weather"}, + {"type": "function_call", "call_id": "call_empty", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_empty", "output": output}, + ], + ) + request_body: Final = _JSON_OBJECT.validate_json(route.calls.last.request.content) + + assert isinstance(response, ResponsesAPIResponse) + assert request_body["contents"] == [ + {"role": "user", "parts": [{"text": "Check the weather"}]}, + {"role": "model", "parts": [{"function_call": {"name": "get_weather", "args": {}}}]}, + {"role": "user", "parts": [{"function_response": {"name": "get_weather", "response": {"content": ""}}}]}, + ] + + @pytest.mark.asyncio @pytest.mark.parametrize( ("assistant_content", "assistant_parts"), From 840c83b4034a266cdd9e0858b9a568e6c348cf22 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 30 Oct 2026 11:19:08 -0700 Subject: [PATCH 42/64] test: bedrock claude prompt caching on /v1/messages reports and charges cache tokens (#49265) * test: bedrock claude prompt caching on /v1/messages reports and charges cache tokens * ci: re-run after ci-coverage and assert-shard-coverage were cancelled by the GitHub Actions incident --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...st_bedrock_messages_prompt_caching_wire.py | 191 ++++++++++++++++++ 1 file changed, 191 insertions(+) create mode 100644 tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py diff --git a/tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py b/tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py new file mode 100644 index 00000000000..f2622da1bf6 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py @@ -0,0 +1,191 @@ +import base64 +import json +import uuid +from typing import Final + +import anthropic +import httpx +import pytest +from integration._support import claude_code as cc +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_MODEL_ID: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" +_TOKEN: Final = "synthetic-bedrock-bearer" +_PROMPT: Final = "Reply with exactly the word PONG" +_INPUT_TOKENS: Final = 10 +_OUTPUT_TOKENS: Final = 4 +_CACHE_READ_TOKENS: Final = 900 +_CACHE_WRITE_TOKENS: Final = 300 +_INPUT_RATE: Final = 0.001 +_OUTPUT_RATE: Final = 0.002 +_CACHE_READ_RATE: Final = 0.0001 +_CACHE_WRITE_RATE: Final = 0.0015 +_EXPECTED_SPEND: Final = ( + _INPUT_TOKENS * _INPUT_RATE + + _CACHE_READ_TOKENS * _CACHE_READ_RATE + + _CACHE_WRITE_TOKENS * _CACHE_WRITE_RATE + + _OUTPUT_TOKENS * _OUTPUT_RATE +) +_PROMPT_TOKENS: Final = _INPUT_TOKENS + _CACHE_READ_TOKENS + _CACHE_WRITE_TOKENS +_USAGE: Final = { + "input_tokens": _INPUT_TOKENS, + "output_tokens": _OUTPUT_TOKENS, + "cache_read_input_tokens": _CACHE_READ_TOKENS, + "cache_creation_input_tokens": _CACHE_WRITE_TOKENS, +} +_USER_CONTENT: Final = [{"type": "text", "text": _PROMPT, "cache_control": cc.CACHE}] +_CACHED_SYSTEM_BLOCKS: Final = [block for block in cc.system_blocks() if "cache_control" in block] + + +def _assert_claude_code_breakpoints_reach_bedrock(request: Request, target: str) -> dict[str, JsonValue]: + assert request.method == "POST" and request.target == target, request.target + assert request.headers["authorization"] == f"Bearer {_TOKEN}" + body: Final = cc.JSON_OBJECT.validate_json(request.body) + assert body["anthropic_version"] == "bedrock-2023-05-31", body + assert body["system"] == _CACHED_SYSTEM_BLOCKS, f"billing-header block dropped, breakpoints kept: {body}" + assert body["messages"] == [{"role": "user", "content": _USER_CONTENT}], body + return body + + +def _invoke_chunk(event: dict[str, JsonValue]) -> bytes: + return _aws_event_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}, "sc", "u") + + +def _stream_frames(message_id: str) -> tuple[bytes, ...]: + return ( + _invoke_chunk( + { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": "claude-sonnet-4-5-20250929", + "stop_reason": None, + "stop_sequence": None, + "usage": {**_USAGE, "output_tokens": 1}, + }, + } + ), + _invoke_chunk({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "PONG"}}), + _invoke_chunk({"type": "content_block_stop", "index": 0}), + _invoke_chunk( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": _OUTPUT_TOKENS}, + } + ), + _invoke_chunk({"type": "message_stop"}), + ) + + +def _spend_row(message_id: str) -> dict[str, object]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (message_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0] + + +def _client(gateway: Gateway, transport: httpx.Client) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=transport + ) + + +def test_bedrock_claude_cache_control_reaches_invoke_and_cache_tokens_are_reported_and_charged( + gateway: Gateway, +) -> None: + message_id: Final = f"msg_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + _assert_claude_code_breakpoints_reach_bedrock(request, f"/model/{_MODEL_ID}/invoke") + return Reply( + body=cc.message_reply( + message_id, "claude-sonnet-4-5-20250929", ({"type": "text", "text": "PONG"},), dict(_USAGE) + ) + ) + + with ( + wire_server(respond) as wire, + gateway.scenario() as scenario, + httpx.Client(timeout=15, trust_env=False) as transport, + ): + model: Final = scenario.model( + model=f"bedrock/invoke/{_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + api_base=wire.url, + input_cost_per_token=_INPUT_RATE, + output_cost_per_token=_OUTPUT_RATE, + cache_read_input_token_cost=_CACHE_READ_RATE, + cache_creation_input_token_cost=_CACHE_WRITE_RATE, + ) + final: Final = _client(gateway, transport).messages.create( + model=model, + max_tokens=32, + system=list(cc.system_blocks()), + messages=[{"role": "user", "content": _USER_CONTENT}], + ) + assert len(wire.drain()) == 1 + assert final.id == message_id, final + assert [(block.type, block.text) for block in final.content] == [("text", "PONG")], final + assert final.usage.model_dump(exclude_none=True, include=set(_USAGE)) == _USAGE, final.usage + row: Final = _spend_row(message_id) + + assert float(str(row["spend"])) == pytest.approx(_EXPECTED_SPEND), row + assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, _OUTPUT_TOKENS), row + + +def test_bedrock_claude_stream_cache_tokens_from_message_start_are_reported_and_charged(gateway: Gateway) -> None: + message_id: Final = f"msg_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _assert_claude_code_breakpoints_reach_bedrock( + request, f"/model/{_MODEL_ID}/invoke-with-response-stream" + ) + assert "stream" not in body, body + return Reply(chunks=_stream_frames(message_id), content_type="application/vnd.amazon.eventstream") + + with ( + wire_server(respond) as wire, + gateway.scenario() as scenario, + httpx.Client(timeout=15, trust_env=False) as transport, + ): + model: Final = scenario.model( + model=f"bedrock/invoke/{_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + api_base=wire.url, + input_cost_per_token=_INPUT_RATE, + output_cost_per_token=_OUTPUT_RATE, + cache_read_input_token_cost=_CACHE_READ_RATE, + cache_creation_input_token_cost=_CACHE_WRITE_RATE, + ) + with _client(gateway, transport).messages.stream( + model=model, + max_tokens=32, + system=list(cc.system_blocks()), + messages=[{"role": "user", "content": _USER_CONTENT}], + ) as stream: + deltas: Final = tuple(stream.text_stream) + final: Final = stream.get_final_message() + assert len(wire.drain()) == 1 + assert deltas == ("PONG",) + assert final.id == message_id, final + assert final.usage.model_dump(exclude_none=True, include=set(_USAGE)) == _USAGE, final.usage + row: Final = _spend_row(message_id) + + assert float(str(row["spend"])) == pytest.approx(_EXPECTED_SPEND), row + assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, _OUTPUT_TOKENS), row From d23af8bdd08943fe1a0e1d9e528fcb12000b287a Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 30 Oct 2026 11:24:49 -0700 Subject: [PATCH 43/64] fix(cost): bill the margin once on rebuilt streams and streamed Responses through the chat bridge (#49254) * fix(cost): bill the margin once on rebuilt streams and streamed Responses through the chat bridge The last usage chunk and the Responses terminal event carry litellm's bill in usage.cost, and two readers took that bill for the upstream's own charge and margined it again: completion_cost on a stream_chunk_builder result, and stamp_responses_usage_cost on the chat bridge. The stamp now keeps the upstream's usage.cost under the hidden upstream_usage_cost key, the chunk builder carries it onto the rebuilt response, and the Responses stamp prices that saved figure * fix(cost): hand the Responses stamp the upstream charge through upstream_usage_cost only, never the response-cost header response_cost_calculator returns the response-cost header as the final bill without the margin, so a priced copy carrying the upstream's raw usage.cost there billed the client the unmargined figure. The hidden upstream_usage_cost key already feeds billed_usage_cost, which margins it once. The Responses test now prices through a real Logging object instead of a mocked calculator * fix(cost): leave stamp_responses_usage_cost as on main The Responses chat bridge copies the inner stream's final usage chunk, so the wrapper stamp keeping upstream_usage_cost already prices it once. Changing the Responses stamp was not needed for that and re-priced xai MCP rounds whose usage.cost is already litellm's bill (0.00512 for 0.00312 in test_xai_stream_auto_execute_bills_each_round_at_its_own_price_with_one_fixed_margin) --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 18 ++++++++++- .../litellm_core_utils/streaming_handler.py | 18 ++++++++--- .../test_streaming_chunk_builder_utils.py | 24 ++++++++++++++ .../test_streaming_handler.py | 31 +++++++++++++++++++ 4 files changed, 86 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 574f1103c1a..0ddfa1210c2 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -14,12 +14,18 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import ( PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS, + UPSTREAM_USAGE_COST, chunk_usage, finite_cost, provider_reported_usage_cost, stamp_litellm_bill_on_usage, ) -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.hidden_params import ( + HIDDEN_PARAMS_ATTR, + get_hidden_params, + set_hidden_param, + set_hidden_params, +) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import partition_annotations from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -1444,6 +1450,14 @@ def _set_stream_builder_response_cost( _set_hidden_response_cost(response, response_cost) +def _carry_upstream_usage_cost(response: ModelResponse, chunks: Sequence[object]) -> None: + chunk_hidden_params: Final = (get_hidden_params(chunk) for chunk in reversed(chunks)) + stamped: Final = (hidden for hidden in chunk_hidden_params if hidden is not None and UPSTREAM_USAGE_COST in hidden) + for hidden in stamped: + set_hidden_param(response, UPSTREAM_USAGE_COST, hidden[UPSTREAM_USAGE_COST]) + return + + def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None: if logging_obj is None: return @@ -1794,6 +1808,7 @@ def stream_chunk_builder( response, logging_obj, ) + _carry_upstream_usage_cost(response, chunks) # pyright: ignore[reportUnknownArgumentType] # chunks is an untyped list param _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost( response, @@ -1836,6 +1851,7 @@ def stream_chunk_builder( ) break + _carry_upstream_usage_cost(response, chunks) # pyright: ignore[reportUnknownArgumentType] # chunks is an untyped list param _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost( response, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b1b3e20f218..9710956c97e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -24,7 +24,12 @@ import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost +from litellm.litellm_core_utils.core_helpers import ( + UPSTREAM_USAGE_COST, + chunk_usage, + provider_reported_usage_cost, + stamp_litellm_bill_on_usage, +) from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, get_hidden_params, @@ -98,12 +103,17 @@ def _stamp_usage_cost( usage_chunk: ModelResponse | TextCompletionResponse | ModelResponseStream, complete_response: ModelResponse | TextCompletionResponse | ModelResponseStream, ) -> None: - bill: Final = (get_hidden_params(complete_response) or {}).get("response_cost") + complete_hidden: Final = get_hidden_params(complete_response) or {} + bill: Final = complete_hidden.get("response_cost") usage: Final[object] = getattr(usage_chunk, "usage", None) if usage is None or isinstance(bill, bool) or not isinstance(bill, (int, float)): return - if bill > 0 or isinstance(getattr(usage, "cost", None), (int, float)): - setattr(usage, "cost", float(bill)) # noqa: B010 # cost is an optional extra on the usage model + if bill <= 0 and not isinstance(getattr(usage, "cost", None), (int, float)): + return + chunk_hidden: Final = get_hidden_params(usage_chunk) + if chunk_hidden is not None and UPSTREAM_USAGE_COST in complete_hidden: + chunk_hidden[UPSTREAM_USAGE_COST] = complete_hidden[UPSTREAM_USAGE_COST] + stamp_litellm_bill_on_usage(usage_chunk, usage, float(bill)) def _stamp_served_service_tier(response: ModelResponseStream, complete_streaming_response: ModelResponse) -> None: diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index c8fa9c94c17..9df45b92686 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1077,6 +1077,30 @@ def test_cost_field_in_usage_chunks(): assert usage.completion_tokens == 5 +def test_the_rebuilt_response_keeps_the_upstream_cost_a_stamped_usage_chunk_saved(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "cost_margin_config", {"global": 0.5}) + text_chunk = ModelResponseStream( + id="chatcmpl-1", + model="gpt-4.1", + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="Hi"))], + ) + usage_chunk = ModelResponseStream( + id="chatcmpl-1", + model="gpt-4.1", + choices=[], + usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.0015), + ) + usage_chunk._hidden_params = {"response_cost": 0.0015, "upstream_usage_cost": 0.001} + + rebuilt = stream_chunk_builder([text_chunk, usage_chunk], messages=[{"role": "user", "content": "hi"}]) + + assert rebuilt is not None + assert rebuilt._hidden_params["upstream_usage_cost"] == 0.001 + assert litellm.completion_cost(rebuilt, custom_llm_provider="openai") == pytest.approx(0.0015), ( + "the upstream's 0.001 gets the 50% margin once, litellm's 0.0015 bill is not margined again" + ) + + def _perplexity_usage_chunk(cost: Mapping[str, float], finish_reason: str | None) -> ModelResponseStream: return ModelResponseStream( diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 148b7193ef6..6df4731e335 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6847,6 +6847,37 @@ def test_the_usage_chunk_carries_the_bill_the_builder_priced_for_the_spend_row() assert usage_chunk.usage is not None and usage_chunk.usage.cost == 0.00032, "the client sees what the row bills" +@pytest.mark.parametrize("upstream_cost", [None, 0.001]) +def test_rebuilding_a_priced_stream_bills_the_margin_once(monkeypatch: pytest.MonkeyPatch, upstream_cost: float | None) -> None: + from litellm.litellm_core_utils.core_helpers import stamp_litellm_bill_on_usage + from litellm.litellm_core_utils.streaming_handler import ( + _stamp_usage_cost, # pyright: ignore[reportPrivateUsage] # the helper under test + ) + + monkeypatch.setattr(litellm, "cost_margin_config", {"global": 0.5}) + model: Final = "anthropic/claude-sonnet-4-5" + shared_usage: Final = Usage(prompt_tokens=9, completion_tokens=16, total_tokens=25) + if upstream_cost is not None: + setattr(shared_usage, "cost", upstream_cost) + reference: Final = ModelResponse(model=model, usage=shared_usage.model_copy()) + billed_once: Final = litellm.completion_cost(reference, custom_llm_provider="anthropic") + complete: Final = ModelResponse(model=model, usage=shared_usage) + stamp_litellm_bill_on_usage(complete, shared_usage, billed_once) + text_chunk: Final = ModelResponseStream( + id="chunk-1", model=model, choices=[StreamingChoices(index=0, delta=Delta(content="hi"), finish_reason="stop")] + ) + usage_chunk: Final = ModelResponseStream(id="chunk-1", model=model, choices=[], usage=shared_usage) + + _stamp_usage_cost(usage_chunk, complete) + rebuilt: Final = litellm.stream_chunk_builder([text_chunk, usage_chunk], messages=[{"role": "user", "content": "hi"}]) + + assert rebuilt is not None and rebuilt.usage is not None # pyright: ignore[reportAttributeAccessIssue] # ModelResponse carries usage at runtime + assert usage_chunk.usage is not None and usage_chunk.usage.cost == pytest.approx(billed_once) + assert litellm.completion_cost(rebuilt, custom_llm_provider="anthropic") == pytest.approx(billed_once), ( + "pricing the rebuilt stream must not margin litellm's own bill again" + ) + + def test_chunk_creator_keeps_provider_model_private_across_stream(): from litellm.router_utils.add_retry_fallback_headers import ( get_hidden_params_dict, From a1c238ff4a1e1e8fc9cf2e9520c3132418e400f6 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 30 Oct 2026 11:31:04 -0700 Subject: [PATCH 44/64] refactor(rust): extract SecretValue and use it for Redis credentials (#49186) * refactor(rust): move SecretValue into litellm-secret-value * docs(rust): add rust-redaction skill * refactor(redis-client): carry credentials as SecretValue --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../.agents/skills/rust-redaction/SKILL.md | 43 +++++++++ litellm-rust/AGENTS.md | 2 + litellm-rust/Cargo.lock | 50 ++++++++--- litellm-rust/Cargo.toml | 1 + litellm-rust/crates/auth-azure/Cargo.toml | 1 + litellm-rust/crates/auth-azure/src/native.rs | 6 +- litellm-rust/crates/auth-azure/src/resolve.rs | 9 +- litellm-rust/crates/auth-azure/src/types.rs | 4 +- litellm-rust/crates/auth-gcp/Cargo.toml | 1 + litellm-rust/crates/auth-gcp/src/auth.rs | 2 +- litellm-rust/crates/auth-gcp/src/config.rs | 4 +- litellm-rust/crates/auth-types/Cargo.toml | 2 +- .../crates/auth-types/src/credential.rs | 5 +- litellm-rust/crates/auth-types/src/lib.rs | 2 - litellm-rust/crates/auth-types/src/policy.rs | 3 +- litellm-rust/crates/auth-types/src/token.rs | 3 +- litellm-rust/crates/auth-types/tests/token.rs | 6 +- litellm-rust/crates/auth/Cargo.toml | 3 + litellm-rust/crates/auth/tests/facade.rs | 3 +- litellm-rust/crates/cache-gcs/Cargo.toml | 1 + litellm-rust/crates/cache-gcs/src/token.rs | 4 +- .../crates/callbacks-legacy-python/Cargo.toml | 2 +- .../callbacks-legacy-python/src/adapter.rs | 2 +- litellm-rust/crates/config/Cargo.toml | 2 +- litellm-rust/crates/config/src/mcp.rs | 2 +- litellm-rust/crates/config/src/redis.rs | 2 +- litellm-rust/crates/config/src/settings.rs | 2 +- litellm-rust/crates/gateway-auth/Cargo.toml | 2 +- .../crates/gateway-auth/src/authentication.rs | 2 +- litellm-rust/crates/gateway-auth/src/http.rs | 2 +- litellm-rust/crates/gateway-auth/src/keys.rs | 2 +- litellm-rust/crates/gateway-auth/src/ui.rs | 2 +- .../crates/gateway-auth/tests/auth.rs | 2 +- .../crates/gateway-auth/tests/keys.rs | 2 +- litellm-rust/crates/gateway-auth/tests/ui.rs | 2 +- .../crates/gateway-inference/Cargo.toml | 1 + .../crates/gateway-inference/src/ocr.rs | 2 +- .../gateway-inference/tests/support/mod.rs | 4 +- .../crates/gateway-management/Cargo.toml | 2 +- .../crates/gateway-management/src/keys.rs | 2 +- .../crates/gateway-management/tests/keys.rs | 2 +- litellm-rust/crates/gateway-mcp/Cargo.toml | 2 +- .../crates/gateway-mcp/src/configured.rs | 2 +- litellm-rust/crates/gateway-ui/Cargo.toml | 2 +- .../crates/gateway-ui/tests/routes.rs | 2 +- litellm-rust/crates/gateway/Cargo.toml | 2 +- litellm-rust/crates/gateway/src/redis.rs | 10 +-- litellm-rust/crates/gateway/src/secrets.rs | 2 +- litellm-rust/crates/gateway/src/settings.rs | 2 +- litellm-rust/crates/gateway/tests/mcp.rs | 2 +- litellm-rust/crates/gateway/tests/server.rs | 2 +- litellm-rust/crates/host/Cargo.toml | 2 +- litellm-rust/crates/host/src/interceptors.rs | 2 +- litellm-rust/crates/inference-chat/Cargo.toml | 1 + .../crates/inference-chat/src/handler.rs | 2 +- .../crates/inference-chat/tests/caching.rs | 4 +- .../crates/inference-messages/Cargo.toml | 1 + .../crates/inference-messages/src/handler.rs | 4 +- .../inference-messages/tests/caching.rs | 4 +- litellm-rust/crates/inference-ocr/Cargo.toml | 1 + .../crates/inference-ocr/src/handler.rs | 2 +- .../inference-ocr/src/provider_config.rs | 16 ++-- .../crates/inference-ocr/src/types.rs | 4 +- litellm-rust/crates/inference-ocr/src/wire.rs | 4 +- .../inference-ocr/tests/ocr/azure_ai.rs | 6 +- .../inference-ocr/tests/ocr/lifecycle.rs | 2 +- .../crates/inference-ocr/tests/ocr/main.rs | 2 +- .../crates/inference-responses/Cargo.toml | 1 + .../inference-responses/tests/caching.rs | 4 +- .../crates/inference-testing/Cargo.toml | 1 + .../crates/inference-testing/src/lib.rs | 3 +- litellm-rust/crates/llms/Cargo.toml | 1 + .../llms/src/anthropic/chat/transformation.rs | 2 +- .../crates/llms/src/anthropic/common_utils.rs | 4 +- .../src/azure_ai/messages/transformation.rs | 2 +- .../document_intelligence/transformation.rs | 2 +- .../llms/src/azure_ai/ocr/transformation.rs | 4 +- litellm-rust/crates/llms/src/base_llm/auth.rs | 4 +- .../llms/src/base_llm/ocr/transformation.rs | 4 +- .../bedrock/chat/converse_transformation.rs | 4 +- .../messages/transformation.rs | 4 +- .../anthropic_claude_transformation.rs | 4 +- .../bedrock_mantle/messages/transformation.rs | 4 +- .../llms/src/cohere/ocr/transformation.rs | 2 +- .../src/deepseek/messages/transformation.rs | 4 +- .../src/edenai/messages/transformation.rs | 4 +- .../src/minimax/messages/transformation.rs | 4 +- .../llms/src/mistral/ocr/transformation.rs | 2 +- .../src/openai/responses/transformation.rs | 3 +- .../src/openai_like/chat/transformation.rs | 4 +- .../src/openrouter/messages/transformation.rs | 4 +- .../openrouter/responses/transformation.rs | 4 +- .../llms/src/reducto/ocr/transformation.rs | 4 +- .../src/tencent/messages/transformation.rs | 4 +- .../llms/src/vertex_ai/ocr/transformation.rs | 2 +- litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/credentials.rs | 4 +- .../python-bridge/src/routes/ocr/project.rs | 2 +- .../python-bridge/src/secrets/callback.rs | 4 +- .../python-bridge/src/secrets/config.rs | 4 +- .../python-bridge/src/secrets/operations.rs | 8 +- .../python-bridge/src/secrets/provider.rs | 4 +- .../python-bridge/src/secrets/python.rs | 4 +- .../python-bridge/src/secrets/resolved.rs | 5 +- .../python-bridge/src/secrets/runtime.rs | 4 +- .../src/secrets/vault/operation.rs | 7 +- litellm-rust/crates/redis-client/Cargo.toml | 1 + .../crates/redis-client/src/client.rs | 11 ++- .../crates/redis-client/src/config.rs | 90 ++++++------------- .../crates/redis-client/src/connection.rs | 4 +- .../crates/redis-client/src/test_support.rs | 3 +- litellm-rust/crates/router-types/Cargo.toml | 1 + .../crates/router-types/src/litellm_params.rs | 4 +- .../router-types/tests/litellm_params.rs | 4 +- litellm-rust/crates/secret-value/Cargo.toml | 14 +++ litellm-rust/crates/secret-value/src/lib.rs | 5 ++ .../src/secret_value.rs} | 6 +- litellm-rust/crates/secrets-aws/Cargo.toml | 1 + .../crates/secrets-aws/src/secret_manager.rs | 5 +- .../secrets-aws/tests/secret_manager.rs | 4 +- litellm-rust/crates/secrets-azure/Cargo.toml | 1 + .../crates/secrets-azure/src/key_vault.rs | 4 +- .../crates/secrets-azure/tests/key_vault.rs | 4 +- .../crates/secrets-cyberark/Cargo.toml | 1 + .../secrets-cyberark/src/secret_manager.rs | 5 +- .../secrets-cyberark/tests/secret_manager.rs | 4 +- litellm-rust/crates/secrets-google/Cargo.toml | 1 + .../crates/secrets-google/src/auth.rs | 2 +- litellm-rust/crates/secrets-google/src/kms.rs | 2 +- .../secrets-google/src/secret_manager.rs | 6 +- .../crates/secrets-hashicorp/Cargo.toml | 1 + .../crates/secrets-hashicorp/src/config.rs | 2 +- .../secrets-hashicorp/src/secret_manager.rs | 5 +- .../secrets-hashicorp/tests/secret_manager.rs | 6 +- litellm-rust/crates/secrets-types/Cargo.toml | 2 +- .../secrets-types/src/base_secret_manager.rs | 3 +- .../crates/secrets-types/src/config.rs | 2 +- .../crates/secrets-types/src/context.rs | 3 +- litellm-rust/crates/secrets-types/src/lib.rs | 1 - .../crates/secrets-types/src/value.rs | 2 +- .../crates/secrets-types/tests/config.rs | 6 +- .../crates/secrets-types/tests/context.rs | 5 +- .../crates/secrets-types/tests/rotation.rs | 6 +- litellm-rust/crates/secrets/Cargo.toml | 1 + .../crates/secrets/src/compatibility.rs | 3 +- litellm-rust/crates/secrets/src/handler.rs | 2 +- litellm-rust/crates/secrets/src/lib.rs | 4 +- litellm-rust/crates/secrets/src/oidc.rs | 3 +- litellm-rust/crates/secrets/src/resolver.rs | 3 +- litellm-rust/crates/secrets/src/source.rs | 3 +- litellm-rust/crates/secrets/tests/aws.rs | 4 +- litellm-rust/crates/secrets/tests/azure.rs | 5 +- .../secrets/tests/common_read_contract.rs | 4 +- litellm-rust/crates/secrets/tests/cyberark.rs | 4 +- litellm-rust/crates/secrets/tests/google.rs | 4 +- litellm-rust/crates/secrets/tests/oidc.rs | 6 +- .../crates/secrets/tests/resolution.rs | 3 +- 157 files changed, 439 insertions(+), 254 deletions(-) create mode 100644 litellm-rust/.agents/skills/rust-redaction/SKILL.md create mode 100644 litellm-rust/crates/secret-value/Cargo.toml create mode 100644 litellm-rust/crates/secret-value/src/lib.rs rename litellm-rust/crates/{auth-types/src/secret.rs => secret-value/src/secret_value.rs} (95%) diff --git a/litellm-rust/.agents/skills/rust-redaction/SKILL.md b/litellm-rust/.agents/skills/rust-redaction/SKILL.md new file mode 100644 index 00000000000..2e906e10607 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-redaction/SKILL.md @@ -0,0 +1,43 @@ +--- +name: rust-redaction +description: Keep secrets out of Debug output in litellm-rust with SecretValue and veil. Use when a struct or enum holds a password, token, key, credential URL or credential provider, when writing or reviewing a Debug impl, or when a type with secrets gets logged or traced +metadata: + internal: true +--- + +# Rust redaction + +Every secret in a Debug output is hidden the same way, so a reviewer can tell at a glance that a field is covered and a new field cannot slip through a hand-written impl + +## Where SecretValue lives + +`SecretValue` lives in `crates/secret-value` and is imported as `litellm_secret_value::SecretValue` everywhere. No crate re-exports it, so there is one import path. Depend on `litellm-secret-value` directly. Do not pull in `litellm-auth-types` or `litellm-auth` just to get it + +`crates/secret-value` holds only the type. Credential acquisition (`TokenProvider`, `CredentialPlan`, provider param tables) stays in `litellm-auth-types`, and generic clients such as `litellm-redis-client` must not depend on it + +## Rules + +- Never hand-write `impl Debug` to hide a value, and never invent markers like `""`, `""` or `.map(|_| "...")`. A hand-written impl silently prints any field added later +- A secret that crosses a crate boundary or lives in config is typed as `SecretValue`. It redacts itself, compares in constant time and only gives up the plaintext through `expose()`. A type whose only secrets are `SecretValue` fields derives plain `Debug` + ```rust + #[derive(Clone, Debug, PartialEq)] + pub struct HostEndpoint { + pub host: String, + pub port: u16, + pub username: Option, + pub password: Option, + pub tls: bool, + pub db: i64, + } + ``` +- Keep a secret as `SecretValue` until the last conversion into the third-party type (redis-rs, reqwest, an SDK) and call `expose()` there. Calling `expose()` early spreads plaintext `String`s through our own types +- A secret that has to stay a raw type, because a serde wire shape or a third-party struct owns the field, is covered by deriving `veil::Redact` on the holder and marking the field `#[redact(with = "[REDACTED]")]`. `Redact` replaces `Debug`, so do not derive both +- A URL that can carry userinfo is a secret as a whole. Do not parse it to show the host. Keep it as a `SecretValue` +- A bag of unparsed settings that may hold credentials (`additional_fields`, raw kwargs) is redacted as a whole with `#[redact(with = "[REDACTED]")]` +- An opaque handle over `Arc` whose trait has no `Debug` gets a newtype whose `Debug` prints only the type name with `finish_non_exhaustive()`. That impl shows nothing, so it is not a redaction impl and the first rule does not apply. A type holding such a handle can derive `Debug` +- Never put a secret in a tracing field or an error message. Log the type with `?value`, or log only non-secret fields. See the rust-tracing and rust-error skills + +## Tests + +- One test per type that holds a secret, built the way production builds it, asserting the plaintext is absent from `format!("{value:?}")` and `[REDACTED]` is present +- Use a distinctive plaintext (`"credential-value"`) so the absence check cannot pass by accident diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 9e7a55eb232..5eaf876e3b2 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,5 +1,7 @@ # Rust workspace rules +For keeping secrets out of Debug output, follow [.agents/skills/rust-redaction/SKILL.md](.agents/skills/rust-redaction/SKILL.md) + For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md) For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e959f77917b..f91deaaa401 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3380,6 +3380,7 @@ dependencies = [ "litellm-auth-azure", "litellm-auth-gcp", "litellm-auth-types", + "litellm-secret-value", ] [[package]] @@ -3416,6 +3417,7 @@ dependencies = [ "azure_identity", "litellm-auth-types", "litellm-error", + "litellm-secret-value", "moka", "rstest", "serde_json", @@ -3435,6 +3437,7 @@ dependencies = [ "http 1.4.2", "litellm-auth-types", "litellm-error", + "litellm-secret-value", "moka", "rstest", "serde", @@ -3450,10 +3453,10 @@ name = "litellm-auth-types" version = "0.1.0" dependencies = [ "litellm-error", + "litellm-secret-value", "rstest", "serde", "serde_json", - "subtle", "thiserror 2.0.19", "tokio", "veil", @@ -3523,6 +3526,7 @@ dependencies = [ "litellm-cache-testing", "litellm-error", "litellm-http", + "litellm-secret-value", "percent-encoding", "reqwest", "rstest", @@ -3668,10 +3672,10 @@ dependencies = [ name = "litellm-callbacks-legacy-python" version = "0.1.0" dependencies = [ - "litellm-auth", "litellm-error", "litellm-host", "litellm-host-python", + "litellm-secret-value", "proptest", "pyo3", "rstest", @@ -3692,9 +3696,9 @@ dependencies = [ name = "litellm-config" version = "0.1.0" dependencies = [ - "litellm-auth-types", "litellm-error", "litellm-router-types", + "litellm-secret-value", "rstest", "serde", "serde_json", @@ -3804,7 +3808,6 @@ dependencies = [ "envy", "futures-util", "http-body-util", - "litellm-auth-types", "litellm-circuit-breaker", "litellm-config", "litellm-error", @@ -3817,6 +3820,7 @@ dependencies = [ "litellm-llms", "litellm-redis-client", "litellm-redis-testing", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "rstest", @@ -3839,9 +3843,9 @@ dependencies = [ "axum", "axum-login", "futures-util", - "litellm-auth-types", "litellm-config", "litellm-error", + "litellm-secret-value", "litellm-secrets", "rstest", "serde", @@ -3879,6 +3883,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-router", + "litellm-secret-value", "litellm-secrets", "rstest", "serde", @@ -3895,9 +3900,9 @@ name = "litellm-gateway-management" version = "0.1.0" dependencies = [ "base64 0.22.1", - "litellm-auth-types", "litellm-error", "litellm-gateway-auth", + "litellm-secret-value", "rand 0.8.7", "rstest", "thiserror 2.0.19", @@ -3912,9 +3917,9 @@ dependencies = [ "base64 0.22.1", "futures-util", "http 1.4.2", - "litellm-auth-types", "litellm-config", "litellm-error", + "litellm-secret-value", "litellm-secrets", "moka", "rmcp", @@ -3942,9 +3947,9 @@ dependencies = [ "futures-util", "governor", "jsonwebtoken", - "litellm-auth-types", "litellm-error", "litellm-gateway-auth", + "litellm-secret-value", "rand 0.8.7", "rstest", "serde", @@ -3964,8 +3969,8 @@ name = "litellm-host" version = "0.1.0" dependencies = [ "futures-util", - "litellm-auth", "litellm-coroutine", + "litellm-secret-value", "rstest", "serde_json", "thiserror 2.0.19", @@ -4094,6 +4099,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "rstest", @@ -4124,6 +4130,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "reqwest", @@ -4154,6 +4161,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "mime_guess", @@ -4188,6 +4196,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "reqwest", @@ -4207,6 +4216,7 @@ dependencies = [ "litellm-auth", "litellm-http", "litellm-inference", + "litellm-secret-value", "litellm-secrets", "tokio", ] @@ -4254,6 +4264,7 @@ dependencies = [ "litellm-llms-types", "litellm-python-compat", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "reqwest", "rstest", @@ -4347,6 +4358,7 @@ dependencies = [ "litellm-llms-types", "litellm-model-catalog", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "litellm-secrets-aws", "litellm-secrets-types", @@ -4396,6 +4408,7 @@ dependencies = [ "futures-util", "litellm-circuit-breaker", "litellm-redis-testing", + "litellm-secret-value", "redis", "rstest", "thiserror 2.0.19", @@ -4432,12 +4445,23 @@ name = "litellm-router-types" version = "0.1.0" dependencies = [ "litellm-auth-types", + "litellm-secret-value", "rstest", "serde", "serde_json", "serde_with", ] +[[package]] +name = "litellm-secret-value" +version = "0.1.0" +dependencies = [ + "rstest", + "serde", + "subtle", + "veil", +] + [[package]] name = "litellm-secrets" version = "0.1.0" @@ -4451,6 +4475,7 @@ dependencies = [ "litellm-error", "litellm-http", "litellm-python-compat", + "litellm-secret-value", "litellm-secrets-aws", "litellm-secrets-azure", "litellm-secrets-cyberark", @@ -4479,6 +4504,7 @@ dependencies = [ "base64 0.22.1", "litellm-auth-aws", "litellm-core-utils", + "litellm-secret-value", "litellm-secrets-types", "litellm-tracing", "rstest", @@ -4498,6 +4524,7 @@ dependencies = [ "litellm-auth-types", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "percent-encoding", "reqwest", @@ -4519,6 +4546,7 @@ dependencies = [ "base64 0.22.1", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "litellm-tracing", "moka", @@ -4549,6 +4577,7 @@ dependencies = [ "litellm-auth-types", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "moka", "percent-encoding", @@ -4570,6 +4599,7 @@ dependencies = [ "litellm-core-utils", "litellm-error", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "reqwest", "rstest", @@ -4589,7 +4619,7 @@ dependencies = [ name = "litellm-secrets-types" version = "0.1.0" dependencies = [ - "litellm-auth-types", + "litellm-secret-value", "moka", "rstest", "serde", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 49a33e08456..75cac2a5d63 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -37,6 +37,7 @@ litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } litellm-framer = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } litellm-auth-types = { path = "crates/auth-types" } +litellm-secret-value = { path = "crates/secret-value" } litellm-auth-aws = { path = "crates/auth-aws" } litellm-auth-azure = { path = "crates/auth-azure" } litellm-auth-gcp = { path = "crates/auth-gcp" } diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index f99972f813b..7dcb0e513a7 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true moka.workspace = true diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index 9fda67746d1..f40745ab7a2 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -13,7 +13,8 @@ use azure_identity::{ use sha2::{Digest, Sha256}; use litellm_auth_types::{Error, ErrorDetail}; -use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced}; +use litellm_auth_types::{InputSource, ResolvedCredential, Sourced}; +use litellm_secret_value::SecretValue; use super::credential_provider_cache::{ AzureCredentialProviderCache, AzureCredentialProviderCacheKey, @@ -498,7 +499,8 @@ mod tests { use azure_core::{Bytes, Result}; use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; - use litellm_auth_types::{InputSource, SecretValue, Sourced}; + use litellm_auth_types::{InputSource, Sourced}; + use litellm_secret_value::SecretValue; fn deployment(value: T) -> Sourced { Sourced::new(value, InputSource::Deployment) diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 25b71741342..a92630af26a 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,7 +1,9 @@ use litellm_auth_types::{ - CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, - SecretValue, Sourced, TokenProviderHandle, + CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced, + TokenProviderHandle, }; +use litellm_secret_value::SecretValue; + use litellm_auth_types::{Error, ErrorDetail}; use std::future::Future; use std::pin::Pin; @@ -500,8 +502,9 @@ mod tests { use litellm_auth_types::ResolvedCredential; use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef, - CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced, + CredentialResolver, CredentialResolverHandle, InputSource, Sourced, }; + use litellm_secret_value::SecretValue; #[derive(Debug)] struct FileResolver; diff --git a/litellm-rust/crates/auth-azure/src/types.rs b/litellm-rust/crates/auth-azure/src/types.rs index a042937a047..6f8290645cf 100644 --- a/litellm-rust/crates/auth-azure/src/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -1,8 +1,10 @@ use std::collections::BTreeMap; use litellm_auth_types::{ - CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle, + CredentialResolverHandle, Error, InputSource, Sourced, TokenProviderHandle, }; +use litellm_secret_value::SecretValue; + use serde_json::{Map, Value}; use strum::EnumString; diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index 7844b35ecc2..ad60f0db817 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true google-sdk = ["dep:google-cloud-auth", "dep:http"] [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true base64.workspace = true diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs index 8ac1e0f54ba..650d5831796 100644 --- a/litellm-rust/crates/auth-gcp/src/auth.rs +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -258,7 +258,7 @@ fn auth_acquisition_error(error: gcp_auth::Error) -> Error { #[cfg(test)] mod tests { - use litellm_auth_types::SecretValue; + use litellm_secret_value::SecretValue; use super::*; diff --git a/litellm-rust/crates/auth-gcp/src/config.rs b/litellm-rust/crates/auth-gcp/src/config.rs index 5444897aaaf..133e1c26903 100644 --- a/litellm-rust/crates/auth-gcp/src/config.rs +++ b/litellm-rust/crates/auth-gcp/src/config.rs @@ -4,7 +4,9 @@ use base64::{ Engine, alphabet, engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig}, }; -use litellm_auth_types::{Error, ErrorDetail, InputSource, SecretValue, Sourced}; +use litellm_auth_types::{Error, ErrorDetail, InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use serde_json::{Map, Value}; use litellm_auth_types::VertexParams; diff --git a/litellm-rust/crates/auth-types/Cargo.toml b/litellm-rust/crates/auth-types/Cargo.toml index 400f2640221..f607f0d1da0 100644 --- a/litellm-rust/crates/auth-types/Cargo.toml +++ b/litellm-rust/crates/auth-types/Cargo.toml @@ -7,9 +7,9 @@ repository.workspace = true [dependencies] litellm-error.workspace = true +litellm-secret-value.workspace = true serde.workspace = true serde_json.workspace = true -subtle.workspace = true thiserror.workspace = true veil.workspace = true diff --git a/litellm-rust/crates/auth-types/src/credential.rs b/litellm-rust/crates/auth-types/src/credential.rs index 4f69c6fffc8..36e21f7794a 100644 --- a/litellm-rust/crates/auth-types/src/credential.rs +++ b/litellm-rust/crates/auth-types/src/credential.rs @@ -5,7 +5,8 @@ use std::sync::Arc; use veil::Redact; -use crate::{Error, ResolvedCredential, SecretValue, TokenProviderHandle}; +use crate::{Error, ResolvedCredential, TokenProviderHandle}; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, PartialEq, Eq)] pub enum CredentialFileRef { @@ -104,8 +105,8 @@ mod tests { CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, }; - use crate::SecretValue; use crate::{Error, ErrorDetail}; + use litellm_secret_value::SecretValue; #[derive(Debug)] struct HostResolver; diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs index 99651977867..ab3bd3febbf 100644 --- a/litellm-rust/crates/auth-types/src/lib.rs +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -5,7 +5,6 @@ mod error; pub mod http; mod params; mod policy; -mod secret; mod token; use serde::{Deserialize, Serialize}; @@ -55,5 +54,4 @@ pub use error::{Error, ErrorDetail}; pub use http::CredentialPlacement; pub use params::{AwsParams, ParamKind, ParamSpec, VertexParams}; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; -pub use secret::SecretValue; pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth-types/src/policy.rs b/litellm-rust/crates/auth-types/src/policy.rs index 8fff59e284a..2232350965a 100644 --- a/litellm-rust/crates/auth-types/src/policy.rs +++ b/litellm-rust/crates/auth-types/src/policy.rs @@ -71,7 +71,8 @@ impl ProviderAuthPolicy { #[cfg(test)] mod tests { use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; - use crate::{CredentialPlacement, ResolvedCredential, SecretValue}; + use crate::{CredentialPlacement, ResolvedCredential}; + use litellm_secret_value::SecretValue; const RULES: &[CredentialRule] = &[CredentialRule { kind: CredentialPlanKind::Static, diff --git a/litellm-rust/crates/auth-types/src/token.rs b/litellm-rust/crates/auth-types/src/token.rs index a3126052065..b948dda44fc 100644 --- a/litellm-rust/crates/auth-types/src/token.rs +++ b/litellm-rust/crates/auth-types/src/token.rs @@ -5,7 +5,8 @@ use std::time::SystemTime; use veil::Redact; -use crate::{Error, SecretValue}; +use crate::Error; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, PartialEq, Eq)] pub enum ResolvedCredential { diff --git a/litellm-rust/crates/auth-types/tests/token.rs b/litellm-rust/crates/auth-types/tests/token.rs index 7f7027942b0..39908ae7b8a 100644 --- a/litellm-rust/crates/auth-types/tests/token.rs +++ b/litellm-rust/crates/auth-types/tests/token.rs @@ -8,9 +8,9 @@ use std::{ time::{Duration, SystemTime}, }; -use litellm_auth_types::{ - Error, ErrorDetail, ResolvedCredential, SecretValue, TokenProviderHandle, -}; +use litellm_auth_types::{Error, ErrorDetail, ResolvedCredential, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use rstest::rstest; fn credential(index: usize, access_token: bool) -> ResolvedCredential { diff --git a/litellm-rust/crates/auth/Cargo.toml b/litellm-rust/crates/auth/Cargo.toml index ee4900ebcc2..674ae65a59f 100644 --- a/litellm-rust/crates/auth/Cargo.toml +++ b/litellm-rust/crates/auth/Cargo.toml @@ -16,3 +16,6 @@ litellm-auth-types.workspace = true litellm-auth-aws = { workspace = true, optional = true } litellm-auth-azure = { workspace = true, optional = true } litellm-auth-gcp = { workspace = true, optional = true } + +[dev-dependencies] +litellm-secret-value.workspace = true diff --git a/litellm-rust/crates/auth/tests/facade.rs b/litellm-rust/crates/auth/tests/facade.rs index f1092b15def..5d29bac1072 100644 --- a/litellm-rust/crates/auth/tests/facade.rs +++ b/litellm-rust/crates/auth/tests/facade.rs @@ -1,7 +1,8 @@ use litellm_auth::{ CredentialPlacement, CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, - ProviderAuthPolicy, ResolvedCredential, SecretValue, + ProviderAuthPolicy, ResolvedCredential, }; +use litellm_secret_value::SecretValue; const RULES: &[CredentialRule] = &[CredentialRule { kind: CredentialPlanKind::Static, diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index dbd8c9ac908..878db4b8010 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-http.workspace = true futures-util.workspace = true litellm-auth-gcp.workspace = true diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs index 079ec44dbce..852d8bbf035 100644 --- a/litellm-rust/crates/cache-gcs/src/token.rs +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -1,7 +1,9 @@ use std::{future::Future, pin::Pin}; use litellm_auth_gcp::{VertexAuth, VertexConfig}; -use litellm_auth_types::{InputSource, SecretValue, Sourced}; +use litellm_auth_types::{InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use litellm_cache::Error; pub trait TokenSource: Send + Sync + 'static { diff --git a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml index 1cd05710366..0133af16b80 100644 --- a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -16,7 +16,7 @@ serde_json.workspace = true thiserror.workspace = true [dev-dependencies] -litellm-auth.workspace = true +litellm-secret-value.workspace = true proptest.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 493e91bb5ea..9431b9d1326 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -925,11 +925,11 @@ assert all(value is failure for name, value in logger.calls if name.endswith('_h mod payload_tests { use std::ffi::CStr; - use litellm_auth::SecretValue; use litellm_host::hooks::CallHooks; use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest}; use litellm_host::lifecycle::ExecutionEvent; use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned, to_py}; + use litellm_secret_value::SecretValue; use proptest::prelude::*; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index 5b0cd22332f..278aea6200e 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth-types.workspace = true +litellm-secret-value.workspace = true litellm-router-types.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/config/src/mcp.rs b/litellm-rust/crates/config/src/mcp.rs index cbe992f9956..00480fa5be1 100644 --- a/litellm-rust/crates/config/src/mcp.rs +++ b/litellm-rust/crates/config/src/mcp.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, fmt}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use crate::Object; diff --git a/litellm-rust/crates/config/src/redis.rs b/litellm-rust/crates/config/src/redis.rs index 5674f993025..7a91691c54b 100644 --- a/litellm-rust/crates/config/src/redis.rs +++ b/litellm-rust/crates/config/src/redis.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, fmt, time::Duration}; -use litellm_auth_types::SecretValue; use litellm_router_types::Spelled; +use litellm_secret_value::SecretValue; use serde::{Deserialize, Deserializer, de::Error as _}; use crate::{AdditionalFields, Error, Object, Value}; diff --git a/litellm-rust/crates/config/src/settings.rs b/litellm-rust/crates/config/src/settings.rs index 9571c4fa5de..1bfc19983f3 100644 --- a/litellm-rust/crates/config/src/settings.rs +++ b/litellm-rust/crates/config/src/settings.rs @@ -1,6 +1,6 @@ use std::fmt; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use crate::{ diff --git a/litellm-rust/crates/gateway-auth/Cargo.toml b/litellm-rust/crates/gateway-auth/Cargo.toml index 27f20203843..ba9097e85d8 100644 --- a/litellm-rust/crates/gateway-auth/Cargo.toml +++ b/litellm-rust/crates/gateway-auth/Cargo.toml @@ -6,9 +6,9 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum = { workspace = true, features = ["matched-path"] } axum-login.workspace = true -litellm-auth-types.workspace = true litellm-config.workspace = true litellm-error.workspace = true litellm-secrets.workspace = true diff --git a/litellm-rust/crates/gateway-auth/src/authentication.rs b/litellm-rust/crates/gateway-auth/src/authentication.rs index 4ecfaa375b4..349913619f5 100644 --- a/litellm-rust/crates/gateway-auth/src/authentication.rs +++ b/litellm-rust/crates/gateway-auth/src/authentication.rs @@ -1,8 +1,8 @@ use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime}; use axum::http::{Request, request::Parts}; -use litellm_auth_types::SecretValue; use litellm_config::Config; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; diff --git a/litellm-rust/crates/gateway-auth/src/http.rs b/litellm-rust/crates/gateway-auth/src/http.rs index ddbf240509a..ad1718161c0 100644 --- a/litellm-rust/crates/gateway-auth/src/http.rs +++ b/litellm-rust/crates/gateway-auth/src/http.rs @@ -4,7 +4,7 @@ use axum::{ middleware::Next, response::Response, }; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{AccessRequest, Auth, AuthenticatedRequest, Credential, Error}; diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index 1fe8fa60e0e..0deaf3f7a82 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -1,6 +1,6 @@ use std::{future::Future, pin::Pin, time::SystemTime}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{Error, hash_token}; diff --git a/litellm-rust/crates/gateway-auth/src/ui.rs b/litellm-rust/crates/gateway-auth/src/ui.rs index 3d655bde187..ee4dd670aa7 100644 --- a/litellm-rust/crates/gateway-auth/src/ui.rs +++ b/litellm-rust/crates/gateway-auth/src/ui.rs @@ -3,7 +3,7 @@ use std::{convert::Infallible, sync::Arc}; use axum::extract::FromRequestParts; use axum::http::{header::AUTHORIZATION, request::Parts}; use axum_login::{AuthUser, AuthnBackend, UserId}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; diff --git a/litellm-rust/crates/gateway-auth/tests/auth.rs b/litellm-rust/crates/gateway-auth/tests/auth.rs index 781ba385381..54180e30504 100644 --- a/litellm-rust/crates/gateway-auth/tests/auth.rs +++ b/litellm-rust/crates/gateway-auth/tests/auth.rs @@ -8,9 +8,9 @@ use axum::{ routing::get, }; use futures_util::future::BoxFuture; -use litellm_auth_types::SecretValue; use litellm_config::Config; use litellm_gateway_auth::{Auth, RequireMasterKey, hash_token}; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use rstest::{fixture, rstest}; use tower::ServiceExt; diff --git a/litellm-rust/crates/gateway-auth/tests/keys.rs b/litellm-rust/crates/gateway-auth/tests/keys.rs index 4d16d3573bc..1f6fc461a5b 100644 --- a/litellm-rust/crates/gateway-auth/tests/keys.rs +++ b/litellm-rust/crates/gateway-auth/tests/keys.rs @@ -1,10 +1,10 @@ use std::time::{Duration, SystemTime}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ Error, keys::{KeyHash, KeyLookup, KeyStatus, LookupFuture, verify_key}, }; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; struct Lookup { diff --git a/litellm-rust/crates/gateway-auth/tests/ui.rs b/litellm-rust/crates/gateway-auth/tests/ui.rs index 2e5f72d605d..8d465e1198d 100644 --- a/litellm-rust/crates/gateway-auth/tests/ui.rs +++ b/litellm-rust/crates/gateway-auth/tests/ui.rs @@ -1,6 +1,6 @@ use axum_login::{AuthUser, AuthnBackend}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{UiBackend, UiCredentials}; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; #[fixture] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 034c3a9826e..69ca561e4b5 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-cache-response.workspace = true axum = { workspace = true, features = ["json", "multipart", "original-uri"] } base64.workspace = true diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs index f200fd7b141..b8e6ef4ff4f 100644 --- a/litellm-rust/crates/gateway-inference/src/ocr.rs +++ b/litellm-rust/crates/gateway-inference/src/ocr.rs @@ -2,10 +2,10 @@ use litellm_gateway_auth::AuthenticatedRequest; use std::sync::Arc; use axum::{Json, extract::State, http::HeaderMap, response::IntoResponse}; -use litellm_auth::SecretValue; use litellm_inference_ocr::types::{LiteLLMOcrRequest, OcrConnectionInputs, OcrDocumentInput}; use litellm_llms::base_llm::ocr::transformation::decode_request_value; use litellm_llms_types::formats::ocr::OcrDocument; +use litellm_secret_value::SecretValue; use serde_json::Value; use crate::{ diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 737dc0c830c..185cbc8594d 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -13,7 +13,9 @@ use futures_util::future::BoxFuture; use litellm_gateway_inference::{Deployment, Gateway, router}; use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; use litellm_inference::resources::CoreResources; -use litellm_secrets::{SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::source::SecretSource; + use serde_json::Value; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; diff --git a/litellm-rust/crates/gateway-management/Cargo.toml b/litellm-rust/crates/gateway-management/Cargo.toml index 502760f8fbe..1ff5dd5319d 100644 --- a/litellm-rust/crates/gateway-management/Cargo.toml +++ b/litellm-rust/crates/gateway-management/Cargo.toml @@ -6,8 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true base64.workspace = true -litellm-auth-types.workspace = true litellm-error.workspace = true litellm-gateway-auth.workspace = true rand.workspace = true diff --git a/litellm-rust/crates/gateway-management/src/keys.rs b/litellm-rust/crates/gateway-management/src/keys.rs index 45d915b7b3d..feeb40b2403 100644 --- a/litellm-rust/crates/gateway-management/src/keys.rs +++ b/litellm-rust/crates/gateway-management/src/keys.rs @@ -1,8 +1,8 @@ use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime}; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::keys::{KeyHash, KeyStatus}; +use litellm_secret_value::SecretValue; use rand::{RngCore, rngs::OsRng}; use crate::Error; diff --git a/litellm-rust/crates/gateway-management/tests/keys.rs b/litellm-rust/crates/gateway-management/tests/keys.rs index bf0221cf1be..8a970bb407c 100644 --- a/litellm-rust/crates/gateway-management/tests/keys.rs +++ b/litellm-rust/crates/gateway-management/tests/keys.rs @@ -7,7 +7,6 @@ use std::{ }; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ hash_token, keys::{KeyHash, KeyLookup, KeyStatus, verify_key}, @@ -16,6 +15,7 @@ use litellm_gateway_management::{ Error, keys::{Keys, Revocation}, }; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; use support::{MemoryStore, Unavailable}; diff --git a/litellm-rust/crates/gateway-mcp/Cargo.toml b/litellm-rust/crates/gateway-mcp/Cargo.toml index bed48bad58d..65077eec0ea 100644 --- a/litellm-rust/crates/gateway-mcp/Cargo.toml +++ b/litellm-rust/crates/gateway-mcp/Cargo.toml @@ -7,11 +7,11 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum = { workspace = true, features = ["json", "query"] } futures-util.workspace = true litellm-config.workspace = true litellm-secrets.workspace = true -litellm-auth-types.workspace = true litellm-error.workspace = true base64.workspace = true http.workspace = true diff --git a/litellm-rust/crates/gateway-mcp/src/configured.rs b/litellm-rust/crates/gateway-mcp/src/configured.rs index cc7a52fc029..33cbd243530 100644 --- a/litellm-rust/crates/gateway-mcp/src/configured.rs +++ b/litellm-rust/crates/gateway-mcp/src/configured.rs @@ -8,8 +8,8 @@ use std::{ use base64::{Engine, engine::general_purpose::STANDARD}; use futures_util::future::try_join_all; use http::{HeaderName, HeaderValue}; -use litellm_auth_types::SecretValue; use litellm_config::{McpAuth, McpServer, McpTransport}; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use rmcp::{ RoleClient, ServiceExt, diff --git a/litellm-rust/crates/gateway-ui/Cargo.toml b/litellm-rust/crates/gateway-ui/Cargo.toml index 1c65665be72..a6104455e96 100644 --- a/litellm-rust/crates/gateway-ui/Cargo.toml +++ b/litellm-rust/crates/gateway-ui/Cargo.toml @@ -23,10 +23,10 @@ tower-http = { version = "0.6.11", features = ["fs", "set-header"] } tower-sessions.workspace = true [dev-dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true serde_json.workspace = true tower = { version = "0.5", features = ["util"] } -litellm-auth-types.workspace = true rstest.workspace = true tempfile.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/gateway-ui/tests/routes.rs b/litellm-rust/crates/gateway-ui/tests/routes.rs index 95246f33a35..86b7db63ed9 100644 --- a/litellm-rust/crates/gateway-ui/tests/routes.rs +++ b/litellm-rust/crates/gateway-ui/tests/routes.rs @@ -16,8 +16,8 @@ use tower_sessions::{ session::{Id, Record}, }; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::UiBackend; +use litellm_secret_value::SecretValue; struct App { router: Router, diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 4d831aa2314..6eeaa43e013 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum.workspace = true envy = "0.4.2" http-body-util = "0.1" @@ -17,7 +18,6 @@ thiserror.workspace = true futures-util.workspace = true litellm-gateway-auth.workspace = true litellm-gateway-ui.workspace = true -litellm-auth-types.workspace = true litellm-config.workspace = true litellm-error.workspace = true litellm-circuit-breaker.workspace = true diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis.rs index 197f40e7988..c40b289029b 100644 --- a/litellm-rust/crates/gateway/src/redis.rs +++ b/litellm-rust/crates/gateway/src/redis.rs @@ -152,12 +152,12 @@ async fn reachable(client: RedisClient) -> Option { fn client(params: Params, breaker: Option) -> Result { tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured"); let endpoint = match params.endpoint { - RedisEndpoint::Url(url) => Endpoint::Url(url.expose().to_owned()), + RedisEndpoint::Url(url) => Endpoint::Url(url), RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint { host: host.host, port: host.port, - username: host.username.map(|value| value.expose().to_owned()), - password: host.password.map(|value| value.expose().to_owned()), + username: host.username, + password: host.password, tls: host.tls, db: host.db, }), @@ -170,8 +170,8 @@ fn client(params: Params, breaker: Option) -> Result futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async { Ok(None) }) } diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index c876c1cbcbb..7de7a97c870 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -159,7 +159,7 @@ async fn mounts_ui_without_exposing_credentials_or_authorizing_inference(inferen std::fs::write(assets.path().join("index.html"), "dashboard").unwrap(); let backend = litellm_gateway_auth::UiBackend::new( "admin".into(), - litellm_auth_types::SecretValue::new("ui-password"), + litellm_secret_value::SecretValue::new("ui-password"), ) .unwrap(); let ui = litellm_gateway_ui::router( diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index d3354fa8622..e0ccab4e0ad 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -6,8 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true -litellm-auth.workspace = true litellm-coroutine.workspace = true serde_json.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index 3cff781792f..39af7c410e0 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -22,7 +22,7 @@ pub struct RequestContext { /// Optional-param names that carry credentials and must be redacted when logged. pub secret_fields: Vec, /// The credential the route resolved for the provider call. - pub api_key: Option, + pub api_key: Option, } #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/litellm-rust/crates/inference-chat/Cargo.toml b/litellm-rust/crates/inference-chat/Cargo.toml index 8a779d44976..bcde3c8d30f 100644 --- a/litellm-rust/crates/inference-chat/Cargo.toml +++ b/litellm-rust/crates/inference-chat/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-cache-response.workspace = true litellm-core-utils.workspace = true diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index 1c724a47956..f948e8fffaf 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -58,7 +58,7 @@ pub(super) async fn execute( custom_llm_provider: provider_name.to_owned(), optional_params: Value::Object(call.optional_params), secret_fields: Vec::new(), - api_key: call.api_key.map(litellm_auth::SecretValue::new), + api_key: call.api_key.map(litellm_secret_value::SecretValue::new), }; let authenticated = resolve_auth(auth, environment, &|key| secrets.get(key)).await?; let identity = litellm_host::interceptors::ProviderIdentity { diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index 4b3e32f04bc..dd229d3d657 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -140,7 +140,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -154,7 +154,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index 694d9b4de31..ac52cb65a0d 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true bytes.workspace = true futures-util.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 9d0bf09197e..2e647688a1a 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -2,7 +2,9 @@ use std::time::Duration; use bytes::Bytes; use futures_util::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream}; -use litellm_auth::{AuthServices, SecretValue}; +use litellm_auth::AuthServices; +use litellm_secret_value::SecretValue; + use litellm_core_utils::{ dot_notation_indexing::delete_nested_value, get_provider_specific_headers::get_provider_specific_headers, settings::Lookup, diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 480b3c4f3b7..663c87453a9 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -107,7 +107,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -121,7 +121,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-ocr/Cargo.toml b/litellm-rust/crates/inference-ocr/Cargo.toml index bbbdaddb050..ad53126b9fe 100644 --- a/litellm-rust/crates/inference-ocr/Cargo.toml +++ b/litellm-rust/crates/inference-ocr/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true base64.workspace = true bytes.workspace = true futures-util.workspace = true diff --git a/litellm-rust/crates/inference-ocr/src/handler.rs b/litellm-rust/crates/inference-ocr/src/handler.rs index 426d4297f28..c126a192153 100644 --- a/litellm-rust/crates/inference-ocr/src/handler.rs +++ b/litellm-rust/crates/inference-ocr/src/handler.rs @@ -128,7 +128,7 @@ mod tests { decode_request(OcrWireRequest { model: model.into(), document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_key: Some(litellm_secret_value::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/inference-ocr/src/provider_config.rs b/litellm-rust/crates/inference-ocr/src/provider_config.rs index 19e50af0dd1..e76d291254b 100644 --- a/litellm-rust/crates/inference-ocr/src/provider_config.rs +++ b/litellm-rust/crates/inference-ocr/src/provider_config.rs @@ -1,4 +1,6 @@ -use litellm_auth::{InputSource, SecretValue, Sourced}; +use litellm_auth::{InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::get_llm_provider_logic::{ CustomLlmProvider, LlmProviders, get_custom_llm_provider, }; @@ -407,7 +409,7 @@ mod tests { fn connection_resolution_preserves_dynamic_precedence_and_input_sources() { let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("explicit-key"), + litellm_secret_value::SecretValue::new("explicit-key"), InputSource::Deployment, )), api_base: Some(Sourced::new( @@ -415,7 +417,7 @@ mod tests { InputSource::Deployment, )), dynamic_api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("dynamic-key"), + litellm_secret_value::SecretValue::new("dynamic-key"), InputSource::Environment, )), dynamic_api_base: Some(Sourced::new( @@ -455,7 +457,7 @@ mod tests { ) { let dynamic_key = dynamic_value.map(|value| { Sourced::new( - litellm_auth::SecretValue::new(value), + litellm_secret_value::SecretValue::new(value), InputSource::Environment, ) }); @@ -463,7 +465,7 @@ mod tests { dynamic_value.map(|value| Sourced::new(value.into(), InputSource::Environment)); let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("explicit-key"), + litellm_secret_value::SecretValue::new("explicit-key"), InputSource::Deployment, )), api_base: Some(Sourced::new( @@ -502,14 +504,14 @@ mod tests { OcrCredentialInputs { api_key: explicit_key.map(|value| { Sourced::new( - litellm_auth::SecretValue::new(value), + litellm_secret_value::SecretValue::new(value), InputSource::Deployment, ) }), api_base: explicit_base .map(|value| Sourced::new(value.into(), InputSource::Deployment)), dynamic_api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("dynamic-key"), + litellm_secret_value::SecretValue::new("dynamic-key"), InputSource::Environment, )), dynamic_api_base: Some(Sourced::new( diff --git a/litellm-rust/crates/inference-ocr/src/types.rs b/litellm-rust/crates/inference-ocr/src/types.rs index 98c609a158f..1b2dbb52d77 100644 --- a/litellm-rust/crates/inference-ocr/src/types.rs +++ b/litellm-rust/crates/inference-ocr/src/types.rs @@ -1,7 +1,9 @@ use std::{collections::BTreeMap, path::PathBuf, time::Duration}; use bytes::Bytes; -use litellm_auth::{InputSource, SecretValue, TokenProviderHandle}; +use litellm_auth::{InputSource, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::call_arguments::CallArguments; use litellm_llms::base_llm::ocr::{ error::Error, diff --git a/litellm-rust/crates/inference-ocr/src/wire.rs b/litellm-rust/crates/inference-ocr/src/wire.rs index 56d22be6523..45a2865823f 100644 --- a/litellm-rust/crates/inference-ocr/src/wire.rs +++ b/litellm-rust/crates/inference-ocr/src/wire.rs @@ -1,6 +1,8 @@ use std::{collections::BTreeMap, time::Duration}; -use litellm_auth::{InputSource, SecretValue}; +use litellm_auth::InputSource; +use litellm_secret_value::SecretValue; + use litellm_llms::base_llm::ocr::{error::Error, transformation::decode_request_value}; use litellm_llms_types::formats::ocr::OcrDocument; use serde::Deserialize; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs index 6f591af73e0..65c7f309b82 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs @@ -3,9 +3,9 @@ use std::sync::{ atomic::{AtomicUsize, Ordering}, }; -use litellm_auth::{ - ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, -}; +use litellm_auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use rstest::rstest; use super::*; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs index 75a91092daf..9f59bb9a2a2 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs @@ -239,7 +239,7 @@ impl litellm_host_native::services::HostCallHandler for CallerTokenHost { OcrOp::AcquireAzureAdToken(reply) => { self.trace.lock().unwrap().push("token".into()); reply.send(litellm_auth::ResolvedCredential::Static( - litellm_auth::SecretValue::new("caller-token"), + litellm_secret_value::SecretValue::new("caller-token"), )); Ok(()) } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs index 5a48bba9e1a..6871d1ca1f6 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs @@ -80,7 +80,7 @@ fn wire(model: &str, base: &str, document: Value, options: Value) -> OcrWireRequ OcrWireRequest { model: model.into(), document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_key: Some(litellm_secret_value::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/inference-responses/Cargo.toml b/litellm-rust/crates/inference-responses/Cargo.toml index 083179c50c8..f90651634c5 100644 --- a/litellm-rust/crates/inference-responses/Cargo.toml +++ b/litellm-rust/crates/inference-responses/Cargo.toml @@ -25,6 +25,7 @@ tokio-tungstenite.workspace = true tracing.workspace = true [dev-dependencies] +litellm-secret-value.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache.workspace = true litellm-cache-memory.workspace = true diff --git a/litellm-rust/crates/inference-responses/tests/caching.rs b/litellm-rust/crates/inference-responses/tests/caching.rs index 462461695c3..56e25fbcb98 100644 --- a/litellm-rust/crates/inference-responses/tests/caching.rs +++ b/litellm-rust/crates/inference-responses/tests/caching.rs @@ -98,7 +98,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -112,7 +112,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-testing/Cargo.toml b/litellm-rust/crates/inference-testing/Cargo.toml index 748b1df0d3b..5f690532d9c 100644 --- a/litellm-rust/crates/inference-testing/Cargo.toml +++ b/litellm-rust/crates/inference-testing/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true publish = false [dependencies] +litellm-secret-value.workspace = true litellm-auth.workspace = true futures-util.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/inference-testing/src/lib.rs b/litellm-rust/crates/inference-testing/src/lib.rs index 2e20d5a9e91..f330e02e99a 100644 --- a/litellm-rust/crates/inference-testing/src/lib.rs +++ b/litellm-rust/crates/inference-testing/src/lib.rs @@ -5,7 +5,8 @@ use litellm_http::{ ClientVariant, HttpClientConfig, HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver, }; -use litellm_secrets::{SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::source::SecretSource; use litellm_inference::resources::CoreResources; diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 962f61ef4ec..79ee17e8b97 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true test-support = ["litellm-http/test-support"] [dependencies] +litellm-secret-value.workspace = true litellm-llms-types.workspace = true litellm-core-utils.workspace = true litellm-error.workspace = true diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index eb796a900f5..164dbce1264 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -1,4 +1,3 @@ -use litellm_auth::SecretValue; use litellm_core_utils::decline::Decline; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, @@ -10,6 +9,7 @@ use litellm_llms_types::{ }, providers::anthropic::DEFAULT_HEADERS, }; +use litellm_secret_value::SecretValue; use serde::Deserialize; use serde_json::{Map, Value, json}; diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 7991a57353e..aec5d9ebc27 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1,5 +1,7 @@ use crate::base_llm::messages::context::MessagesModelCapabilities; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{ has_header, header_value, header_values, with_header, without_headers, diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index d578acea68c..a2ba2cee0d3 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -1,10 +1,10 @@ -use litellm_auth::SecretValue; use litellm_http::request::{has_bearer_auth, has_header}; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH}, }; use litellm_router_types::LitellmParams; +use litellm_secret_value::SecretValue; use crate::{ Error, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 36dfe2c890e..870306b1a83 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -1250,7 +1250,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 2f7957fe56e..25f93846967 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -211,7 +211,7 @@ mod tests { #[fixture] fn connection() -> OcrConnection { OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_base: Some("https://example.com".into()), ..Default::default() } @@ -312,7 +312,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/base_llm/auth.rs b/litellm-rust/crates/llms/src/base_llm/auth.rs index de271a269d3..ab2cf00614e 100644 --- a/litellm-rust/crates/llms/src/base_llm/auth.rs +++ b/litellm-rust/crates/llms/src/base_llm/auth.rs @@ -5,7 +5,9 @@ //! that split: `validate_environment` shapes the forwarded headers and names the credential //! as an [`AuthScheme`], and [`resolve_auth`] turns the scheme into headers and a signer. -use litellm_auth::{AuthServices, CredentialPlacement, SecretValue, TokenProviderHandle}; +use litellm_auth::{AuthServices, CredentialPlacement, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{AwsCredentialSource, SigV4Signer}; use litellm_auth_gcp::VertexConfig; use litellm_http::request::with_header; diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index c90fd99809e..9574e99d357 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -5,7 +5,9 @@ use serde::{ use serde_json::{Map, Value}; use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; -use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; +use litellm_auth::{InputSource, Sourced, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::{call_arguments::CallArguments, settings::ProcessEnvironment}; use litellm_http::outbound::RequestSigner; use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat}; diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index 976c9161dd2..6cf0167a30f 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -1,5 +1,7 @@ use litellm_auth::AwsParams; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 79c30df9b19..5b89345a1c5 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -1,6 +1,8 @@ use std::sync::LazyLock; -use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec, SecretValue}; +use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{AwsCredentialSource, resolve_aws_region}; use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{header_value, with_header, without_headers}; diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs index 3ec60f12021..1d380c917f6 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs @@ -7,7 +7,9 @@ use crate::{ }; use futures_util::StreamExt; use litellm_auth::AwsParams; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, constants::{ diff --git a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs index ae13704f4a0..f139b93e259 100644 --- a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs @@ -1,6 +1,8 @@ use std::sync::LazyLock; -use litellm_auth::{AwsParams, CredentialPlacement, SecretValue}; +use litellm_auth::{AwsParams, CredentialPlacement}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::AwsCredentialSource; use litellm_auth_aws::bedrock_model_id_and_region; use litellm_core_utils::settings::resolve_non_empty; diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index 76ff27b56fc..5bad06adda2 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -735,7 +735,7 @@ mod tests { assert!(matches!( CohereParseConfig.resolve_headers( &OcrConnection { - api_key: Some(litellm_auth::SecretValue::new(" ")), + api_key: Some(litellm_secret_value::SecretValue::new(" ")), ..Default::default() }, &|_| None, diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 168c3e63920..0dbcea01439 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool}, diff --git a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs index 2c3a7eaef1d..d5944484522 100644 --- a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::formats::messages::{ MessagesOptionalParams, MessagesRequest, MessagesResponse, diff --git a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs index 012a657fceb..dafe5fbd16b 100644 --- a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index fa4e2254e86..679aa65f7fd 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -219,7 +219,7 @@ mod tests { #[default(vec![])] extra_headers: Vec<(String, String)>, ) -> OcrConnection { OcrConnection { - api_key: api_key.map(litellm_auth::SecretValue::new), + api_key: api_key.map(litellm_secret_value::SecretValue::new), extra_headers, ..OcrConnection::default() } diff --git a/litellm-rust/crates/llms/src/openai/responses/transformation.rs b/litellm-rust/crates/llms/src/openai/responses/transformation.rs index cf74871894b..deb163daa07 100644 --- a/litellm-rust/crates/llms/src/openai/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openai/responses/transformation.rs @@ -3,7 +3,8 @@ use litellm_llms_types::formats::responses::{ }; use serde_json::{Map, Value}; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; use crate::{ Error, ErrorDetail, diff --git a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs index e998f8785cb..83fb8ac3fa2 100644 --- a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs @@ -4,7 +4,9 @@ //! the `max_completion_tokens` -> `max_tokens` rename and the usage //! `*_tokens` null-to-zero sanitize. -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::core_helpers::unix_now; use litellm_llms_types::formats::chat_completions::{ ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, diff --git a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs index 5b4e4f77f2a..eeb3d1cb7f8 100644 --- a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{has_bearer_auth, without_headers}; use litellm_llms_types::{ diff --git a/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs b/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs index 900fba9e916..76155ba55e5 100644 --- a/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::formats::responses::ResponsesApiResponse; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index bf0ddc1c7a8..beda6899d04 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -690,7 +690,7 @@ mod tests { #[test] fn explicit_key_precedes_environment_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("passed-key")), + api_key: Some(litellm_secret_value::SecretValue::new("passed-key")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some("env-key".into())).unwrap(); @@ -700,7 +700,7 @@ mod tests { #[test] fn blank_explicit_key_uses_environment_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new(" ")), + api_key: Some(litellm_secret_value::SecretValue::new(" ")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some(" env-key ".into())).unwrap(); diff --git a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs index a70336748ad..621a09f4fb0 100644 --- a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index 77881cd2e78..b902967d9dc 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -137,7 +137,7 @@ impl VertexAiOcrConfig { connection .api_key .as_ref() - .map(litellm_auth::SecretValue::expose), + .map(litellm_secret_value::SecretValue::expose), config, &|name: &str| connection.secret(name), ) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index a4e1c57ab1e..e3ef45f7876 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,6 +19,7 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] +litellm-secret-value.workspace = true fancy-regex.workspace = true litellm-tracing.workspace = true litellm-spend-clickhouse.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/credentials.rs b/litellm-rust/crates/python-bridge/src/credentials.rs index a4bb064aae0..adc52a7e77a 100644 --- a/litellm-rust/crates/python-bridge/src/credentials.rs +++ b/litellm-rust/crates/python-bridge/src/credentials.rs @@ -2,7 +2,9 @@ //! keyword arguments and acquired on the host's own thread when the call asks for one. use crate::callable::wrap_failure; -use litellm_auth::{ResolvedCredential, SecretValue}; +use litellm_auth::ResolvedCredential; +use litellm_secret_value::SecretValue; + use pyo3::{ exceptions::PyTypeError, gc::{PyTraverseError, PyVisit}, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index c2c3fe51600..73a2922804a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,4 +1,3 @@ -use litellm_auth::SecretValue; use litellm_auth::VertexParams; use litellm_host_python::{from_py, present}; use litellm_inference_ocr::{ @@ -6,6 +5,7 @@ use litellm_inference_ocr::{ wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, }; use litellm_llms::base_llm::ocr::{error::Error, transformation::DROP_PARAMS_PARAM}; +use litellm_secret_value::SecretValue; use pyo3::{prelude::*, types::PyDict}; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index 669c10e602a..af6485a56e1 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -2,9 +2,11 @@ use std::{future::Future, pin::Pin, sync::Arc}; use litellm_core_utils::settings::Lookup; use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_secret_value::SecretValue; use litellm_secrets::{ - Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, + Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, }; + use pyo3::{ exceptions::PyException, prelude::*, diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs index 7afd06e2373..c1203dcbd89 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/config.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/config.rs @@ -1,7 +1,9 @@ use std::sync::Arc; +use litellm_secret_value::SecretValue; use litellm_secrets::{SecretManager, SecretManagerState}; -use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, SecretValue}; +use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem}; + use pyo3::prelude::*; use serde_json::Value; diff --git a/litellm-rust/crates/python-bridge/src/secrets/operations.rs b/litellm-rust/crates/python-bridge/src/secrets/operations.rs index b458e3f34e8..027da1e4718 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/operations.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/operations.rs @@ -81,7 +81,7 @@ impl From for PythonMutationError { pub(super) async fn write_python_provider( manager: &SecretManager, name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, ) -> Result { match manager { SecretManager::Cyberark(client) => { @@ -128,7 +128,7 @@ pub(super) async fn rotate_python_provider( manager: &SecretManager, current_name: &str, new_name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, ) -> Result { match manager { SecretManager::Cyberark(client) => { @@ -186,7 +186,7 @@ pub(super) enum PythonMutationResponse { pub(super) async fn write_python_provider_with_context( manager: &SecretManager, name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, context: &litellm_secrets_types::SecretWriteContext, ) -> Result { if let (SecretManager::HashicorpVault(client), SecretOperationContext::Hashicorp(operation)) = @@ -235,7 +235,7 @@ pub(super) async fn rotate_python_provider_with_context( manager: &SecretManager, current_name: &str, new_name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, context: &SecretOperationContext, ) -> Result { if let (SecretManager::HashicorpVault(client), SecretOperationContext::Hashicorp(context)) = diff --git a/litellm-rust/crates/python-bridge/src/secrets/provider.rs b/litellm-rust/crates/python-bridge/src/secrets/provider.rs index ff10f596311..55343fb63bb 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/provider.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/provider.rs @@ -1,7 +1,9 @@ use std::time::Duration; use super::operations::PythonReadRequest; -use litellm_secrets::{KeyManagementSystem, SecretValue}; +use litellm_secret_value::SecretValue; +use litellm_secrets::KeyManagementSystem; + use litellm_secrets_types::{ AwsOperationContext, CyberarkOperationContext, GoogleOperationContext, HashicorpOperationContext, SecretOperationContext, diff --git a/litellm-rust/crates/python-bridge/src/secrets/python.rs b/litellm-rust/crates/python-bridge/src/secrets/python.rs index 2f040bff773..c5d9e3565b3 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/python.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/python.rs @@ -2,7 +2,9 @@ use std::sync::Arc; use futures_util::future::BoxFuture; use litellm_host_python::{PythonContext, attach_blocking}; -use litellm_secrets::{Error, SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::{Error, source::SecretSource}; + use pyo3::prelude::*; use super::error::external_error; diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs index deb07175e7a..752b533886e 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -4,10 +4,9 @@ use futures_util::future::BoxFuture; use litellm_core_utils::settings::ProcessEnvironment; use litellm_host_python::PythonContext; use litellm_http::Client; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; -use litellm_secrets::{ - Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver, SecretValue, -}; +use litellm_secrets::{Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver}; use super::config::SecretManagerSnapshot; diff --git a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs index b956ec297b6..0d52b0aa629 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -307,7 +307,7 @@ impl NativeSecretManager { super::operations::write_python_provider_with_context( &backend, &secret_name, - &litellm_secrets::SecretValue::new(secret_value), + &litellm_secret_value::SecretValue::new(secret_value), &context, ) .await, @@ -365,7 +365,7 @@ impl NativeSecretManager { &backend, ¤t_secret_name, &new_secret_name, - &litellm_secrets::SecretValue::new(new_secret_value), + &litellm_secret_value::SecretValue::new(new_secret_value), &context, ) .await, diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs index 41658f0baf8..b99dff7aeac 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs @@ -1,9 +1,8 @@ use std::collections::HashMap; -use litellm_secrets::{ - SecretValue, - hashicorp::{HashicorpVault, RawOperationError}, -}; +use litellm_secret_value::SecretValue; +use litellm_secrets::hashicorp::{HashicorpVault, RawOperationError}; + use litellm_secrets_types::{HashicorpOperationContext, SecretWriteContext}; use serde_json::value::RawValue; diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml index eab57b36423..3e9447a6bfa 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-circuit-breaker.workspace = true +litellm-secret-value.workspace = true futures-util.workspace = true redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] } thiserror.workspace = true diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index a02b0847663..48dbc2b2dfb 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -29,7 +29,7 @@ impl RedisClient { pub fn new(config: RedisConfig) -> Result { let connection = match &config.endpoint { Endpoint::Url(url) => { - Connection::standalone(url.as_str().into_connection_info()?, config.timeout)? + Connection::standalone(url.expose().into_connection_info()?, config.timeout)? } Endpoint::Host(host) => { Connection::standalone(host.connection_info()?, config.timeout)? @@ -204,6 +204,7 @@ mod tests { use std::time::{Duration, Instant}; use litellm_redis_testing::RedisServer; + use litellm_secret_value::SecretValue; use rstest::rstest; use litellm_circuit_breaker::{Breaker, Config, State}; @@ -630,8 +631,12 @@ mod tests { db: 3, })) .unwrap(); - let in_db_0 = connect(Endpoint::Url(server.url().to_owned())).unwrap(); - let via_url = connect(Endpoint::Url(format!("{}/3", server.url()))).unwrap(); + let in_db_0 = connect(Endpoint::Url(SecretValue::new(server.url().to_owned()))).unwrap(); + let via_url = connect(Endpoint::Url(SecretValue::new(format!( + "{}/3", + server.url() + )))) + .unwrap(); in_db_3 .set("k", "3", Duration::from_secs(60)) diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs index 64934c68891..c40b9eb6b36 100644 --- a/litellm-rust/crates/redis-client/src/config.rs +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -1,8 +1,9 @@ -use std::{fmt, time::Duration}; +use std::time::Duration; use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult}; use litellm_circuit_breaker::Breaker; +use litellm_secret_value::SecretValue; use crate::Namespace; @@ -15,28 +16,28 @@ pub struct RedisConfig { pub breaker: Option, } -#[derive(Clone, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub enum Endpoint { - Url(String), + Url(SecretValue), Host(HostEndpoint), Cluster(ClusterEndpoint), } -#[derive(Clone, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub struct HostEndpoint { pub host: String, pub port: u16, - pub username: Option, - pub password: Option, + pub username: Option, + pub password: Option, pub tls: bool, pub db: i64, } -#[derive(Clone, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub struct ClusterEndpoint { pub nodes: Vec, - pub username: Option, - pub password: Option, + pub username: Option, + pub password: Option, pub tls: bool, } @@ -58,11 +59,11 @@ impl HostEndpoint { pub(crate) fn connection_info(&self) -> RedisResult { let settings = RedisConnectionInfo::default().set_db(self.db); let settings = match &self.username { - Some(username) => settings.set_username(username), + Some(username) => settings.set_username(username.expose()), None => settings, }; let settings = match &self.password { - Some(password) => settings.set_password(password), + Some(password) => settings.set_password(password.expose()), None => settings, }; Ok(addr(&self.host, self.port, self.tls) @@ -83,70 +84,35 @@ fn addr(host: &str, port: u16, tls: bool) -> ConnectionAddr { } } -impl fmt::Debug for Endpoint { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Url(url) => formatter - .debug_tuple("Url") - .field(&redacted_url(url)) - .finish(), - Self::Host(host) => formatter - .debug_struct("Host") - .field("host", &host.host) - .field("port", &host.port) - .field("username", &host.username) - .field("password", &host.password.as_ref().map(|_| "")) - .field("tls", &host.tls) - .field("db", &host.db) - .finish(), - Self::Cluster(cluster) => formatter - .debug_struct("Cluster") - .field("nodes", &cluster.nodes) - .field("username", &cluster.username) - .field("password", &cluster.password.as_ref().map(|_| "")) - .field("tls", &cluster.tls) - .finish(), - } - } -} - -fn redacted_url(url: &str) -> String { - match url - .split_once("://") - .and_then(|(scheme, rest)| Some((scheme, rest.rsplit_once('@')?.1))) - { - Some((scheme, host)) => format!("{scheme}://@{host}"), - None => url.to_owned(), - } -} - #[cfg(test)] mod tests { + use litellm_secret_value::SecretValue; use rstest::rstest; - use super::{Endpoint, HostEndpoint}; + use super::{ClusterEndpoint, Endpoint, HostEndpoint}; #[rstest] - #[case::url_with_password(Endpoint::Url("redis://user:secret@cache:6379/0".to_owned()), "redis://@cache:6379/0")] - #[case::url_with_at_in_password(Endpoint::Url("rediss://:p@ss@cache:6380".to_owned()), "rediss://@cache:6380")] - #[case::host_password( + #[case::url(Endpoint::Url(SecretValue::new("redis://credential-value@cache:6379/0")))] + #[case::host( Endpoint::Host(HostEndpoint { host: "cache".to_owned(), port: 6379, - username: None, - password: Some("secret".to_owned()), + username: Some(SecretValue::new("credential-value")), + password: Some(SecretValue::new("credential-value")), tls: false, db: 0, - }), - "" + }) )] - fn debug_output_never_shows_the_password(#[case] endpoint: Endpoint, #[case] shown: &str) { + #[case::cluster(Endpoint::Cluster(ClusterEndpoint { + nodes: Vec::new(), + username: Some(SecretValue::new("credential-value")), + password: Some(SecretValue::new("credential-value")), + tls: false, + }))] + fn debug_output_never_shows_the_secret(#[case] endpoint: Endpoint) { let debug = format!("{endpoint:?}"); - assert!( - !debug.contains("secret") && !debug.contains("p@ss"), - "{debug}" - ); - assert!(debug.contains(shown), "{debug}"); + assert!(!debug.contains("credential-value"), "{debug}"); + assert!(debug.contains("[REDACTED]"), "{debug}"); } } diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs index aaf26f105da..41729bfd107 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -33,11 +33,11 @@ impl Connection { .response_timeout(timeout) .retries(RETRIES); let builder = match &endpoint.username { - Some(username) => builder.username(username), + Some(username) => builder.username(username.expose()), None => builder, }; let builder = match &endpoint.password { - Some(password) => builder.password(password), + Some(password) => builder.password(password.expose()), None => builder, }; let config = ClusterConfig::new() diff --git a/litellm-rust/crates/redis-client/src/test_support.rs b/litellm-rust/crates/redis-client/src/test_support.rs index 64a174b1543..e503739a942 100644 --- a/litellm-rust/crates/redis-client/src/test_support.rs +++ b/litellm-rust/crates/redis-client/src/test_support.rs @@ -2,6 +2,7 @@ use std::time::Duration; use litellm_circuit_breaker::Breaker; use litellm_redis_testing::{RedisCluster, RedisServer}; +use litellm_secret_value::SecretValue; use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr}; use crate::{ @@ -30,7 +31,7 @@ impl Server { pub(crate) fn endpoint(&self) -> Endpoint { match self { - Self::Standalone(server) => Endpoint::Url(server.url().to_owned()), + Self::Standalone(server) => Endpoint::Url(SecretValue::new(server.url().to_owned())), Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { nodes: cluster .nodes() diff --git a/litellm-rust/crates/router-types/Cargo.toml b/litellm-rust/crates/router-types/Cargo.toml index 2e9fec28685..f0470ce8a8e 100644 --- a/litellm-rust/crates/router-types/Cargo.toml +++ b/litellm-rust/crates/router-types/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true description = "The types of litellm/types/router.py: a deployment's litellm_params as a config or a caller spells them" [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/router-types/src/litellm_params.rs b/litellm-rust/crates/router-types/src/litellm_params.rs index 186a1d98036..aa2c47f6e3c 100644 --- a/litellm-rust/crates/router-types/src/litellm_params.rs +++ b/litellm-rust/crates/router-types/src/litellm_params.rs @@ -1,6 +1,8 @@ use std::{collections::BTreeMap, fmt, ops::Deref}; -use litellm_auth_types::{AwsParams, ParamSpec, SecretValue, VertexParams}; +use litellm_auth_types::{AwsParams, ParamSpec, VertexParams}; +use litellm_secret_value::SecretValue; + use serde::Deserialize; use serde_json::Value; use serde_with::{OneOrMany, formats::PreferMany, serde_as}; diff --git a/litellm-rust/crates/router-types/tests/litellm_params.rs b/litellm-rust/crates/router-types/tests/litellm_params.rs index 7d021fd857a..c5c0d16ee8f 100644 --- a/litellm-rust/crates/router-types/tests/litellm_params.rs +++ b/litellm-rust/crates/router-types/tests/litellm_params.rs @@ -1,4 +1,6 @@ -use litellm_auth_types::{AwsParams, SecretValue, VertexParams}; +use litellm_auth_types::{AwsParams, VertexParams}; +use litellm_secret_value::SecretValue; + use litellm_router_types::{LitellmParams, Spelled}; use rstest::rstest; use serde_json::{Map, Value, json}; diff --git a/litellm-rust/crates/secret-value/Cargo.toml b/litellm-rust/crates/secret-value/Cargo.toml new file mode 100644 index 00000000000..467f852a278 --- /dev/null +++ b/litellm-rust/crates/secret-value/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "litellm-secret-value" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +serde.workspace = true +subtle.workspace = true +veil.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/secret-value/src/lib.rs b/litellm-rust/crates/secret-value/src/lib.rs new file mode 100644 index 00000000000..f11d22d3701 --- /dev/null +++ b/litellm-rust/crates/secret-value/src/lib.rs @@ -0,0 +1,5 @@ +#![forbid(unsafe_code)] + +mod secret_value; + +pub use secret_value::SecretValue; diff --git a/litellm-rust/crates/auth-types/src/secret.rs b/litellm-rust/crates/secret-value/src/secret_value.rs similarity index 95% rename from litellm-rust/crates/auth-types/src/secret.rs rename to litellm-rust/crates/secret-value/src/secret_value.rs index 7e6789deef7..27e4227080d 100644 --- a/litellm-rust/crates/auth-types/src/secret.rs +++ b/litellm-rust/crates/secret-value/src/secret_value.rs @@ -32,9 +32,11 @@ impl Hash for SecretValue { #[cfg(test)] mod tests { + use rstest::rstest; + use super::SecretValue; - #[test] + #[rstest] fn debug_redacts_plaintext() { let debug = format!("{:?}", SecretValue::new("credential-value")); @@ -42,7 +44,7 @@ mod tests { assert!(debug.contains("REDACTED")); } - #[test] + #[rstest] fn equality_compares_plaintext_values() { assert_eq!(SecretValue::new("same"), SecretValue::new("same")); assert_ne!(SecretValue::new("same"), SecretValue::new("different")); diff --git a/litellm-rust/crates/secrets-aws/Cargo.toml b/litellm-rust/crates/secrets-aws/Cargo.toml index 5d3bd413484..d17dfcfeec2 100644 --- a/litellm-rust/crates/secrets-aws/Cargo.toml +++ b/litellm-rust/crates/secrets-aws/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth-aws.workspace = true litellm-secrets-types.workspace = true litellm-core-utils.workspace = true diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager.rs b/litellm-rust/crates/secrets-aws/src/secret_manager.rs index 2220e387838..b25872ab3f6 100644 --- a/litellm-rust/crates/secrets-aws/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/src/secret_manager.rs @@ -21,11 +21,12 @@ use litellm_auth_aws::constants::{ AWS_ACCESS_KEY_ID, AWS_REGION, AWS_REGION_NAME, AWS_SECRET_ACCESS_KEY, }; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ AwsOperationContext, BaseSecretManager, KeyManagementSettings, RotationError, Secret, - SecretDeleter, SecretRotator, SecretValue, SecretWriteContext, SecretWriter, - async_rotate_secret, + SecretDeleter, SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, }; + use serde_json::Value; use crate::{Error, auth}; diff --git a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs index c482a168090..56999633a78 100644 --- a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs @@ -10,11 +10,13 @@ use aws_sdk_secretsmanager::{ Client, config::{BehaviorVersion, Credentials, Region, retry::RetryConfig}, }; +use litellm_secret_value::SecretValue; use litellm_secrets_aws::{AwsSecretsManagerV2, Error, RotationResponse}; use litellm_secrets_types::{ AwsOperationContext, BaseSecretManager, KeyManagementSettings, Secret, SecretDeleter, - SecretValue, SecretWriteContext, SecretWriter, + SecretWriteContext, SecretWriter, }; + use rstest::{fixture, rstest}; use serde_json::json; use wiremock::{ diff --git a/litellm-rust/crates/secrets-azure/Cargo.toml b/litellm-rust/crates/secrets-azure/Cargo.toml index 2aff6129c37..c4bfb24898c 100644 --- a/litellm-rust/crates/secrets-azure/Cargo.toml +++ b/litellm-rust/crates/secrets-azure/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-http.workspace = true tokio.workspace = true litellm-auth-azure.workspace = true diff --git a/litellm-rust/crates/secrets-azure/src/key_vault.rs b/litellm-rust/crates/secrets-azure/src/key_vault.rs index d97b7a22cfc..40028828298 100644 --- a/litellm-rust/crates/secrets-azure/src/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/src/key_vault.rs @@ -3,7 +3,9 @@ use std::sync::Arc; use litellm_auth_azure::{AzureAuthInputs, AzureAuthService, ConfigValue}; use litellm_auth_types::{InputSource, Sourced}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::{AzureOperationContext, BaseSecretManager, Secret, SecretValue}; +use litellm_secret_value::SecretValue; +use litellm_secrets_types::{AzureOperationContext, BaseSecretManager, Secret}; + use percent_encoding::{AsciiSet, NON_ALPHANUMERIC}; use serde::Deserialize; diff --git a/litellm-rust/crates/secrets-azure/tests/key_vault.rs b/litellm-rust/crates/secrets-azure/tests/key_vault.rs index a0ba9e6ddf8..c3ba01672a8 100644 --- a/litellm-rust/crates/secrets-azure/tests/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/tests/key_vault.rs @@ -1,7 +1,9 @@ use std::sync::Arc; +use litellm_secret_value::SecretValue; use litellm_secrets_azure::{AzureKeyVault, Error}; -use litellm_secrets_types::{Secret, SecretValue}; +use litellm_secrets_types::Secret; + use rstest::{fixture, rstest}; use serde::Deserialize; use wiremock::{ diff --git a/litellm-rust/crates/secrets-cyberark/Cargo.toml b/litellm-rust/crates/secrets-cyberark/Cargo.toml index 2ae2c1ea5c3..643f322f7fa 100644 --- a/litellm-rust/crates/secrets-cyberark/Cargo.toml +++ b/litellm-rust/crates/secrets-cyberark/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-secrets-types.workspace = true litellm-core-utils.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs index c5cb0308b13..cf077a2bbc8 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -9,11 +9,12 @@ use litellm_core_utils::settings::Lookup; use litellm_http::{ Client, ClientIdentity, ClientVariant, HttpClientConfig, HttpClientPool, TlsSource, Verify, }; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ BaseSecretManager, CyberarkOperationContext, RotationError, SecretCache, SecretDeleter, - SecretRotator, SecretValue, SecretWriteContext, SecretWriter, async_rotate_secret, - validate_secret_name, + SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, validate_secret_name, }; + use moka::future::Cache; use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode}; diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs index 783ba6ff67a..e9b480e026b 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs @@ -9,8 +9,10 @@ use std::{ use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_core_utils::settings::Lookup; use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; +use litellm_secret_value::SecretValue; use litellm_secrets_cyberark::{CyberArkSecretManager, DeleteOutcome, Error}; -use litellm_secrets_types::{BaseSecretManager, CyberarkOperationContext, SecretValue}; +use litellm_secrets_types::{BaseSecretManager, CyberarkOperationContext}; + use rstest::{fixture, rstest}; use serde::Deserialize; use wiremock::{ diff --git a/litellm-rust/crates/secrets-google/Cargo.toml b/litellm-rust/crates/secrets-google/Cargo.toml index 06b3747def5..3c5c2130788 100644 --- a/litellm-rust/crates/secrets-google/Cargo.toml +++ b/litellm-rust/crates/secrets-google/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-http.workspace = true moka.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/secrets-google/src/auth.rs b/litellm-rust/crates/secrets-google/src/auth.rs index 45fc99d8d5f..7022ecc5ae4 100644 --- a/litellm-rust/crates/secrets-google/src/auth.rs +++ b/litellm-rust/crates/secrets-google/src/auth.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use litellm_auth_gcp::{GoogleCredentials, VertexConfig}; use litellm_auth_types::{InputSource, Sourced}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::SecretValue; +use litellm_secret_value::SecretValue; pub(crate) fn credentials( project: Option, diff --git a/litellm-rust/crates/secrets-google/src/kms.rs b/litellm-rust/crates/secrets-google/src/kms.rs index 78ebae008f7..5622a3a4d56 100644 --- a/litellm-rust/crates/secrets-google/src/kms.rs +++ b/litellm-rust/crates/secrets-google/src/kms.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use google_cloud_kms_v1::client::KeyManagementService; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{Error, auth}; diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs index 463d1f6fd73..577a2ccbd4d 100644 --- a/litellm-rust/crates/secrets-google/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -2,9 +2,9 @@ use std::{sync::Arc, time::Duration}; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::{ - BaseSecretManager, GoogleOperationContext, Secret, SecretCache, SecretValue, -}; +use litellm_secret_value::SecretValue; +use litellm_secrets_types::{BaseSecretManager, GoogleOperationContext, Secret, SecretCache}; + use serde::Deserialize; use litellm_auth_gcp::GoogleCredentials; diff --git a/litellm-rust/crates/secrets-hashicorp/Cargo.toml b/litellm-rust/crates/secrets-hashicorp/Cargo.toml index 8731a8deb7a..4fcd496931d 100644 --- a/litellm-rust/crates/secrets-hashicorp/Cargo.toml +++ b/litellm-rust/crates/secrets-hashicorp/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-error.workspace = true litellm-core-utils.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/secrets-hashicorp/src/config.rs b/litellm-rust/crates/secrets-hashicorp/src/config.rs index 4232667ae16..a63161b17b2 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/config.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/config.rs @@ -1,7 +1,7 @@ use std::{path::PathBuf, time::Duration}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::Error; diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index 447beacda85..dae14439630 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -12,11 +12,12 @@ use std::{ }; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ BaseSecretManager, HashicorpOperationContext, RotationError, SecretCache, SecretDeleter, - SecretRotator, SecretValue, SecretWriteContext, SecretWriter, async_rotate_secret, - validate_secret_name, + SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, validate_secret_name, }; + use rustify::errors::ClientError as RustifyClientError; use serde_json::Value; use tokio::sync::Mutex; diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs index 668871d3fc7..d431e472400 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs @@ -1,11 +1,13 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets_hashicorp::{Error, HashicorpVault, HashicorpVaultConfig}; use litellm_secrets_types::{ - BaseSecretManager, HashicorpOperationContext, RotationError, SecretDeleter, SecretValue, - SecretWriteContext, SecretWriter, + BaseSecretManager, HashicorpOperationContext, RotationError, SecretDeleter, SecretWriteContext, + SecretWriter, }; + use rstest::{fixture, rstest}; use serde::Deserialize; use serde_json::json; diff --git a/litellm-rust/crates/secrets-types/Cargo.toml b/litellm-rust/crates/secrets-types/Cargo.toml index 6dd847ec989..b45c41bcf1c 100644 --- a/litellm-rust/crates/secrets-types/Cargo.toml +++ b/litellm-rust/crates/secrets-types/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth-types.workspace = true +litellm-secret-value.workspace = true moka.workspace = true tokio = { workspace = true, features = ["sync"] } serde.workspace = true diff --git a/litellm-rust/crates/secrets-types/src/base_secret_manager.rs b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs index 737fad28662..41dadd1baa6 100644 --- a/litellm-rust/crates/secrets-types/src/base_secret_manager.rs +++ b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs @@ -1,4 +1,5 @@ -use crate::{Error, SecretValue, SecretWriteContext}; +use crate::{Error, SecretWriteContext}; +use litellm_secret_value::SecretValue; pub fn validate_secret_name(name: &str) -> Result<(), Error> { if name.split('/').any(|segment| segment == "..") diff --git a/litellm-rust/crates/secrets-types/src/config.rs b/litellm-rust/crates/secrets-types/src/config.rs index 9b6e82c3a6b..a45acb4df11 100644 --- a/litellm-rust/crates/secrets-types/src/config.rs +++ b/litellm-rust/crates/secrets-types/src/config.rs @@ -3,7 +3,7 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; use strum::IntoStaticStr; -use crate::SecretValue; +use litellm_secret_value::SecretValue; #[derive( Clone, diff --git a/litellm-rust/crates/secrets-types/src/context.rs b/litellm-rust/crates/secrets-types/src/context.rs index 1ec493edf42..b71195e2d21 100644 --- a/litellm-rust/crates/secrets-types/src/context.rs +++ b/litellm-rust/crates/secrets-types/src/context.rs @@ -1,6 +1,7 @@ use std::{collections::BTreeMap, time::Duration}; -use crate::{Error, KeyManagementSystem, SecretValue}; +use crate::{Error, KeyManagementSystem}; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, Default, Eq, PartialEq)] pub enum SecretOperationContext { diff --git a/litellm-rust/crates/secrets-types/src/lib.rs b/litellm-rust/crates/secrets-types/src/lib.rs index 4d7eaf8313e..2563ccbe39d 100644 --- a/litellm-rust/crates/secrets-types/src/lib.rs +++ b/litellm-rust/crates/secrets-types/src/lib.rs @@ -18,5 +18,4 @@ pub use context::{ HashicorpOperationContext, SecretOperationContext, SecretWriteContext, }; pub use error::Error; -pub use litellm_auth_types::SecretValue; pub use value::{PythonSecretRead, Secret}; diff --git a/litellm-rust/crates/secrets-types/src/value.rs b/litellm-rust/crates/secrets-types/src/value.rs index 11f57ef4448..8c92902edad 100644 --- a/litellm-rust/crates/secrets-types/src/value.rs +++ b/litellm-rust/crates/secrets-types/src/value.rs @@ -1,4 +1,4 @@ -use crate::SecretValue; +use litellm_secret_value::SecretValue; #[derive(Clone, PartialEq, Eq, veil::Redact)] pub enum Secret { diff --git a/litellm-rust/crates/secrets-types/tests/config.rs b/litellm-rust/crates/secrets-types/tests/config.rs index 45fd9d0ee03..e7ba63343da 100644 --- a/litellm-rust/crates/secrets-types/tests/config.rs +++ b/litellm-rust/crates/secrets-types/tests/config.rs @@ -1,6 +1,6 @@ -use litellm_secrets_types::{ - AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, -}; +use litellm_secret_value::SecretValue; +use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, Secret}; + use rstest::{fixture, rstest}; use serde_json::json; diff --git a/litellm-rust/crates/secrets-types/tests/context.rs b/litellm-rust/crates/secrets-types/tests/context.rs index 8f8eeabdb62..f8d754564dc 100644 --- a/litellm-rust/crates/secrets-types/tests/context.rs +++ b/litellm-rust/crates/secrets-types/tests/context.rs @@ -1,10 +1,11 @@ use std::{collections::BTreeMap, time::Duration}; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ AwsOperationContext, AzureOperationContext, CyberarkOperationContext, GoogleOperationContext, - HashicorpOperationContext, KeyManagementSystem, SecretOperationContext, SecretValue, - SecretWriteContext, + HashicorpOperationContext, KeyManagementSystem, SecretOperationContext, SecretWriteContext, }; + use rstest::{fixture, rstest}; #[fixture] diff --git a/litellm-rust/crates/secrets-types/tests/rotation.rs b/litellm-rust/crates/secrets-types/tests/rotation.rs index 36c1a065ffa..089890c60bd 100644 --- a/litellm-rust/crates/secrets-types/tests/rotation.rs +++ b/litellm-rust/crates/secrets-types/tests/rotation.rs @@ -1,10 +1,12 @@ use std::sync::atomic::{AtomicUsize, Ordering}; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ BaseSecretManager, Error, HashicorpOperationContext, RotationError, SecretDeleter, - SecretOperationContext, SecretRotator, SecretValue, SecretWriteContext, SecretWriter, - async_rotate_secret, validate_secret_name, + SecretOperationContext, SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, + validate_secret_name, }; + use rstest::{fixture, rstest}; struct Manager { diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index 023d61136e4..21ee1bdc760 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -14,6 +14,7 @@ azure = ["dep:litellm-secrets-azure"] cyberark = ["dep:litellm-secrets-cyberark"] [dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true litellm-python-compat.workspace = true litellm-secrets-types.workspace = true diff --git a/litellm-rust/crates/secrets/src/compatibility.rs b/litellm-rust/crates/secrets/src/compatibility.rs index 368c9931b51..3d6ba854325 100644 --- a/litellm-rust/crates/secrets/src/compatibility.rs +++ b/litellm-rust/crates/secrets/src/compatibility.rs @@ -3,9 +3,10 @@ use litellm_python_compat::{Value, literal::literal_eval}; use litellm_secrets_types::PythonSecretRead; use crate::{ - Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretManager, SecretValue, + Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretManager, get_secret_from_manager, }; +use litellm_secret_value::SecretValue; pub async fn get_secret_from_python_manager( manager: &SecretManager, diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index f42e09c6bdb..2d2b81d197f 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -5,7 +5,7 @@ use litellm_core_utils::settings::Lookup; use crate::{Error, KeyManagementSettings, KeyManagementSystem, Secret}; #[cfg(any(feature = "aws", feature = "google"))] -use crate::SecretValue; +use litellm_secret_value::SecretValue; #[cfg(any( feature = "google", diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index 8d0f0561608..e4459576e94 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -12,9 +12,7 @@ mod state; pub use compatibility::{get_secret_from_python_manager, read_secret_from_python_manager}; pub use error::Error; pub use handler::{ExternalSecretManager, SecretManager, get_secret_from_manager}; -pub use litellm_secrets_types::{ - AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, -}; +pub use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, Secret}; pub use native::load_native_manager; pub use oidc::{OidcProvider, OidcReference, OidcResolver}; pub use resolver::{FailurePolicy, SecretResolver, normalize_nonempty_secret_str}; diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs index 24e6912c797..18fc9479ef1 100644 --- a/litellm-rust/crates/secrets/src/oidc.rs +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -9,7 +9,8 @@ use litellm_http::Client; use moka::future::Cache; use serde::Deserialize; -use crate::{Error, SecretValue}; +use crate::Error; +use litellm_secret_value::SecretValue; const GOOGLE_TOKEN_MAX_TTL: Duration = Duration::from_secs(3540); const GITHUB_TOKEN_TTL: Duration = Duration::from_secs(295); diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs index 830e7c22cd5..631b60cba10 100644 --- a/litellm-rust/crates/secrets/src/resolver.rs +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -4,7 +4,8 @@ use crate::compatibility::python_manager_string; use litellm_core_utils::{serde_compat::parse_str_bool, settings::Lookup}; use crate::state::{LookupTarget, normalize_secret_name}; -use crate::{Error, OidcResolver, Secret, SecretManagerState, SecretValue}; +use crate::{Error, OidcResolver, Secret, SecretManagerState}; +use litellm_secret_value::SecretValue; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum FailurePolicy { diff --git a/litellm-rust/crates/secrets/src/source.rs b/litellm-rust/crates/secrets/src/source.rs index 3c86fd25e9a..1ee224b763a 100644 --- a/litellm-rust/crates/secrets/src/source.rs +++ b/litellm-rust/crates/secrets/src/source.rs @@ -3,7 +3,8 @@ use std::{collections::HashMap, sync::Arc}; use futures_util::future::{BoxFuture, try_join_all}; use litellm_core_utils::settings::Lookup; -use crate::{Error, SecretResolver, SecretValue}; +use crate::{Error, SecretResolver}; +use litellm_secret_value::SecretValue; pub type Secrets = Arc; diff --git a/litellm-rust/crates/secrets/tests/aws.rs b/litellm-rust/crates/secrets/tests/aws.rs index abad04f31c5..6f30d1b5161 100644 --- a/litellm-rust/crates/secrets/tests/aws.rs +++ b/litellm-rust/crates/secrets/tests/aws.rs @@ -2,11 +2,13 @@ use std::sync::Arc; +use litellm_secret_value::SecretValue; use litellm_secrets::{ AccessMode, Error, FailurePolicy, KeyManagementSettings, OidcResolver, Secret, SecretManager, - SecretManagerState, SecretResolver, SecretValue, aws::AwsSecretsManagerV2, + SecretManagerState, SecretResolver, aws::AwsSecretsManagerV2, secret_manager_would_be_consulted, }; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; fn state(server: &MockServer, settings: KeyManagementSettings) -> SecretManagerState { diff --git a/litellm-rust/crates/secrets/tests/azure.rs b/litellm-rust/crates/secrets/tests/azure.rs index 6e81017b0f6..201e31f44b8 100644 --- a/litellm-rust/crates/secrets/tests/azure.rs +++ b/litellm-rust/crates/secrets/tests/azure.rs @@ -68,10 +68,11 @@ async fn azure_handler_reads_missing_and_failed_secrets() { async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_null( #[case] body: serde_json::Value, ) { + use litellm_secret_value::SecretValue; use litellm_secrets::{ - OidcResolver, SecretManager, SecretManagerState, SecretResolver, SecretValue, - azure::AzureKeyVault, + OidcResolver, SecretManager, SecretManagerState, SecretResolver, azure::AzureKeyVault, }; + use std::sync::Arc; use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; let server = MockServer::start().await; diff --git a/litellm-rust/crates/secrets/tests/common_read_contract.rs b/litellm-rust/crates/secrets/tests/common_read_contract.rs index e6928038799..4cf3513fadf 100644 --- a/litellm-rust/crates/secrets/tests/common_read_contract.rs +++ b/litellm-rust/crates/secrets/tests/common_read_contract.rs @@ -9,8 +9,9 @@ use std::sync::Arc; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets::{ - KeyManagementSettings, SecretManager, SecretValue, + KeyManagementSettings, SecretManager, aws::AwsSecretsManagerV2, azure::AzureKeyVault, cyberark::CyberArkSecretManager, @@ -18,6 +19,7 @@ use litellm_secrets::{ google::GoogleSecretManager, hashicorp::{HashicorpVault, HashicorpVaultConfig}, }; + use rstest::rstest; use serde_json::json; use wiremock::{ diff --git a/litellm-rust/crates/secrets/tests/cyberark.rs b/litellm-rust/crates/secrets/tests/cyberark.rs index 0a6c394937e..09342234118 100644 --- a/litellm-rust/crates/secrets/tests/cyberark.rs +++ b/litellm-rust/crates/secrets/tests/cyberark.rs @@ -4,10 +4,12 @@ async fn cyberark_handler_reads_values_and_surfaces_errors() { use std::time::Duration; + use litellm_secret_value::SecretValue; use litellm_secrets::{ - Error, KeyManagementSettings, SecretManager, SecretValue, cyberark::CyberArkSecretManager, + Error, KeyManagementSettings, SecretManager, cyberark::CyberArkSecretManager, get_secret_from_manager, }; + use wiremock::{ Mock, MockServer, ResponseTemplate, matchers::{body_string, path}, diff --git a/litellm-rust/crates/secrets/tests/google.rs b/litellm-rust/crates/secrets/tests/google.rs index 5542c864dfd..a75e34335de 100644 --- a/litellm-rust/crates/secrets/tests/google.rs +++ b/litellm-rust/crates/secrets/tests/google.rs @@ -7,10 +7,12 @@ use std::sync::Arc; #[case::failure(503)] #[tokio::test] async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) { + use litellm_secret_value::SecretValue; use litellm_secrets::{ Error, FailurePolicy, KeyManagementSettings, OidcResolver, SecretManager, - SecretManagerState, SecretResolver, SecretValue, google::GoogleSecretManager, + SecretManagerState, SecretResolver, google::GoogleSecretManager, }; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; let server = MockServer::start().await; Mock::given(method("GET")) diff --git a/litellm-rust/crates/secrets/tests/oidc.rs b/litellm-rust/crates/secrets/tests/oidc.rs index 867a662cbbb..2d250e1de9d 100644 --- a/litellm-rust/crates/secrets/tests/oidc.rs +++ b/litellm-rust/crates/secrets/tests/oidc.rs @@ -315,7 +315,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] Box< dyn std::future::Future< Output = Result< - litellm_secrets::SecretValue, + litellm_secret_value::SecretValue, litellm_secrets::azure::Error, >, > + Send @@ -331,7 +331,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] if self.0 { Err(litellm_secrets::azure::Error::MissingCredentials) } else { - Ok(litellm_secrets::SecretValue::new("azure-token")) + Ok(litellm_secret_value::SecretValue::new("azure-token")) } }) } @@ -346,7 +346,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] let result = resolver .get_secret_str( "oidc/azure/api://audience/path", - Some(litellm_secrets::SecretValue::new("fallback")), + Some(litellm_secret_value::SecretValue::new("fallback")), ) .await; if failed { diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs index 8435e051315..185f2256780 100644 --- a/litellm-rust/crates/secrets/tests/resolution.rs +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -1,9 +1,10 @@ use std::{future::Future, pin::Pin, sync::Arc}; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets::{ Error, ExternalSecretManager, FailurePolicy, KeyManagementSettings, KeyManagementSystem, - OidcResolver, Secret, SecretManager, SecretManagerState, SecretResolver, SecretValue, + OidcResolver, Secret, SecretManager, SecretManagerState, SecretResolver, secret_manager_would_be_consulted, }; From 24b8889965a6f86bee6f743d61c77da0896eaa15 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:35:18 -0700 Subject: [PATCH 45/64] chore(cost-map): add openai gpt-6-rosalind from the pricing page and changelog (#49268) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 56 +++++++++++++++++++ model_prices_and_context_window.json | 56 +++++++++++++++++++ 2 files changed, 112 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ea9cacbdabe..dac7af6ab7a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -81156,5 +81156,61 @@ "output_cost_per_token": 0.0, "source": "https://inference.baseten.co/v1/models", "supports_vision": true + }, + "gpt-6-rosalind": { + "cache_creation_input_token_cost": 1.875e-05, + "cache_creation_input_token_cost_above_12hr": 3.75e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 7.5e-05, + "cache_creation_input_token_cost_above_272k_tokens": 3.75e-05, + "cache_read_input_token_cost": 1.5e-06, + "cache_read_input_token_cost_above_272k_tokens": 3e-06, + "input_cost_per_token": 1.5e-05, + "input_cost_per_token_above_272k_tokens": 3e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "output_cost_per_token_above_272k_tokens": 0.0001125, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "source": "https://developers.openai.com/api/docs/pricing", + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6-rosalind": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ea9cacbdabe..dac7af6ab7a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -81156,5 +81156,61 @@ "output_cost_per_token": 0.0, "source": "https://inference.baseten.co/v1/models", "supports_vision": true + }, + "gpt-6-rosalind": { + "cache_creation_input_token_cost": 1.875e-05, + "cache_creation_input_token_cost_above_12hr": 3.75e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 7.5e-05, + "cache_creation_input_token_cost_above_272k_tokens": 3.75e-05, + "cache_read_input_token_cost": 1.5e-06, + "cache_read_input_token_cost_above_272k_tokens": 3e-06, + "input_cost_per_token": 1.5e-05, + "input_cost_per_token_above_272k_tokens": 3e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "output_cost_per_token_above_272k_tokens": 0.0001125, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "source": "https://developers.openai.com/api/docs/pricing", + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6-rosalind": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true } } From 1ccb970d520aa24a078d5b2a40ccc7f5249a770c Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 30 Oct 2026 11:43:13 -0700 Subject: [PATCH 46/64] refactor(ui): remove the deprecated Old Usage page (#49284) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/ui/fixtures/menuMappings.ts | 1 - tests/e2e/ui/fixtures/migratedPages.ts | 6 - tests/e2e/ui/fixtures/pages.ts | 1 - .../e2e/ui/tests/logs/logsPagination.spec.ts | 2 +- tests/e2e/ui/tests/usage/usagePage.spec.ts | 2 +- ui/litellm-dashboard/eslint-suppressions.json | 28 - .../app/(dashboard)/legacyPageRoutes.test.ts | 2 +- .../src/app/(dashboard)/legacyPageRoutes.ts | 2 +- .../old-usage/_components/usage.test.tsx | 257 ----- .../old-usage/_components/usage.tsx | 950 ------------------ .../src/app/(dashboard)/old-usage/page.tsx | 22 - .../components/UsagePageView.test.tsx | 4 - .../src/components/leftnav.test.tsx | 26 - .../src/components/leftnav.tsx | 8 - .../src/components/page_metadata.ts | 1 - .../src/components/view_user_spend.tsx | 144 --- .../tests/CreateKeyPage.expiredToken.test.tsx | 1 - 17 files changed, 4 insertions(+), 1453 deletions(-) delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/old-usage/page.tsx delete mode 100644 ui/litellm-dashboard/src/components/view_user_spend.tsx diff --git a/tests/e2e/ui/fixtures/menuMappings.ts b/tests/e2e/ui/fixtures/menuMappings.ts index d6e7ea86982..4cf90c00cd0 100644 --- a/tests/e2e/ui/fixtures/menuMappings.ts +++ b/tests/e2e/ui/fixtures/menuMappings.ts @@ -32,7 +32,6 @@ export const menuLabelToPage: Record = { Budgets: Page.Budgets, "API Playground": Page.TransformRequest, "Tag Management": Page.TagManagement, - "Old Usage": Page.Usage, // Tools submenu items "MCP Servers": Page.McpServers, "Vector Stores": Page.VectorStores, diff --git a/tests/e2e/ui/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts index 5fa1f9d2444..6fa8d0c2587 100644 --- a/tests/e2e/ui/fixtures/migratedPages.ts +++ b/tests/e2e/ui/fixtures/migratedPages.ts @@ -118,12 +118,6 @@ export const MIGRATED_E2E_PAGES: Readonly> = { content: { role: "heading", name: "AI Hub" }, }, new_usage: { segment: "usage", linkName: "Usage", content: { role: "heading", name: "Usage" } }, - usage: { - segment: "old-usage", - linkName: "Old Usage", - group: "Experimental", - content: { role: "tab", name: "All Up" }, - }, agents: { segment: "agents", linkName: "Agents", group: "Agentic", content: { role: "heading", name: "Agents" } }, "router-settings": { segment: "router-settings", diff --git a/tests/e2e/ui/fixtures/pages.ts b/tests/e2e/ui/fixtures/pages.ts index 8210334c166..a850dcb4755 100644 --- a/tests/e2e/ui/fixtures/pages.ts +++ b/tests/e2e/ui/fixtures/pages.ts @@ -30,5 +30,4 @@ export enum Page { TagManagement = "tag-management", VectorStores = "vector-stores", NewUsage = "new_usage", - Usage = "usage", } diff --git a/tests/e2e/ui/tests/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts index 585eb31c37c..a611f760e8b 100644 --- a/tests/e2e/ui/tests/logs/logsPagination.spec.ts +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -84,7 +84,7 @@ test.describe("Logs pagination with Aria in the header", () => { await expect(aria).toBeFocused(); await visibleTestId(page, "pagination-first").click(); await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); - for (const destination of [Page.Usage, Page.LlmPlayground, Page.ApiKeys]) { + for (const destination of [Page.NewUsage, Page.LlmPlayground, Page.ApiKeys]) { await navigateToPage(page, destination); await aria.click(); await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); diff --git a/tests/e2e/ui/tests/usage/usagePage.spec.ts b/tests/e2e/ui/tests/usage/usagePage.spec.ts index d532042c7af..2a58a1ef59d 100644 --- a/tests/e2e/ui/tests/usage/usagePage.spec.ts +++ b/tests/e2e/ui/tests/usage/usagePage.spec.ts @@ -11,7 +11,7 @@ import { waitForSpendLog, } from "../../helpers/traffic"; -/** Covers /ui/usage. The legacy /ui/old-usage view is deprecated and deliberately not covered. */ +/** Covers /ui/usage. */ const topKeysCard = (page: PlaywrightPage): Locator => page.locator("section").filter({ has: page.getByRole("heading", { name: "Top Virtual Keys", exact: true }) }); diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 9fce93782b2..c14f3e6f036 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -625,23 +625,6 @@ "count": 1 } }, - "src/app/(dashboard)/old-usage/_components/usage.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "max-lines": { - "count": 1 - }, - "prefer-const": { - "count": 6 - }, - "react-hooks/immutability": { - "count": 1 - }, - "react-hooks/purity": { - "count": 1 - } - }, "src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx": { "react-hooks/set-state-in-effect": { "count": 2 @@ -2090,17 +2073,6 @@ "count": 1 } }, - "src/components/view_user_spend.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "prefer-const": { - "count": 3 - }, - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, "src/contexts/AuthContext.tsx": { "react-hooks/set-state-in-effect": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.test.ts index 6b02b79b04c..315d988db54 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.test.ts @@ -10,7 +10,7 @@ describe("legacyPageRedirectHref", () => { expect(redirect("page=models")).toBe("/ui/models-and-endpoints"); expect(redirect("page=llm-playground")).toBe("/ui/playground"); expect(redirect("page=new_usage")).toBe("/ui/usage"); - expect(redirect("page=usage")).toBe("/ui/old-usage"); + expect(redirect("page=usage")).toBe("/ui/usage"); }); it("keeps the older aliases for renamed pages", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts index e3025266ddb..fe054c7b39a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts @@ -36,7 +36,7 @@ const LEGACY_PAGE_ROUTES: ReadonlyMap = new Map( "logging-and-alerts": "logging-and-alerts", "model-hub-table": "model-hub-table", new_usage: "usage", - usage: "old-usage", + usage: "usage", "cost-optimization": "cost-optimization", "model-insights": "model-insights", agents: "agents", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx deleted file mode 100644 index 0e4455ee912..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx +++ /dev/null @@ -1,257 +0,0 @@ -import React from "react"; -import { describe, it, expect, vi, beforeEach } from "vitest"; -import { act, screen, waitFor, within } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../../../tests/test-utils"; -import UsagePage from "./usage"; - -const networking = vi.hoisted(() => ({ - adminSpendLogsCall: vi.fn(), - adminTopKeysCall: vi.fn(), - adminTopModelsCall: vi.fn(), - adminTopEndUsersCall: vi.fn(), - teamSpendLogsCall: vi.fn(), - tagsSpendLogsCall: vi.fn(), - allTagNamesCall: vi.fn(), - adminspendByProvider: vi.fn(), - adminGlobalActivity: vi.fn(), - adminGlobalActivityPerModel: vi.fn(), - getProxyUISettings: vi.fn(), - modelAvailableCall: vi.fn(), - keyInfoV1Call: vi.fn(), -})); - -vi.mock("@/components/networking", () => networking); -vi.mock("../../../../components/networking", () => networking); - -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => ({ - accessToken: "sk-test", - token: "tok", - userRole: "Admin", - userId: "u1", - premiumUser: true, - }), -})); - -const UNLIMITED_SETTINGS = { DISABLE_EXPENSIVE_DB_QUERIES: false, NUM_SPEND_LOGS_ROWS: 10 }; - -const renderUsage = (overrides: Partial> = {}) => - renderWithProviders( - , - ); - -// Width of this window is guarded by "proves the flush window is wide enough". -const flushPendingRequests = async () => { - await act(async () => { - await new Promise((resolve) => setTimeout(resolve, 0)); - await new Promise((resolve) => setTimeout(resolve, 0)); - }); -}; - -beforeEach(() => { - vi.clearAllMocks(); - networking.getProxyUISettings.mockResolvedValue(UNLIMITED_SETTINGS); - networking.adminSpendLogsCall.mockResolvedValue([{ date: "2026-07-01", spend: 12.5 }]); - networking.adminTopKeysCall.mockResolvedValue([ - { api_key: "sk-abcdefghijk", key_alias: "prod-key", total_spend: 9.5 }, - ]); - networking.adminTopModelsCall.mockResolvedValue([{ model: "gpt-5.1", total_spend: 7.25 }]); - networking.adminTopEndUsersCall.mockResolvedValue([ - { end_user: "customer-alpha", total_spend: 3.5, total_count: 42 }, - ]); - networking.teamSpendLogsCall.mockResolvedValue({ - daily_spend: [{ date: "2026-07-01", "team-a": 5 }], - teams: ["team-a"], - total_spend_per_team: [{ team_id: "team-a", total_spend: 5 }], - }); - networking.tagsSpendLogsCall.mockResolvedValue({ spend_per_tag: [{ name: "prod", spend: 4 }] }); - networking.allTagNamesCall.mockResolvedValue({ tag_names: ["prod", "staging"] }); - networking.adminspendByProvider.mockResolvedValue([{ provider: "openai", spend: 6.75 }]); - networking.adminGlobalActivity.mockResolvedValue({ - sum_api_requests: 120, - sum_total_tokens: 4500, - daily_data: [{ date: "2026-07-01", api_requests: 120, total_tokens: 4500 }], - }); - networking.adminGlobalActivityPerModel.mockResolvedValue([]); - networking.modelAvailableCall.mockResolvedValue({ data: [] }); - networking.keyInfoV1Call.mockResolvedValue({ info: {} }); -}); - -describe("old usage page", () => { - describe("when the proxy has disabled expensive DB queries", () => { - beforeEach(() => { - networking.getProxyUISettings.mockResolvedValue({ - DISABLE_EXPENSIVE_DB_QUERIES: true, - NUM_SPEND_LOGS_ROWS: 2500000, - }); - }); - - it("shows the database query limit warning instead of the usage dashboard", async () => { - renderUsage(); - - expect(await screen.findByText("Database Query Limit Reached")).toBeInTheDocument(); - expect(screen.getByText(/SpendLogs in DB has/)).toHaveTextContent("2500000"); - expect(screen.getByText(/Please follow our guide to view usage when SpendLogs has more than 1M rows/i)); - expect(screen.queryByRole("tab", { name: "All Up" })).not.toBeInTheDocument(); - }); - - it("links to the cost tracking guide in a new tab", async () => { - renderUsage(); - - const link = await screen.findByRole("link", { name: "View Usage Guide" }); - expect(link).toHaveAttribute("href", "https://docs.litellm.ai/docs/proxy/cost_tracking"); - expect(link).toHaveAttribute("target", "_blank"); - }); - - it("skips every expensive usage query", async () => { - renderUsage(); - - await screen.findByText("Database Query Limit Reached"); - await waitFor(() => expect(networking.getProxyUISettings).toHaveBeenCalled()); - - expect(networking.adminSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.adminspendByProvider).not.toHaveBeenCalled(); - expect(networking.adminTopKeysCall).not.toHaveBeenCalled(); - expect(networking.adminTopModelsCall).not.toHaveBeenCalled(); - expect(networking.adminGlobalActivity).not.toHaveBeenCalled(); - expect(networking.adminGlobalActivityPerModel).not.toHaveBeenCalled(); - expect(networking.teamSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.adminTopEndUsersCall).not.toHaveBeenCalled(); - expect(networking.tagsSpendLogsCall).not.toHaveBeenCalled(); - }); - }); - - describe("as an admin", () => { - it("renders the admin tabs", async () => { - renderUsage(); - - expect(await screen.findByRole("tab", { name: "All Up" })).toBeInTheDocument(); - expect(screen.getByRole("tab", { name: "Team Based Usage" })).toBeInTheDocument(); - expect(screen.getByRole("tab", { name: "Customer Usage" })).toBeInTheDocument(); - expect(screen.getByRole("tab", { name: "Tag Based Usage" })).toBeInTheDocument(); - }); - - it("renders the cost panel cards", async () => { - renderUsage(); - - expect(await screen.findByText("Monthly Spend")).toBeInTheDocument(); - expect(screen.getByText("Top Virtual Keys")).toBeInTheDocument(); - expect(screen.getByText("Top Models")).toBeInTheDocument(); - expect(screen.getByText("Spend by Provider")).toBeInTheDocument(); - }); - - it("lists spend by provider in a table", async () => { - renderUsage(); - - const providerCell = await screen.findByText("openai"); - const row = providerCell.closest("tr"); - expect(row).not.toBeNull(); - expect(within(row as HTMLElement).getByText("$6.75")).toBeInTheDocument(); - expect(screen.getByRole("columnheader", { name: "Provider" })).toBeInTheDocument(); - }); - - it("shows the customer usage table when its tab is selected", async () => { - const user = userEvent.setup(); - renderUsage(); - - await user.click(await screen.findByRole("tab", { name: "Customer Usage" })); - - const customerCell = await screen.findByText("customer-alpha"); - const row = customerCell.closest("tr"); - expect(row).not.toBeNull(); - expect(within(row as HTMLElement).getByText("$3.50")).toBeInTheDocument(); - expect(within(row as HTMLElement).getByText("42")).toBeInTheDocument(); - expect(screen.getByRole("columnheader", { name: "Total Events" })).toBeInTheDocument(); - }); - - it("shows the tag spend panel when its tab is selected", async () => { - const user = userEvent.setup(); - renderUsage(); - - await user.click(await screen.findByRole("tab", { name: "Tag Based Usage" })); - - expect(await screen.findByText("Spend Per Tag")).toBeInTheDocument(); - }); - - it("shows the team spend panel when its tab is selected", async () => { - const user = userEvent.setup(); - renderUsage(); - - await user.click(await screen.findByRole("tab", { name: "Team Based Usage" })); - - expect(await screen.findByText("Total Spend Per Team")).toBeInTheDocument(); - expect(screen.getByText("Daily Spend Per Team")).toBeInTheDocument(); - }); - }); - - // org_admin is an organization membership role; those users reach the UI as "Internal User". - describe.each(["Internal User", "Internal Viewer", "internal_user", "internal_user_viewer", "Org Admin"])( - "as %s", - (userRole) => { - it("shows the admin-only notice instead of the usage dashboard", async () => { - renderUsage({ userRole }); - - expect(await screen.findByText(/Proxy-wide usage is only available to admin users/i)).toBeInTheDocument(); - expect(screen.queryByRole("tab", { name: "All Up" })).not.toBeInTheDocument(); - }); - - it("fires no /global/spend or /global/activity request", async () => { - renderUsage({ userRole }); - - await screen.findByText(/Proxy-wide usage is only available to admin users/i); - await flushPendingRequests(); - - expect(networking.getProxyUISettings).not.toHaveBeenCalled(); - expect(networking.adminSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.adminTopKeysCall).not.toHaveBeenCalled(); - expect(networking.adminTopModelsCall).not.toHaveBeenCalled(); - expect(networking.adminTopEndUsersCall).not.toHaveBeenCalled(); - expect(networking.teamSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.tagsSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.allTagNamesCall).not.toHaveBeenCalled(); - expect(networking.adminspendByProvider).not.toHaveBeenCalled(); - expect(networking.adminGlobalActivity).not.toHaveBeenCalled(); - expect(networking.adminGlobalActivityPerModel).not.toHaveBeenCalled(); - }); - }, - ); - - describe("the admin-only gate", () => { - it("proves the flush window is wide enough to catch a leaked request", async () => { - renderUsage({ userRole: "Admin" }); - - await flushPendingRequests(); - - expect(networking.getProxyUISettings).toHaveBeenCalled(); - expect(networking.adminSpendLogsCall).toHaveBeenCalled(); - expect(networking.tagsSpendLogsCall).toHaveBeenCalled(); - expect(networking.adminGlobalActivity).toHaveBeenCalled(); - }); - - it("still lets an admin through, so the notice is a real gate and not a dead branch", async () => { - renderUsage({ userRole: "Admin" }); - - expect(await screen.findByRole("tab", { name: "All Up" })).toBeInTheDocument(); - expect(screen.queryByText(/Proxy-wide usage is only available to admin users/i)).not.toBeInTheDocument(); - await waitFor(() => expect(networking.adminSpendLogsCall).toHaveBeenCalled()); - }); - - it("does not put the session token in the provider spend query", async () => { - renderUsage({ userRole: "Admin", token: "session-jwt-value" }); - - await waitFor(() => expect(networking.adminspendByProvider).toHaveBeenCalled()); - const callArgs = networking.adminspendByProvider.mock.calls[0]; - expect(callArgs).not.toContain("session-jwt-value"); - expect(callArgs[0]).toBe("sk-test"); - }); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx deleted file mode 100644 index 15b2a30b50c..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx +++ /dev/null @@ -1,950 +0,0 @@ -import React, { useState, useEffect } from "react"; - -import ViewUserSpend from "@/components/view_user_spend"; -import { ProxySettings } from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; -import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; -import { Button } from "@/components/ui/button"; -import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; -import { - Combobox, - ComboboxChip, - ComboboxChips, - ComboboxChipsInput, - ComboboxContent, - ComboboxEmpty, - ComboboxItem, - ComboboxList, - ComboboxValue, - useComboboxAnchor, -} from "@/components/ui/combobox"; -import { Meter, MeterIndicator, MeterTrack } from "@/components/shared/Meter"; -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import { - NUMERIC_CELL_CLASS, - Table, - TableBody, - TableCell, - TableHead, - TableHeader, - TableRow, -} from "@/components/ui/table"; -import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import { AreaChart, BarChart, DonutChart } from "@/components/shared/charts"; - -import { - adminSpendLogsCall, - adminTopKeysCall, - adminTopModelsCall, - adminTopEndUsersCall, - teamSpendLogsCall, - tagsSpendLogsCall, - allTagNamesCall, - adminspendByProvider, - adminGlobalActivity, - adminGlobalActivityPerModel, - getProxyUISettings, -} from "@/components/networking"; -import TopKeyView from "@/components/UsagePage/components/EntityUsage/TopKeyView"; -import { MoneyCell } from "@/components/shared/table_cells"; -import { hasCapability } from "@/utils/capabilities"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; - -interface UsagePageProps { - accessToken: string | null; - token: string | null; - userRole: string | null; - userID: string | null; - keys: any[] | null; - premiumUser: boolean; -} - -interface GlobalActivityData { - sum_api_requests: number; - sum_total_tokens: number; - daily_data: { date: string; api_requests: number; total_tokens: number }[]; -} - -const EMPTY_GLOBAL_ACTIVITY: GlobalActivityData = { sum_api_requests: 0, sum_total_tokens: 0, daily_data: [] }; - -type UsageDateRange = { from?: Date; to?: Date }; - -type TeamSpendTotal = { name: string; value: number }; - -type TagOption = { value: string; label: string; disabled: boolean }; - -const ALL_TAGS = "all-tags"; - -const isAdminOrAdminViewer = (role: string | null): boolean => { - if (role === null) return false; - return role === "Admin" || role === "Admin Viewer"; -}; - -const TeamSpendBarList: React.FC<{ data: TeamSpendTotal[] }> = ({ data }) => { - const max = Math.max(0, ...data.map((team) => team.value)); - - return ( -

- {data.map((team) => ( -
-

{team.name}

- - - - - -

- {formatNumberWithCommas(team.value, 2)} -

-
- ))} -
- ); -}; - -const UsagePage: React.FC = ({ accessToken, token, userRole, userID, keys, premiumUser }) => { - const anchor = useComboboxAnchor(); - const canViewGlobalSpend = hasCapability(userRole, "viewGlobalSpend"); - const currentDate = new Date(); - const [keySpendData, setKeySpendData] = useState([]); - const [topKeys, setTopKeys] = useState([]); - const [topModels, setTopModels] = useState([]); - const [topUsers, setTopUsers] = useState([]); - const [teamSpendData, setTeamSpendData] = useState([]); - const [topTagsData, setTopTagsData] = useState([]); - const [allTagNames, setAllTagNames] = useState([]); - const [uniqueTeamIds, setUniqueTeamIds] = useState([]); - const [totalSpendPerTeam, setTotalSpendPerTeam] = useState([]); - const [spendByProvider, setSpendByProvider] = useState([]); - const [globalActivity, setGlobalActivity] = useState(EMPTY_GLOBAL_ACTIVITY); - const [globalActivityPerModel, setGlobalActivityPerModel] = useState([]); - const [selectedKeyToken, setSelectedKeyToken] = useState(null); - const [selectedTags, setSelectedTags] = useState([ALL_TAGS]); - const [dateValue, setDateValue] = useState({ - from: new Date(Date.now() - 7 * 24 * 60 * 60 * 1000), - to: new Date(), - }); - const [proxySettings, setProxySettings] = useState(null); - const [totalMonthlySpend, setTotalMonthlySpend] = useState(0); - - const firstDay = new Date(currentDate.getFullYear(), currentDate.getMonth(), 1); - const lastDay = new Date(currentDate.getFullYear(), currentDate.getMonth() + 1, 0); - - let startTime = formatDate(firstDay); - let endTime = formatDate(lastDay); - - const selectableKeys: { token: string; alias: string }[] = (keys ?? []) - .filter((key: any) => key && typeof key["key_alias"] === "string" && key["key_alias"].length > 0) - .map((key: any) => ({ token: String(key["token"]), alias: String(key["key_alias"]) })); - - const tagOptions: TagOption[] = [ - { value: ALL_TAGS, label: "All Tags", disabled: false }, - ...allTagNames - .filter((tag) => tag !== ALL_TAGS) - .map((tag) => ({ - value: tag, - label: premiumUser ? tag : `✨ ${tag} (Enterprise only Feature)`, - disabled: !premiumUser, - })), - ]; - - function valueFormatterNumbers(number: number) { - const formatter = new Intl.NumberFormat("en-US", { - maximumFractionDigits: 0, - notation: "compact", - compactDisplay: "short", - }); - - return formatter.format(number); - } - - const fetchProxySettings = async () => { - if (accessToken) { - try { - const proxy_settings: ProxySettings = await getProxyUISettings(accessToken); - return proxy_settings; - } catch (error) { - console.error("Error fetching proxy settings:", error); - } - } - }; - - useEffect(() => { - if (!canViewGlobalSpend) { - return; - } - updateTagSpendData(dateValue.from, dateValue.to); - }, [canViewGlobalSpend, dateValue, selectedTags]); - - const updateEndUserData = async ( - startTime: Date | undefined, - endTime: Date | undefined, - uiSelectedKey: string | null, - ) => { - if (!startTime || !endTime || !accessToken) { - return; - } - - let newTopUserData = await adminTopEndUsersCall( - accessToken, - uiSelectedKey, - startTime.toISOString(), - endTime.toISOString(), - ); - setTopUsers(newTopUserData); - }; - - const updateTagSpendData = async (startTime: Date | undefined, endTime: Date | undefined) => { - if (!startTime || !endTime || !accessToken) { - return; - } - - // we refetch because the state variable can be None when the user refreshes the page - const proxy_settings: ProxySettings | undefined = await fetchProxySettings(); - - if (proxy_settings?.DISABLE_EXPENSIVE_DB_QUERIES) { - return; // Don't run expensive DB queries - return out when SpendLogs has more than 1M rows - } - - let top_tags = await tagsSpendLogsCall( - accessToken, - startTime.toISOString(), - endTime.toISOString(), - selectedTags.length === 0 ? undefined : selectedTags, - ); - setTopTagsData(top_tags.spend_per_tag); - }; - - function formatDate(date: Date) { - const year = date.getFullYear(); - let month = date.getMonth() + 1; // JS month index starts from 0 - let day = date.getDate(); - - // Pad with 0 if month or day is less than 10 - const monthStr = month < 10 ? "0" + month : month; - const dayStr = day < 10 ? "0" + day : day; - - return `${year}-${monthStr}-${dayStr}`; - } - - const valueFormatter = (number: number) => `$ ${formatNumberWithCommas(number, 2)}`; - - const fetchAndSetData = async ( - fetchFunction: () => Promise, - setStateFunction: React.Dispatch>, - errorMessage: string, - ) => { - try { - const data = await fetchFunction(); - setStateFunction(data); - } catch (error) { - console.error(errorMessage, error); - // Optionally, update UI to reflect error state for this specific data - } - }; - - // Update the fillMissingDates function to handle different date formats - const fillMissingDates = (data: any[], startDate: Date, endDate: Date, categories: string[]) => { - const filledData = []; - const currentDate = new Date(startDate); - - // Helper function to standardize date format - const standardizeDate = (dateStr: string) => { - if (dateStr.includes("-")) { - // Already in YYYY-MM-DD format - return dateStr; - } else { - // Convert "Jan 06" format - const [month, day] = dateStr.split(" "); - const year = new Date().getFullYear(); - const monthIndex = new Date(`${month} 01 2024`).getMonth(); - const fullDate = new Date(year, monthIndex, parseInt(day)); - return fullDate.toISOString().split("T")[0]; - } - }; - - // Create a map of existing dates for quick lookup - const existingDates = new Map( - data.map((item) => { - const standardizedDate = standardizeDate(item.date); - return [ - standardizedDate, - { - ...item, - date: standardizedDate, // Store standardized date format - }, - ]; - }), - ); - - // Iterate through each date in the range - while (currentDate <= endDate) { - const dateStr = currentDate.toISOString().split("T")[0]; - - if (existingDates.has(dateStr)) { - // Use existing data if we have it - filledData.push(existingDates.get(dateStr)); - } else { - // Create an entry with zero values - const emptyEntry: any = { - date: dateStr, - api_requests: 0, - total_tokens: 0, - }; - - // Add zero values for each model/team if needed - categories.forEach((category) => { - if (!emptyEntry[category]) { - emptyEntry[category] = 0; - } - }); - - filledData.push(emptyEntry); - } - - // Move to next day - currentDate.setDate(currentDate.getDate() + 1); - } - - return filledData; - }; - - // Update the fetchOverallSpend function - const fetchOverallSpend = async () => { - if (!accessToken) { - return; - } - try { - const data = await adminSpendLogsCall(accessToken); - - // Get the first and last day of the current month - const now = new Date(); - const firstDay = new Date(now.getFullYear(), now.getMonth(), 1); - const lastDay = new Date(now.getFullYear(), now.getMonth() + 1, 0); - - // Fill in missing dates - const filledData = fillMissingDates(data, firstDay, lastDay, []); - - // Calculate total spend for the month and round to 2 decimal places - const monthlyTotal = Number(filledData.reduce((sum, day) => sum + (day.spend || 0), 0).toFixed(2)); - setTotalMonthlySpend(monthlyTotal); - - setKeySpendData(filledData); - } catch (error) { - console.error("Error fetching overall spend:", error); - } - }; - - const fetchProviderSpend = () => - fetchAndSetData( - () => (accessToken ? adminspendByProvider(accessToken, startTime, endTime) : Promise.reject("No access token")), - setSpendByProvider, - "Error fetching provider spend", - ); - - const fetchTopKeys = async () => { - if (!accessToken) return; - await fetchAndSetData( - async () => { - const top_keys = await adminTopKeysCall(accessToken); - return top_keys.map((k: any) => ({ - key: k["api_key"].substring(0, 10), - api_key: k["api_key"], - key_alias: k["key_alias"], - spend: Number(k["total_spend"].toFixed(2)), - })); - }, - setTopKeys, - "Error fetching top keys", - ); - }; - - const fetchTopModels = async () => { - if (!accessToken) return; - await fetchAndSetData( - async () => { - const top_models = await adminTopModelsCall(accessToken); - return top_models.map((k: any) => ({ - key: k["model"], - spend: formatNumberWithCommas(k["total_spend"], 2), - })); - }, - setTopModels, - "Error fetching top models", - ); - }; - - // Update the fetchTeamSpend function - const fetchTeamSpend = async () => { - if (!accessToken) return; - await fetchAndSetData( - async () => { - const teamSpend = await teamSpendLogsCall(accessToken); - - // Get the first and last day of the current month - const now = new Date(); - const firstDay = new Date(now.getFullYear(), now.getMonth(), 1); - const lastDay = new Date(now.getFullYear(), now.getMonth() + 1, 0); - - // Fill in missing dates with zero values for all teams - const filledData = fillMissingDates(teamSpend.daily_spend, firstDay, lastDay, teamSpend.teams); - - setTeamSpendData(filledData); - setUniqueTeamIds(teamSpend.teams); - return teamSpend.total_spend_per_team.map((tspt: any) => ({ - name: tspt["team_id"] || "", - value: Number(tspt["total_spend"] || 0), - })); - }, - setTotalSpendPerTeam, - "Error fetching team spend", - ); - }; - - const fetchTagNames = () => { - if (!accessToken) return; - fetchAndSetData( - async () => { - const all_tag_names = await allTagNamesCall(accessToken); - return all_tag_names.tag_names; - }, - setAllTagNames, - "Error fetching tag names", - ); - }; - - const fetchTopTags = () => { - if (!accessToken) return; - fetchAndSetData( - () => tagsSpendLogsCall(accessToken, dateValue.from?.toISOString(), dateValue.to?.toISOString(), undefined), - (data) => setTopTagsData(data.spend_per_tag), - "Error fetching top tags", - ); - }; - - const fetchTopEndUsers = () => { - if (!accessToken) return; - fetchAndSetData( - () => adminTopEndUsersCall(accessToken, null, undefined, undefined), - setTopUsers, - "Error fetching top end users", - ); - }; - - // Update the fetchGlobalActivity function - const fetchGlobalActivity = async () => { - if (!accessToken) return; - try { - const data = await adminGlobalActivity(accessToken, startTime, endTime); - - // Get the date range from the current month - const now = new Date(); - const firstDay = new Date(now.getFullYear(), now.getMonth(), 1); - const lastDay = new Date(now.getFullYear(), now.getMonth() + 1, 0); - - // Fill in missing dates for daily_data - const filledDailyData = fillMissingDates(data.daily_data || [], firstDay, lastDay, [ - "api_requests", - "total_tokens", - ]); - - setGlobalActivity({ - ...data, - daily_data: filledDailyData, - }); - } catch (error) { - console.error("Error fetching global activity:", error); - } - }; - - // Update the fetchGlobalActivityPerModel function - const fetchGlobalActivityPerModel = async () => { - if (!accessToken) return; - try { - const data = await adminGlobalActivityPerModel(accessToken, startTime, endTime); - - // Get the date range from the current month - const now = new Date(); - const firstDay = new Date(now.getFullYear(), now.getMonth(), 1); - const lastDay = new Date(now.getFullYear(), now.getMonth() + 1, 0); - - // Fill in missing dates for each model's daily data - const filledModelData = data.map((modelData: any) => ({ - ...modelData, - daily_data: fillMissingDates(modelData.daily_data || [], firstDay, lastDay, ["api_requests", "total_tokens"]), - })); - - setGlobalActivityPerModel(filledModelData); - } catch (error) { - console.error("Error fetching global activity per model:", error); - } - }; - - useEffect(() => { - const initlizeUsageData = async () => { - if (!canViewGlobalSpend) { - return; - } - if (accessToken && token && userRole && userID) { - const proxy_settings: ProxySettings | undefined = await fetchProxySettings(); - if (proxy_settings) { - setProxySettings(proxy_settings); // saved in state so it can be used when rendering UI - if (proxy_settings?.DISABLE_EXPENSIVE_DB_QUERIES) { - return; // Don't run expensive UI queries - return out of initlizeUsageData at this point - } - } - - fetchOverallSpend(); - fetchProviderSpend(); - fetchTopKeys(); - fetchTopModels(); - fetchGlobalActivity(); - fetchGlobalActivityPerModel(); - - if (isAdminOrAdminViewer(userRole)) { - fetchTeamSpend(); - fetchTagNames(); - fetchTopTags(); - fetchTopEndUsers(); - } - } - }; - - initlizeUsageData(); - }, [canViewGlobalSpend, accessToken, token, userRole, userID, startTime, endTime]); - - if (!canViewGlobalSpend) { - return ( -
- - - Usage - - -

- Proxy-wide usage is only available to admin users. Your own usage is on the Usage page. -

- - -
- ); - } - - if (proxySettings?.DISABLE_EXPENSIVE_DB_QUERIES) { - return ( -
- - - Database Query Limit Reached - - -

- SpendLogs in DB has {proxySettings.NUM_SPEND_LOGS_ROWS} rows. -

- Please follow our guide to view usage when SpendLogs has more than 1M rows. -

-
- ); - } - - return ( -
- - - All Up - - {isAdminOrAdminViewer(userRole) && ( - <> - Team Based Usage - Customer Usage - Tag Based Usage - - )} - - - - - - Cost - Activity - - - -
-
-

- Project Spend {new Date().toLocaleString("default", { month: "long" })} 1 -{" "} - {new Date(new Date().getFullYear(), new Date().getMonth() + 1, 0).getDate()} -

- -
-
- - - Monthly Spend - - - - - -
-
- - - Top Virtual Keys - - - {}} /> - - -
-
- - - Top Models - - - `$${formatNumberWithCommas(value, 2)}`} - /> - - -
-
-
- - - Spend by Provider - - -
-
- `$${formatNumberWithCommas(value, 2)}`} - /> -
-
- - - - Provider - Spend - - - - {spendByProvider.map((provider) => ( - - {provider.provider} - - - - - ))} - -
-
-
- - -
-
- - - -
- - - All Up - - -
-
-

- API Requests {valueFormatterNumbers(globalActivity.sum_api_requests)} -

- -
-
-

- Tokens {valueFormatterNumbers(globalActivity.sum_total_tokens)} -

- -
-
- - - - {globalActivityPerModel.map((globalActivity, index) => ( - - - {globalActivity.model} - - -
-
-

- API Requests {valueFormatterNumbers(globalActivity.sum_api_requests)} -

- -
-
-

- Tokens {valueFormatterNumbers(globalActivity.sum_total_tokens)} -

- -
-
- - - ))} -
- - - - - -
-
- - - Total Spend Per Team - - - - - - - - Daily Spend Per Team - - - - - -
-
- - - -

- Customers of your LLM API calls. Tracked when a `user` param is passed in your LLM calls{" "} - - docs here - -

-
-
- { - setDateValue(value); - updateEndUserData(value.from, value.to, null); - }} - /> -
-
-

Select Key

- -
-
- - - -
- - - - Customer - Spend - Total Events - - - - - {topUsers?.map((user: any, index: number) => ( - - {user.end_user} - - - - {user.total_count} - - ))} - -
-
- - - - - -
-
- { - setDateValue(value); - updateTagSpendData(value.from, value.to); - }} - /> -
- -
- selectedTags.includes(option.value))} - onValueChange={(options: TagOption[]) => setSelectedTags(options.map((option) => option.value))} - isItemEqualToValue={(a: TagOption, b: TagOption) => a.value === b.value} - itemToStringLabel={(option: TagOption) => option.label} - > - }> - - {(options: TagOption[]) => - options.map((option) => ( - - {option.label} - - )) - } - - - - - No tags found - - {(option: TagOption) => ( - - {option.label} - - )} - - - -
-
-
-
- - - Spend Per Tag - - -

- Get Started by Tracking cost per tag{" "} - - here - -

- - - -
-
- - -
- ); -}; - -export default UsagePage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/page.tsx deleted file mode 100644 index 138dd97e5e8..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/page.tsx +++ /dev/null @@ -1,22 +0,0 @@ -"use client"; - -import Usage from "./_components/usage"; -import { DeprecationBanner } from "@/components/DeprecationBanner"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -export default function OldUsagePage() { - const { accessToken, token, userRole, userId: userID, premiumUser } = useAuthorized(); - return ( - <> - - - - ); -} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx index 9d39e9c080c..f7ffd6d52cf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx @@ -102,10 +102,6 @@ vi.mock("@/components/activity_metrics", () => ({ processActivityData: vi.fn(), })); -vi.mock("@/components/view_user_spend", () => ({ - default: () =>
View User Spend
, -})); - vi.mock("@/components/UsagePage/components/EntityUsage/TopKeyView", () => ({ default: () =>
Top Keys
, })); diff --git a/ui/litellm-dashboard/src/components/leftnav.test.tsx b/ui/litellm-dashboard/src/components/leftnav.test.tsx index 6cf05af1ed1..f3bff85e2d5 100644 --- a/ui/litellm-dashboard/src/components/leftnav.test.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.test.tsx @@ -444,30 +444,6 @@ describe("Sidebar (leftnav)", () => { }); expect(screen.queryByText("Prompts")).not.toBeInTheDocument(); }); - - it("should hide Old Usage from internal users while keeping other Experimental children", async () => { - mockUseAuthorized.mockReturnValue(internalAuth); - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByText("Experimental")); - }); - await waitFor(() => { - expect(screen.getByText("API Playground")).toBeInTheDocument(); - }); - expect(screen.queryByText("Old Usage")).not.toBeInTheDocument(); - }); - - it("should show Old Usage to admins", async () => { - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByText("Experimental")); - }); - await waitFor(() => { - expect(screen.getByText("Old Usage")).toBeInTheDocument(); - }); - }); }); // Workflow Runs, Memory and Guardrails Monitor render a shell and then 401 @@ -704,7 +680,6 @@ describe("Sidebar (leftnav)", () => { expectHref("Models + Endpoints", "/ui/models-and-endpoints"); expectHref("Usage", "/ui/usage"); expectHref("API Reference", "/ui/api-reference"); - expectHref("Old Usage", "/ui/old-usage"); }); it("never links a leaf to the legacy ?page= switch", () => { @@ -760,7 +735,6 @@ describe("getBreadcrumb", () => { it("resolves routes whose segment differs from the sidebar page id", () => { expect(getBreadcrumb("/ui/models-and-endpoints")).toEqual({ section: "AI Gateway", title: "Models + Endpoints" }); expect(getBreadcrumb("/ui/usage")).toEqual({ section: "Observability", title: "Usage" }); - expect(getBreadcrumb("/ui/old-usage")).toEqual({ section: "Developer Tools", title: "Old Usage" }); }); it("titles the dashboard root as Home", () => { diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 6795cf1c7d1..2cf6f8219db 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -414,14 +414,6 @@ const menuGroups: MenuGroup[] = [ icon: , roles: all_admin_roles, }, - { - key: "4", - page: "usage", - route: "old-usage", - label: "Old Usage", - icon: , - roles: rolesWithCapability("viewGlobalSpend"), - }, ], }, ], diff --git a/ui/litellm-dashboard/src/components/page_metadata.ts b/ui/litellm-dashboard/src/components/page_metadata.ts index 18f136c3b28..97220a0ef76 100644 --- a/ui/litellm-dashboard/src/components/page_metadata.ts +++ b/ui/litellm-dashboard/src/components/page_metadata.ts @@ -41,7 +41,6 @@ export const pageDescriptions: Record = { "tag-management": "Organize resources with tags", prompts: "Manage and version prompt templates", skills: "Browse and manage Claude Code skills", - usage: "View legacy usage dashboard", "router-settings": "Configure routing and load balancing settings", telemetry: "Choose which anonymous usage telemetry this proxy collects and see the last report", "logging-and-alerts": "Set up logging and alert configurations", diff --git a/ui/litellm-dashboard/src/components/view_user_spend.tsx b/ui/litellm-dashboard/src/components/view_user_spend.tsx deleted file mode 100644 index 655fc092673..00000000000 --- a/ui/litellm-dashboard/src/components/view_user_spend.tsx +++ /dev/null @@ -1,144 +0,0 @@ -"use client"; -import React, { useEffect, useState } from "react"; -import { modelAvailableCall } from "./networking"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -// Define the props type -interface ViewUserSpendProps { - userSpend: number | null; - userMaxBudget: number | null; - selectedTeam: any | null; -} -const ViewUserSpend: React.FC = ({ userSpend, userMaxBudget, selectedTeam }) => { - const { accessToken, userRole, userId: userID } = useAuthorized(); - let [spend, setSpend] = useState(userSpend !== null ? userSpend : 0.0); - const [maxBudget, setMaxBudget] = useState( - selectedTeam ? Number(formatNumberWithCommas(selectedTeam.max_budget, 4)) : null, - ); - useEffect(() => { - if (selectedTeam) { - if (selectedTeam.team_alias === "Default Team") { - setMaxBudget(userMaxBudget); - } else { - let setMaxBudgetFlag = false; - if (selectedTeam.team_memberships) { - /** - * What 'team_memberships' looks like: - * "team_memberships": [ - * { - * "user_id": "2c315de3-e7ce-4269-b73e-b039a06187b1", - * "team_id": "test-team_515e6f42-ded2-4f0d-8919-0a1f43c5a45f", - * "budget_id": "0880769f-716a-4149-ab19-7f7651ad4db5", - * "litellm_budget_table": { - "soft_budget": null, - "max_budget": 20.0, - "max_parallel_requests": null, - "tpm_limit": null, - "rpm_limit": null, - "model_max_budget": null, - "budget_duration": null - } - */ - for (const member of selectedTeam.team_memberships) { - if ( - member.user_id === userID && - "max_budget" in member.litellm_budget_table && - member.litellm_budget_table.max_budget !== null - ) { - setMaxBudget(member.litellm_budget_table.max_budget); - setMaxBudgetFlag = true; - } - } - } - if (!setMaxBudgetFlag) { - setMaxBudget(selectedTeam.max_budget); - } - } - } else { - setMaxBudget(userMaxBudget); - } - }, [selectedTeam, userMaxBudget]); - const [userModels, setUserModels] = useState([]); - useEffect(() => { - const fetchData = async () => { - if (!accessToken || !userID || !userRole) { - return; - } - }; - const fetchUserModels = async () => { - try { - if (userID === null || userRole === null) { - return; - } - - if (accessToken !== null) { - const model_available = await modelAvailableCall(accessToken, userID, userRole); - let available_model_names = model_available["data"].map((element: { id: string }) => element.id); - setUserModels(available_model_names); - } - } catch (error) { - console.error("Error fetching user models:", error); - } - }; - - fetchUserModels(); - fetchData(); - }, [userRole, accessToken, userID]); - - useEffect(() => { - if (userSpend !== null) { - setSpend(userSpend); - } - }, [userSpend]); - - // logic to decide what models to display - let modelsToDisplay = []; - if (selectedTeam && selectedTeam.models) { - modelsToDisplay = selectedTeam.models; - } - - // check if "all-proxy-models" is in modelsToDisplay - if (modelsToDisplay && modelsToDisplay.includes("all-proxy-models")) { - modelsToDisplay = userModels; - } else if (modelsToDisplay && modelsToDisplay.includes("all-team-models")) { - modelsToDisplay = selectedTeam.models; - } else if (modelsToDisplay && modelsToDisplay.length === 0) { - modelsToDisplay = userModels; - } - - const displayMaxBudget = maxBudget !== null ? `$${formatNumberWithCommas(Number(maxBudget), 4)} limit` : "No limit"; - - const roundedSpend = spend !== undefined ? formatNumberWithCommas(spend, 4) : null; - - return ( -
-
-
-

Total Spend

-

${roundedSpend}

-
-
-

Max Budget

-

{displayMaxBudget}

-
-
- {/*
- - Team Models - - - {modelsToDisplay.map((model: string) => ( - - {model} - - ))} - - - -
*/} -
- ); -}; - -export default ViewUserSpend; diff --git a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx index 58d7309c078..e5efd4c8e9f 100644 --- a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx +++ b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx @@ -127,7 +127,6 @@ vi.mock("@/components/new_usage", () => ({ default: stub("new-usage") })); vi.mock("@/components/api_ref", () => ({ default: stub("api-ref") })); vi.mock("@/components/chat_ui/ChatUI", () => ({ default: stub("chat-ui") })); vi.mock("@/components/leftnav", () => ({ default: stub("sidebar") })); -vi.mock("@/app/(dashboard)/old-usage/_components/usage", () => ({ default: stub("usage") })); vi.mock("@/components/cache_dashboard", () => ({ default: stub("cache-dashboard") })); vi.mock("@/app/(dashboard)/guardrails/_components", () => ({ default: stub("guardrails") })); vi.mock("@/components/prompts", () => ({ default: stub("prompts") })); From e752528ac8f9e1399e0e5bff98b7f404bffda1b4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:53:06 -0700 Subject: [PATCH 47/64] feat(keys): allow assigning a project to an unassigned key on /key/update (#41700) * feat(keys): allow assigning a project to an unassigned key on /key/update Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): suppress LIT002 on project assignment error payloads Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): format project assignment error detail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(keys): assert accepted project assignment flows into update payload Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): drop stale LIT002 suppressions removed by main Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(keys): mock writer_db team lookup used by check_db_only path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(keys): write project assignment only while key row is still unassigned Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): satisfy type gate on guarded project assignment write Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(keys): default soft-budget update path to unguarded write Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(keys): run guarded project assignment in one writer transaction Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(keys): include object permission link in assignment guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(keys): return permission-linked payload without mutating the update dict Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): fix type errors and stale noqa after main merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): use Sequence in assignment guard model annotation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): type update fields as a mapping in the key update path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): place cast reason on the cast call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(keys): honor explicit null max_budget when assigning a project * test(keys): assert explicit null max_budget flows into update payload * fix(keys): match NULL models column in guarded project assignment * fix(keys): authorize caller before project lookup on /key/update * fix(keys): drop explicit null object_permission before project assignment write --------- Co-authored-by: jesus Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 249 +++++++- .../management_endpoints/key_requests.py | 2 +- .../test_key_management_endpoints.py | 530 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 4 files changed, 759 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3021b21de51..61da5599886 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -162,7 +162,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # _is_master_key, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export is_master_key, ) -from litellm.proxy.utils import ( # noqa: F401 # legacy module exports +from litellm.proxy.utils import ( PrismaClient, ProxyLogging, _hash_token_if_needed, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -280,6 +280,10 @@ class _BudgetRowSoftBudgetCreate(TypedDict): updated_by: ReadOnly[str] +class _KeyRowDumpable(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + class _KeyUpdateTx(Protocol): @property def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ... @@ -287,6 +291,9 @@ class _KeyUpdateTx(Protocol): @property def litellm_budgettable(self) -> "TableActions[prisma_models.LiteLLM_BudgetTable]": ... + @property + def litellm_objectpermissiontable(self) -> "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]": ... + class _ConfigTableActions(Protocol): """Config table surface this module needs; the shared repository seam exposes no ``update``.""" @@ -2072,6 +2079,56 @@ async def _check_project_key_limits( ) +async def _validate_project_assignment( + data: UpdateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> None: + """Validate assigning a project to a key that does not have one yet.""" + if data.project_id is None or data.project_id == existing_key_row.project_id: + return + if existing_key_row.project_id is not None: + raise HTTPException( + status_code=400, detail="Project reassignment is not supported. Use null to detach the key." + ) + project_obj: Final = await get_project_object( + project_id=data.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if project_obj is None: + raise HTTPException( + status_code=404, + detail={"error": f"Project not found, project_id={data.project_id}"}, + ) + team: Final = data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id + if team is None or team != project_obj.team_id: + raise HTTPException( + status_code=400, + detail={ + "error": f"Project {data.project_id} belongs to team {project_obj.team_id}. Keys can only be assigned to a project owned by their own team (key team: {team})" + }, + ) + await _check_project_key_limits( + project_id=data.project_id, + data=data.model_copy( + update={ + "models": ( + data.models + if "models" in data.model_fields_set + else cast( # cast-ok: LiteLLM_VerificationToken.models is a bare list + "list[str] | None", existing_key_row.models + ) + ), + "max_budget": data.max_budget if "max_budget" in data.model_fields_set else existing_key_row.max_budget, + } + ), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + async def _check_key_project_team( project_id: str, key_team_id: str | None, @@ -2771,6 +2828,115 @@ async def _apply_soft_budget_update( return remaining +async def _apply_object_permission_update( + data_json: Mapping[str, object], + existing_object_permission_id: str | None, + prisma_client: PrismaClient, + table: "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]", +) -> Mapping[str, object]: + """Upsert the requested object permission row inside the caller's writer transaction.""" + new_object_permission: Final = data_json.get("object_permission") + if new_object_permission is None: + return MappingProxyType({k: v for k, v in data_json.items() if k != "object_permission"}) + loaded: Final[object] = ( + json.loads(new_object_permission) if isinstance(new_object_permission, str) else new_object_permission + ) + grants: Final[Mapping[str, object]] = ( + cast( # cast-ok: object_permission payloads are grant-name to grant-list mappings + "Mapping[str, object]", loaded + ) + if isinstance(loaded, dict) + else MappingProxyType({}) + ) + upsert: Final = await prepare_object_permission_upsert( + new_object_permission=grants, + existing_object_permission_id=existing_object_permission_id, + prisma_client=prisma_client, + ) + row: Final = await table.upsert( + where={"object_permission_id": upsert.object_permission_id}, + data={"create": upsert.record, "update": upsert.record}, + ) + return MappingProxyType( + { + **{k: v for k, v in data_json.items() if k != "object_permission"}, + "object_permission_id": row.object_permission_id, + } + ) + + +async def _write_guarded_project_assignment( + table: "TableActions[_KeyRowDumpable]", + hashed_token: str, + existing_key_row: LiteLLM_VerificationToken, + data: Mapping[str, object], +) -> "_KeyRowDumpable | None": + validated_models: Final[Sequence[str]] = ( + cast( # cast-ok: LiteLLM_VerificationToken.models is a bare list + "Sequence[str] | None", existing_key_row.models + ) + or [] + ) + models_guard: Final[Mapping[str, object]] = ( + {"OR": [{"models": {"equals": list[str]()}}, {"models": {"equals": None}}]} + if not validated_models + else {"models": {"equals": validated_models}} + ) + updated_count: Final = await table.update_many( + where={ + "token": hashed_token, + "project_id": None, + "team_id": existing_key_row.team_id, + **models_guard, + "max_budget": existing_key_row.max_budget, + "object_permission_id": existing_key_row.object_permission_id, + }, + data=data, + ) + if updated_count == 0: + raise HTTPException( + status_code=400, + detail="Project assignment failed because the key was updated concurrently. Retry the request.", + ) + return await table.find_unique( + where={"token": hashed_token}, + include={"object_permission": True}, + ) + + +async def _update_key_row_assigning_project( + prisma_client: PrismaClient, + key: str, + non_default_values: Mapping[str, object], + existing_key_row: LiteLLM_VerificationToken, +) -> _KeyUpdateResult: + """Assign a project in one writer transaction, guarded on the key's validated state.""" + hashed_token: Final = _hash_token_if_needed(key) + async with prisma_client.tx() as tx_ctx: + tx: Final[_KeyUpdateTx] = cast( # cast-ok: the transaction object exposes the same table actions + "_KeyUpdateTx", tx_ctx + ) + update_values: Final[Mapping[str, object]] = await _apply_object_permission_update( + data_json=non_default_values, + existing_object_permission_id=existing_key_row.object_permission_id, + prisma_client=prisma_client, + table=tx.litellm_objectpermissiontable, + ) + updated_row: Final = await _write_guarded_project_assignment( + table=tx.litellm_verificationtoken, + hashed_token=hashed_token, + existing_key_row=existing_key_row, + data=with_settings_updated_at( + prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) + ), + ) + updated_data: Final[Mapping[str, object]] = ( + updated_row.model_dump() if updated_row is not None else MappingProxyType({}) + ) + result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data} + return result + + async def _update_key_row_with_soft_budget( prisma_client: PrismaClient, key: str, @@ -2778,25 +2944,48 @@ async def _update_key_row_with_soft_budget( non_default_values: Mapping[str, object], existing_key_row: LiteLLM_VerificationToken, changed_by: str, + expect_unassigned_project: bool = False, ) -> _KeyUpdateResult: hashed_token: Final = hash_token_if_needed(key) key_where: Final[_KeyRowWhere] = {"token": hashed_token} - tx: _KeyUpdateTx - async with prisma_client.tx() as tx: + async with prisma_client.tx() as tx_ctx: + tx: Final[_KeyUpdateTx] = cast( # cast-ok: the transaction object exposes the same table actions + "_KeyUpdateTx", tx_ctx + ) + update_input: Final[Mapping[str, object]] = ( + await _apply_object_permission_update( + data_json=non_default_values, + existing_object_permission_id=existing_key_row.object_permission_id, + prisma_client=prisma_client, + table=tx.litellm_objectpermissiontable, + ) + if expect_unassigned_project + else non_default_values + ) update_values: Final = await _apply_soft_budget_update( data=data, - non_default_values=non_default_values, + non_default_values=update_input, db=tx, existing_key_row=existing_key_row, changed_by=changed_by, ) include_object_permission: Final[prisma.types.LiteLLM_VerificationTokenInclude] = {"object_permission": True} - updated_row: Final = await tx.litellm_verificationtoken.update( - where=key_where, - data=with_settings_updated_at( - prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) - ), - include=include_object_permission, + update_payload: Final = with_settings_updated_at( + prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) + ) + updated_row: Final = ( + await _write_guarded_project_assignment( + table=tx.litellm_verificationtoken, + hashed_token=hashed_token, + existing_key_row=existing_key_row, + data=update_payload, + ) + if expect_unassigned_project + else await tx.litellm_verificationtoken.update( + where=key_where, + data=update_payload, + include=include_object_permission, + ) ) updated_data: Final[Mapping[str, object]] = ( updated_row.model_dump() if updated_row is not None else MappingProxyType({}) @@ -3563,10 +3752,6 @@ async def _validate_update_key_data( user_api_key_dict=user_api_key_dict, ) - if data.project_id is not None and data.project_id != existing_key_row.project_id: - raise HTTPException( - status_code=400, detail="Project reassignment is not supported. Use null to detach the key." - ) is_project_change: Final = "project_id" in data.model_fields_set and data.project_id != existing_key_row.project_id acting_as_team_admin: Final = await _acting_as_team_admin_for_key_update( @@ -3698,6 +3883,12 @@ async def _validate_update_key_data( user_api_key_cache=user_api_key_cache, route=("/key/update (max_budget/spend)" if _is_budget_change else "/key/update"), ) + await _validate_project_assignment( + data=data, + existing_key_row=existing_key_row, + prisma_client=checked_prisma_client, + user_api_key_cache=user_api_key_cache, + ) check_denied_passthrough_routes_caller_permission( data, user_api_key_dict, @@ -3847,7 +4038,7 @@ async def update_key_fn( - user_id: Optional[str] - User ID associated with key - team_id: Optional[str] - Team ID associated with key - agent_id: Optional[str] - The agent id associated with the key. - - project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected. + - project_id: Optional[str] - Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is rejected. - organization_id: Optional[str] - The organization id of the key. - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. - end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it. @@ -3995,23 +4186,31 @@ async def update_key_fn( if prisma_client is None: raise Exception("Not connected to DB!") + is_project_assignment: Final = data.project_id is not None and existing_key_row.project_id is None + update_fields: Final[Mapping[str, object]] = cast( # cast-ok: prepare_key_update_data returns a bare dict + "Mapping[str, object]", non_default_values + ) object_permission_upsert: Final = await _prepare_key_update_object_permission( - object_permission_data=non_default_values.get("object_permission"), + object_permission_data=update_fields.get("object_permission"), existing_key_row=existing_key_row, prisma_client=prisma_client, ) key_update_values: Final = MappingProxyType( - {field: value for field, value in non_default_values.items() if field != "object_permission"} + {field: value for field, value in update_fields.items() if field != "object_permission"} ) await _check_key_project_team_on_mutation( data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, ) - update_values: Final = await _write_prepared_key_update_object_permission( - data_json=key_update_values, - upsert=object_permission_upsert, - prisma_client=prisma_client, + update_values: Final = ( + update_fields + if is_project_assignment + else await _write_prepared_key_update_object_permission( + data_json=key_update_values, + upsert=object_permission_upsert, + prisma_client=prisma_client, + ) ) changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name response: Final = ( @@ -4022,8 +4221,16 @@ async def update_key_fn( non_default_values=update_values, existing_key_row=existing_key_row, changed_by=changed_by, + expect_unassigned_project=is_project_assignment, ) if "soft_budget" in data.model_fields_set + else await _update_key_row_assigning_project( + prisma_client=prisma_client, + key=key, + non_default_values=update_fields, + existing_key_row=existing_key_row, + ) + if is_project_assignment else await _update_key_row( prisma_client=prisma_client, key=key, diff --git a/litellm/types/proxy/management_endpoints/key_requests.py b/litellm/types/proxy/management_endpoints/key_requests.py index b49cb63c8e0..f42a412cf2d 100644 --- a/litellm/types/proxy/management_endpoints/key_requests.py +++ b/litellm/types/proxy/management_endpoints/key_requests.py @@ -118,7 +118,7 @@ class UpdateKeyRequest(KeyRequestBase): project_id: str | None = Field( default=None, - description="Omit to retain the project, or send null to detach. Assigning a different project is not supported.", + description="Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is not supported.", ) @model_validator(mode="before") diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index c26a4bf19b8..0198f582be7 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -67,6 +67,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _enforce_upperbound_key_params, _execute_virtual_key_regeneration, _get_and_validate_existing_key, + _hash_token_if_needed, _is_team_key, _key_generation_required_param_check, _member_key_request, @@ -77,6 +78,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _team_admin_requests_member_key, _team_key_generation_check, _transform_verification_tokens_to_deleted_records, + _update_key_row_assigning_project, _validate_end_user_budget_id_change, _validate_max_budget, _validate_reset_spend_value, @@ -20975,12 +20977,20 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( async def _cache_with_project( - project_id: str, project_models: list[str], team_id: str | None = "team-lit-5823" + project_id: str, project_models: list[str], team_id: str | None = "team-lit-5823", + project_max_budget: float | None = None, ) -> UserApiKeyCache: user_api_key_cache = UserApiKeyCache() await user_api_key_cache.async_set_cache( key=project_cache_key(project_id), - value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id=team_id, models=project_models), + value=LiteLLM_ProjectTableCachedObj( + project_id=project_id, + team_id=team_id, + models=project_models, + litellm_budget_table=( + LiteLLM_BudgetTable(max_budget=project_max_budget) if project_max_budget is not None else None + ), + ), model_type=LiteLLM_ProjectTableCachedObj, ) return user_api_key_cache @@ -21413,6 +21423,487 @@ async def test_project_detachment_uses_effective_project_for_validation_on_unown assert expected in str(exc.value.detail) +def _project_assignment_prisma(stored_row: Mapping[str, object] | None = None) -> MagicMock: + database = MagicMock() + team_row: Final = LiteLLM_TeamTable(team_id="team-lit-5823", members=[]) + database.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + database.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + if stored_row is None: + return database + + database.jsonify_object = lambda data: dict(data) + tx: Final = database.tx.return_value.__aenter__.return_value + + def _where_matches(where: Mapping[str, object], row: Mapping[str, object]) -> bool: + for key, clause in where.items(): + if key == "OR": + if not any(_where_matches(part, row) for part in clause): + return False + elif key == "models": + equals = clause["equals"] + if equals is None: + if row[key] is not None: + return False + elif row[key] is None or list(row[key]) != list(equals): + return False + elif row.get(key) != clause: + return False + return True + + async def _update_many(*, where, data): + if "object_permission" in data: + raise AssertionError("object_permission must not reach update_many") + return 1 if _where_matches(where, stored_row) else 0 + + tx.litellm_verificationtoken.update_many = AsyncMock(side_effect=_update_many) + written = MagicMock() + written.model_dump = MagicMock(return_value={**stored_row, "project_id": "project-orbit"}) + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=written) + return database + + +@pytest.mark.asyncio +async def test_project_assignment_to_unassigned_key_on_same_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + data: Final = UpdateKeyRequest(key=existing.token, project_id="project-orbit") + + await _validate_update_key_data( + data, existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + + update: Final = await prepare_key_update_data(data=data, existing_key_row=existing) + assert update["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_project_on_another_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-other", models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "team" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_key_with_no_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id=None, models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "team" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_nonexistent_project(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + database: Final = _project_assignment_prisma() + database.db.litellm_projecttable.find_unique = AsyncMock(return_value=None) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-ghost"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, database, UserApiKeyCache(), + ) + assert exc.value.status_code == 404 + assert "Project not found" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_validates_existing_models_against_project(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-elsewhere"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "not in project's allowed models" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_with_explicit_null_max_budget_ignores_stored_budget(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=[], max_budget=100.0 + ) + cache: Final = await _cache_with_project("project-orbit", [], project_max_budget=10.0) + + request: Final = UpdateKeyRequest(key=existing.token, project_id="project-orbit", max_budget=None) + await _validate_update_key_data( + request, existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + + result: Final = await prepare_key_update_data(data=request, existing_key_row=existing) + assert result["max_budget"] is None + assert result["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_without_max_budget_uses_stored_budget(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=[], max_budget=100.0 + ) + cache: Final = await _cache_with_project("project-orbit", [], project_max_budget=10.0) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "exceeds project's max_budget" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_includes_new_project_assignment(): + existing: Final = LiteLLM_VerificationToken(token="project-assign-token", project_id=None) + + result: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing_key_row=existing, + ) + + assert result["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_write_requires_row_still_unassigned(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=1) + row = MagicMock() + row.model_dump = MagicMock(return_value={"project_id": "project-orbit"}) + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=row) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + update_where: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["where"] + assert update_where["token"] == result["token"] + assert result["data"] == {"project_id": "project-orbit"} + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_row_gained_project(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=0) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + tx.litellm_verificationtoken.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_project_assignment_write_upserts_object_permission_in_same_tx(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + database.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=1) + permission_row = MagicMock() + permission_row.object_permission_id = "perm-lit-01" + tx.litellm_objectpermissiontable.upsert = AsyncMock(return_value=permission_row) + row = MagicMock() + row.model_dump = MagicMock(return_value={"project_id": "project-orbit"}) + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=row) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823" + ) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit", "object_permission": {"agents": ["agent-a"]}}, + existing_key_row=existing, + ) + + tx.litellm_objectpermissiontable.upsert.assert_awaited_once() + update_data: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["data"] + assert update_data["object_permission_id"] == "perm-lit-01" + assert "object_permission" not in update_data + assert result["data"] == {"project_id": "project-orbit"} + + +@pytest.mark.asyncio +async def test_project_assignment_write_matches_null_models_column(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": None, + "max_budget": None, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=[] + ) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert result["data"]["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_write_strips_explicit_null_object_permission(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": None, + "max_budget": None, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=[] + ) + update_values: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key="sk-assign", project_id="project-orbit", object_permission=None), + existing_key_row=existing, + ) + assert "object_permission" in update_values + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values=update_values, + existing_key_row=existing, + ) + + assert result["data"]["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_max_budget_changed(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": 50.0, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", + models=["model-orbit"], max_budget=100.0, + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_object_permission_changed(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": "perm-new", + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_team_changed(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-moved", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + + +def _assignment_project_update_prisma(stored_row: Mapping[str, object]) -> MagicMock: + database: Final = _project_assignment_prisma(stored_row=stored_row) + project_row: Final = LiteLLM_ProjectTable(project_id="project-orbit", team_id="team-lit-5823") + database.db.litellm_projecttable.find_unique = AsyncMock(return_value=project_row) + database.writer_db.litellm_projecttable.find_unique = AsyncMock(return_value=project_row) + return database + + +def _setup_assignment_update_mocks(monkeypatch: pytest.MonkeyPatch, database: MagicMock, cache: UserApiKeyCache) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", database) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.store_audit_logs", False) + + +@pytest.mark.asyncio +async def test_update_key_fn_assigns_project_through_guarded_write(monkeypatch: pytest.MonkeyPatch): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": None, + } + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + database: Final = _assignment_project_update_prisma(stored) + database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + _setup_assignment_update_mocks(monkeypatch, database, cache) + + result: Final = await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key="sk-assign", project_id="project-orbit"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + litellm_changed_by=None, + ) + + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many.assert_awaited_once() + update_where: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["where"] + assert update_where["project_id"] is None + tx.litellm_verificationtoken.update.assert_not_called() + assert result["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_update_key_fn_soft_budget_runs_guarded_assignment(monkeypatch: pytest.MonkeyPatch): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": None, + } + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + database: Final = _assignment_project_update_prisma(stored) + database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_budgettable.create = AsyncMock(return_value=SimpleNamespace(budget_id="budget-new")) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + _setup_assignment_update_mocks(monkeypatch, database, cache) + + result: Final = await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key="sk-assign", project_id="project-orbit", soft_budget=5.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + litellm_changed_by=None, + ) + + tx.litellm_budgettable.create.assert_awaited_once() + tx.litellm_verificationtoken.update_many.assert_awaited_once() + update_where: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["where"] + assert update_where["project_id"] is None + tx.litellm_verificationtoken.update.assert_not_called() + assert result["project_id"] == "project-orbit" + + @pytest.mark.asyncio async def test_key_creator_cannot_detach_project_without_admin_access(): existing: Final = LiteLLM_VerificationToken( @@ -21430,6 +21921,41 @@ async def test_key_creator_cannot_detach_project_without_admin_access(): assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) +@pytest.mark.asyncio +async def test_unauthorized_caller_gets_403_not_project_lookup_result(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", + user_id="user-orbit", created_by="user-orbit", + ) + database: Final = _project_assignment_prisma() + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-missing"), existing, + UserAPIKeyAuth(user_id="user-other", user_role=LitellmUserRoles.INTERNAL_USER), + None, False, database, UserApiKeyCache(), + ) + assert exc.value.status_code == 403 + assert "Project not found" not in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_unauthorized_caller_gets_403_not_project_team_leak(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", + user_id="user-orbit", created_by="user-orbit", + ) + cache: Final = await _cache_with_project("project-orbit", [], team_id="team-other") + database: Final = _project_assignment_prisma() + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_id="user-other", user_role=LitellmUserRoles.INTERNAL_USER), + None, False, database, cache, + ) + assert exc.value.status_code == 403 + assert "belongs to team" not in str(exc.value.detail) + + _OWNED_PROJECT: Final = "proj-owned-1" _OWNERSHIP_PROJECT_TEAM: Final = "ownership-project-team" _OWNERSHIP_KEY_TEAM: Final = "ownership-key-team" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index aafa732276b..96c7c34621b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8953,7 +8953,7 @@ export interface paths { * - user_id: Optional[str] - User ID associated with key * - team_id: Optional[str] - Team ID associated with key * - agent_id: Optional[str] - The agent id associated with the key. - * - project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected. + * - project_id: Optional[str] - Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is rejected. * - organization_id: Optional[str] - The organization id of the key. * - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. * - end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it. @@ -48704,7 +48704,7 @@ export interface components { policies?: string[] | null; /** * Project Id - * @description Omit to retain the project, or send null to detach. Assigning a different project is not supported. + * @description Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is not supported. */ project_id?: string | null; /** Prompts */ From b20acb48699e91fbe4d64a8878eaa0c92078a895 Mon Sep 17 00:00:00 2001 From: Maxime Boudier <113669264+MaximeBOUDIER@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:10:13 +0100 Subject: [PATCH 48/64] fix(proxy)!: persist sent model_info fields on POST /model/update (#46902) * fix(proxy): persist sent model_info fields on POST /model/update POST /model/update rebuilt litellm_params from the request but dropped the request's model_info, so changing base_model (or mode, tier, etc.) on an existing DB model was silently ignored. Clients that only use this endpoint, like the Terraform provider, could never fix a deployment's base_model without deleting and recreating it. Persist the model_info fields the caller actually sent, on top of the stored ones. ModelInfo defaults never overwrite stored values, server-owned, team, PTU and pricing fields stay on their dedicated write paths, and values echoed from a /model/info response are not pinned. * fix(proxy): keep echoed and blank model_info values out of POST /model/update writes Blank strings from typed clients no longer overwrite stored values. Values echoed from /model/info (the stored base model's cost map entry and discovered limits) are not pinned, and an echoed catalog value resets a differing stored override like PATCH does. direct_access and access_via_team_ids stay display-only. The tests stop patching the auth check class and give the team lookup a real row instead. * fix(proxy): lock the model row and skip unchanged model_info on POST /model/update * fix(proxy): treat only keyed /model/update bodies as echoes and write the stored model_info as saved A typed model_info value that matched the cost map, or the previous base model's entry, was dropped and could erase a stored override. Echo detection now follows PATCH, only a body carrying the /model/info key counts, compared against the base model the row has after the write. The locked write patches the stored JSON instead of dumping the hydrated Deployment, so litellm_params pricing is no longer copied into model_info. Tests run the real Router with discovery over a mock upstream * docs(terraform): note that model_info updates now take effect and overwrite values set outside Terraform * test: expect POST /model/update to save the cache point flag, and wrap long test lines test_x06 pinned the old contract, where POST /model/update answered 200 and ignored model_info. This PR makes POST save model_info the way PATCH does, so x06 now asserts what its PATCH sibling x01 asserts, the stored flag flips to true and cache points come back on. * fix(proxy): answer 400 when POST /model/update gets an access_groups that is not a list of strings * fix(proxy): read discovered model info through the router registry entry point --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../model_management_endpoints.py | 167 ++++++++- terraform/provider/CHANGELOG.md | 3 +- .../provider/litellm/resource_model_crud.go | 2 +- .../test_bedrock_kimi_k3_cache_point_wire.py | 18 +- .../test_model_management_endpoints.py | 345 +++++++++++++++++- 5 files changed, 506 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 74fa8e3b3c1..dbd77e45d57 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -486,6 +486,7 @@ async def _raise_on_invalid_credential_name( AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301 _CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" +_MODEL_ROW_LOCK_SQL: Final = 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = $1 FOR UPDATE' _STORED_LITELLM_PARAMS_SQL: Final = ( "(CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END)" ) @@ -633,6 +634,7 @@ async def _auto_router_capability_slot( effective_params: Mapping[str, object], model_id: str | None, member_write: MemberAutoRouterWrite | None = None, + lock_model_row: bool = False, ) -> AsyncGenerator[_ProxyModelTable, None]: """Hand out the model table to write through while the row's claim on a licensed capability is settled. @@ -649,6 +651,9 @@ async def _auto_router_capability_slot( A heuristic-v1 router whose tuning has moved off its recorded baseline is judged the same way under the same lock, against the DB rows plus this proxy's config.yaml routers. + + ``lock_model_row`` runs the write in a transaction holding the row's ``FOR UPDATE`` lock, + so a caller that re-reads the row through the yielded table writes onto its committed state. """ from litellm.proxy.proxy_server import ( _license_check, # pyright: ignore[reportPrivateUsage] # existing capability slot reads the proxy license singleton @@ -662,7 +667,8 @@ async def _auto_router_capability_slot( baselines: Final = heuristic_v1_tuning_baselines tuning_candidate: Final = _tuning_candidate(effective_params, model_id=model_id) judges_tuning: Final = baselines is not None and is_mutable_tuned_candidate(tuning_candidate, baselines) - if member_write is None and (limit is None or (not capabilities and not judges_tuning)): + checks_license: Final = member_write is not None or (limit is not None and (bool(capabilities) or judges_tuning)) + if not checks_license and not lock_model_row: yield _proxy_model_table(prisma_client) return transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db) @@ -673,14 +679,15 @@ async def _auto_router_capability_slot( ) async with transaction as tx_ctx: tables: Final[_TxModelTables] = tx_ctx - await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) + if checks_license: + await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) + locked_model_id: Final = ( + member_write.model_id if member_write is not None else (model_id if lock_model_row else None) + ) + if locked_model_id is not None: + await tx_ctx.query_raw(_MODEL_ROW_LOCK_SQL, locked_model_id) config_rows: Final = () if llm_router is None else tuple(llm_router.registry().deployments(source="config")) if member_write is not None: - if member_write.model_id is not None: - await tx_ctx.query_raw( - 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = $1 FOR UPDATE', - member_write.model_id, - ) pinned_client: Final = _TransactionClient(tx_ctx) team_where: Final[prisma_types.LiteLLM_TeamTableWhereUniqueInput] = {"team_id": member_write.team_id} team_include: Final[prisma_types.LiteLLM_TeamTableInclude] = {"litellm_model_table": True} @@ -752,7 +759,8 @@ async def _auto_router_capability_slot( prisma_client=pinned_client, llm_router=llm_router, ) - for capability in capabilities: + judged_capabilities: Final = capabilities if checks_license else () + for capability in judged_capabilities: if ( violation := await _claimed_capability_violation( tables, capability=capability, config_rows=config_rows, model_id=model_id, limit=limit @@ -761,7 +769,7 @@ async def _auto_router_capability_slot( raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}" ) - if judges_tuning and baselines is not None: + if checks_license and judges_tuning and baselines is not None: model_rows: Final = await ModelRepository(_TransactionClient(tx_ctx)).find_all_except(model_id or "") _raise_on_tuning_quota_violation( candidate=tuning_candidate, @@ -1049,6 +1057,16 @@ def _raise_if_ptu_cost_attribution_disabled(incoming_model_info: Mapping[str, ob ) +_ACCESS_GROUPS: Final[TypeAdapter[tuple[str, ...] | None]] = TypeAdapter(tuple[str, ...] | None) + + +def _raise_if_access_groups_malformed(model_info_writes: Mapping[str, object]) -> None: + try: + _ACCESS_GROUPS.validate_python(model_info_writes.get("access_groups")) + except ValidationError: + raise HTTPException(status_code=400, detail="model_info.access_groups must be a list of strings") from None + + def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: """Enforce the PTU cross-field invariant on the effective model_info. @@ -1291,6 +1309,109 @@ def _loaded_catalog_entry( return loaded_catalog().get(catalog_key, MappingProxyType({})) +_LEGACY_UPDATE_IGNORED_MODEL_INFO_FIELDS: Final = frozenset( + { + "id", + "db_model", + "created_at", + "created_by", + "updated_at", + "updated_by", + "blocked", + "member_auto_router", + "team_id", + "team_public_model_name", + "direct_access", + "access_via_team_ids", + COST_MAP_LOOKUP_KEY, + *PTU_MODEL_INFO_FIELDS, + } +) + + +def legacy_update_model_info_changes( + db_model: Deployment, + incoming_model_info: ModelInfo, + discovered_model_info: Mapping[str, object] = MappingProxyType({}), + loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, +) -> Mapping[str, object]: + """The ``model_info`` fields ``POST /model/update`` persists from its request. + + Only fields the caller sent with a value count, so ``ModelInfo`` defaults and the blank + strings typed clients send for unset fields never overwrite stored values. Server-owned, + display-only, team, PTU and pricing fields keep their dedicated write paths. As on + ``update_db_model``, only a body carrying the ``key`` that ``/model/info`` emits is an echo: + its values equal to the cost map entry of the base model the row has after the write, the + entry ``key`` names, or a limit discovered from the upstream stay display-only, and an echoed + value that differs from the stored override resets that override (``None``). + """ + stored: Final = db_model.model_info.model_dump(exclude_none=True) + requested: Final = incoming_model_info.model_dump(exclude_unset=True, exclude_none=True) + sent: Final = {k: v for k, v in requested.items() if k not in _LEGACY_UPDATE_IGNORED_MODEL_INFO_FIELDS and v != ""} + echoed: Final = frozenset( + echoed_cost_map_fields( + requested, + _cost_map_entry(db_model, {**stored, **sent}), + _loaded_catalog_entry(requested, loaded_catalog), + discovered_model_info, + ) + ) + kept: Final = without_server_derived_pricing(sent) + return MappingProxyType( + { + **{k: getattr(incoming_model_info, k) for k in kept if k not in echoed}, + **{k: None for k, v in kept.items() if k in echoed and stored.get(k) not in (None, v)}, + } + ) + + +def _legacy_update_model_info_writes( + current: Deployment, + incoming_model_info: ModelInfo, + discovered_model_info: Mapping[str, object], + member_marker: bool | None, +) -> Mapping[str, object]: + changes: Final = { + **legacy_update_model_info_changes(current, incoming_model_info, discovered_model_info), + **({"member_auto_router": member_marker} if member_marker is not None else {}), + } + return MappingProxyType({k: v for k, v in changes.items() if getattr(current.model_info, k, None) != v}) + + +async def _locked_legacy_update_model_info_json( + table: _ProxyModelTable, + model_id: str, + incoming_model_info: ModelInfo, + discovered_model_info: Mapping[str, object], + member_marker: bool | None, +) -> str | None: + row: Final = await table.find_unique(where={"model_id": model_id}) + if row is None: + return None + payload: Final = row.model_dump() + writes: Final = _legacy_update_model_info_writes( + Deployment(**payload), incoming_model_info, discovered_model_info, member_marker + ) + if not writes: + return None + stored_model_info: Final[object] = payload.get("model_info") + stored: Final = ( + TypeAdapter(dict[str, object]).validate_json(stored_model_info) + if isinstance(stored_model_info, str) + else TypeAdapter(dict[str, object]).validate_python(stored_model_info or {}) + ) + return ( + TypeAdapter(dict[str, object]) + .dump_json( + { + **{k: v for k, v in stored.items() if k not in writes}, + **{k: v for k, v in writes.items() if v is not None}, + } + ) + .decode() + ) + + def _param_keys_to_remove(patch: updateDeployment) -> tuple[str, ...]: return patch.litellm_params_to_remove if isinstance(patch, PatchModelRequest) else () @@ -3279,22 +3400,32 @@ async def update_model( renamed_update: Final[PrismaCompatibleUpdateDBModel] = ( {**base_update, "model_name": renamed_to} if renamed_to is not None else base_update ) - _data: Final[PrismaCompatibleUpdateDBModel] = ( - { - **renamed_update, - "model_info": deployment.model_info.model_copy( - update=MappingProxyType({"member_auto_router": member_marker}) - ).model_dump_json(exclude_none=True), - } - if member_marker is not None - else renamed_update + discovered_model_info: Final[Mapping[str, int]] = ( + llm_router.registry().discovered_model_info(_model_id) + if llm_router is not None + else MappingProxyType({}) ) + requested_model_info: Final = _legacy_update_model_info_writes( + deployment, _model_info, discovered_model_info, member_marker + ) + _raise_if_access_groups_malformed(requested_model_info) async with _auto_router_capability_slot( prisma_client, effective_params=effective_params, model_id=_model_id, member_write=member_write, + lock_model_row=bool(requested_model_info), ) as table: + model_info_json: Final = ( + await _locked_legacy_update_model_info_json( + table, _model_id, _model_info, discovered_model_info, member_marker + ) + if requested_model_info + else None + ) + _data: Final[PrismaCompatibleUpdateDBModel] = ( + {**renamed_update, "model_info": model_info_json} if model_info_json is not None else renamed_update + ) model_response: Final = await table.update( where={"model_id": _model_id}, data=_data, diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 597e0c6d4d8..4a1838a7b60 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -18,7 +18,7 @@ longer signal it. - **credential, model, mcp_server**: Write-only secret attributes so secrets never land in Terraform state (Terraform 1.11 or later). `litellm_credential` gains `credential_values_wo` (a JSON object, `jsonencode({...})`, since the SDK cannot mark a map write-only) with `credential_values_wo_version`, and `credential_values` becomes optional with exactly one of the two required. `litellm_model` gains `model_api_key_wo`, `aws_secret_access_key_wo` and `vertex_credentials_wo`, each with a `_wo_version` companion and each conflicting with its in-state twin. `litellm_mcp_server` gains `client_id`, `client_secret_wo` and `auth_value_wo` (with `_wo_version` companions), sent to the proxy as `credentials`, and `auth_type` now accepts every value the proxy does (`oauth2`, `api_key`, `bearer_token`, ...). Every `_wo` value is sent on each apply from the configuration and produces no diff, so bump its `_wo_version` to push a rotated secret. `litellm_vector_store.litellm_params` warns at plan time when it carries a secret-looking key such as `api_key` - **team**: Optional `team_member_model_max_budget` argument on `litellm_team`, a JSON string of per-model budgets applied to every team member and sent as `team_member_model_max_budget` on `/team/new` and `/team/update`; read back from `team_member_budget_table.model_max_budget` on `/team/info` so out-of-band changes show as drift, with `budget_limit`/`time_period` treated as equal to the `max_budget`/`budget_duration` spelling the proxy stores; removing it sends `{}`, which clears every per-model member budget; `tpm_limit`/`rpm_limit` are rejected at plan time because the proxy rejects them on team budgets. Needs a proxy that includes [#43887](https://github.com/BerriAI/litellm/pull/43887): v1.105.0 and v1.106.0-rc.1 accept the field but silently drop it, so every plan shows a diff -- **model**: Optional `display_name` argument on `litellm_model`, sent as `model_info.display_name` and returned as `display_name` by `/v1/models`, so client model pickers show a readable name; changes are persisted through `/model/{id}/update` since `/model/update` ignores `model_info`; also exported by the `litellm_model` and `litellm_models` data sources +- **model**: Optional `display_name` argument on `litellm_model`, sent as `model_info.display_name` and returned as `display_name` by `/v1/models`, so client model pickers show a readable name; changes are also persisted through `/model/{id}/update`, since proxies older than this provider drop the `model_info` sent to `/model/update`; also exported by the `litellm_model` and `litellm_models` data sources - **key**: Computed `server_metadata` attribute on `litellm_key` exposing every metadata entry the proxy stores, so metadata created outside Terraform is visible in state and drift on it shows on refresh, while `metadata` keeps tracking only the declared entries and updates keep preserving undeclared ones - **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them - **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement @@ -70,6 +70,7 @@ longer signal it. ### Changed +- **model** (breaking): on a proxy that includes [#46902](https://github.com/BerriAI/litellm/pull/46902), `litellm_model` updates now save the `model_info` the provider sends to `/model/update`, so `pricing_base_model`, `tier` and `mode` changes take effect. The first apply after upgrading the proxy overwrites values set outside Terraform with the configuration's, including `tier`, which defaults to `free`, and `base_model`, which falls back to the deployment name when `pricing_base_model` is unset. Run `terraform plan` before that apply and copy any value to keep into the configuration - **model** (breaking): a key named in `additional_litellm_params.additional_drop_params` is no longer deleted from the deployment's `litellm_params`. A configuration that also sets that key, such as `vertex_location`, stores it on the deployment at its next apply, and the proxy applies the drop list to requests instead. For `vertex_location` this changes routing. `/v1/messages` calls now go to the configured region, while chat completions drop a listed `vertex_location` and use the default region. To keep such a key off a deployment, remove its attribute and replace the model with `terraform apply -replace`, which gives it a new model id, since an update keeps a stored value the provider leaves out. `model`, `custom_llm_provider` and `merge_reasoning_content_in_choices` are always sent. A value that is not a JSON list of names, such as a bare `"store"` or an empty string, now fails the plan, and with it `terraform destroy`, instead of being stored as a string the proxy ignores. Rewrite it as `jsonencode([...])`, or remove the model without fixing it through `terraform destroy -refresh=false` - **key** (breaking): `model_max_budget` on `litellm_key` is now a JSON string of per-model budget objects (`jsonencode({"gpt-4o-mini" = {budget_limit = 50, time_period = "30d"}})`), matching `litellm_user`, `litellm_budget` and `litellm_tag`. The old `map(number)` form sent bare numbers to `/key/generate`, which the proxy rejects with a 500 (`'int' object is not iterable`), so every key with a non-empty `model_max_budget` failed to apply. Existing state upgrades automatically (schema version 1) and the attribute is refilled from the proxy on the next read; configurations still using the map form must be rewritten - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index 3db6f68760d..1f5aa62560f 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -357,7 +357,7 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e return retryModelRead(d, m, 5) } -// /model/update only merges litellm_params, so model_info changes go through the PATCH endpoint. +// Older proxies drop the model_info sent to /model/update, so display_name also goes through the PATCH endpoint. func patchModelDisplayName(client *Client, modelID, displayName string) error { resp, err := MakeRequest(client, "PATCH", fmt.Sprintf(endpointModelPatch, url.PathEscape(modelID)), ModelInfoPatch{ ModelInfo: ModelInfoPatchFields{ID: modelID, DisplayName: displayName}, diff --git a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py index 0772e928b35..4441582fb6b 100644 --- a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py +++ b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py @@ -1325,7 +1325,9 @@ def test_x01_flipping_the_flag_to_true_through_a_patch_update_turns_cache_points @pytest.mark.timeout(600) -def test_x06_the_legacy_post_update_answers_200_and_leaves_the_stored_flag_alone(gateway: Gateway) -> None: +def test_x06_flipping_the_flag_to_true_through_the_legacy_post_update_turns_cache_points_back_on( + gateway: Gateway, +) -> None: arn: Final = _arn("legacy") with wire_server(_peer()) as wire, gateway.scenario() as scenario: name: Final = _deployment( @@ -1333,6 +1335,7 @@ def test_x06_the_legacy_post_update_answers_200_and_leaves_the_stored_flag_alone ) before, before_received = _observe_at(gateway, wire, "chat", _chat_body(name, _prompt(), _SYSTEM_AND_USER)) _assert_answered(before) + assert before_received.cache_points == 0, before_received.body identity: Final = _model_id(name) updated: Final = gateway.request( "POST", @@ -1344,13 +1347,14 @@ def test_x06_the_legacy_post_update_answers_200_and_leaves_the_stored_flag_alone }, ) assert updated.status_code == 200, updated.text - assert _stored_flag(identity) is False - _settled(gateway, name, wire) - after, after_received = _observe_at(gateway, wire, "chat", _chat_body(name, _prompt(), _SYSTEM_AND_USER)) - _assert_answered(after) - assert after_received.cache_points == before_received.cache_points, (before_received.body, after_received.body) + assert _stored_flag(identity) is True + points: Final = eventually( + lambda: _six_cache_point_counts(gateway, wire, name), + lambda counts: len(counts) == 6 and all(count == 2 for count in counts), + seconds=60, + ) + assert points == (2,) * 6, points _success_row(before.response_id) - _success_row(after.response_id) @dataclass(frozen=True, slots=True) diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 872d1607351..abaa802de2c 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -3,17 +3,19 @@ import contextlib import inspect import json import socket -from collections.abc import Iterator, Mapping -from types import SimpleNamespace +from collections.abc import AsyncIterator, Iterator, Mapping +from types import MappingProxyType, SimpleNamespace from typing import Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException from fastapi.encoders import jsonable_encoder from fastapi.testclient import TestClient from litellm._uuid import uuid +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.models.credentials import CredentialItem from litellm.models.project import LiteLLM_ProjectTable @@ -1639,11 +1641,350 @@ class TestDeletedAutoRouterAvailability: } +class _ContendedModelDb: + def __init__(self, row: LiteLLM_ProxyModelTable, other_writers: tuple[Mapping[str, object], ...]) -> None: + self.row = row + self.writes: tuple[Mapping[str, object], ...] = () + self._other_writers = other_writers + self._row_locked = False + self.db = self + self.litellm_proxymodeltable = self + self.litellm_teamtable = SimpleNamespace(find_unique=self._find_team) + + async def _find_team(self, *, where: Mapping[str, str]) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=where["team_id"]) + + async def find_unique(self, *, where: Mapping[str, object]) -> LiteLLM_ProxyModelTable: + read: Final = self.row + if not self._row_locked and self._other_writers: + self._commit_next_other_writer() + return read + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> LiteLLM_ProxyModelTable: + self.writes = (*self.writes, data) + self.row = self.row.model_copy( + update={k: json.loads(v) if k in ("litellm_params", "model_info") else v for k, v in data.items()} + ) + return self.row + + async def query_raw(self, query: str, *args: object) -> tuple[()]: + self._row_locked = self._row_locked or "FOR UPDATE" in query + return () + + @contextlib.asynccontextmanager + async def tx(self, timeout: object = None) -> AsyncIterator["_ContendedModelDb"]: + yield self + self._row_locked = False + while self._other_writers: + self._commit_next_other_writer() + + def _commit_next_other_writer(self) -> None: + self.row = self.row.model_copy( + update={"model_info": {**(self.row.model_info or {}), **self._other_writers[0]}} + ) + self._other_writers = self._other_writers[1:] + + +_AZURE_DEPLOYMENT_PARAMS: Final = MappingProxyType({"model": "azure/my-deployment", "api_key": "sk-existing"}) +_LOCAL_DEPLOYMENT_PARAMS: Final = MappingProxyType( + {"model": "hosted_vllm/local-model", "api_base": "https://local.test/v1", "api_key": "local-key"} +) + + +async def _discover_from_upstream(router: Router, model_card: Mapping[str, object]) -> None: + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"data": [dict(model_card)]})) + ) as client: + handler.client = client + await router.arefresh_model_info(client=handler) + + class TestUpdateModel: """ Tests for the update_model (POST /model/update) handler. """ + @staticmethod + async def _update_model( + stored_model_info: Mapping[str, object], + model_info: ModelInfo, + litellm_params: Mapping[str, object] = _AZURE_DEPLOYMENT_PARAMS, + upstream_model_card: Mapping[str, object] | None = None, + other_writers: tuple[Mapping[str, object], ...] = (), + ) -> "_ContendedModelDb": + db: Final = _ContendedModelDb( + LiteLLM_ProxyModelTable( + model_id="m1", + model_name="my-model", + litellm_params=dict(litellm_params), + model_info=dict(stored_model_info), + ), + other_writers, + ) + router: Final = Router( + model_list=[ + {"model_name": "my-model", "litellm_params": dict(litellm_params), "model_info": {"id": "m1"}} + ] + ) + if upstream_model_card is not None: + await _discover_from_upstream(router, upstream_model_card) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", db), + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( # test-quality-ok: [TQ008] the route reloads the proxy-global router with no seam to inject + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await update_model( + model_params=updateDeployment(litellm_params=updateLiteLLMParams(), model_info=model_info), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + return db + + @pytest.mark.asyncio + async def test_update_model_persists_sent_model_info_over_stored_fields(self): + stored: Final = { + "id": "m1", + "db_model": True, + "base_model": "my-deployment", + "tier": "paid", + "team_id": "team-a", + } + + db: Final = await self._update_model( + stored, + ModelInfo(id="m1", base_model="test-provider/test-model", mode="chat", db_model=False, team_id="team-b"), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert persisted["mode"] == "chat" + assert {k: persisted[k] for k in ("db_model", "tier", "team_id")} == { + "db_model": True, + "tier": "paid", + "team_id": "team-a", + } + + @pytest.mark.asyncio + async def test_update_model_with_only_an_id_leaves_model_info_untouched(self): + db: Final = await self._update_model({"id": "m1", "base_model": "my-deployment"}, ModelInfo(id="m1")) + + assert tuple("model_info" in write for write in db.writes) == (False,) + + @pytest.mark.asyncio + async def test_update_model_resending_the_stored_model_info_does_not_rewrite_it(self): + db: Final = await self._update_model( + {"id": "m1", "db_model": True, "tier": "paid", "mode": "chat"}, + ModelInfo(id="m1", tier="paid", mode="chat"), + ) + + assert tuple("model_info" in write for write in db.writes) == (False,) + + @pytest.mark.asyncio + async def test_update_model_keeps_model_info_fields_a_concurrent_writer_stored(self): + db: Final = await self._update_model( + {"id": "m1", "db_model": True, "mode": "embedding"}, + ModelInfo(id="m1", mode="chat"), + other_writers=(MappingProxyType({"access_groups": ["group-a"]}), MappingProxyType({"tier": "paid"})), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert {k: persisted.get(k) for k in ("mode", "access_groups", "tier")} == { + "mode": "chat", + "access_groups": ["group-a"], + "tier": "paid", + } + + @pytest.mark.asyncio + async def test_update_model_does_not_pin_values_echoed_from_model_info(self): + import litellm + + catalog_entry: Final = litellm.get_model_info(model="azure/gpt-4o") + echoed: Final = ModelInfo( + id="m1", + base_model="azure/gpt-4o", + key=catalog_entry["key"], + max_tokens=catalog_entry["max_tokens"], + input_cost_per_token=catalog_entry["input_cost_per_token"], + ) + + db: Final = await self._update_model({"id": "m1"}, echoed) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "azure/gpt-4o" + assert not {"key", "max_tokens", "input_cost_per_token"} & persisted.keys() + + @pytest.mark.asyncio + async def test_update_model_saves_typed_values_that_equal_the_catalog(self): + import litellm + + catalog_entry: Final = litellm.get_model_info(model="azure/gpt-4o") + typed: Final = { + "max_tokens": catalog_entry["max_tokens"], + "max_input_tokens": catalog_entry["max_input_tokens"], + "supports_vision": catalog_entry["supports_vision"], + } + + db: Final = await self._update_model( + {"id": "m1", "max_input_tokens": catalog_entry["max_input_tokens"] // 2}, + ModelInfo(id="m1", base_model="azure/gpt-4o", **typed), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert {k: persisted.get(k) for k in typed} == typed + + @pytest.mark.asyncio + async def test_update_model_keeps_stored_value_when_a_field_is_sent_blank(self): + db: Final = await self._update_model( + {"id": "m1", "base_model": "my-deployment", "mode": "chat"}, + ModelInfo(id="m1", base_model="test-provider/test-model", mode=""), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert persisted["mode"] == "chat" + + @pytest.mark.asyncio + async def test_update_model_ignores_ptu_blocked_and_pricing_fields(self): + db: Final = await self._update_model( + {"id": "m1"}, + ModelInfo( + id="m1", base_model="test-provider/test-model", ptu_count=5, blocked=True, input_cost_per_token=0.5 + ), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert not {"ptu_count", "blocked", "input_cost_per_token"} & persisted.keys() + + @pytest.mark.asyncio + @pytest.mark.parametrize("access_groups", [5, "group-a", ["group-a", 5], {"group": "group-a"}]) + async def test_update_model_rejects_access_groups_that_are_not_a_list_of_strings(self, access_groups: object): + with pytest.raises(ProxyException) as error: + await self._update_model({"id": "m1"}, ModelInfo.model_validate({"id": "m1", "access_groups": access_groups})) + + assert error.value.code == "400" + + @pytest.mark.asyncio + async def test_update_model_saves_access_groups_sent_as_a_list_of_strings(self): + db: Final = await self._update_model( + {"id": "m1"}, ModelInfo.model_validate({"id": "m1", "access_groups": ["group-a", "group-b"]}) + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["access_groups"] == ["group-a", "group-b"] + + def test_update_model_does_not_pin_the_previous_base_models_echoed_values(self): + from litellm.proxy.management_endpoints.model_management_endpoints import legacy_update_model_info_changes + + catalog: Final = MappingProxyType( + {"test-provider/previous": MappingProxyType({"supported_openai_params": ["previous-only-param"]})} + ) + stored: Final = Deployment( + model_name="my-model", + litellm_params=LiteLLM_Params(model="test-provider/my-deployment"), + model_info=ModelInfo(id="m1", base_model="test-provider/previous"), + ) + + changes: Final = legacy_update_model_info_changes( + stored, + ModelInfo( + id="m1", + base_model="test-provider/next", + key="test-provider/previous", + supported_openai_params=["previous-only-param"], + ), + loaded_catalog=lambda: catalog, + ) + + assert dict(changes) == {"base_model": "test-provider/next"} + + @pytest.mark.asyncio + async def test_update_model_saves_a_value_the_previous_base_model_has_in_the_catalog(self): + import litellm + + previous: Final = litellm.get_model_info(model="azure/gpt-4o") + + db: Final = await self._update_model( + {"id": "m1", "base_model": "azure/gpt-4o", "max_input_tokens": previous["max_input_tokens"] // 2}, + ModelInfo(id="m1", base_model="azure/gpt-4.1", max_input_tokens=previous["max_input_tokens"]), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert {k: persisted.get(k) for k in ("base_model", "max_input_tokens")} == { + "base_model": "azure/gpt-4.1", + "max_input_tokens": previous["max_input_tokens"], + } + + @pytest.mark.asyncio + async def test_update_model_does_not_pin_discovered_limits_or_team_access_echoed_from_model_info(self): + db: Final = await self._update_model( + {"id": "m1"}, + ModelInfo( + id="m1", + base_model="test-provider/test-model", + key="test-provider/test-model", + max_input_tokens=128000, + max_output_tokens=64000, + direct_access=True, + access_via_team_ids=["team-a"], + ), + litellm_params=_LOCAL_DEPLOYMENT_PARAMS, + upstream_model_card=MappingProxyType( + {"id": "local-model", "max_model_len": 128000, "max_output_tokens": 32000} + ), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert persisted["max_output_tokens"] == 64000 + assert not {"max_input_tokens", "direct_access", "access_via_team_ids"} & persisted.keys() + + @pytest.mark.asyncio + async def test_update_model_does_not_copy_litellm_params_pricing_into_model_info(self): + db: Final = await self._update_model( + {"id": "m1"}, + ModelInfo(id="m1", tier="paid"), + litellm_params={**_AZURE_DEPLOYMENT_PARAMS, "input_cost_per_token": 0.000003}, + ) + + persisted: Final = db.row.model_info + assert persisted == {"id": "m1", "tier": "paid"} + + @pytest.mark.asyncio + async def test_update_model_resets_a_stored_override_echoed_back_at_its_catalog_value(self): + import litellm + + catalog_entry: Final = litellm.get_model_info(model="azure/gpt-4o") + stored_override: Final = catalog_entry["max_tokens"] // 2 + + db: Final = await self._update_model( + {"id": "m1", "base_model": "azure/gpt-4o", "max_tokens": stored_override}, + ModelInfo( + id="m1", base_model="azure/gpt-4o", key=catalog_entry["key"], max_tokens=catalog_entry["max_tokens"] + ), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert "max_tokens" not in persisted + @pytest.mark.asyncio async def test_update_model_clears_cache_after_db_write(self): """ From 6ad4f9cfdf3f2c4a8db3e39acb02d3d7730d2dce Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Fri, 30 Oct 2026 12:14:17 -0700 Subject: [PATCH 49/64] ci: cache Rust compilation with sccache (#49117) This should make compilation much faster --- .github/workflows/test-rust-db.yml | 10 ++++++++++ .github/workflows/test-rust.yml | 20 ++++++++++++++++++++ .github/workflows/test-unit.yml | 12 ++++++++++++ 3 files changed, 42 insertions(+) diff --git a/.github/workflows/test-rust-db.yml b/.github/workflows/test-rust-db.yml index bb88b593d5b..80ec256fa6f 100644 --- a/.github/workflows/test-rust-db.yml +++ b/.github/workflows/test-rust-db.yml @@ -42,6 +42,11 @@ concurrency: env: CARGO_TERM_COLOR: always + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" jobs: db-queries: @@ -87,6 +92,11 @@ jobs: workspaces: litellm-rust cache-on-failure: true + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Install Python dependencies working-directory: . run: uv sync --frozen --extra extra-proxy diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 04af770b24f..cf9fed192f6 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -78,6 +78,11 @@ concurrency: env: CARGO_TERM_COLOR: always + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" jobs: rust-lint: @@ -101,6 +106,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - run: cargo clippy --workspace --all-targets --locked -- -D warnings rust-test: @@ -143,6 +153,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Check generated trace contracts working-directory: . run: uv run scripts/generate_trace_types.py --check @@ -189,6 +204,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - run: uv build --wheel --out-dir dist - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 68823d99519..5ec5975d83b 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -73,9 +73,21 @@ jobs: timeout-minutes: 5 uses: ./.github/actions/cache-cargo-build + - name: Set up sccache + if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true' + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Build editable Rust bridge if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true' timeout-minutes: 15 + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" run: | mkdir -p rust-bridge-dist uv run --no-project --with maturin==1.15.0 python -c 'import maturin; maturin.build_editable("rust-bridge-dist")' From 5898a42f5aac37fc9fa2392f3c7fe16644e3e958 Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Fri, 30 Oct 2026 12:19:51 -0700 Subject: [PATCH 50/64] fix(openrouter)!: list OpenRouter's live catalog for openrouter/* when check_provider_endpoint is on (#48820) * fix(openrouter): list OpenRouter's live catalog for openrouter/* when check_provider_endpoint is on * fix(openrouter): list models off the event loop and keep only chat-callable catalog ids The model-list builders are sync, and with check_provider_endpoint on they call OpenRouter's catalog with a 10 second timeout and no caching of failures. They ran on the event loop inside /v1/models, /model/info, /model_group/info and the per-user listing, so one slow catalog froze every request on the worker. Those calls now run in a worker thread with asyncio.to_thread File, batch and RAG credential lookups only need each name's deployment, and a wildcard that is itself a deployment sits in the list ahead of its expansions and resolves its own credentials. get_complete_model_list takes expand_deployed_wildcard_routes, which that lookup sets to False, so a deployed wildcard never fetches its provider catalog there. A key wildcard wider than any deployment, like anthropic/* over a deployed anthropic/claude-*, still expands so it finds the matching deployment's credentials. Every other caller keeps the default The pattern router's model lookup tolerates a pattern removed between its two reads of the pattern table, which a deployment delete on the event loop can now do while a listing runs in a thread Listing now reads only OpenRouter's models/user catalog. OpenRouter's own /models list carries no embedding models, the OpenAI-format /v1/models has no mode field so chat pickers cannot tell embeddings apart, /model_group/info labels them mode chat, and one failing catalog emptied the whole list. Ids ending in :batch are dropped because chat completions answers 404 for them, while :free and every other variant stay. An api_base that already ends in /chat/completions, which chat accepts as is, has that suffix stripped before the catalog route is joined * fix(proxy): keep model listing off the shared thread pool and out of credential lookups Listings run on their own thread pool so many slow catalogs cannot starve the default executor that chat and Responses calls share. File and batch credential lookups match deployed names against a key's wildcard grants instead of listing any provider catalog. The token counter keeps OpenRouter deployments on their previous path, and a catch-all "*" over provider/* lists bare ids that route back through it * test(proxy): give the off-loop listing tests a generous deadlock guard and typed ids * fix(proxy): keep routed models in wildcard model info, skip key-only OpenRouter catalogs, cache failed listings briefly * fix(utils): cache failed model listings under their own key so a late failure never hides a good listing * fix(proxy): remember only timed-out provider listings so a provider that failed fast is listed again * fix(proxy): ask each provider catalog once per worker, keep deployed names out of catch-all ids, and match file credentials by name Listings that start together for one deployment now wait for the first catalog request. A catch-all wildcard no longer lists a name a deployment serves by name. A key wildcard only finds file and batch credentials on deployments named under it, so a bare name keeps the default credentials as before. * fix(proxy): keep the listing failure log line the same for timed-out listings * refactor(proxy): fold the cached listing lookup back into the provider listing function * fix(proxy): list each model id once and keep a named deployment's row over a wildcard's * fix(proxy): drop a wildcard's model info row when named deployments serve every name it lists --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/openrouter/common_utils.py | 87 +++++++ litellm/proxy/auth/model_checks.py | 74 +++++- .../openai_files_endpoints/common_utils.py | 53 +++-- litellm/proxy/proxy_server.py | 59 +++-- litellm/proxy/utils.py | 55 +++-- .../router_utils/pattern_match_deployments.py | 4 +- litellm/utils.py | 81 +++++-- .../test_openrouter_common_utils.py | 159 +++++++++++++ tests/unit/proxy/auth/test_model_checks.py | 219 ++++++++++++++++++ .../test_files_common_utils.py | 112 +++++++++ .../proxy_server/test_routes_model_info.py | 12 +- .../proxy/proxy_server/test_routes_models.py | 73 ++++++ tests/unit/proxy/test_proxy_token_counter.py | 33 +++ .../proxy/utils/helpers/test_model_access.py | 67 +++++- tests/unit/test_utils.py | 71 ++++++ 15 files changed, 1055 insertions(+), 104 deletions(-) create mode 100644 tests/unit/llms/openrouter/test_openrouter_common_utils.py diff --git a/litellm/llms/openrouter/common_utils.py b/litellm/llms/openrouter/common_utils.py index 96e53a5aae7..64caf83d745 100644 --- a/litellm/llms/openrouter/common_utils.py +++ b/litellm/llms/openrouter/common_utils.py @@ -1,5 +1,92 @@ +from collections.abc import Mapping, Sequence +from typing import Final + +import httpx +from pydantic import BaseModel, ConfigDict + +import litellm +from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues + +_DEFAULT_API_BASE: Final = "https://openrouter.ai/api/v1" +_CHAT_CATALOG_ROUTE: Final = "models/user" +_CHAT_COMPLETIONS_PATH: Final = "/chat/completions" +_BATCH_ONLY_SUFFIX: Final = ":batch" +_CATALOG_TIMEOUT: Final = httpx.Timeout(10.0) class OpenRouterException(BaseLLMException): pass + + +class _CatalogModel(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + + +class _Catalog(BaseModel): + model_config = ConfigDict(frozen=True) + + data: tuple[_CatalogModel, ...] + + +def _chat_catalog_url(api_base: str) -> str: + return join_api_base(api_base.rstrip("/").removesuffix(_CHAT_COMPLETIONS_PATH), _CHAT_CATALOG_ROUTE) + + +class OpenRouterModelInfo(BaseLLMModelInfo): + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return ( + api_key + or litellm.api_key + or litellm.openrouter_key + or get_secret_str("OPENROUTER_API_KEY") + or get_secret_str("OR_API_KEY") + ) + + @staticmethod + def get_api_base(api_base: str | None = None) -> str: + return api_base or litellm.api_base or get_secret_str("OPENROUTER_API_BASE") or _DEFAULT_API_BASE + + @staticmethod + def get_base_model(model: str) -> str: + return model.removeprefix("openrouter/") + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseLLMModelInfo.validate_environment returns a dict + resolved_key: Final = self.get_api_key(api_key) + if resolved_key is None: + return {**headers} + return {**headers, "Authorization": f"Bearer {resolved_key}"} + + def get_models( + self, api_key: str | None = None, api_base: str | None = None + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + resolved_key: Final = self.get_api_key(api_key) + if resolved_key is None: + raise ValueError("Set OPENROUTER_API_KEY or the deployment's api_key to list OpenRouter models") + headers: Final = self.validate_environment( + headers={}, model="", messages=(), optional_params={}, litellm_params={}, api_key=resolved_key + ) + response: Final = litellm.module_level_client.get( # pyright: ignore[reportUnknownMemberType] # legacy HTTP wrapper has untyped query parameters + url=_chat_catalog_url(self.get_api_base(api_base)), headers=headers, timeout=_CATALOG_TIMEOUT + ) + response.raise_for_status() + return [ + f"openrouter/{model.id}" + for model in _Catalog.model_validate_json(response.content).data + if not model.id.endswith(_BATCH_ONLY_SUFFIX) + ] diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index b21315529eb..cfc0d2750bd 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -1,12 +1,15 @@ # What is this? ## Common checks for /v1/models and `/model/info` +import asyncio +import contextvars import copy import functools import itertools import re from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from types import MappingProxyType -from typing import Final +from typing import Final, TypeVar import litellm from litellm._logging import verbose_proxy_logger @@ -30,6 +33,17 @@ _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) +_MODEL_LISTING_POOL: Final = ThreadPoolExecutor(thread_name_prefix="litellm-model-listing") + +_PROVIDERS_LISTED_ONLY_THROUGH_A_DEPLOYMENT: Final = frozenset({LlmProviders.OPENROUTER.value}) + +_Listed: Final = TypeVar("_Listed") + + +async def run_model_listing(listing: Callable[[], _Listed]) -> _Listed: + context: Final = contextvars.copy_context() + return await asyncio.get_running_loop().run_in_executor(_MODEL_LISTING_POOL, context.run, listing) + def _check_wildcard_routing(model: str) -> bool: """ @@ -209,6 +223,7 @@ def get_complete_model_list( include_model_access_groups: bool | None = False, only_model_access_groups: bool | None = False, team_id: str | None = None, + expand_wildcard_routes: bool = True, ) -> list[str]: """Logic for returning complete model list for a given key + team pair""" @@ -254,11 +269,10 @@ def get_complete_model_list( return_wildcard_routes=return_wildcard_routes, llm_router=llm_router, team_id=team_id, + expand_wildcard_routes=expand_wildcard_routes, ) - complete_model_list: Final = unique_models + all_wildcard_models - - return complete_model_list + return list(dict.fromkeys(unique_models + all_wildcard_models)) def _hydrate_litellm_credential_name( @@ -425,6 +439,15 @@ def _cached_wildcard_routing(llm_router: Router | None) -> Callable[[str, str | return functools.cache(wildcard_routed_models) +def _deployed_by_name(llm_router: Router | None) -> Callable[[str, str | None], bool]: + if llm_router is None: + return lambda _model, _team_id: False + registry: Final = llm_router.registry() + deployed_names: Final = frozenset(registry.names(kind="group_set")) + team_deployed_names: Final = frozenset(registry.index("team_model")) + return lambda model, team_id: model in deployed_names or (team_id, model) in team_deployed_names + + def _routes_to_model( routed_models: Callable[[str, str | None], frozenset[str]] | None, team_id: str | None ) -> Callable[[str, str], bool]: @@ -455,7 +478,11 @@ def get_known_models_from_wildcard( wildcard_provider_prefix, wildcard_suffix = wildcard_model_to_expand.split("/", 1) except ValueError: # safely fail return [] - deployment_pattern_naming_listed_ids: Final = _deployment_pattern_naming_listed_ids(wildcard_suffix, litellm_params) + deployment_pattern_naming_listed_ids: Final = ( + wildcard_model_to_expand + if wildcard_model_to_expand != wildcard_model + else _deployment_pattern_naming_listed_ids(wildcard_suffix, litellm_params) + ) # Use provider from litellm_params when available, otherwise from wildcard prefix # (e.g., "openai" from "openai/*" - needed for BYOK where wildcard isn't in router) @@ -514,6 +541,27 @@ def _under_wildcard_prefix(model: str, wildcard_provider_prefix: str, listed_pro return f"{wildcard_provider_prefix}/{model_suffix if leading_is_listed_providers_own_prefix else model}" +@functools.lru_cache(maxsize=256) +def _wildcard_name_regex(name_pattern: str) -> re.Pattern[str]: + return re.compile(PatternMatchRouter().pattern_to_regex(name_pattern)) + + +def _routed_deployment_model(name_pattern: str, listed_name: str, deployment_model: str) -> str: + name_match: Final = _wildcard_name_regex(name_pattern).match(listed_name) + if name_match is None: + return listed_name + return PatternMatchRouter.set_deployment_model_name( + matched_pattern=name_match, litellm_deployment_litellm_model=deployment_model + ) + + +def _names_no_deployment_serves_by_name( + deployment: Mapping[str, object], names: Sequence[str], deployed_by_name: Callable[[str, str | None], bool] +) -> tuple[str, ...]: + team_id: Final = _deployment_team_id(deployment) + return tuple(name for name in names if not deployed_by_name(name, team_id)) + + def expand_wildcard_deployments_for_model_info( deployments: list[dict[str, object]], llm_router: Router | None = None, @@ -527,6 +575,7 @@ def expand_wildcard_deployments_for_model_info( """ expanded: Final[list[dict[str, object]]] = [] routed_models: Final = _cached_wildcard_routing(llm_router) + deployed_by_name: Final = _deployed_by_name(llm_router) for deployment in deployments: model_name = str(deployment.get("model_name") or "") raw_params = deployment.get("litellm_params") @@ -560,22 +609,29 @@ def expand_wildcard_deployments_for_model_info( expanded.append(deployment) continue - for name in expanded_names: + for name in _names_no_deployment_serves_by_name(deployment, expanded_names, deployed_by_name): row = copy.deepcopy(deployment) row["model_name"] = name params = row.get("litellm_params") if isinstance(params, dict): - params["model"] = name + params["model"] = _routed_deployment_model(model_name, name, litellm_model) expanded.append(row) return expanded +def _known_models_without_a_deployment(wildcard_model: str) -> Sequence[str]: + if wildcard_model.partition("/")[0] in _PROVIDERS_LISTED_ONLY_THROUGH_A_DEPLOYMENT: + return () + return get_known_models_from_wildcard(wildcard_model=wildcard_model, litellm_params=None) + + def _get_wildcard_models( unique_models: list[str], return_wildcard_routes: bool | None = False, llm_router: Router | None = None, team_id: str | None = None, + expand_wildcard_routes: bool = True, ) -> list[str]: models_to_remove: Final = set() routes_to_model: Final = _routes_to_model(_cached_wildcard_routing(llm_router), team_id) @@ -586,6 +642,8 @@ def _get_wildcard_models( all_wildcard_models.append(model) models_to_remove.add(model) + if not expand_wildcard_routes: + continue model_list = llm_router.registry().deployments(model, team_id=team_id) if llm_router else None if model_list: @@ -598,7 +656,7 @@ def _get_wildcard_models( ) ) else: - all_wildcard_models.extend(get_known_models_from_wildcard(wildcard_model=model, litellm_params=None)) + all_wildcard_models.extend(_known_models_without_a_deployment(model)) for model in models_to_remove: unique_models.remove(model) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 907c2ca303e..c0d81abeba7 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1,4 +1,5 @@ import base64 +import itertools import logging import mimetypes import re @@ -541,7 +542,8 @@ def get_team_provider_credentials( listings on the team's own provider account/key instead of a shared global one. 2. Fallback: any deployment the caller is granted access to for this - provider, expanding wildcard routes and the all-proxy-models sentinel. + provider, matching wildcard grants against deployed model names without + listing any provider's catalog, and expanding the all-proxy-models sentinel. Skipped when ``include_accessible_deployments`` is False. Credential lookup is scoped to both the team's allowlist and the key's own @@ -559,6 +561,7 @@ def get_team_provider_credentials( return None from litellm.proxy._types import SpecialModelNames + from litellm.proxy.auth.auth_checks import is_model_allowed_by_pattern from litellm.proxy.auth.model_checks import get_complete_model_list, get_key_models team_id: Final = user_api_key_dict.team_id @@ -645,28 +648,36 @@ def get_team_provider_credentials( grants_all_models: Final = SpecialModelNames.all_proxy_models.value in team_models effective_team_models: Final = [] if grants_all_models else team_models - models_to_try: Final = list( - dict.fromkeys( - get_complete_model_list( - key_models=list(key_model_allowlist), - team_models=effective_team_models, - proxy_model_list=proxy_model_list, - user_model=None, - infer_model_from_keys=False, - return_wildcard_routes=True, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=True, - team_id=team_id, - ) + granted_models: Final = get_complete_model_list( + key_models=list(key_model_allowlist), + team_models=effective_team_models, + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=True, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=True, + team_id=team_id, + expand_wildcard_routes=False, + ) + deployments_granted_by_wildcard: Final = ( + model_name + for model_name in proxy_model_list + if any( + is_model_allowed_by_pattern(model=model_name, allowed_model_pattern=granted) for granted in granted_models ) ) - for model_name in models_to_try: - credentials = _provider_credentials(model_name) - if credentials is not None: - return credentials - - return None + return next( + ( + credentials + for credentials in map( + _provider_credentials, dict.fromkeys(itertools.chain(granted_models, deployments_granted_by_wildcard)) + ) + if credentials is not None + ), + None, + ) HTTP_HEADER_VALUE: Final = re.compile(r"[\x21-\x7e]+(?:[ \t]+[\x21-\x7e]+)*") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f078267c6b2..7ea1c246da5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -199,6 +199,7 @@ from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( PRICING_OVERRIDES_KEY, CountContextManagement, + LlmProviders, ModelResponse, ModelResponseStream, TextCompletionResponse, @@ -443,6 +444,7 @@ from litellm.proxy.auth.model_checks import ( get_key_models, get_mcp_server_ids, get_team_models, + run_model_listing, ) from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401, RUF100 # legacy module exports @@ -12487,20 +12489,23 @@ async def model_list( # Get complete model list including wildcard routes if requested from litellm.proxy.auth.model_checks import get_complete_model_list - all_models = get_complete_model_list( - key_models=[], - team_models=[], - proxy_model_list=[ - *proxy_model_list, - *(typesafe_allowlist_model_names() if asks_for_allowlist_names else ()), - ], - user_model=None, - infer_model_from_keys=False, - return_wildcard_routes=include_wildcard_routes, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups or False, - only_model_access_groups=only_model_access_groups or False, + all_models = await run_model_listing( + partial( + get_complete_model_list, + key_models=[], + team_models=[], + proxy_model_list=[ + *proxy_model_list, + *(typesafe_allowlist_model_names() if asks_for_allowlist_names else ()), + ], + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=include_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups or False, + only_model_access_groups=only_model_access_groups or False, + ) ) expanded_undiscoverable_names: Final = undiscoverable_model_names( @@ -14747,6 +14752,8 @@ from litellm.responses.litellm_completion_transformation.count_request import ( with_provider_file_ids, ) +_PROVIDERS_WITH_LISTING_ONLY_MODEL_INFO: Final = frozenset({LlmProviders.OPENROUTER}) + def _with_client_oauth_token( counter: BaseTokenCounter | None, client_oauth_token: str | None @@ -14782,14 +14789,13 @@ def _get_provider_token_counter( ) # Switch case pattern using existing get_provider_model_info - from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager # Convert string provider to LlmProviders enum llm_provider_enum: Final = LlmProviders(custom_llm_provider) # Add more provider mappings as needed - if llm_provider_enum: + if llm_provider_enum not in _PROVIDERS_WITH_LISTING_ONLY_MODEL_INFO: provider_model_info: Final = ProviderConfigManager.get_provider_model_info( model=full_model, provider=llm_provider_enum ) @@ -17941,7 +17947,7 @@ async def model_info_v1( alias_models: Final = copy.deepcopy(registry.deployments(source="alias")) all_models.extend(alias_models) - all_models = expand_wildcard_deployments_for_model_info(all_models, llm_router) + all_models = await run_model_listing(partial(expand_wildcard_deployments_for_model_info, all_models, llm_router)) allowed_model_names: Final = await _get_v1_model_info_allowed_model_names_with_access_groups( user_api_key_dict=user_api_key_dict, @@ -18279,14 +18285,17 @@ async def model_group_info( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ) all_models_str: Final = ( - get_complete_model_list( - key_models=(), - team_models=(), - proxy_model_list=llm_router.registry().names(), - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), - return_wildcard_routes=False, - llm_router=llm_router, + await run_model_listing( + partial( + get_complete_model_list, + key_models=(), + team_models=(), + proxy_model_list=llm_router.registry().names(), + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + return_wildcard_routes=False, + llm_router=llm_router, + ) ) if is_proxy_admin else await get_available_models_for_user( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ddaa8d2e69d..00dbcd0f9a8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -9500,7 +9500,7 @@ async def _agent_access_group_visible_models( ) -> frozenset[str] | None: """Models an agent key may still list once its attached access groups cap it, ``None`` when nothing caps it, so ``/v1/models`` never advertises a model the same key would be denied on.""" - from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models + from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models, run_model_listing if not user_api_key_dict.agent_id: return None @@ -9521,17 +9521,20 @@ async def _agent_access_group_visible_models( if not granted: return frozenset() return frozenset( - get_complete_model_list( - key_models=granted, - team_models=(), - proxy_model_list=proxy_model_list, - user_model=None, - infer_model_from_keys=False, - return_wildcard_routes=return_wildcard_routes, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - team_id=team_id, + await run_model_listing( + partial( + get_complete_model_list, + key_models=granted, + team_models=(), + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=return_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + team_id=team_id, + ) ) ) @@ -9576,6 +9579,7 @@ async def get_available_models_for_user( get_complete_model_list, get_key_models, get_team_models, + run_model_listing, ) # Get proxy model list and access groups @@ -9656,18 +9660,21 @@ async def get_available_models_for_user( ) # Get complete model list - complete_models: Final = get_complete_model_list( - key_models=granted_key_models, - team_models=granted_team_models, - proxy_model_list=proxy_model_list, - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), - return_wildcard_routes=return_wildcard_routes, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - only_model_access_groups=only_model_access_groups, - team_id=effective_team_id, + complete_models: Final = await run_model_listing( + partial( + get_complete_model_list, + key_models=granted_key_models, + team_models=granted_team_models, + proxy_model_list=proxy_model_list, + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + return_wildcard_routes=return_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + only_model_access_groups=only_model_access_groups, + team_id=effective_team_id, + ) ) all_models: Final = ( list( diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 4dd217cf59b..2b8d1e9c11b 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -279,7 +279,9 @@ class PatternMatchRouter: PatternMatchRouter.set_deployment_model_name( matched_pattern=pattern_match, litellm_deployment_litellm_model=deployment["litellm_params"]["model"] ) - for deployment in self.patterns[pattern] + for deployment in self.patterns.get( # pyright: ignore[reportUnknownMemberType] # untyped deployment dicts + pattern, () + ) ) diff --git a/litellm/utils.py b/litellm/utils.py index 2ca8a49f968..560d0c31e4a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8724,11 +8724,13 @@ class AvailableModelsCache(InMemoryCache): cache_key: Final = self._get_cache_key(custom_llm_provider, litellm_params) - result: Final = cast(list[str] | None, self.get_cache(cache_key)) + result: Final = self._cached(cache_key) - if result is not None: - return copy.deepcopy(result) - return result + if isinstance(result, Sequence): + return list(result) + if self._cached(_failed_listing_key(cache_key)) is not None: + return [] + return None def set_cached_model_info( self, @@ -8738,12 +8740,46 @@ class AvailableModelsCache(InMemoryCache): ): """Set cached model info""" cache_key: Final = self._get_cache_key(custom_llm_provider, litellm_params) - self.set_cache(cache_key, copy.deepcopy(available_models)) + self._store(cache_key, copy.deepcopy(available_models)) + + def set_failed_model_listing( + self, custom_llm_provider: str, litellm_params: LiteLLM_Params | None, ttl_seconds: float + ) -> None: + cache_key: Final = self._get_cache_key(custom_llm_provider, litellm_params) + self._store(_failed_listing_key(cache_key), True, ttl=ttl_seconds) + + def listing_lock(self, custom_llm_provider: str, litellm_params: LiteLLM_Params | None) -> threading.RLock: + return _provider_listing_lock(self._get_cache_key(custom_llm_provider, litellm_params)) + + def _cached(self, key: str) -> Sequence[str] | bool | None: + return cast(Sequence[str] | bool | None, self.get_cache(key)) + + def _store(self, key: str, value: object, ttl: float | None = None) -> None: + self.set_cache(key, value, ttl=ttl) + + +def _failed_listing_key(cache_key: str) -> str: + return f"{cache_key}:failed" + + +_PROVIDER_LISTING_LOCK_GUARD: Final = threading.Lock() + + +@lru_cache(maxsize=1024) +def _cached_provider_listing_lock(cache_key: str) -> threading.RLock: + return threading.RLock() + + +def _provider_listing_lock(cache_key: str) -> threading.RLock: + with _PROVIDER_LISTING_LOCK_GUARD: + return _cached_provider_listing_lock(cache_key) # Global cache instance _model_cache: Final = AvailableModelsCache() +_FAILED_PROVIDER_LISTING_TTL_SECONDS: Final = 30 + def _infer_valid_provider_from_env_vars( custom_llm_provider: str | None = None, @@ -8791,20 +8827,27 @@ def _get_valid_models_from_provider_api( custom_llm_provider: str, litellm_params: LiteLLM_Params | None = None, ) -> list[str]: - try: - cached_result: Final = _model_cache.get_cached_model_info(custom_llm_provider, litellm_params) + with _model_cache.listing_lock(custom_llm_provider, litellm_params): + try: + cached_result: Final = _model_cache.get_cached_model_info(custom_llm_provider, litellm_params) - if cached_result is not None: - return cached_result - models: Final = provider_config.discover_models( - litellm_params=_provider_listing_params(custom_llm_provider, litellm_params) - ) + if cached_result is not None: + return cached_result + models: Final = provider_config.discover_models( + litellm_params=_provider_listing_params(custom_llm_provider, litellm_params) + ) - _model_cache.set_cached_model_info(custom_llm_provider, litellm_params, models) - return models - except Exception as e: - verbose_logger.warning("Error getting valid models: %s", e) - return [] + _model_cache.set_cached_model_info(custom_llm_provider, litellm_params, models) + return models + except httpx.TimeoutException as e: + verbose_logger.warning("Error getting valid models: %s", e) + _model_cache.set_failed_model_listing( + custom_llm_provider, litellm_params, ttl_seconds=_FAILED_PROVIDER_LISTING_TTL_SECONDS + ) + return [] + except Exception as e: + verbose_logger.warning("Error getting valid models: %s", e) + return [] def bundled_provider_models(provider: str) -> tuple[str, ...]: @@ -10470,6 +10513,10 @@ class ProviderConfigManager: from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo return AzureFoundryModelInfo(model=model) + elif LlmProviders.OPENROUTER == provider: + from litellm.llms.openrouter.common_utils import OpenRouterModelInfo + + return OpenRouterModelInfo() return None @staticmethod diff --git a/tests/unit/llms/openrouter/test_openrouter_common_utils.py b/tests/unit/llms/openrouter/test_openrouter_common_utils.py new file mode 100644 index 00000000000..087c05561b5 --- /dev/null +++ b/tests/unit/llms/openrouter/test_openrouter_common_utils.py @@ -0,0 +1,159 @@ +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.llms.openrouter.common_utils import OpenRouterModelInfo + +DEFAULT_BASE: Final = "https://openrouter.ai/api/v1" +KEY_SOURCES: Final = ("api_key", "litellm.api_key", "litellm.openrouter_key", "OPENROUTER_API_KEY", "OR_API_KEY") +BASE_SOURCES: Final = ("api_base", "litellm.api_base", "OPENROUTER_API_BASE") + + +@pytest.fixture(autouse=True) +def _no_ambient_openrouter_settings(monkeypatch: pytest.MonkeyPatch): + for env_var in ("OPENROUTER_API_BASE", "OPENROUTER_API_KEY", "OR_API_KEY"): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.setattr(litellm, "openrouter_key", None) + + +def _catalog(*model_ids: str) -> httpx.Response: + return httpx.Response(200, json={"data": [{"id": model_id, "name": model_id} for model_id in model_ids]}) + + +def _configure(monkeypatch: pytest.MonkeyPatch, source: str, value: str) -> None: + if source.startswith("litellm."): + monkeypatch.setattr(litellm, source.removeprefix("litellm."), value) + else: + monkeypatch.setenv(source, value) + + +def test_get_models_lists_only_the_chat_catalog_under_openrouter_keeping_the_org_namespace(): + with respx.mock() as respx_mock: + respx_mock.get(f"{DEFAULT_BASE}/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini", "anthropic/claude-haiku-5-5") + ) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-test") == [ + "openrouter/openai/gpt-5.4-mini", + "openrouter/anthropic/claude-haiku-5-5", + ] + assert [call.request.url.path for call in respx_mock.calls] == ["/api/v1/models/user"] + + +@respx.mock +def test_get_models_drops_batch_only_ids_and_keeps_every_other_variant(): + respx.get(f"{DEFAULT_BASE}/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini:batch", "openai/gpt-5.4-mini", "qwen/qwen3-coder:free") + ) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-test") == [ + "openrouter/openai/gpt-5.4-mini", + "openrouter/qwen/qwen3-coder:free", + ] + + +@pytest.mark.parametrize( + "api_base", + [ + "https://gateway.example/openrouter/api/v1/", + "https://gateway.example/openrouter/api/v1/chat/completions", + "https://gateway.example/openrouter/api/v1/chat/completions/", + ], +) +@respx.mock +def test_get_models_keeps_the_api_base_path_and_sends_the_deployment_key(api_base: str): + chat_route: Final = respx.get("https://gateway.example/openrouter/api/v1/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini") + ) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-deployment", api_base=api_base) == [ + "openrouter/openai/gpt-5.4-mini" + ] + assert chat_route.calls.last.request.headers["authorization"] == "Bearer sk-or-deployment" + + +@pytest.mark.parametrize("winner", range(len(KEY_SOURCES)), ids=KEY_SOURCES) +@respx.mock +def test_get_models_sends_the_first_configured_key_in_the_chat_paths_order( + monkeypatch: pytest.MonkeyPatch, winner: int +): + for source in KEY_SOURCES[winner + 1 :]: + _configure(monkeypatch, source, f"sk-or-from-{source}") + if winner > 0: + _configure(monkeypatch, KEY_SOURCES[winner], "sk-or-winner") + chat_route: Final = respx.get(f"{DEFAULT_BASE}/models/user").mock(return_value=_catalog("qwen/qwen3-coder")) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-winner" if winner == 0 else None) == [ + "openrouter/qwen/qwen3-coder" + ] + assert chat_route.calls.last.request.headers["authorization"] == "Bearer sk-or-winner" + + +@pytest.mark.parametrize("winner", range(len(BASE_SOURCES) + 1), ids=(*BASE_SOURCES, "default")) +@respx.mock +def test_get_models_calls_the_first_configured_base_in_the_chat_paths_order( + monkeypatch: pytest.MonkeyPatch, winner: int +): + for source in BASE_SOURCES[winner + 1 :]: + _configure(monkeypatch, source, f"https://{source.lower().replace('.', '-')}.example/api/v1") + if 0 < winner < len(BASE_SOURCES): + _configure(monkeypatch, BASE_SOURCES[winner], "https://winner.example/api/v1") + expected_base: Final = "https://winner.example/api/v1" if winner < len(BASE_SOURCES) else DEFAULT_BASE + respx.get(f"{expected_base}/models/user").mock(return_value=_catalog("qwen/qwen3-coder")) + + assert OpenRouterModelInfo().get_models( + api_key="sk-or-test", api_base="https://winner.example/api/v1" if winner == 0 else None + ) == ["openrouter/qwen/qwen3-coder"] + + +def test_get_models_without_any_key_raises_before_calling_openrouter(): + with respx.mock(assert_all_called=False) as respx_mock: + chat_route: Final = respx_mock.get(f"{DEFAULT_BASE}/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini") + ) + + with pytest.raises(ValueError, match="OPENROUTER_API_KEY"): + OpenRouterModelInfo().get_models() + + assert not chat_route.called + + +@respx.mock +def test_get_models_raises_when_the_catalog_answers_an_error(): + respx.get(f"{DEFAULT_BASE}/models/user").mock(return_value=httpx.Response(401, json={"error": {"code": 401}})) + + with pytest.raises(httpx.HTTPStatusError): + OpenRouterModelInfo().get_models(api_key="sk-or-revoked") + + +def test_validate_environment_adds_the_resolved_key_and_keeps_the_callers_headers(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-env") + headers: Final = {"X-Title": "my-app"} + + assert OpenRouterModelInfo().validate_environment( + headers=headers, model="openrouter/openai/gpt-5.4-mini", messages=(), optional_params={}, litellm_params={} + ) == {"X-Title": "my-app", "Authorization": "Bearer sk-or-env"} + assert headers == {"X-Title": "my-app"} + + +def test_validate_environment_without_any_key_sends_no_authorization(): + assert OpenRouterModelInfo().validate_environment( + headers={"X-Title": "my-app"}, + model="openrouter/openai/gpt-5.4-mini", + messages=(), + optional_params={}, + litellm_params={}, + ) == {"X-Title": "my-app"} + + +@pytest.mark.parametrize( + ("model", "base_model"), + [("openrouter/openai/gpt-5.4-mini", "openai/gpt-5.4-mini"), ("openai/gpt-5.4-mini", "openai/gpt-5.4-mini")], +) +def test_get_base_model_drops_only_the_openrouter_prefix(model: str, base_model: str): + assert OpenRouterModelInfo.get_base_model(model) == base_model diff --git a/tests/unit/proxy/auth/test_model_checks.py b/tests/unit/proxy/auth/test_model_checks.py index c67610c24a1..d0156ff0195 100644 --- a/tests/unit/proxy/auth/test_model_checks.py +++ b/tests/unit/proxy/auth/test_model_checks.py @@ -7,6 +7,8 @@ import httpx import pytest import respx +from litellm import Router + def test_get_team_models_for_all_models_and_team_only_models(): from litellm.proxy.auth.model_checks import get_team_models @@ -1261,6 +1263,223 @@ def test_partial_bedrock_wildcard_lists_nothing_when_no_discovered_id_carries_it assert get_known_models_from_wildcard("bedrock/anthropic.*", deployment) == [] +@pytest.mark.parametrize( + ("wildcard", "deployment_model", "listed_ids"), + [ + ( + "openrouter/*", + "openrouter/*", + ["openrouter/openai/gpt-5.4-mini", "openrouter/anthropic/claude-haiku-5-5"], + ), + ("or-team/*", "openrouter/*", ["or-team/openai/gpt-5.4-mini", "or-team/anthropic/claude-haiku-5-5"]), + ("openrouter/openai/*", "openrouter/openai/*", ["openrouter/openai/gpt-5.4-mini"]), + ("*", "openrouter/*", ["openai/gpt-5.4-mini", "anthropic/claude-haiku-5-5"]), + ], +) +@respx.mock +def test_openrouter_wildcard_lists_the_live_catalog_with_each_vendor_namespace( + monkeypatch: pytest.MonkeyPatch, wildcard: str, deployment_model: str, listed_ids: list[str] +) -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.delenv("OPENROUTER_API_BASE", raising=False) + respx.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response( + 200, json={"data": [{"id": "openai/gpt-5.4-mini"}, {"id": "anthropic/claude-haiku-5-5"}]} + ) + ) + + assert ( + get_known_models_from_wildcard(wildcard, LiteLLM_Params(model=deployment_model, api_key=f"sk-or-{wildcard}")) + == listed_ids + ) + + +@respx.mock +def test_key_granted_openrouter_wildcard_without_an_openrouter_deployment_lists_no_environment_catalog( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + from litellm import Router + from litellm.proxy.auth.model_checks import get_complete_model_list + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-environment-only") + catalog: Final = respx.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]}) + ) + router: Final = Router( + model_list=[{"model_name": "or-team/*", "litellm_params": {"model": "openrouter/*", "api_key": "sk-or-team"}}] + ) + + listed: Final = get_complete_model_list( + key_models=["openrouter/*"], + team_models=[], + proxy_model_list=router.model_names, + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + + assert (listed, catalog.call_count) == ([], 0) + + +@contextmanager +def _openai_catalog_behind_a_catch_all( + monkeypatch: pytest.MonkeyPatch, concrete_name: str, concrete_model_info: Mapping[str, str] +) -> Generator[Router, None, None]: + import litellm + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.get("https://api.openai.com/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "gpt-4o"}, {"id": "gpt-4o-mini"}]}) + ) + yield Router( + model_list=[ + {"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-openai-catch-all"}}, + { + "model_name": concrete_name, + "litellm_params": {"model": "openrouter/openai/gpt-4o", "api_key": "sk-or-gpt-4o"}, + "model_info": dict(concrete_model_info), + }, + ] + ) + + +@pytest.mark.parametrize( + ("concrete_name", "concrete_model_info", "team_id"), + [ + ("gpt-4o", {}, None), + ("model_name_team-x_gpt-4o", {"team_id": "team-x", "team_public_model_name": "gpt-4o"}, "team-x"), + ], +) +def test_catch_all_wildcard_lists_a_name_a_concrete_deployment_serves_once( + monkeypatch: pytest.MonkeyPatch, concrete_name: str, concrete_model_info: Mapping[str, str], team_id: str | None +) -> None: + from litellm.proxy.auth.model_checks import get_complete_model_list + + with _openai_catalog_behind_a_catch_all(monkeypatch, concrete_name, concrete_model_info) as router: + listed: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=router.get_model_names(team_id=team_id), + user_model=None, + infer_model_from_keys=False, + llm_router=router, + team_id=team_id, + ) + + assert sorted(listed) == ["gpt-4o", "gpt-4o-mini"] + + +def test_catch_all_wildcard_model_info_keeps_only_the_concrete_row_for_a_deployed_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info + + with _openai_catalog_behind_a_catch_all(monkeypatch, "gpt-4o", {}) as router: + rows: Final = expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or ()], llm_router=router + ) + + assert sorted((row["model_name"], row["litellm_params"]["model"]) for row in rows) == [ + ("gpt-4o", "openrouter/openai/gpt-4o"), + ("gpt-4o-mini", "openai/gpt-4o-mini"), + ] + + +@contextmanager +def _openrouter_catalog_beside_a_named_deployment( + monkeypatch: pytest.MonkeyPatch, + wildcard_name: str, + named: str, + catalog: tuple[str, ...] = ("openai/gpt-5.4-mini", "anthropic/claude-haiku-5-5"), + wildcard_key: str | None = None, +) -> Generator[Router, None, None]: + import litellm + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.delenv("OPENROUTER_API_BASE", raising=False) + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response(200, json={"data": [{"id": catalog_id} for catalog_id in catalog]}) + ) + yield Router( + model_list=[ + { + "model_name": wildcard_name, + "litellm_params": { + "model": "openrouter/*", + "api_key": wildcard_key or f"sk-or-wildcard-{wildcard_name}", + }, + }, + { + "model_name": named, + "litellm_params": {"model": "openrouter/openai/gpt-5.4-mini", "api_key": "sk-or-named"}, + }, + ] + ) + + +@pytest.mark.parametrize( + ("wildcard_name", "named"), + [("openrouter/*", "openrouter/openai/gpt-5.4-mini"), ("or-team/*", "or-team/openai/gpt-5.4-mini")], +) +def test_provider_wildcard_lists_a_name_a_named_deployment_serves_once( + monkeypatch: pytest.MonkeyPatch, wildcard_name: str, named: str +) -> None: + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list + + with _openrouter_catalog_beside_a_named_deployment(monkeypatch, wildcard_name, named) as router: + listed: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=list(router.model_names), + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + rows: Final = expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or ()], llm_router=router + ) + + vendor_prefix: Final = wildcard_name.removesuffix("*") + assert sorted(listed) == [f"{vendor_prefix}anthropic/claude-haiku-5-5", named] + assert sorted((row["model_name"], row["litellm_params"]["api_key"]) for row in rows) == [ + (f"{vendor_prefix}anthropic/claude-haiku-5-5", f"sk-or-wildcard-{wildcard_name}"), + (named, "sk-or-named"), + ] + + +@pytest.mark.parametrize( + ("wildcard_name", "named"), + [("openrouter/*", "openrouter/openai/gpt-5.4-mini"), ("or-team/*", "or-team/openai/gpt-5.4-mini")], +) +def test_provider_wildcard_whose_whole_catalog_is_deployed_by_name_shows_only_the_named_row( + monkeypatch: pytest.MonkeyPatch, wildcard_name: str, named: str +) -> None: + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info + + with _openrouter_catalog_beside_a_named_deployment( + monkeypatch, + wildcard_name, + named, + catalog=("openai/gpt-5.4-mini",), + wildcard_key=f"sk-or-one-model-{wildcard_name}", + ) as router: + rows: Final = expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or ()], llm_router=router + ) + + assert [(row["model_name"], row["litellm_params"]["api_key"]) for row in rows] == [(named, "sk-or-named")] + + @respx.mock def test_custom_prefix_that_starts_a_discovered_id_still_prefixes_every_listed_model( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py index 82523572b28..238f1fa6495 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py @@ -3,10 +3,13 @@ from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +import respx from fastapi import HTTPException from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles +import litellm from litellm import Router from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth @@ -852,3 +855,112 @@ def test_a_team_keeps_its_first_config_file_openai_deployment_over_older_db_ones assert credentials is not None assert credentials["api_key"] == "deployment-z-key" + + +@pytest.mark.parametrize( + ("custom_llm_provider", "openrouter_key", "api_key"), + [("openai", "sk-or-files-openai", "openai-key"), ("openrouter", "sk-or-files-openrouter", "sk-or-files-openrouter")], +) +def test_a_provider_credential_lookup_never_lists_an_openrouter_wildcards_live_catalog( + monkeypatch: pytest.MonkeyPatch, custom_llm_provider: str, openrouter_key: str, api_key: str +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + router: Final = Router( + model_list=[ + { + "model_name": "openrouter/*", + "litellm_params": { + "model": "openrouter/*", + "api_key": openrouter_key, + "api_base": "https://openrouter.ai/api/v1", + }, + }, + {"model_name": "gpt-5.4-mini", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "openai-key"}}, + ] + ) + + with respx.mock(assert_all_called=False) as respx_mock: + catalog_route: Final = respx_mock.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response(200, json={"data": [{"id": "openai/gpt-5.4-mini"}]}) + ) + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user"), + custom_llm_provider=custom_llm_provider, + include_accessible_deployments=True, + ) + + assert credentials is not None + assert (credentials["custom_llm_provider"], credentials["api_key"]) == (custom_llm_provider, api_key) + assert catalog_route.call_count == 0 + + +@pytest.mark.parametrize( + ("key_models", "deployments", "custom_llm_provider", "api_key"), + [ + ( + ["anthropic/*"], + (("anthropic/claude-*", "anthropic/claude-*", "anthropic-narrow-key"),), + "anthropic", + "anthropic-narrow-key", + ), + (["anthropic/claude-3*"], (("anthropic/*", "anthropic/*", "anthropic-wide-key"),), "anthropic", "anthropic-wide-key"), + (["openrouter/*"], (("openrouter/openai/*", "openrouter/openai/*", "sk-or-narrow"),), "openrouter", "sk-or-narrow"), + ( + ["team/*"], + (("team/*", "openrouter/*", "sk-or-team"), ("team/openai/*", "openai/*", "openai-team-key")), + "openai", + "openai-team-key", + ), + ], +) +def test_a_key_wildcard_finds_the_deployment_it_grants_without_listing_any_provider_catalog( + monkeypatch: pytest.MonkeyPatch, + key_models: list[str], + deployments: tuple[tuple[str, str, str], ...], + custom_llm_provider: str, + api_key: str, +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + router: Final = Router( + model_list=[ + {"model_name": model_name, "litellm_params": {"model": model, "api_key": deployment_key}} + for model_name, model, deployment_key in deployments + ] + ) + + with respx.mock(assert_all_called=False) as respx_mock: + catalog_requests: Final = respx_mock.route().mock(return_value=httpx.Response(200, json={"data": []})) + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", models=key_models), + custom_llm_provider=custom_llm_provider, + include_accessible_deployments=True, + ) + + assert credentials is not None + assert (credentials["custom_llm_provider"], credentials["api_key"]) == (custom_llm_provider, api_key) + assert catalog_requests.call_count == 0 + + +@pytest.mark.parametrize( + ("deployed_name", "api_key"), [("gpt-4o-mini", None), ("openai/gpt-4o-mini", "openai-deployment-key")] +) +def test_a_provider_wildcard_grant_finds_only_deployments_named_under_it(deployed_name: str, api_key: str | None): + router: Final = Router( + model_list=[ + { + "model_name": deployed_name, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "openai-deployment-key"}, + } + ] + ) + + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", models=["openai/*"]), + custom_llm_provider="openai", + include_accessible_deployments=True, + ) + + assert (None if credentials is None else credentials["api_key"]) == api_key diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index 86364d93ed0..bffa4612964 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -550,15 +550,12 @@ async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monke deployment = { "model_name": "*", - "litellm_params": {"model": "openai/*"}, + "litellm_params": {"model": "openai/*", "api_key": "sk-star"}, } - router = route_entry_points_to_stubs(MagicMock()) - router.get_model_access_groups = MagicMock(return_value={}) - router.get_model_names = MagicMock(return_value=["*"]) - router.get_model_list = MagicMock(return_value=[deployment]) + router = litellm.Router(model_list=[deployment]) monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) - expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info([deployment]) + expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info(router.get_model_list(), router) allowed_model_names = await proxy_server._get_v1_model_info_allowed_model_names_with_access_groups( user_api_key_dict=UserAPIKeyAuth( api_key="sk-test", @@ -575,7 +572,8 @@ async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monke allowed_model_names=allowed_model_names, ) - assert [model["model_name"] for model in result] == ["openai/gpt-4o"] + assert [(model["model_name"], model["litellm_params"]["model"]) for model in result] == [("gpt-4o", "openai/gpt-4o")] + assert router.registry().names("gpt-4o", kind="wildcard") == frozenset({"openai/gpt-4o"}) # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/proxy_server/test_routes_models.py b/tests/unit/proxy/proxy_server/test_routes_models.py index ffd23fa3a6f..874f471f3e7 100644 --- a/tests/unit/proxy/proxy_server/test_routes_models.py +++ b/tests/unit/proxy/proxy_server/test_routes_models.py @@ -9,10 +9,16 @@ Pins (PR2): from __future__ import annotations +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field from typing import Final from unittest.mock import MagicMock +import httpx import pytest +import respx import litellm from litellm import Router @@ -541,6 +547,73 @@ def test_codex_service_tiers_resolved_via_internal_team_key(client, auth_as, pat assert (entry["slug"], [tier["id"] for tier in entry["service_tiers"]]) == ("gpt-4-team", ["ultrafast"]) +OPENROUTER_BASE: Final = "https://openrouter.ai/api/v1" + + +@dataclass(frozen=True, slots=True) +class _CatalogHeldUntilTheLoopRuns: + requested: threading.Event = field(default_factory=threading.Event) + released: threading.Event = field(default_factory=threading.Event) + + def answer(self, _request: httpx.Request) -> httpx.Response: + self.requested.set() + if not self.released.wait(timeout=30): + return httpx.Response(504) + return httpx.Response(200, json={"data": [{"id": "openai/gpt-5.4-mini"}]}) + + async def release_from_the_loop(self) -> None: + assert await asyncio.to_thread(self.requested.wait, 30) + self.released.set() + + +@pytest.mark.parametrize("scope", [None, "expand"]) +def test_v1_models_fetches_a_live_openrouter_catalog_without_holding_the_event_loop_or_its_default_thread_pool( + monkeypatch: pytest.MonkeyPatch, scope: str | None +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "openrouter/*", + "litellm_params": { + "model": "openrouter/*", + "api_key": f"sk-or-v1-models-off-loop-{scope}", + "api_base": OPENROUTER_BASE, + }, + } + ] + ), + ) + catalog: Final = _CatalogHeldUntilTheLoopRuns() + + async def listed_ids_on_a_one_thread_default_pool() -> list[str]: + asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) + listing: Final = asyncio.ensure_future( + proxy_server.model_list( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ), + scope=scope, + ) + ) + await catalog.release_from_the_loop() + response: Final = await listing + return [str(model["id"]) for model in response["data"]] + + with respx.mock() as respx_mock: + catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer) + listed_ids: Final = asyncio.run(listed_ids_on_a_one_thread_default_pool()) + + assert catalog_route.calls.last.response.status_code == 200 + assert listed_ids == ["openrouter/openai/gpt-5.4-mini"] + + def _priced_jev_models() -> frozenset[str]: return frozenset(name.removeprefix("typesafe/") for name in litellm.model_cost if name.startswith("typesafe/")) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index c182cd1ec75..68eb05e60dd 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -179,6 +179,39 @@ async def test_token_counting_model_not_in_model_list(): assert response.model_used == "special-alias" +@pytest.mark.asyncio +async def test_asking_the_provider_endpoint_leaves_an_openrouter_deployments_count_unchanged( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "or-gpt", + "litellm_params": { + "model": "openai/gpt-4o", + "custom_llm_provider": "openrouter", + "api_key": "sk-or-token-count", + }, + } + ] + ), + ) + request: Final = TokenCountRequest( + model="or-gpt", messages=[{"role": "user", "content": "Count these tokens, ünïcödé 你好 🙂"}] + ) + + via_provider_endpoint: Final = await token_counter(request=request, call_endpoint=True) + counted_locally: Final = await token_counter(request=request, call_endpoint=False) + + assert (via_provider_endpoint.total_tokens, via_provider_endpoint.model_used) == ( + counted_locally.total_tokens, + counted_locally.model_used, + ) + + @pytest.mark.asyncio async def test_gpt_token_counting(): """ diff --git a/tests/unit/proxy/utils/helpers/test_model_access.py b/tests/unit/proxy/utils/helpers/test_model_access.py index 324360cc2ee..30d24e9eddc 100644 --- a/tests/unit/proxy/utils/helpers/test_model_access.py +++ b/tests/unit/proxy/utils/helpers/test_model_access.py @@ -1,10 +1,17 @@ +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Final from unittest.mock import MagicMock +import httpx import pytest +import respx from fastapi import HTTPException import litellm -from litellm import ModelResponse +from litellm import ModelResponse, Router from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ( create_model_info_response, @@ -650,3 +657,61 @@ async def test_listing_is_unchanged_without_an_agent_or_without_attached_groups( resolve_agent_ceiling=_agent_ceiling(None), ) assert (plain_key, agent_without_groups) == (["model-a", "model-b"], ["model-a", "model-b"]) + + +OPENROUTER_BASE: Final = "https://openrouter.ai/api/v1" + + +@dataclass(frozen=True, slots=True) +class _CatalogHeldUntilTheLoopRuns: + requested: threading.Event = field(default_factory=threading.Event) + released: threading.Event = field(default_factory=threading.Event) + + def answer(self, _request: httpx.Request) -> httpx.Response: + self.requested.set() + if not self.released.wait(timeout=30): + return httpx.Response(504) + return httpx.Response(200, json={"data": [{"id": "openai/gpt-5.4-mini"}]}) + + async def release_from_the_loop(self) -> None: + assert await asyncio.to_thread(self.requested.wait, 30) + self.released.set() + + +def test_get_available_models_for_user_fetches_a_live_catalog_without_holding_the_event_loop_or_its_default_pool( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + catalog: Final = _CatalogHeldUntilTheLoopRuns() + router: Final = Router( + model_list=[ + { + "model_name": "openrouter/*", + "litellm_params": { + "model": "openrouter/*", + "api_key": "sk-or-available-models-off-loop", + "api_base": OPENROUTER_BASE, + }, + } + ] + ) + + async def list_models_on_a_one_thread_default_pool() -> list[str]: + asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) + listing: Final = asyncio.ensure_future( + get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-key", user_id="user-1"), + llm_router=router, + general_settings={}, + user_model=None, + ) + ) + await catalog.release_from_the_loop() + return await listing + + with respx.mock() as respx_mock: + catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer) + models: Final = asyncio.run(list_models_on_a_one_thread_default_pool()) + + assert catalog_route.calls.last.response.status_code == 200 + assert models == ["openrouter/openai/gpt-5.4-mini"] diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 0f471b1de71..5cf872f4ea7 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -9269,6 +9269,77 @@ def reset_mock_cache() -> None: _model_cache.flush_cache() +@respx.mock +def test_get_valid_models_does_not_ask_a_timed_out_provider_again_on_the_next_listing(reset_mock_cache: None) -> None: + catalog: Final = respx.get("https://openrouter.ai/api/v1/models/user").mock( + side_effect=httpx.ReadTimeout("catalog hung") + ) + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-during-an-outage") + + listings: Final = tuple( + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + for _ in range(2) + ) + + assert (listings, catalog.call_count) == (([], []), 1) + + +@respx.mock +def test_get_valid_models_lists_a_provider_again_right_after_a_fast_failure(reset_mock_cache: None) -> None: + respx.get("https://openrouter.ai/api/v1/models/user").mock( + side_effect=(httpx.Response(503), httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]})) + ) + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-recovers") + + listings: Final = tuple( + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + for _ in range(2) + ) + + assert listings == ([], ["openrouter/anthropic/claude-haiku-5-5"]) + + +@respx.mock +def test_get_valid_models_keeps_a_good_listing_when_an_overlapping_listing_fails_after_it( + reset_mock_cache: None, +) -> None: + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-overlapping-listings") + first_request: Final = iter((True, False)) + + def answer_after_an_overlapping_listing(request: httpx.Request) -> httpx.Response: + if next(first_request): + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + raise httpx.ReadTimeout("catalog hung", request=request) + return httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]}) + + respx.get("https://openrouter.ai/api/v1/models/user").mock(side_effect=answer_after_an_overlapping_listing) + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + + assert get_valid_models( + check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment + ) == ["openrouter/anthropic/claude-haiku-5-5"] + + +@respx.mock +def test_get_valid_models_asks_a_provider_once_when_many_listings_start_together(reset_mock_cache: None) -> None: + def slow_catalog(request: httpx.Request) -> httpx.Response: + threading.Event().wait(0.2) + return httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]}) + + catalog: Final = respx.get("https://openrouter.ai/api/v1/models/user").mock(side_effect=slow_catalog) + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-many-listings") + + def list_models(_: int) -> tuple[str, ...]: + return tuple( + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + ) + + with ThreadPoolExecutor(max_workers=8) as pool: + listings: Final = frozenset(pool.map(list_models, range(8))) + + assert (listings, catalog.call_count) == (frozenset({("openrouter/anthropic/claude-haiku-5-5",)}), 1) + + @pytest.fixture def isolated_openai_model_sets(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr( From 4f7150f832b3180a6dfcce2be5a052038de45eac Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 12:21:30 -0700 Subject: [PATCH 51/64] fix(ollama): forward assistant tool call ids in ollama_chat requests (#49278) * test(ollama): cover assistant tool call ids in ollama_chat requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ollama): forward assistant tool call ids in ollama_chat requests Co-authored-by: Kenneth Bingham Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ollama): format tool call id regression assertion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ollama): keep tool call ID handling type safe Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ollama): format tool-call id helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ollama): preserve typed tool call arguments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Kenneth Bingham --- litellm/llms/ollama/chat/transformation.py | 24 +++-- litellm/types/llms/ollama.py | 3 + .../ollama/test_ollama_chat_transformation.py | 93 +++++++++++++++++++ 3 files changed, 110 insertions(+), 10 deletions(-) diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 755d647da1d..168b562b358 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -235,6 +235,19 @@ class OllamaChatConfig(BaseConfig): litellm_params["function_name"] = function_name tools: Final = optional_params.pop("tools", None) + def transform_tool_call(typed_tool: ChatCompletionAssistantToolCall) -> OllamaToolCall: + function: Final = OllamaToolCallFunction( + name=typed_tool["function"].get("name") or "", + arguments=json.loads(typed_tool["function"]["arguments"]) + if "arguments" in typed_tool["function"] + else {}, + ) + tool_id: Final = typed_tool.get("id") + ollama_tool_call: Final = ( + OllamaToolCall(function=function) if tool_id is None else OllamaToolCall(id=tool_id, function=function) + ) + return ollama_tool_call + new_messages: Final = [] for m in messages: if isinstance( @@ -248,16 +261,7 @@ class OllamaChatConfig(BaseConfig): for tool in tool_calls: typed_tool = ChatCompletionAssistantToolCall(**tool) if typed_tool["type"] == "function": - arguments = {} - if "arguments" in typed_tool["function"]: - arguments = json.loads(typed_tool["function"]["arguments"]) - ollama_tool_call = OllamaToolCall( - function=OllamaToolCallFunction( - name=typed_tool["function"].get("name") or "", - arguments=arguments, - ) - ) - new_tools.append(ollama_tool_call) + new_tools.append(transform_tool_call(typed_tool)) reasoning_content, parsed_content = extract_reasoning_content(cast(dict, m)) content_str = convert_content_list_to_str(cast(AllMessageValues, m)) images = extract_images_from_message(cast(AllMessageValues, m)) diff --git a/litellm/types/llms/ollama.py b/litellm/types/llms/ollama.py index 4783eac9a8d..fb9d9e8c363 100644 --- a/litellm/types/llms/ollama.py +++ b/litellm/types/llms/ollama.py @@ -1,4 +1,6 @@ from typing_extensions import ( + NotRequired, + ReadOnly, Required, TypedDict, ) @@ -12,6 +14,7 @@ class OllamaToolCallFunction( class OllamaToolCall(TypedDict): + id: NotRequired[ReadOnly[str]] function: OllamaToolCallFunction diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index af1b85cf325..6a80f29ac43 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -856,6 +856,99 @@ class TestOllamaReasoningContentStreaming: class TestOllamaToolCallTransformation: + @pytest.mark.parametrize("stream", [False, True]) + def test_transform_request_keeps_parallel_call_ids_paired_with_results(self, stream: bool): + config: Final = OllamaChatConfig() + messages: Final = cast( + list[AllMessageValues], + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_first", + "type": "function", + "function": { + "name": "read_fixture", + "arguments": '{"path": "first.txt"}', + }, + }, + { + "id": "call_second", + "type": "function", + "function": { + "name": "read_fixture", + "arguments": '{"path": "second.txt"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_second", + "content": "second result", + }, + { + "role": "tool", + "tool_call_id": "call_first", + "content": "first result", + }, + ], + ) + + result: Final = config.transform_request( + model="ollama_chat/test-model", + messages=messages, + optional_params={"stream": stream}, + litellm_params={}, + headers={}, + ) + + assert result["messages"][0]["tool_calls"] == [ + { + "id": "call_first", + "function": {"name": "read_fixture", "arguments": {"path": "first.txt"}}, + }, + { + "id": "call_second", + "function": {"name": "read_fixture", "arguments": {"path": "second.txt"}}, + }, + ] + assert [message.get("tool_call_id") for message in result["messages"][1:]] == [ + "call_second", + "call_first", + ] + + def test_transform_request_does_not_invent_an_id_for_an_unidentified_call(self): + config: Final = OllamaChatConfig() + messages: Final = cast( + list[AllMessageValues], + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": None, + "type": "function", + "function": {"name": "read_fixture", "arguments": "{}"}, + } + ], + } + ], + ) + + result: Final = config.transform_request( + model="ollama_chat/test-model", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert result["messages"][0]["tool_calls"] == [{"function": {"name": "read_fixture", "arguments": {}}}] + def test_transform_request_preserves_tool_calls(self): """ tool_calls on assistant messages must survive transform_request. From b2466f8d7b2c783c04d424009b4b5b1066e06d23 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 12:31:16 -0700 Subject: [PATCH 52/64] chore(cost-map): flag Vertex Gemini computer use and add Cohere Rerank 5 (#49297) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 33 +++++++++++++++++++ model_prices_and_context_window.json | 33 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dac7af6ab7a..1d18903bb07 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29561,6 +29561,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30167,6 +30168,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_pdf_input": true, @@ -30210,6 +30212,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30274,6 +30277,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30334,6 +30338,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30401,6 +30406,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -32405,6 +32411,7 @@ "text" ], "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -32479,6 +32486,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32543,6 +32551,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32603,6 +32612,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32670,6 +32680,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -47624,6 +47635,24 @@ "output_cost_per_token": 0.0, "source": "https://cohere.com/pricing" }, + "rerank-v5.0-fast": { + "input_cost_per_token": 2e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, + "rerank-v5.0-pro": { + "input_cost_per_token": 5e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -52713,6 +52742,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76498,6 +76528,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76543,6 +76574,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76654,6 +76686,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dac7af6ab7a..1d18903bb07 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29561,6 +29561,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30167,6 +30168,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_pdf_input": true, @@ -30210,6 +30212,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30274,6 +30277,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30334,6 +30338,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30401,6 +30406,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -32405,6 +32411,7 @@ "text" ], "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -32479,6 +32486,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32543,6 +32551,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32603,6 +32612,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32670,6 +32680,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -47624,6 +47635,24 @@ "output_cost_per_token": 0.0, "source": "https://cohere.com/pricing" }, + "rerank-v5.0-fast": { + "input_cost_per_token": 2e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, + "rerank-v5.0-pro": { + "input_cost_per_token": 5e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -52713,6 +52742,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76498,6 +76528,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76543,6 +76574,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76654,6 +76686,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, From 0d5cf9b2f7d8cba315cd6c550bd98aa9a2fa88e1 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 12:32:03 -0700 Subject: [PATCH 53/64] fix(headroom): forward tenant-scoped session affinity (#49255) --- .../litellm_core_utils/provider_affinity.py | 42 ++++- .../guardrail_hooks/headroom/headroom.py | 91 ++++----- litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/types/litellm_params.py | 2 + .../guardrails/guardrail_hooks/headroom.py | 12 +- .../test_provider_affinity.py | 16 +- .../guardrail_hooks/test_headroom.py | 174 +++++++++--------- tests/unit/types/test_litellm_params.py | 3 + 8 files changed, 209 insertions(+), 132 deletions(-) diff --git a/litellm/litellm_core_utils/provider_affinity.py b/litellm/litellm_core_utils/provider_affinity.py index c831e92c09d..767eab962c8 100644 --- a/litellm/litellm_core_utils/provider_affinity.py +++ b/litellm/litellm_core_utils/provider_affinity.py @@ -2,7 +2,12 @@ import re from collections.abc import Mapping from typing import Final -from litellm.constants import SESSION_ID_FROM_TRACE_ID_PARAM_KEY, SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, + SESSION_ID_FROM_TRACE_ID_PARAM_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping _HTTP_HEADER_NAME_PATTERN: Final = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") _FORBIDDEN_AFFINITY_HEADERS: Final = frozenset( @@ -98,6 +103,41 @@ def get_stable_session_id(litellm_params: object | None) -> str | None: return _non_generated_litellm_session_id(litellm_params) or _metadata_session_id(litellm_params) +def get_caller_session_id(litellm_params: object | None) -> str | None: + """Resolve a caller conversation without accepting a generated session or a promoted trace id.""" + if litellm_params is None: + return None + metadata_values: Final = _metadata_values(litellm_params) + preserved_session: Final = next( + ( + value + for metadata in metadata_values + if isinstance(value := _get_value(metadata, CALLER_SESSION_ID_METADATA_KEY), str) and value + ), + None, + ) + if preserved_session: + return preserved_session + direct_session: Final = _get_value(litellm_params, "session_id") + if isinstance(direct_session, str) and direct_session: + return direct_session + + proxy_request: Final = _get_value(litellm_params, "proxy_server_request") + headers: Final = as_str_mapping(_get_value(proxy_request, "headers")) + trace_header: Final = next( + (value for key, value in (headers or {}).items() if key.lower() == "x-litellm-trace-id"), None + ) + candidates: Final = ( + get_caller_litellm_session_id(litellm_params), + *( + _get_value(metadata, "session_id") + for metadata in metadata_values + if not _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) + ), + ) + return next((value for value in candidates if isinstance(value, str) and value and value != trace_header), None) + + def add_provider_affinity_header( headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None ) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index dac802a1ddc..7f48766bf26 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -1,10 +1,10 @@ from __future__ import annotations +import hashlib import json import math import re import time -import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar, Final, Literal, Protocol, TypeGuard @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( group_tool_exchanges, has_tool_with_name, ) +from litellm.litellm_core_utils.provider_affinity import get_caller_session_id from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] httpxSpecialProvider, @@ -39,6 +40,7 @@ from litellm.proxy.guardrails.guardrail_hooks.content_text import ( merge_rewritten_text_parts, ) from litellm.proxy.spend_tracking.compression_savings import HEADROOM_GUARDRAIL_PROVIDER +from litellm.router_utils.common_utils import request_user_api_key_auth from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.integrations.custom_logger import ( @@ -47,6 +49,11 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopRequestPatch, as_converted_stream, ) +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HEADROOM_REQUEST_CONTEXT_KEY, + HeadroomRequestContext, +) from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -357,15 +364,18 @@ def _build_headroom_retrieve_tool() -> dict[str, object]: } -def _resolve_call_id(logging_obj: object, request_state: dict[str, object]) -> str | None: - """Resolve the litellm_call_id shared by a request's pre-call hook and its - agentic-loop hooks, so CCR hash validation can be scoped per call instead - of trusting any hash-shaped string that shows up in message text.""" - logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None) - if isinstance(logging_call_id, str) and logging_call_id: - return logging_call_id - kwargs_call_id: Final = request_state.get("litellm_call_id") - return kwargs_call_id if isinstance(kwargs_call_id, str) else None +def _session_affinity(request_data: Mapping[str, object]) -> str | None: + auth: Final = request_user_api_key_auth(request_data) + session_id: Final = get_caller_session_id(request_data) + if not isinstance(auth, UserAPIKeyAuth) or not session_id: + return None + tenant: Final = ( + ("key", auth.api_key) if auth.api_key else ("principal", auth.org_id, auth.team_id, auth.user_id, auth.agent_id) + ) + if not any(tenant[1:]): + return None + identity: Final = json.dumps(["headroom-affinity-v1", tenant, session_id], separators=(",", ":")) + return hashlib.sha256(identity.encode("utf-8")).hexdigest() def has_headroom_retrieve_tool(tools: object) -> bool: @@ -521,7 +531,6 @@ class HeadroomGuardrail(CustomGuardrail): self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, ) - self._issued_hashes_by_call_id: dict[str, tuple[frozenset[str], float]] = {} super().__init__( # pyright: ignore[reportUnknownMemberType] guardrail_name=guardrail_name, event_hook=event_hook, @@ -540,10 +549,12 @@ class HeadroomGuardrail(CustomGuardrail): value: Final = headers.get(BYPASS_HEADER) return str(value).lower() == "true" - def _request_headers(self) -> dict[str, str]: + def _request_headers(self, affinity: str | None) -> dict[str, str]: headers: Final[dict[str, str]] = {"Content-Type": "application/json"} if self.headroom_api_key: headers["Authorization"] = f"Bearer {self.headroom_api_key}" + if affinity is not None: + headers["x-headroom-affinity"] = affinity return headers @staticmethod @@ -569,14 +580,6 @@ class HeadroomGuardrail(CustomGuardrail): super().update_in_memory_litellm_params(litellm_params) self.timeout = self._resolve_timeout(litellm_params.timeout) - def _prune_expired_hashes(self) -> None: - now: Final = time.monotonic() - self._issued_hashes_by_call_id = { - call_id: (hashes, expiry) - for call_id, (hashes, expiry) in self._issued_hashes_by_call_id.items() - if expiry > now - } - def _handle_compress_failure( self, messages: list[dict[str, object]], @@ -596,16 +599,19 @@ class HeadroomGuardrail(CustomGuardrail): self, messages: list[dict[str, object]], model: str | None, + affinity: str | None, ) -> _CompressResult: payload: Final[dict[str, object]] = {"messages": messages} if model: payload["model"] = model + if self.ccr_retrieval: + payload["config"] = {"mode": "ccr"} try: raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped url=f"{self.headroom_api_base}/v1/compress", json=payload, - headers=self._request_headers(), + headers=self._request_headers(affinity), timeout=self.timeout, ) except httpx.HTTPStatusError as e: @@ -734,7 +740,7 @@ class HeadroomGuardrail(CustomGuardrail): stats["tokens_saved"] = tokens_before - tokens_after return _CompressResult(filtered, True, stats, _read_ccr_hashes(body)) - async def _call_retrieve(self, hash_value: str, query: str | None = None) -> str: + async def _call_retrieve(self, hash_value: str, affinity: str | None, query: str | None = None) -> str: params: Final[dict[str, str]] = {} if query: params["query"] = query @@ -743,7 +749,7 @@ class HeadroomGuardrail(CustomGuardrail): raw_response: HttpxResponse = await self.async_handler.get( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.get is untyped url=f"{self.headroom_api_base}/v1/retrieve/{hash_value}", params=params, - headers=self._request_headers(), + headers=self._request_headers(affinity), timeout=self.timeout, ) except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e: @@ -777,13 +783,14 @@ class HeadroomGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: if input_type != "request": return inputs + request_data.pop(HEADROOM_REQUEST_CONTEXT_KEY, None) if self._should_bypass(request_data): verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER) return inputs @@ -809,9 +816,7 @@ class HeadroomGuardrail(CustomGuardrail): # /v1/compress grows a field for sending the live turn as the retrieval # query without compressing it: query-aware compression reads the newest # user message, so it is withheld here at some cost to history ranking. - # request_data is a bare dict on the base signature; narrow it before - # reading the untranslated messages so long tool names can be recovered. - raw_messages: Final = _REQUEST_DATA_ADAPTER.validate_python(request_data).get("messages") + raw_messages: Final = request_data.get("messages") raw_retrieve_call_ids: Final = _raw_retrieve_call_ids(raw_messages) protected_indices: Final = _protected_indices(messages, raw_retrieve_call_ids) compressible: Final = [m for i, m in enumerate(messages) if i not in protected_indices] @@ -820,9 +825,11 @@ class HeadroomGuardrail(CustomGuardrail): model: Final = self.headroom_model or request_data.get("model") start_time: Final = time.time() + affinity: Final = _session_affinity(request_data) result: Final = await self._call_compress( messages=_flatten_messages_for_compression(compressible), model=model if isinstance(model, str) else None, + affinity=affinity, ) end_time: Final = time.time() @@ -868,12 +875,9 @@ class HeadroomGuardrail(CustomGuardrail): if not hashes: return {**inputs, "structured_messages": compressed} # pyright: ignore[reportReturnType] - self._prune_expired_hashes() - call_id = _resolve_call_id(logging_obj, request_data) - if not call_id: - call_id = str(uuid.uuid4()) - request_data["litellm_call_id"] = call_id - self._issued_hashes_by_call_id[call_id] = (frozenset(hashes), time.monotonic() + _HASH_CACHE_TTL_SECONDS) + request_data[HEADROOM_REQUEST_CONTEXT_KEY] = HeadroomRequestContext( + frozenset(hashes), affinity, time.monotonic() + _HASH_CACHE_TTL_SECONDS + ) existing_tools: Final = inputs.get("tools") retrieve_tool: Final = _build_headroom_retrieve_tool() @@ -934,9 +938,11 @@ class HeadroomGuardrail(CustomGuardrail): ) -> AgenticLoopPlan: tool_calls: Final[list[dict[str, object]]] = tools.get("tool_calls", []) - self._prune_expired_hashes() - call_id: Final = _resolve_call_id(logging_obj, kwargs) - valid_hashes = self._issued_hashes_by_call_id.get(call_id, (frozenset(), 0.0))[0] if call_id else frozenset() + context: Final = _REQUEST_DATA_ADAPTER.validate_python(kwargs).get(HEADROOM_REQUEST_CONTEXT_KEY) + valid_context: Final = ( + context if isinstance(context, HeadroomRequestContext) and context.expires_at > time.monotonic() else None + ) + valid_hashes: Final = valid_context.issued_hashes if valid_context is not None else frozenset() retrieved: Final[list[tuple[dict[str, object], str]]] = [] for tc in tool_calls: @@ -944,11 +950,6 @@ class HeadroomGuardrail(CustomGuardrail): raw_hash = arguments.get("hash", "") if isinstance(arguments, dict) else "" hash_value = str(raw_hash).lower() query = arguments.get("query") if isinstance(arguments, dict) else None - # A hash is only honored if it was issued by *this request's own* - # Headroom /v1/compress call, scoped by litellm_call_id. Scoping by - # message text alone is forgeable -- an attacker can plant a - # hash-shaped string in their own prompt, and a hash issued for one - # request would validate for any other request that echoes it back. if hash_value not in valid_hashes: verbose_proxy_logger.warning( "Headroom CCR: rejecting hash=%s not produced by current request compression", @@ -958,6 +959,7 @@ class HeadroomGuardrail(CustomGuardrail): else: content = await self._call_retrieve( hash_value=hash_value, + affinity=valid_context.affinity if valid_context is not None else None, query=str(query) if query else None, ) verbose_proxy_logger.debug("Headroom CCR: retrieved hash=%s (%d chars)", hash_value, len(content)) @@ -996,7 +998,12 @@ class HeadroomGuardrail(CustomGuardrail): max_tokens=max_tokens, optional_params=optional_params_without_max_tokens, kwargs={ - k: v for k, v in kwargs.items() if not k.startswith("_headroom") and k != "litellm_logging_obj" + **{ + k: v + for k, v in _REQUEST_DATA_ADAPTER.validate_python(kwargs).items() + if not k.startswith("_headroom") and k != "litellm_logging_obj" + }, + **({HEADROOM_REQUEST_CONTEXT_KEY: valid_context} if valid_context is not None else {}), }, ), metadata={"tool_type": "headroom_ccr"}, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9b52f2f8a8d..ae9a6d4b959 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2566,6 +2566,7 @@ def refresh_proxy_server_request_body_snapshot( VECTOR_STORE_CALLER_METADATA_KWARG, STREAM_USAGE_INJECTED_FIELD, "_litellm_strip_stream_usage", + "_headroom_request_context", } ) | _TRANSPORT_ONLY_CREDENTIAL_KEYS diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 2065093f295..3c2c636b5ed 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets from litellm.types.caching import DynamicCacheControl from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionUserMessage + from litellm.types.proxy.guardrails.guardrail_hooks.headroom import HeadroomRequestContext from litellm.types.proxy.litellm_pre_call_utils import SecretFields from litellm.types.router import ConfigurableClientsideParamsCustomAuth, DeploymentTypedDict, RetryPolicy from litellm.types.router_weights import RouterWeights @@ -352,6 +353,7 @@ class CallState: @dataclass(frozen=True, slots=True, kw_only=True) class AgenticLoopState: + headroom_context: "HeadroomRequestContext | None" = field(default=None, metadata=wire("_headroom_request_context")) depth: int | None = field(default=None, metadata=wire("_agentic_loop_depth")) fingerprints: Sequence[str] | None = field(default=None, metadata=wire("_agentic_loop_fingerprints")) api_surface: Literal["chat_completions", "responses"] | None = field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py index 0df0f0c80ef..76e806e3899 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py @@ -1,9 +1,19 @@ -from typing import Literal +from dataclasses import dataclass +from typing import Final, Literal from pydantic import BaseModel, Field from .base import GuardrailConfigModel +HEADROOM_REQUEST_CONTEXT_KEY: Final = "_headroom_request_context" + + +@dataclass(frozen=True, slots=True) +class HeadroomRequestContext: + issued_hashes: frozenset[str] + affinity: str | None + expires_at: float + class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]): api_base: str | None = Field( diff --git a/tests/unit/litellm_core_utils/test_provider_affinity.py b/tests/unit/litellm_core_utils/test_provider_affinity.py index edf4f5169b5..6ccd6b17c01 100644 --- a/tests/unit/litellm_core_utils/test_provider_affinity.py +++ b/tests/unit/litellm_core_utils/test_provider_affinity.py @@ -1,12 +1,16 @@ +from collections.abc import Callable + import pytest from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.litellm_core_utils.provider_affinity import ( add_provider_affinity_header, + get_caller_session_id, get_stable_session_id, ) +@pytest.mark.parametrize("resolve", [get_stable_session_id, get_caller_session_id]) @pytest.mark.parametrize( ("litellm_params", "expected"), [ @@ -16,18 +20,22 @@ from litellm.litellm_core_utils.provider_affinity import ( ({"litellm_metadata": {"session_id": "litellm-metadata-session"}}, "litellm-metadata-session"), ], ) -def test_get_stable_session_id_uses_explicit_session_sources(litellm_params: dict, expected: str): - assert get_stable_session_id(litellm_params) == expected +def test_session_ids(litellm_params: dict[str, object], expected: str, resolve: Callable[[object], str | None]) -> None: + assert resolve(litellm_params) == expected def test_get_stable_session_id_does_not_use_trace_id(): assert get_stable_session_id({"litellm_trace_id": "per-request-trace"}) is None + assert get_caller_session_id( + {"litellm_session_id": "trace", "proxy_server_request": {"headers": {"x-litellm-trace-id": "trace"}}} + ) is None +@pytest.mark.parametrize("resolve", [get_stable_session_id, get_caller_session_id]) @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) -def test_get_stable_session_id_ignores_proxy_generated_session(metadata_key: str): +def test_generated_sessions_are_ignored(metadata_key: str, resolve: Callable[[object], str | None]) -> None: assert ( - get_stable_session_id( + resolve( { "litellm_session_id": "generated-session", metadata_key: { diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index 741fea5f2b9..0cdcb2e2d07 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -23,7 +23,6 @@ Tests cover: import datetime import json -import time from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -33,17 +32,22 @@ import respx from fastapi import HTTPException import litellm - from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( - HeadroomGuardrail, - has_headroom_retrieve_tool, HEADROOM_RETRIEVE_TOOL_NAME, + HeadroomGuardrail, + _session_affinity, + has_headroom_retrieve_tool, ) from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, ) from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, HEADROOM_STREAM_OPTIONS_KEY +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HEADROOM_REQUEST_CONTEXT_KEY, + HeadroomRequestContext, +) from litellm.types.utils import ( CallTypes, GenericGuardrailAPIInputs, @@ -148,6 +152,10 @@ def guardrail() -> HeadroomGuardrail: return _make_guardrail() +def _issued_context(*hashes: str, affinity: str | None = None) -> HeadroomRequestContext: + return HeadroomRequestContext(frozenset(hashes), affinity, float("inf")) + + def _recorded_guardrail_entries(request_data: dict) -> list: for container_key in ("metadata", "litellm_metadata"): container = request_data.get(container_key) @@ -175,20 +183,22 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages( structured_messages=ORIGINAL_MESSAGES, ) mock_response = _make_compress_response(COMPRESSED_MESSAGES) - request_data = {"model": "gpt-4o"} + auth: Final = UserAPIKeyAuth(api_key="tenant") + request_data: Final = {"session_id": "conversation", "metadata": {"user_api_key_auth": auth}} with patch.object( guardrail.async_handler, "post", new_callable=AsyncMock, return_value=mock_response, - ): + ) as mock_post: result = await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, input_type="request", ) + assert mock_post.call_args.kwargs["headers"]["x-headroom-affinity"] == _session_affinity(request_data) assert result.get("structured_messages") == EXPECTED_MESSAGES entries = _recorded_guardrail_entries(request_data) @@ -510,10 +520,6 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( tool_id="call_abc123", ) messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -530,11 +536,12 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac", affinity="original")}, ) assert plan.run_agentic_loop is True assert plan.request_patch is not None + assert HEADROOM_REQUEST_CONTEXT_KEY in plan.request_patch.kwargs follow_up = plan.request_patch.messages assert follow_up is not None @@ -545,6 +552,7 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( assert tool_result_message["tool_call_id"] == "call_abc123" mock_get.assert_called_once() + assert mock_get.call_args.kwargs["headers"]["x-headroom-affinity"] == "original" call_url = mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0] assert "b573993006976af767214fac" in call_url @@ -576,10 +584,6 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( "content": "Retrieve more: hash=deadbeef000000000000dead", } ] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"deadbeef000000000000dead"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -596,7 +600,7 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("deadbeef000000000000dead")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -606,9 +610,13 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( @pytest.mark.asyncio +@pytest.mark.parametrize("context", [ + None, {"issued_hashes": ["deadbeef000000000000dead"]}, + HeadroomRequestContext(frozenset({"deadbeef000000000000dead"}), None, 0.0), +]) async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( - guardrail: HeadroomGuardrail, -): + guardrail: HeadroomGuardrail, context: object, +) -> None: """A hash-shaped string planted in message text must not be honored when this guardrail has no record of ever issuing it, even if it's echoed back in the current request's own messages (e.g. via prompt injection).""" @@ -625,8 +633,6 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( arguments={"hash": "deadbeef000000000000dead"}, tool_id="call_xyz", ) - assert not guardrail._issued_hashes_by_call_id - with patch.object( guardrail.async_handler, "get", @@ -641,7 +647,7 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-unknown"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: context}, ) mock_get.assert_not_called() @@ -653,17 +659,11 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( @pytest.mark.asyncio -async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_call( +async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_request( guardrail: HeadroomGuardrail, ): - """A hash issued for one request must not be retrievable by a different - request just because the second request echoes that hash-shaped string - back in its own messages -- retrieval must be scoped per litellm_call_id, - not derived by re-scanning attacker-controlled message text.""" - guardrail._issued_hashes_by_call_id["call-A"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) + issuer_data: Final[dict[str, object]] = {"model": "gpt-4o", "litellm_call_id": "call-A"} + await _apply(guardrail, ORIGINAL_MESSAGES, ["b573993006976af767214fac"], request_data=issuer_data) tool_calls = [ { @@ -693,7 +693,7 @@ async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_c anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-B"}, + kwargs={"litellm_call_id": "call-A"}, ) mock_get.assert_not_called() @@ -736,10 +736,6 @@ async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_ } ] messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -756,7 +752,7 @@ async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_ anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -806,7 +802,6 @@ async def test_responses_retrieval_followup_reuses_the_model_the_caller_requeste "arguments": json.dumps({"hash": retrieve_hash}), } ] - guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({retrieve_hash}), time.monotonic() + 999) with patch.object( guardrail.async_handler, @@ -832,7 +827,7 @@ async def test_responses_retrieval_followup_reuses_the_model_the_caller_requeste anthropic_messages_optional_request_params={}, logging_obj=logging_obj, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(retrieve_hash)}, ) assert plan.request_patch is not None @@ -876,10 +871,6 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag } ] messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -896,7 +887,7 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -929,16 +920,20 @@ HASH_SHAPED_HISTORY = [ ] -async def _apply(guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None) -> dict: - request_data = {"model": "claude-sonnet-5"} +async def _apply( + guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None, + request_data: dict[str, object] | None = None, +) -> dict: + data: Final = request_data if request_data is not None else {"model": "claude-sonnet-5"} def _echo(**kwargs): + assert kwargs["json"].get("config") == ({"mode": "ccr"} if guardrail.ccr_retrieval else None) return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"])), ccr_hashes=ccr_hashes) with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): return await guardrail.apply_guardrail( inputs=GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))), - request_data=request_data, + request_data=data, input_type="request", ) @@ -950,10 +945,11 @@ async def test_hash_shaped_text_in_history_never_injects_retrieve_tool(guardrail tool must not appear and no hash may be registered as issued. Covers a service that omits ccr_hashes, returns it empty, or returns a non-list.""" for ccr_hashes in (None, [], "b573993006976af767214fac"): - result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes) + request_data: Final[dict[str, object]] = {} + result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes, request_data=request_data) assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id + assert HEADROOM_REQUEST_CONTEXT_KEY not in request_data @pytest.mark.asyncio @@ -961,15 +957,18 @@ async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardr """Only the hashes the service reports in ccr_hashes are honored, in the service's own 12 to 24 hex grammar; anything else is dropped because each entry is interpolated into the /v1/retrieve URL.""" + request_data: Final[dict[str, object]] = {} result = await _apply( guardrail, HASH_SHAPED_HISTORY, ccr_hashes=["98CA69107318", "b573993006976af767214fac", "../../etc/passwd", "tooshort", 42], + request_data=request_data, ) assert has_headroom_retrieve_tool(result.get("tools") or []) - ((issued, _expiry),) = guardrail._issued_hashes_by_call_id.values() - assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"}) + context: Final = request_data[HEADROOM_REQUEST_CONTEXT_KEY] + assert isinstance(context, HeadroomRequestContext) + assert context.issued_hashes == frozenset({"98ca69107318", "b573993006976af767214fac"}) @pytest.mark.asyncio @@ -987,7 +986,6 @@ async def test_ccr_retrieval_disabled_ignores_service_declared_hashes(monkeypatc assert result["structured_messages"][-1] == HASH_SHAPED_HISTORY[-1] assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id @pytest.mark.asyncio @@ -1498,6 +1496,7 @@ async def test_apply_guardrail_fail_open_does_not_register_hashes_from_original_ ] guardrail = _make_guardrail(unreachable_fallback="fail_open") + request_data: Final = {HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("deadbeef000000000000dead")} inputs = GenericGuardrailAPIInputs( texts=["hello"], structured_messages=messages_with_fake_hash, @@ -1511,13 +1510,13 @@ async def test_apply_guardrail_fail_open_does_not_register_hashes_from_original_ ): result = await guardrail.apply_guardrail( inputs=inputs, - request_data={}, + request_data=request_data, input_type="request", ) assert result["structured_messages"] == messages_with_fake_hash assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id + assert HEADROOM_REQUEST_CONTEXT_KEY not in request_data @pytest.mark.asyncio @@ -2164,18 +2163,18 @@ def _openai_completion_payload(message: dict, finish_reason: str) -> dict: } -def _openai_tool_call_payload() -> dict: +def _openai_tool_call_payload(hash_value: str = CCR_HASH) -> dict: return _openai_completion_payload( { "role": "assistant", "content": None, "tool_calls": [ { - "id": "call_ccr", + "id": f"call_{hash_value}", "type": "function", "function": { "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": CCR_HASH}), + "arguments": json.dumps({"hash": hash_value}), }, } ], @@ -2215,14 +2214,14 @@ def _openai_responses_payload(output_item: dict) -> dict: } -def _openai_responses_retrieve_call_payload() -> dict: +def _openai_responses_retrieve_call_payload(hash_value: str = CCR_HASH) -> dict: return _openai_responses_payload( { "type": "function_call", "id": "fc_ccr", - "call_id": "call_ccr", + "call_id": f"call_{hash_value}", "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": CCR_HASH}), + "arguments": json.dumps({"hash": hash_value}), "status": "completed", } ) @@ -2381,10 +2380,6 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( to a client that never declared the tool.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "callbacks", [guardrail]) @@ -2392,6 +2387,7 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( upstream = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response(200, json=_openai_tool_call_payload()), + httpx.Response(200, json=_openai_tool_call_payload("a" * 24)), httpx.Response(200, json=_openai_text_payload(final_answer)), ] ) @@ -2408,20 +2404,22 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( tools=[_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH, "a" * 24, affinity="original-affinity")}, ) chunks = [chunk async for chunk in response] streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) assert streamed_text == final_answer assert not any(chunk.choices and chunk.choices[0].delta.tool_calls for chunk in chunks) - mock_get.assert_called_once() - assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + assert mock_get.call_count == 2 + assert mock_get.call_args.kwargs["url"].endswith("a" * 24) + assert all(call.kwargs["headers"]["x-headroom-affinity"] == "original-affinity" for call in mock_get.call_args_list) - assert len(upstream.calls) == 2 - followup_body = json.loads(upstream.calls[1].request.content) + assert len(upstream.calls) == 3 + followup_body = json.loads(upstream.calls[2].request.content) assert not followup_body.get("stream") assert original_content in json.dumps(followup_body["messages"]) - assert not any(key.startswith("_headroom_interception") for key in followup_body) + assert not any(key.startswith("_headroom") for key in followup_body) @pytest.mark.asyncio @@ -2435,10 +2433,6 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( instead of streaming a headroom_retrieve function_call to the client.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "callbacks", [guardrail]) @@ -2446,6 +2440,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( side_effect=[ httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_retrieve_call_payload("a" * 24)), httpx.Response(200, json=_openai_responses_text_payload(final_answer)), ] ) @@ -2462,6 +2457,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( tools=[_responses_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH, "a" * 24, affinity="original-affinity")}, ) events = [event async for event in response] @@ -2471,14 +2467,15 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( assert streamed_text == final_answer assert not any("function_call" in str(getattr(event, "type", "")) for event in events) assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events) - mock_get.assert_called_once() - assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + assert mock_get.call_count == 2 + assert mock_get.call_args.kwargs["url"].endswith("a" * 24) + assert all(call.kwargs["headers"]["x-headroom-affinity"] == "original-affinity" for call in mock_get.call_args_list) - assert len(upstream.calls) == 2 - followup_body = json.loads(upstream.calls[1].request.content) + assert len(upstream.calls) == 3 + followup_body = json.loads(upstream.calls[2].request.content) assert not followup_body.get("stream") assert original_content in json.dumps(followup_body["input"]) - assert not any(key.startswith("_headroom_interception") for key in followup_body) + assert not any(key.startswith("_headroom") for key in followup_body) def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( @@ -2491,10 +2488,6 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( completed response object.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -2518,6 +2511,7 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( tools=[_responses_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH)}, ) events = list(response) @@ -2917,7 +2911,6 @@ def _anthropic_response_with_text_and_tool_call() -> dict: async def _plan_for(guardrail: HeadroomGuardrail, response, messages: list): - guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({"h" * 24}), time.monotonic() + 60) logging_obj = MagicMock() logging_obj.litellm_call_id = "call-1" logging_obj.model_call_details = {} @@ -2938,7 +2931,7 @@ async def _plan_for(guardrail: HeadroomGuardrail, response, messages: list): anthropic_messages_optional_request_params={}, logging_obj=logging_obj, stream=False, - kwargs={}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("h" * 24)}, ) @@ -3033,7 +3026,7 @@ async def test_retrieve_call_passes_bounded_timeout(guardrail: HeadroomGuardrail new_callable=AsyncMock, return_value=_make_retrieve_response("original"), ) as mock_get: - result = await guardrail._call_retrieve("a" * 24) + result = await guardrail._call_retrieve("a" * 24, affinity=None) assert result == "original" timeout = _timeout_of(mock_get.call_args) @@ -3143,3 +3136,16 @@ def test_unusable_timeout_falls_back_to_the_default(configured: float): assert guardrail.timeout.read == 60.0 assert guardrail.timeout.connect == 5.0 + + +@pytest.mark.parametrize("session", [None, "conversation", "other"]) +@pytest.mark.parametrize("scope", [None, {}, {"api_key": "key"}, {"api_key": "other"}, {"user_id": "user"}]) +def test_affinity_requires_trusted_identity_and_caller_session( + scope: dict[str, str] | None, session: str | None +) -> None: + baseline: Final = {"session_id": "conversation", "metadata": {"user_api_key_auth": UserAPIKeyAuth(api_key="key")}} + auth: Final = UserAPIKeyAuth(**scope) if scope else scope + candidate: Final = {"session_id": session, "metadata": {"user_api_key_auth": auth}} + affinity: Final = _session_affinity(candidate) + assert (affinity is None) is (session is None or not isinstance(auth, UserAPIKeyAuth)) + assert (affinity == _session_affinity(baseline)) is (candidate == baseline) diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index e369b7e93ab..22eb9929e7f 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -33,6 +33,7 @@ from litellm.types.litellm_params import ( wire_names, ) from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionUserMessage +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import HeadroomRequestContext from litellm.types.proxy.litellm_pre_call_utils import SecretFields from litellm.types.router import ( ConfigurableClientsideParamsCustomAuth, @@ -231,6 +232,7 @@ OPTION_NAMES: Final = ( ) AGENTIC_LOOP_STATE_NAMES: Final = ( + "_headroom_request_context", "_agentic_loop_depth", "_agentic_loop_fingerprints", "_agentic_loop_api_surface", @@ -529,6 +531,7 @@ TYPE_HINT_NAMESPACE: Final[Mapping[str, object]] = { "Logging": Logging, "SecretFields": SecretFields, "CompactionState": CompactionState, + "HeadroomRequestContext": HeadroomRequestContext, "RouterWeights": RouterWeights, "AttemptedFallbackTargets": AttemptedFallbackTargets, } From 7b259690e7eac9720ee21067a71152912f1080ef Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 12:36:11 -0700 Subject: [PATCH 54/64] fix(proxy): don't send llm_exceptions alert for BudgetExceededError (#49290) * fix(proxy): don't send llm_exceptions alert for BudgetExceededError * test(proxy): verify budget alert suppression through failure hooks --------- Co-authored-by: Arvind Kumar --- litellm/proxy/utils.py | 2 + ..._utils_model_creation_and_error_logging.py | 49 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 00dbcd0f9a8..cc72dd03a40 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -633,6 +633,8 @@ def _is_client_error_exception(exc: Exception) -> bool: return exc.status_code < 500 if isinstance(exc, ProxyException): return not (exc.code.isdigit() and int(exc.code) >= 500) + if isinstance(exc, litellm.BudgetExceededError): + return exc.status_code < 500 return False diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 3cdd2edb2fb..0a483840336 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1844,6 +1844,26 @@ class TestPostCallFailureHookLLMExceptionAlerting: alerting_handler = await self._alerted(HTTPException(status_code=429, detail="rate limited")) assert alerting_handler.called is False + @pytest.mark.asyncio + @pytest.mark.parametrize("status_code, should_alert", ((400, False), (422, False), (429, False), (500, True))) + async def test_budget_exceeded_alerts_only_for_server_status( + self, status_code: int, should_alert: bool, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "budget_exceeded_status_code", status_code) + exc: Final = litellm.BudgetExceededError(current_cost=1.0, max_budget=0.5) + assert (await self._alerted(exc)).called is should_alert + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "exc", + ( + litellm.AuthenticationError(message="invalid provider key", model="test-model", llm_provider="openai"), + litellm.RateLimitError(message="provider quota exceeded", model="test-model", llm_provider="openai"), + ), + ) + async def test_provider_client_status_still_alerts(self, exc: Exception) -> None: + assert (await self._alerted(exc)).called is True + class TestPostCallFailureHookProxyExceptionLogging: """A guardrail block raises a ProxyException; on an LLM route it must still @@ -4109,3 +4129,32 @@ def test_handle_exception_on_proxy_logs_bug_report_only_for_unmapped_500(caplog) assert provider_result.code == internal_result.code == "500" assert ISSUE_URL_BASE in caplog.text assert ISSUE_URL_BASE not in internal_result.message + + +def test_is_client_error_exception_budget_exceeded_is_client_error() -> None: + from litellm.exceptions import BudgetExceededError + from litellm.proxy.utils import _is_client_error_exception + + exc: Final = BudgetExceededError(current_cost=100.56, max_budget=100.0) + assert _is_client_error_exception(exc) is True + + +def test_is_client_error_exception_budget_exceeded_5xx_is_not_client_error() -> None: + from litellm.exceptions import BudgetExceededError + from litellm.proxy.utils import _is_client_error_exception + + exc: Final = BudgetExceededError(current_cost=100.56, max_budget=100.0) + exc.status_code = 500 + assert _is_client_error_exception(exc) is False + + +def test_is_client_error_exception_existing_behaviour_unchanged() -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.utils import _is_client_error_exception + + assert _is_client_error_exception(HTTPException(status_code=400, detail="x")) is True + assert _is_client_error_exception(HTTPException(status_code=500, detail="x")) is False + assert _is_client_error_exception( + ProxyException(message="x", type="x", param=None, code="429") + ) is True + assert _is_client_error_exception(ValueError("x")) is False From 67954b5c3a4b00a1da3c86b18bca8452194002ac Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 12:38:21 -0700 Subject: [PATCH 55/64] fix(proxy): apply key and team router_settings to /v1/images/generations (#49286) * fix(proxy): apply key and team router_settings to /v1/images/generations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): format router settings helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): resolve key model_group_alias once on image generations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 93 ++++++---- litellm/proxy/image_endpoints/endpoints.py | 18 ++ ...st_key_router_settings_image_generation.py | 162 ++++++++++++++++++ .../proxy/image_endpoints/test_endpoints.py | 148 +++++++++++++++- 4 files changed, 389 insertions(+), 32 deletions(-) create mode 100644 tests/integration/routing/test_key_router_settings_image_generation.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 762f552dadc..9b2813a8cd8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -151,7 +151,7 @@ from litellm.router_utils.common_utils import resolve_model_group_alias, resolve from litellm.types.extraction import ExtractionResponse from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError -from litellm.types.router_weights import validate_router_weights +from litellm.types.router_weights import RouterWeights, validate_router_weights _LateResponseT = TypeVar("_LateResponseT", bound=Response) _LlmCallT = TypeVar("_LlmCallT") @@ -894,6 +894,61 @@ async def _resolve_per_request_model_group_alias( return target +async def hierarchical_router_settings_updates( + requested_model: object, + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router, + proxy_config: ProxyConfig, + proxy_logging_obj: ProxyLogging, + apply_model_group_alias: bool, +) -> Mapping[str, object]: + # Apply hierarchical router_settings (Key > Team) + # Global router_settings are already on the Router object itself. + # This sits with the other alias rewrites, and ahead of the guardrail + # merge and the pre-call hooks, so everything that keys off the model + # group -- model-level guardrails, per-model budgets and rate limits, + # the logging object -- sees the group that will actually serve. + from litellm.proxy.proxy_server import prisma_client + + router_settings: Final[Mapping[str, object] | None] = _ROUTER_SETTINGS_ADAPTER.validate_python( + await proxy_config._get_hierarchical_router_settings( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + ) + if router_settings is None: + return MappingProxyType({}) + + # If router_settings found (from key or team), apply them + # Pass settings as per-request overrides instead of creating a new Router + # This avoids expensive Router instantiation on each request + def validated_router_weights() -> RouterWeights | None: + try: + return validate_router_weights(router_settings.get("weights")) + except ValidationError: + verbose_proxy_logger.warning("Ignoring invalid saved router weights; update team/key router_settings") + return None + + router_weights: Final = validated_router_weights() + alias_target: Final = ( + await _resolve_per_request_model_group_alias( + requested_model=requested_model, + router_settings=router_settings, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + ) + if apply_model_group_alias + else None + ) + settings_updates: Final[Mapping[str, object]] = { + "router_settings_override": router_settings, + "_router_weights": router_weights, + **({"model": alias_target} if alias_target is not None else {}), + } + return MappingProxyType(settings_updates) + + _REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None) _DECODED_JSON: Final = TypeAdapter(object) _ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) @@ -2495,41 +2550,17 @@ class ProxyBaseLLMRequestProcessing: ): self.data["model"] = user_api_key_dict.aliases[self.data["model"]] - # Apply hierarchical router_settings (Key > Team) - # Global router_settings are already on the Router object itself. - # This sits with the other alias rewrites, and ahead of the guardrail - # merge and the pre-call hooks, so everything that keys off the model - # group -- model-level guardrails, per-model budgets and rate limits, - # the logging object -- sees the group that will actually serve. if llm_router is not None and proxy_config is not None: - from litellm.proxy.proxy_server import prisma_client - - router_settings: Final = await proxy_config._get_hierarchical_router_settings( - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - ) - - # If router_settings found (from key or team), apply them - # Pass settings as per-request overrides instead of creating a new Router - # This avoids expensive Router instantiation on each request - if router_settings is not None: - self.data["router_settings_override"] = router_settings - try: - self.data["_router_weights"] = validate_router_weights(router_settings.get("weights")) - except ValidationError: - self.data["_router_weights"] = None - verbose_proxy_logger.warning( - "Ignoring invalid saved router weights; update team/key router_settings" - ) - alias_target: Final = await _resolve_per_request_model_group_alias( + self.data.update( + await hierarchical_router_settings_updates( requested_model=self.data.get("model"), - router_settings=router_settings, user_api_key_dict=user_api_key_dict, llm_router=llm_router, + proxy_config=proxy_config, + proxy_logging_obj=proxy_logging_obj, + apply_model_group_alias=True, ) - if alias_target is not None: - self.data["model"] = alias_target + ) aliased_model: Final = self.data.get("model") if llm_router is not None and route_type in ROUTER_SCHEDULED_ROUTE_TYPES and isinstance(aliased_model, str): diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 28fb9baf2aa..6244e5025ce 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -8,8 +8,10 @@ from typing import Final, get_type_hints import orjson from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, UploadFile, status from fastapi.responses import ORJSONResponse +from pydantic import TypeAdapter import litellm +from litellm.constants import MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -17,6 +19,7 @@ from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, + hierarchical_router_settings_updates, log_llm_api_exception, resolve_litellm_call_id, ) @@ -36,6 +39,7 @@ from litellm.types.images.main import ImageEditRequestParams from litellm.types.llms.openai import ChatCompletionUserMessage router: Final = APIRouter() +_IMAGE_GENERATION_MODEL_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object) IMAGE_EDIT_NUMERIC_FORM_FIELDS: Final = numeric_form_fields(get_type_hints(ImageEditRequestParams)) IMAGE_EDIT_OPTIONAL_FIELD_DEFAULTS: Final = MappingProxyType({"prompt": None, "image": None}) @@ -132,6 +136,19 @@ async def image_generation( if data["model"] in litellm.model_alias_map: data["model"] = litellm.model_alias_map[data["model"]] + if llm_router is not None: + requested_model: Final = _IMAGE_GENERATION_MODEL_ADAPTER.validate_python(data.get("model")) + data.update( + await hierarchical_router_settings_updates( + requested_model=requested_model, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + proxy_config=proxy_config, + proxy_logging_obj=proxy_logging_obj, + apply_model_group_alias=request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is not True, + ) + ) + ### CALL HOOKS ### - modify incoming data / reject request before calling the model prompt_value: Final = data.get("prompt") if prompt_value is not None: @@ -156,6 +173,7 @@ async def image_generation( route_type="aimage_generation", llm_router=llm_router, user_model=user_model, + user_api_key_dict=user_api_key_dict, ) response = await llm_call diff --git a/tests/integration/routing/test_key_router_settings_image_generation.py b/tests/integration/routing/test_key_router_settings_image_generation.py new file mode 100644 index 00000000000..84415c72f2f --- /dev/null +++ b/tests/integration/routing/test_key_router_settings_image_generation.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import json +import uuid +from types import MappingProxyType +from typing import Final + +from integration._support.client import Gateway, Scenario, gateway_from_environment +from integration._support.wire import Reply, Request, Wire, wire_server + +FRESH_CONNECTION: Final = MappingProxyType({"Connection": "close"}) +IMAGE_RESPONSE: Final = { + "created": 1, + "background": None, + "data": [ + { + "b64_json": None, + "revised_prompt": None, + "url": "https://images.example/generated.png", + "provider_specific_fields": None, + } + ], + "output_format": None, + "quality": None, + "size": None, + "usage": { + "total_tokens": 0, + "input_tokens": 0, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "output_tokens": 0, + }, +} + + +def _primary_unavailable(_: Request) -> Reply: + return Reply( + status=503, + body=json.dumps({"error": {"message": "primary unavailable", "type": "server_error"}}).encode(), + ) + + +def _image_reply(_: Request) -> Reply: + return Reply( + body=json.dumps( + { + "created": 1, + "data": [{"url": "https://images.example/generated.png"}], + } + ).encode() + ) + + +def _image_model(scenario: Scenario, group: str, wire: Wire) -> str: + return scenario.model( + name=group, + model="openai/dall-e-3", + api_base=f"{wire.url}/v1", + api_key="integration-provider-key", + num_retries=0, + ) + + +def _assert_image_fallback( + gateway: Gateway, + key: str, + primary: Wire, + fallback: Wire, + primary_model: str, +) -> None: + response: Final = gateway.request( + "POST", + "/v1/images/generations", + {"model": primary_model, "prompt": "draw a tree"}, + key=key, + headers=FRESH_CONNECTION, + ) + + assert response.status_code == 200, response.text + assert response.json() == IMAGE_RESPONSE, response.text + assert tuple((request.method, request.target) for request in primary.drain()) == ( + ("POST", "/v1/images/generations"), + ), response.text + assert tuple((request.method, request.target) for request in fallback.drain()) == ( + ("POST", "/v1/images/generations"), + ), response.text + + +def test_key_router_settings_fallback_routes_image_generation() -> None: + with ( + wire_server(_primary_unavailable) as primary_wire, + wire_server(_image_reply) as fallback_wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire) + fallback_model: Final = _image_model(scenario, f"image-fallback-{uuid.uuid4().hex}", fallback_wire) + key: Final = scenario.key(router_settings={"fallbacks": [{primary_model: [fallback_model]}], "num_retries": 0}) + + _assert_image_fallback(gateway, key, primary_wire, fallback_wire, primary_model) + + +def test_team_router_settings_fallback_routes_image_generation() -> None: + with ( + wire_server(_primary_unavailable) as primary_wire, + wire_server(_image_reply) as fallback_wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire) + fallback_model: Final = _image_model(scenario, f"image-fallback-{uuid.uuid4().hex}", fallback_wire) + team: Final = scenario.team( + router_settings={"fallbacks": [{primary_model: [fallback_model]}], "num_retries": 0} + ) + key: Final = scenario.key(team_id=team) + + _assert_image_fallback(gateway, key, primary_wire, fallback_wire, primary_model) + + +def test_key_model_group_alias_resolves_once_for_image_generation() -> None: + with ( + wire_server(_image_reply) as primary_wire, + wire_server(_image_reply) as target_wire, + wire_server(_image_reply) as chained_target_wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire) + target_model: Final = _image_model(scenario, f"image-target-{uuid.uuid4().hex}", target_wire) + chained_target_model: Final = _image_model( + scenario, f"image-chained-target-{uuid.uuid4().hex}", chained_target_wire + ) + alias_name: Final = f"image-alias-{uuid.uuid4().hex}" + key: Final = scenario.key( + router_settings={ + "model_group_alias": { + alias_name: target_model, + target_model: chained_target_model, + } + } + ) + + response: Final = gateway.request( + "POST", + "/v1/images/generations", + {"model": alias_name, "prompt": "draw a tree"}, + key=key, + headers=FRESH_CONNECTION, + ) + primary_requests: Final = tuple((request.method, request.target) for request in primary_wire.drain()) + target_requests: Final = tuple((request.method, request.target) for request in target_wire.drain()) + chained_target_requests: Final = tuple( + (request.method, request.target) for request in chained_target_wire.drain() + ) + + assert alias_name not in {primary_model, target_model, chained_target_model}, response.text + assert response.status_code == 200, response.text + assert response.json() == IMAGE_RESPONSE, response.text + assert primary_requests == (), response.text + assert target_requests == (("POST", "/v1/images/generations"),), ( + f"{response.text}; chained target requests: {chained_target_requests}" + ) + assert chained_target_requests == (), response.text diff --git a/tests/unit/proxy/image_endpoints/test_endpoints.py b/tests/unit/proxy/image_endpoints/test_endpoints.py index f4aebecc11e..6a2c7e3ed41 100644 --- a/tests/unit/proxy/image_endpoints/test_endpoints.py +++ b/tests/unit/proxy/image_endpoints/test_endpoints.py @@ -3,7 +3,7 @@ import copy import logging from collections.abc import Iterator, Mapping from types import SimpleNamespace -from typing import Any, Dict +from typing import Any, Dict, Final, cast import orjson import pytest @@ -12,10 +12,17 @@ from fastapi.testclient import TestClient from starlette.requests import Request from starlette.responses import Response +import litellm from litellm._logging import verbose_proxy_logger +from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.image_endpoints import endpoints +from litellm.proxy.proxy_server import PrismaClient, ProxyConfig +from litellm.proxy.utils import ProxyLogging +from litellm.types.utils import ImageObject, ImageResponse +from litellm.utils import custom_llm_setup @pytest.mark.asyncio @@ -410,3 +417,142 @@ async def test_failure_before_the_provider_call_bills_the_callers_litellm_call_i assert raised.value.headers["x-litellm-call-id"] == call_id assert [data["litellm_call_id"] for data in hook_request_data] == [call_id] + + +class _ImageRouterHandler(litellm.CustomLLM): + def __init__(self) -> None: + self.calls: list[str] = [] + + async def aimage_generation( + self, + model: str, + prompt: str, + model_response: ImageResponse, + api_key: str | None, + api_base: str | None, + optional_params: Mapping[str, object], + logging_obj: object, + timeout: object | None = None, + client: object | None = None, + ) -> ImageResponse: + self.calls.append(model) + if model.endswith("failing"): + raise litellm.RateLimitError(message="primary failed", llm_provider="image-regression", model=model) + return ImageResponse(created=1, data=[ImageObject(url="https://images.example/generated.png")]) + + +def _image_generation_client( + monkeypatch: pytest.MonkeyPatch, + user_api_key_dict: UserAPIKeyAuth, + team_settings: Mapping[str, object] | None = None, +) -> tuple[TestClient, _ImageRouterHandler]: + from litellm.proxy import proxy_server + + handler: Final = _ImageRouterHandler() + cache: Final = UserApiKeyCache() + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list]) + monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers]) + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "image-regression", "custom_handler": handler}]) + custom_llm_setup() + if team_settings is not None: + team_id: Final = user_api_key_dict.team_id + assert team_id is not None + cache.set_cache( + key=f"team_id:{team_id}", + value=LiteLLM_TeamTableCachedObj(team_id=team_id, router_settings=dict(team_settings)), + local_only=True, + model_type=LiteLLM_TeamTableCachedObj, + ) + monkeypatch.setattr(proxy_server, "prisma_client", cast(PrismaClient, object())) + + router: Final = litellm.Router( + model_list=[ + {"model_name": "A", "litellm_params": {"model": "image-regression/failing"}}, + {"model_name": "B", "litellm_params": {"model": "image-regression/working"}}, + ], + num_retries=0, + ) + monkeypatch.setattr(litellm, "model_alias_map", {}) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "proxy_config", ProxyConfig()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "version", "test-version") + + app: Final = FastAPI() + app.include_router(endpoints.router) + app.dependency_overrides[user_api_key_auth] = lambda: user_api_key_dict + return TestClient(app, raise_server_exceptions=False), handler + + +def test_image_generation_uses_key_router_settings_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(router_settings={"fallbacks": [{"A": ["B"]}]}), + ) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 200, response.text + assert response.json() == ImageResponse( + created=1, + data=[ImageObject(url="https://images.example/generated.png")], + ).model_dump(), response.text + assert handler.calls == ["failing", "working"], response.text + + +def test_image_generation_uses_team_router_settings_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None: + team_id: Final = "image-regression-team" + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(team_id=team_id), + {"fallbacks": [{"A": ["B"]}]}, + ) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 200, response.text + assert response.json() == ImageResponse( + created=1, + data=[ImageObject(url="https://images.example/generated.png")], + ).model_dump(), response.text + assert handler.calls == ["failing", "working"], response.text + + +def test_image_generation_without_router_settings_does_not_use_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth()) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 500, response.text + assert handler.calls == ["failing"], response.text + + +def test_image_generation_uses_key_model_group_alias(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(router_settings={"model_group_alias": {"alias": "B"}}), + ) + + response: Final = client.post("/v1/images/generations", json={"model": "alias", "prompt": "draw a tree"}) + + assert response.status_code == 200, response.text + assert response.json() == ImageResponse( + created=1, + data=[ImageObject(url="https://images.example/generated.png")], + ).model_dump(), response.text + assert handler.calls == ["working"], response.text + + +def test_image_generation_uses_key_num_retries(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(router_settings={"num_retries": 2}), + ) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 500, response.text + assert handler.calls == ["failing", "failing", "failing"], response.text From 8f968b6c8e2ed2a7c3a1bdbc4d75dc733308e14e Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 12:51:48 -0700 Subject: [PATCH 56/64] feat(openrouter): record upstream provider in logging payload and spend logs (#49287) * test: cover OpenRouter upstream provider logging Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: scope streaming provider metadata guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(openrouter): preserve upstream provider in logs Co-authored-by: BeWayllar <124474371+BeWayllar@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): omit upstream_provider from spend metadata when unset * refactor(openrouter): set upstream provider hidden params without mutation --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: BeWayllar <124474371+BeWayllar@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 34 ++- .../litellm_core_utils/streaming_handler.py | 23 +- .../llms/openrouter/chat/transformation.py | 30 ++- litellm/proxy/_types.py | 1 + .../spend_tracking/spend_tracking_utils.py | 21 ++ litellm/types/utils.py | 1 + .../test_openrouter_upstream_provider.py | 216 ++++++++++++++++++ .../test_litellm_logging.py | 7 +- .../test_streaming_chunk_builder_utils.py | 65 ++++++ .../test_openrouter_chat_transformation.py | 117 ++++++++++ .../test_spend_tracking_utils.py | 64 +++++- 11 files changed, 564 insertions(+), 15 deletions(-) create mode 100644 tests/integration/spend/test_openrouter_upstream_provider.py diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0ddfa1210c2..4857973d51d 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -329,6 +329,20 @@ class ChunkProcessor: models[0] if models else None, ) + @staticmethod + def _get_upstream_provider(chunks: Sequence[object]) -> str | None: + def _provider_from_chunk(chunk: object) -> str | None: + hidden_params: Final = get_hidden_params(chunk) + if hidden_params is None: + return None + provider: Final[object] = hidden_params.get("upstream_provider") + return provider if isinstance(provider, str) and provider else None + + return next( + (provider for chunk in chunks if (provider := _provider_from_chunk(chunk)) is not None), + None, + ) + @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, @@ -506,10 +520,22 @@ class ChunkProcessor: chunks, first_chunk_model, ) - if provider_response_model is not None: - response._hidden_params = dict( # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter - response._hidden_params, # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params getter - provider_response_model=provider_response_model, + upstream_provider: Final = self._get_upstream_provider(chunks) + provider_hidden_params: Final[Mapping[str, str]] = MappingProxyType( + { + key: value + for key, value in ( + ("provider_response_model", provider_response_model), + ("upstream_provider", upstream_provider), + ) + if value is not None + } + ) + if provider_hidden_params: + response_hidden_params: Final[Mapping[str, builtins.object]] = get_hidden_params(response) or {} + set_hidden_params( + response, + dict(response_hidden_params, **provider_hidden_params), ) return response diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 9710956c97e..484edd869cb 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -277,6 +277,7 @@ class _VertexChunkLike(Protocol): class _ParsedChunkHiddenParams(LiteLLMBaseModel): provider_specific_fields: Mapping[str, object] | None = None + upstream_provider: str | None = None def _provider_response_model(chunk: object) -> str | None: @@ -293,9 +294,17 @@ def _parsed_provider_hidden_params(hidden: object) -> _ParsedChunkHiddenParams | return None +def _provider_upstream_provider(chunk: object) -> str | None: + parsed: Final = _parsed_provider_hidden_params(get_hidden_params(chunk)) + if parsed is None or not parsed.upstream_provider: + return None + return parsed.upstream_provider + + def _provider_hidden_params( chunk: object, provider_response_model: str | None, + upstream_provider: str | None, ) -> Mapping[str, object] | None: hidden: Final = get_hidden_params(chunk) parsed: Final = _parsed_provider_hidden_params(hidden) @@ -308,6 +317,10 @@ def _provider_hidden_params( for key, value in ( ("provider_response_model", provider_response_model), ("provider_specific_fields", provider_specific_fields), + ( + "upstream_provider", + parsed.upstream_provider if parsed is not None and parsed.upstream_provider else upstream_provider, + ), ) if value is not None } @@ -360,6 +373,7 @@ class CustomStreamWrapper: self.system_fingerprint: str | None = None self._provider_response_model: str | None = None + self._upstream_provider: str | None = None self.received_finish_reason: str | None = None self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.special_tokens = [ @@ -1653,10 +1667,17 @@ class CustomStreamWrapper: provider_response_model: Final = _provider_response_model(chunk) if provider_response_model is not None: self._provider_response_model = provider_response_model + upstream_provider: Final = _provider_upstream_provider(chunk) + if self._upstream_provider is None and upstream_provider is not None: + self._upstream_provider = upstream_provider prebuilt_delta: Final = self._prebuilt_plain_delta(chunk) plain_delta_created: bool = prebuilt_delta is not None # rebind-ok: the first chunk builds its delta below model_response = self.model_response_creator( - hidden_params=_provider_hidden_params(chunk, self._provider_response_model), + hidden_params=_provider_hidden_params( + chunk, + self._provider_response_model, + self._upstream_provider, + ), delta=prebuilt_delta, ) response_obj: dict[str, Any] = {} diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index cca6815399f..69f3d76c034 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -14,7 +14,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.core_helpers import finite_cost -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, set_hidden_params from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -212,19 +212,27 @@ class OpenrouterConfig(OpenAIGPTConfig): # OpenRouter returns cost information in the usage object when usage.include=true try: response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) + response_provider: Final = response_json.get("provider") + if isinstance(response_provider, str) and response_provider: + set_hidden_params( + model_response, + {**(get_hidden_params(model_response) or {}), "upstream_provider": response_provider}, + ) if "usage" in response_json and response_json["usage"]: response_cost: Final = _JSON_OBJECT.validate_python(response_json["usage"]).get("cost") if response_cost is not None: # Store cost in hidden params for the cost calculator to use if not hasattr(model_response, HIDDEN_PARAMS_ATTR): set_hidden_params(model_response, {}) - hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic response metadata - dict[str, object], getattr(model_response, HIDDEN_PARAMS_ATTR) + cost_hidden_params: Final = ( + cast( # cast-ok: preserve mapping operations on dynamic response metadata + dict[str, object], getattr(model_response, HIDDEN_PARAMS_ATTR) + ) ) - if "additional_headers" not in hidden_params: - hidden_params["additional_headers"] = {} + if "additional_headers" not in cost_hidden_params: + cost_hidden_params["additional_headers"] = {} additional_headers: Final = cast( # cast-ok: preserve mapping operations on response metadata - dict[str, object], hidden_params["additional_headers"] + dict[str, object], cost_hidden_params["additional_headers"] ) reported_cost: Final = ( finite_cost(float(response_cost)) if isinstance(response_cost, str | int | float) else None @@ -282,7 +290,7 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): for choice in chunk["choices"]: choice["delta"]["reasoning_content"] = choice["delta"].get("reasoning") new_choices.append(choice) - return ModelResponseStream( + model_response: Final = ModelResponseStream( id=chunk["id"], object="chat.completion.chunk", created=chunk["created"], @@ -290,6 +298,12 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): model=chunk["model"], choices=new_choices, ) + typed_chunk: Final[Mapping[str, object]] = cast( # cast-ok: stream chunks are decoded JSON objects + Mapping[str, object], chunk + ) + upstream_provider: Final[object] = typed_chunk.get("provider") + if isinstance(upstream_provider, str) and upstream_provider: + set_hidden_params(model_response, {"upstream_provider": upstream_provider}) except KeyError as e: raise OpenRouterException( message=f"KeyError: {e}, Got unexpected response from OpenRouter: {chunk}", @@ -298,3 +312,5 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): ) except Exception as e: raise e + else: + return model_response diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f165355da5a..73b12399af4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3758,6 +3758,7 @@ class SpendLogsMetadata(TypedDict): kubernetes_pod_routing: ReadOnly[StandardLoggingKubernetesPodRouting | None] internal_call_origin: InternalCallOrigin | None billed_retry_attempt: ReadOnly[NotRequired[bool | None]] + upstream_provider: ReadOnly[NotRequired[str | None]] litellm_roi_estimator: ReadOnly[NotRequired[bool | None]] is_silent_experiment: ReadOnly[NotRequired[bool | None]] guardrail_information: list[StandardLoggingGuardrailInformation] | None diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index f8d5cba4229..306e70f7029 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -11,6 +11,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, cast, runtime_checkable from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -194,6 +195,7 @@ _STAMPED_METADATA_KEYS: Final = frozenset( "autorouter_savings_estimate", "autorouter_baseline_observation", "used_client_oauth_token", + "upstream_provider", KUBERNETES_POD_ROUTING_KEY, "litellm_roi_estimator", "is_silent_experiment", @@ -203,6 +205,10 @@ _STAMPED_METADATA_KEYS: Final = frozenset( ) +class _UpstreamProviderMetadata(TypedDict, total=False): + upstream_provider: ReadOnly[str] + + def _stored_user_email(metadata: Mapping[str, object]) -> str | None: email: Final = metadata.get("user_api_key_user_email") return email if isinstance(email, str) and email.strip() else None @@ -221,6 +227,7 @@ def _get_spend_logs_metadata( model_map_information: StandardLoggingModelInformation | None = None, cold_storage_object_key: str | None = None, litellm_overhead_time_ms: float | None = None, + upstream_provider: str | None = None, cost_breakdown: CostBreakdown | None = None, litellm_call_id: str | None = None, autorouter_savings: float | None = None, @@ -232,6 +239,9 @@ def _get_spend_logs_metadata( kubernetes_pod_routing: StandardLoggingKubernetesPodRouting | None = None, is_silent_experiment: bool = False, ) -> SpendLogsMetadata: + upstream_provider_metadata: Final[_UpstreamProviderMetadata] = ( + {"upstream_provider": upstream_provider} if upstream_provider is not None else {} + ) if metadata is None: return SpendLogsMetadata( user_api_key=None, @@ -265,6 +275,7 @@ def _get_spend_logs_metadata( eval_information=None, cold_storage_object_key=cold_storage_object_key, litellm_overhead_time_ms=None, + **upstream_provider_metadata, attempted_retries=None, max_retries=None, attempted_fallbacks=None, @@ -290,6 +301,7 @@ def _get_spend_logs_metadata( **MappingProxyType( {key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key not in _STAMPED_METADATA_KEYS} ), + **upstream_provider_metadata, autorouter_savings=autorouter_savings, autorouter_savings_estimate=autorouter_savings_estimate, autorouter_baseline_observation=autorouter_baseline_observation, @@ -706,6 +718,14 @@ def get_logging_payload( if standard_logging_payload is not None: hidden_params: Final = standard_logging_payload.get("hidden_params", {}) litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms") + upstream_provider_value: Final[object] = ( + standard_logging_payload["hidden_params"]["upstream_provider"] + if standard_logging_payload is not None + and "hidden_params" in standard_logging_payload + and "upstream_provider" in standard_logging_payload["hidden_params"] + else None + ) + upstream_provider: Final = upstream_provider_value if isinstance(upstream_provider_value, str) else None logged_provider: Final = ( kwargs.get("custom_llm_provider") @@ -807,6 +827,7 @@ def get_logging_payload( else None ), litellm_overhead_time_ms=litellm_overhead_time_ms, + upstream_provider=upstream_provider, cost_breakdown=( standard_logging_payload.get("cost_breakdown", None) if standard_logging_payload is not None else None ), diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f8115571dfd..db865c7488a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3460,6 +3460,7 @@ class StandardLoggingHiddenParams(TypedDict): batch_models: list[str] | None batch_successful_requests: ReadOnly[int | None] batch_failed_requests: ReadOnly[int | None] + upstream_provider: ReadOnly[NotRequired[str | None]] litellm_model_name: str | None # the model name sent to the provider by litellm usage_object: dict | None diff --git a/tests/integration/spend/test_openrouter_upstream_provider.py b/tests/integration/spend/test_openrouter_upstream_provider.py new file mode 100644 index 00000000000..fa17a364aea --- /dev/null +++ b/tests/integration/spend/test_openrouter_upstream_provider.py @@ -0,0 +1,216 @@ +import json +import uuid +from collections.abc import Callable +from typing import Final, Literal + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.shares_worker_state + +_BACKEND: Final = "openai/gpt-4o-mini" +_PROMPT: Final = "upstream provider trace" +_PROVIDER: Final = "serving-site-a" +_ANSWER: Final = "hello" +_USAGE: Final[dict[str, JsonValue]] = { + "prompt_tokens": 5, + "completion_tokens": 1, + "total_tokens": 6, + "cost": 0.00015, +} + + +def _reply( + request_id: str, + *, + response_provider: bool, + stream_provider_mode: Literal["all", "last", "none"], +) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target.split("?", 1)[0].endswith("/chat/completions"), request.target + body: Final = JSON_OBJECT.validate_json(request.body) + assert body["messages"] == [{"role": "user", "content": _PROMPT}], body + if body.get("stream") is True: + common: Final = { + "id": request_id, + "object": "chat.completion.chunk", + "created": 1, + "model": _BACKEND, + } + all_provider: Final[dict[str, JsonValue]] = {"provider": _PROVIDER} + first_provider: Final = all_provider if stream_provider_mode == "all" else {} + finish_provider: Final = ( + all_provider if stream_provider_mode in ("all", "last") else {} + ) + usage_provider: Final = all_provider if stream_provider_mode == "all" else {} + chunks: Final[tuple[dict[str, JsonValue], ...]] = ( + { + **common, + **first_provider, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": _ANSWER}, "finish_reason": None}], + }, + { + **common, + **finish_provider, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + {**common, **usage_provider, "choices": [], "usage": _USAGE}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(chunk)}\n\n".encode() for chunk in chunks) + + (b"data: [DONE]\n\n",), + ) + + provider: Final[dict[str, JsonValue]] = {"provider": _PROVIDER} if response_provider else {} + return Reply( + body=json.dumps( + { + "id": request_id, + "object": "chat.completion", + "created": 1, + "model": _BACKEND, + **provider, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": _ANSWER}, + "finish_reason": "stop", + } + ], + "usage": _USAGE, + } + ).encode() + ) + + return respond + + +def _stream_chunks(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]: + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _send( + gateway: Gateway, + scenario: Scenario, + wire_url: str, + request_id: str, + *, + openrouter: bool, + stream: bool, +) -> httpx.Response: + model: Final = scenario.model( + model=f"{'openrouter' if openrouter else 'openai'}/{_BACKEND}", + api_base=wire_url, + api_key="synthetic-provider-key", + ) + key: Final = scenario.key(models=[model]) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "stream": stream, + **({"stream_options": {"include_usage": True}} if stream else {}), + }, + key=key, + ) + assert response.status_code == 200, response.text + returned_id: Final = ( + string_value(_stream_chunks(response)[0]["id"]) + if stream + else string_value(JSON_OBJECT.validate_json(response.content)["id"]) + ) + assert returned_id == request_id, response.text + return response + + +def _spend_metadata(request_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return object_value(rows[0]["metadata"]) + + +@pytest.mark.parametrize( + ("stream", "stream_provider_mode"), + [ + pytest.param(False, "none", id="nonstream"), + pytest.param(True, "all", id="stream-every-chunk"), + pytest.param(True, "last", id="stream-later-chunk"), + ], +) +def test_openrouter_provider_is_persisted_in_spend_metadata_regression( + gateway: Gateway, stream: bool, stream_provider_mode: Literal["all", "last", "none"] +) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=True, + stream_provider_mode=stream_provider_mode, + ) + ) as wire, + gateway.scenario() as scenario, + ): + _send(gateway, scenario, wire.url, request_id, openrouter=True, stream=stream) + metadata: Final = _spend_metadata(request_id) + + assert metadata.get("upstream_provider") == _PROVIDER, metadata + + +def test_openrouter_client_response_keeps_provider_in_body_guard(gateway: Gateway) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=True, + stream_provider_mode="none", + ) + ) as wire, + gateway.scenario() as scenario, + ): + response: Final = _send(gateway, scenario, wire.url, request_id, openrouter=True, stream=False) + + body: Final = JSON_OBJECT.validate_json(response.content) + choices: Final = body["choices"] + assert isinstance(choices, list) + message: Final = object_value(object_value(choices[0])["message"]) + assert body["provider"] == _PROVIDER + assert "upstream_provider" not in body + assert message["content"] == _ANSWER + + +def test_non_openrouter_spend_metadata_has_no_provider_guard(gateway: Gateway) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=False, + stream_provider_mode="none", + ) + ) as wire, + gateway.scenario() as scenario, + ): + _send(gateway, scenario, wire.url, request_id, openrouter=False, stream=False) + metadata: Final = _spend_metadata(request_id) + + assert "upstream_provider" not in metadata, metadata diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index d0e31cdd57a..f7c2082b683 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -12574,9 +12574,10 @@ def test_get_hidden_params(): assert result["api_base"] is None assert result["response_cost"] is None assert result["additional_headers"] is None + assert result.get("upstream_provider") is None # assert all fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + assert all(field in result for field in StandardLoggingHiddenParams.__annotations__ if field != "upstream_provider") # Test with valid params hidden_params = { @@ -12584,6 +12585,7 @@ def test_get_hidden_params(): "cache_key": "test-cache", "api_base": "https://api.test.com", "response_cost": 0.001, + "upstream_provider": "serving-site-a", "additional_headers": { "x-ratelimit-limit-requests": "2000", "x-ratelimit-remaining-requests": "1999", @@ -12596,8 +12598,9 @@ def test_get_hidden_params(): assert result["response_cost"] == 0.001 assert result["additional_headers"] is not None assert result["additional_headers"]["x_ratelimit_limit_requests"] == 2000 + assert result.get("upstream_provider") == "serving-site-a" # assert all fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + assert all(field in result for field in StandardLoggingHiddenParams.__annotations__ if field != "upstream_provider") @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 9df45b92686..bbf1897b3a7 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -910,6 +910,71 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" +def test_stream_chunk_builder_finds_upstream_provider_in_later_chunk_regression() -> None: + first: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + later: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" world", role=None), + ) + ], + ) + first_hidden_params: Final = {"origin": "first"} + later_hidden_params: Final = {"upstream_provider": "serving-site-a"} + chunks: Final = [ + {**first.model_dump(), "_hidden_params": first_hidden_params}, + {**later.model_dump(), "_hidden_params": later_hidden_params}, + ] + + response: Final = stream_chunk_builder(chunks=chunks) + + assert response is not None + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + assert first_hidden_params.get("upstream_provider") is None + assert later_hidden_params.get("upstream_provider") == "serving-site-a" + + +def test_stream_chunk_builder_preserves_source_provider_metadata_guard() -> None: + chunk: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + hidden_params: Final = {"upstream_provider": "serving-site-a"} + chunk_dict: Final = {**chunk.model_dump(), "_hidden_params": hidden_params} + + response: Final = stream_chunk_builder(chunks=[chunk_dict]) + + assert response is not None + assert chunk_dict["_hidden_params"]["upstream_provider"] == "serving-site-a" + + def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): """Vertex AI metadata on streaming chunks must appear on assembled response.""" grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] diff --git a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py index 888c47505a4..3d659641f54 100644 --- a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py +++ b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py @@ -872,3 +872,120 @@ def test_openrouter_streaming_reasoning_is_preserved() -> None: assembled: Final = litellm.stream_chunk_builder(chunks=list(stream_chunks), messages=messages) assert assembled is not None assert assembled.choices[0].message.reasoning_content == "Consider the evidence. Apply the rule." + + +def _transform_openrouter_response(body: dict[str, object]) -> litellm.ModelResponse: + return OpenrouterConfig().transform_response( + model="openrouter/openai/gpt-4o-mini", + raw_response=httpx.Response(200, json=body), + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_openrouter_nonstream_provider_without_usage_is_preserved_regression() -> None: + response: Final = _transform_openrouter_response( + { + "id": "gen-provider-no-usage", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-4o-mini", + "provider": "serving-site-a", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + } + ) + + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + + +@pytest.mark.parametrize( + ("provider_fields", "case"), + [ + pytest.param({}, "missing", id="missing"), + pytest.param({"provider": None}, "null", id="null"), + pytest.param({"provider": ""}, "empty", id="empty"), + pytest.param({"provider": 123}, "non-string", id="non-string"), + ], +) +def test_openrouter_invalid_provider_keeps_reported_cost_guard( + provider_fields: dict[str, object], case: str +) -> None: + response: Final = _transform_openrouter_response( + { + "id": f"gen-invalid-provider-{case}", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-4o-mini", + **provider_fields, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30, + "cost": 0.00015, + }, + } + ) + + assert response._hidden_params.get("upstream_provider") is None + assert response._hidden_params["additional_headers"][OPENROUTER_COST_HEADER] == 0.00015 + + +def test_openrouter_streaming_chunk_preserves_provider_regression() -> None: + chunk: Final = { + "id": "gen-provider-stream", + "object": "chat.completion.chunk", + "created": 1, + "model": "openai/gpt-4o-mini", + "provider": "serving-site-a", + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + } + + response: Final = OpenRouterChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ).chunk_parser(chunk) + + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + + +@pytest.mark.parametrize( + ("provider_fields", "case"), + [ + pytest.param({}, "missing", id="missing"), + pytest.param({"provider": None}, "null", id="null"), + pytest.param({"provider": ""}, "empty", id="empty"), + pytest.param({"provider": 123}, "non-string", id="non-string"), + ], +) +def test_openrouter_streaming_chunk_ignores_invalid_provider_guard(provider_fields: dict[str, object], case: str) -> None: + chunk: Final = { + "id": f"gen-invalid-provider-{case}", + "object": "chat.completion.chunk", + "created": 1, + "model": "openai/gpt-4o-mini", + **provider_fields, + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + } + + response: Final = OpenRouterChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ).chunk_parser(chunk) + + assert response._hidden_params.get("upstream_provider") is None diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 1ea445fda0d..a91cdd512fe 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2394,6 +2394,65 @@ def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): ) +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_carries_upstream_provider_from_hidden_params_regression() -> None: + standard_logging_payload: Final = _make_standard_logging_payload_with_usage_object(usage_object={}) + payload_with_provider: Final = { + **dict(standard_logging_payload), + "hidden_params": { + **dict(standard_logging_payload["hidden_params"]), + "upstream_provider": "serving-site-a", + }, + } + payload: Final = get_logging_payload( + kwargs={ + "model": "gpt-5.6", + "call_type": "responses", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "standard_logging_object": payload_with_provider, + }, + response_obj={ + "id": "resp-test", + "usage": { + "input_tokens": 1000, + "output_tokens": 10, + "total_tokens": 1010, + }, + }, + start_time=datetime.datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1, tzinfo=timezone.utc), + ) + + assert json.loads(payload["metadata"]).get("upstream_provider") == "serving-site-a" + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_omits_upstream_provider_when_hidden_param_absent_guard() -> None: + standard_logging_payload: Final = _make_standard_logging_payload_with_usage_object(usage_object={}) + payload: Final = get_logging_payload( + kwargs={ + "model": "gpt-5.6", + "call_type": "responses", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "standard_logging_object": standard_logging_payload, + }, + response_obj={ + "id": "resp-test", + "usage": { + "input_tokens": 1000, + "output_tokens": 10, + "total_tokens": 1010, + }, + }, + start_time=datetime.datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1, tzinfo=timezone.utc), + ) + + assert "upstream_provider" not in json.loads(payload["metadata"]) + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_strips_null_bytes_from_request_tags(): @@ -6911,7 +6970,10 @@ def test_spend_logs_payload(model_id: str | None): ) assert len(payload["request_id"]) > 0 - expected_metadata_keys: Final = SpendLogsMetadata.__annotations__.keys() - {"is_silent_experiment"} + expected_metadata_keys: Final = SpendLogsMetadata.__annotations__.keys() - { + "is_silent_experiment", + "upstream_provider", + } assert "metadata" in payload assert isinstance(payload["metadata"], str) From f4d1f45d87dcba757edd1feae71300146fbb1570 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 30 Oct 2026 13:05:12 -0700 Subject: [PATCH 57/64] fix(proxy): raise live budget counters for failure spend rows (#47741) * fix(proxy): raise live budget counters for failure spend rows A stream that breaks mid-flight writes its recovered cost to the spend log, but the failure hook had already released the reservation and never touched the live counters, so an over-budget key kept passing until the counter was next repaired from the database Every charged failure row with a cost now raises the same counters the success path raises. The error waits at most a second for them, so a hung Redis cannot hold it back Fixes #35568 * test(proxy): cover failure spend counters across budgets, endpoints and outages Integration cells for team, user, end-user, tag and window budgets, concurrent streams, billed retries, native /v1/messages, Bedrock guardrail cost, a Redis stall on two workers and disable_spend_updates * test(proxy): check stalled failure counters and use the real spend writer - The Redis stall cell now reads the burst key's live counter and requires at least the full failure spend - Unit tests drop the charging writer double, bare dicts, list mutation and wall-clock timestamps - Partial-stream test models are frozen * test(proxy): move the stalled counter bound reason into the assert message --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../proxy/hooks/proxy_track_cost_callback.py | 124 ++++++- litellm/proxy/proxy_server.py | 2 +- tests/integration/spend/_partial_stream.py | 173 +++++++++ .../test_failure_spend_live_counters_chaos.py | 239 ++++++++++++ .../spend/test_partial_stream_live_budget.py | 350 ++++++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 225 +++++++++++ 6 files changed, 1103 insertions(+), 10 deletions(-) create mode 100644 tests/integration/spend/_partial_stream.py create mode 100644 tests/integration/spend/test_failure_spend_live_counters_chaos.py create mode 100644 tests/integration/spend/test_partial_stream_live_budget.py diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0d333d4fd08..2951bc02d8c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -12,6 +12,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final +from litellm.caching.redis_batch import flush_post_call_redis_batches from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( @@ -134,6 +135,8 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( ) ) +FAILURE_SPEND_COUNTER_WAIT_SECONDS: Final = 1.0 + _OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) @@ -235,6 +238,12 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter: return proxy_logging_obj.db_spend_update_writer +def _proxy_spend_counters() -> "_IncrementSpendCounters": + from litellm.proxy.proxy_server import increment_spend_counters + + return increment_spend_counters + + DeploymentLookup: TypeAlias = Callable[[str], Deployment | None] @@ -361,11 +370,16 @@ class _ProxyDBLogger(CustomLogger): message_logging: bool = True, spend_writer: Callable[[], DBSpendUpdateWriter] = _proxy_spend_writer, deployment_lookup: DeploymentLookup = _proxy_deployment, + spend_counters: Callable[[], "_IncrementSpendCounters"] = _proxy_spend_counters, + failure_counter_wait_seconds: float = FAILURE_SPEND_COUNTER_WAIT_SECONDS, ) -> None: super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging) self.spend_event_producer = spend_event_producer self._spend_writer: Final = spend_writer self._deployment_lookup: Final = deployment_lookup + self._spend_counters: Final = spend_counters + self._failure_counter_wait_seconds: Final = failure_counter_wait_seconds + self._pending_increments: Final[set[asyncio.Task[None]]] = set() # mutable-ok: tasks leave when done async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime @@ -590,9 +604,11 @@ class _ProxyDBLogger(CustomLogger): completed_cost: bool = False, request_tags: list[str] | None = None, ) -> None: - """One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair.""" - for row_cost, row_kwargs in spend_rows: - if completed_cost: + """One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair, then the + live spend counters raised by every charged row that cost something. A completed cost settles through the + success path's write, which raises the counters against the kept reservation.""" + if completed_cost: + for row_cost, row_kwargs in spend_rows: from litellm.proxy.proxy_server import increment_spend_counters, proxy_logging_obj await _update_database_and_spend_counters( @@ -623,12 +639,16 @@ class _ProxyDBLogger(CustomLogger): router=get_llm_router(), ), ) - continue + return + + token: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) + end_user_id: Final = None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id + charged: Final = [ await self._spend_writer().update_database( - token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), + token=token, response_cost=row_cost, user_id=user_api_key_dict.user_id, - end_user_id=None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id, + end_user_id=end_user_id, team_id=user_api_key_dict.team_id, kwargs=dict(row_kwargs), completion_response=original_exception, @@ -637,6 +657,91 @@ class _ProxyDBLogger(CustomLogger): org_id=user_api_key_dict.org_id, project_id=user_api_key_dict.project_id, ) + for row_cost, row_kwargs in spend_rows + ] + counted_rows: Final = tuple( + row for row, row_charged in zip(spend_rows, charged, strict=True) if row_charged and row[0] > 0 + ) + if counted_rows and not ProxyUpdateSpend.disable_spend_updates(): + await self._increment_failure_spend_counters( + user_api_key_dict, token, end_user_id, counted_rows, start_time + ) + + async def _increment_failure_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + counted_rows: Sequence[tuple[float, Mapping[str, object]]], + start_time: datetime, + ) -> None: + """The counters the success path raises for the same spend, sent before the client sees the error. A + counter store slower than ``failure_counter_wait_seconds`` finishes them in the background instead of + holding the error back. The reservation was already released, so every counter takes the full row cost.""" + increment: Final = asyncio.create_task( + self._raise_failure_spend_counters(user_api_key_dict, token, end_user_id, counted_rows, start_time) + ) + self._pending_increments.add(increment) + increment.add_done_callback(self._pending_increments.discard) + finished, _ = await asyncio.wait((increment,), timeout=self._failure_counter_wait_seconds) + if not finished: + verbose_proxy_logger.warning( + "Spend counters for failure spend rows took over %ss, finishing them in the background", + self._failure_counter_wait_seconds, + ) + + async def _raise_failure_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + counted_rows: Sequence[tuple[float, Mapping[str, object]]], + start_time: datetime, + ) -> None: + for row_cost, row_kwargs in counted_rows: + await self._raise_row_spend_counters( + user_api_key_dict, token, end_user_id, row_cost, row_kwargs, start_time + ) + await flush_post_call_redis_batches() + + async def _raise_row_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + row_cost: float, + row_kwargs: Mapping[str, object], + start_time: datetime, + ) -> None: + standard_logging_object: Final = row_kwargs.get("standard_logging_object") + sl_object: Final = ( + cast("StandardLoggingPayload", standard_logging_object) # cast-ok: the failure snapshot dict + if isinstance(standard_logging_object, dict) + else None + ) + metadata: Final = _str_keyed_fields(_str_keyed_fields(row_kwargs.get("litellm_params")).get("metadata")) + served_model_id: Final = _str_keyed_fields(metadata.get("model_info")).get("id") or _str_keyed_fields( + standard_logging_object + ).get("model_id") + try: + await self._spend_counters()( + token=token, + team_id=user_api_key_dict.team_id, + user_id=user_api_key_dict.user_id, + response_cost=row_cost, + org_id=user_api_key_dict.org_id, + end_user_id=end_user_id, + tags=_get_request_tags_for_cost_tracking(sl_object=sl_object, metadata=dict(metadata)), + request_started_at=start_time, + model_access_groups=debitable_model_access_groups( + attributed=get_request_model_access_groups(row_kwargs), + served_model_id=served_model_id if isinstance(served_model_id, str) else None, + router=get_llm_router(), + ), + project_id=user_api_key_dict.project_id, + ) + except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing + verbose_proxy_logger.exception("Failed to raise spend counters for a failure spend row") async def _PROXY_track_cost_callback( self, @@ -1060,7 +1165,7 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]: return tuple(sorted(str(key) for key in metadata)) -def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: +def _get_budget_reservation_from_metadata(metadata: dict) -> dict[str, object] | None: return budget_reservation_from_metadata(metadata) @@ -1095,6 +1200,7 @@ class _IncrementSpendCounters(Protocol): tags: list[str] | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, ) -> None: ... @@ -1111,7 +1217,7 @@ async def _update_database_and_spend_counters( start_time: datetime | None, end_time: datetime | None, response_cost: float, - budget_reservation: dict | None, + budget_reservation: dict[str, object] | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, @@ -1174,7 +1280,7 @@ async def _update_database_and_spend_counters_in_batch( start_time: datetime | None, end_time: datetime | None, response_cost: float, - budget_reservation: dict | None, + budget_reservation: dict[str, object] | None, request_tags: list[str] | None, model_access_groups: Sequence[str] | None, project_id: str | None, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7ea1c246da5..ea6b51b228b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3510,7 +3510,7 @@ async def increment_spend_counters( user_id: str | None, response_cost: float | None, org_id: str | None = None, - budget_reservation: dict | None = None, + budget_reservation: dict[str, object] | None = None, end_user_id: str | None = None, tags: list[str] | None = None, request_started_at: datetime | None = None, diff --git a/tests/integration/spend/_partial_stream.py b/tests/integration/spend/_partial_stream.py new file mode 100644 index 00000000000..e35320a8d8e --- /dev/null +++ b/tests/integration/spend/_partial_stream.py @@ -0,0 +1,173 @@ +import json +import threading +import uuid +from collections.abc import Callable, Iterable, Iterator +from hashlib import sha256 +from typing import Final + +import httpx +from integration._support.client import Gateway, Scenario, eventually, held +from integration._support.wire import Reply, Request, Wire +from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +BACKEND: Final = "integration-partial-stream" +PIECE: Final = "partial stream billed tokens " +DELIVERED_PIECES: Final = 12 +_FAILURE_ROWS: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND status=%s ORDER BY "startTime"' +ROWS: Final = 'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"' + + +class _SpendRow(BaseModel): + model_config = ConfigDict(frozen=True) + + spend: float + + +class StatusRow(BaseModel): + model_config = ConfigDict(frozen=True) + + status: str + spend: float + + +class _BudgetError(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str + + +class _ErrorResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + error: _BudgetError + + +def _chunk(stream_id: str, delta: dict[str, JsonValue], finish_reason: str | None = None) -> bytes: + body: Final = { + "id": stream_id, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": BACKEND, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(body)}\n\n".encode() + + +def completion() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1700000000, + "model": BACKEND, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ).encode() + + +def peer(delivered: int | None, gate: threading.Event | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if JSON_OBJECT.validate_json(request.body).get("stream") is True: + stream_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + return Reply( + content_type="text/event-stream", + chunks=( + *(_chunk(stream_id, {"role": "assistant", "content": PIECE}) for _ in range(DELIVERED_PIECES)), + _chunk(stream_id, {}, "stop"), + b"data: [DONE]\n\n", + ), + abort_after=delivered, + pause_between_chunks=0.02, + gate_before_first=gate, + ) + return Reply(body=completion()) + + return respond + + +def priced_model(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"openai/{BACKEND}", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + + +def stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 5, + "stream": True, + "messages": [{"role": "user", "content": "partial stream"}], + **fields, + }, + key=key, + ) + + +def broken_stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> None: + response: Final = stream(gateway, model, key, **fields) + assert response.status_code == 200, response.text + events: Final = tuple( + line.removeprefix("data: ") + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + assert "error" in JSON_OBJECT.validate_json(events[-1]), response.text + delivered: Final = "".join(_streamed_text(ChatCompletionChunk.model_validate_json(event) for event in events[:-1])) + assert delivered == PIECE * DELIVERED_PIECES, response.text + + +def _streamed_text(chunks: Iterable[ChatCompletionChunk]) -> Iterator[str]: + for chunk in chunks: + yield from (choice.delta.content or "" for choice in chunk.choices) + + +def chat(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 5, + "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}], + **fields, + }, + key=key, + ) + + +def failure_spend(gateway: Gateway, key: str, rows: int = 1) -> tuple[float, ...]: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda found: len(found) >= rows, seconds=70) + found: Final = held( + lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda latest: len(latest) == rows, holding=2 + ) + return tuple(_SpendRow.model_validate(row).spend for row in found) + + +def chats_sent(wire: Wire) -> int: + return sum(1 for request in wire.drain() if request.method == "POST") + + +def recovered_cost(gateway: Gateway, scenario: Scenario, model: str) -> float: + key: Final = scenario.key(models=[model], max_budget=100) + broken_stream(gateway, model, key) + (spend,) = failure_spend(gateway, key) + assert spend > 0 + return spend + + +def assert_budget_refusal(response: httpx.Response) -> None: + assert response.status_code == 422, response.text + assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text diff --git a/tests/integration/spend/test_failure_spend_live_counters_chaos.py b/tests/integration/spend/test_failure_spend_live_counters_chaos.py new file mode 100644 index 00000000000..ecdd5aadb8e --- /dev/null +++ b/tests/integration/spend/test_failure_spend_live_counters_chaos.py @@ -0,0 +1,239 @@ +import json +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.process import owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.timing import assert_not_delayed_by, timed +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.spend._partial_stream import ( + DELIVERED_PIECES, + assert_budget_refusal, + broken_stream, + chat, + chats_sent, + failure_spend, + peer, + priced_model, + recovered_cost, +) +from pydantic import JsonValue +from redis import Redis + +from litellm.constants import REDIS_SOCKET_TIMEOUT + +_BURST: Final = 12 +_SPREAD: Final = 6 +_ATTACK_MARKER: Final = "synthetic-attack-marker" +_POLICY_ID: Final = "synthetic-priced-policy" +_POLICY_APPLY_PATH: Final = f"/guardrail/{_POLICY_ID}/version/DRAFT/apply" + + +def _proxy_config(directory: Path, name: str, sections: dict[str, JsonValue]) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = directory / f"{name}.yaml" + path.write_text(yaml.safe_dump({**base, **sections})) + return path + + +def _redis_counter_config(directory: Path) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + return _proxy_config( + directory, + "redis-counters", + { + "router_settings": { + **base["router_settings"], + "redis_host": "os.environ/REDIS_HOST", + "redis_port": "os.environ/REDIS_PORT", + } + }, + ) + + +def _spread_requests(gateway: Gateway, model: str, key: str) -> tuple[httpx.Response, ...]: + return tuple( + gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "max_tokens": 5, "messages": [{"role": "user", "content": f"spread {uuid.uuid4().hex}"}]}, + key=key, + headers={"connection": "close"}, + ) + for _ in range(_SPREAD) + ) + + +@pytest.mark.timeout(600) +def test_broken_streams_through_a_redis_stall_log_once_and_the_budget_holds_on_every_worker( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + with ( + owned_redis(tmp_path) as counters, + wire_server(peer(DELIVERED_PIECES)) as wire, + wire_server(peer(DELIVERED_PIECES, gate=release)) as held_wire, + owned_proxy_process( + gateway, + tmp_path, + {"REDIS_HOST": counters.host, "REDIS_PORT": str(counters.port)}, + config=_redis_counter_config(tmp_path), + workers=2, + ) as owned, + owned.gateway.scenario() as scenario, + ): + scenario.cleanups.callback(release.set) + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + stalled_model: Final = priced_model(scenario, held_wire) + cost: Final = recovered_cost(candidate, scenario, model) + control: Final = timed(lambda: broken_stream(candidate, model, scenario.key(models=[model]))) + burst_key: Final = scenario.key(models=[stalled_model]) + chats_sent(held_wire) + + with ThreadPoolExecutor(max_workers=_BURST) as pool: + pending: Final = tuple( + pool.submit(broken_stream, candidate, stalled_model, burst_key) for _ in range(_BURST) + ) + eventually(held_wire.received.qsize, lambda held: held >= _BURST, seconds=30) + with counters.paused(): + release.set() + drained: Final = timed(lambda: tuple(future.result() for future in pending)) + + assert_not_delayed_by((drained.seconds,), control=(control.seconds,), injected_delay=REDIS_SOCKET_TIMEOUT) + assert failure_spend(candidate, burst_key, rows=_BURST) == pytest.approx((cost,) * _BURST) + spent: Final = cost * _BURST + with Redis(host=counters.host, port=counters.port) as live_counters: + landed: Final = eventually( + lambda: float(live_counters.get(f"spend:key:{sha256(burst_key.encode()).hexdigest()}") or 0), + lambda counted: counted >= spent - 1e-9, + seconds=30, + return_last_on_timeout=True, + ) + assert spent - 1e-9 <= landed <= 2 * spent + 1e-9, ( + f"one broken stream costs {cost}. The cold counter is seeded from the database, and a spend flush" + " during the stall can land these rows there first, so up to twice their cost is expected, as for" + " completed streams held past a flush" + ) + + budgeted: Final = scenario.key(models=[model], max_budget=cost / 2) + broken_stream(candidate, model, budgeted) + before_kill: Final = _spread_requests(candidate, model, budgeted) + owned.kill_worker(min(owned.worker_pids())) + after_kill: Final = _spread_requests(candidate, model, budgeted) + + for refused in (*before_kill, *after_kill): + assert_budget_refusal(refused) + + +def _policy(request: Request) -> Reply: + if request.method != "POST" or request.target != _POLICY_APPLY_PATH: + return Reply(status=404) + if _ATTACK_MARKER.encode() not in request.body: + return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "blocked by synthetic policy"}], + "assessments": [ + {"contentPolicy": {"filters": [{"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"}]}} + ], + "usage": {"contentPolicyUnits": 1}, + } + ).encode() + ) + + +def _policy_config(directory: Path, policy: Wire) -> Path: + return _proxy_config( + directory, + "priced-policy", + { + "guardrails": [ + { + "guardrail_name": "priced-bedrock-policy", + "litellm_params": { + "guardrail": "bedrock", + "mode": "pre_call", + "default_on": True, + "guardrailIdentifier": _POLICY_ID, + "guardrailVersion": "DRAFT", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIASYNTHETICGUARDRAIL", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_bedrock_runtime_endpoint": policy.url, + }, + } + ] + }, + ) + + +def _attack(gateway: Gateway, model: str, key: str) -> httpx.Response: + blocked: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"{_ATTACK_MARKER} {uuid.uuid4().hex}"}]}, + key=key, + ) + assert blocked.status_code == 400, blocked.text + return blocked + + +@pytest.mark.timeout(480) +def test_guardrail_cost_on_blocked_requests_counts_toward_the_budget(gateway: Gateway, tmp_path: Path) -> None: + with ( + wire_server(_policy) as policy, + wire_server(peer(None)) as wire, + owned_proxy_process(gateway, tmp_path, {}, config=_policy_config(tmp_path, policy), workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + meter: Final = scenario.key(models=[model], max_budget=100) + _attack(candidate, model, meter) + (cost,) = failure_spend(candidate, meter) + assert cost > 0 + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + chats_sent(wire) + + _attack(candidate, model, key) + _attack(candidate, model, key) + after: Final = chat(candidate, model, key) + + assert_budget_refusal(after) + assert chats_sent(wire) == 0 + + +@pytest.mark.timeout(480) +def test_broken_streams_leave_the_budget_alone_while_spend_updates_are_disabled( + gateway: Gateway, tmp_path: Path +) -> None: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = _proxy_config( + tmp_path, "spend-updates-off", {"general_settings": {**base["general_settings"], "disable_spend_updates": True}} + ) + with ( + wire_server(peer(DELIVERED_PIECES)) as wire, + gateway.scenario() as metering, + owned_proxy_process(gateway, tmp_path, {}, config=config) as owned, + owned.gateway.scenario() as scenario, + ): + cost: Final = recovered_cost(gateway, metering, priced_model(metering, wire)) + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + key: Final = scenario.key(models=[model], max_budget=cost / 2) + + broken_stream(candidate, model, key) + after: Final = chat(candidate, model, key) + + assert after.status_code == 200, after.text diff --git a/tests/integration/spend/test_partial_stream_live_budget.py b/tests/integration/spend/test_partial_stream_live_budget.py new file mode 100644 index 00000000000..89fb66c9bbf --- /dev/null +++ b/tests/integration/spend/test_partial_stream_live_budget.py @@ -0,0 +1,350 @@ +import json +import threading +import uuid +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, held +from integration._support.wire import Reply, Request, wire_server +from integration.spend._partial_stream import ( + BACKEND, + DELIVERED_PIECES, + JSON_OBJECT, + PIECE, + ROWS, + StatusRow, + assert_budget_refusal, + broken_stream, + chat, + chats_sent, + completion, + failure_spend, + peer, + priced_model, + recovered_cost, + stream, +) +from pydantic import JsonValue + +_CONCURRENT_STREAMS: Final = 6 +_MESSAGES_BACKEND: Final = "claude-integration-partial-stream" +_UNCACHED: Final[dict[str, JsonValue]] = {"cache": {"no-cache": True}} + + +@pytest.mark.timeout(240) +def test_a_broken_stream_that_spends_past_the_budget_blocks_the_next_request(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost / 2) + chats_sent(wire) + + broken_stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert failure_spend(gateway, key, rows=2) == pytest.approx((cost, 0.0)) + assert chats_sent(wire) == 1 + + +@pytest.mark.timeout(240) +def test_broken_stream_spend_counts_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + + broken_stream(gateway, model, key) + broken_stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert failure_spend(gateway, key, rows=3) == pytest.approx((cost, cost, 0.0)) + + +@pytest.mark.timeout(240) +def test_a_stream_that_breaks_before_any_chunk_leaves_the_budget_untouched(gateway: Gateway) -> None: + with wire_server(peer(0)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + key: Final = scenario.key(models=[model], max_budget=0.05) + + refused: Final = stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert refused.status_code == 500, refused.text + + assert after.status_code == 200, after.text + assert failure_spend(gateway, key) == (0.0,) + + +def _broken_stream_then_refusal( + gateway: Gateway, model: str, spender: str, other: str, **fields: JsonValue +) -> httpx.Response: + broken_stream(gateway, model, spender, **fields) + return chat(gateway, model, other, **fields) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_on_one_team_key_blocks_the_teams_other_keys(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + team: Final = scenario.team(models=[model], max_budget=cost / 2) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model]) + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_on_one_of_a_users_keys_blocks_their_other_keys(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + user: Final = scenario.user(max_budget=cost / 2) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model]) + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_for_an_end_user_blocks_their_next_request_on_any_key(gateway: Gateway) -> None: + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + gateway.post("/customer/new", {"user_id": end_user, "max_budget": cost / 2}) + scenario.cleanups.callback(gateway.post, "/customer/delete", {"user_ids": [end_user]}) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(models=[model]), scenario.key(models=[model]), user=end_user + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_under_a_tag_blocks_the_next_request_with_that_tag(gateway: Gateway) -> None: + tag: Final = f"integration-partial-stream-{uuid.uuid4().hex}" + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + gateway.post("/tag/new", {"name": tag, "max_budget": cost / 2}) + scenario.cleanups.callback(gateway.post, "/tag/delete", {"name": tag}) + key: Final = scenario.key(models=[model]) + + after: Final = _broken_stream_then_refusal(gateway, model, key, key, metadata={"tags": [tag]}) + untagged: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert untagged.status_code == 200, untagged.text + + +@pytest.mark.timeout(240) +def test_a_broken_stream_fills_the_keys_budget_window(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], budget_limits=[{"budget_duration": "1d", "max_budget": cost / 2}]) + + after: Final = _broken_stream_then_refusal(gateway, model, key, key) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_concurrent_broken_streams_each_count_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost * (_CONCURRENT_STREAMS + 0.5)) + + with ThreadPoolExecutor(max_workers=_CONCURRENT_STREAMS) as pool: + tuple(pool.map(lambda _: broken_stream(gateway, model, key), range(_CONCURRENT_STREAMS))) + within_budget: Final = chat(gateway, model, key) + broken_stream(gateway, model, key) + over_budget: Final = chat(gateway, model, key) + + assert within_budget.status_code == 200, within_budget.text + assert_budget_refusal(over_budget) + assert sorted(failure_spend(gateway, key, rows=_CONCURRENT_STREAMS + 2)) == pytest.approx( + [0.0] + [cost] * (_CONCURRENT_STREAMS + 1) + ) + + +@pytest.mark.timeout(240) +def test_completed_streams_count_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(None)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + meter: Final = scenario.key(models=[model], max_budget=100) + assert stream(gateway, model, meter, **_UNCACHED).status_code == 200 + (metered,) = eventually( + lambda: gateway.rows(ROWS, (sha256(meter.encode()).hexdigest(),)), + lambda found: len(found) == 1, + seconds=70, + ) + cost: Final = StatusRow.model_validate(metered).spend + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + + first: Final = stream(gateway, model, key, **_UNCACHED) + second: Final = stream(gateway, model, key, **_UNCACHED) + assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text) + assert _rows(gateway, key, 2) == (StatusRow(status="success", spend=cost),) * 2 + third: Final = stream(gateway, model, key, **_UNCACHED) + + assert_budget_refusal(third) + + +def _hung_then_served(release: threading.Event, hung_model: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if JSON_OBJECT.validate_json(request.body).get("model") == hung_model: + return Reply(chunks=(completion(),), gate_before_first=release) + return Reply(body=completion()) + + return respond + + +def _timed_out_spend(gateway: Gateway, model: str, key: str) -> float: + timed_out: Final = chat(gateway, model, key, messages=[{"role": "user", "content": "hung"}]) + assert timed_out.status_code in (408, 500), timed_out.text + digest: Final = sha256(key.encode()).hexdigest() + landed: Final = eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= 3, seconds=70) + rows: Final = tuple( + StatusRow.model_validate(row) + for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == len(landed), holding=2) + ) + assert {row.status for row in rows} == {"failure"}, rows + return sum(row.spend for row in rows) + + +@pytest.mark.timeout(240) +def test_retries_the_provider_billed_before_a_timeout_count_toward_the_budget(gateway: Gateway) -> None: + hung_backend: Final = f"{BACKEND}-hung" + release: Final = threading.Event() + with wire_server(_hung_then_served(release, hung_backend)) as wire, gateway.scenario() as scenario: + scenario.cleanups.callback(release.set) + hung: Final = scenario.model( + model=f"openai/{hung_backend}", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + timeout=1, + num_retries=2, + ) + served: Final = priced_model(scenario, wire) + billed: Final = _timed_out_spend(gateway, hung, scenario.key(models=[hung, served], max_budget=100)) + assert billed > 0 + key: Final = scenario.key(models=[hung, served], max_budget=billed / 2) + + assert _timed_out_spend(gateway, hung, key) == pytest.approx(billed) + after: Final = chat(gateway, served, key) + + assert_budget_refusal(after) + + +def _messages_event(kind: str, fields: dict[str, JsonValue]) -> bytes: + return f"event: {kind}\ndata: {json.dumps({'type': kind, **fields})}\n\n".encode() + + +def _messages_peer(request: Request) -> Reply: + if request.method != "POST" or request.target != "/v1/messages": + return Reply(status=404) + message: Final[dict[str, JsonValue]] = { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": _MESSAGES_BACKEND, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + } + if JSON_OBJECT.validate_json(request.body).get("stream") is not True: + return Reply( + body=json.dumps( + {**message, "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn"} + ).encode() + ) + return Reply( + content_type="text/event-stream", + chunks=( + _messages_event("message_start", {"message": {**message, "content": []}}), + _messages_event("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}), + *( + _messages_event("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": PIECE}}) + for _ in range(DELIVERED_PIECES) + ), + _messages_event("content_block_stop", {"index": 0}), + _messages_event( + "message_delta", + {"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 60}}, + ), + _messages_event("message_stop", {}), + ), + abort_after=DELIVERED_PIECES + 2, + pause_between_chunks=0.02, + ) + + +def _messages(gateway: Gateway, model: str, key: str, *, streamed: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 5, + "stream": streamed, + "messages": [{"role": "user", "content": f"partial stream {uuid.uuid4().hex}"}], + }, + key=key, + ) + + +def _broken_messages_stream(gateway: Gateway, model: str, key: str) -> None: + response: Final = _messages(gateway, model, key, streamed=True) + assert response.status_code == 200, response.text + assert response.text.count(PIECE) == DELIVERED_PIECES, response.text + assert "message_stop" not in response.text, response.text + + +def _rows(gateway: Gateway, key: str, count: int) -> tuple[StatusRow, ...]: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= count, seconds=70) + return tuple( + StatusRow.model_validate(row) + for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == count, holding=2) + ) + + +@pytest.mark.timeout(240) +def test_a_broken_messages_stream_counts_toward_the_budget(gateway: Gateway) -> None: + with wire_server(_messages_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_MESSAGES_BACKEND}", + api_base=wire.url, + api_key="synthetic-anthropic-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + meter: Final = scenario.key(models=[model], max_budget=100) + _broken_messages_stream(gateway, model, meter) + (metered,) = _rows(gateway, meter, 1) + assert metered.spend > 0, metered + key: Final = scenario.key(models=[model], max_budget=metered.spend / 2) + + _broken_messages_stream(gateway, model, key) + after: Final = _messages(gateway, model, key, streamed=False) + + assert_budget_refusal(after) + assert _rows(gateway, key, 2) == (metered, StatusRow(status="failure", spend=0.0)) diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 8d4f8cd87ad..e482b6552a4 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -14,6 +14,7 @@ from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import ProxyException, SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer @@ -36,6 +37,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payloa from litellm.proxy.utils import ( REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, + BilledFailedAttempt, ProxyUpdateSpend, ) from litellm.types.router import Deployment, LiteLLM_Params @@ -2676,6 +2678,229 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): assert mock_update_database.call_args[1]["response_cost"] == 3.5e-05 +@pytest.fixture() +def live_spend_counters(monkeypatch: pytest.MonkeyPatch) -> DualCache: + counters: Final = DualCache() + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", counters) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + return counters + + +def _failure_with_billed_attempts() -> Mapping[str, object]: + usage: Final = Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) + return { + "combined_usage_object": usage, + "response_cost": 0.5, + "billed_failed_attempts": ( + BilledFailedAttempt(attempted_retries=0, deployment_id="dep-a", usage=usage, response_cost=0.25), + BilledFailedAttempt(attempted_retries=1, deployment_id="dep-b", usage=usage, response_cost=0.5), + ), + } + + +_LIVE_COUNTERS: Final = { + "spend:key:hashed-key": 0.2, + "spend:team:team-1": 0.3, + "spend:team_member:user-1:team-1": 0.35, + "spend:user:user-1": 0.4, + "spend:end_user:end-user-1": 0.45, + "spend:org:org-1": 0.5, + "spend:project:project-1": 0.55, + "spend:tag:batch-jobs": 0.6, + "spend:model_access_group:premium": 0.65, +} + + +async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> None: + await logger.async_post_call_failure_hook( + request_data={ + "model": "gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"tags": ["batch-jobs"], MODEL_ACCESS_GROUP_METADATA_KEY: ["premium"]}, + "proxy_server_request": {"request_id": "rid"}, + **failure, + }, + original_exception=Exception("stream broke mid-flight"), + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-key", + team_id="team-1", + user_id="user-1", + end_user_id="end-user-1", + org_id="org-1", + project_id="project-1", + ), + ) + + +async def _counter_values(counters: DualCache) -> dict[str, float]: + return {key: await counters.async_get_cache(key) for key in _LIVE_COUNTERS} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "charged"), + [ + ( + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + }, + 0.125, + ), + (_failure_with_billed_attempts(), 1.25), + ({}, 0.0), + ], +) +async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost( + live_spend_counters: DualCache, failure: Mapping[str, object], charged: float +): + """#35568: budget checks read the live counters, so a failure row's cost has to land there as well as in the + spend log, or an over-budget key keeps passing until the counter is next repaired from the database.""" + for key, spend in _LIVE_COUNTERS.items(): + await live_spend_counters.async_set_cache(key, spend) + + await _fail_request(ProxyDBLogger(), failure) + + assert await _counter_values(live_spend_counters) == pytest.approx( + {key: spend + charged for key, spend in _LIVE_COUNTERS.items()} + ) + + +@pytest.mark.asyncio +async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: DualCache): + """A failure whose cost was already finalized settles through the success path's write, which raises the + counters itself.""" + for key, spend in _LIVE_COUNTERS.items(): + await live_spend_counters.async_set_cache(key, spend) + + await _fail_request( + ProxyDBLogger(), {"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}} + ) + + assert await _counter_values(live_spend_counters) == pytest.approx( + {key: spend + 0.125 for key, spend in _LIVE_COUNTERS.items()} + ) + + +@pytest.mark.asyncio +async def test_failure_spend_rows_leave_the_live_spend_counters_alone_while_spend_updates_are_disabled( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_spend_updates": True}) + for key, spend in _LIVE_COUNTERS.items(): + await live_spend_counters.async_set_cache(key, spend) + + await _fail_request(ProxyDBLogger(), _failure_with_billed_attempts()) + + assert await _counter_values(live_spend_counters) == pytest.approx(_LIVE_COUNTERS) + + +@pytest.mark.asyncio +async def test_a_failing_spend_counter_store_is_still_sent_every_failure_spend_row(): + counters: Final = AsyncMock(side_effect=RuntimeError("redis unavailable")) + + await _fail_request(ProxyDBLogger(spend_counters=lambda: counters), _failure_with_billed_attempts()) + + assert [call.kwargs["response_cost"] for call in counters.call_args_list] == [0.5, 0.25, 0.5] + + +@pytest.mark.asyncio +@pytest.mark.timeout(30) +async def test_a_slow_spend_counter_store_lets_the_failure_out_and_raises_the_counters_once_it_answers(): + store_answers: Final = asyncio.Event() + raised_cost: Final[asyncio.Future[object]] = asyncio.get_running_loop().create_future() + + async def slow_counter_store(**counter_args: object) -> None: + await store_answers.wait() + raised_cost.set_result(counter_args["response_cost"]) + + logger: Final = ProxyDBLogger(spend_counters=lambda: slow_counter_store, failure_counter_wait_seconds=0) + + await _fail_request( + logger, + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + }, + ) + assert not raised_cost.done() + + store_answers.set() + assert await raised_cost == 0.125 + + +@pytest.mark.asyncio +async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_request_start( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + """The window row's one-time seed from the spend log stops at the request start, which keeps the failure row + it already logged from being counted twice.""" + from litellm.proxy.proxy_server import proxy_logging_obj + + window_queue: Final = proxy_logging_obj.db_spend_update_writer.window_spend_update_queue + await window_queue.flush_and_get_aggregated_window_spend_transactions() + key_cache: Final = DualCache() + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", key_cache) + request_start: Final = datetime(2026, 10, 20, 12, 0, tzinfo=timezone.utc) + reset_at: Final = request_start + timedelta(days=10) + await key_cache.async_set_cache( + "hashed-key", + UserAPIKeyAuth( + token="hashed-key", + budget_limits=[{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}], + ), + ) + + await _fail_request( + ProxyDBLogger(), + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + "start_time": request_start, + }, + ) + + assert [ + (row["entity_type"], row["entity_id"], row["spend"], row["started_at"]) + for row in await window_queue.flush_and_get_aggregated_window_spend_transactions() + ] == [("key", "hashed-key", pytest.approx(0.125), "2026-10-20T12:00:00.000000")] + + +@pytest.mark.asyncio +async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_attributed_access_group( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + """Its spend row names the dropped deployment, which narrows nothing, so the counters must not narrow it to the + final attempt's groups either.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + litellm.Router( + model_list=[ + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "sk-test"}, + "model_info": {"id": "dep-b", "access_groups": ["pool"]}, + } + ] + ), + ) + + await _fail_request( + ProxyDBLogger(), + { + **_failure_with_billed_attempts(), + "metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool", "premium"], "model_info": {"id": "dep-b"}}, + }, + ) + + assert { + group: await live_spend_counters.async_get_cache(f"spend:model_access_group:{group}") + for group in ("premium", "pool") + } == pytest.approx({"premium": 0.25, "pool": 1.25}) + + @pytest.mark.asyncio async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): """MCP tool calls may only carry user_api_key; user/team rollups still need user_id.""" From 632f71b850d157efbfa75aa4c1f7a2066576446c Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 13:19:58 -0700 Subject: [PATCH 58/64] fix(router): recover streamed context window errors before output (#49293) --- litellm/responses/streaming_iterator.py | 2 +- litellm/router.py | 81 ++++++- ..._mantle_context_overflow_fallbacks_wire.py | 10 +- .../test_streaming_iterator_error_events.py | 19 +- tests/unit/test_router/test_router.py | 15 +- .../unit/test_router/test_router_fallbacks.py | 201 ++++++++++++++++++ 6 files changed, 303 insertions(+), 25 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index fb91ba6a618..d42f1af380e 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -312,7 +312,7 @@ def map_stream_error_to_exception( error_body: Final = {"message": error_message, "type": error_type, "code": error_code} provider_exception: Final = BaseLLMException( status_code=status_code, - message=f"Error code: {status_code} - {{'error': {error_body}}}", + message=f"Error code: {status_code} - {json.dumps({'error': error_body})}", body=error_body, headers=headers, ) diff --git a/litellm/router.py b/litellm/router.py index 233e7581a40..538cd751fb7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -897,6 +897,23 @@ def _stream_pre_content_error( ) +def _context_window_stream_fallback_error( + error: "MidStreamFallbackError | litellm.ContextWindowExceededError", + model: object, + llm_provider: str, + has_output: bool, +) -> "MidStreamFallbackError": + if isinstance(error, litellm.ContextWindowExceededError): + if has_output: + raise error + return _stream_pre_content_error(error, model if isinstance(model, str) else "", llm_provider) + if isinstance(error.original_exception, litellm.ContextWindowExceededError) and ( + has_output or error.generated_content or error.emitted_disqualifying_content + ): + raise error.original_exception + return error + + def _deployment_num_retries(deployment: "Deployment | None") -> int | None: """The deployment's own num_retries litellm_param, an int or a digit string the way the config loader leaves it.""" configured: Final = getattr(deployment.litellm_params, "num_retries", None) if deployment is not None else None @@ -3845,7 +3862,19 @@ class Router(DeprecatedRouterMembers): try: async for item in _achunks_with_first(model_response, first_chunk): yield item - except MidStreamFallbackError as e: + except (MidStreamFallbackError, litellm.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + getattr(model_response, "model", ""), + model_response.custom_llm_provider or "", + bool(getattr(model_response, "response_uptil_now", "")) + or _stream_chunks_have_generated_content(model_response.chunks), + ) + if isinstance(stream_error, litellm.ContextWindowExceededError) + or isinstance(stream_error.original_exception, litellm.ContextWindowExceededError) + else stream_error + ) with anyio.CancelScope(shield=True): await close_model_response() await held_slot.aclose() @@ -3917,7 +3946,9 @@ class Router(DeprecatedRouterMembers): initial_kwargs["messages"] = messages self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = await self.async_function_with_fallbacks_common_utils( - e=e, + e=e.original_exception + if isinstance(e.original_exception, litellm.ContextWindowExceededError) + else e, disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, @@ -3985,10 +4016,10 @@ class Router(DeprecatedRouterMembers): if fallback_response is not None and hasattr(fallback_response, "aclose"): try: await fallback_response.aclose() - except BaseException as e: + except BaseException as close_error: verbose_router_logger.debug( "stream_with_fallbacks: error closing fallback_response: %s", - e, + close_error, ) wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks()) @@ -4348,6 +4379,10 @@ class Router(DeprecatedRouterMembers): def __aiter__(self): return self + @property + def yielded_first_chunk(self) -> bool: + return self._yielded_first_chunk + async def __anext__(self): try: chunk: Final = await self._async_generator.__anext__() @@ -4375,6 +4410,7 @@ class Router(DeprecatedRouterMembers): # the very key that created the container (#30210). if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: self.completed_response = chunk + self._yielded_first_chunk = True return chunk async def aclose(self): @@ -4401,7 +4437,18 @@ class Router(DeprecatedRouterMembers): yield item for held_event in held_lifecycle_events: yield held_event - except MidStreamFallbackError as e: + except (MidStreamFallbackError, litellm.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + source_iterator.model or "", + source_iterator.custom_llm_provider or "", + wrapper.yielded_first_chunk, + ) + if isinstance(stream_error, litellm.ContextWindowExceededError) + or isinstance(stream_error.original_exception, litellm.ContextWindowExceededError) + else stream_error + ) if response_id_exposed: self._raise_released_stream_error(e, initial_kwargs) async with contextlib.aclosing( @@ -4572,7 +4619,11 @@ class Router(DeprecatedRouterMembers): # The content-policy dispatch branch matches on the trigger's own type, so a refusal's # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. fallback_trigger: Final[Exception] = ( - e.original_exception if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e + e.original_exception + if isinstance( + e.original_exception, (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError) + ) + else e ) fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success e=fallback_trigger, @@ -4658,7 +4709,19 @@ class Router(DeprecatedRouterMembers): try: for item in _chunks_with_first(model_response, first_chunk): yield item - except MidStreamFallbackError as e: + except (MidStreamFallbackError, litellm.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + getattr(model_response, "model", ""), + model_response.custom_llm_provider or "", + bool(getattr(model_response, "response_uptil_now", "")) + or _stream_chunks_have_generated_content(model_response.chunks), + ) + if isinstance(stream_error, litellm.ContextWindowExceededError) + or isinstance(stream_error.original_exception, litellm.ContextWindowExceededError) + else stream_error + ) if fallbacks_disabled_for_request(initial_kwargs) or ( not e.is_pre_first_chunk and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) @@ -4683,7 +4746,9 @@ class Router(DeprecatedRouterMembers): router_self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = run_async_function( router_self.async_function_with_fallbacks_common_utils, - e=e, + e=e.original_exception + if isinstance(e.original_exception, litellm.ContextWindowExceededError) + else e, disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, diff --git a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py index 37aa96c2467..23b75a5a8cf 100644 --- a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py @@ -419,12 +419,11 @@ def test_context_window_fallback_is_not_consulted_on_streamed_messages(cwf_rig: _assert_no_fallback(cwf_rig, "overflow", prompt) -def test_context_window_fallback_does_not_fire_on_streamed_responses(cwf_rig: Rig) -> None: +def test_context_window_fallback_fires_on_streamed_responses(cwf_rig: Rig) -> None: prompt: Final = _prompt("cwf responses stream") response, text = _consumed(cwf_rig.gateway, "/v1/responses", _responses(cwf_rig.models["cwf"], prompt, True)) assert response.status_code == 200, text - assert _GENERIC in str(_only_failed(text)["message"]), text - _assert_no_fallback(cwf_rig, "overflow", prompt) + _assert_fell_back(cwf_rig, "overflow", prompt, text) def test_plain_fallback_fires_for_the_overflow_on_chat_completions(fb_rig: Rig) -> None: @@ -467,12 +466,11 @@ def test_streamed_messages_overflow_on_a_plain_fallback_deployment_returns_400_w _assert_no_fallback(fb_rig, "overflow", prompt) -def test_plain_fallback_does_not_fire_on_streamed_responses(fb_rig: Rig) -> None: +def test_plain_fallback_fires_on_streamed_responses(fb_rig: Rig) -> None: prompt: Final = _prompt("fb responses stream") response, text = _consumed(fb_rig.gateway, "/v1/responses", _responses(fb_rig.models["fb"], prompt, True)) assert response.status_code == 200, text - assert _GENERIC in str(_only_failed(text)["message"]), text - _assert_no_fallback(fb_rig, "overflow", prompt) + _assert_fell_back(fb_rig, "overflow", prompt, text) def test_streamed_messages_overflow_on_an_openai_fallback_deployment_returns_the_error_event_without_falling_back( diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 7445070bd3a..14ecb185f06 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -166,21 +166,24 @@ def test_maybe_raise_for_error_event_maps_context_length_code_to_400(): assert not isinstance(exc_info.value, MidStreamFallbackError) -def test_maybe_raise_for_error_event_raises_context_window_exceeded_directly(): +@pytest.mark.parametrize( + "message", + ( + "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", + "Your input exceeds the context window of this model. Please adjust your input and try again.", + ), +) +def test_maybe_raise_for_error_event_raises_context_window_exceeded_directly(message: str) -> None: """A context-length error event maps to ContextWindowExceededError exactly like the non-streaming path and, being a non-retriable client error, is raised directly rather than wrapped for mid-stream fallback, preserving the direct-SDK 400 contract from issue #15785.""" - iterator = _make_iterator() - chunk = _make_error_chunk( - "invalid_request_error", - "context_length_exceeded", - "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", - ) + iterator: Final = _make_iterator() + chunk: Final = _make_error_chunk("invalid_request_error", "context_length_exceeded", message) with pytest.raises(litellm.ContextWindowExceededError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) - assert "maximum context length" in str(exc_info.value) + assert message in str(exc_info.value) CONTENT_POLICY_MESSAGE = "This content was flagged for possible cybersecurity risk. The response was halted mid-stream." diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index a6436e47507..96df8d3e1f0 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -5140,7 +5140,10 @@ async def test_set_response_headers_is_the_only_complexity_header_source_for_pro @pytest.mark.asyncio -async def test_acompletion_streaming_iterator_adopts_fallback_response_headers(): +@pytest.mark.parametrize("close_fails", [False, True]) +async def test_acompletion_streaming_iterator_adopts_fallback_response_headers( + close_fails: bool, +) -> None: """LIT-6767: after a successful pre-first-chunk fallback, the wrapper must describe the deployment that served the stream, with no value left over from the attempt that failed.""" @@ -5187,6 +5190,7 @@ async def test_acompletion_streaming_iterator_adopts_fallback_response_headers() class FallbackStream: def __init__(self): + self.close_attempted: bool = False self._response_headers = {"x-request-id": "req-FALLBACK"} self._hidden_params = { "model_id": "fallback-deployment", @@ -5204,6 +5208,11 @@ async def test_acompletion_streaming_iterator_adopts_fallback_response_headers() except StopIteration: raise StopAsyncIteration from None + async def aclose(self) -> None: + self.close_attempted = True + if close_fails: + raise RuntimeError("fallback transport failed during close") + fallback_stream = FallbackStream() with patch.object( router, @@ -5220,12 +5229,14 @@ async def test_acompletion_streaming_iterator_adopts_fallback_response_headers() # the very first chunk the fallback produces must already be published under # the fallback's identity: the proxy commits response headers once that chunk # is buffered, so adopting any later is adopting too late - await result.__anext__() + first_chunk: Final = await result.__anext__() + assert first_chunk.choices[0].delta.content == "OK" assert result._response_headers == {"x-request-id": "req-FALLBACK"} assert result._hidden_params["model_id"] == "fallback-deployment" async for _ in result: pass + assert fallback_stream.close_attempted assert result._response_headers == {"x-request-id": "req-FALLBACK"} assert result._hidden_params["model_id"] == "fallback-deployment" assert result._hidden_params["api_base"] == "https://fallback.example" diff --git a/tests/unit/test_router/test_router_fallbacks.py b/tests/unit/test_router/test_router_fallbacks.py index 88d89696297..bccd6290f44 100644 --- a/tests/unit/test_router/test_router_fallbacks.py +++ b/tests/unit/test_router/test_router_fallbacks.py @@ -1006,3 +1006,204 @@ async def test_unknown_default_fallback_raises_after_primary_failure( assert route.call_count == 1 assert "missing" in str(exc_info.value) + + +def _context_error_response_stream( + *, event_type: str = "error", preceding_event: dict[str, object] | None = None, completed: bool = False +) -> httpx.Response: + response: Final = { + "id": "resp_backup" if completed else "resp_primary", + "object": "response", + "created_at": 1700000000, + "model": "backup" if completed else "primary", + "status": "completed" if completed else "in_progress", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "temperature": 1.0, + "top_p": 1.0, + } + error: Final = { + "type": "invalid_request_error", + "code": "context_length_exceeded", + "message": "Your input exceeds the context window of this model. Please adjust your input and try again.", + } + terminal: Final = ( + {"type": "response.completed", "response": response, "sequence_number": 2} + if completed + else {"type": event_type, "error": error, "sequence_number": 2} + if event_type == "error" + else {"type": event_type, "response": {**response, "status": "failed", "error": error}, "sequence_number": 2} + ) + events: Final = ( + {"type": "response.created", "response": response, "sequence_number": 0}, + *((preceding_event,) if preceding_event is not None else ()), + terminal, + ) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode(), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_type", ("error", "response.failed")) +@pytest.mark.parametrize("context_fallback", (False, True)) +async def test_responses_context_error_uses_the_configured_fallback_before_output( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, event_type: str, context_fallback: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _fallback_router(fallbacks=[{"primary": ["missing"]}] if context_fallback else None) + route: Final = respx_mock.post("https://fallback-migration.local/v1/responses").mock( + side_effect=( + _context_error_response_stream(event_type=event_type), + _context_error_response_stream(completed=True), + ) + ) + stream: Final = await router.aresponses( + model="primary", + input="oversized prompt", + stream=True, + context_window_fallbacks=[{"primary": ["backup"]}] if context_fallback else None, + ) + events: Final = [event async for event in stream] + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "backup") + assert tuple(event.type for event in events) == ("response.created", "response.completed") + assert events[0].response.id == events[-1].response.id + assert events[-1].response.model == "backup" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("barrier", ("text", "tool", "disabled", "unconfigured")) +async def test_responses_context_error_does_not_restart_committed_or_disabled_streams( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, barrier: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _fallback_router(fallbacks=[] if barrier == "unconfigured" else None) + preceding: Final = { + "text": { + "type": "response.output_text.delta", + "delta": "already visible", + "item_id": "msg_primary", + "output_index": 0, + "content_index": 0, + "sequence_number": 1, + }, + "tool": { + "type": "response.output_item.added", + "item": { + "type": "function_call", "id": "fc_primary", "call_id": "call_primary", + "name": "lookup", "arguments": "", + }, + "output_index": 0, + "sequence_number": 1, + }, + }.get(barrier) + route: Final = respx_mock.post("https://fallback-migration.local/v1/responses").mock( + return_value=_context_error_response_stream(preceding_event=preceding) + ) + stream: Final = await router.aresponses( + model="primary", input="oversized prompt", stream=True, disable_fallbacks=barrier == "disabled" + ) + with pytest.raises(litellm.ContextWindowExceededError): + async for _ in stream: + pass + assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", (False, True)) +@pytest.mark.parametrize("emission", ("role", "text", "tool")) +async def test_chat_context_error_falls_back_after_role_only_but_not_after_output( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool, emission: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _fallback_router(fallbacks=[{"primary": ["missing"]}]) + delta: Final = ( + {"role": "assistant", "tool_calls": [{ + "index": 0, "id": "call_primary", "type": "function", "function": {"name": "lookup", "arguments": ""} + }]} + if emission == "tool" + else {"role": "assistant", "content": "visible" if emission == "text" else ""} + ) + first_chunk: Final = { + "id": "chatcmpl-primary", "object": "chat.completion.chunk", "created": 1700000000, "model": "primary", + "choices": [{"index": 0, "delta": delta}], + } + error: Final = {"error": { + "type": "invalid_request_error", "code": "context_length_exceeded", + "message": "This model's maximum context length is 128000 tokens. Your messages resulted in 130000 tokens.", + }} + backup: Final = { + "id": "chatcmpl-backup", "object": "chat.completion.chunk", "created": 1700000000, "model": "backup", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "backup answer"}, "finish_reason": "stop"}], + } + route: Final = respx_mock.post("https://fallback-migration.local/v1/chat/completions").mock( + side_effect=( + httpx.Response(200, content=f"data: {json.dumps(first_chunk)}\n\ndata: {json.dumps(error)}\n\n".encode()), + httpx.Response(200, content=f"data: {json.dumps(backup)}\n\ndata: [DONE]\n\n".encode()), + ) + ) + request: Final = { + "model": "primary", "messages": [{"role": "user", "content": "oversized prompt"}], "stream": True, + "context_window_fallbacks": [{"primary": ["backup"]}], + } + stream: Final = router.completion(**request) if sync_mode else await router.acompletion(**request) + if emission != "role": + with pytest.raises(litellm.ContextWindowExceededError): + tuple(stream) if sync_mode else tuple([chunk async for chunk in stream]) + assert route.call_count == 1 + return + chunks: Final = tuple(stream) if sync_mode else tuple([chunk async for chunk in stream]) + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "backup") + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "backup answer" + + +@pytest.mark.parametrize( + "wrapped,has_output,generated_content,disqualifying,recoverable", + ( + (False, False, "", False, True), + (False, True, "", False, False), + (True, False, "", False, True), + (True, True, "", False, False), + (True, False, "already visible", False, False), + (True, False, "", True, False), + ), +) +def test_context_stream_fallback_obeys_all_output_receipts( + wrapped: bool, has_output: bool, generated_content: str, disqualifying: bool, recoverable: bool +) -> None: + from litellm.exceptions import MidStreamFallbackError + from litellm.router import _context_window_stream_fallback_error + + original: Final = litellm.ContextWindowExceededError(message="too long", model="primary", llm_provider="openai") + error: Final = ( + MidStreamFallbackError( + message="too long", model="primary", llm_provider="openai", original_exception=original, + is_pre_first_chunk=True, generated_content=generated_content, emitted_disqualifying_content=disqualifying, + ) + if wrapped + else original + ) + if not recoverable: + with pytest.raises(litellm.ContextWindowExceededError) as raised: + _context_window_stream_fallback_error(error, "primary", "openai", has_output) + assert raised.value is original + return + recovered: Final = _context_window_stream_fallback_error(error, "primary", "openai", has_output) + assert recovered.original_exception is original + assert recovered.is_pre_first_chunk is True + if wrapped: + assert recovered is error + + +def test_context_stream_fallback_preserves_other_errors_for_existing_continuation_policy() -> None: + from litellm.exceptions import MidStreamFallbackError + from litellm.router import _context_window_stream_fallback_error + + error: Final = MidStreamFallbackError( + message="provider unavailable", model="primary", llm_provider="openai", generated_content="already visible" + ) + assert _context_window_stream_fallback_error(error, "primary", "openai", True) is error From fbd126e5182040a5ed8b97d54d41b5febc94452b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 30 Oct 2026 13:25:47 -0700 Subject: [PATCH 59/64] test: fix integration tests that fail on test order, races and intended changes (#49310) Seven integration tests fail on main for reasons that are not product bugs. Two are order or race bugs in the tests, one reads rows other tests wrote, one passed on an error body, and three still expect behavior that #48644, #39864 and #39540 changed on purpose --- .../compatibility/test_missing_body_param_status.py | 2 +- .../database/migrations/test_request_log_indexes.py | 9 ++++++--- .../integration/observability/test_guardrail_effects.py | 7 +++++-- .../providers/test_gemini_tts_voice_mapping_wire.py | 4 ++-- .../providers/test_hosted_vllm_stream_stall_wire.py | 6 +++++- tests/integration/routing/test_observed_routing.py | 2 +- .../security/test_hide_secrets_tool_compatibility.py | 8 +++++++- 7 files changed, 27 insertions(+), 11 deletions(-) diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index 005e581c643..d47503e7f08 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -1199,7 +1199,7 @@ def test_anthropic_messages_explicit_null_reaches_upstream(gateway: Gateway) -> "anthropic_messages", deployment_params={"max_tokens": 32}, provider_model="openai/gpt-4o-mini", - response_route="acompletion", + response_route="aresponses", ) response: Final = _post( gateway, diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index d06f0b13cff..5741c85470b 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -1033,13 +1033,14 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does( create_index: Final = sql.SQL('CREATE INDEX "entry_probe_idx" ON "entry_probe" ({})').format( sql.SQL(", ").join(sql.Identifier(column) for column in index_columns) ) + fitting_key: Final = secrets.token_hex(longest_fitting_key)[:longest_fitting_key] with psycopg.connect(scratch_database, autocommit=True) as conn: conn.execute('CREATE TABLE "entry_probe" ("s" TEXT, "k" TEXT, "t" TIMESTAMP(3))') conn.execute( 'INSERT INTO "entry_probe" ("s", "k", "t") VALUES (%s, %s, %s), (%s, %s, %s)', ( s_value, - secrets.token_hex(longest_fitting_key)[:longest_fitting_key], + fitting_key, t_value, s_value, secrets.token_hex(longest_fitting_key)[: longest_fitting_key + 1], @@ -1051,7 +1052,8 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does( ) with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="exceeds btree version 4 maximum 2704"): conn.execute(create_index) - conn.execute('DELETE FROM "entry_probe" WHERE octet_length("k") > %s', (longest_fitting_key,)) + conn.execute('TRUNCATE "entry_probe"') + conn.execute('INSERT INTO "entry_probe" ("s", "k", "t") VALUES (%s, %s, %s)', (s_value, fitting_key, t_value)) conn.execute(create_index) (finding,) = findings @@ -1075,7 +1077,8 @@ def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random ) with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"): conn.execute(create_index) - conn.execute('DELETE FROM "compress_probe" WHERE "k" <> %s', ("a" * 3000,)) + conn.execute('TRUNCATE "compress_probe"') + conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s)', ("a" * 3000,)) conn.execute(create_index) stored: Final = conn.execute('SELECT octet_length("k"), pg_column_size("k") FROM "compress_probe"').fetchone() diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 4f5c04bddcc..1220a65528c 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -23,7 +23,7 @@ import yaml from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.mcp import mcp_peer, register_mcp, tool_names -from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process +from integration._support.process import OwnedProxy, OwnedState, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from openai import APIError, AsyncOpenAI, OpenAI @@ -3263,7 +3263,10 @@ def test_client_disconnect_during_the_body_upload_runs_no_endpoint_and_writes_no with ( wire_server(_openai_upstream) as upstream, owned_proxy_process( - gateway, tmp_path, {"OPENAI_BASE_URL": upstream.url + "/v1", "OPENAI_API_KEY": "synthetic-openai-key"} + gateway, + tmp_path, + {"OPENAI_BASE_URL": upstream.url + "/v1", "OPENAI_API_KEY": "synthetic-openai-key"}, + state=OwnedState.PRIVATE, ) as owned, owned.gateway.scenario() as scenario, ): diff --git a/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py b/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py index 90ff32e01dd..20472dc2399 100644 --- a/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py +++ b/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py @@ -501,7 +501,7 @@ def test_vertex_batch_file_rows_carry_the_mapped_voice(gateway: Gateway) -> None assert tuple(text_in(row) for row in rows) == (first, second) -def test_vertex_live_default_mode_sends_one_setup_without_the_client_voice(gateway: Gateway) -> None: +def test_vertex_live_sends_one_setup_carrying_the_mapped_session_update_voice(gateway: Gateway) -> None: with gateway.scenario() as scenario: handle: Final = live_scenario(scenario) key: Final = scenario.key() @@ -511,7 +511,7 @@ def test_vertex_live_default_mode_sends_one_setup_without_the_client_voice(gatew assert types[0] == "session.created" and types[-1] == "response.done", frames setups: Final = live_setups(gateway, handle.scenario_id) assert len(setups) == 1, setups - assert voice_in(setups[0]) == NO_VOICE, setups[0] + assert voice_in(setups[0]) == "Umbriel", setups[0] @dataclass(frozen=True, slots=True) diff --git a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py index fad7898871d..018c6a653d1 100644 --- a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py +++ b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py @@ -279,7 +279,11 @@ def cooldown_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Cooldown upstream, directory, {"PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": "2"}, - config=cooldown_config(directory, router_settings(cache, allowed_fails=0), prometheus=False), + config=cooldown_config( + directory, + router_settings(cache, allowed_fails=0, allowed_fails_policy={"TimeoutErrorAllowedFails": 0}), + prometheus=False, + ), ) ) yield _CooldownRig(owned.gateway, cache) diff --git a/tests/integration/routing/test_observed_routing.py b/tests/integration/routing/test_observed_routing.py index ca5f248e0de..64ab799fe3a 100644 --- a/tests/integration/routing/test_observed_routing.py +++ b/tests/integration/routing/test_observed_routing.py @@ -593,7 +593,7 @@ def test_health_routing_fallbacks_preserve_error_specific_targets_and_request_co "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}], }, response_text assert {key: events[1]["error"][key] for key in events[1]["error"] if key != "message"} == { - "type": None, + "type": "internal_server_error", "param": None, "code": "500", }, response_text diff --git a/tests/integration/security/test_hide_secrets_tool_compatibility.py b/tests/integration/security/test_hide_secrets_tool_compatibility.py index d8b900de6d0..a89678aef8e 100644 --- a/tests/integration/security/test_hide_secrets_tool_compatibility.py +++ b/tests/integration/security/test_hide_secrets_tool_compatibility.py @@ -140,9 +140,15 @@ def _key(rig: Rig, *, permissions: dict[str, JsonValue] | None = None) -> Iterat ) +def _without_cooldowns(configuration: dict[str, JsonValue]) -> None: + settings: Final = configuration["router_settings"] + assert isinstance(settings, dict) + settings["disable_cooldowns"] = True + + @pytest.fixture(scope="module") def caller(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Caller]: - with tool_rig(tmp_path_factory.mktemp("hide-secrets-compat"), _peer) as rig: + with tool_rig(tmp_path_factory.mktemp("hide-secrets-compat"), _peer, configure=_without_cooldowns) as rig: with _key(rig) as key: yield Caller(rig, key) From 0382f134e74e027d923c0cde7ad013174990b9b4 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 13:39:01 -0700 Subject: [PATCH 60/64] fix(ui): add Meta Model API to the Add Model provider dropdown (#49314) * test(ui): cover separate Meta Model API provider onboarding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): add Meta Model API provider onboarding Co-authored-by: lostmartian Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): align Meta coverage with provider keys and listed registry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): show Meta model placeholder Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: lostmartian --- .../provider_create_fields.json | 28 +++++++++++++++++++ .../public_endpoints/test_public_endpoints.py | 23 ++++++++++++++- .../components/provider_info_helpers.test.tsx | 16 +++++++++++ .../src/components/provider_info_helpers.tsx | 6 +++- 4 files changed, 71 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index df3af4ad487..d849b9fbaac 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2377,6 +2377,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "META", + "provider_display_name": "Meta Model API", + "litellm_provider": "meta", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": null, + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": false, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "meta/muse-spark-1.1" + }, { "provider": "MILVUS", "provider_display_name": "Milvus", diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 48672c7a305..3fe5345d008 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -28,6 +28,28 @@ from litellm.types.proxy.public_endpoints.public_endpoints import ProviderCreate from litellm.types.utils import LlmProviders +def test_meta_model_api_fields_are_separate_from_meta_llama() -> None: + app: Final = FastAPI() + app.include_router(router) + response: Final = TestClient(app).get("/public/providers/fields") + assert response.status_code == 200 + providers: Final = response.json() + meta: Final = next((p for p in providers if p["litellm_provider"] == "meta"), None) + llama: Final = next(p for p in providers if p["litellm_provider"] == "meta_llama") + assert meta is not None + assert meta["provider"] == "META" + assert meta["provider_display_name"] == "Meta Model API" + assert meta["default_model_placeholder"] == "meta/muse-spark-1.1" + fields: Final = {field["key"]: field for field in meta["credential_fields"]} + assert set(fields) == {"api_base", "api_key"} + assert fields["api_base"]["field_type"] == "text" + assert fields["api_key"]["field_type"] == "password" + assert fields["api_base"]["required"] is False + assert fields["api_key"]["required"] is False + assert llama["provider"] == "LLAMA" + assert llama["provider_display_name"] == "Meta Llama" + + def test_bigmodel_fields_require_private_key_and_share_default_base_with_adapter() -> None: from litellm.llms.bigmodel.common_utils import BIGMODEL_API_BASE @@ -568,7 +590,6 @@ ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset( "libertai", "litellm_agent", "manus", - "meta", "modelscope", "mongodb", "nano-gpt", diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 02c31ecc082..4c483434831 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -20,6 +20,22 @@ const BUNDLED_MODEL_MAP: Record = JSON.pa ); describe("provider_info_helpers", () => { + it("resolves Meta Model API separately from Meta Llama and isolates their bundled models", () => { + expect(resolveLitellmProviderSlug("META")).toBe("meta"); + expect(getPlaceholder(Providers.META)).toBe("meta/muse-spark-1.1"); + expect(getProviderLogoAndName("meta").displayName).toBe("Meta Model API"); + const metaModels = Object.keys(BUNDLED_MODEL_MAP).filter( + (key) => BUNDLED_MODEL_MAP[key]?.litellm_provider === "meta", + ); + const llamaModels = Object.keys(BUNDLED_MODEL_MAP).filter( + (key) => BUNDLED_MODEL_MAP[key]?.litellm_provider === "meta_llama", + ); + expect(metaModels).toContain("meta/muse-spark-1.1"); + expect(llamaModels.length).toBeGreaterThan(0); + expect(getProviderModels("META", BUNDLED_MODEL_MAP).sort()).toEqual(metaModels.sort()); + expect(getProviderModels("LLAMA", BUNDLED_MODEL_MAP).sort()).toEqual(llamaModels.sort()); + }); + it("resolves BigModel metadata, name, logo and model placeholder", () => { const metadata = PROVIDER_CREATE_FIELDS.find((provider) => provider.provider === "BIGMODEL"); expect(provider_map.BIGMODEL).toBe(metadata?.litellm_provider); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 6448bd936cc..a8491f656a0 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -140,6 +140,7 @@ export enum Providers { LLAMAFILE = "Llamafile", LM_STUDIO = "Lm Studio", LLAMA = "Meta Llama", + META = "Meta Model API", MARITALK = "Maritalk", MICROSOFT_365_COPILOT = "Microsoft 365 Copilot", MiniMax = "MiniMax", @@ -270,6 +271,7 @@ export const provider_map: Record = { LEMONADE: "lemonade", LLAMAFILE: "llamafile", LLAMA: "meta_llama", + META: "meta", LM_STUDIO: "lm_studio", MICROSOFT_365_COPILOT: "microsoft_365_copilot", MARITALK: "maritalk", @@ -337,7 +339,7 @@ export const provider_map: Record = { ZeroGPU: "zerogpu", }; -const standaloneSubproviderSlugs = new Set(["bedrock_mantle"]); +const standaloneSubproviderSlugs = new Set(["bedrock_mantle", "meta_llama"]); export const resolveLitellmProviderSlug = (provider: string): string => provider_map[provider] ?? provider.toLowerCase(); @@ -393,6 +395,7 @@ export const providerLogoMap: Partial> = { [Providers.LAMBDA_AI]: lambdaLogo.src, [Providers.LM_STUDIO]: lmstudioLogo.src, [Providers.LLAMA]: metaLlamaLogo.src, + [Providers.META]: metaLlamaLogo.src, [Providers.MiniMax]: minimaxLogo.src, [Providers.MICROSOFT_365_COPILOT]: microsoft365CopilotLogo.src, [Providers.MistralAI]: mistralLogo.src, @@ -491,6 +494,7 @@ const providerPlaceholderMap: Partial> = { [Providers.FalAI]: "fal_ai/fal-ai/flux-pro/v1.1-ultra", [Providers.Google_AI_Studio]: "gemini-pro", [Providers.JinaAI]: "jina_ai/", + [Providers.META]: "meta/muse-spark-1.1", [Providers.MICROSOFT_365_COPILOT]: "microsoft_365_copilot/chat", [Providers.NVIDIA_RIVA]: "nvidia_riva/nvidia/parakeet-ctc-1_1b-asr", [Providers.Oracle]: "oci/xai.grok-4", From 516c796e86a5b1028b2073cd74ca341a174fa289 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 13:46:27 -0700 Subject: [PATCH 61/64] fix(rerank): preserve token usage in spend logs (#49308) --- litellm/litellm_core_utils/litellm_logging.py | 11 +- litellm/llms/bedrock/rerank/transformation.py | 6 +- .../llms/hosted_vllm/rerank/transformation.py | 4 +- .../llms/together_ai/rerank/transformation.py | 4 +- litellm/types/rerank.py | 46 +++++++- .../test_litellm_logging.py | 100 ++++++++++++++++++ .../test_hosted_vllm_rerank_transformation.py | 35 ++++++ .../jina_ai/rerank/test_transformation.py | 5 +- 8 files changed, 200 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 26949d9bbfb..924603dc7e8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -165,7 +165,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.mcp import MCPPostCallResponseObject from litellm.types.prompts.init_prompts import PromptSpec -from litellm.types.rerank import RerankResponse +from litellm.types.rerank import RerankResponse, get_rerank_usage from litellm.types.router import RouterRateLimitError, SilentExperimentFlag from litellm.types.utils import ( DEPLOYMENT_SCOPED_PRICING_FIELDS, @@ -355,6 +355,7 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys _CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) _METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_RESPONSE_MAPPING: Final = TypeAdapter(Mapping[str, object]) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS _UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( ("user_api_key_auth", "user_api_key_budget_reservation") @@ -6817,7 +6818,11 @@ class StandardLoggingPayloadSetup: total_tokens=0, ) - usage: Final = response_obj.get("usage", None) or {} + usage: Final = response_obj.get("usage", None) + if usage is None: + rerank_usage: Final = get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj)) + if rerank_usage is not None: + return Usage(**rerank_usage) if usage is None or (not isinstance(usage, dict) and not isinstance(usage, Usage)): return Usage( prompt_tokens=0, @@ -6853,7 +6858,7 @@ class StandardLoggingPayloadSetup: return _empty _raw: Final = response_obj.get("usage", None) if _raw is None: - return _empty + return dict(get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj)) or _empty) if isinstance(_raw, ResponseAPIUsage): return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(_raw).model_dump() if isinstance(_raw, dict): diff --git a/litellm/llms/bedrock/rerank/transformation.py b/litellm/llms/bedrock/rerank/transformation.py index 442903fc05e..c2f5035528d 100644 --- a/litellm/llms/bedrock/rerank/transformation.py +++ b/litellm/llms/bedrock/rerank/transformation.py @@ -6,6 +6,8 @@ Why separate file? Make it easy to see how transformation works from typing import Final +from pydantic import TypeAdapter + from litellm._uuid import uuid from litellm.types.llms.bedrock import ( BedrockRerankBedrockRerankingConfiguration, @@ -90,7 +92,9 @@ class BedrockRerankConfig: example input: {"results":[{"index":0,"relevanceScore":0.6847912669181824},{"index":1,"relevanceScore":0.5980774760246277}]} """ - _billed_units = RerankBilledUnits(**response.get("usage", {"search_units": 1})) # by default 1 search unit + _billed_units: Final = TypeAdapter(RerankBilledUnits).validate_python( + response.get("usage", {"search_units": 1}) + ) _tokens: Final = RerankTokens(**response.get("usage", {})) rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 0b77a04c69a..1b7da6f6635 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -244,5 +244,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): return RerankResponse( id=response.get("id") or str(uuid.uuid4()), results=rerank_results, - meta=rerank_meta, + meta=TypeAdapter(RerankResponseMeta).validate_python(response["meta"]) + if "meta" in response and response["meta"] is not None + else rerank_meta, ) diff --git a/litellm/llms/together_ai/rerank/transformation.py b/litellm/llms/together_ai/rerank/transformation.py index f1c3ca9e87d..eb7ec53416d 100644 --- a/litellm/llms/together_ai/rerank/transformation.py +++ b/litellm/llms/together_ai/rerank/transformation.py @@ -6,6 +6,8 @@ Why separate file? Make it easy to see how transformation works from typing import Final +from pydantic import TypeAdapter + from litellm._uuid import uuid from litellm.types.rerank import ( RerankBilledUnits, @@ -19,7 +21,7 @@ from litellm.types.rerank import ( class TogetherAIRerankConfig: def _transform_response(self, response: dict) -> RerankResponse: - _billed_units: Final = RerankBilledUnits(**response.get("usage", {})) + _billed_units: Final = TypeAdapter(RerankBilledUnits).validate_python(response.get("usage", {})) _tokens: Final = RerankTokens(**response.get("usage", {})) rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index ffdcf53fd7e..ca2767a7174 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -4,10 +4,10 @@ https://docs.cohere.com/reference/rerank """ -from collections.abc import Sequence -from typing import Literal, overload +from collections.abc import Mapping, Sequence +from typing import Final, Literal, overload -from pydantic import ConfigDict, PrivateAttr +from pydantic import ConfigDict, PrivateAttr, TypeAdapter, ValidationError from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMBaseModel @@ -57,6 +57,8 @@ class OptionalRerankParams(TypedDict, total=False): class RerankBilledUnits(TypedDict, total=False): search_units: int | None total_tokens: int | None + input_tokens: ReadOnly[int | None] + output_tokens: ReadOnly[int | None] class RerankTokens(TypedDict, total=False): @@ -65,11 +67,47 @@ class RerankTokens(TypedDict, total=False): class RerankResponseMeta(TypedDict, total=False): - api_version: dict | None + api_version: dict[str, object] | None billed_units: RerankBilledUnits | None tokens: RerankTokens | None +class RerankUsage(TypedDict): + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + total_tokens: ReadOnly[int] + + +_RERANK_META: Final = TypeAdapter(RerankResponseMeta) + + +def get_rerank_usage(response: Mapping[str, object]) -> RerankUsage | None: + if "results" not in response or response.get("meta") is None: + return None + try: + meta: Final = _RERANK_META.validate_python(response.get("meta")) + except ValidationError: + return None + tokens: Final = meta.get("tokens") or RerankTokens() + billed: Final = meta.get("billed_units") or RerankBilledUnits() + input_tokens: Final = next( + (count for count in (tokens.get("input_tokens"), billed.get("input_tokens")) if count is not None), None + ) + output_tokens: Final = next( + (count for count in (tokens.get("output_tokens"), billed.get("output_tokens")) if count is not None), None + ) + total_tokens: Final = billed.get("total_tokens") + if input_tokens is None and output_tokens is None and total_tokens is None: + return None + completion_tokens: Final = output_tokens or 0 + prompt_tokens: Final = input_tokens if input_tokens is not None else max((total_tokens or 0) - completion_tokens, 0) + return RerankUsage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens if total_tokens is not None else prompt_tokens + completion_tokens, + ) + + class RerankResponseDocument(TypedDict): text: str diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index f7c2082b683..e952127a32c 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -59,6 +59,7 @@ from litellm.proxy.hooks.max_iterations_limiter import PROXY_MaxIterationsHandle from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse +from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CallTypes, CostBreakdown, @@ -4514,6 +4515,105 @@ def test_logging_preserves_missing_usage_without_accepting_request_metadata( assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (0, 0, 0) +@pytest.mark.parametrize( + ("meta", "expected"), + [ + ({"tokens": {"input_tokens": 7, "output_tokens": 3}}, (7, 3, 10)), + ({"tokens": {}, "billed_units": {"total_tokens": 42}}, (42, 0, 42)), + ({"tokens": {"output_tokens": 3}, "billed_units": {"total_tokens": 12}}, (9, 3, 12)), + ({"tokens": {"input_tokens": 7, "output_tokens": 3}, "billed_units": {"total_tokens": 0}}, (7, 3, 0)), + ({"billed_units": {"input_tokens": 7, "output_tokens": 3}}, (7, 3, 10)), + ( + { + "tokens": {"input_tokens": 7, "output_tokens": 3}, + "billed_units": {"input_tokens": 90, "output_tokens": 10, "total_tokens": 12}, + }, + (7, 3, 12), + ), + ( + { + "tokens": {"input_tokens": 0, "output_tokens": 0}, + "billed_units": {"input_tokens": 90, "output_tokens": 10, "total_tokens": 12}, + }, + (0, 0, 12), + ), + ( + { + "tokens": {"input_tokens": None, "output_tokens": None}, + "billed_units": {"input_tokens": 7, "output_tokens": 3, "total_tokens": None}, + }, + (7, 3, 10), + ), + ({"tokens": {"input_tokens": 0}}, (0, 0, 0)), + (None, None), + ({}, None), + ({"tokens": None, "billed_units": None}, None), + ({"billed_units": {"search_units": 3}}, None), + ], +) +def test_rerank_metadata_reaches_standard_and_spend_logs( + logging_obj: Logging, meta: dict[str, object] | None, expected: tuple[int, int, int] | None +) -> None: + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0) + native: Final = {"id": "rerank-usage", "results": [{"index": 0, "relevance_score": 0.9}], "meta": meta} + response: Final = RerankResponse.model_validate(native) + counts: Final = expected if expected is not None else (0, 0, 0) + expected_usage: Final = dict(zip(("prompt_tokens", "completion_tokens", "total_tokens"), counts)) + assert response.model_dump() == native + usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response.model_dump()) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == counts + assert StandardLoggingPayloadSetup.get_usage_as_dict(response.model_dump()) == ( + expected_usage if expected is not None else {} + ) + kwargs: Final = {"model": "rerank-model", "call_type": "arerank", "litellm_params": {"metadata": {}}} + standard: Final = get_standard_logging_object_payload( + kwargs=kwargs, + init_response_obj=response, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + assert standard is not None + assert (standard["prompt_tokens"], standard["completion_tokens"], standard["total_tokens"]) == counts + assert standard["metadata"]["usage_object"] == (expected_usage if expected is not None else {}) + spend: Final = get_logging_payload( + kwargs={**kwargs, "standard_logging_object": standard}, response_obj=response, start_time=now, end_time=now + ) + assert (spend["prompt_tokens"], spend["completion_tokens"], spend["total_tokens"]) == counts + + +@pytest.mark.parametrize("combined", [None, Usage(prompt_tokens=11, completion_tokens=2, total_tokens=13)]) +@pytest.mark.parametrize("explicit_usage", [{}, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}]) +def test_rerank_metadata_does_not_override_explicit_or_combined_usage( + combined: Usage | None, explicit_usage: dict[str, int] +) -> None: + response: Final = {"results": [], "meta": {"tokens": {"input_tokens": 42}}, "usage": explicit_usage} + expected: Final = combined.model_dump() if combined is not None else explicit_usage + assert StandardLoggingPayloadSetup.get_usage_as_dict(response, combined) == expected + usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response, combined) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + expected.get("prompt_tokens", 0), + expected.get("completion_tokens", 0), + expected.get("total_tokens", 0), + ) + + +@pytest.mark.parametrize( + "response", + [ + {"meta": {"tokens": {"input_tokens": 42}}}, + {"results": [], "meta": {"tokens": {"input_tokens": "invalid"}}}, + ], +) +def test_rerank_metadata_ignores_unrelated_and_invalid_responses(response: dict[str, object]) -> None: + assert StandardLoggingPayloadSetup.get_usage_as_dict(response) == {} + usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (0, 0, 0) + + def test_append_system_prompt_messages(): """ Test append_system_prompt_messages prepends system message from kwargs to messages list. diff --git a/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index 9a62fcf6f0f..7461611e0fa 100644 --- a/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -148,6 +148,41 @@ class TestHostedVLLMRerankTransform: self.config._transform_response(response_dict) +@pytest.mark.parametrize( + "meta", + [ + { + "api_version": {"version": "2"}, + "tokens": {"input_tokens": 7, "output_tokens": 3}, + "billed_units": {"total_tokens": 12}, + }, + {"tokens": {"input_tokens": 0}, "billed_units": {"input_tokens": 0, "output_tokens": 0}}, + {"tokens": None, "billed_units": None}, + {}, + ], +) +@pytest.mark.parametrize("legacy_usage", [{}, {"usage": {"total_tokens": 99}}]) +def test_hosted_vllm_rerank_preserves_native_metadata_over_legacy_usage( + meta: dict[str, object], legacy_usage: dict[str, object] +) -> None: + response: Final = HostedVLLMRerankConfig()._transform_response( + {"id": "rerank-native", "results": [{"index": 0, "relevance_score": 0.9}], "meta": meta, **legacy_usage} + ) + assert response.meta == meta + assert response.model_dump() == { + "id": "rerank-native", + "results": [{"index": 0, "relevance_score": 0.9}], + "meta": meta, + } + + +def test_hosted_vllm_rerank_null_metadata_uses_legacy_usage() -> None: + response: Final = HostedVLLMRerankConfig()._transform_response( + {"results": [], "meta": None, "usage": {"total_tokens": 42}} + ) + assert response.meta == {"billed_units": {"total_tokens": 42}, "tokens": {"input_tokens": 42}} + + class TestGetOptionalRerankParamsInstruction: """`instruction` is threaded through get_optional_rerank_params only when set.""" diff --git a/tests/unit/llms/jina_ai/rerank/test_transformation.py b/tests/unit/llms/jina_ai/rerank/test_transformation.py index 9ec02a4548b..b752f7b4649 100644 --- a/tests/unit/llms/jina_ai/rerank/test_transformation.py +++ b/tests/unit/llms/jina_ai/rerank/test_transformation.py @@ -46,7 +46,10 @@ def test_transform_rerank_response_maps_results_and_usage(): ({"results": [], "usage": {}}, {"billed_units": {}, "tokens": {}}), ( {"results": [], "usage": {"total_tokens": 12, "input_tokens": 7, "output_tokens": 3, "unknown": 1}}, - {"billed_units": {"total_tokens": 12}, "tokens": {"input_tokens": 7, "output_tokens": 3}}, + { + "billed_units": {"total_tokens": 12, "input_tokens": 7, "output_tokens": 3}, + "tokens": {"input_tokens": 7, "output_tokens": 3}, + }, ), ], ) From a42c1c138a83cb15cf4586f967f105ff90c5c3e0 Mon Sep 17 00:00:00 2001 From: Graham Neubig Date: Fri, 30 Oct 2026 16:46:37 -0400 Subject: [PATCH 62/64] fix(exceptions): classify SGLang context window overflow (#48837) * fix(exceptions): classify SGLang context window overflow Co-authored-by: openhands * test(exceptions): use SGLang's real top-level error body in overflow regression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: openhands Co-authored-by: Moe Khalil Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../exception_mapping_utils.py | 1 + .../test_exception_mapping_utils.py | 61 +++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 33b6d4e1857..85c232ad3f4 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -98,6 +98,7 @@ class ExceptionCheckers: known_exception_substrings: Final = [ "exceed context limit", "this model's maximum context length is", + "requested token count exceeds the model's maximum context length", "string too long. expected a string with maximum length", "model's maximum context limit", "is longer than the model's context length", diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index efb9648cf13..2b1e4a5eed3 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -440,6 +440,67 @@ def test_gemini_context_window_error_mapping(error_message, should_raise_context ) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("streaming", [True, False]) +@pytest.mark.parametrize( + ("error_message", "expected_type"), + [ + ( + ( + "Requested token count exceeds the model's maximum context length of 131072 tokens. " + "You requested a total of 140000 tokens: 130000 tokens from the input messages and " + "10000 tokens for the completion. Please reduce the number of tokens in the input " + "messages or the completion to fit within the limit." + ), + litellm.ContextWindowExceededError, + ), + ( + ( + "requested token count exceeds the model's maximum context length of 8192 tokens. " + "you requested a total of 9000 tokens" + ), + litellm.ContextWindowExceededError, + ), + ("max_tokens must be a positive integer", litellm.BadRequestError), + ], + ids=["sglang-context-overflow", "lowercase-context-overflow", "unrelated-bad-request"], +) +@pytest.mark.asyncio +async def test_sglang_context_window_error_mapping( + sync_mode: bool, + streaming: bool, + error_message: str, + expected_type: type[litellm.BadRequestError], +) -> None: + def reject_request(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + request=request, + json={"object": "error", "message": error_message, "type": "BadRequestError", "param": None, "code": 400}, + ) + + transport: Final = httpx.MockTransport(reject_request) + if sync_mode: + with ( + httpx.Client(transport=transport) as http_client, + OpenAI(api_key="sk-test", max_retries=0, http_client=http_client) as client, + pytest.raises(litellm.BadRequestError) as caught, + ): + _call_openai_api_sync("openai/sglang-model", "chat_completion", streaming, client) + else: + async with ( + httpx.AsyncClient(transport=transport) as async_http_client, + AsyncOpenAI(api_key="sk-test", max_retries=0, http_client=async_http_client) as async_client, + ): + with pytest.raises(litellm.BadRequestError) as caught: + await _call_openai_api_async("openai/sglang-model", "chat_completion", streaming, async_client) + + assert type(caught.value) is expected_type + assert caught.value.status_code == 400 + assert error_message in str(caught.value) + assert caught.value.llm_provider == "openai" + + def test_lemonade_context_window_error_mapping(): """Lemonade's llama.cpp backend should map context overflows to LiteLLM's standard error.""" From 467847b2e9da67dd1c4ef43ff7d95ca56c013fa3 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 13:58:04 -0700 Subject: [PATCH 63/64] fix(ui): restore Z.AI credential fields in provider catalog (#49325) Carry the Z.AI catalog fix from #40088 and move its regression into the mapped unit test file. Co-authored-by: Chaitanya Laxman Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../provider_create_fields.json | 28 ++++++++++++++++++ .../public_endpoints/test_public_endpoints.py | 29 ++++++++++++++++++- 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index d849b9fbaac..01a822c5d93 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -4285,6 +4285,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "ZAI", + "provider_display_name": "Z.AI (Zhipu AI)", + "litellm_provider": "zai", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.z.ai/api/paas/v4", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "zai/glm-4.5" + }, { "provider": "ZeroGPU", "provider_display_name": "ZeroGPU", diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 3fe5345d008..72755633376 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -408,6 +408,34 @@ def test_cognition_provider_fields(): assert fields_by_key["api_base"]["required"] is False +def test_zai_provider_fields(): + from litellm.llms.zai.chat.transformation import ZAI_API_BASE + + app_instance: Final = FastAPI() + app_instance.include_router(router) + test_client: Final = TestClient(app_instance) + + response: Final = test_client.get("/public/providers/fields") + assert response.status_code == 200 + providers: Final = response.json() + + zai: Final = next((p for p in providers if p["provider"] == "ZAI"), None) + assert zai is not None, "Z.AI provider entry not found" + + assert zai["provider_display_name"] == "Z.AI (Zhipu AI)" + assert zai["litellm_provider"] == "zai" + assert zai["default_model_placeholder"].startswith("zai/") + + fields_by_key: Final = {f["key"]: f for f in zai["credential_fields"]} + + assert fields_by_key["api_key"]["required"] is True + assert fields_by_key["api_key"]["field_type"] == "password" + + assert fields_by_key["api_base"]["required"] is False + assert fields_by_key["api_base"]["field_type"] == "text" + assert fields_by_key["api_base"]["placeholder"] == ZAI_API_BASE + + def test_zerogpu_add_model_form_defaults_resolve_to_a_priced_zerogpu_model(): import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -611,7 +639,6 @@ ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset( "transcribe", "valkey", "xiaomi_mimo", - "zai", } ) From 7e37c207482dde6b17c19783e9d8aeafadaecd11 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Fri, 30 Oct 2026 14:12:51 -0700 Subject: [PATCH 64/64] fix(responses): run the guardrails a Responses WebSocket frame selects (#48725) * fix(responses): run the guardrails a Responses WebSocket frame selects * fix(responses): name the field a malformed frame guardrail selection got wrong A frame whose guardrails or guardrail_config has the wrong shape now gets an error event naming the field, instead of pydantic's raw validation dump * fix(responses): accept every guardrail selection shape a Responses WebSocket frame sends, and run during_call guardrails A frame can pick guardrails as a list of names or name-to-config dicts, inside metadata or litellm_metadata, with policies validated as a list of strings. during_call guardrails now run before the turn goes upstream, and selection fields sent next to a nested response are rejected instead of dropped * fix(proxy): keep a key's configured guardrails when policies add guardrails A policy's guardrails were merged with dict.fromkeys, so a configured guardrail entry ({name: {extra_body}}) in key metadata raised unhashable type: 'dict' and every request naming policies failed with a 500 on HTTP and on the Responses WebSocket. The merge now keeps existing entries and skips policy guardrails already selected by name. Also types the WebSocket test guardrail's hooks. * fix(responses): run streaming-scoped guardrails on socket turns, keep guardrail fields on blocks, merge guardrail lists by name * fix(responses): check frame guardrail picks against the team permission, honor picks beside a nested response, prefer a guardrail entry with settings * style(responses): format the frame guard and give the pass-through guardrail test an enterprise license * fix(proxy): keep pass-through guardrail settings, gate frame policy picks, skip moderation-only hooks on socket turns * fix(responses): keep string metadata guardrail values as client metadata on socket turns * fix(responses): ask no team permission for string metadata named like a guardrail pick on socket turns * fix(proxy): socket guardrail picks read the key's team metadata when the proxy has no database --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 43 +- litellm/proxy/auth/auth_checks.py | 81 +-- litellm/proxy/litellm_pre_call_utils.py | 67 +- .../passthrough_guardrails.py | 30 +- .../proxy/response_api_endpoints/endpoints.py | 294 ++++++++- litellm/proxy/utils.py | 6 +- litellm/responses/streaming_iterator.py | 36 +- .../integrations/test_custom_guardrail.py | 13 + ...st_auth_checks_object_access_and_lookup.py | 6 +- .../test_passthrough_guardrails.py | 19 + .../response_api_endpoints/test_endpoints.py | 615 +++++++++++++++++- .../unit/proxy/test_litellm_pre_call_utils.py | 82 +++ .../utils/helpers/test_guardrail_merge.py | 18 +- 13 files changed, 1134 insertions(+), 176 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index a9fb5c29bad..736ebd58130 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -3,8 +3,9 @@ import copy import hashlib import os import secrets -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args @@ -208,6 +209,46 @@ def guardrail_request_data_with_streaming( } +_GUARDRAIL_ENTRY_WITH_SETTINGS: Final = TypeAdapter(Mapping[str, object]) + + +def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]: + """Merges guardrail lists so each guardrail is named once: the first ``{name: settings}`` entry naming it wins over + any plain name in either order, else its first plain name stays""" + entries: Final = tuple(chain.from_iterable(entry_lists)) + settings_names: Final = tuple(_settings_entry_names(entry) for entry in entries) + named_with_settings: Final = frozenset[str]().union(*settings_names) + return [ + entry + for index, entry in enumerate(entries) + if not _superseded_guardrail_entry(entries, settings_names, named_with_settings, index) + ] + + +def guardrail_entry_names(entry: object) -> tuple[str, ...]: + if isinstance(entry, str): + return (entry,) + if isinstance(entry, dict): + return tuple(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) + return () + + +def _settings_entry_names(entry: object) -> frozenset[str]: + return frozenset(guardrail_entry_names(entry)) if isinstance(entry, dict) else frozenset() + + +def _superseded_guardrail_entry( + entries: tuple[object, ...], + settings_names: tuple[frozenset[str], ...], + named_with_settings: frozenset[str], + index: int, +) -> bool: + entry: Final = entries[index] + if isinstance(entry, str): + return entry in named_with_settings or entry in entries[:index] + return bool(settings_names[index]) and settings_names[index] <= frozenset[str]().union(*settings_names[:index]) + + def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool: if event_type in _REALTIME_STREAMING_HOOKS: return True diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 79a46ef355e..7d379a6f042 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -937,53 +937,58 @@ _GUARDRAIL_MODIFICATION_KEYS: Final[tuple] = ( ) -def _guardrail_modification_check(request_body: Mapping[str, object], team_object: LiteLLM_TeamTable | None) -> None: +def requests_guardrail_modification(request_body: Mapping[str, object]) -> bool: + """Whether a request body sets the plural ``guardrails`` list or a per-request toggle that influences whether + default-on guardrails run (``disable_global_guardrails``, ``disable_global_guardrail`` singular, and + ``opted_out_global_guardrails``), at the top level or in either metadata key""" + return ( + _requests_guardrail_modification_in(request_body.get("metadata")) + or _requests_guardrail_modification_in(request_body.get("litellm_metadata")) + or _requests_guardrail_modification_in(request_body) + ) + + +def _requests_guardrail_modification_in(container: object) -> bool: + coerced: Final = _guardrail_container_as_dict(container) + if coerced is None: + return False + return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) + + +def _guardrail_container_as_dict(container: object) -> Mapping[str, object] | None: + """Accept dict or JSON-string (from multipart/form-data or extra_body). + + Without this, an attacker can smuggle guardrail keys past the check by + sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` — + ``isinstance(dict)`` on the string returns False, the check returns + no-modification, and ``add_litellm_data_to_request`` parses the string + to a dict downstream. + """ + if isinstance(container, Mapping): + return container + if isinstance(container, str): + parsed: Final = _safe_json_loads_obj(container) + return parsed if isinstance(parsed, dict) else None + return None + + +def guardrail_modification_check(request_body: Mapping[str, object], team_object: LiteLLM_TeamTable | None) -> None: """ Reject user-supplied metadata flags that would modify guardrail behavior - unless the team has explicit permission. Checked keys include the plural - ``guardrails`` list plus the per-request toggles that influence whether - default-on guardrails run (``disable_global_guardrails``, - ``disable_global_guardrail`` singular, and ``opted_out_global_guardrails``). + unless the team has explicit permission. User-supplied values for the bypass toggles are also silently ignored by ``_get_admin_metadata`` at read time; this check adds defense in depth by failing loudly at the auth layer so operators see an explicit 403 instead of a confusing silent-ignore. """ + if requests_guardrail_modification(request_body): + require_guardrail_modification_permission(team_object) + + +def require_guardrail_modification_permission(team_object: LiteLLM_TeamTable | None) -> None: from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails - def _coerce_to_dict(container: object) -> dict | None: - """Accept dict or JSON-string (from multipart/form-data or extra_body). - - Without this, an attacker can smuggle guardrail keys past the check by - sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` — - ``isinstance(dict)`` on the string returns False, the check returns - no-modification, and ``add_litellm_data_to_request`` parses the string - to a dict downstream. - """ - if isinstance(container, dict): - return container - if isinstance(container, str): - parsed: Final = _safe_json_loads_obj(container) - return parsed if isinstance(parsed, dict) else None - return None - - def _user_requested_modification(container: object) -> bool: - coerced: Final = _coerce_to_dict(container) - if coerced is None: - return False - return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) - - # Check both metadata keys — callers can populate either depending on the - # endpoint. Cover the top-level too so root-level injection is rejected. - modifies: Final = ( - _user_requested_modification(request_body.get("metadata")) - or _user_requested_modification(request_body.get("litellm_metadata")) - or _user_requested_modification(request_body) - ) - if not modifies: - return - if not can_modify_guardrails(team_object): raise HTTPException( status_code=403, @@ -1485,7 +1490,7 @@ async def common_checks( _enforce_user_param_check(general_settings, request, request_body, route) _global_proxy_budget_check(global_proxy_spend, skip_all_budget_checks, route) - _guardrail_modification_check(_typed_request_body(request_body), team_object) + guardrail_modification_check(_typed_request_body(request_body), team_object) # 10 [OPTIONAL] Organization RBAC checks organization_role_based_access_check(user_object=user_object, route=route, request_body=request_body) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index ae9a6d4b959..a9cffcc6566 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -48,7 +48,7 @@ from litellm.constants import ( VECTOR_STORE_CALLER_METADATA_KWARG, X_LITELLM_DISABLE_CALLBACKS, ) -from litellm.integrations.custom_guardrail import without_server_streaming_classification +from litellm.integrations.custom_guardrail import merged_guardrail_entries, without_server_streaming_classification from litellm.litellm_core_utils.core_helpers import is_codex_user_agent, normalize_drop_params from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -3948,30 +3948,14 @@ def add_guardrails_from_key_or_team_metadata( """ from litellm.proxy.utils import premium_user_check - # Initialize guardrails set (avoiding duplicates) - combined_guardrails: Final = set() - - # Add key-level guardrails first - if key_metadata and "guardrails" in key_metadata: - if isinstance(key_metadata["guardrails"], list) and len(key_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(key_metadata["guardrails"]) - - # Add team-level guardrails (set automatically handles duplicates) - if team_metadata and "guardrails" in team_metadata: - if isinstance(team_metadata["guardrails"], list) and len(team_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(team_metadata["guardrails"]) - - # Add project-level guardrails (set automatically handles duplicates) - if project_metadata and "guardrails" in project_metadata: - if isinstance(project_metadata["guardrails"], list) and len(project_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(project_metadata["guardrails"]) - - # Set combined guardrails in metadata as list - if combined_guardrails: - data[metadata_variable_name]["guardrails"] = list(combined_guardrails) + configured_guardrails: Final = tuple( + metadata["guardrails"] + for metadata in (key_metadata, team_metadata, project_metadata) + if metadata and isinstance(metadata.get("guardrails"), list) and len(metadata["guardrails"]) > 0 + ) + if configured_guardrails: + premium_user_check() + data[metadata_variable_name]["guardrails"] = merged_guardrail_entries(*configured_guardrails) _add_guardrails_from_key_or_team_metadata: Final = add_guardrails_from_key_or_team_metadata @@ -4066,14 +4050,10 @@ def _add_guardrails_from_policies_in_metadata( if metadata_variable_name not in data: data[metadata_variable_name] = {} - existing_guardrails = data[metadata_variable_name].get("guardrails", []) - if not isinstance(existing_guardrails, list): - existing_guardrails = [] - - # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined: Final = set(existing_guardrails) - combined.update(resolved_guardrails) - data[metadata_variable_name]["guardrails"] = list(combined) + existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) + data[metadata_variable_name]["guardrails"] = merged_guardrail_entries( + existing_guardrails if isinstance(existing_guardrails, list) else (), resolved_guardrails + ) # Store applied policies in metadata for tracking if "applied_policies" not in data[metadata_variable_name]: @@ -4308,10 +4288,11 @@ def _apply_resolved_guardrails_to_metadata( return existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) - existing_guardrails_list: Final = existing_guardrails if isinstance(existing_guardrails, list) else [] + existing_guardrails_list: Final = ( + _REQUEST_ENTRIES.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else () + ) - # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined: Final = list(dict.fromkeys((*existing_guardrails_list, *resolved_guardrails))) + combined: Final = merged_guardrail_entries(existing_guardrails_list, resolved_guardrails) data[metadata_variable_name]["guardrails"] = combined verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) @@ -4328,7 +4309,7 @@ def _policy_match_tags(request_data: Mapping[str, object], metadata_variable_nam async def add_guardrails_from_policy_engine( - data: dict, + data: dict[str, object], # mutable-ok: pops the request's policies and records the resolved guardrails in place metadata_variable_name: str, user_api_key_dict: UserAPIKeyAuth, ) -> None: @@ -4371,11 +4352,13 @@ async def add_guardrails_from_policy_engine( all_tags: Final = _policy_match_tags(data, metadata_variable_name) _team_alias: Final = user_api_key_dict.team_alias _key_alias: Final = user_api_key_dict.key_alias - context: Final = PolicyMatchContext( - team_alias=_team_alias if isinstance(_team_alias, str) else None, - key_alias=_key_alias if isinstance(_key_alias, str) else None, - model=data.get("model"), - tags=all_tags, + context: Final = PolicyMatchContext.model_validate( + { + "team_alias": _team_alias if isinstance(_team_alias, str) else None, + "key_alias": _key_alias if isinstance(_key_alias, str) else None, + "model": data.get("model"), + "tags": all_tags, + } ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index 1c5d270542f..bd8390d704e 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -10,6 +10,7 @@ Handles guardrail execution for passthrough endpoints with: from typing import Any, Final from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import merged_guardrail_entries from litellm.proxy._types import ( PassThroughGuardrailsConfig, PassThroughGuardrailSettings, @@ -205,7 +206,7 @@ class PassthroughGuardrailHandler: def collect_guardrails( user_api_key_dict: UserAPIKeyAuth, passthrough_guardrails_config: PassThroughGuardrailsConfigInput | None, - ) -> dict[str, bool] | None: + ) -> list[object] | None: """ Collect guardrails for a passthrough endpoint. @@ -225,7 +226,7 @@ class PassthroughGuardrailHandler: passthrough_guardrails_config: List or Dict of guardrail names/settings Returns: - Dict of guardrail names to run (format: {guardrail_name: True}), or None + Guardrail entries to run, each a name or a ``{name: settings}`` entry, or None """ from litellm.proxy.litellm_pre_call_utils import ( add_guardrails_from_key_or_team_metadata, @@ -242,15 +243,6 @@ class PassthroughGuardrailHandler: verbose_proxy_logger.debug("Passthrough guardrails config is empty, skipping") return None - # Passthrough is enabled - collect guardrails - guardrails_to_run: Final[dict[str, bool]] = {} - - # Add passthrough-specific guardrails - for guardrail_name in normalized_config: - guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug("Added passthrough-specific guardrail: %s", guardrail_name) - - # Add org/team/key level guardrails using shared helper temp_data: Final[dict[str, Any]] = {"metadata": {}} add_guardrails_from_key_or_team_metadata( key_metadata=user_api_key_dict.metadata, @@ -258,20 +250,12 @@ class PassthroughGuardrailHandler: data=temp_data, metadata_variable_name="metadata", ) + inherited_guardrails: Final[list[object]] = temp_data["metadata"].get("guardrails", []) + guardrails_to_run: Final = merged_guardrail_entries(normalized_config, inherited_guardrails) - # Merge inherited guardrails into guardrails_to_run - inherited_guardrails: Final = temp_data["metadata"].get("guardrails", []) - for guardrail_name in inherited_guardrails: - if guardrail_name not in guardrails_to_run: - guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug("Added inherited guardrail (key/team level): %s", guardrail_name) + verbose_proxy_logger.debug("Collected guardrails for passthrough endpoint: %s", guardrails_to_run) - verbose_proxy_logger.debug( - "Collected guardrails for passthrough endpoint: %s", - list(guardrails_to_run.keys()), - ) - - return guardrails_to_run if guardrails_to_run else None + return guardrails_to_run or None @staticmethod def get_field_targeted_text( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index ef59deb2e5e..089cb75a795 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,11 +1,13 @@ import asyncio import contextlib +import copy import json import uuid from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from enum import Enum from functools import partial +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -35,7 +37,7 @@ from litellm.anthropic_interface.exceptions.exception_mapping_utils import Anthr from litellm.caching.dual_cache import DualCache from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES from litellm.exceptions import EndpointNotSupportedByModelError -from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.custom_guardrail import ModifyResponseException, guardrail_request_data_with_streaming from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -47,6 +49,11 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( build_blocked_response, ) from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import ( + get_team_object, + requests_guardrail_modification, + require_guardrail_modification_permission, +) from litellm.proxy.auth.auth_utils import carries_client_side_overrides from litellm.proxy.auth.user_api_key_auth import ( UserAPIKeyAuth, @@ -57,6 +64,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, cost_headers, create_response, + proxy_exception_from_route_error, resolve_counted_tool_fields, resolve_litellm_call_id, ) @@ -79,6 +87,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry from litellm.proxy.litellm_pre_call_utils import ( _CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list _UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list + add_guardrails_from_policy_engine, is_chat_completions_body, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -89,7 +98,12 @@ from litellm.proxy.utils import count_request_input_tokens, handle_exception_on_ from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD from litellm.responses.retrieve_stream_query import RETRIEVE_INCLUDE_QUERY_KEYS, RETRIEVE_STREAM_QUERY_KEYS -from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task +from litellm.responses.streaming_iterator import ( + GuardedResponsesWsFrame, + ResponsesWsFrameGuard, + ResponsesWsFrameRejection, + spawn_logging_task, +) from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -2149,11 +2163,15 @@ _PII_TOKENS: Final = TypeAdapter(dict[str, str]) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" _GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") _CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") -# the guardrail list is merged from key, team and the admission frame; a later frame cannot shrink it _CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") +_METADATA_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config") +_PROXY_ONLY_FRAME_FIELDS: Final = (*_CLIENT_GUARDRAIL_SELECTORS, "include_guardrail_response") +_UNFORWARDED_FRAME_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS, "litellm_metadata") +_OBJECTS: Final = TypeAdapter(list[object]) +_TURN_COPIED_CONTAINERS: Final = (list, dict) -class _PreCallGuardrailWalk(Protocol): +class _GuardrailWalk(Protocol): async def pre_call_hook( self, *, @@ -2163,15 +2181,49 @@ class _PreCallGuardrailWalk(Protocol): guardrails_only: bool, ) -> dict[str, object]: ... # mutable-ok: ProxyLogging.pre_call_hook returns the guardrails' own dict + async def during_call_hook( + self, + *, + data: dict[str, object], # mutable-ok: ProxyLogging.during_call_hook hands the guardrails a dict they edit + user_api_key_dict: UserAPIKeyAuth | None, + call_type: Literal["aresponses"], + guardrails_only: bool, + ) -> object: ... + + +async def _key_team_object(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + team_id: Final = user_api_key_dict.team_id + if team_id is None or team_id == UI_TEAM_ID: + return None + team_from_key: Final = LiteLLM_TeamTable(team_id=team_id, metadata=user_api_key_dict.team_metadata) + if prisma_client is None: + return team_from_key + try: + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + return team_from_key + def _responses_ws_frame_guard( admitted: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, - proxy_logging_obj: _PreCallGuardrailWalk, + proxy_logging_obj: _GuardrailWalk, record_refused_turn: Callable[[tuple[Mapping[str, object], ...], Exception], None], + *, + team_object: Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]] = _key_team_object, ) -> ResponsesWsFrameGuard: """Socket admission skips guardrails (the first frame's ``input`` is not part of the admitted request), so every - ``response.create`` frame gets the pre_call guardrail walk the HTTP ``/v1/responses`` path runs on its body.""" + ``response.create`` frame gets the pre_call and during_call guardrail walks the HTTP ``/v1/responses`` path runs on + its body. during_call runs before the frame is forwarded, since a forwarded frame cannot be recalled, and like HTTP, + which sends the model call beside during_call, the frame carries the pre_call output and not during_call rewrites.""" metadata_key: Final = get_metadata_variable_name_from_kwargs(admitted) admitted_bucket: Final = admitted.get(metadata_key) bucket: Final = { @@ -2186,32 +2238,57 @@ def _responses_ws_frame_guard( # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None}) + async def rejection( + exception: Exception, + data: dict[str, object], # mutable-ok: billed + ) -> ResponsesWsFrameRejection: + await _bill_rejected_frame(rejected_frame_auth, exception, data) + proxy_exception: Final = proxy_exception_from_route_error(exception) + return ResponsesWsFrameRejection(proxy_exception.message, proxy_exception.provider_specific_fields) + async def guard(frame: str) -> GuardedResponsesWsFrame: frame_obj: Final = _WS_FRAME.validate_json(frame) nested: Final = frame_obj.get("response") - if isinstance(nested, dict) and any(key in frame_obj for key in _GUARDED_FRAME_FIELDS): - raise ValueError( - "response.create carries input, instructions or tools both at the top level and under response" - ) body: Final = _WS_FRAME.validate_python(nested) if isinstance(nested, dict) else frame_obj - frame_data: Final = { - **{ - key: _without_untrusted_metadata(key, value) - for key, value in body.items() - if key != "type" and key not in _CLIENT_GUARDRAIL_SELECTORS + selection_sources: Final = (frame_obj, body) if isinstance(nested, dict) else (body,) + turn_data: Final = guardrail_request_data_with_streaming( + { + **{ + key: _without_untrusted_metadata(key, value) + for key, value in body.items() + if key != "type" and key not in _PROXY_ONLY_FRAME_FIELDS + }, + "model": model, + "litellm_call_id": str(uuid.uuid4()), + metadata_key: {**bucket}, }, - "model": model, - "litellm_call_id": str(uuid.uuid4()), - # shallow: the bucket carries proxy-owned objects (the OTEL parent span) that cannot be deep-copied - metadata_key: {**bucket}, - } + is_streaming=True, + ) + try: + _reject_guarded_field_beside_nested_response(frame_obj, nested) + await _check_frame_may_modify_guardrails(selection_sources, user_api_key_dict, team_object) + frame_data: Final = { + **turn_data, + metadata_key: await _metadata_with_frame_guardrails( + bucket, selection_sources, model, metadata_key, user_api_key_dict + ), + } + except Exception as e: + raise await rejection(e, turn_data) from e try: guarded_data: Final = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=frame_data, call_type="aresponses", guardrails_only=True ) + # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them + guarded_fields: Final = { + key: copy.deepcopy(guarded_data[key]) for key in _GUARDED_FRAME_FIELDS if key in guarded_data + } + pii_tokens: Final = _recorded_pii_tokens(guarded_data.get("metadata")) + await proxy_logging_obj.during_call_hook( + data=guarded_data, user_api_key_dict=user_api_key_dict, call_type="aresponses", guardrails_only=True + ) except Exception as e: - await _bill_rejected_frame(rejected_frame_auth, e, frame_data) - raise + raise await rejection(e, frame_data) from e if guarded_data.get("model") != model: rerouted: Final = ValueError( f"a guardrail rerouted this turn to {guarded_data.get('model')!r}; an open Responses WebSocket" @@ -2219,25 +2296,186 @@ def _responses_ws_frame_guard( ) await _bill_rejected_frame(rejected_frame_auth, rerouted, frame_data) raise rerouted - # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them - guarded_fields: Final = {key: guarded_data[key] for key in _GUARDED_FRAME_FIELDS if key in guarded_data} - kept: Final = {key: value for key, value in body.items() if key not in _GUARDED_FRAME_FIELDS} rewritten: Final = ( - {**frame_obj, "response": {**kept, **guarded_fields}} + {**_forwarded_frame_fields(frame_obj), "response": {**_forwarded_frame_fields(body), **guarded_fields}} if isinstance(nested, dict) - else {**kept, **guarded_fields} + else {**_forwarded_frame_fields(body), **guarded_fields} ) guardrail_information: Final = _recorded_guardrail_information(guarded_data.get(metadata_key)) return GuardedResponsesWsFrame( frame=json.dumps(rewritten), guardrail_information=guardrail_information, - pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")), + pii_tokens=pii_tokens, bill_refusal=partial(record_refused_turn, guardrail_information), ) return guard +def _reject_guarded_field_beside_nested_response(frame: Mapping[str, object], nested: object) -> None: + """Only the nested ``response`` body runs through the guardrails, so a guarded field beside it would go upstream + unguarded""" + stray: Final = ( + next((key for key in _GUARDED_FRAME_FIELDS if key in frame), None) if isinstance(nested, dict) else None + ) + if stray is not None: + raise ValueError(f"response.create sends {stray} at the top level next to response; send it inside response") + + +async def _check_frame_may_modify_guardrails( + sources: tuple[Mapping[str, object], ...], + user_api_key_dict: UserAPIKeyAuth, + team_object: Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]], +) -> None: + """A frame that picks guardrails or policies needs the team's permission to modify guardrails""" + if any( + "policies" in source or requests_guardrail_modification(_without_client_metadata(source)) for source in sources + ): + require_guardrail_modification_permission(await team_object(user_api_key_dict)) + + +def _without_client_metadata(source: Mapping[str, object]) -> Mapping[str, object]: + metadata: Final = source.get("metadata") + return {**source, "metadata": _metadata_picks("metadata", metadata)} if isinstance(metadata, dict) else source + + +def _forwarded_frame_fields(source: Mapping[str, object]) -> dict[str, object]: + return { + key: _without_guardrail_selection(value) if key == "metadata" else value + for key, value in source.items() + if key not in _UNFORWARDED_FRAME_FIELDS + } + + +class _FrameGuardrailSelection(TypedDict, total=False): + guardrails: ReadOnly[Sequence[str | Mapping[str, Mapping[str, object]]]] + guardrail_config: ReadOnly[Mapping[str, object]] + policies: ReadOnly[Sequence[str]] + + +_FRAME_GUARDRAIL_SELECTION: Final = TypeAdapter(_FrameGuardrailSelection) + + +async def _metadata_with_frame_guardrails( + bucket: Mapping[str, object], + sources: tuple[Mapping[str, object], ...], + model: object, + metadata_key: str, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + """A frame selects guardrails the way an HTTP ``/v1/responses`` body does, adding to the key, team and policy + guardrails the socket was admitted with, so a frame can never shrink that list. A nested frame selects at its top + level and in its ``response`` body, the body last so its ``guardrail_config`` wins""" + selections: Final = tuple(chain.from_iterable(_frame_guardrail_selections(source) for source in sources)) + picked_policies: Final = tuple(selection["policies"] for selection in selections if "policies" in selection) + metadata: Final = ( + await _metadata_with_frame_policies( + bucket, tuple(chain.from_iterable(picked_policies)), model, metadata_key, user_api_key_dict + ) + if picked_policies + else bucket + ) + selected_guardrails: Final = tuple(chain.from_iterable(selection.get("guardrails", ()) for selection in selections)) + selected_config: Final = dict( + chain.from_iterable( + selection["guardrail_config"].items() for selection in selections if "guardrail_config" in selection + ) + ) + existing_guardrails: Final = metadata.get("guardrails") + existing_config: Final = metadata.get("guardrail_config") + return { + **metadata, + **( + { + "guardrails": [ + *(_OBJECTS.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else ()), + *selected_guardrails, + ] + } + if selected_guardrails + else {} + ), + **( + { + "guardrail_config": { + **(_WS_FRAME.validate_python(existing_config) if isinstance(existing_config, dict) else {}), + **selected_config, + } + } + if selected_config + else {} + ), + } + + +def _frame_guardrail_selections(body: Mapping[str, object]) -> tuple[_FrameGuardrailSelection, ...]: + """Like HTTP ``/v1/responses``, a frame selects guardrails at its top level and in its ``metadata`` and + ``litellm_metadata``, the top level last so its ``guardrail_config`` wins""" + return ( + *( + _frame_guardrail_selection((key,), _metadata_picks(key, value), _METADATA_GUARDRAIL_SELECTORS) + for key, value in body.items() + if key in _CLIENT_METADATA_KEYS and isinstance(value, dict) + ), + _frame_guardrail_selection((), body, _CLIENT_GUARDRAIL_SELECTORS), + ) + + +def _metadata_picks(key: str, metadata: object) -> Mapping[str, object]: + """OpenAI ``metadata`` values are strings, so a string there is the client's own metadata and never a pick""" + validated: Final = _WS_FRAME.validate_python(metadata) + if key != "metadata": + return validated + return { + name: item + for name, item in validated.items() + if name not in _METADATA_GUARDRAIL_SELECTORS or not isinstance(item, str) + } + + +def _frame_guardrail_selection( + location: tuple[str, ...], source: Mapping[str, object], selectors: tuple[str, ...] +) -> _FrameGuardrailSelection: + try: + return _FRAME_GUARDRAIL_SELECTION.validate_python({key: source[key] for key in selectors if key in source}) + except ValidationError as e: + raise ValueError(f"response.create {first_validation_error(e, prefix=location)}") from e + + +def _without_guardrail_selection(metadata: object) -> object: + if not isinstance(metadata, dict): + return metadata + return { + name: item + for name, item in _WS_FRAME.validate_python(metadata).items() + if name not in _METADATA_GUARDRAIL_SELECTORS or isinstance(item, str) + } + + +async def _metadata_with_frame_policies( + bucket: Mapping[str, object], + policies: Sequence[str], + model: object, + metadata_key: str, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + selection: Final = { + "model": model, + "policies": list(policies), + metadata_key: {key: _turn_copy(value) for key, value in bucket.items()}, + } + await add_guardrails_from_policy_engine( + data=selection, metadata_variable_name=metadata_key, user_api_key_dict=user_api_key_dict + ) + return _WS_FRAME.validate_python(selection[metadata_key]) + + +def _turn_copy(value: object) -> object: + """Shallow, since the bucket carries proxy-owned objects (the OTEL parent span) that cannot be copied, and the + policy engine writes its lists in place""" + return copy.copy(value) if type(value) in _TURN_COPIED_CONTAINERS else value + + def _without_untrusted_metadata(key: str, value: object) -> object: """Like the HTTP path, drop the proxy-owned ``user_api_key_*`` and guardrail control keys a client may forge.""" if key not in _CLIENT_METADATA_KEYS or not isinstance(value, dict): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cc72dd03a40..0b0698faa9a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -147,6 +147,7 @@ from litellm.exceptions import ( from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, + merged_guardrail_entries, ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger @@ -3506,6 +3507,7 @@ class ProxyLogging: data: dict, user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, + guardrails_only: bool = False, ): caps: Final = ProxyLogging.callback_capabilities() if not caps.has_guardrail and not caps.has_moderation_override: @@ -3517,6 +3519,7 @@ class ProxyLogging: if ( isinstance(callback, CustomLogger) and not isinstance(callback, CustomGuardrail) + and (not guardrails_only or callback.enforces_request_content) and _overrides_moderation_hook(callback) and user_api_key_dict is not None ): @@ -9100,8 +9103,7 @@ def _merge_guardrails_with_existing(data: dict, model_level_guardrails: object) if not isinstance(model_level_guardrails, list): model_level_guardrails = [model_level_guardrails] if model_level_guardrails else [] - # Combine existing and model-level guardrails - metadata["guardrails"] = list(set(existing_guardrails + model_level_guardrails)) + metadata["guardrails"] = merged_guardrail_entries(existing_guardrails, model_level_guardrails) return modified_data diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d42f1af380e..a8e55a5e94d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2277,6 +2277,29 @@ class GuardedResponsesWsFrame: ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] +class ResponsesWsFrameRejection(ValueError): + def __init__(self, message: str, provider_specific_fields: Mapping[str, object] | None = None) -> None: + super().__init__(message) + self.provider_specific_fields: Final = provider_specific_fields + + +def _frame_rejection_event(rejection: Exception) -> str: + provider_fields: Final = ( + rejection.provider_specific_fields if isinstance(rejection, ResponsesWsFrameRejection) else None + ) + return json.dumps( + { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": str(rejection), + **({"provider_specific_fields": dict(provider_fields)} if provider_fields else {}), + }, + }, + default=str, + ) + + async def _reject_guarded_turn( websocket: ResponsesClientWebSocket, guarded: GuardedResponsesWsFrame, rejection: Exception ) -> None: @@ -2562,7 +2585,8 @@ class ResponsesWebSocketStreaming: return await self.frame_guard(message) except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event verbose_logger.info("Responses WebSocket: pre_call guardrail rejected a response.create frame: %s", e) - await self._send_error_event("invalid_request_error", str(e)) + with contextlib.suppress(Exception): + await self.websocket.send_text(_frame_rejection_event(e)) return None async def _forward_guarded_frame( @@ -2995,14 +3019,6 @@ class ResponsesWebSocketStreaming: except Exception: # noqa: BLE001, S110 # client may already be gone pass - async def _send_error_event(self, error_type: str, message: str) -> None: - try: - await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) - ) - except Exception: # noqa: BLE001, S110 # client may already be gone - pass - async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: @@ -3416,7 +3432,7 @@ class ManagedResponsesWebSocketHandler: return await self.frame_guard(raw_message) except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event verbose_logger.info("ManagedResponsesWS: pre_call guardrail rejected a response.create frame: %s", e) - await self._send_error(str(e), error_type="invalid_request_error") + await self._send_frame(_frame_rejection_event(e)) return None @staticmethod diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 0483888533e..c32c1b71f75 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3800,3 +3800,16 @@ def test_default_on_guardrail() -> None: ) is True ) + + +def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first_entry_with_settings(): + from litellm.integrations.custom_guardrail import merged_guardrail_entries + + assert merged_guardrail_entries(["a"], [{"a": {}, "b": {}}, {"a": {"x": 1}}], ["b", "c", "c"]) == [ + {"a": {}, "b": {}}, + "c", + ] + assert merged_guardrail_entries(["pii", {"pii": {"extra_body": {"threshold": 1}}}], ["model-guard"]) == [ + {"pii": {"extra_body": {"threshold": 1}}}, + "model-guard", + ] diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 6b59f6d40dc..dc7bd8cb312 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -5884,7 +5884,7 @@ async def test_team_member_budget_check_reads_from_spend_counter(): class TestGuardrailModificationCheck: - """Defense-in-depth: `_guardrail_modification_check` must 403 when the + """Defense-in-depth: `guardrail_modification_check` must 403 when the caller's metadata attempts to modify any guardrail-related key and the team lacks the `modify_guardrails` permission. Checks both the historically-covered `guardrails` list and the bypass toggles that @@ -5892,11 +5892,11 @@ class TestGuardrailModificationCheck: """ def _call(self, request_body): - from litellm.proxy.auth.auth_checks import _guardrail_modification_check + from litellm.proxy.auth.auth_checks import guardrail_modification_check team_object = MagicMock() team_object.metadata = {} # no permission - return _guardrail_modification_check(request_body=request_body, team_object=team_object) + return guardrail_modification_check(request_body=request_body, team_object=team_object) def test_noop_when_no_guardrail_keys_present(self): # no-op — should return silently diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py index 40403de846f..ef19482b30a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py @@ -73,6 +73,25 @@ class TestPassthroughGuardrailHandlerGetGuardrailNames: assert set(result) == {"pii-detection", "content-moderation"} +class TestPassthroughGuardrailHandlerCollectGuardrails: + def test_inherits_a_key_guardrail_with_its_settings(self, monkeypatch: pytest.MonkeyPatch): + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + key = UserAPIKeyAuth( + api_key="hashed", metadata={"guardrails": [{"pii": {"extra_body": {"threshold": 1}}}, "toxic"]} + ) + + guardrails = PassthroughGuardrailHandler.collect_guardrails(key, ["passthrough-guard", "toxic"]) + + assert guardrails == ["passthrough-guard", "toxic", {"pii": {"extra_body": {"threshold": 1}}}] + request_data = {"metadata": {"guardrails": guardrails}} + assert CustomGuardrail(guardrail_name="pii").get_guardrail_dynamic_request_body_params(request_data) == { + "threshold": 1 + } + + class TestPassthroughGuardrailHandlerNormalizeConfig: """Tests for PassthroughGuardrailHandler.normalize_config method.""" diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 1e4dded801e..07f85cb9628 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -7,6 +7,7 @@ import contextlib import dataclasses import datetime import json +import re import unittest from collections.abc import Callable, Mapping from typing import Any, Final, Literal, cast @@ -22,12 +23,17 @@ from fastapi.testclient import TestClient from httpx import Response import litellm +from litellm.caching.caching import DualCache from litellm.constants import EMPTY_MAPPING -from litellm.proxy._types import ProxyException +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.proxy_server import app +from litellm.proxy.utils import ProxyLogging from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterErrors +from litellm.types.utils import CallTypesLiteral from tests._master_key import MASTER_KEY @@ -2213,6 +2219,59 @@ class TestResponsesWSFirstFrameValidation: ws.close.assert_awaited_once_with(code=1011, reason="Internal server error") +class _FrameGuardrail(CustomGuardrail): + def __init__( + self, + name: str, + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, + rewrites: bool = False, + blocks: bool = False, + ): + super().__init__(guardrail_name=name, event_hook=event_hook, default_on=False) + self.inputs: list[object] = [] + self.configs: list[object] = [] + self.rewrites = rewrites + self.blocks = blocks + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + return self._scan(data) + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + data["input"] = self._scan(data)["input"] + + def _scan(self, data: dict[str, object]) -> dict[str, object]: + self.inputs.append(data["input"]) + self.configs.append(cast(dict[str, object], data["litellm_metadata"]).get("guardrail_config")) + if self.blocks: + raise HTTPException( + status_code=400, + detail={"error": f"Content blocked by {self.guardrail_name}", "guardrail_name": self.guardrail_name}, + ) + return {**data, "input": f"{data['input']} [{self.guardrail_name}]"} if self.rewrites else data + + +@contextlib.contextmanager +def _guardrails_on(*guardrails: CustomGuardrail): + with patch.object(litellm, "callbacks", list(guardrails)): + ProxyLogging._callback_capabilities_cache.clear() + try: + yield ProxyLogging(user_api_key_cache=DualCache()) + finally: + ProxyLogging._callback_capabilities_cache.clear() + + +def _key() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="hashed") + + class TestResponsesWSFirstFrameModelAuth: @pytest.mark.asyncio async def test_endpoint_enforces_auth_after_model_from_first_frame(self): @@ -2344,7 +2403,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_name": "turn"}] return {**data, "input": data["input"] + " [guarded]"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) admitted = { "model": "gpt-4o-mini", @@ -2375,7 +2434,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] raise ValueError("blocked") - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) proxy_logging_obj.post_call_failure_hook = AsyncMock() spend_logger = MagicMock(spec=ProxyDBLogger) @@ -2413,7 +2472,7 @@ class TestResponsesWSFirstFrameModelAuth: async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): return {**data, "input": data["input"] + " [guarded]", "tools": data["tools"][:1]} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2432,10 +2491,10 @@ class TestResponsesWSFirstFrameModelAuth: async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( - {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() ) metadata = { "user_api_key_metadata": {"disable_global_guardrails": True}, @@ -2454,17 +2513,18 @@ class TestResponsesWSFirstFrameModelAuth: [ {"type": "response.create", "input": "forbidden", "guardrails": []}, {"type": "response.create", "input": "forbidden", "metadata": {"guardrails": []}}, + {"type": "response.create", "input": "forbidden", "litellm_metadata": {"guardrails": []}}, {"type": "response.create", "response": {"input": "forbidden", "guardrails": [], "policies": []}}, ], ) async def test_frame_guard_keeps_the_admitted_guardrail_list_when_a_frame_tries_to_replace_it(self, frame): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, - MagicMock(), + _key(), proxy_logging_obj, MagicMock(), ) @@ -2476,13 +2536,258 @@ class TestResponsesWSFirstFrameModelAuth: assert "guardrails" not in (scanned.get("metadata") or {}) assert scanned["litellm_metadata"]["guardrails"] == ["key-required"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "selected", + [ + {"guardrails": ["frame-pick"]}, + {"guardrails": [{"frame-pick": {"extra_body": {"threshold": 0.8}}}]}, + {"metadata": {"guardrails": ["frame-pick"], "tag": "kept"}}, + {"litellm_metadata": {"guardrails": ["frame-pick"]}}, + {"response": {"guardrails": ["frame-pick"]}}, + ], + ids=["top-level", "with-extra-body", "metadata", "litellm-metadata", "nested"], + ) + async def test_frame_guard_runs_the_guardrails_a_frame_selects_on_top_of_the_admitted_ones(self, selected: dict[str, object]): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick", rewrites=True) + unselected = _FrameGuardrail("unselected") + admitted = {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}} + nested = "response" in selected + body = {"input": "hi", "store": False, **(selected["response"] if nested else selected)} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + with _guardrails_on(key_required, frame_pick, unselected) as proxy_logging_obj: + guard = _responses_ws_frame_guard(admitted, _key(), proxy_logging_obj, MagicMock()) + selected_turn = await guard(json.dumps(frame)) + plain_turn = await guard(json.dumps({"type": "response.create", "input": "again"})) + + kept_metadata = {"metadata": {"tag": "kept"}} if "metadata" in selected else {} + forwarded = {"input": "hi [frame-pick]", "store": False, **kept_metadata} + assert json.loads(selected_turn.frame) == ( + {"type": "response.create", "response": forwarded} if nested else {"type": "response.create", **forwarded} + ) + assert json.loads(plain_turn.frame) == {"type": "response.create", "input": "again"} + assert key_required.inputs == ["hi", "again"] + assert frame_pick.inputs == ["hi"] + assert unselected.inputs == [] + assert admitted["litellm_metadata"] == {"guardrails": ["key-required"]} + + @pytest.mark.asyncio + async def test_frame_guard_forwards_string_metadata_named_like_a_pick_without_running_it(self): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick") + metadata = {"guardrails": "frame-pick", "guardrail_config": "strict", "trace": "t1"} + + async def locked_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": False}}) + + with _guardrails_on(key_required, frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + team_object=locked_team, + ) + guarded = await guard(json.dumps({"type": "response.create", "input": "hi", "metadata": metadata})) + + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi", "metadata": metadata} + assert key_required.inputs == ["hi"] + assert frame_pick.inputs == [] + + @pytest.mark.asyncio + async def test_frame_guard_hands_a_selected_guardrail_its_config_with_the_top_level_winning(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + frame = { + "type": "response.create", + "input": "hi", + "metadata": {"guardrail_config": {"frame-pick": {"threshold": 1}, "metadata-only": {"threshold": 3}}}, + "guardrails": ["frame-pick"], + "guardrail_config": {"frame-pick": {"threshold": 2}}, + } + + with _guardrails_on(frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrail_config": {"key-required": {"threshold": 0}}}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard(json.dumps(frame)) + + assert frame_pick.configs == [ + {"key-required": {"threshold": 0}, "frame-pick": {"threshold": 2}, "metadata-only": {"threshold": 3}} + ] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi", "metadata": {}} + + @pytest.mark.asyncio + async def test_frame_guard_runs_streaming_scoped_guardrails_and_skips_non_streaming_ones_on_a_socket_turn(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + streaming_only = _FrameGuardrail("streaming-only") + streaming_only.apply_stream_scope("streaming") + streaming_during_call = _FrameGuardrail("streaming-during-call", event_hook=GuardrailEventHooks.during_call) + streaming_during_call.apply_stream_scope("streaming") + non_streaming_only = _FrameGuardrail("non-streaming-only") + non_streaming_only.apply_stream_scope("non_streaming") + frame = { + "type": "response.create", + "input": "hi", + "guardrails": ["streaming-only", "streaming-during-call", "non-streaming-only"], + } + + with _guardrails_on(streaming_only, streaming_during_call, non_streaming_only) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + guarded = await guard(json.dumps(frame)) + + assert streaming_only.inputs == ["hi"] + assert streaming_during_call.inputs == ["hi"] + assert non_streaming_only.inputs == [] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + + @pytest.mark.asyncio + async def test_frame_guard_blocks_a_turn_with_the_message_a_selected_during_call_guardrail_raises(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + during_pick = _FrameGuardrail("during-pick", event_hook=GuardrailEventHooks.during_call, blocks=True) + + with _guardrails_on(during_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + unselected_turn = await guard(json.dumps({"type": "response.create", "input": "plain"})) + with pytest.raises(ValueError, match=r"^Content blocked by during-pick$"): + await guard(json.dumps({"type": "response.create", "input": "hi", "guardrails": ["during-pick"]})) + + assert json.loads(unselected_turn.frame) == {"type": "response.create", "input": "plain"} + assert during_pick.inputs == ["hi"] + + @pytest.mark.asyncio + async def test_frame_guard_runs_only_the_moderation_hooks_that_enforce_request_content(self): + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + class _ModerationHook(CustomLogger): + def __init__(self, enforces_request_content: bool) -> None: + super().__init__() + self.enforces_request_content = enforces_request_content + self.inputs: list[object] = [] + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + self.inputs.append(data.get("input")) + + vendor_moderation = _ModerationHook(enforces_request_content=False) + content_enforcer = _ModerationHook(enforces_request_content=True) + + with _guardrails_on(vendor_moderation, content_enforcer) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + assert vendor_moderation.inputs == [] + assert content_enforcer.inputs == ["hi"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("selection", "message"), + [ + ( + {"guardrails": "frame-pick"}, + "response.create guardrails: 'str' instances are not allowed as a Sequence value", + ), + ({"guardrails": [123]}, "response.create guardrails.0.str: Input should be a valid string"), + ({"guardrail_config": []}, "response.create guardrail_config: Input should be a valid dictionary"), + ( + {"policies": "frame-policy"}, + "response.create policies: 'str' instances are not allowed as a Sequence value", + ), + ({"policies": [1]}, "response.create policies.0: Input should be a valid string"), + ( + {"litellm_metadata": {"guardrails": "frame-pick"}}, + "response.create litellm_metadata.guardrails: 'str' instances are not allowed as a Sequence value", + ), + ( + {"litellm_metadata": {"guardrail_config": []}}, + "response.create litellm_metadata.guardrail_config: Input should be a valid dictionary", + ), + ], + ) + async def test_frame_guard_refuses_and_bills_a_frame_whose_guardrail_selection_is_malformed( + self, selection: dict[str, object], message: str + ): + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + + with _guardrails_on(key_required, spend_logger) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + await guard(json.dumps({"type": "response.create", "input": "hi", **selection})) + + assert key_required.inputs == [] + spend_logger.async_post_call_failure_hook.assert_awaited_once() + + @pytest.mark.asyncio + async def test_frame_guard_runs_the_guardrails_of_the_policies_a_frame_selects(self): + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + key_required = _FrameGuardrail("key-required") + policy_pick = _FrameGuardrail("policy-pick") + policy_registry = get_policy_registry() + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + try: + policy_registry.add_policy("frame-policy", Policy(guardrails=PolicyGuardrails(add=["policy-pick"]))) + with _guardrails_on(key_required, policy_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard( + json.dumps({"type": "response.create", "input": "hi", "policies": ["frame-policy"]}) + ) + finally: + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + + assert key_required.inputs == ["hi"] + assert policy_pick.inputs == ["hi"] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + @pytest.mark.asyncio async def test_frame_guard_copies_an_admitted_bucket_that_holds_the_otel_parent_span(self): import threading from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) span = threading.RLock() guard = _responses_ws_frame_guard( @@ -2510,7 +2815,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] return {**data, "model": "onprem-llama"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock() @@ -2543,7 +2848,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] raise ValueError("blocked") - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock() @@ -2571,7 +2876,7 @@ class TestResponsesWSFirstFrameModelAuth: async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): return {key: value for key, value in data.items() if key != "instructions"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2596,7 +2901,7 @@ class TestResponsesWSFirstFrameModelAuth: data["metadata"]["pii_tokens"] = {"": "alice@example.com"} return {**data, "input": "email please"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2617,19 +2922,279 @@ class TestResponsesWSFirstFrameModelAuth: assert dict(guarded.pii_tokens) == {"": "alice@example.com"} @pytest.mark.asyncio - async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self): + @pytest.mark.parametrize("field", ["input", "instructions", "tools"]) + async def test_frame_guard_rejects_and_bills_a_guarded_field_beside_a_nested_response(self, field: str): + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + ) + message = f"response.create sends {field} at the top level next to response; send it inside response" + + with ( + patch.object(litellm, "callbacks", [spend_logger]), + pytest.raises(ValueError, match=f"^{re.escape(message)}$"), + ): + await guard(json.dumps({"type": "response.create", field: "unguarded", "response": {"input": "hi"}})) + + proxy_logging_obj.pre_call_hook.assert_not_awaited() + assert spend_logger.async_post_call_failure_hook.await_args.kwargs["request_data"]["input"] == "hi" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "top_level", + [ + {"guardrails": ["frame-pick"]}, + {"metadata": {"guardrails": ["frame-pick"], "tag": "kept"}}, + {"litellm_metadata": {"guardrails": ["frame-pick"]}}, + ], + ids=["top-level", "metadata", "litellm-metadata"], + ) + async def test_frame_guard_runs_the_guardrails_a_nested_frame_selects_beside_its_response(self, top_level: dict[str, object]): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick", rewrites=True) + unselected = _FrameGuardrail("unselected") + + with _guardrails_on(key_required, frame_pick, unselected) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard(json.dumps({"type": "response.create", **top_level, "response": {"input": "hi"}})) + + kept_metadata = {"metadata": {"tag": "kept"}} if "metadata" in top_level else {} + assert json.loads(guarded.frame) == { + "type": "response.create", + **kept_metadata, + "response": {"input": "hi [frame-pick]"}, + } + assert key_required.inputs == ["hi"] + assert frame_pick.inputs == ["hi"] + assert unselected.inputs == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}, + {"type": "response.create", "input": "hi", "metadata": {"guardrails": ["frame-pick"]}}, + {"type": "response.create", "input": "hi", "metadata": {"disable_global_guardrails": True}}, + {"type": "response.create", "input": "hi", "metadata": {"disable_global_guardrails": "true"}}, + {"type": "response.create", "response": {"input": "hi", "guardrails": ["frame-pick"]}}, + {"type": "response.create", "litellm_metadata": {"guardrails": ["frame-pick"]}, "response": {"input": "hi"}}, + {"type": "response.create", "input": "hi", "policies": ["frame-policy"]}, + {"type": "response.create", "response": {"input": "hi", "policies": ["frame-policy"]}}, + ], + ids=[ + "top-level", + "metadata", + "global-toggle", + "string-global-toggle", + "nested", + "beside-nested", + "policies", + "nested-policies", + ], + ) + async def test_frame_guard_refuses_guardrail_and_policy_picks_from_a_team_without_permission(self, frame: dict[str, object]): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick") + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + team_lookups: list[UserAPIKeyAuth] = [] + key = _key() + + async def locked_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + team_lookups.append(user_api_key_dict) + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": False}}) + + with _guardrails_on(key_required, frame_pick, spend_logger) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + key, + proxy_logging_obj, + MagicMock(), + team_object=locked_team, + ) + plain_turn = await guard(json.dumps({"type": "response.create", "input": "plain"})) + with pytest.raises(ValueError, match=r"^Your team does not have permission to modify guardrails\.$"): + await guard(json.dumps(frame)) + + assert json.loads(plain_turn.frame) == {"type": "response.create", "input": "plain"} + assert key_required.inputs == ["plain"] + assert frame_pick.inputs == [] + assert team_lookups == [key] + spend_logger.async_post_call_failure_hook.assert_awaited_once() + + @pytest.mark.asyncio + async def test_frame_guard_runs_a_frame_s_picks_for_a_team_allowed_to_modify_guardrails(self): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + + async def open_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": True}}) + + with _guardrails_on(frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, + _key(), + proxy_logging_obj, + MagicMock(), + team_object=open_team, + ) + await guard(json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]})) + + assert frame_pick.inputs == ["hi"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("modify_guardrails", "picked_inputs"), + [(True, ["hi"]), (False, [])], + ids=["open", "locked"], + ) + @pytest.mark.parametrize( + "team_lookup", + [ + lambda: patch( + "litellm.proxy.response_api_endpoints.endpoints.get_team_object", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist"})), + ), + lambda: patch("litellm.proxy.proxy_server.prisma_client", None), + ], + ids=["team-row-gone", "no-database"], + ) + async def test_frame_guard_reads_the_key_s_team_metadata_when_the_team_cannot_be_loaded( + self, + modify_guardrails: bool, + picked_inputs: list[str], + team_lookup: Callable[[], contextlib.AbstractContextManager[object]], + ): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + key = UserAPIKeyAuth( + api_key="hashed", + team_id="deleted-team", + team_metadata={"guardrails": {"modify_guardrails": modify_guardrails}}, + ) + frame = json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}) + + with _guardrails_on(frame_pick) as proxy_logging_obj, team_lookup(): + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, key, proxy_logging_obj, MagicMock() + ) + if modify_guardrails: + await guard(frame) + else: + with pytest.raises(ValueError, match=r"^Your team does not have permission to modify guardrails\.$"): + await guard(frame) + + assert frame_pick.inputs == picked_inputs + + @pytest.mark.asyncio + async def test_frame_guard_forwards_the_pre_call_output_and_not_a_during_call_rewrite_like_http(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + pre_pick = _FrameGuardrail("pre-pick", rewrites=True) + during_pick = _FrameGuardrail("during-pick", event_hook=GuardrailEventHooks.during_call, rewrites=True) + + with _guardrails_on(pre_pick, during_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + guarded = await guard( + json.dumps({"type": "response.create", "input": "hi", "guardrails": ["pre-pick", "during-pick"]}) + ) + + assert during_pick.inputs == ["hi [pre-pick]"] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi [pre-pick]"} + + @pytest.mark.asyncio + async def test_frame_guard_forwards_a_nested_turn_without_top_level_fields_that_pick_no_guardrail(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() ) - with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"): - await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}})) + with patch.object(litellm, "callbacks", []): + guarded = await guard( + json.dumps( + { + "type": "response.create", + "include_guardrail_response": True, + "guardrail_config": {"frame-pick": {"threshold": 1}}, + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + ) + ) - proxy_logging_obj.pre_call_hook.assert_not_awaited() + assert json.loads(guarded.frame) == { + "type": "response.create", + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + + @pytest.mark.asyncio + async def test_a_blocked_socket_turn_answers_with_the_guardrail_fields_http_returns(self): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + frame_pick = _FrameGuardrail("frame-pick", blocks=True) + websocket = MagicMock(send_text=AsyncMock()) + + with _guardrails_on(frame_pick) as proxy_logging_obj: + handler = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-haiku-5-5", + logging_obj=Logging( + model="anthropic/claude-haiku-5-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + frame_guard=_responses_ws_frame_guard( + {"model": "claude-managed", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ), + ) + await handler._process_response_create( + json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}) + ) + + assert json.loads(websocket.send_text.await_args.args[0]) == { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "Content blocked by frame-pick", + "provider_specific_fields": { + "error": "Content blocked by frame-pick", + "guardrail_name": "frame-pick", + "guardrail_mode": "pre_call", + }, + }, + } @pytest.mark.asyncio async def test_frame_guard_keeps_the_guardrail_error_when_billing_the_rejected_frame_fails(self): @@ -2637,7 +3202,7 @@ class TestResponsesWSFirstFrameModelAuth: from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=ValueError("blocked")) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock(side_effect=RuntimeError("db down")) @@ -2679,7 +3244,7 @@ class TestResponsesWSFirstFrameModelAuth: async def fake_llm_call(): return failure if provider_rejected else None - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() @@ -2768,7 +3333,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() row_written = asyncio.Event() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) proxy_logging_obj.post_call_failure_hook = AsyncMock(side_effect=lambda **kwargs: row_written.set()) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) @@ -2858,7 +3423,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: data ) @@ -2948,7 +3513,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: data ) @@ -3022,7 +3587,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: { **data, @@ -3083,7 +3648,7 @@ class TestResponsesWSFirstFrameModelAuth: rejection = litellm.RateLimitError( message="origin deployment is cooling down", model="gpt-4o-mini", llm_provider="openai" ) - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c5136ff6e4a..3194124ad1c 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5360,6 +5360,88 @@ async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_ assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"] +@pytest.mark.asyncio +async def test_add_guardrails_from_policy_engine_keeps_configured_guardrails_and_skips_ones_already_selected(): + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"guardrails": [configured, "audit"]}, + } + policy_registry = get_policy_registry() + policy_registry._policies = { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "audit", "hipaa_audit"])), + } + policy_registry._initialized = True + attachment_registry = get_attachment_registry() + attachment_registry._attachments = [PolicyAttachment(policy="baseline", scope="*")] + attachment_registry._initialized = True + + try: + await add_guardrails_from_policy_engine( + data=data, + metadata_variable_name="metadata", + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] + + +def test_add_guardrails_from_auth_metadata_lists_each_key_and_team_guardrail_once_keeping_the_key_s_settings(): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = {"model": "gpt-4", "metadata": {}} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + metadata={"guardrails": [configured, "audit"]}, + team_metadata={"guardrails": ["audit", "pii_blocker", "toxic"]}, + ) + + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_dict, data=data, metadata_variable_name="metadata" + ) + + assert data["metadata"]["guardrails"] == [configured, "audit", "toxic"] + + +def test_add_guardrails_from_auth_metadata_keeps_a_selected_guardrail_s_settings_when_a_key_policy_adds_it(): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = {"model": "gpt-4", "metadata": {"guardrails": [configured, "audit"]}} + policy_registry = get_policy_registry() + policy_registry._policies = { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "audit", "hipaa_audit"])), + } + policy_registry._initialized = True + + try: + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", metadata={"policies": ["baseline"]}), + data=data, + metadata_variable_name="metadata", + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + + assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_applies_inherited_parent_guardrail_when_child_condition_misses(): from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry diff --git a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py index e2c1a7c2c09..dd933c508cb 100644 --- a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py +++ b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py @@ -362,7 +362,17 @@ def test_merge_guardrails_with_existing_creates_metadata_when_missing(): } -def test_merge_guardrails_with_existing_raises_on_unhashable_guardrail(): - data = {"metadata": {"guardrails": [{"unhashable": True}]}} - with pytest.raises(TypeError): - _merge_guardrails_with_existing(data, ["g1"]) +def test_merge_guardrails_with_existing_keeps_a_selected_guardrail_s_settings_over_the_model_s_plain_entry(): + configured = {"pii-redact": {"extra_body": {"threshold": 0.9}}} + + result = _merge_guardrails_with_existing({"metadata": {"guardrails": [configured]}}, ["pii-redact", "toxic-filter"]) + + assert result["metadata"]["guardrails"] == [configured, "toxic-filter"] + + +def test_merge_guardrails_with_existing_keeps_a_selected_guardrail_s_settings_listed_after_the_key_s_plain_entry(): + configured = {"pii-redact": {"extra_body": {"threshold": 0.9}}} + + result = _merge_guardrails_with_existing({"metadata": {"guardrails": ["pii-redact", configured]}}, ["toxic-filter"]) + + assert result["metadata"]["guardrails"] == [configured, "toxic-filter"]