From 2ba923e18c76d053c11888454bd70feae97f5769 Mon Sep 17 00:00:00 2001 From: Acacian Date: Mon, 10 Aug 2026 22:20:32 +0900 Subject: [PATCH 01/14] fix(xai): bill from the cost xAI reports instead of recomputing it xAI states the amount it charged in usage.cost_in_usd_ticks, at 10^10 ticks to the dollar, and that figure covers tokens and every server-side tool invocation together. The xAI chat and responses transformations restate it in USD on usage.cost, the field litellm already carries a provider-stated cost in, and the xAI cost calculator bills from it the way the perplexity calculator does Routing it through usage.cost rather than a private field means the streaming chunk assembler carries it too, and no provider-neutral file has to learn about an xAI wire field Only a finite, non-negative amount is trusted, so an endpoint a caller can point litellm at cannot report a negative amount to subtract from its own recorded spend, and cannot report a NaN, which Usage stores unvalidated and which compares false against every budget threshold, disabling enforcement for the key rather than mispricing one request. Absent a usable figure nothing changes: the existing token math and the $5 per 1,000 web search calls fallback both run as before The web search surcharge is suppressed once the reported total applies, since that total already covers the search calls --- litellm/llms/xai/chat/transformation.py | 31 +++- litellm/llms/xai/common_utils.py | 24 +++ litellm/llms/xai/cost_calculator.py | 53 ++++++- litellm/llms/xai/responses/transformation.py | 51 ++++++- .../test_xai_responses_transformation.py | 74 +++++++++- .../llms/xai/test_xai_chat_transformation.py | 119 ++++++++++++++- .../llms/xai/test_xai_cost_calculator.py | 139 ++++++++++++++++++ 7 files changed, 478 insertions(+), 13 deletions(-) diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index ae5849812bf..1a8e54882f8 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -11,7 +11,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, strip_name_from_messages, ) -from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd from litellm.llms.xai.cost_calculator import ( apply_server_side_tool_usage_details_to_usage, ) @@ -30,6 +30,33 @@ from ...openai.chat.gpt_transformation import ( ) +def _adopt_cost_reported_by_xai(usage: Usage | dict[str, Any] | None) -> None: # mutable-ok: streaming dict write + """Bill what xAI charged instead of repricing the request locally. + + xAI reports the amount on ``cost_in_usd_ticks``; restate it in USD on ``cost``, + the field litellm already carries a provider stated cost in and the one + ``llms/xai/cost_calculator.py`` prices from. When xAI reported nothing usable, + ``cost`` is left alone and the request falls back to token pricing. + + Accepts a ``Usage`` (non-streaming) or a raw usage ``dict`` (streaming chunk), + matching ``_fold_reasoning_tokens_into_completion``, so both paths stay in sync. + Streaming needs the dict form because chunk aggregation rebuilds usage from the + fields it models plus ``cost``, dropping everything else xAI sent. + """ + if usage is None: + return + + if isinstance(usage, dict): + chunk_cost: Final = xai_reported_cost_in_usd(usage.get("cost_in_usd_ticks")) + if chunk_cost is not None: + usage["cost"] = chunk_cost + return + + reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) + if reported_cost is not None: + usage.cost = reported_cost + + class XAIChatConfig(OpenAIGPTConfig): @property def custom_llm_provider(self) -> str | None: @@ -283,6 +310,7 @@ class XAIChatConfig(OpenAIGPTConfig): self._fold_reasoning_tokens_into_completion(response) self._normalize_openai_compatible_usage_totals(getattr(response, "usage", None)) + _adopt_cost_reported_by_xai(getattr(response, "usage", None)) return response @staticmethod @@ -410,5 +438,6 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): if "usage" in chunk and chunk["usage"] is not None: XAIChatConfig._fold_reasoning_tokens_into_completion(chunk["usage"]) XAIChatConfig._normalize_openai_compatible_usage_totals(chunk["usage"]) + _adopt_cost_reported_by_xai(chunk["usage"]) return super().chunk_parser(chunk) diff --git a/litellm/llms/xai/common_utils.py b/litellm/llms/xai/common_utils.py index f122098332e..248b440b8f0 100644 --- a/litellm/llms/xai/common_utils.py +++ b/litellm/llms/xai/common_utils.py @@ -8,6 +8,30 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ProviderSpecificModelInfo +USD_TICKS_PER_DOLLAR: Final = 10_000_000_000 + + +def xai_reported_cost_in_usd(cost_in_usd_ticks: object) -> float | None: + """ + Convert the amount xAI says it charged into USD, or None when it reported nothing usable. + + xAI states what it billed in ``usage.cost_in_usd_ticks``, at ``USD_TICKS_PER_DOLLAR`` + ticks to the dollar: https://docs.x.ai/developers/cost-tracking + That single figure covers the whole request, tokens and every server side tool + invocation together, so whoever bills from it must not add anything on top. + + The value arrives on an untyped field of a response body that a caller able to set + api_base controls, so only the documented shape is accepted: a non-negative integer, + with bool refused since it is an int subclass. Anything else yields None and the + request is priced from tokens instead, which stops such an endpoint from reporting a + negative amount to subtract from its own recorded spend. + """ + if not isinstance(cost_in_usd_ticks, int) or isinstance(cost_in_usd_ticks, bool): + return None + if cost_in_usd_ticks < 0: + return None + return cost_in_usd_ticks / USD_TICKS_PER_DOLLAR + class XAIModelInfo(BaseLLMModelInfo): def get_provider_info( diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index dd77b8d5d09..65c642b5f3f 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -1,9 +1,11 @@ """ Helper util for handling XAI-specific cost calculation +- Prefers the cost xAI reports on the response over recomputing it locally - Uses the generic cost calculator which already handles tiered pricing correctly - Handles XAI-specific reasoning token billing (billed as part of completion tokens) """ +import math from collections.abc import Mapping from typing import TYPE_CHECKING, Final @@ -36,10 +38,42 @@ def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping usage.prompt_tokens_details = prompt_tokens_details # rebind-ok: write details onto caller usage +def _cost_reported_by_xai(usage: "Usage") -> float | None: + """ + Return what xAI billed for the request in USD, or None if it reported nothing usable. + + The xAI transformations restate ``usage.cost_in_usd_ticks`` as ``usage.cost``, the + field litellm already carries a provider stated cost in and the same one + ``llms/perplexity/cost_calculator.py`` bills from. That figure is the total for the + whole request, tokens and every server side tool invocation together, so nothing may + be added on top of it. + + A negative amount is refused rather than billed: a caller who can point litellm at an + api_base they control also controls the response body, and a negative cost would + subtract from their own recorded spend. Those requests are priced from tokens instead. + + NaN and the infinities are refused for the same reason and are the worse case, because + ``Usage`` stores a provider supplied ``cost`` without validating it and NaN compares + false against every budget threshold. Billing one would disable spend enforcement for + the key rather than mispricing a single request. + """ + reported_cost: Final[object] = getattr(usage, "cost", None) + if not isinstance(reported_cost, (int, float)) or isinstance(reported_cost, bool): + return None + if not math.isfinite(reported_cost): + return None + if reported_cost < 0: + return None + return float(reported_cost) + + def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ - Calculates the cost per token for a given XAI model, prompt tokens, and completion tokens. - Uses the generic cost calculator for all pricing logic, with XAI-specific reasoning token handling. + Prefers the amount xAI reported for the request, matching how the perplexity + calculator treats a provider-stated cost. That total is returned as completion + cost because xAI does not break it down by direction. Without one, falls back to + the generic cost calculator for all pricing logic, with XAI-specific reasoning + token handling. Input: - model: str, the model name without provider prefix @@ -48,6 +82,10 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ + reported_cost: Final = _cost_reported_by_xai(usage) + if reported_cost is not None: + return 0.0, reported_cost + # XAI-specific completion cost: completion is billed as visible + reasoning # tokens. Detect when the transformation layer already folded them so we # don't double-count; fall back to raw xAI shape for callers that bypass @@ -108,10 +146,15 @@ def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> floa """ Calculate the cost of web search requests for X.AI models. - Counts invocations from usage.server_side_tool_usage_details.web_search_calls. - Per-call rate comes from model_info.search_context_cost_per_query when set, - otherwise the default xAI tools rate ($5 / 1k calls). + When xAI reports what it billed, that figure already covers the server-side + search calls and ``cost_per_token`` has returned it, so there is nothing to add + here. Otherwise price the invocations from + usage.server_side_tool_usage_details.web_search_calls at the per-call rate + (model_info.search_context_cost_per_query when set, else the default $5 / 1k). """ + if _cost_reported_by_xai(usage) is not None: + return 0.0 + details: Final = getattr(usage, "server_side_tool_usage_details", None) if not isinstance(details, Mapping): return 0.0 diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index d79e7d4c146..045ea7dcbe9 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -1,17 +1,31 @@ -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final + +import httpx import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig -from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +from litellm.types.llms.openai import ( + ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, +) from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as _LiteLLMLoggingObj, + ) + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): """ @@ -250,6 +264,37 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return f"{api_base}/responses" + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIResponse: + """ + Bill what xAI charged instead of repricing the request locally. + + xAI reports the amount on ``usage.cost_in_usd_ticks``; restate it in USD on + ``usage.cost``, which ``ResponseAPILoggingUtils`` already copies onto the chat + Usage that ``llms/xai/cost_calculator.py`` prices, so /v1/responses bills the + reported figure the same way /v1/chat/completions does. When xAI reported nothing + usable, ``cost`` is left alone and the request falls back to token pricing. + """ + response: Final = super().transform_response_api_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) + + usage: Final = response.usage + if usage is None: + return response + + reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) + if reported_cost is not None: + usage.cost = reported_cost + + return response + def supports_native_websocket(self) -> bool: """XAI does not support native WebSocket for Responses API""" return False diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py index befd4c5ffbd..905db3d0f2a 100644 --- a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py +++ b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py @@ -7,12 +7,13 @@ transformations for the Responses API. Source: litellm/llms/xai/responses/transformation.py """ -from unittest.mock import MagicMock - +from unittest.mock import MagicMock, Mock +import httpx import pytest import litellm +from litellm.llms.xai.cost_calculator import cost_per_token from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.llms.openai import ( @@ -400,3 +401,72 @@ class TestXAIResponsesWebSearchBilling: bridged = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(event.response.usage) assert getattr(bridged, "server_side_tool_usage_details") == self._TOOL_DETAILS + + +class TestXAIResponsesReportedCost: + """xAI reports what it charged; the transformation moves it to where litellm bills from. + + ``ResponseAPILoggingUtils`` copies ``usage.cost`` onto the chat Usage that cost + tracking prices, so restating ``cost_in_usd_ticks`` there is what makes /v1/responses + bill the reported figure. At 10^10 ticks to the dollar, 37756000 ticks is $0.0037756. + """ + + @staticmethod + def _transformed_usage(usage: dict) -> ResponseAPIUsage | None: + raw_response = httpx.Response( + status_code=200, + json={ + "id": "resp_xai", + "object": "response", + "created_at": 0, + "model": "grok-4-latest", + "status": "completed", + "output": [], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": usage, + }, + ) + + response = XAIResponsesAPIConfig().transform_response_api_response( + model="grok-4-latest", + raw_response=raw_response, + logging_obj=Mock(), + ) + return response.usage + + def test_reported_cost_reaches_the_cost_calculator(self): + usage = self._transformed_usage( + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + } + ) + + assert usage.cost == 0.0037756 + + chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756) + + def test_usage_without_a_reported_cost_is_left_alone(self): + usage = self._transformed_usage( + {"input_tokens": 100, "output_tokens": 200, "total_tokens": 300} + ) + + assert usage.cost is None + + def test_negative_reported_cost_is_not_carried(self): + """A caller who can set api_base must not be able to report negative spend.""" + usage = self._transformed_usage( + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": -37756000, + } + ) + + assert usage.cost is None diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index e5e853ec82f..c67dca11a56 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -1,9 +1,14 @@ +from unittest.mock import Mock - +import httpx import pytest import litellm -from litellm.llms.xai.chat.transformation import XAIChatConfig +from litellm.llms.xai.chat.transformation import ( + XAIChatCompletionStreamingHandler, + XAIChatConfig, +) +from litellm.llms.xai.cost_calculator import cost_per_token from litellm.types.utils import ( CompletionTokensDetailsWrapper, ModelResponse, @@ -195,3 +200,113 @@ class TestXAIChatWebSearchBilling: ) assert with_search - without_search == pytest.approx(3 * 5.0 / 1000.0) + + +class TestXAIReportedCost: + """xAI reports what it charged; the transformation moves it to where litellm bills from. + + ``cost`` is the field litellm already carries a provider stated cost in, so restating + ``cost_in_usd_ticks`` there is what lets ``llms/xai/cost_calculator.py`` bill the + reported figure. At 10^10 ticks to the dollar, 37756000 ticks is $0.0037756. + """ + + @staticmethod + def _transformed_usage(usage: dict) -> Usage: + raw_response = httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-xai", + "object": "chat.completion", + "created": 0, + "model": "grok-4-latest", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": usage, + }, + ) + + response = XAIChatConfig().transform_response( + model="grok-4-latest", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + return response.usage + + def test_reported_cost_reaches_the_cost_calculator(self): + usage = self._transformed_usage( + { + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + } + ) + + assert usage.cost == 0.0037756 + assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0037756) + + def test_usage_without_a_reported_cost_is_left_alone(self): + usage = self._transformed_usage( + {"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300} + ) + + assert getattr(usage, "cost", None) is None + + def test_negative_reported_cost_is_not_carried(self): + """A caller who can set api_base must not be able to report negative spend.""" + usage = self._transformed_usage( + { + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": -37756000, + } + ) + + assert getattr(usage, "cost", None) is None + + def test_streamed_reported_cost_survives_chunk_aggregation(self): + """Streamed spend only matches if the conversion happens on the chunk. + + Chunk aggregation rebuilds usage from the fields it models plus ``cost``, so a + chunk still carrying only ``cost_in_usd_ticks`` loses the reported amount. + """ + handler = XAIChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + + parsed = handler.chunk_parser( + { + "id": "chatcmpl-xai", + "object": "chat.completion.chunk", + "created": 0, + "model": "grok-4-latest", + "choices": [], + "usage": { + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + }, + } + ) + + assert parsed.usage.cost == 0.0037756 + + assembled = litellm.stream_chunk_builder(chunks=[parsed]) + assert assembled.usage.cost == 0.0037756 + assert cost_per_token(model="grok-4-latest", usage=assembled.usage) == ( + 0.0, + 0.0037756, + ) diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 92e76fd18ab..6556adcae2d 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -361,6 +361,145 @@ class TestXAICostCalculator: response_object=object(), usage=usage ) + def test_reported_cost_is_preferred_over_token_math(self): + """The amount xAI reported, carried on usage.cost by the transformation, is billed. + + It lands entirely on completion cost because xAI does not split its total by + direction, the same shape the perplexity calculator returns. + """ + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + total_tokens=300, + cost=0.0037756, + ) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost == 0.0 + assert math.isclose(completion_cost, 0.0037756, rel_tol=1e-10) + + def test_reported_cost_suppresses_web_search_surcharge(self): + """The reported total already covers server-side tool calls. + + Without the suppression these 3 searches would be billed a second time on + top of the total xAI already charged. + """ + usage = Usage( + prompt_tokens=100, + completion_tokens=50, + total_tokens=150, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=100, + web_search_requests=3, + ), + cost=0.0037756, + ) + + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + + def test_web_search_surcharge_suppressed_through_the_dispatcher(self): + """The suppression has to hold on the path cost tracking actually uses. + + Legacy behaviour stays intact when xAI reports no cost. + """ + from litellm.llms import get_cost_for_web_search_request + + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3}) + + assert get_cost_for_web_search_request("xai", usage, {}) > 0.0 + + reported = Usage( + prompt_tokens=100, completion_tokens=50, total_tokens=150, cost=0.0037756 + ) + setattr(reported, "server_side_tool_usage_details", {"web_search_calls": 3}) + assert get_cost_for_web_search_request("xai", reported, {}) == 0.0 + + def test_no_reported_cost_falls_back_to_token_math(self): + """Absent the provider figure, nothing changes for existing callers.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + + def test_malformed_reported_cost_falls_back_to_token_math(self): + """A junk value must not fail the request, fall back to calculating.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + setattr(usage, "cost", "not-a-number") + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + + def test_boolean_reported_cost_falls_back_to_token_math(self): + """True is an int in python and would otherwise be billed as $1.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + setattr(usage, "cost", True) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + assert completion_cost != 1.0 + + def test_negative_reported_cost_is_rejected(self): + """A negative amount must never reach spend tracking. + + A caller who can set api_base controls the response body, so trusting a + negative figure would let them subtract from their own recorded spend and + slip past a budget. Fall back to token pricing instead, and keep charging + the web search surcharge, since no trustworthy total was reported. + """ + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + total_tokens=300, + cost=-0.0037756, + ) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3}) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + assert cost_per_web_search_request(usage=usage, model_info={}) > 0.0 + + def test_non_finite_reported_cost_is_rejected(self): + """NaN compares false against every budget threshold. + + Usage stores a provider supplied cost without validating it, so a caller who + controls the response body could report NaN and leave spend >= max_budget + false for the life of the key rather than mispricing one request. The + infinities are refused alongside it. Fall back to token pricing and keep + charging the web search surcharge, since no trustworthy total was reported. + """ + for reported_cost in (float("nan"), float("inf"), float("-inf")): + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + total_tokens=300, + cost=reported_cost, + ) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3}) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert math.isfinite(prompt_cost), reported_cost + assert math.isfinite(completion_cost), reported_cost + assert prompt_cost > 0.0, reported_cost + assert completion_cost > 0.0, reported_cost + assert cost_per_web_search_request(usage=usage, model_info={}) > 0.0, reported_cost + + def test_zero_reported_cost_is_honoured(self): + """A reported zero is a real answer, not a missing value.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300, cost=0.0) + + assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0) + def test_grok_4_20_beta_reasoning_cost_calculation(self): """Test cost calculation for grok-4.20-beta-0309-reasoning model.""" usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) From 1a1d459701b949ad95c7831be25d11cd63d2870a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:13:31 -0700 Subject: [PATCH 02/14] fix(xai): keep streamed and custom-priced billing inside the cost calculator Restate xAI's usage.cost_in_usd_ticks as usage.cost on chat and responses replies, streamed ones included, then let the cost calculator own the figure: a deployment with its own input_cost_per_token and output_cost_per_token keeps that price, cost margins apply on chat streams as they already did on non-streamed calls, and only OpenRouter's usage cost becomes the llm_provider-x-litellm-response-cost header, so xAI streams no longer skip the calculator through the header or the stream_chunk_builder hidden response_cost. --- litellm/cost_calculator.py | 12 ++- .../litellm_core_utils/streaming_handler.py | 17 ++-- litellm/llms/xai/chat/transformation.py | 40 +++----- litellm/llms/xai/common_utils.py | 15 +-- litellm/llms/xai/cost_calculator.py | 33 +------ litellm/llms/xai/responses/transformation.py | 51 +++++++---- litellm/main.py | 7 +- .../test_streaming_handler.py | 91 ++++++++++++++++--- .../test_xai_responses_transformation.py | 54 +++++++---- .../llms/xai/test_xai_cost_calculator.py | 45 +++++++++ tests/test_litellm/test_main.py | 19 ++++ 11 files changed, 257 insertions(+), 127 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 3adc1c25dfd..b1c55a4c280 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -4,6 +4,7 @@ import logging import time from collections.abc import Mapping, Sequence from functools import lru_cache +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response @@ -1164,6 +1165,12 @@ def _store_cost_breakdown_in_logging_obj( # Don't fail the main cost calculation if breakdown storage fails +def _without_provider_stated_cost(usage: Usage | None) -> Usage | None: + if usage is None or getattr(usage, "cost", None) is None: + return usage + return usage.model_copy(update=MappingProxyType({"cost": None})) + + def completion_cost( completion_response: object | None = None, model: str | None = None, @@ -1243,7 +1250,10 @@ def completion_cost( cache_creation_input_tokens: int | None = None cache_read_input_tokens: int | None = None audio_transcription_file_duration: float = 0.0 - cost_per_token_usage_object: Final[Usage | None] = _get_usage_object(completion_response=completion_response) + provider_usage_object: Final = _get_usage_object(completion_response=completion_response) + cost_per_token_usage_object: Final[Usage | None] = ( + _without_provider_stated_cost(provider_usage_object) if custom_pricing else provider_usage_object + ) rerank_billed_units: RerankBilledUnits | None = None # Extract service_tier from optional_params if not provided directly diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f5f671e7585..ac1ebccff8e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -54,6 +54,7 @@ FUNCTION_CALL_ATTRIBUTE: Final = "function_call" _SYNC_ITER_EXHAUSTED: Final = object() _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) +_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value}) def _next_sync_or_exhausted(it: Any) -> object: @@ -1884,8 +1885,8 @@ class CustomStreamWrapper: @staticmethod def _resolve_provider_reported_cost(usage_cost: object) -> float | None: """ - Providers report usage.cost either as a number or, for Perplexity, as a - breakdown object whose total lives under ``total_cost``. + Providers report usage.cost either as a number or as a breakdown object + whose total lives under ``total_cost``. """ if isinstance(usage_cost, bool): return None @@ -1898,12 +1899,10 @@ class CustomStreamWrapper: @staticmethod def _propagate_usage_cost_to_hidden_params( response: "ModelResponse", + custom_llm_provider: str | None, ) -> None: - """ - If the assembled response carries a provider-reported cost on - usage.cost, copy it into _hidden_params so litellm's cost - calculator uses it instead of a token-based estimate. - """ + if custom_llm_provider not in _USAGE_COST_HEADER_PROVIDERS: + return _usage: Final[Usage | None] = getattr(response, "usage", None) _cost: Final = CustomStreamWrapper._resolve_provider_reported_cost(getattr(_usage, "cost", None)) if _cost is not None: @@ -2018,7 +2017,7 @@ class CustomStreamWrapper: response = self.model_response_creator() if complete_streaming_response is not None: - self._propagate_usage_cost_to_hidden_params(complete_streaming_response) + self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) setattr( response, @@ -2268,7 +2267,7 @@ class CustomStreamWrapper: response: Final = self.model_response_creator() if complete_streaming_response is not None: - self._propagate_usage_cost_to_hidden_params(complete_streaming_response) + self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) setattr( response, diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 1a8e54882f8..09af8544b28 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -1,4 +1,5 @@ from collections.abc import AsyncIterator, Iterator, Mapping +from types import MappingProxyType from typing import Any, Final import httpx @@ -30,31 +31,11 @@ from ...openai.chat.gpt_transformation import ( ) -def _adopt_cost_reported_by_xai(usage: Usage | dict[str, Any] | None) -> None: # mutable-ok: streaming dict write - """Bill what xAI charged instead of repricing the request locally. - - xAI reports the amount on ``cost_in_usd_ticks``; restate it in USD on ``cost``, - the field litellm already carries a provider stated cost in and the one - ``llms/xai/cost_calculator.py`` prices from. When xAI reported nothing usable, - ``cost`` is left alone and the request falls back to token pricing. - - Accepts a ``Usage`` (non-streaming) or a raw usage ``dict`` (streaming chunk), - matching ``_fold_reasoning_tokens_into_completion``, so both paths stay in sync. - Streaming needs the dict form because chunk aggregation rebuilds usage from the - fields it models plus ``cost``, dropping everything else xAI sent. - """ - if usage is None: - return - - if isinstance(usage, dict): - chunk_cost: Final = xai_reported_cost_in_usd(usage.get("cost_in_usd_ticks")) - if chunk_cost is not None: - usage["cost"] = chunk_cost - return - +def _usage_restated_from_xai_ticks(usage: Usage | None) -> Usage | None: reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) - if reported_cost is not None: - usage.cost = reported_cost + if usage is None or reported_cost is None: + return None + return usage.model_copy(update=MappingProxyType({"cost": reported_cost})) class XAIChatConfig(OpenAIGPTConfig): @@ -310,7 +291,9 @@ class XAIChatConfig(OpenAIGPTConfig): self._fold_reasoning_tokens_into_completion(response) self._normalize_openai_compatible_usage_totals(getattr(response, "usage", None)) - _adopt_cost_reported_by_xai(getattr(response, "usage", None)) + restated_usage: Final = _usage_restated_from_xai_ticks(getattr(response, "usage", None)) + if restated_usage is not None: + response.usage = restated_usage return response @staticmethod @@ -438,6 +421,9 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): if "usage" in chunk and chunk["usage"] is not None: XAIChatConfig._fold_reasoning_tokens_into_completion(chunk["usage"]) XAIChatConfig._normalize_openai_compatible_usage_totals(chunk["usage"]) - _adopt_cost_reported_by_xai(chunk["usage"]) - return super().chunk_parser(chunk) + parsed_chunk: Final = super().chunk_parser(chunk) + restated_usage: Final = _usage_restated_from_xai_ticks(getattr(parsed_chunk, "usage", None)) + if restated_usage is not None: + parsed_chunk.usage = restated_usage + return parsed_chunk diff --git a/litellm/llms/xai/common_utils.py b/litellm/llms/xai/common_utils.py index 248b440b8f0..cf76a851a86 100644 --- a/litellm/llms/xai/common_utils.py +++ b/litellm/llms/xai/common_utils.py @@ -12,20 +12,7 @@ USD_TICKS_PER_DOLLAR: Final = 10_000_000_000 def xai_reported_cost_in_usd(cost_in_usd_ticks: object) -> float | None: - """ - Convert the amount xAI says it charged into USD, or None when it reported nothing usable. - - xAI states what it billed in ``usage.cost_in_usd_ticks``, at ``USD_TICKS_PER_DOLLAR`` - ticks to the dollar: https://docs.x.ai/developers/cost-tracking - That single figure covers the whole request, tokens and every server side tool - invocation together, so whoever bills from it must not add anything on top. - - The value arrives on an untyped field of a response body that a caller able to set - api_base controls, so only the documented shape is accepted: a non-negative integer, - with bool refused since it is an int subclass. Anything else yields None and the - request is priced from tokens instead, which stops such an endpoint from reporting a - negative amount to subtract from its own recorded spend. - """ + """xAI bills in ticks of a dollar: https://docs.x.ai/developers/cost-tracking""" if not isinstance(cost_in_usd_ticks, int) or isinstance(cost_in_usd_ticks, bool): return None if cost_in_usd_ticks < 0: diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index 65c642b5f3f..164568451d4 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -39,24 +39,6 @@ def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping def _cost_reported_by_xai(usage: "Usage") -> float | None: - """ - Return what xAI billed for the request in USD, or None if it reported nothing usable. - - The xAI transformations restate ``usage.cost_in_usd_ticks`` as ``usage.cost``, the - field litellm already carries a provider stated cost in and the same one - ``llms/perplexity/cost_calculator.py`` bills from. That figure is the total for the - whole request, tokens and every server side tool invocation together, so nothing may - be added on top of it. - - A negative amount is refused rather than billed: a caller who can point litellm at an - api_base they control also controls the response body, and a negative cost would - subtract from their own recorded spend. Those requests are priced from tokens instead. - - NaN and the infinities are refused for the same reason and are the worse case, because - ``Usage`` stores a provider supplied ``cost`` without validating it and NaN compares - false against every budget threshold. Billing one would disable spend enforcement for - the key rather than mispricing a single request. - """ reported_cost: Final[object] = getattr(usage, "cost", None) if not isinstance(reported_cost, (int, float)) or isinstance(reported_cost, bool): return None @@ -69,11 +51,8 @@ def _cost_reported_by_xai(usage: "Usage") -> float | None: def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ - Prefers the amount xAI reported for the request, matching how the perplexity - calculator treats a provider-stated cost. That total is returned as completion - cost because xAI does not break it down by direction. Without one, falls back to - the generic cost calculator for all pricing logic, with XAI-specific reasoning - token handling. + Calculates the cost per token for a given XAI model, prompt tokens, and completion tokens. + Uses the generic cost calculator for all pricing logic, with XAI-specific reasoning token handling. Input: - model: str, the model name without provider prefix @@ -146,11 +125,9 @@ def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> floa """ Calculate the cost of web search requests for X.AI models. - When xAI reports what it billed, that figure already covers the server-side - search calls and ``cost_per_token`` has returned it, so there is nothing to add - here. Otherwise price the invocations from - usage.server_side_tool_usage_details.web_search_calls at the per-call rate - (model_info.search_context_cost_per_query when set, else the default $5 / 1k). + Counts invocations from usage.server_side_tool_usage_details.web_search_calls. + Per-call rate comes from model_info.search_context_cost_per_query when set, + otherwise the default xAI tools rate ($5 / 1k calls). """ if _cost_reported_by_xai(usage) is not None: return 0.0 diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 045ea7dcbe9..acf03d88911 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -1,3 +1,4 @@ +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx @@ -10,8 +11,13 @@ from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfi from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, + ResponsesAPIStreamingResponse, ) from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool from litellm.types.router import GenericLiteLLMParams @@ -27,6 +33,13 @@ else: LiteLLMLoggingObj = Any +def _usage_restated_from_xai_ticks(usage: ResponseAPIUsage | None) -> ResponseAPIUsage | None: + reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) + if usage is None or reported_cost is None: + return None + return usage.model_copy(update=MappingProxyType({"cost": reported_cost})) + + class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): """ Configuration for XAI's Responses API. @@ -270,31 +283,35 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, ) -> ResponsesAPIResponse: - """ - Bill what xAI charged instead of repricing the request locally. - - xAI reports the amount on ``usage.cost_in_usd_ticks``; restate it in USD on - ``usage.cost``, which ``ResponseAPILoggingUtils`` already copies onto the chat - Usage that ``llms/xai/cost_calculator.py`` prices, so /v1/responses bills the - reported figure the same way /v1/chat/completions does. When xAI reported nothing - usable, ``cost`` is left alone and the request falls back to token pricing. - """ response: Final = super().transform_response_api_response( model=model, raw_response=raw_response, logging_obj=logging_obj, ) - usage: Final = response.usage - if usage is None: - return response - - reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) - if reported_cost is not None: - usage.cost = reported_cost - + restated_usage: Final = _usage_restated_from_xai_ticks(response.usage) + if restated_usage is not None: + response.usage = restated_usage return response + def transform_streaming_response( + self, + model: str, + parsed_chunk: dict, # mutable-ok: overrides the base class signature + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIStreamingResponse: + event: Final = super().transform_streaming_response( + model=model, + parsed_chunk=parsed_chunk, + logging_obj=logging_obj, + ) + if not isinstance(event, (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent)): + return event + restated_usage: Final = _usage_restated_from_xai_ticks(event.response.usage) + if restated_usage is not None: + event.response.usage = restated_usage + return event + def supports_native_websocket(self) -> bool: """XAI does not support native WebSocket for Responses API""" return False diff --git a/litellm/main.py b/litellm/main.py index c341db08155..724ffe76f81 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8592,10 +8592,11 @@ def stream_chunk_builder_text_completion(chunks: list, messages: list | None = N def _stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> float | None: usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None) - if isinstance(usage_cost, (int, float)): - return float(usage_cost) + numeric_usage_cost: Final = float(usage_cost) if isinstance(usage_cost, (int, float)) else None if logging_obj is not None: - return None + return numeric_usage_cost if litellm.include_cost_in_streaming_usage else None + if numeric_usage_cost is not None: + return numeric_usage_cost provider_hint: Final = response._hidden_params.get( # pyright: ignore[reportPrivateUsage] # no public accessor "custom_llm_provider" ) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 7f54fbfb4c2..7194557c5f0 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.streaming_handler import ( from litellm.types.utils import ( CompletionTokensDetailsWrapper, Delta, + ModelResponse, ModelResponseStream, PromptTokensDetailsWrapper, StandardLoggingPayload, @@ -1750,7 +1751,7 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): assert complete_response.usage.cost == 0.00025 # Use the real propagation method from CustomStreamWrapper - CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openrouter") assert "additional_headers" in complete_response._hidden_params assert ( @@ -1769,14 +1770,12 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): assert provider_cost == 0.00025 -def test_perplexity_streaming_dict_cost_propagates_to_hidden_params(): - """ - Regression: Perplexity reports usage.cost as a breakdown object, which used to - blow up the end of the stream with - `float() argument must be a string or a real number, not 'dict'`. - """ +def test_perplexity_streaming_dict_cost_bills_through_its_own_calculator(): import litellm - from litellm.cost_calculator import get_response_cost_from_hidden_params + from litellm.cost_calculator import ( + get_response_cost_from_hidden_params, + response_cost_calculator, + ) chunks = [ ModelResponseStream( @@ -1828,13 +1827,81 @@ def test_perplexity_streaming_dict_cost_propagates_to_hidden_params(): assert complete_response is not None - CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "perplexity") - assert ( - get_response_cost_from_hidden_params(complete_response._hidden_params) - == 0.00503 + assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None + assert response_cost_calculator( + response_object=complete_response, + model="perplexity/sonar", + custom_llm_provider="perplexity", + call_type="completion", + optional_params={}, + ) == pytest.approx(0.00503) + + +def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map(): + import litellm + from litellm.cost_calculator import ( + get_response_cost_from_hidden_params, + response_cost_calculator, ) + model = "openai/streams-cost-in-nanodollars" + litellm.register_model( + { + model: { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "litellm_provider": "openai", + "mode": "chat", + } + } + ) + complete_response = ModelResponse( + id="chatcmpl-openai-compatible", + model=model, + choices=[], + usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=3_144_000), + ) + + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openai") + + assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None + assert response_cost_calculator( + response_object=complete_response, + model=model, + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) == pytest.approx(2e-5) + + +def test_xai_streaming_reported_cost_still_takes_the_margin(monkeypatch): + import litellm + from litellm.cost_calculator import ( + get_response_cost_from_hidden_params, + response_cost_calculator, + ) + + complete_response = ModelResponse( + id="chatcmpl-xai", + model="grok-4-latest", + choices=[], + usage=Usage(completion_tokens=353, prompt_tokens=198, total_tokens=551, cost=0.0009956), + ) + + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "xai") + + assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": 0.5}) + assert response_cost_calculator( + response_object=complete_response, + model="xai/grok-4-latest", + custom_llm_provider="xai", + call_type="completion", + optional_params={}, + ) == pytest.approx(0.0009956 * 1.5) + def test_provider_reported_cost_ignores_unusable_shapes(): assert CustomStreamWrapper._resolve_provider_reported_cost(None) is None diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py index 905db3d0f2a..4cff5c76b9e 100644 --- a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py +++ b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py @@ -412,22 +412,22 @@ class TestXAIResponsesReportedCost: """ @staticmethod - def _transformed_usage(usage: dict) -> ResponseAPIUsage | None: - raw_response = httpx.Response( - status_code=200, - json={ - "id": "resp_xai", - "object": "response", - "created_at": 0, - "model": "grok-4-latest", - "status": "completed", - "output": [], - "parallel_tool_calls": False, - "tool_choice": "auto", - "tools": [], - "usage": usage, - }, - ) + def _response_body(usage: dict) -> dict: + return { + "id": "resp_xai", + "object": "response", + "created_at": 0, + "model": "grok-4-latest", + "status": "completed", + "output": [], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": usage, + } + + def _transformed_usage(self, usage: dict) -> ResponseAPIUsage | None: + raw_response = httpx.Response(status_code=200, json=self._response_body(usage)) response = XAIResponsesAPIConfig().transform_response_api_response( model="grok-4-latest", @@ -451,6 +451,28 @@ class TestXAIResponsesReportedCost: chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756) + def test_streamed_reported_cost_reaches_the_cost_calculator(self): + event = XAIResponsesAPIConfig().transform_streaming_response( + model="grok-4-latest", + parsed_chunk={ + "type": "response.completed", + "sequence_number": 7, + "response": self._response_body( + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + } + ), + }, + logging_obj=Mock(), + ) + + assert isinstance(event, ResponseCompletedEvent) + chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(event.response.usage) + assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756) + def test_usage_without_a_reported_cost_is_left_alone(self): usage = self._transformed_usage( {"input_tokens": 100, "output_tokens": 200, "total_tokens": 300} diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 6556adcae2d..6503e956a51 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -7,7 +7,10 @@ import os import litellm from litellm.types.utils import ( + Choices, CompletionTokensDetailsWrapper, + Message, + ModelResponse, PromptTokensDetailsWrapper, Usage, ) @@ -576,6 +579,48 @@ class TestXAICostCalculator: assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + def test_custom_pricing_beats_the_reported_cost(self): + response = ModelResponse( + id="chatcmpl-xai", + model="grok-4-latest", + choices=[Choices(index=0, message=Message(role="assistant", content="x"), finish_reason="stop")], + usage=Usage(prompt_tokens=198, completion_tokens=353, total_tokens=551, cost=0.0009956), + ) + + billed = litellm.completion_cost( + completion_response=response, + model="xai/grok-4-latest", + custom_llm_provider="xai", + custom_cost_per_token={"input_cost_per_token": 0.001, "output_cost_per_token": 0.001}, + custom_pricing=True, + ) + + assert math.isclose(billed, 0.551, rel_tol=1e-10) + + def test_deployment_custom_pricing_beats_the_reported_cost(self, monkeypatch): + deployment_id = "xai-deployment-priced-by-the-operator" + monkeypatch.setitem( + litellm.model_cost, + deployment_id, + {"input_cost_per_token": 0.001, "output_cost_per_token": 0.001, "litellm_provider": "xai", "mode": "chat"}, + ) + response = ModelResponse( + id="chatcmpl-xai", + model="grok-4-latest", + choices=[Choices(index=0, message=Message(role="assistant", content="x"), finish_reason="stop")], + usage=Usage(prompt_tokens=198, completion_tokens=353, total_tokens=551, cost=0.0009956), + ) + + billed = litellm.completion_cost( + completion_response=response, + model="xai/grok-4-latest", + custom_llm_provider="xai", + custom_pricing=True, + router_model_id=deployment_id, + ) + + assert math.isclose(billed, 0.551, rel_tol=1e-10) + class TestXAIWebSearchCostHelpers: """Focused coverage for web_search / tool-usage helpers in cost_calculator.py.""" diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 8cf878d05d9..8befa51fb9d 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3181,3 +3181,22 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( assert response is not None assert response._hidden_params.get("response_cost") is None + + +def test_stream_chunk_builder_defers_provider_reported_cost_to_logging_obj(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") + usage_chunk.usage = Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7, cost=0.42) + 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 + assert response.usage.cost == 0.42 + assert response._hidden_params.get("response_cost") is None From f81928f7aec870adc4491fd19025c16bef4a8cf6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:22:12 -0700 Subject: [PATCH 03/14] test(vector-store): accept embedding_executor in the Bedrock KB hook fake handler --- tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 98045725177..6b162ea90f1 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -370,6 +370,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters( custom_llm_provider, litellm_params, logging_obj, + embedding_executor=None, extra_headers=None, extra_body=None, timeout=None, From f77b3b2b5234b51aefc84d64dcb9a82fc3cc0f52 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:26:31 -0700 Subject: [PATCH 04/14] refactor(s3_vectors): embed search queries through the shared vector store executor S3 Vectors now subclasses BaseQueryEmbeddingVectorStoreConfig, so its query embedding runs through the Router executor with the request metadata instead of a private router lookup. embedding_model stays accepted as an alias of litellm_embedding_model. The router kwarg is gone from the search handler and every provider transform now that nothing but the executor fallback read it. --- .../azure_ai/vector_stores/transformation.py | 7 +- .../base_llm/vector_store/transformation.py | 13 +- .../bedrock/vector_stores/transformation.py | 2 - litellm/llms/custom_httpx/llm_http_handler.py | 8 - .../gemini/vector_stores/transformation.py | 2 - .../milvus/vector_stores/transformation.py | 7 +- .../openai/vector_stores/transformation.py | 2 - .../pg_vector/vector_stores/transformation.py | 2 - .../ragflow/vector_stores/transformation.py | 2 - .../vector_stores/transformation.py | 209 +++++-------- .../vector_stores/rag_api/transformation.py | 2 - .../search_api/transformation.py | 2 - litellm/vector_stores/main.py | 1 - .../test_s3_vectors_transformation.py | 282 +++++++++--------- tests/test_litellm/vector_stores/test_main.py | 25 +- 15 files changed, 241 insertions(+), 325 deletions(-) diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index db1a0fc89a3..7638229bc32 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -24,7 +24,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -121,11 +120,10 @@ class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: query_text: Final = self.query_text(query) - query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor, router) + query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor) return self._search_request( vector_store_id, query_text, @@ -145,11 +143,10 @@ class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: query_text: Final = self.query_text(query) - query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor, router) + query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor) return self._search_request( vector_store_id, query_text, diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 9624a721870..a3b8bcc499c 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -153,7 +153,6 @@ class BaseVectorStoreConfig: litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: Router | None = None, ) -> tuple[str, dict]: pass @@ -166,7 +165,6 @@ class BaseVectorStoreConfig: litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: Router | None = None, ) -> tuple[str, dict]: """ Optional async version of transform_search_vector_store_request. @@ -182,7 +180,6 @@ class BaseVectorStoreConfig: litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, extra_body=extra_body, - router=router, ) @abstractmethod @@ -271,7 +268,6 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: pass @@ -285,7 +281,6 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: return self.transform_search_vector_store_request( @@ -296,7 +291,6 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, extra_body=extra_body, - router=router, embedding_executor=embedding_executor, ) @@ -338,11 +332,10 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): query_text: str, litellm_params: Mapping[str, object], embedding_executor: VectorStoreEmbeddingExecutor | None, - router: Router | None = None, ) -> Sequence[float]: model: Final = self.query_embedding_model(litellm_params) configuration: Final = self.query_embedding_configuration(litellm_params) - executor: Final = self.query_embedding_executor(embedding_executor, router) + executor: Final = self.query_embedding_executor(embedding_executor, None) try: response: Final = executor.embed(model, query_text, configuration) except Exception as e: @@ -354,11 +347,10 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): query_text: str, litellm_params: Mapping[str, object], embedding_executor: VectorStoreEmbeddingExecutor | None, - router: Router | None = None, ) -> Sequence[float]: model: Final = self.query_embedding_model(litellm_params) configuration: Final = self.query_embedding_configuration(litellm_params) - executor: Final = self.query_embedding_executor(embedding_executor, router) + executor: Final = self.query_embedding_executor(embedding_executor, None) try: response: Final = await executor.aembed(model, query_text, configuration) except Exception as e: @@ -408,7 +400,6 @@ class BaseDirectVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, ) -> NoReturn: raise NotImplementedError("Direct vector store providers execute the search themselves; no HTTP request shape") diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index bad17a2181d..2d72db0cdba 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -27,7 +27,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -197,7 +196,6 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict]: if isinstance(query, list): query = " ".join(query) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 0f6966b0ae2..71a598a6fe7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -184,7 +184,6 @@ if TYPE_CHECKING: AnthropicMessagesStreamingResponse, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig - from litellm.router import Router from litellm.types.llms.openai_evals import ( CancelEvalResponse, CancelRunResponse, @@ -9709,7 +9708,6 @@ class BaseLLMHTTPHandler: timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - router: "Router | None" = None, ) -> VectorStoreSearchResponse: if isinstance(vector_store_provider_config, BaseDirectVectorStoreConfig): self._pre_call_direct_vector_store_search( @@ -9760,7 +9758,6 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, - router=router, embedding_executor=embedding_executor, ) else: @@ -9775,7 +9772,6 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, - router=router, ) all_optional_params: Final[dict[str, object]] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) @@ -9828,7 +9824,6 @@ class BaseLLMHTTPHandler: timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - router: "Router | None" = None, ) -> VectorStoreSearchResponse | Coroutine[object, object, VectorStoreSearchResponse]: if _is_async: return self.async_vector_store_search_handler( @@ -9844,7 +9839,6 @@ class BaseLLMHTTPHandler: extra_body=extra_body, timeout=timeout, client=client, - router=router, ) if isinstance(vector_store_provider_config, BaseDirectVectorStoreConfig): @@ -9893,7 +9887,6 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, - router=router, embedding_executor=embedding_executor, ) else: @@ -9908,7 +9901,6 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, - router=router, ) all_optional_params: Final[dict[str, object]] = dict(litellm_params) diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index 82586b1f638..f6525a449b6 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -33,7 +33,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -169,7 +168,6 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Mapping[str, object] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict]: """ Transform search request to Gemini's generateContent format. diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 4f3c366d8c1..70d3649debd 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -24,7 +24,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -129,11 +128,10 @@ class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: query_text: Final = self.query_text(query) - query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor, router) + query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor) return self._search_request( vector_store_id, query_text, @@ -153,11 +151,10 @@ class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: query_text: Final = self.query_text(query) - query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor, router) + query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor) return self._search_request( vector_store_id, query_text, diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index 4e925494039..f6c093f2e2a 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -21,7 +21,6 @@ from litellm.utils import add_openai_metadata if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -100,7 +99,6 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict]: encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url: Final = f"{api_base}/{encoded_vector_store_id}/search" diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index 9de1f589ae4..e4b06c36bf4 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -8,7 +8,6 @@ from litellm.types.vector_stores import VectorStoreSearchOptionalRequestParams if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -81,7 +80,6 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict]: encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url: Final = f"{api_base}/{encoded_vector_store_id}/search" diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index ffa6c9e1076..282cb7a92a7 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -17,7 +17,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -93,7 +92,6 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict]: """RAGFlow vector stores are management-only, search is not supported.""" raise NotImplementedError("RAGFlow vector stores support dataset management only, not search/retrieval") diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 733358381fe..a9902a0d27c 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -1,9 +1,12 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx -from litellm.caching._embedding_router import resolve_embedding_router -from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import ( + BaseQueryEmbeddingVectorStoreConfig, + VectorStoreEmbeddingExecutor, +) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -18,16 +21,18 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.router import Router else: LiteLLMLoggingObj = Any +_DEFAULT_QUERY_EMBEDDING_MODEL: Final = "text-embedding-3-small" +_DEFAULT_TOP_K: Final = 5 -class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): + +class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM): """Vector store configuration for AWS S3 Vectors.""" def __init__(self) -> None: - BaseVectorStoreConfig.__init__(self) + BaseQueryEmbeddingVectorStoreConfig.__init__(self) BaseAWSLLM.__init__(self) def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: @@ -59,141 +64,94 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): return headers def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str: - # Resolve region the same way the ingestion path does: - # dynamic param -> AWS_REGION_NAME -> AWS_REGION -> default (us-west-2) aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(litellm_params.get("aws_region_name")) return f"https://s3vectors.{aws_region_name}.api.aws" - def _resolve_query_embedding_router(self, embedding_model: str, router: "Router | None") -> "Router | None": - """Return the router iff it serves ``embedding_model`` as a deployment.""" - if router is None: - return None - model_list: Final = [ - dict(m) for m in (router.get_model_list() or ()) - ] # mutable-ok: resolve_embedding_router requires list[dict] - return resolve_embedding_router(embedding_model=embedding_model, llm_router=router, llm_model_list=model_list) + @staticmethod + def query_embedding_model(litellm_params: Mapping[str, object]) -> str: + configured: Final = litellm_params.get("litellm_embedding_model") or litellm_params.get("embedding_model") + return configured if isinstance(configured, str) and configured else _DEFAULT_QUERY_EMBEDDING_MODEL + + @staticmethod + def _query_target(vector_store_id: str, litellm_params: Mapping[str, object]) -> tuple[str, str]: + if ":" in vector_store_id: + bucket_name, index_name = vector_store_id.split(":", 1) + return bucket_name, index_name + bucket_name_from_params: Final = litellm_params.get("vector_bucket_name") + if not isinstance(bucket_name_from_params, str) or not bucket_name_from_params: + raise ValueError( + "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " + "or vector_bucket_name must be provided in litellm_params" + ) + return bucket_name_from_params, vector_store_id + + @staticmethod + def _query_request( + bucket_name: str, + index_name: str, + query_text: str, + query_vector: Sequence[float], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + ) -> tuple[str, dict[str, object]]: + litellm_logging_obj.model_call_details["query"] = query_text + return f"{api_base}/QueryVectors", { + "vectorBucketName": bucket_name, + "indexName": index_name, + "queryVector": {"float32": list(query_vector)}, + "topK": vector_store_search_optional_params.get("max_num_results", _DEFAULT_TOP_K), + "returnDistance": True, + "returnMetadata": True, + } def transform_search_vector_store_request( self, vector_store_id: str, - query: str | list[str], + query: str | Sequence[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, - ) -> tuple[str, dict]: - """Sync version - generates embedding synchronously.""" - # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name - # If not in that format, try to construct it from litellm_params - bucket_name: str - index_name: str - - if ":" in vector_store_id: - bucket_name, index_name = vector_store_id.split(":", 1) - else: - # Try to get bucket_name from litellm_params - bucket_name_from_params: Final = litellm_params.get("vector_bucket_name") - if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): - raise ValueError( - "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " - "or vector_bucket_name must be provided in litellm_params" - ) - bucket_name = bucket_name_from_params - index_name = vector_store_id - - if isinstance(query, list): - query = " ".join(query) - - # Generate embedding for the query - embedding_model: Final = litellm_params.get("embedding_model", "text-embedding-3-small") - embedding_router: Final = self._resolve_query_embedding_router(embedding_model=embedding_model, router=router) - - import litellm as litellm_module - - embedding_input: Final = [query] # mutable-ok: the embedding API takes list input - embedding_response: Final = ( - embedding_router.embedding(model=embedding_model, input=embedding_input) - if embedding_router is not None - else litellm_module.embedding(model=embedding_model, input=embedding_input) + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + bucket_name, index_name = self._query_target(vector_store_id, litellm_params) + query_text: Final = self.query_text(query) + query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor) + return self._query_request( + bucket_name, + index_name, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, ) - query_embedding: Final = embedding_response.data[0]["embedding"] - - url: Final = f"{api_base}/QueryVectors" - - request_body: Final[dict[str, Any]] = { - "vectorBucketName": bucket_name, - "indexName": index_name, - "queryVector": {"float32": query_embedding}, - "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 - "returnDistance": True, - "returnMetadata": True, - } - - litellm_logging_obj.model_call_details["query"] = query - return url, request_body async def atransform_search_vector_store_request( self, vector_store_id: str, - query: str | list[str], + query: str | Sequence[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, - ) -> tuple[str, dict]: - """Async version - generates embedding asynchronously.""" - # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name - # If not in that format, try to construct it from litellm_params - bucket_name: str - index_name: str - - if ":" in vector_store_id: - bucket_name, index_name = vector_store_id.split(":", 1) - else: - # Try to get bucket_name from litellm_params - bucket_name_from_params: Final = litellm_params.get("vector_bucket_name") - if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): - raise ValueError( - "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " - "or vector_bucket_name must be provided in litellm_params" - ) - bucket_name = bucket_name_from_params - index_name = vector_store_id - - if isinstance(query, list): - query = " ".join(query) - - # Generate embedding for the query asynchronously - embedding_model: Final = litellm_params.get("embedding_model", "text-embedding-3-small") - embedding_router: Final = self._resolve_query_embedding_router(embedding_model=embedding_model, router=router) - - import litellm as litellm_module - - embedding_input: Final = [query] # mutable-ok: the embedding API takes list input - embedding_response: Final = ( - await embedding_router.aembedding(model=embedding_model, input=embedding_input) - if embedding_router is not None - else await litellm_module.aembedding(model=embedding_model, input=embedding_input) + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + bucket_name, index_name = self._query_target(vector_store_id, litellm_params) + query_text: Final = self.query_text(query) + query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor) + return self._query_request( + bucket_name, + index_name, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, ) - query_embedding: Final = embedding_response.data[0]["embedding"] - - url: Final = f"{api_base}/QueryVectors" - - request_body: Final[dict[str, Any]] = { - "vectorBucketName": bucket_name, - "indexName": index_name, - "queryVector": {"float32": query_embedding}, - "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 - "returnDistance": True, - "returnMetadata": True, - } - - litellm_logging_obj.model_call_details["query"] = query - return url, request_body def sign_request( self, @@ -226,21 +184,13 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): if not source_text: continue - # Extract file information from metadata chunk_index = metadata.get("chunk_index", "0") file_id = f"s3-vectors-chunk-{chunk_index}" filename = metadata.get("filename", f"document-{chunk_index}") - # S3 Vectors returns distance, convert to similarity score (0-1) - # Lower distance = higher similarity - # We'll normalize using 1 / (1 + distance) to get a 0-1 score distance = item.get("distance") score = None if distance is not None: - # Convert distance to similarity score between 0 and 1 - # For cosine distance: similarity = 1 - distance - # For euclidean: use 1 / (1 + distance) - # Assuming cosine distance here score = max(0.0, min(1.0, 1.0 - float(distance))) results.append( @@ -265,7 +215,6 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): headers=response.headers, ) - # Vector store creation is not yet implemented def transform_create_vector_store_request( self, vector_store_create_optional_params, diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 36b57e7c995..5c250fc1a7e 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -21,7 +21,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -162,7 +161,6 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Mapping[str, object] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict[str, object]]: """ Transform search request for Vertex AI RAG API diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index f0812e3ed9f..0bcf16ee06f 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -25,7 +25,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -246,7 +245,6 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Mapping[str, object] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict[str, object]]: """ Transform a search request for the Vertex AI Search (Discovery Engine) API. diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 636bdd4b52e..2fe1965a192 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -482,7 +482,6 @@ def search( timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), - router=router, ) return response diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 4b58d220623..e2b02ea2151 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -1,40 +1,72 @@ +from collections.abc import Mapping from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest +from litellm.llms.base_llm.vector_store.transformation import ( + RouterVectorStoreEmbeddingExecutor, +) from litellm.llms.s3_vectors.vector_stores.transformation import ( S3VectorsVectorStoreConfig, ) +from litellm.types.utils import EmbeddingResponse from litellm.types.vector_stores import VectorStoreSearchResponse +QUERY_VECTOR = [0.1, 0.2, 0.3] -def _mock_router(model_names, sync=False): - """Router mock serving the given embedding model names.""" - router = MagicMock() - router.get_model_list.return_value = [{"model_name": name} for name in model_names] - embedding_response = Mock(data=[{"embedding": [0.1, 0.2, 0.3]}]) - if sync: - router.embedding = MagicMock(return_value=embedding_response) - else: - router.aembedding = AsyncMock(return_value=embedding_response) - return router + +def _embedding_response(vector): + return EmbeddingResponse(data=[{"embedding": vector, "index": 0, "object": "embedding"}]) + + +class _RecordingExecutor: + """Executor double recording every (model, query, configuration) it was asked to embed.""" + + def __init__(self, vector=QUERY_VECTOR): + self.vector = vector + self.calls = [] + + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + self.calls.append((model, query, dict(configuration))) + return _embedding_response(self.vector) + + async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + self.calls.append((model, query, dict(configuration))) + return _embedding_response(self.vector) + + +def _logging_obj(): + logging_obj = Mock() + logging_obj.model_call_details = {} + return logging_obj + + +def _search_kwargs(**overrides): + kwargs = { + "vector_store_id": "test-bucket:test-index", + "query": "test query", + "vector_store_search_optional_params": {}, + "api_base": "https://s3vectors.us-west-2.api.aws", + "litellm_logging_obj": _logging_obj(), + "litellm_params": {}, + "extra_body": None, + } + kwargs.update(overrides) + return kwargs class TestS3VectorsVectorStoreConfig: def test_init(self): - """Test that S3VectorsVectorStoreConfig initializes correctly""" config = S3VectorsVectorStoreConfig() assert config is not None def test_get_supported_openai_params(self): - """Test that supported OpenAI params are returned""" config = S3VectorsVectorStoreConfig() params = config.get_supported_openai_params("test-model") assert "max_num_results" in params def test_get_complete_url(self): - """Test URL generation for S3 Vectors""" config = S3VectorsVectorStoreConfig() litellm_params = {"aws_region_name": "us-west-2"} url = config.get_complete_url(None, litellm_params) @@ -57,180 +89,149 @@ class TestS3VectorsVectorStoreConfig: assert url == "https://s3vectors.eu-west-1.api.aws" def test_get_complete_url_invalid_region_format(self): - """Invalid region format raises""" config = S3VectorsVectorStoreConfig() with pytest.raises(ValueError, match="Invalid AWS region format"): config.get_complete_url(None, {"aws_region_name": "Bad_Region!"}) def test_transform_search_request(self): - """Full request-body transformation with a router-injected embedding""" + """Full request-body transformation with the query embedded through the injected executor""" config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} - router = _mock_router(["text-embedding-3-small"], sync=True) + logging_obj = _logging_obj() + executor = _RecordingExecutor() url, request_body = config.transform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={"max_num_results": 7}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={}, - extra_body=None, - router=router, + **_search_kwargs( + vector_store_search_optional_params={"max_num_results": 7}, + litellm_logging_obj=logging_obj, + embedding_executor=executor, + ) ) assert url == "https://s3vectors.us-west-2.api.aws/QueryVectors" assert request_body == { "vectorBucketName": "test-bucket", "indexName": "test-index", - "queryVector": {"float32": [0.1, 0.2, 0.3]}, + "queryVector": {"float32": QUERY_VECTOR}, "topK": 7, "returnDistance": True, "returnMetadata": True, } - assert mock_logging_obj.model_call_details["query"] == "test query" + assert executor.calls == [("text-embedding-3-small", "test query", {})] + assert logging_obj.model_call_details["query"] == "test query" + + @pytest.mark.parametrize( + ("litellm_params", "expected_model"), + [ + ({}, "text-embedding-3-small"), + ({"embedding_model": ""}, "text-embedding-3-small"), + ({"embedding_model": "my-embedding-model"}, "my-embedding-model"), + ({"litellm_embedding_model": "shared-key-model"}, "shared-key-model"), + ( + {"litellm_embedding_model": "shared-key-model", "embedding_model": "legacy-alias"}, + "shared-key-model", + ), + ], + ) + def test_query_embedding_model_accepts_embedding_model_alias(self, litellm_params, expected_model): + assert S3VectorsVectorStoreConfig.query_embedding_model(litellm_params) == expected_model @pytest.mark.asyncio - async def test_atransform_search_uses_router_for_virtual_model(self): - """Regression: router-served embedding models must resolve via the router, - not a bare litellm.aembedding call (which has no deployment credentials).""" + async def test_atransform_search_embeds_alias_and_store_config_through_executor(self): + """The store's embedding_model alias and litellm_embedding_config reach the executor unchanged""" config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} - router = _mock_router(["my-embedding-model"]) + executor = _RecordingExecutor(vector=[0.4, 0.5]) - with patch("litellm.aembedding", new=AsyncMock()) as mock_bare_aembedding: # test-quality-ok: guards that the bare-embedding path is not taken; dispatch seam is the behavior under test - url, request_body = await config.atransform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={"embedding_model": "my-embedding-model"}, - extra_body=None, - router=router, + _, request_body = await config.atransform_search_vector_store_request( + **_search_kwargs( + query=["test", "query"], + litellm_params={ + "embedding_model": "my-embedding-model", + "litellm_embedding_config": {"api_key": "store-key"}, + }, + embedding_executor=executor, ) + ) - router.aembedding.assert_awaited_once_with(model="my-embedding-model", input=["test query"]) - mock_bare_aembedding.assert_not_awaited() - assert request_body["queryVector"]["float32"] == [0.1, 0.2, 0.3] - assert request_body["topK"] == 5 # default - - @pytest.mark.asyncio - async def test_atransform_search_falls_back_when_router_does_not_serve_model(self): - """Router present but embedding_model is not a router deployment -> - bare litellm.aembedding keeps working (provider-prefixed + env creds stores).""" - config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} - router = _mock_router(["some-other-model"]) - - mock_bare = AsyncMock(return_value=Mock(data=[{"embedding": [0.4, 0.5]}])) - with patch("litellm.aembedding", new=mock_bare): # test-quality-ok: stubs the bare-embedding fallback whose request body the test asserts on - _, request_body = await config.atransform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={"embedding_model": "azure/text-embedding-3-small"}, - extra_body=None, - router=router, - ) - - mock_bare.assert_awaited_once_with(model="azure/text-embedding-3-small", input=["test query"]) - router.aembedding.assert_not_awaited() + assert executor.calls == [("my-embedding-model", "test query", {"api_key": "store-key"})] assert request_body["queryVector"]["float32"] == [0.4, 0.5] + assert request_body["topK"] == 5 @pytest.mark.asyncio - async def test_atransform_search_without_router_uses_bare_embedding(self): - """Backward compat: no router -> bare litellm.aembedding as before""" + async def test_atransform_search_router_executor_carries_request_metadata(self): + """Regression (LIT-6750): a bare Router alias resolves through the Router with the request's + team metadata on the embedding call, so the embedding is attributed to the calling key and team.""" config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} + router = MagicMock() + router.aembedding = AsyncMock(return_value=_embedding_response(QUERY_VECTOR)) + request_metadata = {"user_api_key_team_id": "team-a", "user_api_key": "hashed-key"} - mock_bare = AsyncMock(return_value=Mock(data=[{"embedding": [0.6, 0.7]}])) - with patch("litellm.aembedding", new=mock_bare): # test-quality-ok: stubs the bare-embedding fallback whose request body the test asserts on - _, request_body = await config.atransform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={}, - extra_body=None, + _, request_body = await config.atransform_search_vector_store_request( + **_search_kwargs( + litellm_params={"embedding_model": "team-embeddings"}, + embedding_executor=RouterVectorStoreEmbeddingExecutor(router=router, metadata=request_metadata), ) + ) + + router.aembedding.assert_awaited_once_with( + model="team-embeddings", input=["test query"], metadata=request_metadata + ) + assert request_body["queryVector"]["float32"] == QUERY_VECTOR + + @pytest.mark.asyncio + async def test_atransform_search_without_executor_uses_bare_embedding(self): + """Backward compat: SDK callers without an executor keep embedding through litellm.aembedding""" + config = S3VectorsVectorStoreConfig() + + mock_bare = AsyncMock(return_value=_embedding_response([0.6, 0.7])) + with patch("litellm.aembedding", new=mock_bare): # test-quality-ok: stubs the bare-embedding fallback whose request body the test asserts on + _, request_body = await config.atransform_search_vector_store_request(**_search_kwargs()) mock_bare.assert_awaited_once_with(model="text-embedding-3-small", input=["test query"]) assert request_body["queryVector"]["float32"] == [0.6, 0.7] - def test_transform_search_uses_router_for_virtual_model_sync(self): - """Sync twin: router-served embedding model resolves via router.embedding""" + def test_transform_search_without_executor_uses_bare_embedding_sync(self): + """Sync twin: no executor -> bare litellm.embedding as before""" config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} - router = _mock_router(["my-embedding-model"], sync=True) - with patch("litellm.embedding", new=MagicMock()) as mock_bare_embedding: # test-quality-ok: guards that the bare-embedding path is not taken; dispatch seam is the behavior under test - _, request_body = config.transform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={"embedding_model": "my-embedding-model"}, - extra_body=None, - router=router, - ) - - router.embedding.assert_called_once_with(model="my-embedding-model", input=["test query"]) - mock_bare_embedding.assert_not_called() - assert request_body["queryVector"]["float32"] == [0.1, 0.2, 0.3] - - def test_transform_search_without_router_uses_bare_embedding_sync(self): - """Sync twin: no router -> bare litellm.embedding as before""" - config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} - - mock_bare = MagicMock(return_value=Mock(data=[{"embedding": [0.8, 0.9]}])) + mock_bare = MagicMock(return_value=_embedding_response([0.8, 0.9])) with patch("litellm.embedding", new=mock_bare): # test-quality-ok: stubs the bare-embedding fallback whose request body the test asserts on _, request_body = config.transform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={}, - extra_body=None, + **_search_kwargs(litellm_params={"embedding_model": "my-embedding-model"}) ) - mock_bare.assert_called_once_with(model="text-embedding-3-small", input=["test query"]) + mock_bare.assert_called_once_with(model="my-embedding-model", input=["test query"]) assert request_body["queryVector"]["float32"] == [0.8, 0.9] def test_transform_search_request_invalid_vector_store_id(self): - """Test that invalid vector_store_id format raises error""" + """An unparseable vector_store_id raises before any embedding is generated""" config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} + executor = _RecordingExecutor() with pytest.raises( ValueError, match="vector_store_id must be in format 'bucket_name:index_name'", ): config.transform_search_vector_store_request( - vector_store_id="invalid-format", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={}, - extra_body=None, + **_search_kwargs(vector_store_id="invalid-format", embedding_executor=executor) ) + assert executor.calls == [] + + def test_transform_search_request_bucket_from_litellm_params(self): + config = S3VectorsVectorStoreConfig() + + _, request_body = config.transform_search_vector_store_request( + **_search_kwargs( + vector_store_id="only-index", + litellm_params={"vector_bucket_name": "params-bucket"}, + embedding_executor=_RecordingExecutor(), + ) + ) + + assert request_body["vectorBucketName"] == "params-bucket" + assert request_body["indexName"] == "only-index" + def test_transform_search_response(self): - """Test search response transformation""" config = S3VectorsVectorStoreConfig() mock_logging_obj = Mock() mock_logging_obj.model_call_details = {"query": "test query"} @@ -239,7 +240,7 @@ class TestS3VectorsVectorStoreConfig: mock_response.json.return_value = { "vectors": [ { - "distance": 0.05, # S3 Vectors returns distance, not score + "distance": 0.05, "metadata": { "source_text": "This is test content", "chunk_index": "0", @@ -258,23 +259,18 @@ class TestS3VectorsVectorStoreConfig: mock_response.status_code = 200 mock_response.headers = {} - result = config.transform_search_vector_store_response( - mock_response, mock_logging_obj - ) + result = config.transform_search_vector_store_response(mock_response, mock_logging_obj) - # VectorStoreSearchResponse is a TypedDict, so check structure instead of isinstance assert result["object"] == "vector_store.search_results.page" assert result["search_query"] == "test query" assert len(result["data"]) == 2 - # Score should be 1 - distance (cosine similarity) - assert result["data"][0]["score"] == 0.95 # 1 - 0.05 + assert result["data"][0]["score"] == 0.95 assert result["data"][0]["content"][0]["text"] == "This is test content" assert result["data"][0]["filename"] == "test.pdf" - assert result["data"][1]["score"] == 0.85 # 1 - 0.15 + assert result["data"][1]["score"] == 0.85 assert result["data"][1]["content"][0]["text"] == "More test content" def test_map_openai_params(self): - """Test OpenAI parameter mapping""" config = S3VectorsVectorStoreConfig() non_default_params = {"max_num_results": 5} optional_params = {} diff --git a/tests/test_litellm/vector_stores/test_main.py b/tests/test_litellm/vector_stores/test_main.py index d01e696906a..234e0b01094 100644 --- a/tests/test_litellm/vector_stores/test_main.py +++ b/tests/test_litellm/vector_stores/test_main.py @@ -2,14 +2,17 @@ Tests for litellm/vector_stores/main.py. Pins the router threading contract for vector store search: the router is an -explicit named parameter that reaches the HTTP handler, and it must never leak -into litellm_params/kwargs where logging would model_dump() it (the #19550 -serialization trap). +explicit named parameter that reaches the HTTP handler wrapped in the embedding +executor, and it must never leak into litellm_params/kwargs where logging would +model_dump() it (the #19550 serialization trap). """ from unittest.mock import MagicMock, patch import litellm.vector_stores.main as vector_stores_main +from litellm.llms.base_llm.vector_store.transformation import ( + RouterVectorStoreEmbeddingExecutor, +) from litellm.vector_stores.main import search MOCK_SEARCH_RESPONSE = { @@ -19,17 +22,18 @@ MOCK_SEARCH_RESPONSE = { } -def test_search_threads_router_to_handler(): - """search() must pass its router param through to the HTTP handler""" +def test_search_wraps_router_into_the_handler_embedding_executor(): + """search() hands the HTTP handler a Router-backed embedding executor carrying the + request metadata, and no bare router kwarg (LIT-6750)""" mock_router = MagicMock() logger = MagicMock() with ( - patch( # test-quality-ok: stubs provider config resolution; the seam under test is the router kwarg threading + patch( # test-quality-ok: stubs provider config resolution; the seam under test is the executor threading "litellm.vector_stores.main.ProviderConfigManager.get_provider_vector_stores_config", return_value=MagicMock(), ), - patch.object( # test-quality-ok: the handler call is the observable boundary for the router kwarg contract + patch.object( # test-quality-ok: the handler call is the observable boundary for the executor contract vector_stores_main.base_llm_http_handler, "vector_store_search_handler", return_value=MOCK_SEARCH_RESPONSE, @@ -41,11 +45,16 @@ def test_search_threads_router_to_handler(): custom_llm_provider="s3_vectors", router=mock_router, litellm_logging_obj=logger, + litellm_metadata={"user_api_key_team_id": "team-a"}, ) assert response == MOCK_SEARCH_RESPONSE mock_handler.assert_called_once() - assert mock_handler.call_args.kwargs["router"] is mock_router + assert "router" not in mock_handler.call_args.kwargs + executor = mock_handler.call_args.kwargs["embedding_executor"] + assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) + assert executor.router is mock_router + assert dict(executor.metadata) == {"user_api_key_team_id": "team-a"} def test_search_router_not_in_litellm_params(): From bcd3e2d94d3faf06cdb70c4648d872e8ddd86f5b Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:50:08 -0700 Subject: [PATCH 05/14] feat(rust-python-harness): wire existing e2e SDK tests into the matrix (#39463) Adds chat_completions and transcription as SDK function columns, backed by the existing rust_bridge test files. Adds a fourth strategy folder, existing_e2e_test_sdk, that points at already-existing live-API SDK tests (tests/ocr_tests/ as a whole folder, plus chat completion and Whisper transcription tests) instead of writing new parity tests. Extends selector_matches_node with trailing-slash folder selectors so a whole test folder can back one matrix cell. --- tests/rust-python-harness/README.md | 9 ++++++--- tests/rust-python-harness/cli.py | 9 ++++----- .../e2e_fuzz_tests/strategy.json | 4 +++- .../existing_e2e_test_sdk/README.md | 3 +++ .../existing_e2e_test_sdk/strategy.json | 14 ++++++++++++++ tests/rust-python-harness/models.py | 2 +- tests/rust-python-harness/runner.py | 2 ++ tests/rust-python-harness/ui.py | 2 +- .../unit_tests_rust/strategy.json | 4 +++- .../validate_sub_methods/strategy.json | 4 +++- tests/test_rust_python_harness.py | 16 +++++++++++++--- 11 files changed, 53 insertions(+), 16 deletions(-) create mode 100644 tests/rust-python-harness/existing_e2e_test_sdk/README.md create mode 100644 tests/rust-python-harness/existing_e2e_test_sdk/strategy.json diff --git a/tests/rust-python-harness/README.md b/tests/rust-python-harness/README.md index e94ac87c3b3..a34e1a1ab73 100644 --- a/tests/rust-python-harness/README.md +++ b/tests/rust-python-harness/README.md @@ -8,14 +8,17 @@ The matrix always has these SDK columns: - `messages / amessages` - `responses / aresponses` - `count_tokens` +- `chat_completions / acompletion` +- `transcription / atranscription` -The harness has three deliberately broad test-strategy folders: +The harness has four deliberately broad test-strategy folders: | Strategy | Folder | | --- | --- | | Public SDK parity over generated and recorded inputs | [`e2e_fuzz_tests/`](e2e_fuzz_tests/) | | Focused tests of Rust-owned behavior | [`unit_tests_rust/`](unit_tests_rust/) | | Isolated transform and Python-to-Rust helper coverage | [`validate_sub_methods/`](validate_sub_methods/) | +| Already-existing live-API SDK tests | [`existing_e2e_test_sdk/`](existing_e2e_test_sdk/) | ## Run it @@ -112,7 +115,7 @@ The initial end-to-end entries deliberately show `◐`: the repository has Rust ## Attach parity tests -Each of the three folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list: +Each of the four folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list: ```json { @@ -123,7 +126,7 @@ Each of the three folders contains a concise `README.md` and a `strategy.json`. } ``` -Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice. +Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family; a selector ending in `/` aggregates every test in that folder, recursively. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice. Use these coverage values: diff --git a/tests/rust-python-harness/cli.py b/tests/rust-python-harness/cli.py index f9e97d7ad43..c996b68846c 100644 --- a/tests/rust-python-harness/cli.py +++ b/tests/rust-python-harness/cli.py @@ -6,14 +6,13 @@ from collections.abc import Sequence from pathlib import Path from .catalog import load_catalog -from .models import HarnessCase, Strategy +from .models import SDK_FUNCTIONS, HarnessCase, Strategy from .runner import run_pytest from .ui import make_dashboard from .strategies.unit_tests.mapping_validator import FunctionReport, build_function_report REPO_ROOT = Path(__file__).resolve().parents[2] COVERAGE_ROOT = REPO_ROOT / "target" / "rust-python-harness" -SDK_FUNCTION_CHOICES = ("ocr", "messages", "responses", "count_tokens") def _parser() -> argparse.ArgumentParser: @@ -42,7 +41,7 @@ def _parser() -> argparse.ArgumentParser: action="append", default=[], dest="sdk_functions", - choices=SDK_FUNCTION_CHOICES, + choices=SDK_FUNCTIONS, help="run only this SDK function", ) parser.add_argument( @@ -110,7 +109,7 @@ def _interactive_filters(strategies: Sequence[Strategy]) -> tuple[set[str], set[ ) sdk_functions = _pick_values( "SDK functions", - [(name, name) for name in SDK_FUNCTION_CHOICES], + [(name, name) for name in SDK_FUNCTIONS], ) return strategy_ids, sdk_functions @@ -166,7 +165,7 @@ def _print_function_report(report: FunctionReport) -> None: def _validate_ledger(sdk_functions: set[str]) -> int: - functions = sdk_functions or set(SDK_FUNCTION_CHOICES) + functions = sdk_functions or set(SDK_FUNCTIONS) reports = tuple(build_function_report(function) for function in sorted(functions)) for report in reports: _print_function_report(report) diff --git a/tests/rust-python-harness/e2e_fuzz_tests/strategy.json b/tests/rust-python-harness/e2e_fuzz_tests/strategy.json index abeea01d9b5..d838486772d 100644 --- a/tests/rust-python-harness/e2e_fuzz_tests/strategy.json +++ b/tests/rust-python-harness/e2e_fuzz_tests/strategy.json @@ -7,6 +7,8 @@ "ocr": {"coverage": "partial", "selectors": ["tests/test_litellm/ocr/test_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}, "messages": {"coverage": "partial", "selectors": ["tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}, "responses": {"coverage": "partial", "selectors": ["tests/test_litellm/responses/test_rust_bridge_websocket.py"], "note": "Covers the websocket bridge; full responses parity is still being added."}, - "count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."} + "count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."}, + "chat_completions": {"coverage": "partial", "selectors": ["tests/test_litellm/rust_bridge/test_chat_completions.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}, + "transcription": {"coverage": "partial", "selectors": ["tests/test_litellm/test_audio_transcription_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."} } } diff --git a/tests/rust-python-harness/existing_e2e_test_sdk/README.md b/tests/rust-python-harness/existing_e2e_test_sdk/README.md new file mode 100644 index 00000000000..fb84f170703 --- /dev/null +++ b/tests/rust-python-harness/existing_e2e_test_sdk/README.md @@ -0,0 +1,3 @@ +# Existing e2e SDK tests + +Wires already-existing live-API SDK tests into the matrix instead of writing new parity tests. Selectors point at real test files and folders, such as `tests/ocr_tests/`, rather than individual node IDs, so future tests added to those folders are picked up automatically. diff --git a/tests/rust-python-harness/existing_e2e_test_sdk/strategy.json b/tests/rust-python-harness/existing_e2e_test_sdk/strategy.json new file mode 100644 index 00000000000..eefceea1a75 --- /dev/null +++ b/tests/rust-python-harness/existing_e2e_test_sdk/strategy.json @@ -0,0 +1,14 @@ +{ + "order": 40, + "id": "existing_e2e_test_sdk", + "label": "Existing e2e SDK tests", + "description": "Wire already-existing live-API SDK tests into the matrix instead of writing new parity tests.", + "functions": { + "ocr": {"coverage": "partial", "selectors": ["tests/ocr_tests/"], "note": "Existing live OCR provider tests; not yet a frozen Rust/Python oracle comparison."}, + "messages": {"coverage": "planned", "selectors": []}, + "responses": {"coverage": "planned", "selectors": []}, + "count_tokens": {"coverage": "planned", "selectors": []}, + "chat_completions": {"coverage": "partial", "selectors": ["tests/llm_translation/test_anthropic_completion.py", "tests/llm_translation/test_bedrock_completion.py"], "note": "Existing live chat completion tests for providers with confirmed Rust bridge regressions."}, + "transcription": {"coverage": "partial", "selectors": ["tests/audio_tests/test_whisper.py"], "note": "Existing live Whisper transcription test."} + } +} diff --git a/tests/rust-python-harness/models.py b/tests/rust-python-harness/models.py index 21097e0f7d0..a02684dc63d 100644 --- a/tests/rust-python-harness/models.py +++ b/tests/rust-python-harness/models.py @@ -33,7 +33,7 @@ class ConfidenceLevel(str, Enum): LOW = "LOW" -SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens") +SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens", "chat_completions", "transcription") @dataclass(frozen=True) diff --git a/tests/rust-python-harness/runner.py b/tests/rust-python-harness/runner.py index 82393ef234e..286f2c4116d 100644 --- a/tests/rust-python-harness/runner.py +++ b/tests/rust-python-harness/runner.py @@ -15,6 +15,8 @@ UpdateCallback = Callable[[HarnessRun], None] def selector_matches_node(selector: str, nodeid: str) -> bool: normalized_selector = selector.replace("\\", "/") normalized_nodeid = nodeid.replace("\\", "/") + if normalized_selector.endswith("/"): + return normalized_nodeid.startswith(normalized_selector) if "::" in normalized_selector: return normalized_nodeid == normalized_selector or normalized_nodeid.startswith( f"{normalized_selector}[" diff --git a/tests/rust-python-harness/ui.py b/tests/rust-python-harness/ui.py index 57fedd17fa6..bf143a4e100 100644 --- a/tests/rust-python-harness/ui.py +++ b/tests/rust-python-harness/ui.py @@ -119,7 +119,7 @@ class RichDashboard(AbstractContextManager["RichDashboard"]): table = Table(box=box.ROUNDED, expand=True, title="Strategy × SDK function") table.add_column("Strategy", ratio=3) - for label in ("ocr/aocr", "messages", "responses", "count_tokens"): + for label in SDK_FUNCTIONS: table.add_column(label, justify="center", ratio=1) for strategy in self.strategies: cells = [] diff --git a/tests/rust-python-harness/unit_tests_rust/strategy.json b/tests/rust-python-harness/unit_tests_rust/strategy.json index 89e897c872d..bfb2bc0dad0 100644 --- a/tests/rust-python-harness/unit_tests_rust/strategy.json +++ b/tests/rust-python-harness/unit_tests_rust/strategy.json @@ -7,6 +7,8 @@ "ocr": {"coverage": "planned", "selectors": []}, "messages": {"coverage": "planned", "selectors": []}, "responses": {"coverage": "planned", "selectors": []}, - "count_tokens": {"coverage": "planned", "selectors": []} + "count_tokens": {"coverage": "planned", "selectors": []}, + "chat_completions": {"coverage": "planned", "selectors": []}, + "transcription": {"coverage": "planned", "selectors": []} } } diff --git a/tests/rust-python-harness/validate_sub_methods/strategy.json b/tests/rust-python-harness/validate_sub_methods/strategy.json index 6e6381678e0..a26bc2d70ed 100644 --- a/tests/rust-python-harness/validate_sub_methods/strategy.json +++ b/tests/rust-python-harness/validate_sub_methods/strategy.json @@ -7,6 +7,8 @@ "ocr": {"coverage": "planned", "selectors": []}, "messages": {"coverage": "planned", "selectors": []}, "responses": {"coverage": "planned", "selectors": []}, - "count_tokens": {"coverage": "planned", "selectors": []} + "count_tokens": {"coverage": "planned", "selectors": []}, + "chat_completions": {"coverage": "planned", "selectors": []}, + "transcription": {"coverage": "planned", "selectors": []} } } diff --git a/tests/test_rust_python_harness.py b/tests/test_rust_python_harness.py index 6a8fa8d35cc..514446577fd 100644 --- a/tests/test_rust_python_harness.py +++ b/tests/test_rust_python_harness.py @@ -66,13 +66,14 @@ def _manifest() -> dict[str, object]: } -def test_should_load_the_three_harness_strategies_in_order() -> None: +def test_should_load_the_four_harness_strategies_in_order() -> None: strategies = load_catalog() assert [strategy.id for strategy in strategies] == [ "e2e_fuzz_tests", "unit_tests_rust", "validate_sub_methods", + "existing_e2e_test_sdk", ] assert all( tuple(case.sdk_function for case in strategy.cases) == SDK_FUNCTIONS @@ -104,6 +105,8 @@ def test_should_reject_a_manifest_missing_an_sdk_function(tmp_path: Path) -> Non True, ), ("tests/test_parity.py::test_one", "tests/test_parity.py::test_two", False), + ("tests/ocr_tests/", "tests/ocr_tests/test_ocr_mistral.py::test_one", True), + ("tests/ocr_tests/", "tests/other_tests/test_ocr_mistral.py::test_one", False), ], ) def test_should_match_pytest_file_and_node_selectors( @@ -123,6 +126,13 @@ def test_should_only_return_selectors_whose_files_exist(tmp_path: Path) -> None: assert runnable_selectors((case,), tmp_path) == ("tests/test_parity.py",) +def test_should_treat_an_existing_folder_selector_as_runnable(tmp_path: Path) -> None: + (tmp_path / "tests" / "ocr_tests").mkdir(parents=True) + case = _case(selectors=("tests/ocr_tests/",)) + + assert runnable_selectors((case,), tmp_path) == ("tests/ocr_tests/",) + + def test_should_mark_planned_and_not_applicable_cases_without_running() -> None: planned = CaseResult(case=_case(coverage=Coverage.PLANNED)) not_applicable = CaseResult(case=_case(coverage=Coverage.NOT_APPLICABLE)) @@ -239,8 +249,8 @@ def test_should_report_confidence_for_each_sdk_section() -> None: } assert scores["responses"].verified_strategies == 1 - assert scores["responses"].required_strategies == 3 - assert scores["responses"].percentage == 33 + assert scores["responses"].required_strategies == 4 + assert scores["responses"].percentage == 25 assert scores["responses"].level.value == "MEDIUM" assert scores["count_tokens"].percentage == 0 assert scores["count_tokens"].level.value == "LOW" From 099d26204f1998f3c8181e7eb3581ac097a59d38 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 2 Sep 2026 19:54:09 -0700 Subject: [PATCH 06/14] fix(ui): read the preset catalog at runtime in the vitest mock The autoRouterPresets mock imported litellm/proxy/public_endpoints/autorouter_presets.json as a module. That path sits outside ui/litellm-dashboard, the only directory the UI Dockerfile copies, so `next build` type-checking inside the image failed with "Cannot find module" and the ui-image job went red on every PR that touched an image-scan path. Read the file with fs at runtime instead; vitest still derives expectations from the real bundled catalog. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HqEPCNLDrssxsuezhAaL4j --- .../tests/mocks/autoRouterPresets.ts | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts index f73faaa70fe..7356d92d99e 100644 --- a/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts +++ b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts @@ -1,10 +1,18 @@ +import { readFileSync } from "fs"; +import { resolve } from "path"; import { vi } from "vitest"; -import bundledPresets from "../../../../litellm/proxy/public_endpoints/autorouter_presets.json"; import { hydratePresets, type AutoRouterPresetsResponse } from "@/lib/autorouter_presets"; // Derived from the real bundled catalog so a preset edit there flows into test expectations -// instead of redding on a stale copy. Exported as vi.fn so a test can override the query state. -export const BUNDLED_PRESETS = hydratePresets(bundledPresets as AutoRouterPresetsResponse); +// instead of redding on a stale copy. Read at runtime rather than imported as a module: the +// catalog lives outside ui/litellm-dashboard, so a module import fails `next build`'s type +// check inside the UI Docker image, whose build context is only this package. +// Exported as vi.fn so a test can override the query state. +const CATALOG_PATH = resolve(process.cwd(), "../../litellm/proxy/public_endpoints/autorouter_presets.json"); + +export const BUNDLED_PRESETS = hydratePresets( + JSON.parse(readFileSync(CATALOG_PATH, "utf8")) as AutoRouterPresetsResponse, +); export const LOADED_PRESETS_QUERY = { data: BUNDLED_PRESETS, From fcc9b813afcc8faece3dd44d3acfc59b0f503456 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 2 Sep 2026 20:01:35 -0700 Subject: [PATCH 07/14] fix(ui): resolve the preset catalog relative to the mock, not cwd Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HqEPCNLDrssxsuezhAaL4j --- ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts index 7356d92d99e..cff417e6bfa 100644 --- a/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts +++ b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts @@ -4,11 +4,8 @@ import { vi } from "vitest"; import { hydratePresets, type AutoRouterPresetsResponse } from "@/lib/autorouter_presets"; // Derived from the real bundled catalog so a preset edit there flows into test expectations -// instead of redding on a stale copy. Read at runtime rather than imported as a module: the -// catalog lives outside ui/litellm-dashboard, so a module import fails `next build`'s type -// check inside the UI Docker image, whose build context is only this package. -// Exported as vi.fn so a test can override the query state. -const CATALOG_PATH = resolve(process.cwd(), "../../litellm/proxy/public_endpoints/autorouter_presets.json"); +// instead of redding on a stale copy. Exported as vi.fn so a test can override the query state. +const CATALOG_PATH = resolve(__dirname, "../../../../litellm/proxy/public_endpoints/autorouter_presets.json"); export const BUNDLED_PRESETS = hydratePresets( JSON.parse(readFileSync(CATALOG_PATH, "utf8")) as AutoRouterPresetsResponse, From 291d02f8aa90e0ed2cc6a08b38c4976559f9cafe Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 20:32:43 -0700 Subject: [PATCH 08/14] fix(mcp): never exchange the LiteLLM virtual key as the upstream subject token (#39446) --- .../mcp_server/mcp_server_manager.py | 97 ++++-- .../proxy/_experimental/mcp_server/server.py | 9 + .../mcp_server/test_mcp_server.py | 6 +- .../mcp_server/test_mcp_server_manager.py | 311 +++++++++++++++++- 4 files changed, 396 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1f552ff3e13..1434fa5bfea 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -820,20 +820,46 @@ def _should_strip_caller_authorization( if not (mcp_server.is_oauth_passthrough or mcp_server.is_oauth_delegate): return False - normalized_raw_headers: Final = {str(k).lower(): v for k, v in (raw_headers or {}).items() if isinstance(k, str)} - has_explicit_litellm_admission_header: Final = normalized_raw_headers.get("x-litellm-api-key") is not None + has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers) if mcp_server.is_oauth_delegate: return not has_explicit_litellm_admission_header - admission_consumed_authorization_as_litellm_key: Final = ( - user_api_key_auth is not None - and bool(getattr(user_api_key_auth, "api_key", None)) - and not has_explicit_litellm_admission_header - ) - return admission_consumed_authorization_as_litellm_key or ( + return _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth) or ( user_api_key_auth is None and not has_explicit_litellm_admission_header ) +LITELLM_VIRTUAL_KEY_PREFIX: Final = "sk-" + + +def _raw_header_value(raw_headers: Mapping[str, str] | None, name: str) -> str | None: + return next((v for k, v in (raw_headers or {}).items() if isinstance(k, str) and k.lower() == name), None) + + +def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None) -> bool: + """Admission only consumes a non-empty ``x-litellm-api-key``; an empty one falls back to ``Authorization``.""" + return bool(_raw_header_value(raw_headers, "x-litellm-api-key")) + + +def _authorization_is_litellm_admission_credential( + raw_headers: Mapping[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, +) -> bool: + """True when ``Authorization`` carries the LiteLLM key admission validated. + + That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the + same key in both headers. + """ + if user_api_key_auth is None or not user_api_key_auth.api_key: + return False + admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key") + if not admission_header: + return True + authorization: Final = _raw_header_value(raw_headers, "authorization") + return authorization is not None and strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme( + admission_header, "Bearer" + ) + + def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str: """Format a raw BYOK credential for OpenAPI tool ``Authorization`` injection. @@ -3277,8 +3303,8 @@ class MCPServerManager: ######################################################### @staticmethod def _extract_bearer_token( - oauth2_headers: dict[str, str] | None, - raw_headers: dict[str, str] | None, + oauth2_headers: Mapping[str, str] | None, + raw_headers: Mapping[str, str] | None, ) -> str | None: """Extract the bare Bearer token from oauth2_headers or raw_headers. @@ -3298,10 +3324,29 @@ class MCPServerManager: return auth_value return None + @staticmethod + def _extract_subject_token( + oauth2_headers: Mapping[str, str] | None, + raw_headers: Mapping[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, + ) -> str | None: + """The caller's upstream identity token, or ``None`` when the bearer is a LiteLLM key. + + Rejects the key admission validated and, because virtual keys always carry the ``sk-`` prefix, + any other LiteLLM key a client puts in ``Authorization`` next to ``x-litellm-api-key``. + """ + if _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth): + return None + bearer: Final = MCPServerManager._extract_bearer_token(oauth2_headers, raw_headers) + if bearer is not None and bearer.startswith(LITELLM_VIRTUAL_KEY_PREFIX): + return None + return bearer + def _obo_subject_token( self, server: MCPServer, - raw_headers: dict[str, str] | None, + raw_headers: Mapping[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, ) -> str | None: """The caller's bearer as the token_exchange (OBO) subject token, for that mode only. @@ -3311,7 +3356,7 @@ class MCPServerManager: """ if server.auth_type != MCPAuth.oauth2_token_exchange: return None - return self._extract_bearer_token(None, raw_headers) + return self._extract_subject_token(None, raw_headers, user_api_key_auth) def _build_stdio_env( self, @@ -3566,6 +3611,7 @@ class MCPServerManager: server: MCPServer, oauth2_headers: dict[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None, + raw_headers: Mapping[str, str] | None = None, ) -> None: """Run the OBO exchange for a caller-supplied subject at the transport edge. @@ -3577,13 +3623,15 @@ class MCPServerManager: """ if server.auth_type != MCPAuth.oauth2_token_exchange: return - subject_token: Final = self._extract_bearer_token(oauth2_headers, None) - if not subject_token: + if not self._extract_bearer_token(oauth2_headers, None): return resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) spec: Final = to_server_spec(resolved_server) if spec is None or not isinstance(spec.config, TokenExchangeConfig): return + subject_token: Final = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) + if subject_token is None: + raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path()) match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): case Ok(_): return @@ -3851,7 +3899,7 @@ class MCPServerManager: # token (mirrors the call path), not v1's deleted client_credentials fallback. Other modes # never read the inbound bearer, so leave subject_token None to avoid forwarding it. subject_token: Final = ( - self._extract_bearer_token(oauth2_headers, raw_headers) + self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) if server.auth_type == MCPAuth.oauth2_token_exchange else None ) @@ -3931,6 +3979,7 @@ class MCPServerManager: async def get_prompts_from_server( self, server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, add_prefix: bool = True, @@ -3959,7 +4008,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env: Final = self._build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers) + subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client = await self._create_mcp_client( server=server, @@ -3982,6 +4031,7 @@ class MCPServerManager: async def get_resources_from_server( self, server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, add_prefix: bool = True, @@ -4001,7 +4051,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env: Final = self._build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers) + subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client = await self._create_mcp_client( server=server, @@ -4024,6 +4074,7 @@ class MCPServerManager: async def get_resource_templates_from_server( self, server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, add_prefix: bool = True, @@ -4043,7 +4094,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env: Final = self._build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers) + subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client = await self._create_mcp_client( server=server, @@ -4068,6 +4119,7 @@ class MCPServerManager: async def read_resource_from_server( self, server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, url: AnyUrl, mcp_auth_header: str | dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -4084,7 +4136,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env: Final = self._build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers) + subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client: Final = await self._create_mcp_client( server=server, @@ -4099,6 +4151,7 @@ class MCPServerManager: async def get_prompt_from_server( self, server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, prompt_name: str, arguments: dict[str, str] | None = None, mcp_auth_header: str | dict[str, str] | None = None, @@ -4116,7 +4169,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env: Final = self._build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers) + subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client: Final = await self._create_mcp_client( server=server, @@ -5290,7 +5343,7 @@ class MCPServerManager: MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag, ): - subject_token = self._extract_bearer_token(oauth2_headers, raw_headers) + subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) elif mcp_server.auth_type == MCPAuth.oauth2: if mcp_server.has_client_credentials: # For M2M OAuth servers, Authorization must come from token fetch. @@ -5638,7 +5691,7 @@ class MCPServerManager: subject_token: str | None = None if isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)): - subject_token = self._extract_bearer_token(oauth2_headers, raw_headers) + subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) elif isinstance(spec.config, PassthroughConfig): inbound_token, forwarded_headers = _take_forwarded_authorization(forwarded_headers) per_server_token: Final = _passthrough_token_from_mcp_auth_header(mcp_auth_header) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index af3ff6714af..bb075220530 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2189,6 +2189,7 @@ if MCP_AVAILABLE: try: prompts = await global_mcp_server_manager.get_prompts_from_server( server=server, + user_api_key_auth=user_api_key_auth, mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=True, # Always add server prefix @@ -2242,6 +2243,7 @@ if MCP_AVAILABLE: try: resources = await global_mcp_server_manager.get_resources_from_server( server=server, + user_api_key_auth=user_api_key_auth, mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=True, # Always add server prefix @@ -2293,6 +2295,7 @@ if MCP_AVAILABLE: try: resource_templates = await global_mcp_server_manager.get_resource_templates_from_server( server=server, + user_api_key_auth=user_api_key_auth, mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=True, # Always add server prefix @@ -3211,6 +3214,7 @@ if MCP_AVAILABLE: return await global_mcp_server_manager.get_prompt_from_server( server=server, + user_api_key_auth=user_api_key_auth, prompt_name=original_prompt_name, arguments=arguments, mcp_auth_header=server_auth_header, @@ -3261,6 +3265,7 @@ if MCP_AVAILABLE: return await global_mcp_server_manager.read_resource_from_server( server=server, + user_api_key_auth=user_api_key_auth, url=url, mcp_auth_header=server_auth_header, extra_headers=extra_headers, @@ -3723,6 +3728,7 @@ if MCP_AVAILABLE: user_api_key_auth: UserAPIKeyAuth | None, client_ip: str | None, allowed_server_ids: set[str] | None = None, + raw_headers: Mapping[str, str] | None = None, ) -> None: """Fail fast with HTTP 401 for MCP servers that need user auth but didn't receive it on this request. Covers both gateway-managed OAuth2 @@ -3867,6 +3873,7 @@ if MCP_AVAILABLE: server=server, oauth2_headers=oauth2_headers, user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, ) # Pass-through OAuth: when the admin has opted a server into @@ -4195,6 +4202,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, client_ip=_client_ip, allowed_server_ids=toolset_allowed_server_ids, + raw_headers=raw_headers, ) # Pre-flight auth check for pass-through servers. Must run after @@ -4518,6 +4526,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, client_ip=_sse_client_ip, allowed_server_ids=toolset_allowed_server_ids, + raw_headers=raw_headers, ) # Pre-flight auth check for pass-through servers: surface upstream diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 3f6d8f8837c..ff0007f47aa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -914,6 +914,7 @@ async def test_mcp_get_prompt_success(): ) mock_manager.get_prompt_from_server.assert_awaited_once_with( server=server, + user_api_key_auth=user_api_key_auth, prompt_name="hello", arguments={"foo": "bar"}, mcp_auth_header={"Authorization": "token"}, @@ -976,6 +977,7 @@ async def test_mcp_read_resource_success(): ) mock_manager.read_resource_from_server.assert_awaited_once_with( server=server, + user_api_key_auth=user_api_key_auth, url="https://example.com/resource", mcp_auth_header={"Authorization": "token"}, extra_headers={"X-Test": "1"}, @@ -8268,7 +8270,9 @@ class TestOboPreflightScopedToAllowedServers: _, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key) - preflight.assert_awaited_once_with(server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key) + preflight.assert_awaited_once_with( + server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key, raw_headers=None + ) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 5508259273d..5cde2e83f62 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -30,6 +30,7 @@ from mcp.types import ( TextResourceContents, ) from mcp.types import Tool as MCPTool +from pydantic import AnyUrl from litellm.constants import MCP_METADATA_TIMEOUT from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -2270,7 +2271,9 @@ class TestMCPServerManager: """prompts/list on an OBO server must exchange the caller's bearer, not connect with none.""" server = self._token_exchange_server("te-prompts") st = await self._capture_subject_token( - lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"}) + lambda m: m.get_prompts_from_server( + server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"} + ) ) assert st == "subj-jwt" @@ -2279,7 +2282,9 @@ class TestMCPServerManager: """resources/list on an OBO server must exchange the caller's bearer.""" server = self._token_exchange_server("te-resources") st = await self._capture_subject_token( - lambda m: m.get_resources_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"}) + lambda m: m.get_resources_from_server( + server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"} + ) ) assert st == "subj-jwt" @@ -2290,6 +2295,7 @@ class TestMCPServerManager: st = await self._capture_subject_token( lambda m: m.read_resource_from_server( server=server, + user_api_key_auth=None, url="https://up.example.com/r", raw_headers={"authorization": "Bearer subj-jwt"}, ) @@ -2307,7 +2313,9 @@ class TestMCPServerManager: auth_type=MCPAuth.none, ) st = await self._capture_subject_token( - lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"}) + lambda m: m.get_prompts_from_server( + server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"} + ) ) assert st is None @@ -3254,7 +3262,7 @@ class TestMCPServerManager: new_callable=AsyncMock, return_value=mock_client, ): - prompts = await manager.get_prompts_from_server(server, add_prefix=True) + prompts = await manager.get_prompts_from_server(server, user_api_key_auth=None, add_prefix=True) mock_client.list_prompts.assert_awaited_once() assert len(prompts) == 1 @@ -3289,6 +3297,7 @@ class TestMCPServerManager: ): result = await manager.get_prompt_from_server( server=server, + user_api_key_auth=None, prompt_name="hello", arguments={"tone": "casual"}, ) @@ -3334,6 +3343,7 @@ class TestMCPServerManager: ): result = await manager.get_resources_from_server( server=server, + user_api_key_auth=None, mcp_auth_header="auth", extra_headers={"X-Test": "1"}, add_prefix=True, @@ -3391,6 +3401,7 @@ class TestMCPServerManager: ): result = await manager.get_resource_templates_from_server( server=server, + user_api_key_auth=None, mcp_auth_header="auth", extra_headers=None, add_prefix=False, @@ -3441,6 +3452,7 @@ class TestMCPServerManager: ) as mock_create_client: result = await manager.read_resource_from_server( server=server, + user_api_key_auth=None, url="https://example.com/resource", mcp_auth_header="auth", extra_headers={"X-Test": "1"}, @@ -11006,3 +11018,294 @@ class TestOpenApiHandlerRelaysUpstreamAuth: assert result.isError is True assert "upstream returned HTTP 503" in result.content[0].text + + +class TestLitellmAdmissionKeyIsNeverTheSubjectToken: + """The bearer that admitted the request as a LiteLLM key must not be sent to the IdP as the + RFC 8693 subject_token (or ID-JAG assertion). Only ``x-litellm-api-key`` disambiguates: with it + present, ``Authorization`` is the caller's own identity token and is exchanged as before.""" + + _ADMISSION_KEY: Final = "sk-litellm-virtual-key" + _USER_TOKEN: Final = "user-idp-jwt" + + @staticmethod + def _token_exchange_server(server_id: str) -> MCPServer: + return MCPServer( + server_id=server_id, + name=f"{server_id}-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + + @staticmethod + def _id_jag_server(server_id: str) -> MCPServer: + return MCPServer( + server_id=server_id, + name=f"{server_id}-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_id_jag, + client_id="cid", + client_secret="csec", + token_exchange_endpoint="https://idp.example.com/token", + id_jag_resource_token_endpoint="https://resource-as.example.com/token", + ) + + @staticmethod + def _recording_provider() -> MagicMock: + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + provider: Final = MagicMock() + provider.resolve_credentials = AsyncMock( + return_value=Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) + ) + return provider + + @staticmethod + def _subjects_seen_by(provider: MagicMock) -> list[str | None]: + return [ + call.args[0].inbound_token.get_secret_value() if call.args[0].inbound_token else None + for call in provider.resolve_credentials.call_args_list + ] + + @staticmethod + def _manager_with_recording_client() -> MCPServerManager: + manager: Final = MCPServerManager() + client: Final = AsyncMock() + client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) + client.list_prompts = AsyncMock(return_value=[]) + client.read_resource = AsyncMock(return_value=ReadResourceResult(contents=[])) + manager._create_mcp_client = AsyncMock(return_value=client) + return manager + + @staticmethod + def _subject_token_given_to_client(manager: MCPServerManager) -> str | None: + return manager._create_mcp_client.call_args.kwargs["subject_token"] + + async def _call_tool_subject(self, server: MCPServer, oauth2_headers, raw_headers, user_api_key_auth): + manager: Final = self._manager_with_recording_client() + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="tool", + arguments={}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + proxy_logging_obj=None, + user_api_key_auth=user_api_key_auth, + ) + return self._subject_token_given_to_client(manager) + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag]) + async def test_tools_call_with_only_the_litellm_key_has_no_subject(self, auth_type): + server = ( + self._token_exchange_server("te-call") + if auth_type == MCPAuth.oauth2_token_exchange + else self._id_jag_server("jag-call") + ) + subject_token = await self._call_tool_subject( + server, + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_rest_tools_call_with_only_the_litellm_key_has_no_subject(self): + """The REST facade passes no oauth2_headers; the bearer is reached through raw_headers only.""" + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-rest"), + oauth2_headers=None, + raw_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_tools_call_exchanges_the_user_token_when_x_litellm_api_key_admits(self): + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-split"), + oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"}, + raw_headers={ + "X-LiteLLM-API-Key": f"Bearer {self._ADMISSION_KEY}", + "authorization": f"Bearer {self._USER_TOKEN}", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token == self._USER_TOKEN + + @pytest.mark.asyncio + async def test_tools_call_with_an_empty_x_litellm_api_key_has_no_subject(self): + """Admission ignores an empty ``x-litellm-api-key`` and validates ``Authorization`` instead.""" + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-empty-header"), + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + raw_headers={"x-litellm-api-key": "", "authorization": f"Bearer {self._ADMISSION_KEY}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_tools_call_with_the_same_litellm_key_in_both_headers_has_no_subject(self): + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-same-key"), + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + raw_headers={ + "x-litellm-api-key": self._ADMISSION_KEY, + "authorization": f"Bearer {self._ADMISSION_KEY}", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_tools_call_with_a_different_litellm_key_in_authorization_has_no_subject(self): + """A second ``sk-`` virtual key next to ``x-litellm-api-key`` is still a gateway credential.""" + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-second-key"), + oauth2_headers={"Authorization": "Bearer sk-another-virtual-key"}, + raw_headers={ + "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}", + "authorization": "Bearer sk-another-virtual-key", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_tools_call_exchanges_the_bearer_when_jwt_admission_left_api_key_unset(self): + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-jwt"), + oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"}, + raw_headers={"authorization": f"Bearer {self._USER_TOKEN}"}, + user_api_key_auth=UserAPIKeyAuth(api_key=None, user_id="alice"), + ) + assert subject_token == self._USER_TOKEN + + @pytest.mark.asyncio + async def test_tools_list_with_only_the_litellm_key_has_no_subject(self): + manager: Final = self._manager_with_recording_client() + manager._fetch_tools_with_timeout = AsyncMock(return_value=[]) + await manager._get_tools_from_server( + server=self._token_exchange_server("te-list-key"), + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert self._subject_token_given_to_client(manager) is None + + @pytest.mark.asyncio + async def test_prompts_list_with_only_the_litellm_key_has_no_subject(self): + manager: Final = self._manager_with_recording_client() + await manager.get_prompts_from_server( + server=self._token_exchange_server("te-prompts-key"), + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + ) + assert self._subject_token_given_to_client(manager) is None + + @pytest.mark.asyncio + async def test_resource_read_with_only_the_litellm_key_has_no_subject(self): + manager: Final = self._manager_with_recording_client() + await manager.read_resource_from_server( + server=self._token_exchange_server("te-read-key"), + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + url=AnyUrl("file:///notes.txt"), + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + ) + assert self._subject_token_given_to_client(manager) is None + + @pytest.mark.asyncio + async def test_resource_read_exchanges_the_user_token_when_x_litellm_api_key_admits(self): + manager: Final = self._manager_with_recording_client() + await manager.read_resource_from_server( + server=self._token_exchange_server("te-read-split"), + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + url=AnyUrl("file:///notes.txt"), + raw_headers={ + "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}", + "authorization": f"Bearer {self._USER_TOKEN}", + }, + ) + assert self._subject_token_given_to_client(manager) == self._USER_TOKEN + + @pytest.mark.asyncio + async def test_openapi_call_never_hands_the_litellm_key_to_the_exchanger(self): + provider: Final = self._recording_provider() + manager = MCPServerManager(cred_provider=provider) + server = MCPServer( + server_id="te-openapi", + name="te_openapi", + server_name="te_openapi", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + spec_path="https://api.example.com/openapi.json", + ) + user_auth = UserAPIKeyAuth(api_key="hashed-key", user_id="alice") + + await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + mcp_auth_header=None, + user_api_key_auth=user_auth, + forwarded_headers=None, + ) + await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"}, + raw_headers={ + "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}", + "authorization": f"Bearer {self._USER_TOKEN}", + }, + mcp_auth_header=None, + user_api_key_auth=user_auth, + forwarded_headers=None, + ) + assert self._subjects_seen_by(provider) == [None, self._USER_TOKEN] + + @pytest.mark.asyncio + async def test_preflight_challenges_instead_of_exchanging_the_litellm_key(self): + provider: Final = self._recording_provider() + manager = MCPServerManager(cred_provider=provider) + + with pytest.raises(HTTPException) as exc_info: + await manager.preflight_token_exchange( + server=self._token_exchange_server("te-preflight-key"), + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + ) + assert exc_info.value.status_code == 401 + headers = exc_info.value.headers or {} + assert "resource_metadata" in (headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "") + assert self._subjects_seen_by(provider) == [] + + @pytest.mark.asyncio + async def test_preflight_exchanges_the_user_token_when_x_litellm_api_key_admits(self): + provider: Final = self._recording_provider() + manager = MCPServerManager(cred_provider=provider) + + await manager.preflight_token_exchange( + server=self._token_exchange_server("te-preflight-split"), + oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + raw_headers={ + "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}", + "authorization": f"Bearer {self._USER_TOKEN}", + }, + ) + assert self._subjects_seen_by(provider) == [self._USER_TOKEN] From e058aa68c4331ce7ffce2802528eacc2f2333642 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 20:40:00 -0700 Subject: [PATCH 09/14] test: add mistral ocr transformation parity coverage (#39482) * test: cover mistral ocr transformation parity Co-Authored-By: Claude Code * test: map mistral ocr parity contracts Co-Authored-By: Claude Code --------- Co-authored-by: Claude Code --- .../providers/mistral/ocr/transformation.rs | 188 ++- .../ledgers/ocr/ocr_test_ledger.json | 1240 ++++++++++++++--- 2 files changed, 1208 insertions(+), 220 deletions(-) diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index 9648321d7ff..0125886aac1 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -186,68 +186,180 @@ mod tests { use serde_json::json; #[test] - fn supported_params_match_python_mistral_ocr_config() { + fn extract_header_is_a_supported_ocr_param() { + assert!(supported_ocr_params().contains(&"extract_header")); + } + + #[test] + fn extract_footer_is_a_supported_ocr_param() { + assert!(supported_ocr_params().contains(&"extract_footer")); + } + + #[test] + fn existing_ocr_params_remain_supported() { + for param in [ + "pages", + "include_image_base64", + "image_limit", + "image_min_size", + "bbox_annotation_format", + "document_annotation_format", + ] { + assert!(supported_ocr_params().contains(¶m)); + } + } + + #[test] + fn map_ocr_params_forwards_extract_header() { + let params = json!({"extract_header": true}); assert_eq!( - supported_ocr_params(), - &[ - "pages", - "include_image_base64", - "image_limit", - "image_min_size", - "bbox_annotation_format", - "document_annotation_format", - "document_annotation_prompt", - "extract_header", - "extract_footer", - "table_format", - "confidence_scores_granularity", - "include_blocks", - "id", - ] + map_ocr_params(params.as_object().unwrap()), + params.as_object().unwrap().clone() + ); + } + + #[test] + fn map_ocr_params_forwards_extract_footer() { + let params = json!({"extract_footer": true}); + assert_eq!( + map_ocr_params(params.as_object().unwrap()), + params.as_object().unwrap().clone() + ); + } + + #[test] + fn map_ocr_params_forwards_extract_header_and_footer() { + let params = json!({"extract_header": true, "extract_footer": false}); + assert_eq!( + map_ocr_params(params.as_object().unwrap()), + params.as_object().unwrap().clone() ); } #[test] fn map_ocr_params_drops_unknown_params() { - let params = json!({ - "extract_header": true, - "unsupported_param": "value", - "pages": [0, 1] - }); + let params = json!({"extract_header": true, "unsupported_param": "value"}); let mapped = map_ocr_params(params.as_object().unwrap()); - assert_eq!(mapped.get("extract_header"), Some(&json!(true))); - assert_eq!(mapped.get("pages"), Some(&json!([0, 1]))); assert!(!mapped.contains_key("unsupported_param")); } #[test] - fn transform_ocr_request_builds_mistral_body() { + fn new_ocr_params_are_supported() { + for param in [ + "table_format", + "confidence_scores_granularity", + "document_annotation_prompt", + "include_blocks", + "id", + ] { + assert!(supported_ocr_params().contains(¶m)); + } + } + + #[test] + fn map_ocr_params_forwards_new_ocr_params() { + for (param, value) in [ + ("table_format", json!("html")), + ("confidence_scores_granularity", json!("word")), + ( + "document_annotation_prompt", + json!("Extract all invoice line items"), + ), + ("include_blocks", json!(true)), + ("id", json!("req-123")), + ] { + let params = json!({param: value}); + assert_eq!( + map_ocr_params(params.as_object().unwrap()), + params.as_object().unwrap().clone() + ); + } + } + + #[test] + fn transform_ocr_request_includes_each_optional_param() { + let document = json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }); + for (param, value) in [ + ("table_format", json!("html")), + ("confidence_scores_granularity", json!("word")), + ( + "document_annotation_prompt", + json!("Extract all invoice line items"), + ), + ("id", json!("req-123")), + ("extract_header", json!(true)), + ("include_blocks", json!(true)), + ("pages", json!([0, 1])), + ] { + let result = transform_ocr_request( + "mistral-ocr-latest", + document.clone(), + json!({param: value}).as_object().unwrap().clone(), + ) + .expect("request should transform"); + assert_eq!(result.data.get(param), Some(&value)); + assert_eq!(result.data.get("model"), Some(&json!("mistral-ocr-latest"))); + assert_eq!(result.data.get("document"), Some(&document)); + assert_eq!(result.files, None); + } + } + + #[test] + fn transform_ocr_request_includes_multiple_new_params() { let document = json!({ "type": "document_url", "document_url": "https://example.com/doc.pdf" }); let optional_params = json!({ - "include_image_base64": true, - "table_format": "html" + "table_format": "html", + "confidence_scores_granularity": "page", + "extract_header": true }) .as_object() .unwrap() .clone(); - - let result = transform_ocr_request("mistral-ocr-latest", document.clone(), optional_params) + let result = transform_ocr_request("mistral-ocr-latest", document, optional_params) .expect("request should transform"); - + assert_eq!(result.data.get("table_format"), Some(&json!("html"))); assert_eq!( - result.data, - json!({ - "model": "mistral-ocr-latest", - "document": document, - "include_image_base64": true, - "table_format": "html" - }) + result.data.get("confidence_scores_granularity"), + Some(&json!("page")) ); - assert_eq!(result.files, None); + assert_eq!(result.data.get("extract_header"), Some(&json!(true))); + } + + #[test] + fn transform_ocr_response_preserves_blocks_and_confidence_scores() { + let blocks = json!([{"type": "title", "content": "Invoice"}]); + let confidence_scores = json!({"page": 0.98}); + let response = json!({ + "pages": [{"index": 0, "markdown": "# Invoice", "blocks": blocks, "confidence_scores": confidence_scores}], + "model": "mistral-ocr-4-0", + "usage_info": {"pages_processed": 1} + }); + let result = + transform_ocr_response("mistral-ocr-4-0", response).expect("response should transform"); + assert_eq!(result.pages[0].get("blocks"), Some(&blocks)); + assert_eq!( + result.pages[0].get("confidence_scores"), + Some(&confidence_scores) + ); + } + + #[test] + fn transform_ocr_response_preserves_ocr4_page_fields() { + let response = json!({ + "pages": [{"index": 0, "markdown": "table page", "tables": [{"rows": 2, "cols": 3}], "hyperlinks": ["https://example.com"], "header": "Acme Corp", "footer": "Page 1"}], + "model": "mistral-ocr-4-0", + "usage_info": {"pages_processed": 1} + }); + let result = transform_ocr_response("mistral-ocr-4-0", response.clone()) + .expect("response should transform"); + assert_eq!(result.pages[0], response["pages"][0]); } #[test] diff --git a/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json b/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json index 799a1320463..e617ddf8f94 100644 --- a/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json +++ b/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json @@ -20,189 +20,1065 @@ "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs" ], "entries": [ - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_should_encode_azure_document_intelligence_model_id", "status": "unmapped", "reason": "model-id URL percent-encoding has no Rust test; Rust only tests pages/features query building"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_should_reject_dot_segment_azure_document_intelligence_model_id", "status": "unmapped", "reason": "model-id dot-segment validation has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_preserves_azure_native_fields", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_response_normalizes_pages", "justification": "both assert page markdown, dimension (inch-to-pixel) normalization, and usage_info.pages_processed from the same Azure succeeded response shape"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_async_transform_ocr_response_preserves_azure_native_fields", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_response_normalizes_pages", "justification": "async twin of the sync case above, same underlying transform is exercised on the Rust side"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_tolerates_missing_native_fields", "status": "unmapped", "reason": "tables/keyValuePairs absence tolerance is not asserted by the Rust response test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_non_succeeded_status_raises", "status": "unmapped", "reason": "no Rust test asserts on a non-succeeded Azure DI status"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_supported_ocr_params_includes_features", "status": "unmapped", "reason": "supported-params list content has no Rust equivalent for Azure"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_native_format_carries_raw_operation", "status": "unmapped", "reason": "native req_format raw-operation passthrough is not tested in Rust"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_async_transform_ocr_response_native_format_carries_raw_operation", "status": "unmapped", "reason": "native req_format raw-operation passthrough is not tested in Rust"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_default_format_omits_raw_operation", "status": "unmapped", "reason": "req_format gating of raw-operation output has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_passes_through_req_format", "status": "unmapped", "reason": "req_format passthrough in map_ocr_params has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_rejects_unknown_req_format_as_bad_request", "status": "unmapped", "reason": "req_format validation error path has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_omits_req_format_query_param", "status": "unmapped", "reason": "no Rust test asserts req_format is excluded from the built URL"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_features", "status": "unmapped", "reason": "features-string normalization in map_ocr_params has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_empty_features_list_omitted", "status": "unmapped", "reason": "empty-features omission has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_invalid_features_raises", "status": "unmapped", "reason": "features validation error path has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_appends_features_query", "status": "unmapped", "reason": "features query-param construction has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_combines_pages_and_features", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_url_normalizes_zero_based_pages", "justification": "both assert 0-based, duplicate page indices are deduped, sorted, and rewritten 1-based into the request URL"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_validate_environment_uses_subscription_key", "status": "unmapped", "reason": "Python-side header derivation from litellm_params; Rust's poll test only checks the header is present, not how it was resolved"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_validate_environment_falls_back_to_entra_token", "status": "unmapped", "reason": "Entra bearer-token fallback logic has no Rust test"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_doc_intelligence_route", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_documentintelligence_and_is_case_insensitive", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_does_not_match_mistral_route", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_does_not_hijack_doc_intelligence", "status": "unmapped", "reason": "api_base resolution from the secret manager runs before the Rust bridge is called, no Rust test exists for it"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_explicit_api_base_is_honoured_for_doc_intelligence", "status": "unmapped", "reason": "api_base precedence resolution is Python-only"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_still_applies_to_mistral_ocr", "status": "unmapped", "reason": "api_base precedence resolution is Python-only"}, - - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_use_litellm_rust_toggles_flag", "status": "unmapped", "reason": "bridge-plumbing: Python-side feature-flag toggle, no Rust equivalent"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_env_var_enables_rust_ocr", "status": "unmapped", "reason": "bridge-plumbing: Python-side env-var flag gating"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_returns_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: dependency-injection test hook, not provider behavior"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_loader_returns_none_when_extension_absent", "status": "unmapped", "reason": "bridge-plumbing: native-extension import/loader fallback"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_loader_caches_absent_extension", "status": "unmapped", "reason": "bridge-plumbing: loader caching behavior"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_available_reflects_loader", "status": "unmapped", "reason": "bridge-plumbing: loader availability check"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_aocr_returns_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: dependency-injection test hook"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_toggle_without_ocr_arg_preserves_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: injected-impl state retention regression"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_explicit_ocr_none_clears_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: injected-impl clearing behavior"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_none_when_extension_absent", "status": "unmapped", "reason": "bridge-plumbing: degrade path when the native extension is missing"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_uses_compiled_extension", "status": "unmapped", "reason": "bridge-plumbing: native module resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_timeout_to_seconds_handles_float_timeout_and_none", "status": "unmapped", "reason": "bridge-plumbing: Python-side timeout normalization helper"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_bridge_wrapper_forwards_prepared_args_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: wrapper argument forwarding, asserted against a fake bridge not the real Rust code"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: async wrapper argument forwarding"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_prepares_request_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: request preparation and response wrapping in Python"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_resolves_key_via_secret_manager_when_missing", "status": "unmapped", "reason": "secret-manager: API key resolution happens in Python before the bridge is invoked"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_prefers_explicit_key_over_resolver", "status": "unmapped", "reason": "secret-manager: key precedence resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_uses_provider_api_key_env_var", "status": "unmapped", "reason": "secret-manager: provider-specific env var name resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_forwards_vertex_routing_metadata", "status": "unmapped", "reason": "secret-manager: vertex routing metadata merge happens in Python"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager", "status": "unmapped", "reason": "secret-manager: vertex project/location resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager", "status": "unmapped", "reason": "secret-manager: azure_ai api_base resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint", "status": "unmapped", "reason": "secret-manager: doc-intelligence endpoint resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_runs_pre_call_logging", "status": "unmapped", "reason": "bridge-plumbing: Python logging-object pre_call invocation"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_routes_to_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: routing to a fake bridge, not the real Rust transform"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_routes_azure_ai_to_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: provider-prefix stripping before routing"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_rust_path_converts_file_document_before_bridge", "status": "unmapped", "reason": "file-normalization: raw-bytes-to-data-URI conversion happens in Python before the bridge call"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_exception_type_uses_resolved_provider_context", "status": "unmapped", "reason": "bridge-plumbing: Python exception-type mapping on bridge failure"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_aocr_routes_to_async_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: async routing to a fake bridge"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_aocr_exception_type_uses_resolved_provider_context", "status": "unmapped", "reason": "bridge-plumbing: async exception-type mapping on bridge failure"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_does_not_route_to_rust_when_disabled", "status": "unmapped", "reason": "bridge-plumbing: Python control flow for the toggle-disabled branch, no Rust-owned behavior runs"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_falls_back_to_python_when_bridge_unavailable", "status": "unmapped", "reason": "bridge-plumbing: Python-only fallback when the compiled Rust extension is absent, Rust cannot test its own absence"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_forwards_timeout_to_rust", "status": "unmapped", "reason": "bridge-plumbing: asserts the Python call site forwards a timeout kwarg, Rust receives an already-constructed request"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_passes_default_request_timeout_to_rust", "status": "unmapped", "reason": "bridge-plumbing: asserts the Python call site supplies a default timeout kwarg, no Rust equivalent"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_provider_configs_expose_api_key_env_vars", "status": "unmapped", "reason": "asserts per-provider get_api_key_env_var() strings; the closest Rust test (ocr_dispatch_supports_migrated_providers) asserts provider dispatch/param resolution instead, not API key env var names"}, - - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_pdf_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection is Python-only preprocessing before the bridge call"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_png_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_jpg_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_jpeg_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_gif_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_webp_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_tiff_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_tif_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_bmp_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_be_case_insensitive", "status": "unmapped", "reason": "file-normalization: MIME detection case handling"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_fallback_for_unknown_extension", "status": "unmapped", "reason": "file-normalization: MIME detection fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pdf_pathlib_path_to_document_url", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion happens in Python"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_image_pathlib_path_to_image_url", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_reject_bare_str_path", "status": "unmapped", "reason": "file-normalization: arbitrary-file-read guard on bare str paths"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pathlib_path", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes", "status": "unmapped", "reason": "file-normalization: raw-bytes-to-data-URI conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_explicit_mime_type", "status": "unmapped", "reason": "file-normalization: explicit MIME override on raw bytes"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_image_mime_type", "status": "unmapped", "reason": "file-normalization: explicit MIME override on raw bytes"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object", "status": "unmapped", "reason": "file-normalization: file-like-object conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object_with_name", "status": "unmapped", "reason": "file-normalization: file-like-object name-based MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_missing_file_field", "status": "unmapped", "reason": "file-normalization: missing-field validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_nonexistent_pathlib_path", "status": "unmapped", "reason": "file-normalization: missing-file validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_empty_file", "status": "unmapped", "reason": "file-normalization: empty-file validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_unsupported_type", "status": "unmapped", "reason": "file-normalization: unsupported input type validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_invalid_mime_type", "status": "unmapped", "reason": "file-normalization: MIME-type injection validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_override_mime_type_for_pathlib_path", "status": "unmapped", "reason": "file-normalization: explicit MIME override precedence"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_document_url_for_pdf", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_png", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_jpeg", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_octet_stream", "status": "unmapped", "reason": "file-normalization: filename-based MIME fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_none", "status": "unmapped", "reason": "file-normalization: filename-based MIME fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_fallback_to_octet_stream_for_unknown", "status": "unmapped", "reason": "file-normalization: default MIME fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_preserve_base64_content_correctly", "status": "unmapped", "reason": "file-normalization: binary round-trip through base64"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_from_content_type", "status": "unmapped", "reason": "file-normalization: content-type parameter stripping"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_with_multiple_params", "status": "unmapped", "reason": "file-normalization: content-type parameter stripping"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_reject_file_type_document_in_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body file-type guard, a different mechanism than Rust's URL-fetch SSRF guard"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_accept_document_url_type_in_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body parsing"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_raise_on_invalid_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body parsing error path"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_ignore_document_form_field_injection", "status": "unmapped", "reason": "proxy-layer multipart form-field injection guard, a different mechanism than Rust's URL-fetch SSRF guard"}, - - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_extract_header_in_supported_params", "status": "unmapped", "reason": "Rust's fixed-list test checks the full list as one assertion, not this individual param"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_extract_footer_in_supported_params", "status": "unmapped", "reason": "Rust's fixed-list test checks the full list as one assertion, not this individual param"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_existing_params_still_present", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "supported_params_match_python_mistral_ocr_config", "justification": "both assert the full supported_ocr_params list matches the same fixed set of param names"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_header_passed_through", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_drops_unknown_params", "justification": "both assert extract_header survives map_ocr_params filtering unchanged"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_footer_passed_through", "status": "unmapped", "reason": "Rust's map_ocr_params test does not assert on extract_footer specifically"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_header_and_footer_together", "status": "unmapped", "reason": "combined extract_header+extract_footer passthrough is not asserted together in Rust"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_unknown_param_is_dropped", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_drops_unknown_params", "justification": "both assert an unrecognized param key is dropped while a known one is kept"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestNewSupportedParams::test_new_param_in_supported_list", "status": "unmapped", "reason": "OCR4-specific new params (table_format etc) are not individually verified against the Rust fixed-list test"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestNewParamsMapOcr::test_new_param_passed_through", "status": "unmapped", "reason": "OCR4-specific new params are not individually asserted in the Rust map_ocr_params test"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrRequest::test_param_included_in_request_body", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_builds_mistral_body", "justification": "both assert an optional param value ends up in the built request body alongside model/document"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrRequest::test_multiple_new_params_together", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_builds_mistral_body", "justification": "both assert multiple optional params (table_format/include_image_base64) land correctly in the same request body"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrResponseOcr4Fields::test_blocks_and_confidence_scores_preserved", "status": "unmapped", "reason": "OCR4 blocks/confidence_scores fields are not asserted by the Rust response test"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrResponseOcr4Fields::test_ocr4_fields_survive_model_dump", "status": "unmapped", "reason": "OCR4 tables/hyperlinks/header/footer fields are not asserted by the Rust response test"}, - - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_model_info_ocr4_price", "status": "unmapped", "reason": "cost-calc: pricing/model-info lookup is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr4_cost_scales_with_pages", "status": "unmapped", "reason": "cost-calc: per-page pricing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_pricing_entry", "status": "unmapped", "reason": "cost-calc: cost-map JSON entry validation is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_model_info_price", "status": "unmapped", "reason": "cost-calc: pricing/model-info lookup is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_cost_scales_with_pages", "status": "unmapped", "reason": "cost-calc: per-page pricing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_ocr_and_annotation_pages_at_their_own_rates", "status": "unmapped", "reason": "cost-calc: mixed-rate billing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_annotation_only_response", "status": "unmapped", "reason": "cost-calc: annotation-only billing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_annotation_pages_when_pages_processed_missing", "status": "unmapped", "reason": "cost-calc: fallback billing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate", "status": "unmapped", "reason": "cost-calc: fallback billing math is Python-only"}, - - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_rust_ocr_serves_default_format", "status": "unmapped", "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_rust_ocr_skipped_for_native_format", "status": "unmapped", "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_native_format_rejected_for_provider_without_support_as_bad_request", "status": "unmapped", "reason": "provider-support validation for req_format happens in Python"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_unknown_format_rejected_for_provider_without_support_as_bad_request", "status": "unmapped", "reason": "req_format validation error path is Python-only"}, - - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestHandlerDiscovery::test_handler_discovered_for_ocr", "status": "unmapped", "reason": "guardrail-translation handler discovery is a Python proxy-layer concern"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestHandlerDiscovery::test_handler_discovered_for_aocr", "status": "unmapped", "reason": "guardrail-translation handler discovery is a Python proxy-layer concern"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_document_url", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_image_url", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_no_document", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_invalid_document", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_input_blocking_guardrail", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_single_page", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_multiple_pages", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_empty_pages", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_page_with_empty_markdown", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_preserves_page_metadata", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_output_blocking_guardrail", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestPIIMaskingScenario::test_pii_masking_in_ocr_pages", "status": "unmapped", "reason": "PII redaction in the translation handler has no Rust equivalent"}, - - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_read_req_format_from_header", "status": "unmapped", "reason": "proxy-layer header parsing has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_prefer_body_req_format_over_header", "status": "unmapped", "reason": "proxy-layer body-vs-header precedence has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_omit_req_format_when_header_absent", "status": "unmapped", "reason": "proxy-layer parsing has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_reject_unknown_req_format", "status": "unmapped", "reason": "proxy-layer validation has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_return_native_payload_with_litellm_response_headers", "status": "unmapped", "reason": "proxy-layer response construction has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_return_normalized_response_when_no_native_payload", "status": "unmapped", "reason": "proxy-layer response construction has no Rust equivalent"} + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_should_encode_azure_document_intelligence_model_id", + "status": "unmapped", + "reason": "model-id URL percent-encoding has no Rust test; Rust only tests pages/features query building" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_should_reject_dot_segment_azure_document_intelligence_model_id", + "status": "unmapped", + "reason": "model-id dot-segment validation has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_preserves_azure_native_fields", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_response_normalizes_pages", + "justification": "both assert page markdown, dimension (inch-to-pixel) normalization, and usage_info.pages_processed from the same Azure succeeded response shape" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_async_transform_ocr_response_preserves_azure_native_fields", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_response_normalizes_pages", + "justification": "async twin of the sync case above, same underlying transform is exercised on the Rust side" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_tolerates_missing_native_fields", + "status": "unmapped", + "reason": "tables/keyValuePairs absence tolerance is not asserted by the Rust response test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_non_succeeded_status_raises", + "status": "unmapped", + "reason": "no Rust test asserts on a non-succeeded Azure DI status" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_supported_ocr_params_includes_features", + "status": "unmapped", + "reason": "supported-params list content has no Rust equivalent for Azure" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_native_format_carries_raw_operation", + "status": "unmapped", + "reason": "native req_format raw-operation passthrough is not tested in Rust" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_async_transform_ocr_response_native_format_carries_raw_operation", + "status": "unmapped", + "reason": "native req_format raw-operation passthrough is not tested in Rust" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_default_format_omits_raw_operation", + "status": "unmapped", + "reason": "req_format gating of raw-operation output has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_passes_through_req_format", + "status": "unmapped", + "reason": "req_format passthrough in map_ocr_params has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_rejects_unknown_req_format_as_bad_request", + "status": "unmapped", + "reason": "req_format validation error path has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_complete_url_omits_req_format_query_param", + "status": "unmapped", + "reason": "no Rust test asserts req_format is excluded from the built URL" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_features", + "status": "unmapped", + "reason": "features-string normalization in map_ocr_params has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_empty_features_list_omitted", + "status": "unmapped", + "reason": "empty-features omission has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_invalid_features_raises", + "status": "unmapped", + "reason": "features validation error path has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_complete_url_appends_features_query", + "status": "unmapped", + "reason": "features query-param construction has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_complete_url_combines_pages_and_features", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_url_normalizes_zero_based_pages", + "justification": "both assert 0-based, duplicate page indices are deduped, sorted, and rewritten 1-based into the request URL" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_validate_environment_uses_subscription_key", + "status": "unmapped", + "reason": "Python-side header derivation from litellm_params; Rust's poll test only checks the header is present, not how it was resolved" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_validate_environment_falls_back_to_entra_token", + "status": "unmapped", + "reason": "Entra bearer-token fallback logic has no Rust test" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_doc_intelligence_route", + "status": "unmapped", + "reason": "model-route string matching is Python-only dispatch logic" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_documentintelligence_and_is_case_insensitive", + "status": "unmapped", + "reason": "model-route string matching is Python-only dispatch logic" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestIsAzureDocumentIntelligenceModel::test_does_not_match_mistral_route", + "status": "unmapped", + "reason": "model-route string matching is Python-only dispatch logic" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_does_not_hijack_doc_intelligence", + "status": "unmapped", + "reason": "api_base resolution from the secret manager runs before the Rust bridge is called, no Rust test exists for it" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestDocIntelligenceApiBaseResolution::test_explicit_api_base_is_honoured_for_doc_intelligence", + "status": "unmapped", + "reason": "api_base precedence resolution is Python-only" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_still_applies_to_mistral_ocr", + "status": "unmapped", + "reason": "api_base precedence resolution is Python-only" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_use_litellm_rust_toggles_flag", + "status": "unmapped", + "reason": "bridge-plumbing: Python-side feature-flag toggle, no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_env_var_enables_rust_ocr", + "status": "unmapped", + "reason": "bridge-plumbing: Python-side env-var flag gating" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_ocr_returns_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: dependency-injection test hook, not provider behavior" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_native_bridge_loader_returns_none_when_extension_absent", + "status": "unmapped", + "reason": "bridge-plumbing: native-extension import/loader fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_native_bridge_loader_caches_absent_extension", + "status": "unmapped", + "reason": "bridge-plumbing: loader caching behavior" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_native_bridge_available_reflects_loader", + "status": "unmapped", + "reason": "bridge-plumbing: loader availability check" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_aocr_returns_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: dependency-injection test hook" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_toggle_without_ocr_arg_preserves_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: injected-impl state retention regression" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_explicit_ocr_none_clears_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: injected-impl clearing behavior" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_ocr_none_when_extension_absent", + "status": "unmapped", + "reason": "bridge-plumbing: degrade path when the native extension is missing" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_ocr_uses_compiled_extension", + "status": "unmapped", + "reason": "bridge-plumbing: native module resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_timeout_to_seconds_handles_float_timeout_and_none", + "status": "unmapped", + "reason": "bridge-plumbing: Python-side timeout normalization helper" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_bridge_wrapper_forwards_prepared_args_and_wraps_response", + "status": "unmapped", + "reason": "bridge-plumbing: wrapper argument forwarding, asserted against a fake bridge not the real Rust code" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response", + "status": "unmapped", + "reason": "bridge-plumbing: async wrapper argument forwarding" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_prepares_request_and_wraps_response", + "status": "unmapped", + "reason": "bridge-plumbing: request preparation and response wrapping in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_resolves_key_via_secret_manager_when_missing", + "status": "unmapped", + "reason": "secret-manager: API key resolution happens in Python before the bridge is invoked" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_prefers_explicit_key_over_resolver", + "status": "unmapped", + "reason": "secret-manager: key precedence resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_uses_provider_api_key_env_var", + "status": "unmapped", + "reason": "secret-manager: provider-specific env var name resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_forwards_vertex_routing_metadata", + "status": "unmapped", + "reason": "secret-manager: vertex routing metadata merge happens in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager", + "status": "unmapped", + "reason": "secret-manager: vertex project/location resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager", + "status": "unmapped", + "reason": "secret-manager: azure_ai api_base resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint", + "status": "unmapped", + "reason": "secret-manager: doc-intelligence endpoint resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_runs_pre_call_logging", + "status": "unmapped", + "reason": "bridge-plumbing: Python logging-object pre_call invocation" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_routes_to_rust_when_enabled", + "status": "unmapped", + "reason": "bridge-plumbing: routing to a fake bridge, not the real Rust transform" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_routes_azure_ai_to_rust_when_enabled", + "status": "unmapped", + "reason": "bridge-plumbing: provider-prefix stripping before routing" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_rust_path_converts_file_document_before_bridge", + "status": "unmapped", + "reason": "file-normalization: raw-bytes-to-data-URI conversion happens in Python before the bridge call" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_exception_type_uses_resolved_provider_context", + "status": "unmapped", + "reason": "bridge-plumbing: Python exception-type mapping on bridge failure" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_aocr_routes_to_async_rust_when_enabled", + "status": "unmapped", + "reason": "bridge-plumbing: async routing to a fake bridge" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_aocr_exception_type_uses_resolved_provider_context", + "status": "unmapped", + "reason": "bridge-plumbing: async exception-type mapping on bridge failure" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_does_not_route_to_rust_when_disabled", + "status": "unmapped", + "reason": "bridge-plumbing: Python control flow for the toggle-disabled branch, no Rust-owned behavior runs" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_falls_back_to_python_when_bridge_unavailable", + "status": "unmapped", + "reason": "bridge-plumbing: Python-only fallback when the compiled Rust extension is absent, Rust cannot test its own absence" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_forwards_timeout_to_rust", + "status": "unmapped", + "reason": "bridge-plumbing: asserts the Python call site forwards a timeout kwarg, Rust receives an already-constructed request" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_passes_default_request_timeout_to_rust", + "status": "unmapped", + "reason": "bridge-plumbing: asserts the Python call site supplies a default timeout kwarg, no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_provider_configs_expose_api_key_env_vars", + "status": "unmapped", + "reason": "asserts per-provider get_api_key_env_var() strings; the closest Rust test (ocr_dispatch_supports_migrated_providers) asserts provider dispatch/param resolution instead, not API key env var names" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_pdf_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection is Python-only preprocessing before the bridge call" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_png_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_jpg_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_jpeg_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_gif_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_webp_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_tiff_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_tif_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_bmp_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_be_case_insensitive", + "status": "unmapped", + "reason": "file-normalization: MIME detection case handling" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_fallback_for_unknown_extension", + "status": "unmapped", + "reason": "file-normalization: MIME detection fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pdf_pathlib_path_to_document_url", + "status": "unmapped", + "reason": "file-normalization: local-path-to-data-URI conversion happens in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_image_pathlib_path_to_image_url", + "status": "unmapped", + "reason": "file-normalization: local-path-to-data-URI conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_reject_bare_str_path", + "status": "unmapped", + "reason": "file-normalization: arbitrary-file-read guard on bare str paths" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pathlib_path", + "status": "unmapped", + "reason": "file-normalization: local-path-to-data-URI conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes", + "status": "unmapped", + "reason": "file-normalization: raw-bytes-to-data-URI conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_explicit_mime_type", + "status": "unmapped", + "reason": "file-normalization: explicit MIME override on raw bytes" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_image_mime_type", + "status": "unmapped", + "reason": "file-normalization: explicit MIME override on raw bytes" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object", + "status": "unmapped", + "reason": "file-normalization: file-like-object conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object_with_name", + "status": "unmapped", + "reason": "file-normalization: file-like-object name-based MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_missing_file_field", + "status": "unmapped", + "reason": "file-normalization: missing-field validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_nonexistent_pathlib_path", + "status": "unmapped", + "reason": "file-normalization: missing-file validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_empty_file", + "status": "unmapped", + "reason": "file-normalization: empty-file validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_unsupported_type", + "status": "unmapped", + "reason": "file-normalization: unsupported input type validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_invalid_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME-type injection validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_override_mime_type_for_pathlib_path", + "status": "unmapped", + "reason": "file-normalization: explicit MIME override precedence" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_build_document_url_for_pdf", + "status": "unmapped", + "reason": "file-normalization: multipart upload conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_png", + "status": "unmapped", + "reason": "file-normalization: multipart upload conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_jpeg", + "status": "unmapped", + "reason": "file-normalization: multipart upload conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_octet_stream", + "status": "unmapped", + "reason": "file-normalization: filename-based MIME fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_none", + "status": "unmapped", + "reason": "file-normalization: filename-based MIME fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_fallback_to_octet_stream_for_unknown", + "status": "unmapped", + "reason": "file-normalization: default MIME fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_preserve_base64_content_correctly", + "status": "unmapped", + "reason": "file-normalization: binary round-trip through base64" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_from_content_type", + "status": "unmapped", + "reason": "file-normalization: content-type parameter stripping" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_with_multiple_params", + "status": "unmapped", + "reason": "file-normalization: content-type parameter stripping" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_reject_file_type_document_in_json_body", + "status": "unmapped", + "reason": "proxy-layer JSON-body file-type guard, a different mechanism than Rust's URL-fetch SSRF guard" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_accept_document_url_type_in_json_body", + "status": "unmapped", + "reason": "proxy-layer JSON-body parsing" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_raise_on_invalid_json_body", + "status": "unmapped", + "reason": "proxy-layer JSON-body parsing error path" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_ignore_document_form_field_injection", + "status": "unmapped", + "reason": "proxy-layer multipart form-field injection guard, a different mechanism than Rust's URL-fetch SSRF guard" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestGetSupportedOcrParams::test_extract_header_in_supported_params", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "extract_header_is_a_supported_ocr_param", + "justification": "both assert extract_header appears in the Mistral supported OCR params list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestGetSupportedOcrParams::test_extract_footer_in_supported_params", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "extract_footer_is_a_supported_ocr_param", + "justification": "both assert extract_footer appears in the Mistral supported OCR params list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestGetSupportedOcrParams::test_existing_params_still_present", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "existing_ocr_params_remain_supported", + "justification": "both assert the previously supported params are still present in the supported list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_extract_header_passed_through", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_extract_header", + "justification": "both assert extract_header alone survives map_ocr_params unchanged" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_extract_footer_passed_through", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_extract_footer", + "justification": "both assert extract_footer alone survives map_ocr_params unchanged" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_extract_header_and_footer_together", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_extract_header_and_footer", + "justification": "both assert header and footer passed together are both forwarded with their given values" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_unknown_param_is_dropped", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_drops_unknown_params", + "justification": "both assert an unrecognized param key is dropped while a known one is kept" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestNewSupportedParams::test_new_param_in_supported_list", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "new_ocr_params_are_supported", + "justification": "both assert each OCR4 param is in the supported list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestNewParamsMapOcr::test_new_param_passed_through", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_new_ocr_params", + "justification": "both assert each OCR4 param/value pair survives map_ocr_params unchanged" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrRequest::test_param_included_in_request_body", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_request_includes_each_optional_param", + "justification": "both assert each optional param value lands in the built request body alongside model/document with no files" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrRequest::test_multiple_new_params_together", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_request_includes_multiple_new_params", + "justification": "both assert multiple OCR4 params passed together all land in the same request body" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrResponseOcr4Fields::test_blocks_and_confidence_scores_preserved", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_response_preserves_blocks_and_confidence_scores", + "justification": "both assert blocks and confidence_scores survive the OCR response transform on the returned page" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrResponseOcr4Fields::test_ocr4_fields_survive_model_dump", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_response_preserves_ocr4_page_fields", + "justification": "both assert tables, hyperlinks, header and footer survive the OCR response transform on the returned page" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_model_info_ocr4_price", + "status": "unmapped", + "reason": "cost-calc: pricing/model-info lookup is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr4_cost_scales_with_pages", + "status": "unmapped", + "reason": "cost-calc: per-page pricing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_pricing_entry", + "status": "unmapped", + "reason": "cost-calc: cost-map JSON entry validation is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_model_info_price", + "status": "unmapped", + "reason": "cost-calc: pricing/model-info lookup is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_cost_scales_with_pages", + "status": "unmapped", + "reason": "cost-calc: per-page pricing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_bills_ocr_and_annotation_pages_at_their_own_rates", + "status": "unmapped", + "reason": "cost-calc: mixed-rate billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_bills_annotation_only_response", + "status": "unmapped", + "reason": "cost-calc: annotation-only billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_bills_annotation_pages_when_pages_processed_missing", + "status": "unmapped", + "reason": "cost-calc: fallback billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate", + "status": "unmapped", + "reason": "cost-calc: fallback billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_rust_ocr_serves_default_format", + "status": "unmapped", + "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_rust_ocr_skipped_for_native_format", + "status": "unmapped", + "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_native_format_rejected_for_provider_without_support_as_bad_request", + "status": "unmapped", + "reason": "provider-support validation for req_format happens in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_unknown_format_rejected_for_provider_without_support_as_bad_request", + "status": "unmapped", + "reason": "req_format validation error path is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestHandlerDiscovery::test_handler_discovered_for_ocr", + "status": "unmapped", + "reason": "guardrail-translation handler discovery is a Python proxy-layer concern" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestHandlerDiscovery::test_handler_discovered_for_aocr", + "status": "unmapped", + "reason": "guardrail-translation handler discovery is a Python proxy-layer concern" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_document_url", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_image_url", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_no_document", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_invalid_document", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_input_blocking_guardrail", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_single_page", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_multiple_pages", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_empty_pages", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_page_with_empty_markdown", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_preserves_page_metadata", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_output_blocking_guardrail", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestPIIMaskingScenario::test_pii_masking_in_ocr_pages", + "status": "unmapped", + "reason": "PII redaction in the translation handler has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_read_req_format_from_header", + "status": "unmapped", + "reason": "proxy-layer header parsing has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_prefer_body_req_format_over_header", + "status": "unmapped", + "reason": "proxy-layer body-vs-header precedence has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_omit_req_format_when_header_absent", + "status": "unmapped", + "reason": "proxy-layer parsing has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_reject_unknown_req_format", + "status": "unmapped", + "reason": "proxy-layer validation has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_return_native_payload_with_litellm_response_headers", + "status": "unmapped", + "reason": "proxy-layer response construction has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_return_normalized_response_when_no_native_payload", + "status": "unmapped", + "reason": "proxy-layer response construction has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_explicit_false_overrides_process_enable", + "status": "unmapped", + "reason": "bridge-plumbing: asserts the explicit-per-request rust:False override wins over the Python process flag, no Rust-owned behavior runs" + } ], "rust_only_tests": [ - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "rust_custom_logger_reads_failure_payload_for_non_ocr_call_type", "reason": "exercises the non-OCR (acompletion) call-type branch of the logger; the OCR branch is covered separately by rust_custom_logger_reads_success_payload_for_ocr"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "no_callback_fast_path_dispatches_nothing", "reason": "Rust-only fast-path optimization test for when zero callbacks are registered; Python has no equivalent no-op dispatch path"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "with_standard_logging_payload_keeps_top_level_fields_in_sync", "reason": "Rust-internal builder-method invariant, Python has no equivalent internal builder"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "blocks_private_and_metadata_ips", "reason": "SSRF IP-blocking helper has no Python unit test; Python relies on the proxy-layer JSON/form guards instead"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "convert_document_url_rejects_loopback_fetch", "reason": "URL-fetch SSRF protection is Rust-gateway-only"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "convert_document_url_leaves_data_uri_untouched", "reason": "URL-fetch SSRF protection is Rust-gateway-only"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_passes_short_strings_through", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_caps_long_payloads", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_does_not_split_multibyte_chars", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_dispatch_supports_migrated_providers", "reason": "Rust-internal provider-config dispatch table has no equivalent Python unit test"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "string_headers_accepts_string_values", "reason": "Rust-gateway header-coercion helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "auth_header_detection_is_case_insensitive", "reason": "Rust-gateway header-detection helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_runs_pre_during_and_success_hooks", "reason": "full gateway-level guardrail-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_runs_failure_hook_on_provider_error", "reason": "full gateway-level failure-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_pre_call_block_skips_provider_socket", "reason": "full gateway-level pre-call-block-plus-socket-skip test with no Python equivalent at this integration scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_does_not_duplicate_authorization_header_when_header_is_supplied", "reason": "outgoing HTTP header dedup at the Rust gateway has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "document_intelligence_poll_uses_resolved_subscription_key", "reason": "full Azure DI poll-loop integration test with no Python equivalent at this scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "string_headers_rejects_non_string_values", "reason": "Rust-gateway header-coercion error path has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "azure_ai_reuses_mistral_body_transform", "reason": "Rust-internal delegation-to-Mistral-transform implementation detail, no Python test asserts this delegation"}, - {"rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_request_uses_base64_source_for_data_uri", "reason": "no Python test asserts on the base64Source request body shape"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_mistral_url_uses_project_location_and_model", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_mistral_reuses_mistral_body_transform", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_deepseek_request_uses_ocr_endpoint_shape", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_deepseek_response_wraps_markdown_content", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_rejects_non_object_document", "reason": "non-object document rejection has no dedicated Python unit test"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_response_normalizes_mistral_json", "reason": "Python's response tests target OCR4-specific fields only, none asserts the same base normalization this Rust test checks"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "complete_url_defaults_and_dedupes_v1", "reason": "URL-building/defaulting for Mistral has no Python unit test"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "resolve_api_key_prefers_param_then_env", "reason": "API key resolution precedence at the Rust provider-config layer has no Python unit test"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "resolve_api_key_errors_when_absent", "reason": "API key resolution error path at the Rust provider-config layer has no Python unit test"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "rust_custom_logger_reads_success_payload_for_ocr", "reason": "Rust-internal custom-logger dispatch for OCR payloads has no Python unit test at this layer"} + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "rust_custom_logger_reads_failure_payload_for_non_ocr_call_type", + "reason": "exercises the non-OCR (acompletion) call-type branch of the logger; the OCR branch is covered separately by rust_custom_logger_reads_success_payload_for_ocr" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "no_callback_fast_path_dispatches_nothing", + "reason": "Rust-only fast-path optimization test for when zero callbacks are registered; Python has no equivalent no-op dispatch path" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "with_standard_logging_payload_keeps_top_level_fields_in_sync", + "reason": "Rust-internal builder-method invariant, Python has no equivalent internal builder" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "rust_test": "blocks_private_and_metadata_ips", + "reason": "SSRF IP-blocking helper has no Python unit test; Python relies on the proxy-layer JSON/form guards instead" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "rust_test": "convert_document_url_rejects_loopback_fetch", + "reason": "URL-fetch SSRF protection is Rust-gateway-only" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "rust_test": "convert_document_url_leaves_data_uri_untouched", + "reason": "URL-fetch SSRF protection is Rust-gateway-only" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "truncate_error_body_passes_short_strings_through", + "reason": "Rust-gateway error-body truncation helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "truncate_error_body_caps_long_payloads", + "reason": "Rust-gateway error-body truncation helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "truncate_error_body_does_not_split_multibyte_chars", + "reason": "Rust-gateway error-body truncation helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_dispatch_supports_migrated_providers", + "reason": "Rust-internal provider-config dispatch table has no equivalent Python unit test" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "string_headers_accepts_string_values", + "reason": "Rust-gateway header-coercion helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "auth_header_detection_is_case_insensitive", + "reason": "Rust-gateway header-detection helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_lifecycle_runs_pre_during_and_success_hooks", + "reason": "full gateway-level guardrail-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_lifecycle_runs_failure_hook_on_provider_error", + "reason": "full gateway-level failure-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_lifecycle_pre_call_block_skips_provider_socket", + "reason": "full gateway-level pre-call-block-plus-socket-skip test with no Python equivalent at this integration scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_does_not_duplicate_authorization_header_when_header_is_supplied", + "reason": "outgoing HTTP header dedup at the Rust gateway has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "document_intelligence_poll_uses_resolved_subscription_key", + "reason": "full Azure DI poll-loop integration test with no Python equivalent at this scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "string_headers_rejects_non_string_values", + "reason": "Rust-gateway header-coercion error path has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "azure_ai_reuses_mistral_body_transform", + "reason": "Rust-internal delegation-to-Mistral-transform implementation detail, no Python test asserts this delegation" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_request_uses_base64_source_for_data_uri", + "reason": "no Python test asserts on the base64Source request body shape" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_mistral_url_uses_project_location_and_model", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_mistral_reuses_mistral_body_transform", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_deepseek_request_uses_ocr_endpoint_shape", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_deepseek_response_wraps_markdown_content", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_request_rejects_non_object_document", + "reason": "non-object document rejection has no dedicated Python unit test" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_response_normalizes_mistral_json", + "reason": "Python's response tests target OCR4-specific fields only, none asserts the same base normalization this Rust test checks" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "complete_url_defaults_and_dedupes_v1", + "reason": "URL-building/defaulting for Mistral has no Python unit test" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "resolve_api_key_prefers_param_then_env", + "reason": "API key resolution precedence at the Rust provider-config layer has no Python unit test" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "resolve_api_key_errors_when_absent", + "reason": "API key resolution error path at the Rust provider-config layer has no Python unit test" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "rust_custom_logger_reads_success_payload_for_ocr", + "reason": "Rust-internal custom-logger dispatch for OCR payloads has no Python unit test at this layer" + } ] } From 7f7e0d55178c38a5800462c91425a8240b28ed69 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:32:00 -0700 Subject: [PATCH 10/14] fix(vector-store): embed through the SDK when the Router does not serve the query embedding model The Router executor only routed a query embedding when the vector store carried extra embedding configuration, so a store registered with no embedding model at all always went to the Router and 500'd on the s3_vectors default text-embedding-3-small when no deployment served it. Route on whether the Router serves the model, which is the rule the executor had before, and keep the request metadata on the SDK fallback so the embedding stays attributed either way. --- .../base_llm/vector_store/transformation.py | 7 ++-- .../test_bedrock_knowledgebase_hook.py | 2 +- .../test_router_embedding_integration.py | 13 ++++---- .../test_s3_vectors_transformation.py | 33 +++++++++++++++---- 4 files changed, 36 insertions(+), 19 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index a3b8bcc499c..c8d2b7fe522 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -99,12 +99,9 @@ class RouterVectorStoreEmbeddingExecutor: ) return bool(resolved) or model in deployment_models - def _embeds_through_sdk(self, model: str, configuration: Mapping[str, object]) -> bool: - return bool(configuration) and not self._router_serves(model) - def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: embedding_kwargs: Final = self._embedding_kwargs(configuration) - if self._embeds_through_sdk(model, configuration): + if not self._router_serves(model): return LiteLLMVectorStoreEmbeddingExecutor().embed(model, query, embedding_kwargs) return self.router.embedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, @@ -114,7 +111,7 @@ class RouterVectorStoreEmbeddingExecutor: async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: embedding_kwargs: Final = self._embedding_kwargs(configuration) - if self._embeds_through_sdk(model, configuration): + if not self._router_serves(model): return await LiteLLMVectorStoreEmbeddingExecutor().aembed(model, query, embedding_kwargs) return await self.router.aembedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 98045725177..044e1e0de8e 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -375,7 +375,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters( timeout=None, client=None, _is_async=False, - router: "litellm.Router | None" = None, + embedding_executor=None, ): litellm_params_dict = ( litellm_params.model_dump(exclude_none=False) diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index 2cc9914c9b3..e10ba0f0962 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -187,7 +187,7 @@ class TestRouterEmbeddingIntegration: assert _sent(store_route, 1) == ("Bearer store-key", "text-embedding-3-large", ["async query"]) @pytest.mark.asyncio - async def test_router_executor_rejects_unserved_models_without_explicit_config( + async def test_router_executor_embeds_unserved_models_through_the_sdk( self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -198,12 +198,13 @@ class TestRouterEmbeddingIntegration: metadata={"user_api_key_team_id": "team-a"}, ) - with pytest.raises(litellm.BadRequestError): - executor.embed("openai/text-embedding-3-large", "sync query", {}) - with pytest.raises(litellm.BadRequestError): - await executor.aembed("openai/text-embedding-3-large", "async query", {}) + sync_response = executor.embed("text-embedding-3-large", "sync query", {}) + async_response = await executor.aembed("text-embedding-3-large", "async query", {}) - assert openai_route.call_count == 0 + assert sync_response.data[0]["embedding"] == QUERY_VECTOR + assert async_response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(openai_route, 0) == ("Bearer env-key", "text-embedding-3-large", ["sync query"]) + assert _sent(openai_route, 1) == ("Bearer env-key", "text-embedding-3-large", ["async query"]) def test_router_executor_routes_deployment_model_names_through_the_router( self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index e2b02ea2151..e313b749d06 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -21,8 +21,6 @@ def _embedding_response(vector): class _RecordingExecutor: - """Executor double recording every (model, query, configuration) it was asked to embed.""" - def __init__(self, vector=QUERY_VECTOR): self.vector = vector self.calls = [] @@ -94,7 +92,6 @@ class TestS3VectorsVectorStoreConfig: config.get_complete_url(None, {"aws_region_name": "Bad_Region!"}) def test_transform_search_request(self): - """Full request-body transformation with the query embedded through the injected executor""" config = S3VectorsVectorStoreConfig() logging_obj = _logging_obj() executor = _RecordingExecutor() @@ -137,7 +134,6 @@ class TestS3VectorsVectorStoreConfig: @pytest.mark.asyncio async def test_atransform_search_embeds_alias_and_store_config_through_executor(self): - """The store's embedding_model alias and litellm_embedding_config reach the executor unchanged""" config = S3VectorsVectorStoreConfig() executor = _RecordingExecutor(vector=[0.4, 0.5]) @@ -177,9 +173,34 @@ class TestS3VectorsVectorStoreConfig: ) assert request_body["queryVector"]["float32"] == QUERY_VECTOR + @pytest.mark.asyncio + async def test_atransform_search_default_model_falls_back_to_the_sdk(self): + """Regression (LIT-6750): a store that never named an embedding model keeps working on a proxy + whose model list has no text-embedding-3-small, embedding through the SDK instead of erroring.""" + config = S3VectorsVectorStoreConfig() + router = MagicMock() + router.get_model_list.return_value = [ + {"model_name": "team-embeddings", "litellm_params": {"model": "openai/text-embedding-3-small"}} + ] + router.resolved_litellm_models.return_value = [] + router.aembedding = AsyncMock(side_effect=AssertionError("unserved model must not reach the Router")) + request_metadata = {"user_api_key_team_id": "team-a"} + + mock_bare = AsyncMock(return_value=_embedding_response(QUERY_VECTOR)) + with patch("litellm.aembedding", new=mock_bare): # test-quality-ok: stubs the bare-embedding fallback whose call the test asserts on + _, request_body = await config.atransform_search_vector_store_request( + **_search_kwargs( + embedding_executor=RouterVectorStoreEmbeddingExecutor(router=router, metadata=request_metadata) + ) + ) + + mock_bare.assert_awaited_once_with( + model="text-embedding-3-small", input=["test query"], metadata=request_metadata + ) + assert request_body["queryVector"]["float32"] == QUERY_VECTOR + @pytest.mark.asyncio async def test_atransform_search_without_executor_uses_bare_embedding(self): - """Backward compat: SDK callers without an executor keep embedding through litellm.aembedding""" config = S3VectorsVectorStoreConfig() mock_bare = AsyncMock(return_value=_embedding_response([0.6, 0.7])) @@ -190,7 +211,6 @@ class TestS3VectorsVectorStoreConfig: assert request_body["queryVector"]["float32"] == [0.6, 0.7] def test_transform_search_without_executor_uses_bare_embedding_sync(self): - """Sync twin: no executor -> bare litellm.embedding as before""" config = S3VectorsVectorStoreConfig() mock_bare = MagicMock(return_value=_embedding_response([0.8, 0.9])) @@ -203,7 +223,6 @@ class TestS3VectorsVectorStoreConfig: assert request_body["queryVector"]["float32"] == [0.8, 0.9] def test_transform_search_request_invalid_vector_store_id(self): - """An unparseable vector_store_id raises before any embedding is generated""" config = S3VectorsVectorStoreConfig() executor = _RecordingExecutor() From f62e87d28c0bd82da95479db188aa2309384583d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:42:47 -0700 Subject: [PATCH 11/14] fix(ui): read the preset catalog through the shared mock in the lib test --- ui/litellm-dashboard/src/lib/autorouter_presets.test.ts | 5 ++--- ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts | 6 +++--- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index 44ac6973a1c..76996447ac8 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -1,9 +1,8 @@ import { describe, it, expect } from "vitest"; -import bundledPresets from "../../../../litellm/proxy/public_endpoints/autorouter_presets.json"; +import { BUNDLED_PRESETS_RESPONSE } from "../../tests/mocks/autoRouterPresets"; import { hydratePresets, AutoRouterPreset, - AutoRouterPresetsResponse, getRequiredModelsInPreset, getMissingModelsInPreset, getRequiredModels, @@ -21,7 +20,7 @@ import { DEFAULT_ESCALATION_KEYWORDS } from "@/components/add_model/EscalationKe const groupsOnly = (models: Iterable) => buildModelAvailability(models, []); // Hydrated from the real bundled catalog so a catalog edit flows into these expectations. -const PRESETS = hydratePresets(bundledPresets as AutoRouterPresetsResponse); +const PRESETS = hydratePresets(BUNDLED_PRESETS_RESPONSE); const getAllPresets = (): AutoRouterPreset[] => PRESETS; const getPresetByKey = (key: string): AutoRouterPreset | undefined => PRESETS.find((p) => p.key === key); diff --git a/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts index cff417e6bfa..ac6da9bceba 100644 --- a/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts +++ b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts @@ -7,9 +7,9 @@ import { hydratePresets, type AutoRouterPresetsResponse } from "@/lib/autorouter // instead of redding on a stale copy. Exported as vi.fn so a test can override the query state. const CATALOG_PATH = resolve(__dirname, "../../../../litellm/proxy/public_endpoints/autorouter_presets.json"); -export const BUNDLED_PRESETS = hydratePresets( - JSON.parse(readFileSync(CATALOG_PATH, "utf8")) as AutoRouterPresetsResponse, -); +export const BUNDLED_PRESETS_RESPONSE = JSON.parse(readFileSync(CATALOG_PATH, "utf8")) as AutoRouterPresetsResponse; + +export const BUNDLED_PRESETS = hydratePresets(BUNDLED_PRESETS_RESPONSE); export const LOADED_PRESETS_QUERY = { data: BUNDLED_PRESETS, From 534003da03a90a9f9bab5f7ed319ac4b5016fcb4 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 2 Sep 2026 22:33:14 -0700 Subject: [PATCH 12/14] feat(ui): add 1M context auto-router preset (#39490) * feat(ui): add 1M context auto-router preset * feat(ui): use heuristic v2 for 1M preset * fix(ui): keep 1M preset test within lint budget --- .../public_endpoints/autorouter_presets.json | 22 +++++++++++++++ .../public_endpoints/test_public_endpoints.py | 10 +++++++ .../add_model/add_auto_router_tab.test.tsx | 18 +++++++++++-- .../src/lib/autorouter_presets.test.ts | 27 ++++++++++++++++++- 4 files changed, 74 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/public_endpoints/autorouter_presets.json b/litellm/proxy/public_endpoints/autorouter_presets.json index c2b13b81542..6a6642d4211 100644 --- a/litellm/proxy/public_endpoints/autorouter_presets.json +++ b/litellm/proxy/public_endpoints/autorouter_presets.json @@ -1,4 +1,26 @@ { + "1m_context": { + "label": "1M Context", + "description": "Routes across models with 1M-token context windows: Luna for simple queries, Terra for medium, Opus 5 for complex, Opus 5 at high thinking for reasoning.", + "complexity_router_config": { + "tiers": { + "SIMPLE": ["gpt-5.6-luna"], + "MEDIUM": ["gpt-5.6-terra"], + "COMPLEX": ["claude-opus-5"], + "REASONING": ["claude-opus-5"] + }, + "tier_model_configs": { + "REASONING": [{ "model_name": "claude-opus-5", "litellm_params": { "reasoning_effort": "high" } }] + }, + "classifier_type": "heuristic_v2", + "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", + "session_affinity": false, + "modality_routing": false, + "modality_pin_override": false, + "deployment_affinity": true + } + }, "anthropic_family": { "label": "Anthropic Family", "description": "Routes across the Claude model family: Haiku for simple queries, Sonnet for medium, Opus for complex, Opus at high thinking for reasoning.", diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 551e27a18f5..41439f28638 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1104,6 +1104,16 @@ def test_get_autorouter_presets_local_mode_serves_bundled_catalog( assert response.status_code == 200 payload = response.json() assert "anthropic_family" in payload + assert payload["1m_context"]["complexity_router_config"]["classifier_type"] == "heuristic_v2" + assert payload["1m_context"]["complexity_router_config"]["tiers"] == { + "SIMPLE": ["gpt-5.6-luna"], + "MEDIUM": ["gpt-5.6-terra"], + "COMPLEX": ["claude-opus-5"], + "REASONING": ["claude-opus-5"], + } + assert payload["1m_context"]["complexity_router_config"]["tier_model_configs"] == { + "REASONING": [{"model_name": "claude-opus-5", "litellm_params": {"reasoning_effort": "high"}}] + } for preset in payload.values(): assert isinstance(preset["label"], str) assert isinstance(preset["description"], str) diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 27a169524b9..836b4c6ff3d 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -636,7 +636,14 @@ describe("AddAutoRouterTab", () => { const labels = visibleOptions().map((option) => option.querySelector(".font-medium")?.textContent); - expect(labels).toEqual(["Anthropic Family", "Gemini Family", "Lite", "OpenAI Family", "Custom Configuration"]); + expect(labels).toEqual([ + "1M Context", + "Anthropic Family", + "Gemini Family", + "Lite", + "OpenAI Family", + "Custom Configuration", + ]); }); describe("routing test", () => { @@ -1060,7 +1067,14 @@ describe("AddAutoRouterTab", () => { expect(isOptionDisabled(optionByLabel("Anthropic Family")!)).toBe(false); }); const labels = visibleOptions().map((option) => option.querySelector(".font-medium")?.textContent); - expect(labels).toEqual(["Anthropic Family", "Gemini Family", "Lite", "OpenAI Family", "Custom Configuration"]); + expect(labels).toEqual([ + "Anthropic Family", + "1M Context", + "Gemini Family", + "Lite", + "OpenAI Family", + "Custom Configuration", + ]); }); it.each([ diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index 44ac6973a1c..a7333e4704b 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -28,7 +28,13 @@ const getPresetByKey = (key: string): AutoRouterPreset | undefined => PRESETS.fi describe("autorouter_presets", () => { it("hydrates exactly the bundled presets", () => { const presets = getAllPresets(); - expect(presets.map((p) => p.label).sort()).toEqual(["Anthropic Family", "Gemini Family", "Lite", "OpenAI Family"]); + expect(presets.map((p) => p.label).sort()).toEqual([ + "1M Context", + "Anthropic Family", + "Gemini Family", + "Lite", + "OpenAI Family", + ]); // Every preset carries all four fields the UI relies on; a JSON typo dropping one fails here. for (const p of presets) { expect(p).toMatchObject({ key: expect.any(String), label: expect.any(String), description: expect.any(String) }); @@ -202,6 +208,25 @@ describe("autorouter_presets", () => { }); }); + it("pins the 1M context preset to Luna, Terra, and Opus at high thinking", () => { + const preset = getPresetByKey("1m_context")!; + const expectedTiers = { + SIMPLE: ["gpt-5.6-luna"], + MEDIUM: ["gpt-5.6-terra"], + COMPLEX: ["claude-opus-5"], + REASONING: ["claude-opus-5"], + }; + expect(preset.complexity_router_config.classifier_type).toBe("heuristic_v2"); + expect(preset.complexity_router_config.tiers).toEqual(expectedTiers); + expect(preset.complexity_router_config.tier_model_configs).toEqual({ + REASONING: [{ model_name: "claude-opus-5", litellm_params: { reasoning_effort: "high" } }], + }); + const prefill = buildPresetPrefill(preset.complexity_router_config, groupsOnly(getRequiredModelsInPreset(preset))); + expect(prefill.complexityRouterConfig.tier_model_params).toEqual({ + REASONING: { "claude-opus-5": { reasoning_effort: "high" } }, + }); + }); + it("pins the gemini preset to concrete model ids, never Google's hot-swapping -latest aliases", () => { const gemini = getPresetByKey("gemini_family")!; const config = gemini.complexity_router_config; From c841a56e9ab930e86ba6c484444f4e56eb0d7bab Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 2 Sep 2026 22:34:15 -0700 Subject: [PATCH 13/14] fix(ui): stop the create team form resetting organization and models (#39476) * fix(ui): stop the create team form resetting organization and models The organization preselect ran in an effect keyed on the organizations query, so any refetch of that list while the Create Team modal was open overwrote the user's organization pick, which in turn cleared their models pick. The models field was also cleared whenever the available models fetch resolved. Preselect the organization when the modal opens instead, and clear the models only when the user picks a different organization. An org admin whose admin orgs narrow to one while the form is open can still pick, rather than facing a locked empty field. * fix(ui): block team create when the picked organization is no longer available An organization picked in the Create Team form now survives a refetch of the organization list, so it can outlive the admin's access to it. Refuse the create with a message on the field rather than letting the request fail authorization at the proxy. * fix(ui): keep the team create organization field usable when the pick goes stale Locking the field on a single admin organization also locked it while it held a rejected organization, so an admin who lost access could not pick the one organization left. Lock it only while it holds that organization. * test(ui): hoist the created team fixture out of the mock call The inline object pushed the repo past its no-large-inline-object-arg lint budget, which has no headroom. --- ui/litellm-dashboard/eslint-suppressions.json | 3 - .../src/components/Teams.test.tsx | 202 ++++++++++++++++++ ui/litellm-dashboard/src/components/Teams.tsx | 74 ++++--- 3 files changed, 243 insertions(+), 36 deletions(-) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 7de7373b20b..7d7066addcf 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1399,9 +1399,6 @@ }, "prefer-const": { "count": 2 - }, - "react-hooks/set-state-in-effect": { - "count": 1 } }, "src/components/TeamsPage/teamTableColumns.tsx": { diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index bee35ac5c5d..2bceb00aae1 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -15,6 +15,7 @@ import { teamCreateCall, } from "./networking"; import Teams from "./Teams"; +import { chooseSelectOption } from "../../tests/test-utils"; const can = vi.fn(); vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ @@ -1488,3 +1489,204 @@ describe("Teams - the exact bytes the create call sends", () => { expect(teamCreateCall).not.toHaveBeenCalled(); }); }); + +describe("Teams - the create form keeps the organization and models picks while it is open", () => { + const ORGS = [ + { organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }, + { organization_id: "org-2", organization_alias: "Org 2", models: [], members: [] }, + ]; + + const orgField = () => screen.getByRole("combobox", { name: /organization/i }); + const modelsField = () => screen.getByTestId("create-team-models-select"); + + const openCreateModal = async () => { + act(() => { + fireEvent.click(screen.getAllByRole("button", { name: /create team/i })[0]); + }); + await screen.findByLabelText(/team name/i); + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(getDefaultTeamSettings).mockResolvedValue({ values: {} }); + mockUseOrganizations.mockReturnValue({ data: ORGS }); + }); + + it("keeps both picks when the organizations list comes back changed from a refetch", async () => { + const user = userEvent.setup(); + renderWithQueryClient(); + await openCreateModal(); + + await chooseSelectOption(user, orgField(), /Org 1/); + fireEvent.change(modelsField(), { target: { value: "gpt-4" } }); + + mockUseOrganizations.mockReturnValue({ data: ORGS.map((org) => ({ ...org, spend: 1 })) }); + fireEvent.click(screen.getByText("Additional Settings")); + + expect(orgField()).toHaveValue("Org 1"); + expect(modelsField()).toHaveValue("gpt-4"); + }); + + it("keeps models picked before the available models finish loading", async () => { + let resolveModels: (models: string[]) => void = () => {}; + vi.mocked(fetchAvailableModelsForTeamOrKey).mockReturnValue( + new Promise((resolve) => { + resolveModels = resolve; + }), + ); + renderWithQueryClient(); + await openCreateModal(); + + fireEvent.change(modelsField(), { target: { value: "gpt-4" } }); + await act(async () => { + resolveModels(["gpt-4", "gpt-3.5-turbo"]); + }); + + expect(modelsField()).toHaveValue("gpt-4"); + }); + + it("clears the models pick when the organization is changed, since models are org scoped", async () => { + const user = userEvent.setup(); + renderWithQueryClient(); + await openCreateModal(); + + await chooseSelectOption(user, orgField(), /Org 1/); + fireEvent.change(modelsField(), { target: { value: "gpt-4" } }); + await chooseSelectOption(user, orgField(), /Org 2/); + + await waitFor(() => expect(orgField()).toHaveValue("Org 2")); + expect(modelsField()).toHaveValue(""); + }); + + it("keeps the models pick when the same organization is chosen again", async () => { + const user = userEvent.setup(); + renderWithQueryClient(); + await openCreateModal(); + + await chooseSelectOption(user, orgField(), /Org 1/); + fireEvent.change(modelsField(), { target: { value: "gpt-4" } }); + await chooseSelectOption(user, orgField(), /Org 1/); + + expect(orgField()).toHaveValue("Org 1"); + expect(modelsField()).toHaveValue("gpt-4"); + }); + + it("still preselects the only organization an org admin can create teams in", async () => { + mockUseOrganizations.mockReturnValue({ + data: [ + { + organization_id: "org-1", + organization_alias: "Org 1", + models: [], + members: [{ user_id: "user-123", user_role: "org_admin" }], + }, + ], + }); + renderWithQueryClient(); + await openCreateModal(); + + expect(orgField()).toHaveValue("Org 1"); + expect(orgField()).toBeDisabled(); + }); + + it("leaves an org admin able to pick when their admin orgs narrow to one while the form is open", async () => { + const orgAdminOrgs = [ + { + organization_id: "org-1", + organization_alias: "Org 1", + models: [], + members: [{ user_id: "user-123", user_role: "org_admin" }], + }, + { + organization_id: "org-2", + organization_alias: "Org 2", + models: [], + members: [{ user_id: "user-123", user_role: "org_admin" }], + }, + ]; + mockUseOrganizations.mockReturnValue({ data: orgAdminOrgs }); + renderWithQueryClient(); + await openCreateModal(); + expect(orgField()).toHaveValue(""); + + mockUseOrganizations.mockReturnValue({ data: [orgAdminOrgs[0]] }); + fireEvent.click(screen.getByText("Additional Settings")); + + expect(orgField()).toBeEnabled(); + }); + + it("refuses to create the team in an organization the admin has lost access to", async () => { + const user = userEvent.setup(); + const orgAdminOrgs = ORGS.map((org) => ({ ...org, members: [{ user_id: "user-123", user_role: "org_admin" }] })); + mockUseOrganizations.mockReturnValue({ data: orgAdminOrgs }); + renderWithQueryClient(); + await openCreateModal(); + + fireEvent.change(screen.getByTestId("team-name-input"), { target: { value: "Revoked Team" } }); + await chooseSelectOption(user, orgField(), /Org 1/); + + mockUseOrganizations.mockReturnValue({ data: [orgAdminOrgs[1]] }); + fireEvent.click(screen.getByText("Additional Settings")); + + const submitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(submitButtons[submitButtons.length - 1]); + + await screen.findByText(/no longer create teams in this organization/i); + expect(teamCreateCall).not.toHaveBeenCalled(); + }); + + it("lets the admin switch to the one organization left after losing access to their pick", async () => { + const user = userEvent.setup(); + const orgAdminOrgs = ORGS.map((org) => ({ ...org, members: [{ user_id: "user-123", user_role: "org_admin" }] })); + mockUseOrganizations.mockReturnValue({ data: orgAdminOrgs }); + const createdTeam = { + team_id: "new-team-1", + team_alias: "Recovered Team", + models: [], + organization_id: "org-2", + keys: [], + members_with_roles: [], + spend: 0, + }; + vi.mocked(teamCreateCall).mockResolvedValue(createdTeam); + renderWithQueryClient(); + await openCreateModal(); + + fireEvent.change(screen.getByTestId("team-name-input"), { target: { value: "Recovered Team" } }); + await chooseSelectOption(user, orgField(), /Org 1/); + + mockUseOrganizations.mockReturnValue({ data: [orgAdminOrgs[1]] }); + fireEvent.click(screen.getByText("Additional Settings")); + + expect(orgField()).toBeEnabled(); + await chooseSelectOption(user, orgField(), /Org 2/); + const submitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(submitButtons[submitButtons.length - 1]); + + await waitFor(() => + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ team_alias: "Recovered Team", organization_id: "org-2" }), + ), + ); + }); + + it("starts the form clean again when the modal is closed and reopened", async () => { + const user = userEvent.setup(); + renderWithQueryClient(); + await openCreateModal(); + + await chooseSelectOption(user, orgField(), /Org 1/); + fireEvent.change(modelsField(), { target: { value: "gpt-4" } }); + fireEvent.click(screen.getByRole("button", { name: /^close$/i })); + await waitFor(() => expect(screen.queryByLabelText(/team name/i)).not.toBeInTheDocument()); + + await openCreateModal(); + expect(orgField()).toHaveValue(""); + expect(modelsField()).toHaveValue(""); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index e2eda4adb23..ef58237a6aa 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -208,7 +208,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const queryClient = useQueryClient(); const refreshTeams = () => queryClient.invalidateQueries({ queryKey: teamsTableKeys.all }); const [currentOrg] = useState(null); - const [currentOrgForCreateTeam, setCurrentOrgForCreateTeam] = useState(null); const isOrgAdmin = userRole !== "Admin"; const [additionalSettingsOpen, setAdditionalSettingsOpen] = useState(false); @@ -216,17 +215,33 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const [agentSettingsOpen, setAgentSettingsOpen] = useState(false); const [searchToolSettingsOpen, setSearchToolSettingsOpen] = useState(false); + const adminOrgs = useMemo( + () => getAdminOrganizations(userRole, userID, organizations), + [userRole, userID, organizations], + ); + const teamCreateSchema = useMemo( () => teamCreateFieldsSchema.superRefine((values, ctx) => { if (isOrgAdmin && !values.organization_id) { ctx.addIssue({ code: "custom", message: SUPPRESSED_BY_DESCRIPTION, path: ["organization_id"] }); } + const organizationIsStillPickable = + values.organization_id == null || + organizations == null || + adminOrgs.some((org) => org.organization_id === values.organization_id); + if (!organizationIsStillPickable) { + ctx.addIssue({ + code: "custom", + message: "You can no longer create teams in this organization", + path: ["organization_id"], + }); + } if (additionalSettingsOpen && !isParsableJson(values.secret_manager_settings)) { ctx.addIssue({ code: "custom", message: SUPPRESSED_BY_DESCRIPTION, path: ["secret_manager_settings"] }); } }), - [isOrgAdmin, additionalSettingsOpen], + [isOrgAdmin, additionalSettingsOpen, adminOrgs, organizations], ); const form = useZodForm(teamCreateSchema, { defaultValues: EMPTY_TEAM_CREATE_VALUES }); @@ -264,28 +279,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser ? `Default: ${getBudgetDurationLabel(defaultBudgetDuration)} (${defaultBudgetDuration})` : "n/a"; - useEffect(() => { - form.setValue("models", []); - }, [currentOrgForCreateTeam, userModels]); - - // Handle organization preselection when modal opens - useEffect(() => { - if (isTeamModalVisible) { - const adminOrgs = getAdminOrganizations(userRole, userID, organizations); - - // Org admins must scope a team to an org, so with exactly one we preselect it. - // Proxy admins can create org-less teams, so the field stays optional regardless of org count. - if (isOrgAdmin && adminOrgs.length === 1) { - const org = adminOrgs[0]; - form.setValue("organization_id", org.organization_id); - setCurrentOrgForCreateTeam(org); - } else { - form.setValue("organization_id", currentOrg?.organization_id || null); - setCurrentOrgForCreateTeam(currentOrg); - } - } - }, [isTeamModalVisible, isOrgAdmin, userRole, userID, organizations, currentOrg]); - // Add this useEffect to fetch guardrails useEffect(() => { const fetchGuardrails = async () => { @@ -320,6 +313,26 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser if (canViewPolicies) fetchPolicies(); }, [accessToken, canViewPolicies]); + const openCreateTeamModal = () => { + // Org admins must scope a team to an org, so with exactly one we preselect it. + // Proxy admins can create org-less teams, so the field stays optional regardless of org count. + if (isOrgAdmin && adminOrgs.length === 1) { + form.setValue("organization_id", adminOrgs[0].organization_id); + } + setIsTeamModalVisible(true); + }; + + const selectCreateTeamOrganization = ( + next: string, + currentOrganizationId: string | null, + onChange: (organizationId: string | null) => void, + ) => { + const nextOrganizationId = next === "" ? null : next; + if (nextOrganizationId === currentOrganizationId) return; + onChange(nextOrganizationId); + form.setValue("models", []); + }; + const resetCreateForm = () => { form.reset(EMPTY_TEAM_CREATE_VALUES); setAdditionalSettingsOpen(false); @@ -636,7 +649,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser subtitle="Manage teams, members, and their access to models and budgets" primaryAction={ canCreateOrManageTeams(userRole, userID, organizations) ? ( - setIsTeamModalVisible(true)} data-testid="create-team-button"> + Create Team @@ -683,9 +696,9 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser )} {(() => { - const adminOrgs = getAdminOrganizations(userRole, userID, organizations); const isSingleOrg = adminOrgs.length === 1; const hasNoOrgs = adminOrgs.length === 0; + const soleOrganizationId = isSingleOrg ? adminOrgs[0].organization_id ?? null : null; return ( <> @@ -715,18 +728,13 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser label: org.organization_alias ?? "", sublabel: org.organization_id ?? "", }))} - disabled={isOrgAdmin && isSingleOrg} + disabled={isOrgAdmin && soleOrganizationId !== null && value === soleOrganizationId} allowClear={!isOrgAdmin} placeholder={ hasNoOrgs ? "No organizations available" : "Search or select an Organization" } emptyText="No organizations available" - onValueChange={(next) => { - onChange(next === "" ? null : next); - setCurrentOrgForCreateTeam( - adminOrgs.find((org) => org.organization_id === next) ?? null, - ); - }} + onValueChange={(next) => selectCreateTeamOrganization(next, value ?? null, onChange)} /> )} From 3cac5e5cd4c12a782e0afe96218aaff986ef3f60 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 2 Sep 2026 22:44:17 -0700 Subject: [PATCH 14/14] fix(sso): resolve multi-valued role claims to the highest privilege role (#39480) * fix(sso): resolve multi-valued role claims to the highest privilege role A role claim carrying several roles used to resolve to whichever one the IdP listed first, so a user holding both proxy_admin_viewer and internal_user lost org-level spend visibility depending on claim ordering alone. get_litellm_user_role now picks the highest privilege role out of a list-valued claim, and the Entra app_roles path shares that same resolution instead of keeping its own copy of the hierarchy. SAML assertions carrying several role values go through the same path rather than taking the first value. * test(sso): lock ranked-over-unranked role resolution for mixed claims org_admin, team and customer sit outside the privilege ladder. Pin the resolution for a claim that mixes one of them with a ranked role so the asymmetry is covered rather than implicit. * fix(sso): label the claim-sequence cast for the type-discipline gate * fix(sso): resolve claim entries without recursing The repo's recursive-function gate rejects self-recursion here, and a role claim is flat anyway. Pull the single-value lookup into its own helper so the list branch maps over it instead of calling back into itself. --- .../management_endpoints/sso/saml_sso.py | 6 +- litellm/proxy/management_endpoints/types.py | 62 +++++++-- litellm/proxy/management_endpoints/ui_sso.py | 19 +-- ruff-strict-budget.json | 4 +- .../management_endpoints/test_saml_sso.py | 33 +++++ .../proxy/management_endpoints/test_ui_sso.py | 129 +++++++++++++++++- type-discipline-budget.json | 4 +- 7 files changed, 218 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py index 466b100ea1f..12e1f1a03f3 100644 --- a/litellm/proxy/management_endpoints/sso/saml_sso.py +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -443,7 +443,9 @@ class SAMLAuthHandler: last_name: Final = SAMLAuthHandler._attribute_value( attributes, "SAML_ATTRIBUTE_LAST_NAME", _LAST_NAME_ATTRIBUTE_CANDIDATES ) - role_value = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES) + role_values: Final = SAMLAuthHandler._attribute_values( + attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES + ) team_ids: Final = SAMLAuthHandler._attribute_values( attributes, "SAML_ATTRIBUTE_TEAM_IDS", _TEAM_IDS_ATTRIBUTE_CANDIDATES ) @@ -464,7 +466,7 @@ class SAMLAuthHandler: picture=None, provider="saml", team_ids=team_ids, - user_role=get_litellm_user_role(role_value) if role_value else None, + user_role=get_litellm_user_role(role_values), ) except ValidationError as e: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 070df97d09a..4414eed97b2 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -4,12 +4,44 @@ Types for the management endpoints Might include fastapi/proxy requirements.txt related imports """ +from collections.abc import Iterable, Sequence from typing import Any, Final, cast from fastapi_sso.sso.base import OpenID from litellm.proxy._types import LitellmUserRoles +# Ordered highest to lowest privilege +LITELLM_USER_ROLE_HIERARCHY: Final = ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, +) + + +def highest_privilege_role(roles: Iterable[LitellmUserRoles]) -> LitellmUserRoles | None: + """ + Pick the highest privilege role out of the roles an IdP asserted for one user. + + IdPs do not guarantee ordering within a multi-valued role claim, so a user holding + several roles resolves to the most privileged one rather than whichever came first. + Roles the hierarchy does not rank (org_admin, team, customer) resolve by name to stay + deterministic. + + Args: + roles: The roles resolved from the claim + + Returns: + The highest privilege role, or None if `roles` is empty + """ + resolved: Final = frozenset(roles) + if not resolved: + return None + + ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None) + return ranked if ranked is not None else min(resolved, key=lambda role: role.value) + def is_valid_litellm_user_role(role_str: str) -> bool: """ @@ -28,12 +60,22 @@ def is_valid_litellm_user_role(role_str: str) -> bool: return False -def get_litellm_user_role(role_str) -> LitellmUserRoles | None: +def _role_from_claim_value(role_str: object) -> LitellmUserRoles | None: + if not isinstance(role_str, str): + return None + # Use _value2member_map_ for O(1) lookup, case-insensitive + result: Final = LitellmUserRoles._value2member_map_.get(role_str.lower()) + return cast(LitellmUserRoles | None, result) + + +def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None: """ Convert a string (or list of strings) to a LitellmUserRoles enum if valid (case-insensitive). Handles list inputs since some SSO providers (e.g., Keycloak) return roles - as arrays like ["proxy_admin"] instead of plain strings. + as arrays like ["proxy_admin"] instead of plain strings. A claim carrying several + roles resolves to the highest privilege one, so a user does not lose access just + because the IdP listed a weaker role first. Args: role_str: String or list to convert (e.g., "proxy_admin", ["proxy_admin"]) @@ -41,16 +83,12 @@ def get_litellm_user_role(role_str) -> LitellmUserRoles | None: Returns: LitellmUserRoles enum if valid, None otherwise """ - try: - if isinstance(role_str, list): - if len(role_str) == 0: - return None - role_str = role_str[0] - # Use _value2member_map_ for O(1) lookup, case-insensitive - result: Final = LitellmUserRoles._value2member_map_.get(role_str.lower()) - return cast(LitellmUserRoles | None, result) - except Exception: - return None + if isinstance(role_str, (list, tuple)): + entries: Final = cast(Sequence[object], role_str) # cast-ok: isinstance narrows the claim, not its elements + return highest_privilege_role( + role for role in (_role_from_claim_value(entry) for entry in entries) if role is not None + ) + return _role_from_claim_value(role_str) class CustomOpenID(OpenID): diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 1feefa5725d..84150ef7935 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -112,6 +112,7 @@ from litellm.proxy.management_endpoints.sso_helper_utils import ( ) from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add from litellm.proxy.management_endpoints.types import ( + LITELLM_USER_ROLE_HIERARCHY, CustomOpenID, get_litellm_user_role, is_valid_litellm_user_role, @@ -809,15 +810,6 @@ def normalize_email(email: str | None) -> str | None: return email.lower() if isinstance(email, str) else email -# Ordered highest to lowest privilege -LITELLM_USER_ROLE_HIERARCHY: Final = ( - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, -) - - def determine_role_from_groups( user_groups: list[str], role_mappings: "RoleMappings", @@ -4312,14 +4304,7 @@ class MicrosoftSSOHandler: listed first. Roles the hierarchy does not rank (org_admin, team, customer) resolve by name to stay deterministic """ - resolved: Final = frozenset( - role for role in (get_litellm_user_role(role_str) for role_str in app_roles or ()) if role is not None - ) - if not resolved: - return None - - ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None) - return ranked if ranked is not None else min(resolved, key=lambda role: role.value) + return get_litellm_user_role(tuple(app_roles or ())) @staticmethod def get_app_roles_from_id_token(id_token: str | None) -> list[str]: diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 4fcf650a8bc..be2b30fc189 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2985 + "limit": 2984 }, "ANN002": { "limit": 71 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2917 + "limit": 2916 }, "C401": { "limit": 8 diff --git a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py index 36d6414ba9f..635e6332958 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py @@ -516,6 +516,39 @@ async def test_team_ids_extracted_from_groups_attribute(saml_env_idp_initiated): assert result.team_ids == ["team-a", "team-b"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "roles", + [ + ["internal_user", "proxy_admin_viewer"], + ["proxy_admin_viewer", "internal_user"], + ], +) +async def test_multi_valued_role_attribute_resolves_to_highest_privilege(saml_env_idp_initiated, roles): + """An assertion carrying several roles must not depend on the order the IdP emitted them in.""" + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response( + key_pem, + cert_pem, + attributes={ + "email": ["dave@example.com"], + "role": roles, + }, + ) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_assertion_without_role_attribute_has_no_user_role(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem, attributes={"email": ["erin@example.com"]}) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.user_role is None + + @pytest.mark.asyncio async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env): cache = DualCache() diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index dd8c752a868..5dfff53f7c3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -6598,13 +6598,94 @@ def test_get_litellm_user_role_with_invalid_role(): assert result is None -def test_get_litellm_user_role_with_list_multiple_roles(): - """Test that get_litellm_user_role takes the first element from a multi-element list.""" +@pytest.mark.parametrize( + "role_claim", + [ + ["proxy_admin", "internal_user"], + ["internal_user", "proxy_admin"], + ], +) +def test_get_litellm_user_role_picks_highest_privilege_regardless_of_order(role_claim): + """A multi-valued role claim resolves to the most privileged role, not the first one listed.""" from litellm.proxy._types import LitellmUserRoles from litellm.proxy.management_endpoints.types import get_litellm_user_role - result = get_litellm_user_role(["proxy_admin", "internal_user"]) - assert result == LitellmUserRoles.PROXY_ADMIN + assert get_litellm_user_role(role_claim) == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["proxy_admin_viewer", "internal_user"], + ["internal_user", "proxy_admin_viewer"], + ], +) +def test_get_litellm_user_role_keeps_org_spend_visibility_for_mixed_roles(role_claim): + """ + Regression for LIT-6077: a user holding both proxy_admin_viewer and internal_user kept + losing org-level spend visibility whenever the IdP happened to list internal_user first. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + +def test_get_litellm_user_role_ignores_unrecognised_entries(): + """Roles LiteLLM does not know about are skipped rather than swallowing the whole claim.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(["some_idp_group", "internal_user"]) == LitellmUserRoles.INTERNAL_USER + assert get_litellm_user_role(["some_idp_group", "another_group"]) is None + + +def test_get_litellm_user_role_list_lookup_is_case_insensitive(): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(["INTERNAL_USER", "Proxy_Admin"]) == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["org_admin", "team"], + ["team", "org_admin"], + ], +) +def test_get_litellm_user_role_is_deterministic_for_unranked_roles(role_claim): + """Roles outside the privilege hierarchy still resolve the same way in either claim order.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.ORG_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["org_admin", "internal_user"], + ["internal_user", "org_admin"], + ], +) +def test_get_litellm_user_role_prefers_a_ranked_role_over_an_unranked_one(role_claim): + """ + org_admin, team and customer sit outside the privilege ladder, so a claim mixing one of + them with a ranked role settles on the ranked role in either order. Same rule the Entra + app_roles and role_mappings paths already follow. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.INTERNAL_USER + + +def test_get_litellm_user_role_returns_none_for_non_string_claims(): + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(None) is None + assert get_litellm_user_role({"role": "proxy_admin"}) is None # ============================================================================ @@ -6654,6 +6735,46 @@ def test_process_sso_jwt_access_token_extracts_role_from_access_token(): assert result.user_role == LitellmUserRoles.PROXY_ADMIN +@pytest.mark.parametrize( + "role_claim", + [ + ["internal_user", "proxy_admin_viewer"], + ["proxy_admin_viewer", "internal_user"], + ], +) +def test_process_sso_jwt_access_token_resolves_highest_privilege_role(role_claim): + """ + The generic SSO access-token path must land on the same role for a user whose role + claim holds several roles, whichever order the IdP emitted them in. + """ + import jwt as pyjwt + + from litellm.proxy._types import LitellmUserRoles + + access_token_str = pyjwt.encode( + {"sub": "user-123", "email": "mixed@test.com", "litellm_role": role_claim}, + "secret", + algorithm="HS256", + ) + result = CustomOpenID( + id="user-123", + email="mixed@test.com", + display_name="Mixed Role User", + team_ids=[], + user_role=None, + ) + + with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "litellm_role"}): + process_sso_jwt_access_token( + access_token_str=access_token_str, + sso_jwt_handler=None, + result=result, + role_mappings=None, + ) + + assert result.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + def test_process_sso_jwt_access_token_does_not_override_existing_role(): """ Test that process_sso_jwt_access_token does NOT override a role that was diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 5c25b8722ef..d5bf3883be4 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16478 + "limit": 16477 }, "LIT011": { - "limit": 5520 + "limit": 5519 }, "LIT012": { "limit": 4489