From 779441b47ba564696c35e86a66527e683c1fad31 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 12 Aug 2026 00:30:19 +0000 Subject: [PATCH 01/69] fix(bedrock_mantle): source per-request AWS credential params from litellm_params when signing chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/llm_http_handler.py | 26 ++++++- .../test_bedrock_mantle_transformation.py | 77 +++++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 19 +++++ 3 files changed, 120 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 721b9545ac1..dfc234d08e0 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5,7 +5,7 @@ import ssl from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from contextlib import asynccontextmanager from functools import lru_cache -from types import ModuleType +from types import MappingProxyType, ModuleType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, Union, cast, get_type_hints from urllib.parse import parse_qs, urlencode, urlparse, urlunparse @@ -20,6 +20,7 @@ from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -252,6 +253,24 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: return False +def _aws_signing_overrides( + optional_params: Mapping[str, Any], litellm_params: Mapping[str, Any] +) -> Mapping[str, Any]: + """AWS credential params for SigV4 signers that read them off optional_params. + + Only `bedrock`/`sagemaker` keep `aws_*` in optional_params: every other provider + spreads optional_params into the request body, so the params are stripped there + and survive on litellm_params alone. + """ + return MappingProxyType( + { + key: litellm_params[key] + for key in AWS_CREDENTIAL_KWARGS_KEYS + if optional_params.get(key) is None and litellm_params.get(key) is not None + } + ) + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -495,7 +514,10 @@ class BaseLLMHTTPHandler: headers, signed_json_body = provider_config.sign_request( headers=headers, - optional_params=optional_params, + optional_params={ + **optional_params, + **_aws_signing_overrides(optional_params, litellm_params), + }, request_data=data, api_base=api_base, api_key=api_key, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 275fb460b9f..468cc9b9130 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -489,6 +489,83 @@ class TestBedrockMantleChatAuth: assert "/us-east-2/bedrock/aws4_request" in authorization assert requests[0]["url"].startswith("https://bedrock-mantle.us-east-2.api.aws") + def test_completion_per_request_role_reaches_signer_and_not_the_body( + self, monkeypatch + ): + # Per-request aws_role_name/aws_session_name are stripped from optional_params + # for non-bedrock providers, so they must be sourced from litellm_params at + # signing time, and must never be serialized into the provider request body. + from botocore.credentials import Credentials + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + for var in ( + "BEDROCK_MANTLE_API_KEY", + "AWS_BEARER_TOKEN_BEDROCK", + "BEDROCK_MANTLE_API_BASE", + ): + monkeypatch.delenv(var, raising=False) + + credential_calls = [] + + def fake_get_credentials(self, **kwargs): + credential_calls.append(kwargs) + return Credentials( + access_key="ASIAEXAMPLE", + secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk", + token="assumed-session-token", + ) + + monkeypatch.setattr(BaseAWSLLM, "get_credentials", fake_get_credentials) + + requests = [] + + def mock_post(self, url, data=None, headers=None, **kwargs): + raw_body = data.decode("utf-8") if isinstance(data, bytes) else data + requests.append({"headers": headers or {}, "body": json.loads(raw_body or "{}")}) + return httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1733529600, + "model": "google.gemma-4-31b", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + request=httpx.Request("POST", url), + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post + ): + litellm.completion( + model="bedrock_mantle/google.gemma-4-31b", + messages=[{"role": "user", "content": "hello"}], + aws_role_name="arn:aws:iam::000000000000:role/attributed-role", + aws_session_name="user-123", + aws_region_name="us-east-1", + ) + + assert len(credential_calls) == 1 + assert ( + credential_calls[0]["aws_role_name"] + == "arn:aws:iam::000000000000:role/attributed-role" + ) + assert credential_calls[0]["aws_session_name"] == "user-123" + assert requests[0]["headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert not [key for key in requests[0]["body"] if key.startswith("aws_")] + class TestBedrockMantleProjectHeader: def test_validate_environment_sets_openai_project_header(self): diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index fddd8d09dfc..0906b39c514 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2071,3 +2071,22 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques retry_authorization = posts[1]["headers"]["Authorization"] assert retry_authorization.startswith("AWS4-HMAC-SHA256") assert retry_authorization != first_attempt_headers["Authorization"] + + +def test_aws_signing_overrides_only_fills_missing_credentials(): + from litellm.llms.custom_httpx.llm_http_handler import _aws_signing_overrides + + overrides = _aws_signing_overrides( + {"temperature": 0.2, "aws_region_name": "us-west-2"}, + { + "aws_role_name": "arn:aws:iam::000000000000:role/attributed", + "aws_session_name": "user-123", + "aws_region_name": "us-east-1", + "api_key": "not-an-aws-param", + }, + ) + + assert dict(overrides) == { + "aws_role_name": "arn:aws:iam::000000000000:role/attributed", + "aws_session_name": "user-123", + } From 2c7de60692d7a4fcd53964872d4355042d52b90b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 12 Aug 2026 00:53:40 +0000 Subject: [PATCH 02/69] style: ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/llm_http_handler.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index dfc234d08e0..cadf4c701e2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -253,9 +253,7 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: return False -def _aws_signing_overrides( - optional_params: Mapping[str, Any], litellm_params: Mapping[str, Any] -) -> Mapping[str, Any]: +def _aws_signing_overrides(optional_params: Mapping[str, Any], litellm_params: Mapping[str, Any]) -> Mapping[str, Any]: """AWS credential params for SigV4 signers that read them off optional_params. Only `bedrock`/`sagemaker` keep `aws_*` in optional_params: every other provider From 3275459aec0935b03b60950ba5d625854e2d6c9b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:47:26 -0700 Subject: [PATCH 03/69] fix(mcp): cap tools preview and test-connection at the listing timeout and name the unreachable upstream --- .../mcp_server/rest_endpoints.py | 34 +++++--- .../mcp_server/test_rest_endpoints.py | 77 +++++++++++++++++-- 2 files changed, 96 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3efb6429326..d73277f4417 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -4,10 +4,12 @@ from collections.abc import Awaitable, Callable, Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal +import anyio import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from litellm._logging import verbose_logger +from litellm.constants import MCP_TOOL_LISTING_TIMEOUT from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, @@ -68,7 +70,13 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( ) -def _connection_error_message(exc: BaseException) -> str: +def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + if isinstance(exc, TimeoutError): + return ( + f"Failed to connect to MCP server: no response from {url or 'the server'} " + f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " + "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." + ) if isinstance(exc, httpx.LocalProtocolError): return ( "Failed to connect to MCP server: a request header is malformed. " @@ -1136,6 +1144,7 @@ if MCP_AVAILABLE: mcp_auth_header: str | dict[str, str] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, + timeout_seconds: float = MCP_TOOL_LISTING_TIMEOUT, ) -> Mapping[str, object]: """ Create a temporary MCP client from *request*, run *operation*, and return the result. @@ -1151,6 +1160,10 @@ if MCP_AVAILABLE: oauth2_headers: Headers extracted from the incoming request (may contain the litellm API key — must NOT be forwarded for M2M servers). raw_headers: Raw request headers forwarded for stdio env construction. + timeout_seconds: Cap on OAuth discovery, connect, handshake, and *operation* + combined. Defaults to ``MCP_TOOL_LISTING_TIMEOUT`` (30s, below common LB + timeouts) so an unreachable upstream yields this endpoint's JSON error + instead of an opaque load-balancer 504 with an empty body. Returns: The dict returned by *operation*, or an error dict on failure. @@ -1240,15 +1253,16 @@ if MCP_AVAILABLE: static_headers=request.static_headers, ) - client: Final = await global_mcp_server_manager._create_mcp_client( - server=server_model, - mcp_auth_header=mcp_auth_header, - extra_headers=merged_headers, - stdio_env=stdio_env, - cred_provider=preview_cred_provider, - ) + with anyio.fail_after(timeout_seconds): + client: Final = await global_mcp_server_manager._create_mcp_client( + server=server_model, + mcp_auth_header=mcp_auth_header, + extra_headers=merged_headers, + stdio_env=stdio_env, + cred_provider=preview_cred_provider, + ) - return await operation(client) + return await operation(client) except (KeyboardInterrupt, SystemExit, asyncio.CancelledError): raise @@ -1257,7 +1271,7 @@ if MCP_AVAILABLE: return { "status": "error", "error": True, - "message": _connection_error_message(e), + "message": _connection_error_message(e, request.url, timeout_seconds), } async def _preview_openapi_tools(spec_path: str) -> dict: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index ef5631218f3..51b946c11b7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,4 +1,5 @@ import asyncio +import inspect import json import sys from datetime import datetime @@ -13,6 +14,7 @@ import pytest from fastapi import HTTPException from starlette.requests import Request +from litellm.constants import MCP_TOOL_LISTING_TIMEOUT from litellm.proxy._experimental.mcp_server import rest_endpoints from litellm.proxy._experimental.mcp_server.auth import ( user_api_key_auth_mcp as auth_mcp, @@ -109,6 +111,71 @@ class TestExecuteWithMcpClient: assert result["status"] == "error" assert "stack_trace" not in result + @pytest.mark.asyncio + async def test_timeout_caps_hanging_operation_and_names_url(self, monkeypatch): + async def fake_create_client(*args, **kwargs): + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + ) + + async def hanging_operation(client): + await asyncio.Event().wait() + + payload = NewMCPServerRequest( + server_name="example", + url="https://mcp.example.com/mcp/", + auth_type=MCPAuth.none, + ) + + result = await asyncio.wait_for( + rest_endpoints._execute_with_mcp_client(payload, hanging_operation, timeout_seconds=0.05), + timeout=5, + ) + + assert result["error"] is True + assert "https://mcp.example.com/mcp/" in result["message"] + + @pytest.mark.asyncio + async def test_timeout_covers_client_creation(self, monkeypatch): + async def hanging_create_client(*args, **kwargs): + await asyncio.Event().wait() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + hanging_create_client, + ) + + async def unreached_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="example", + url="https://mcp.example.com/mcp/", + auth_type=MCPAuth.none, + ) + + result = await asyncio.wait_for( + rest_endpoints._execute_with_mcp_client(payload, unreached_operation, timeout_seconds=0.05), + timeout=5, + ) + + assert result["error"] is True + assert "https://mcp.example.com/mcp/" in result["message"] + + def test_timeout_defaults_to_tool_listing_timeout(self): + default = inspect.signature(rest_endpoints._execute_with_mcp_client).parameters["timeout_seconds"].default + assert default == MCP_TOOL_LISTING_TIMEOUT + + def test_connection_error_message_timeout_names_url_and_budget(self): + message = rest_endpoints._connection_error_message(TimeoutError(), "https://api.example.com/mcp/", 30.0) + assert "https://api.example.com/mcp/" in message + assert "30s" in message + @pytest.mark.asyncio async def test_forwards_static_headers(self, monkeypatch): """Ensure static_headers are forwarded to the MCP client during test calls. @@ -2881,17 +2948,17 @@ class TestConnectionErrorMessage: secret = "Bearer sk-super-secret-token" exc = httpx.LocalProtocolError(f"Illegal header value b' {secret}'") - message = rest_endpoints._connection_error_message(exc) + message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "header" in message.lower() assert secret not in message def test_connect_error_points_at_reachability(self): - message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed")) + message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0) assert "unreachable" in message.lower() def test_timeout_error_message(self): - message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out")) + message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out"), "https://example.com", 30.0) assert "unreachable" in message.lower() def test_http_status_error_includes_status_code(self): @@ -2901,11 +2968,11 @@ class TestConnectionErrorMessage: request=httpx.Request("POST", "http://x/"), response=response, ) - message = rest_endpoints._connection_error_message(exc) + message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "503" in message def test_unknown_error_falls_back_to_generic(self): - message = rest_endpoints._connection_error_message(RuntimeError("weird")) + message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0) assert "weird" not in message assert "proxy logs" in message.lower() From f4b5449c6a65cf167658fd5bb32695abda9633a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:48:04 -0700 Subject: [PATCH 04/69] fix(openai_like): strip cache_control ttl before forwarding /v1/messages to non-Anthropic providers --- litellm/llms/anthropic/common_utils.py | 33 +++++ litellm/llms/openai_like/README.md | 5 +- .../openai_like/messages/transformation.py | 38 +++++- ..._like_anthropic_messages_transformation.py | 121 ++++++++++++++++++ 4 files changed, 195 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 681a8397f66..695e3a313ef 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1302,6 +1302,39 @@ def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format +def _normalized_cache_control(cache_control: dict) -> dict: # mutable-ok: as sibling sanitizers + cache_type: Final = cache_control.get("type") + return {"type": cache_type if isinstance(cache_type, str) else "ephemeral"} # mutable-ok: JSON wire format + + +def _normalize_cache_control_value(value: object) -> object: + if isinstance(value, dict): + return normalize_cache_control_in_anthropic_payload(value) + if isinstance(value, list): + return [_normalize_cache_control_value(item) for item in value] # mutable-ok: JSON wire format + return value + + +def normalize_cache_control_in_anthropic_payload(payload: dict) -> dict: # mutable-ok: as sibling sanitizers + """ + Return a copy of an Anthropic /v1/messages payload with every + ``cache_control`` entry reduced to ``{"type": }``, + recursing through message content blocks, system blocks, and tools. + + Anthropic itself accepts prompt-caching extensions such as ``ttl``, but + strict non-Anthropic implementations of the Messages API validate the field + literally and reject the whole request (``cache_control.ttl: 1h is not + supported``, ``cache_control.type is required``), which 400s clients like + Claude Code that always send cache hints. Non-dict ``cache_control`` values + are dropped entirely. The caller's payload is never mutated. + """ + return { # mutable-ok: JSON wire format, as sibling sanitizers + key: _normalized_cache_control(value) if key == "cache_control" else _normalize_cache_control_value(value) + for key, value in payload.items() + if key != "cache_control" or isinstance(value, dict) + } + + def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: openai_headers: Final = {} if "anthropic-ratelimit-requests-limit" in headers: diff --git a/litellm/llms/openai_like/README.md b/litellm/llms/openai_like/README.md index e9aaafe48a1..e1409b81c35 100644 --- a/litellm/llms/openai_like/README.md +++ b/litellm/llms/openai_like/README.md @@ -54,7 +54,10 @@ That's it! The provider will be automatically loaded and available. "constraints": { "temperature_max": 1.0, "temperature_min": 0.0, - "temperature_min_with_n_gt_1": 0.3 + "temperature_min_with_n_gt_1": 0.3, + // /v1/messages providers only: keep Anthropic cache_control extensions + // such as ttl instead of stripping them down to {"type": ...} + "cache_control_ttl": true }, // Optional: Special handling flags diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 11dc236064d..29973fe2101 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -1,11 +1,13 @@ from typing import Any, Final import litellm +from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -19,7 +21,9 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): ``"/v1/messages"``. The inbound Anthropic payload (system, cache_control, thinking, tools, ...) is forwarded essentially unchanged to ``{api_base}/v1/messages``, so Anthropic-only features that the - Anthropic->OpenAI translation would otherwise drop are preserved. Response + Anthropic->OpenAI translation would otherwise drop are preserved. The one + exception is ``cache_control``, whose Anthropic-only extensions (``ttl``) + are stripped unless ``supports_cache_control_ttl`` says otherwise. Response parsing and streaming are inherited from the native Anthropic config. """ @@ -53,6 +57,35 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): def should_filter_anthropic_beta_headers(self) -> bool: return False + def supports_cache_control_ttl(self) -> bool: + return False + + def transform_anthropic_messages_request( + self, + model: str, + messages: list[dict], # mutable-ok: matches dict-typed base signature + anthropic_messages_optional_request_params: dict, # mutable-ok: matches dict-typed base signature + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: matches dict-typed base signature + ) -> dict: # mutable-ok: matches dict-typed base signature + """ + Anthropic ignores prompt-caching hints it cannot honor, but strict + non-Anthropic implementations of the Messages API 400 the whole request + on Anthropic-only ``cache_control`` extensions (``cache_control.ttl: 1h + is not supported``), so unless the provider declares ttl support the + hints are reduced to their portable ``{"type": ...}`` core. + """ + request: Final = super().transform_anthropic_messages_request( + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + if self.supports_cache_control_ttl(): + return request + return normalize_cache_control_in_anthropic_payload(request) + def get_complete_url( self, api_base: str | None, @@ -91,6 +124,9 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): def should_strip_billing_metadata(self) -> bool: return True + def supports_cache_control_ttl(self) -> bool: + return bool(self._provider.constraints.get("cache_control_ttl")) + def _resolve_api_key(self, api_key: str | None) -> str | None: return api_key or get_secret_str(self._provider.api_key_env) or litellm.api_key diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 33e677b000e..2cdd969b00f 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -317,3 +317,124 @@ def test_json_provider_messages_config_probes_capabilities_under_provider_slug() ) assert JSONProviderAnthropicMessagesConfig(provider).custom_llm_provider == "exampleprovider" assert OpenAILikeAnthropicMessagesConfig().custom_llm_provider == "anthropic" + + +def _cache_control_request_params() -> tuple[list, dict]: + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "write a regex for a US phone number", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + ] + optional_params = { + "max_tokens": 256, + "system": [ + { + "type": "text", + "text": "You are Claude Code.", + "cache_control": {"type": "ephemeral", "ttl": "5m"}, + } + ], + "tools": [ + { + "name": "lookup", + "input_schema": {"type": "object"}, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + return messages, optional_params + + +def test_request_strips_cache_control_ttl_everywhere(config): + """Regression: Claude Code always sends ``cache_control: {type: ephemeral, + ttl: 1h}``, and strict non-Anthropic /v1/messages validators 400 the whole + request on the ttl extension (``cache_control.ttl: 1h is not supported``).""" + messages, optional_params = _cache_control_request_params() + + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["system"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_request_defaults_missing_cache_control_type_and_drops_non_dict(config): + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "a", "cache_control": {"ttl": "1h"}}, + {"type": "text", "text": "b", "cache_control": None}, + ], + } + ], + anthropic_messages_optional_request_params={"max_tokens": 64}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + blocks = payload["messages"][0]["content"] + assert blocks[0]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in blocks[1] + + +def test_native_anthropic_config_keeps_cache_control_ttl(): + """Anthropic itself accepts ttl, so the normalization must stay scoped to + the OpenAI-like passthrough and never reach the native Anthropic path.""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + messages, optional_params = _cache_control_request_params() + payload = AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-sonnet-4-20250514", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + assert payload["system"][0]["cache_control"] == {"type": "ephemeral", "ttl": "5m"} + + +def test_json_provider_constraint_opts_into_cache_control_ttl(): + from litellm.llms.openai_like.json_loader import SimpleProviderConfig + from litellm.llms.openai_like.messages.transformation import ( + JSONProviderAnthropicMessagesConfig, + ) + + base_data = {"base_url": "https://api.example.com/v1", "api_key_env": "EXAMPLE_API_KEY"} + strict = JSONProviderAnthropicMessagesConfig(SimpleProviderConfig(slug="strictprov", data=base_data)) + lenient = JSONProviderAnthropicMessagesConfig( + SimpleProviderConfig(slug="lenientprov", data={**base_data, "constraints": {"cache_control_ttl": True}}) + ) + + def transform(provider_config): + messages, optional_params = _cache_control_request_params() + return provider_config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert transform(strict)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert transform(lenient)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} From d3268e4e184f8b7cde6cce4473e14b1acedfc751 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:26:04 -0700 Subject: [PATCH 05/69] test(mcp): wrap over-long connection error message calls --- .../proxy/_experimental/mcp_server/test_rest_endpoints.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 51b946c11b7..e66101f6177 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -2954,11 +2954,15 @@ class TestConnectionErrorMessage: assert secret not in message def test_connect_error_points_at_reachability(self): - message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0) + message = rest_endpoints._connection_error_message( + httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0 + ) assert "unreachable" in message.lower() def test_timeout_error_message(self): - message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out"), "https://example.com", 30.0) + message = rest_endpoints._connection_error_message( + httpx.ConnectTimeout("timed out"), "https://example.com", 30.0 + ) assert "unreachable" in message.lower() def test_http_status_error_includes_status_code(self): From 0baf376efd691e139902ca7933b83666ea459a41 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:07:05 -0700 Subject: [PATCH 06/69] fix(openai_like): scope cache_control normalization to Messages API locations Rewrite the sanitizer without recursion (the code-quality gate rejects new recursive functions) and only touch cache_control where the Messages API defines it: the request, system blocks, tools, message content blocks, and tool_result content. Application data such as tool_use.input and tool input_schema is left untouched even when it contains a cache_control key --- litellm/llms/anthropic/common_utils.py | 78 +++++++++++++++---- ..._like_anthropic_messages_transformation.py | 62 +++++++++++++++ 2 files changed, 124 insertions(+), 16 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 695e3a313ef..1ef14362601 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1302,37 +1302,83 @@ def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format -def _normalized_cache_control(cache_control: dict) -> dict: # mutable-ok: as sibling sanitizers +def _normalized_cache_control(cache_control: object) -> dict[str, str] | None: # mutable-ok: JSON wire format + if not isinstance(cache_control, Mapping): + return None cache_type: Final = cache_control.get("type") return {"type": cache_type if isinstance(cache_type, str) else "ephemeral"} # mutable-ok: JSON wire format -def _normalize_cache_control_value(value: object) -> object: - if isinstance(value, dict): - return normalize_cache_control_in_anthropic_payload(value) - if isinstance(value, list): - return [_normalize_cache_control_value(item) for item in value] # mutable-ok: JSON wire format - return value +def _with_portable_cache_control(block: Mapping[str, object]) -> dict[str, object]: # mutable-ok: JSON wire format + if "cache_control" not in block: + return dict(block) # mutable-ok: JSON wire format + normalized: Final = _normalized_cache_control(block["cache_control"]) + rest: Final = {key: value for key, value in block.items() if key != "cache_control"} # mutable-ok: JSON wire format + return rest if normalized is None else {**rest, "cache_control": normalized} # mutable-ok: JSON wire format -def normalize_cache_control_in_anthropic_payload(payload: dict) -> dict: # mutable-ok: as sibling sanitizers +def _with_portable_cache_control_in_blocks(blocks: object) -> object: + if isinstance(blocks, str) or not isinstance(blocks, Sequence): + return blocks + return [ # mutable-ok: JSON wire format + _with_portable_cache_control(block) if isinstance(block, Mapping) else block for block in blocks + ] + + +def _with_portable_cache_control_in_content_block(block: object) -> object: + if not isinstance(block, Mapping): + return block + portable: Final = _with_portable_cache_control(block) + if portable.get("type") != "tool_result" or "content" not in portable: + return portable + return { # mutable-ok: JSON wire format + **portable, + "content": _with_portable_cache_control_in_blocks(portable["content"]), + } + + +def _with_portable_cache_control_in_message(message: object) -> object: + if not isinstance(message, Mapping) or "content" not in message: + return message + content: Final = message["content"] + if isinstance(content, str) or not isinstance(content, Sequence): + return message + return { # mutable-ok: JSON wire format + **message, + "content": [_with_portable_cache_control_in_content_block(block) for block in content], + } + + +def normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire format + payload: Mapping[str, object], +) -> dict[str, object]: """ Return a copy of an Anthropic /v1/messages payload with every - ``cache_control`` entry reduced to ``{"type": }``, - recursing through message content blocks, system blocks, and tools. + ``cache_control`` entry reduced to ``{"type": }`` + at the places the Messages API defines it: the request itself, system + blocks, tools, message content blocks, and ``tool_result`` content blocks. + Application data such as ``tool_use.input`` and tool ``input_schema`` is + never touched, even when it happens to contain a ``cache_control`` key. Anthropic itself accepts prompt-caching extensions such as ``ttl``, but strict non-Anthropic implementations of the Messages API validate the field literally and reject the whole request (``cache_control.ttl: 1h is not supported``, ``cache_control.type is required``), which 400s clients like - Claude Code that always send cache hints. Non-dict ``cache_control`` values - are dropped entirely. The caller's payload is never mutated. + Claude Code that send cache hints. Non-dict ``cache_control`` values are + dropped entirely. The caller's payload is never mutated. """ - return { # mutable-ok: JSON wire format, as sibling sanitizers - key: _normalized_cache_control(value) if key == "cache_control" else _normalize_cache_control_value(value) - for key, value in payload.items() - if key != "cache_control" or isinstance(value, dict) + portable: Final = _with_portable_cache_control(payload) + scoped: Final = { # mutable-ok: JSON wire format + key: ( + _with_portable_cache_control_in_blocks(value) + if key in ("system", "tools") + else [_with_portable_cache_control_in_message(message) for message in value] + if key == "messages" and isinstance(value, Sequence) and not isinstance(value, str) + else value + ) + for key, value in portable.items() } + return scoped def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 2cdd969b00f..d325492914e 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -438,3 +438,65 @@ def test_json_provider_constraint_opts_into_cache_control_ttl(): assert transform(strict)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} assert transform(lenient)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_request_strips_ttl_only_where_the_messages_api_defines_cache_control(config): + """Regression: the sanitizer must only touch ``cache_control`` where the + Messages API defines it (request, system, tools, content blocks, tool_result + content), never application data such as ``tool_use.input`` or a tool's + ``input_schema`` that happens to contain a ``cache_control`` key.""" + tool_input = {"cache_control": {"type": "ephemeral", "ttl": "1h"}, "query": "x"} + input_schema = { + "type": "object", + "properties": {"cache_control": {"type": "string", "ttl": "1h"}}, + } + messages = [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": tool_input}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "content": [ + {"type": "text", "text": "result", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ], + }, + {"type": "text", "text": "plain string content stays", "cache_control": {"ttl": "1h"}}, + ], + }, + {"role": "user", "content": "a plain string message"}, + ] + optional_params = { + "max_tokens": 64, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "tools": [ + { + "name": "lookup", + "input_schema": input_schema, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["input_schema"] == input_schema + assert payload["messages"][0]["content"][0]["input"] == tool_input + tool_result = payload["messages"][1]["content"][0] + assert tool_result["cache_control"] == {"type": "ephemeral"} + assert tool_result["content"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["messages"][1]["content"][1]["cache_control"] == {"type": "ephemeral"} + assert payload["messages"][2] == {"role": "user", "content": "a plain string message"} From c32eb41aad3b7b087c7d0023a71876d3dea6511d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:17:13 -0700 Subject: [PATCH 07/69] feat(openai_like): let a passthrough deployment keep cache_control ttl via model_info.cache_control_ttl The supported_endpoints passthrough had no way to keep ttl for an upstream that honors it, so the deployment now opts in with model_info.cache_control_ttl: true, injected into the config the same way the providers.json constraint is for JSON providers --- .../messages/handler.py | 8 +- .../openai_like/messages/transformation.py | 16 +-- ...erimental_pass_through_messages_handler.py | 97 ++++++++++--------- ..._like_anthropic_messages_transformation.py | 17 ++++ 4 files changed, 84 insertions(+), 54 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 69985bcdaa3..b82903d6f87 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -99,6 +99,10 @@ def _deployment_passes_through_anthropic_messages(model_info: object) -> bool: return isinstance(supported_endpoints, (list, tuple)) and "/v1/messages" in supported_endpoints +def _deployment_supports_cache_control_ttl(model_info: object) -> bool: + return isinstance(model_info, dict) and model_info.get("cache_control_ttl") is True + + ####### ENVIRONMENT VARIABLES ################### # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() @@ -568,7 +572,9 @@ def anthropic_messages_handler( OpenAILikeAnthropicMessagesConfig, ) - anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig() + anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig( + cache_control_ttl=_deployment_supports_cache_control_ttl(kwargs.get("model_info")), + ) if anthropic_messages_provider_config is None: # Route to Responses API for OpenAI / Azure, chat/completions for everything else. if _should_route_to_responses_api(custom_llm_provider, original_model, model): diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 29973fe2101..ac99617521c 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -23,10 +23,15 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): ``{api_base}/v1/messages``, so Anthropic-only features that the Anthropic->OpenAI translation would otherwise drop are preserved. The one exception is ``cache_control``, whose Anthropic-only extensions (``ttl``) - are stripped unless ``supports_cache_control_ttl`` says otherwise. Response - parsing and streaming are inherited from the native Anthropic config. + are stripped unless the deployment opts in with + ``model_info.cache_control_ttl: true``. Response parsing and streaming are + inherited from the native Anthropic config. """ + def __init__(self, cache_control_ttl: bool = False) -> None: + super().__init__() + self._cache_control_ttl: Final = cache_control_ttl + def validate_anthropic_messages_environment( self, headers: dict[str, str], @@ -58,7 +63,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): return False def supports_cache_control_ttl(self) -> bool: - return False + return self._cache_control_ttl def transform_anthropic_messages_request( self, @@ -114,7 +119,7 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): """ def __init__(self, provider: SimpleProviderConfig): - super().__init__() + super().__init__(cache_control_ttl=bool(provider.constraints.get("cache_control_ttl"))) self._provider = provider @property @@ -124,9 +129,6 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): def should_strip_billing_metadata(self) -> bool: return True - def supports_cache_control_ttl(self) -> bool: - return bool(self._provider.constraints.get("cache_control_ttl")) - def _resolve_api_key(self, api_key: str | None) -> str | None: return api_key or get_secret_str(self._provider.api_key_env) or litellm.api_key diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index ad4c3d6bfbb..e819433c269 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -296,21 +296,15 @@ async def test_bedrock_converse_budget_tokens_preserved(): mock_acompletion.assert_called_once() call_kwargs = mock_acompletion.call_args.kwargs - print( - "acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str) - ) + print("acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str)) # Verify thinking parameter is passed through with budget_tokens preserved thinking_param = call_kwargs.get("thinking") - assert ( - thinking_param is not None - ), "thinking parameter should be passed to acompletion" - assert ( - thinking_param.get("type") == "enabled" - ), "thinking.type should be 'enabled'" - assert ( - thinking_param.get("budget_tokens") == 1024 - ), f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" + assert thinking_param is not None, "thinking parameter should be passed to acompletion" + assert thinking_param.get("type") == "enabled", "thinking.type should be 'enabled'" + assert thinking_param.get("budget_tokens") == 1024, ( + f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" + ) def test_openai_model_with_thinking_converts_to_reasoning(): @@ -342,23 +336,18 @@ def test_openai_model_with_thinking_converts_to_reasoning(): call_kwargs = mock_responses.call_args.kwargs # Verify reasoning is set (converted from thinking) - assert ( - "reasoning" in call_kwargs - ), "reasoning should be passed to litellm.responses" + assert "reasoning" in call_kwargs, "reasoning should be passed to litellm.responses" # budget_tokens=1024 -> effort="low" (at the LOW budget threshold) # reasoning_auto_summary is False by default, so no summary key expected_reasoning = {"effort": "low"} assert call_kwargs["reasoning"] == expected_reasoning, ( - f"reasoning should be {expected_reasoning} for budget_tokens=1024, " - f"got {call_kwargs.get('reasoning')}" + f"reasoning should be {expected_reasoning} for budget_tokens=1024, got {call_kwargs.get('reasoning')}" ) assert "summary" not in call_kwargs["reasoning"] # Verify thinking is NOT passed directly to the Responses API - assert ( - "thinking" not in call_kwargs - ), "thinking should NOT be passed directly to litellm.responses" + assert "thinking" not in call_kwargs, "thinking should NOT be passed directly to litellm.responses" class TestThinkingParameterTransformation: @@ -411,9 +400,7 @@ class TestThinkingParameterTransformation: thinking=thinking, model="openai/gpt-5.2", ) - assert result == { - "reasoning_effort": {"effort": "high", "summary": "detailed"} - } + assert result == {"reasoning_effort": {"effort": "high", "summary": "detailed"}} finally: litellm.reasoning_auto_summary = original @@ -611,9 +598,9 @@ class TestThinkingSummaryPreservation: mock_responses.assert_called_once() call_kwargs = mock_responses.call_args.kwargs reasoning = call_kwargs["reasoning"] - assert ( - reasoning["summary"] == "concise" - ), f"Expected summary='concise', got summary='{reasoning.get('summary')}'" + assert reasoning["summary"] == "concise", ( + f"Expected summary='concise', got summary='{reasoning.get('summary')}'" + ) def test_responses_adapter_preserves_summary(self): """translate_thinking_to_reasoning should include summary when user provides it.""" @@ -622,9 +609,7 @@ class TestThinkingSummaryPreservation: ) thinking = {"type": "enabled", "budget_tokens": 5000, "summary": "concise"} - result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( - thinking - ) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) assert result == {"effort": "high", "summary": "concise"} def test_responses_adapter_no_summary_by_default(self): @@ -638,11 +623,7 @@ class TestThinkingSummaryPreservation: try: litellm.reasoning_auto_summary = False thinking = {"type": "enabled", "budget_tokens": 5000} - result = ( - LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( - thinking - ) - ) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) assert result == {"effort": "high"} assert result is not None and "summary" not in result finally: @@ -659,9 +640,7 @@ class TestThinkingSummaryPreservation: thinking=thinking, model="openai/gpt-5.2", ) - assert result == { - "reasoning_effort": {"effort": "high", "summary": "concise"} - } + assert result == {"reasoning_effort": {"effort": "high", "summary": "concise"}} def test_translate_thinking_for_model_disabled_stays_plain_string_when_auto_summary_enabled(self): """Disabled thinking must stay a plain string even when reasoning_auto_summary is on.""" @@ -807,9 +786,7 @@ def test_presanitized_flag_not_leaked_to_provider_params(): def fake_base_handler(*args, **kwargs): captured.update(kwargs) - captured["optional"] = kwargs.get( - "anthropic_messages_optional_request_params", {} - ) + captured["optional"] = kwargs.get("anthropic_messages_optional_request_params", {}) return "stub" with patch.object( @@ -974,6 +951,38 @@ def test_gate_passthrough_skipped_when_only_chat_completions_supported(monkeypat assert "config" not in captured +@pytest.mark.parametrize( + "model_info, expected_ttl_support", + [ + ({"supported_endpoints": ["/v1/messages"]}, False), + ({"supported_endpoints": ["/v1/messages"], "cache_control_ttl": True}, True), + ({"supported_endpoints": ["/v1/messages"], "cache_control_ttl": "yes"}, False), + ], +) +def test_gate_passthrough_forwards_cache_control_ttl_only_when_deployment_opts_in( + monkeypatch, model_info, expected_ttl_support +): + """The passthrough config strips cache_control.ttl unless the deployment sets + model_info.cache_control_ttl to exactly true.""" + from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( + anthropic_messages_handler, + ) + + captured, _ = _gate_stubs(monkeypatch) + + result = anthropic_messages_handler( + max_tokens=100, + messages=[{"role": "user", "content": "Hello"}], + model="openai/some-model", + api_key="sk-test", + api_base="https://host/v1", + model_info=model_info, + ) + + assert result == "native-passthrough" + assert captured["config"].supports_cache_control_ttl() is expected_ttl_support + + def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_flag(): """Regional and provider-prefixed Claude 4.8+/5 entries carry ``supports_mid_conversation_system``, but the bare first-party keys @@ -987,9 +996,7 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys import litellm - cost_map_path = os.path.join( - os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json" - ) + cost_map_path = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json") with open(cost_map_path) as f: cost_map = json.load(f) rules = cost_map["fallback_generalizations"]["rules"] @@ -1028,9 +1035,7 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys ("perplexity/sonar", "sonar", "https://api.perplexity.ai/chat/completions"), ], ) -async def test_messages_strips_provider_prefix_exactly_once( - requested_model, expected_wire_model, expected_url -): +async def test_messages_strips_provider_prefix_exactly_once(requested_model, expected_wire_model, expected_url): """ BerriAI/litellm#37716: only the leading provider segment may be stripped on the way upstream. diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index d325492914e..e33b03afdff 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -414,6 +414,23 @@ def test_native_anthropic_config_keeps_cache_control_ttl(): assert payload["system"][0]["cache_control"] == {"type": "ephemeral", "ttl": "5m"} +def test_deployment_opt_in_keeps_cache_control_ttl(): + config = OpenAILikeAnthropicMessagesConfig(cache_control_ttl=True) + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + } + ], + anthropic_messages_optional_request_params={"max_tokens": 16}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + def test_json_provider_constraint_opts_into_cache_control_ttl(): from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.llms.openai_like.messages.transformation import ( From d4fc54a11d1ac18f10c33741b760b99716faac93 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 08:25:42 +0000 Subject: [PATCH 08/69] chore(techdebt): clear fresh debt from the 2026-08-31 window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 6 +-- litellm/llms/gigachat/authenticator.py | 28 +++++--------- litellm/llms/gigachat/chat/streaming.py | 1 - litellm/llms/gigachat/chat/transformation.py | 38 +++++-------------- .../llms/gigachat/embedding/transformation.py | 14 ++----- .../gigachat/passthrough/transformation.py | 2 - litellm/llms/gigachat/utils.py | 1 - litellm/passthrough/main.py | 3 -- .../llm_passthrough_endpoints.py | 4 +- .../router_strategy/test_complexity_router.py | 6 ++- type-discipline-budget.json | 8 ++-- 11 files changed, 35 insertions(+), 76 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index a07b9352659..d84aacfeaf0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5607 }, "reportMissingTypeArgument": { - "limit": 15310 + "limit": 15308 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38368 + "limit": 38367 }, "reportUnknownParameterType": { "limit": 19633 }, "reportUnknownVariableType": { - "limit": 29908 + "limit": 29906 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index d6b217d5746..73086ba395b 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -8,6 +8,7 @@ Based on official GigaChat SDK authentication flow. import time import uuid from collections.abc import Mapping +from types import MappingProxyType from typing import Final import httpx @@ -32,8 +33,8 @@ GIGACHAT_SCOPE: Final = "GIGACHAT_API_PERS" # Token expiry buffer in milliseconds (refresh token 60s before expiry) TOKEN_EXPIRY_BUFFER_MS: Final = 60000 -# Cache for access tokens _token_cache: Final = InMemoryCache() +_NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) class GigaChatAuthError(BaseLLMException): @@ -80,10 +81,9 @@ def get_access_token( Raises: GigaChatAuthError: If authentication fails """ - if not litellm_params: - litellm_params = {} # mutable-ok: empty dict default; rebind-ok: provide default + params: Final = litellm_params or _NO_LITELLM_PARAMS - access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") + access_token: Final = params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") if access_token: return access_token @@ -94,24 +94,20 @@ def get_access_token( message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", ) - effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope() - effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() + effective_scope: Final = scope or params.get("gigachat_scope") or _get_scope() + effective_auth_url: Final = auth_url or params.get("gigachat_auth_url") or _get_auth_url() - # Check cache cache_key: Final = f"gigachat_token:{effective_credentials[:16]}" cached: Final = _token_cache.get_cache(cache_key) if cached: _token, _expires_at = cached - # Check if token is still valid (with buffer) if time.time() * 1000 < _expires_at - TOKEN_EXPIRY_BUFFER_MS: verbose_logger.debug("Using cached GigaChat access token") return _token - # Request new token new_token, new_expires_at = _request_token_sync(effective_credentials, effective_scope, effective_auth_url) # pyright: ignore[reportArgumentType] # credential keys may be broader than str if new_expires_at: - # Cache token ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) if ttl_seconds > 0: _token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds) @@ -126,10 +122,9 @@ async def get_access_token_async( litellm_params: Mapping[str, object] | None = None, ) -> str: """Async version of get_access_token.""" - if not litellm_params: - litellm_params = {} # mutable-ok: empty dict default; rebind-ok: provide default + params: Final = litellm_params or _NO_LITELLM_PARAMS - access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") + access_token: Final = params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") if access_token: return access_token @@ -140,10 +135,9 @@ async def get_access_token_async( message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", ) - effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope() - effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() + effective_scope: Final = scope or params.get("gigachat_scope") or _get_scope() + effective_auth_url: Final = auth_url or params.get("gigachat_auth_url") or _get_auth_url() - # Check cache cache_key: Final = f"gigachat_token:{effective_credentials[:16]}" cached: Final = _token_cache.get_cache(cache_key) if cached: @@ -152,11 +146,9 @@ async def get_access_token_async( verbose_logger.debug("Using cached GigaChat access token") return _token - # Request new token new_token, new_expires_at = await _request_token_async(effective_credentials, effective_scope, effective_auth_url) # pyright: ignore[reportArgumentType] # credential keys may be broader than str if new_expires_at: - # Cache token ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) if ttl_seconds > 0: _token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds) diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 2875b30232e..0a4cbd8e520 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -52,7 +52,6 @@ class GigaChatModelResponseIterator: tool_use: ChatCompletionToolCallChunk | None = None # rebind-ok: conditionally assigned on function_call finish_reason: str | None = chunk_finish_reason - # Handle function_call in stream raw_function_call: Final = delta.get("function_call") if chunk_finish_reason == "function_call" and isinstance(raw_function_call, Mapping) and raw_function_call: func_call: Final[Mapping[str, object]] = raw_function_call diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 8f23c5175ec..991a93ccb21 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -111,11 +111,9 @@ class GigaChatConfig(BaseConfig): """ Set up headers with OAuth token. """ - # Get access token credentials: Final = api_key or get_secret_str("GIGACHAT_CREDENTIALS") or get_secret_str("GIGACHAT_API_KEY") access_token: Final = get_access_token(credentials=credentials, litellm_params=litellm_params) - # Store credentials for image uploads self._current_credentials = credentials self._current_api_base = api_base @@ -208,18 +206,16 @@ class GigaChatConfig(BaseConfig): def _convert_tools_to_functions(self, tools: Sequence) -> Sequence[dict]: """Convert OpenAI tools format to GigaChat functions format.""" - functions: Final[list[dict]] = [] # mutable-ok: accumulator for building functions list - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "function": - func = tool.get("function", {}) - functions.append( - { - "name": func.get("name", ""), - "description": func.get("description", ""), - "parameters": func.get("parameters", {}), - } - ) - return functions + return [ + { + "name": function.get("name", ""), + "description": function.get("description", ""), + "parameters": function.get("parameters", {}), + } + for function in ( + tool.get("function", {}) for tool in tools if isinstance(tool, dict) and tool.get("type") == "function" + ) + ] def _map_tool_choice(self, tool_choice: str | Mapping[str, object]) -> str | Mapping[str, object] | None: """ @@ -299,7 +295,6 @@ class GigaChatConfig(BaseConfig): if part.get("type") == "text": texts.append(part.get("text", "")) elif part.get("type") == "image_url": - # Extract image URL and upload to GigaChat image_url: object = part.get("image_url", {}) upload_url: str if isinstance(image_url, str): @@ -322,16 +317,13 @@ class GigaChatConfig(BaseConfig): headers: Mapping[str, object], ) -> dict: # mutable-ok: request payload sent to httpx """Transform OpenAI request to GigaChat format.""" - # Transform messages giga_messages: Final = self._transform_messages(messages) - # Build request request_data: Final[dict[str, object]] = { "model": model.replace("gigachat/", ""), "messages": giga_messages, } - # Add optional params for key in [ "temperature", "top_p", @@ -343,7 +335,6 @@ class GigaChatConfig(BaseConfig): if key in optional_params: request_data[key] = optional_params[key] - # Add functions if present if "functions" in optional_params: request_data["functions"] = optional_params["functions"] if "function_call" in optional_params: @@ -358,10 +349,8 @@ class GigaChatConfig(BaseConfig): for i, msg in enumerate(messages): message = dict(msg) - # Remove unsupported fields message.pop("name", None) - # Transform roles role = message.get("role", "user") if role == "developer": message["role"] = "system" @@ -374,18 +363,15 @@ class GigaChatConfig(BaseConfig): if not isinstance(content, str) or not is_valid_json(content): message["content"] = json.dumps(content, ensure_ascii=False) - # Handle None content if message.get("content") is None: message["content"] = "" - # Handle list content (multimodal) - extract text and images content = message.get("content") if isinstance(content, list): message["content"], attachments = self._transform_list_content(content) if attachments: message["attachments"] = attachments - # Transform tool_calls to function_call tool_calls = message.get("tool_calls") if tool_calls and isinstance(tool_calls, list) and len(tool_calls) > 0: tool_call = tool_calls[0] @@ -436,13 +422,11 @@ class GigaChatConfig(BaseConfig): message_data = choice.get("message", {}) finish_reason = choice.get("finish_reason", "stop") - # Transform function_call to tool_calls or content if finish_reason == "function_call" and message_data.get("function_call"): func_call = message_data["function_call"] args = func_call.get("arguments", {}) if is_structured_output: - # Convert to content for structured output if isinstance(args, dict): content = json.dumps(args, ensure_ascii=False) else: @@ -452,7 +436,6 @@ class GigaChatConfig(BaseConfig): message_data.pop("functions_state_id", None) finish_reason = "stop" else: - # Convert to tool_calls format if isinstance(args, dict): args = json.dumps(args, ensure_ascii=False) message_data["tool_calls"] = [ @@ -468,7 +451,6 @@ class GigaChatConfig(BaseConfig): message_data.pop("function_call", None) finish_reason = "tool_calls" - # Clean up GigaChat-specific fields message_data.pop("functions_state_id", None) choices.append( diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index 2ec8324e33c..0db4475be8f 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -112,18 +112,10 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): "input": ["text1", "text2", ...] } """ - # Normalize input to list - if isinstance(input, str): - input_list: list = [input] # rebind-ok: locally scoped conversion - else: - input_list = input - - # Remove gigachat/ prefix from model if present - model = model.removeprefix("gigachat/") # rebind-ok: parameter reassignment for normalization - + normalized_input: Final = [input] if isinstance(input, str) else input # mutable-ok: preserve list API return { - "model": model, - "input": input_list, + "model": model.removeprefix("gigachat/"), + "input": normalized_input, } def transform_embedding_response( diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index a0edc6f5682..e1f73d04275 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -60,7 +60,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): """ Set up headers with OAuth token. """ - # Get access token access_token: Final = get_access_token(credentials=api_key, litellm_params=litellm_params) headers["Authorization"] = f"Bearer {access_token}" # rebind-ok: mutating for OAuth setup @@ -82,7 +81,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): from litellm.types.utils import LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager - # cost tracking only for completions and embeddings if "completions" in endpoint: provider_chat_config: Final = ProviderConfigManager.get_provider_chat_config( provider=LlmProviders(custom_llm_provider), diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index cbb35cd1b57..ce7e848ed7f 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -4,7 +4,6 @@ from typing import Final from litellm.secret_managers.main import get_secret_str from litellm.types.utils import PromptTokensDetailsWrapper, Usage -# GigaChat API endpoint GIGACHAT_BASE_URL: Final = "https://gigachat.devices.sberbank.ru/api/v1" diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 9095cee15a9..689c34b7a88 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -113,10 +113,8 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): ) ) - # Compliant: Save a strong reference to prevent GC self._background_tasks.add(task) - # Remove the task from the set when it finishes to avoid memory leaks task.add_done_callback(self._background_tasks.discard) except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging verbose_logger.exception( @@ -578,7 +576,6 @@ def llm_passthrough_route( else: return response except Exception as e: - # provider_config is guaranteed non-None here due to the earlier guard assert provider_config is not None raise base_llm_http_handler._handle_error( e=e, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 78d8ce296b8..b48b8d81494 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1731,7 +1731,7 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: def get_vertex_pass_through_handler( - call_type: Literal["discovery", "aiplatform"], # noqa: UP037 + call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here ) -> BaseVertexAIPassThroughHandler: if call_type == "discovery": return VertexAIDiscoveryPassThroughHandler() @@ -2961,7 +2961,6 @@ async def handle_gigachat_passthrough_router_model( """ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - # Detect streaming based on request body is_streaming: Final = request_body.get("stream", False) # pyright: ignore[reportUnknownVariableType] # request_body is dict[Unknown, Unknown] data: dict[str, Any] = await _read_request_body( @@ -2997,7 +2996,6 @@ async def handle_gigachat_passthrough_router_model( data["json"] = request_body data["custom_llm_provider"] = "gigachat" - # Remove sensitive keys from data keys: Final = [ # mutable-ok: list of keys to remove from data "gigachat_auth_url", "gigachat_access_token", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 1ec8be88c9b..93803ce1005 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -10268,7 +10268,8 @@ class TestContextWindowEscalation: litellm_router_instance=_windowed_router(_SMALL, _BIG), complexity_router_config=_tier_config(session_affinity=True), ) - session_kwargs = lambda: {"metadata": {"session_id": "s-1", "user_api_key_hash": "k-1"}} # noqa: E731 + def session_kwargs() -> dict[str, object]: + return {"metadata": {"session_id": "s-1", "user_api_key_hash": "k-1"}} first = await router.async_pre_routing_hook( model="test-router", request_kwargs=session_kwargs(), messages=_OVERSIZED_TURNS @@ -10291,7 +10292,8 @@ class TestContextWindowEscalation: litellm_router_instance=_windowed_router(_SMALL, _BIG), complexity_router_config=_tier_config(session_affinity=True), ) - session_kwargs = lambda: {"metadata": {"session_id": "s-2", "user_api_key_hash": "k-2"}} # noqa: E731 + def session_kwargs() -> dict[str, object]: + return {"metadata": {"session_id": "s-2", "user_api_key_hash": "k-2"}} pinned = await router.async_pre_routing_hook( model="test-router", request_kwargs=session_kwargs(), messages=[{"role": "user", "content": "ok continue"}] diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 83c49afb538..f65ebd24599 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,12 +1,12 @@ { "LIT001": { - "limit": 22403 + "limit": 22402 }, "LIT002": { "limit": 26780 }, "LIT003": { - "limit": 269 + "limit": 268 }, "LIT004": { "limit": 40 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16512 + "limit": 16511 }, "LIT011": { - "limit": 5537 + "limit": 5535 }, "LIT012": { "limit": 4495 From ab1161344199539bc8dec51161fb59d6d0acf703 Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 5 Aug 2026 18:01:29 +0000 Subject: [PATCH 09/69] fix(bedrock): strip client_metadata from converse additionalModelRequestFields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/chat/converse_transformation.py | 1 + .../chat/test_converse_transformation.py | 24 +++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 395d99a4caa..9e40cb5ee5b 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1324,6 +1324,7 @@ class AmazonConverseConfig(BaseConfig): ) additional_request_params.pop("parallel_tool_calls", None) + additional_request_params.pop("client_metadata", None) # Only set the topK value in for models that support it additional_request_params.update(self._handle_top_k_value(model, inference_params, drop_params)) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 63f895e1819..bd68857d664 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -979,6 +979,30 @@ def test_config_blocks_do_not_leak_into_inference_config(): assert data["serviceTier"] == {"type": "priority"} +def test_client_metadata_stripped_from_converse_request(): + """``client_metadata`` sent by codex must not reach Bedrock as a passthrough model field. + + Converse forwards ``additionalModelRequestFields`` verbatim to the model, and Anthropic + rejects the request with "client_metadata: Extra inputs are not permitted". + """ + config = AmazonConverseConfig() + + data = config._transform_request_helper( + model="anthropic.claude-opus-4-8", + system_content_blocks=[], + optional_params={ + "maxTokens": 16, + "anthropic_beta": ["computer-use-2025-01-24"], + "client_metadata": {"originator": "codex_cli_rs"}, + }, + messages=None, + ) + + fields = data.get("additionalModelRequestFields", {}) + assert "client_metadata" not in fields + assert fields["anthropic_beta"] == ["computer-use-2025-01-24"] + + def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost From 2063c29f5d95f8dd00eef3fd7dfcbc1df05787b8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 18:45:26 +0000 Subject: [PATCH 10/69] fix(anthropic): upgrade legacy thinking to adaptive on adaptive-only models for chat and Bedrock Converse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/chat/transformation.py | 3 + litellm/llms/anthropic/common_utils.py | 47 +++++++++++- .../messages/transformation.py | 47 +----------- .../bedrock/chat/converse_transformation.py | 3 + .../test_anthropic_chat_transformation.py | 25 ++++++ .../chat/test_converse_transformation.py | 76 +++++++++++++++++++ 6 files changed, 154 insertions(+), 47 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e1387a9068c..319eecfac2c 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1544,6 +1544,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params.pop("thinking", None) else: optional_params["thinking"] = value + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=model, optional_params=optional_params, custom_llm_provider=self._resolved_provider + ) elif param == "reasoning_effort": # Accept both string ("low") and dict ({"effort": "low", # "summary": "concise"}). The Responses->Chat parser keeps the diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 9871001bf66..ede93c6deb2 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -13,7 +13,12 @@ import httpx from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm -from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME +from litellm.constants import ( + DEFAULT_MODEL_CREATED_AT_TIME, + DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, +) from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, ) @@ -490,6 +495,46 @@ class AnthropicModelInfo(BaseLLMModelInfo): ) optional_params.pop("thinking", None) + @staticmethod + def translate_legacy_thinking_for_adaptive_model( + model: str, + optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param, as in maybe_drop_disabled_thinking + custom_llm_provider: str, + ) -> None: + """Translate legacy ``thinking.type=enabled`` to adaptive for the + adaptive-thinking models that reject it (4.7+ and the 5 families). + Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the + legacy shape natively, so it is forwarded verbatim and the caller's + ``budget_tokens`` cap keeps applying. Caller-provided + ``output_config.effort`` is never overridden. + """ + if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider): + return + if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider): + return + thinking: Final = optional_params.get("thinking") + if not isinstance(thinking, dict) or thinking.get("type") != "enabled": + return + + budget: Final = int(thinking.get("budget_tokens") or 0) + if budget >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and ( + AnthropicModelInfo._supports_model_capability(model, "supports_xhigh_reasoning_effort", custom_llm_provider) + ): + effort = "xhigh" + elif budget >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET: + effort = "high" + elif budget >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: + effort = "medium" + else: + effort = "low" + + optional_params["thinking"] = {"type": "adaptive"} + existing_output_config = optional_params.get("output_config") + if not isinstance(existing_output_config, dict): + existing_output_config = {} + existing_output_config.setdefault("effort", effort) + optional_params["output_config"] = existing_output_config + def is_effort_used( self, optional_params: dict | None, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 3d62b8b4784..988f81c9eb4 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -3,11 +3,6 @@ from typing import Any, Final import httpx -from litellm.constants import ( - DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, - DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, - DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, -) from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import verbose_logger @@ -400,46 +395,6 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): existing_output_config.setdefault("effort", mapped_effort) optional_params["output_config"] = existing_output_config - @staticmethod - def _translate_legacy_thinking_for_adaptive_model( - model: str, optional_params: dict, custom_llm_provider: str - ) -> None: - """Translate legacy ``thinking.type=enabled`` to adaptive for the - adaptive-thinking models that reject it (4.7+ and the 5 families). - Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the - legacy shape natively, so it is forwarded verbatim and the caller's - ``budget_tokens`` cap keeps applying. Caller-provided - ``output_config.effort`` is never overridden. - """ - from litellm.llms.anthropic.chat.transformation import AnthropicConfig - - if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider): - return - if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider): - return - thinking: Final = optional_params.get("thinking") - if not isinstance(thinking, dict) or thinking.get("type") != "enabled": - return - - budget: Final = int(thinking.get("budget_tokens") or 0) - if budget >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and ( - AnthropicConfig._supports_effort_level(model, "xhigh", custom_llm_provider) - ): - effort = "xhigh" - elif budget >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET: - effort = "high" - elif budget >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: - effort = "medium" - else: - effort = "low" - - optional_params["thinking"] = {"type": "adaptive"} - existing_output_config = optional_params.get("output_config") - if not isinstance(existing_output_config, dict): - existing_output_config = {} - existing_output_config.setdefault("effort", effort) - optional_params["output_config"] = existing_output_config - @staticmethod def _translate_adaptive_effort_for_non_adaptive_model( model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str @@ -606,7 +561,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): custom_llm_provider=self._resolved_provider, ) - self._translate_legacy_thinking_for_adaptive_model( + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( model=model, optional_params=anthropic_messages_optional_request_params, custom_llm_provider=self._resolved_provider, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 395d99a4caa..0a378dfc11c 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -934,6 +934,9 @@ class AmazonConverseConfig(BaseConfig): litellm.verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, model) else: optional_params["thinking"] = value + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=model, optional_params=optional_params, custom_llm_provider="bedrock" + ) elif param == "reasoning_effort" and isinstance(value, str): self._handle_reasoning_effort_parameter( model=model, reasoning_effort=value, optional_params=optional_params diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 25e2c3cda80..4f30e7d10f0 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -3127,6 +3127,31 @@ def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model( ) +@pytest.mark.parametrize( + "model,budget_tokens,expected", + [ + ("claude-opus-4-8", 4096, ({"type": "adaptive"}, {"effort": "high"})), + ("claude-opus-4-7", 24000, ({"type": "adaptive"}, {"effort": "xhigh"})), + ("claude-opus-4-6", 4096, ({"type": "enabled", "budget_tokens": 4096}, None)), + ("claude-sonnet-4-5-20250929", 4096, ({"type": "enabled", "budget_tokens": 4096}, None)), + ], +) +def test_legacy_thinking_translated_to_adaptive_on_adaptive_only_models(model, budget_tokens, expected): + """Adaptive-only models reject thinking={type: enabled} with a 400, so the + legacy shape must be upgraded to adaptive + output_config.effort on + /chat/completions too, while models that accept it keep the caller's budget.""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget_tokens}, "max_tokens": 64000}, + optional_params={}, + model=model, + drop_params=False, + ) + + assert (result["thinking"], result.get("output_config")) == expected + + @pytest.mark.parametrize( "bad_value", [ diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 63f895e1819..b729c9366cc 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -6269,6 +6269,82 @@ def test_adaptive_thinking_passes_through_on_46_plus_converse(model): assert optional_params.get("thinking") == {"type": "adaptive"} +@pytest.mark.parametrize( + "model,budget_tokens,expected_effort", + [ + ("anthropic.claude-opus-4-8", 4096, "high"), + ("us.anthropic.claude-opus-4-8", 2000, "low"), + ("global.anthropic.claude-opus-4-8", 12000, "xhigh"), + ("us.anthropic.claude-opus-4-7", 3000, "medium"), + ("anthropic.claude-fable-5", 4096, "high"), + ], +) +def test_legacy_thinking_translated_to_adaptive_on_adaptive_only_converse(model, budget_tokens, expected_effort): + """Adaptive-only models (4.7+, 5 families) reject thinking={type: enabled} + with a 400 on Bedrock Converse, so the legacy shape from callers like Claude + Code must be upgraded to thinking={type: adaptive} + output_config.effort + derived from budget_tokens, matching the /v1/messages passthrough.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget_tokens}, "max_tokens": 64000}, + optional_params={}, + model=model, + drop_params=False, + ) + request = config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert request["additionalModelRequestFields"]["thinking"] == {"type": "adaptive"} + assert request["additionalModelRequestFields"]["output_config"] == {"effort": expected_effort} + + +def test_legacy_thinking_translation_keeps_caller_output_config_effort_converse(): + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={ + "output_config": {"effort": "low"}, + "thinking": {"type": "enabled", "budget_tokens": 12000}, + "max_tokens": 64000, + }, + optional_params={}, + model="anthropic.claude-opus-4-8", + drop_params=False, + ) + + assert optional_params["thinking"] == {"type": "adaptive"} + assert optional_params["output_config"] == {"effort": "low"} + + +@pytest.mark.parametrize( + "model", + [ + "us.anthropic.claude-opus-4-6", + "anthropic.claude-3-5-sonnet-20241022-v2:0", + ], +) +def test_legacy_thinking_forwarded_verbatim_when_model_accepts_it_converse(model): + """The 4.6 family and pre-adaptive models accept thinking={type: enabled} + natively, so the caller's budget_tokens cap must keep applying.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 4096}, "max_tokens": 8192}, + optional_params={}, + model=model, + drop_params=False, + ) + + assert optional_params["thinking"] == {"type": "enabled", "budget_tokens": 4096} + assert "output_config" not in optional_params + + def test_adaptive_thinking_dropped_when_max_tokens_too_small_converse(): """When max_tokens can't fit even the minimum thinking budget, the raw adaptive block must be dropped entirely rather than translated, so the From b8dd27a77fdaaa390761a99bf27737a255c37e3a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:00:46 +0000 Subject: [PATCH 11/69] style(anthropic): keep mutable-ok annotation within ruff format width Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/common_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index ede93c6deb2..f2fcbc6c232 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -498,7 +498,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): @staticmethod def translate_legacy_thinking_for_adaptive_model( model: str, - optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param, as in maybe_drop_disabled_thinking + optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param like the sibling helpers custom_llm_provider: str, ) -> None: """Translate legacy ``thinking.type=enabled`` to adaptive for the From ed4343a02645e19590657ae257e6ae43b95f8e48 Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 1 Sep 2026 19:04:32 +0000 Subject: [PATCH 12/69] fix(bedrock): scope client_metadata drop to anthropic converse models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/chat/converse_transformation.py | 8 ++++++- .../chat/test_converse_transformation.py | 24 ++++++++++++++++--- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 9e40cb5ee5b..8d1d905a129 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1324,7 +1324,13 @@ class AmazonConverseConfig(BaseConfig): ) additional_request_params.pop("parallel_tool_calls", None) - additional_request_params.pop("client_metadata", None) + + if base_model.startswith("anthropic") and additional_request_params.pop("client_metadata", None) is not None: + litellm.verbose_logger.debug( + "Bedrock Converse: dropping `client_metadata` for model=%s, Anthropic rejects it with " + "'client_metadata: Extra inputs are not permitted'", + model, + ) # Only set the topK value in for models that support it additional_request_params.update(self._handle_top_k_value(model, inference_params, drop_params)) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index bd68857d664..3600e366b5a 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -979,8 +979,9 @@ def test_config_blocks_do_not_leak_into_inference_config(): assert data["serviceTier"] == {"type": "priority"} -def test_client_metadata_stripped_from_converse_request(): - """``client_metadata`` sent by codex must not reach Bedrock as a passthrough model field. +@pytest.mark.parametrize("model", ["anthropic.claude-opus-4-8", "us.anthropic.claude-opus-4-8"]) +def test_client_metadata_stripped_for_anthropic_converse_request(model): + """``client_metadata`` sent by codex must not reach Anthropic as a passthrough model field. Converse forwards ``additionalModelRequestFields`` verbatim to the model, and Anthropic rejects the request with "client_metadata: Extra inputs are not permitted". @@ -988,7 +989,7 @@ def test_client_metadata_stripped_from_converse_request(): config = AmazonConverseConfig() data = config._transform_request_helper( - model="anthropic.claude-opus-4-8", + model=model, system_content_blocks=[], optional_params={ "maxTokens": 16, @@ -1003,6 +1004,23 @@ def test_client_metadata_stripped_from_converse_request(): assert fields["anthropic_beta"] == ["computer-use-2025-01-24"] +def test_client_metadata_kept_for_non_anthropic_converse_request(): + """Only Anthropic is known to reject ``client_metadata``, so other families keep the passthrough.""" + config = AmazonConverseConfig() + + data = config._transform_request_helper( + model="amazon.nova-pro-v1:0", + system_content_blocks=[], + optional_params={ + "maxTokens": 16, + "client_metadata": {"originator": "codex_cli_rs"}, + }, + messages=None, + ) + + assert data["additionalModelRequestFields"]["client_metadata"] == {"originator": "codex_cli_rs"} + + def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost From b96121efb173965405aa521ffbbba026ce73d3a4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:10:22 +0000 Subject: [PATCH 13/69] refactor(anthropic): build adaptive output_config in one shot in legacy thinking helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/common_utils.py | 37 +++++++++++++++----------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index f2fcbc6c232..17fa8022388 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -516,24 +516,29 @@ class AnthropicModelInfo(BaseLLMModelInfo): if not isinstance(thinking, dict) or thinking.get("type") != "enabled": return - budget: Final = int(thinking.get("budget_tokens") or 0) - if budget >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and ( + effort: Final = AnthropicModelInfo._legacy_budget_to_effort( + model=model, + budget_tokens=int(thinking.get("budget_tokens") or 0), + custom_llm_provider=custom_llm_provider, + ) + existing_output_config: Final = optional_params.get("output_config") + optional_params["thinking"] = {"type": "adaptive"} + optional_params["output_config"] = { + "effort": effort, + **(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})), + } + + @staticmethod + def _legacy_budget_to_effort(model: str, budget_tokens: int, custom_llm_provider: str) -> str: + if budget_tokens >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and ( AnthropicModelInfo._supports_model_capability(model, "supports_xhigh_reasoning_effort", custom_llm_provider) ): - effort = "xhigh" - elif budget >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET: - effort = "high" - elif budget >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: - effort = "medium" - else: - effort = "low" - - optional_params["thinking"] = {"type": "adaptive"} - existing_output_config = optional_params.get("output_config") - if not isinstance(existing_output_config, dict): - existing_output_config = {} - existing_output_config.setdefault("effort", effort) - optional_params["output_config"] = existing_output_config + return "xhigh" + if budget_tokens >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET: + return "high" + if budget_tokens >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: + return "medium" + return "low" def is_effort_used( self, From bba951c5ebe20871caab9848c474585c91fa3535 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 1 Sep 2026 19:26:19 +0000 Subject: [PATCH 14/69] fix(bedrock): drop client_metadata for ARNs that hide the model family Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/chat/converse_transformation.py | 4 +- litellm/llms/bedrock/common_utils.py | 9 ++++ .../chat/test_converse_transformation.py | 43 +++++++++++++++++++ 3 files changed, 55 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 8d1d905a129..99d45bfc94b 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -86,6 +86,7 @@ from litellm.utils import ( from ..common_utils import ( BedrockError, BedrockModelInfo, + bedrock_arn_hides_model_family, bedrock_converse_supports_parallel_tool_use_config, get_anthropic_beta_from_headers, get_bedrock_tool_name, @@ -1325,7 +1326,8 @@ class AmazonConverseConfig(BaseConfig): additional_request_params.pop("parallel_tool_calls", None) - if base_model.startswith("anthropic") and additional_request_params.pop("client_metadata", None) is not None: + drops_client_metadata: Final = base_model.startswith("anthropic") or bedrock_arn_hides_model_family(model) + if drops_client_metadata and additional_request_params.pop("client_metadata", None) is not None: litellm.verbose_logger.debug( "Bedrock Converse: dropping `client_metadata` for model=%s, Anthropic rejects it with " "'client_metadata: Extra inputs are not permitted'", diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9cbceb4880c..3b82d98ceee 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -720,6 +720,15 @@ def get_bedrock_base_model(model: str) -> str: return model +def bedrock_arn_hides_model_family(model: str) -> bool: + """ + True for an ARN-addressed model whose base name carries no ``provider.model`` + id, such as an application inference profile or a provisioned throughput ARN. + Callers that gate behavior on the model family cannot resolve one here. + """ + return "arn:" in model.lower() and "." not in get_bedrock_base_model(model) + + def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool: return any( (litellm.model_cost.get(candidate) or {}).get("supports_parallel_tool_use_config") is True diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 3600e366b5a..423e3a5a7c3 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -1021,6 +1021,49 @@ def test_client_metadata_kept_for_non_anthropic_converse_request(): assert data["additionalModelRequestFields"]["client_metadata"] == {"originator": "codex_cli_rs"} +@pytest.mark.parametrize( + "model", + [ + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456", + "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/abcdef123456", + ], +) +def test_client_metadata_stripped_for_arn_models_converse(model): + """An ARN hides which family serves the request, and pointing one at Claude is how + teams route codex traffic, so the field has to go there too or the 400 comes back.""" + config = AmazonConverseConfig() + + data = config._transform_request_helper( + model=model, + system_content_blocks=[], + optional_params={ + "maxTokens": 16, + "client_metadata": {"originator": "codex_cli_rs"}, + }, + messages=None, + ) + + assert "client_metadata" not in data.get("additionalModelRequestFields", {}) + + +def test_client_metadata_kept_for_arn_naming_a_non_anthropic_family(): + """An inference profile ARN that still spells out the family is resolvable, so a + non-Anthropic one keeps its passthrough.""" + config = AmazonConverseConfig() + + data = config._transform_request_helper( + model="arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.amazon.nova-pro-v1:0", + system_content_blocks=[], + optional_params={ + "maxTokens": 16, + "client_metadata": {"originator": "codex_cli_rs"}, + }, + messages=None, + ) + + assert data["additionalModelRequestFields"]["client_metadata"] == {"originator": "codex_cli_rs"} + + def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost From 692f3b513ca9ac2fafb225c24d58f7ee5152eae6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:43:01 -0700 Subject: [PATCH 15/69] fix(proxy-extras): recover the v2 migration resolver from concurrent migrate deploy deadlocks Two instances racing prisma migrate deploy on one database deadlock on CREATE INDEX CONCURRENTLY: the victim gets P3018 with 40P01 and the survivor then sees the failed ledger row as P3009. Both were treated as unrecoverable, so neither instance came up. Roll the deadlocked migration's ledger row back and retry the deploy on P3018, consult the failed row's logs in _prisma_migrations to do the same on P3009, and retry a deadlock reported without a Prisma error code. Genuinely broken migrations still fail fast. --- .../litellm_proxy_extras/utils.py | 91 ++++++++++- .../tests/test_setup_database_fail_fast.py | 141 ++++++++++++++++++ 2 files changed, 228 insertions(+), 4 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b8032dd0d28..fb948afd200 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -40,6 +40,8 @@ def _get_prisma_env() -> dict: _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") +_MIGRATION_DEADLOCK_MARKER = "deadlock detected" + _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE) _SPEND_LOGS_ARTIFACT_DROP_RE = re.compile( r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE @@ -262,6 +264,48 @@ class ProxyExtrasDBManager: env=prisma_env, ) + @staticmethod + def _roll_back_migration_best_effort(migration_name: str) -> None: + """Mark a migration rolled back, tolerating a concurrent resolver + having already done it.""" + try: + ProxyExtrasDBManager._roll_back_migration(migration_name) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired): + pass + + @staticmethod + def _failed_migration_logs(migration_name: str) -> str: + """Logs recorded on the migration's failed _prisma_migrations row. + + P3009 stderr does not carry the original failure, so this is the only + way to tell a migration that lost a deadlock race against a concurrent + migrate deploy from one whose SQL is genuinely broken. Returns "" when + psycopg is missing, the DB is unreachable, or no failed row exists. + """ + database_url = os.getenv("DATABASE_URL") + if not database_url: + return "" + + try: + import psycopg + except ImportError: + return "" + + cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + try: + with psycopg.connect( + cleaned_url, connect_timeout=10, autocommit=True + ) as conn: + row = conn.execute( + "SELECT logs FROM _prisma_migrations " + "WHERE migration_name = %s AND finished_at IS NULL " + "AND rolled_back_at IS NULL", + (migration_name,), + ).fetchone() + except (psycopg.OperationalError, psycopg.DatabaseError): + return "" + return (row[0] or "") if row else "" + @staticmethod def _resolve_specific_migration(migration_name: str): """Mark a specific migration as applied""" @@ -658,7 +702,8 @@ class ProxyExtrasDBManager: v2 migration resolver (opt-in via --use_v2_migration_resolver). Runs `prisma migrate deploy` and handles standard recovery paths - (P3005 baseline, P3009/P3018 idempotent errors). Critically, it does + (P3005 baseline, P3009/P3018 idempotent errors, deadlocks against a + concurrent migrate deploy). Critically, it does NOT call `_resolve_all_migrations` — the diff-and-force recovery that caused schema thrashing when two LiteLLM versions contended for the same DB during rolling deploys. @@ -764,6 +809,22 @@ class ProxyExtrasDBManager: f"Detail: {resolve_err}" ) from resolve_err continue + if migration_match and _MIGRATION_DEADLOCK_MARKER in ( + ProxyExtrasDBManager._failed_migration_logs( + migration_match.group(1) + ) + ): + logger.info( + "Migration %s lost a deadlock race against a " + "concurrent migrate deploy, rolling its ledger " + "row back and retrying", + migration_match.group(1), + ) + ProxyExtrasDBManager._roll_back_migration_best_effort( + migration_match.group(1) + ) + time.sleep(random.randrange(5, 15)) + continue raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" @@ -809,11 +870,33 @@ class ProxyExtrasDBManager: ) from resolve_err continue + if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr: + logger.info( + "Migration %s deadlocked against a concurrent " + "migrate deploy, rolling its ledger row back " + "and retrying", + migration_match.group(1), + ) + ProxyExtrasDBManager._roll_back_migration_best_effort( + migration_match.group(1) + ) + time.sleep(random.randrange(5, 15)) + continue + raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" ) from e + if _MIGRATION_DEADLOCK_MARKER in stderr: + logger.info( + "prisma migrate deploy attempt %s deadlocked against " + "a concurrent migrate deploy, retrying", + attempt + 1, + ) + time.sleep(random.randrange(5, 15)) + continue + raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" @@ -821,9 +904,9 @@ class ProxyExtrasDBManager: raise RuntimeError( "Database migration failed after 4 attempts (retry loop " - "exhausted by timeouts or repeated idempotent-recovery " - "continues). Check database connectivity, load, and " - "_prisma_migrations ledger state." + "exhausted by timeouts, deadlock retries, or repeated " + "idempotent-recovery continues). Check database connectivity, " + "load, and _prisma_migrations ledger state." ) finally: os.chdir(original_dir) diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 8d66bf872de..c4347a91dce 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -240,3 +240,144 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery" + + +_DEADLOCK_P3018_STDERR = ( + "Error: P3018\n" + "Migration name: 20260415120000_health_check_latest_per_model_index\n" + "Database error code: 40P01\n" + "deadlock detected" +) + + +def _stub_v2_env(monkeypatch, tmp_path): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setattr( + ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setattr("time.sleep", lambda _: None) + + +def _succeed_after(failures: int, stderr: str): + calls = {"n": 0} + + class _OkResult: + stdout = "Applied migration.\n" + stderr = "" + + def _run(*args, **kwargs): + if "deploy" not in args[0]: + return _OkResult() + calls["n"] += 1 + if calls["n"] <= failures: + raise subprocess.CalledProcessError( + returncode=1, cmd=args[0], stderr=stderr, output="" + ) + return _OkResult() + + return _run + + +def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path): + """v2: losing the migrate deploy deadlock race against a concurrent + instance rolls the ledger row back and retries instead of dying.""" + _stub_v2_env(monkeypatch, tmp_path) + + rolled_back = [] + monkeypatch.setattr( + ProxyExtrasDBManager, + "_roll_back_migration", + lambda name: rolled_back.append(name), + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_resolve_specific_migration", + lambda name: pytest.fail("a deadlocked migration must never be marked applied"), + ) + monkeypatch.setattr("subprocess.run", _succeed_after(1, _DEADLOCK_P3018_STDERR)) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] + + +def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path): + """v2: a deadlock on every attempt still fails after the retry budget.""" + _stub_v2_env(monkeypatch, tmp_path) + monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None) + + with patch( + "subprocess.run", + side_effect=_fake_migrate_deploy_failure(1, _DEADLOCK_P3018_STDERR), + ): + with pytest.raises(RuntimeError, match="after 4 attempts"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + +def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_path): + """v2: the surviving instance sees the victim's failed ledger row as P3009. + When that row's logs show a deadlock, roll it back and retry.""" + _stub_v2_env(monkeypatch, tmp_path) + + stderr = ( + "Error: P3009\n" + "migrate found failed migrations in the target database\n" + "The `20260415120000_health_check_latest_per_model_index` migration " + "started at 2026-09-01 18:46:13 UTC failed" + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_failed_migration_logs", + lambda name: "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock", + ) + rolled_back = [] + monkeypatch.setattr( + ProxyExtrasDBManager, + "_roll_back_migration", + lambda name: rolled_back.append(name), + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_resolve_specific_migration", + lambda name: pytest.fail("a deadlocked migration must never be marked applied"), + ) + monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] + + +def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): + """v2: a failed ledger row whose logs show a real SQL error stays fatal.""" + _stub_v2_env(monkeypatch, tmp_path) + + stderr = ( + "Error: P3009\n" + "migrate found failed migrations in the target database\n" + "The `20260101000000_genuinely_broken` migration started at " + "2026-09-01 18:46:13 UTC failed" + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_failed_migration_logs", + lambda name: 'ERROR: syntax error at or near "BRKN"', + ) + + with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + +def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path): + """v2: a deadlock reported without a Prisma error code (the advisory-lock + waiter as victim) is retried, not fatal.""" + _stub_v2_env(monkeypatch, tmp_path) + monkeypatch.setattr( + "subprocess.run", _succeed_after(1, "Database error: deadlock detected") + ) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True From 5192b2162c987f260c9c33700343a73ea4676749 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 14:04:12 -0700 Subject: [PATCH 16/69] fix(proxy-extras): schema-qualify the _prisma_migrations logs lookup for non-public Prisma schemas --- litellm-proxy-extras/litellm_proxy_extras/utils.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index fb948afd200..97b6b1c667c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -292,14 +292,22 @@ class ProxyExtrasDBManager: return "" cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + ledger_table = psycopg.sql.SQL("{}.{}").format( + psycopg.sql.Identifier( + ProxyExtrasDBManager._prisma_schema_param(database_url) or "public" + ), + psycopg.sql.Identifier("_prisma_migrations"), + ) try: with psycopg.connect( cleaned_url, connect_timeout=10, autocommit=True ) as conn: row = conn.execute( - "SELECT logs FROM _prisma_migrations " - "WHERE migration_name = %s AND finished_at IS NULL " - "AND rolled_back_at IS NULL", + psycopg.sql.SQL( + "SELECT logs FROM {} " + "WHERE migration_name = %s AND finished_at IS NULL " + "AND rolled_back_at IS NULL" + ).format(ledger_table), (migration_name,), ).fetchone() except (psycopg.OperationalError, psycopg.DatabaseError): From bfa5eac76b18ae9e3965d5f942b2fd0382e6b796 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 10:54:39 -0700 Subject: [PATCH 17/69] fix(vector-store): resolve embedding aliases for search --- .../proxy/vector_store_endpoints/endpoints.py | 13 ++-- .../management_endpoints.py | 48 +++++++++++--- .../test_vector_store_endpoints.py | 66 ++++++++++++------- 3 files changed, 89 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index a59d7a277cc..e0b6cf8817a 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -70,17 +70,17 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( # time, instead of at row-creation time. The resolved # ``api_key`` / ``api_base`` / ``api_version`` lives only in # this per-request ``data`` dict and is never persisted. - # Legacy rows that already carry a resolved (cleartext) - # ``litellm_embedding_config`` skip the lookup and pass through - # unchanged so the embed call keeps working. + # Legacy rows that carry a resolved config are refreshed when the + # embedding model is an alias so the provider-qualified model is used. embedding_model: Final = litellm_params.get("litellm_embedding_model") - if embedding_model and not litellm_params.get("litellm_embedding_config"): + if embedding_model: from litellm.proxy.proxy_server import prisma_client - resolved_config: Final = await _resolve_embedding_config( + embedding_resolution: Final = await _resolve_embedding_config( embedding_model=embedding_model, prisma_client=prisma_client ) - if resolved_config: + if embedding_resolution: + resolved_model, resolved_config = embedding_resolution # Build a fresh dict via spread instead of mutating # ``litellm_params`` in place — the registry hands back # a reference to its cached object, so an in-place @@ -88,6 +88,7 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( # in-memory cache for the lifetime of the process. litellm_params = { **litellm_params, + "litellm_embedding_model": resolved_model, "litellm_embedding_config": resolved_config, } data.update(litellm_params) diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 183a03cc13c..1930ec4aaba 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -10,7 +10,7 @@ All /vector_store management endpoints import copy import json -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypeAlias from fastapi import APIRouter, Depends, HTTPException @@ -49,6 +49,7 @@ from litellm.types.vector_stores import ( from litellm.vector_stores.vector_store_registry import VectorStoreRegistry router: Final = APIRouter() +EmbeddingResolution: TypeAlias = tuple[str, dict[str, object]] def _vector_store_table(prisma_client: "PrismaClient") -> "TableActions[_VectorStoreRow]": @@ -155,7 +156,19 @@ async def _fetch_and_authorize_vector_store( return typed -def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, object] | None: +def _provider_qualified_embedding_model( + fallback: str, + model: object, + custom_llm_provider: object, +) -> str: + if not isinstance(model, str) or not model: + return fallback + if "/" in model or not isinstance(custom_llm_provider, str) or not custom_llm_provider: + return model + return f"{custom_llm_provider}/{model}" + + +def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> EmbeddingResolution | None: """ Resolve embedding config from router's config-defined models. @@ -168,7 +181,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d llm_router: The LiteLLM router instance Returns: - Dictionary with api_key, api_base, and api_version if model found, None otherwise + Provider-qualified model and its connection config if found, otherwise None """ if not embedding_model or llm_router is None: return None @@ -218,12 +231,21 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d if project_id: embedding_config["project_id"] = project_id + resolved_model: Final = _provider_qualified_embedding_model( + fallback=embedding_model, + model=getattr(litellm_params, "model", None), + custom_llm_provider=getattr(litellm_params, "custom_llm_provider", None), + ) + # Only return config if we have at least api_key or api_base if embedding_config: verbose_proxy_logger.debug( "Resolved embedding config from router model %s: %s", model_name, list(embedding_config.keys()) ) - return embedding_config + return ( + resolved_model, + embedding_config, + ) except Exception as e: verbose_proxy_logger.debug("Error resolving embedding config from router for model %s: %s", model_name, e) continue @@ -233,7 +255,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d async def _resolve_embedding_config_from_db( embedding_model: str, prisma_client: "PrismaClient" -) -> dict[str, object] | None: +) -> EmbeddingResolution | None: """ Resolve embedding config from database model configuration. @@ -246,7 +268,7 @@ async def _resolve_embedding_config_from_db( prisma_client: The Prisma client instance Returns: - Dictionary with api_key, api_base, and api_version if model found, None otherwise + Provider-qualified model and its connection config if found, otherwise None """ if not embedding_model: return None @@ -315,7 +337,15 @@ async def _resolve_embedding_config_from_db( model_name, list(embedding_config.keys()), ) - return embedding_config + resolved_model: Final = _provider_qualified_embedding_model( + fallback=embedding_model, + model=decrypted_params.get("model"), + custom_llm_provider=decrypted_params.get("custom_llm_provider"), + ) + return ( + resolved_model, + embedding_config, + ) except Exception as e: verbose_proxy_logger.debug("Error resolving embedding config for model %s: %s", model_name, e) continue @@ -325,7 +355,7 @@ async def _resolve_embedding_config_from_db( async def _resolve_embedding_config( embedding_model: str, prisma_client: "PrismaClient | None", llm_router: "Router | None" = None -) -> dict[str, object] | None: +) -> EmbeddingResolution | None: """ Resolve embedding config from either router (config-defined) or database models. @@ -343,7 +373,7 @@ async def _resolve_embedding_config( llm_router: The LiteLLM router instance (optional, will be imported if not provided) Returns: - Dictionary with api_key, api_base, and api_version if model found, None otherwise + Provider-qualified model and its connection config if found, otherwise None """ if not embedding_model: return None diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index eae6f90863a..1484adb258f 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -514,7 +514,7 @@ async def test_update_request_data_resolves_embedding_config_at_use_time(): "vector_store_id": "test_store", "custom_llm_provider": "azure_ai", "litellm_params": { - "litellm_embedding_model": "azure/text-embedding-3-large", + "litellm_embedding_model": "multilingual-e5-large", # Note: no litellm_embedding_config persisted }, } @@ -534,24 +534,22 @@ async def test_update_request_data_resolves_embedding_config_at_use_time(): patch.object(litellm, "vector_store_registry", mock_registry), patch( "litellm.proxy.vector_store_endpoints.endpoints._resolve_embedding_config", - new=AsyncMock(return_value=resolved), + new=AsyncMock(return_value=("azure/multilingual-e5-large", resolved)), ), ): result = await _update_request_data_with_litellm_managed_vector_store_registry( data={}, vector_store_id="test_store" ) - assert result["litellm_embedding_model"] == "azure/text-embedding-3-large" + assert result["litellm_embedding_model"] == "azure/multilingual-e5-large" assert result["litellm_embedding_config"] == resolved @pytest.mark.asyncio -async def test_update_request_data_passes_through_legacy_embedding_config(): +async def test_update_request_data_preserves_legacy_embedding_config_when_model_not_resolved(): """A vector store row created by an older proxy version may already carry a fully-resolved ``litellm_embedding_config`` in its persisted - ``litellm_params`` (the very leak this PR closes). Those legacy rows - must still work — the use-time resolver skips re-resolution when - the config is already present so the embed call keeps succeeding.""" + ``litellm_params``. Preserve it when the model cannot be resolved.""" legacy_config = { "api_key": "legacy-cleartext-key", "api_base": "https://legacy-azure.example", @@ -571,7 +569,7 @@ async def test_update_request_data_passes_through_legacy_embedding_config(): mock_vector_store ) - resolve_mock = AsyncMock() + resolve_mock = AsyncMock(return_value=None) with ( patch.object(litellm, "vector_store_registry", mock_registry), @@ -585,7 +583,7 @@ async def test_update_request_data_passes_through_legacy_embedding_config(): ) assert result["litellm_embedding_config"] == legacy_config - resolve_mock.assert_not_awaited() + resolve_mock.assert_awaited_once() class TestCheckVectorStorePermission: @@ -2010,6 +2008,7 @@ async def test_resolve_embedding_config_from_db(): # Mock database model with litellm_params mock_db_model = MagicMock() mock_db_model.litellm_params = { + "model": "openai/text-embedding-3-small", "api_key": "test-api-key", "api_base": "https://api.openai.com", "api_version": "2024-01-01", @@ -2028,9 +2027,11 @@ async def test_resolve_embedding_config_from_db(): ) assert result is not None - assert result["api_key"] == "test-api-key" - assert result["api_base"] == "https://api.openai.com" - assert result["api_version"] == "2024-01-01" + resolved_model, resolved_config = result + assert resolved_model == "openai/text-embedding-3-small" + assert resolved_config["api_key"] == "test-api-key" + assert resolved_config["api_base"] == "https://api.openai.com" + assert resolved_config["api_version"] == "2024-01-01" mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_called_once_with( where={"model_name": "text-embedding-ada-002"} ) @@ -2164,6 +2165,8 @@ def test_resolve_embedding_config_from_router(): mock_litellm_params.api_key = "config-api-key" mock_litellm_params.api_base = "https://config-api-base.com" mock_litellm_params.api_version = "2024-02-01" + mock_litellm_params.model = "text-embedding-3-small" + mock_litellm_params.custom_llm_provider = "openai" mock_deployment = MagicMock(spec=Deployment) mock_deployment.litellm_params = mock_litellm_params @@ -2176,9 +2179,11 @@ def test_resolve_embedding_config_from_router(): ) assert result is not None - assert result["api_key"] == "config-api-key" - assert result["api_base"] == "https://config-api-base.com" - assert result["api_version"] == "2024-02-01" + resolved_model, resolved_config = result + assert resolved_model == "openai/text-embedding-3-small" + assert resolved_config["api_key"] == "config-api-key" + assert resolved_config["api_base"] == "https://config-api-base.com" + assert resolved_config["api_version"] == "2024-02-01" mock_router.get_deployment_by_model_group_name.assert_called_once_with( model_group_name="text-embedding-ada-002" @@ -2197,6 +2202,8 @@ def test_resolve_embedding_config_from_router_with_provider_prefix(): mock_litellm_params.api_key = "azure-api-key" mock_litellm_params.api_base = "https://azure-endpoint.openai.azure.com" mock_litellm_params.api_version = "2024-02-15" + mock_litellm_params.model = "text-embedding-3-large" + mock_litellm_params.custom_llm_provider = "azure" mock_deployment = MagicMock(spec=Deployment) mock_deployment.litellm_params = mock_litellm_params @@ -2209,9 +2216,11 @@ def test_resolve_embedding_config_from_router_with_provider_prefix(): ) assert result is not None - assert result["api_key"] == "azure-api-key" - assert result["api_base"] == "https://azure-endpoint.openai.azure.com" - assert result["api_version"] == "2024-02-15" + resolved_model, resolved_config = result + assert resolved_model == "azure/text-embedding-3-large" + assert resolved_config["api_key"] == "azure-api-key" + assert resolved_config["api_base"] == "https://azure-endpoint.openai.azure.com" + assert resolved_config["api_version"] == "2024-02-15" # Should have tried both the full name and stripped name assert mock_router.get_deployment_by_model_group_name.call_count == 2 @@ -2239,6 +2248,8 @@ def test_resolve_embedding_config_from_router_handles_os_environ(): mock_litellm_params.api_key = "os.environ/OPENAI_API_KEY" mock_litellm_params.api_base = "https://direct-url.com" mock_litellm_params.api_version = None + mock_litellm_params.model = "text-embedding-3-small" + mock_litellm_params.custom_llm_provider = "openai" mock_deployment = MagicMock(spec=Deployment) mock_deployment.litellm_params = mock_litellm_params @@ -2254,9 +2265,11 @@ def test_resolve_embedding_config_from_router_handles_os_environ(): ) assert result is not None - assert result["api_key"] == "resolved-from-env" - assert result["api_base"] == "https://direct-url.com" - assert "api_version" not in result + resolved_model, resolved_config = result + assert resolved_model == "openai/text-embedding-3-small" + assert resolved_config["api_key"] == "resolved-from-env" + assert resolved_config["api_base"] == "https://direct-url.com" + assert "api_version" not in resolved_config mock_get_secret.assert_called_once_with("os.environ/OPENAI_API_KEY") @@ -2274,6 +2287,8 @@ async def test_resolve_embedding_config_tries_router_then_db(): mock_litellm_params.api_key = "router-api-key" mock_litellm_params.api_base = "https://router-api-base.com" mock_litellm_params.api_version = None + mock_litellm_params.model = "text-embedding-3-small" + mock_litellm_params.custom_llm_provider = "openai" mock_deployment = MagicMock(spec=Deployment) mock_deployment.litellm_params = mock_litellm_params @@ -2290,7 +2305,9 @@ async def test_resolve_embedding_config_tries_router_then_db(): ) assert result is not None - assert result["api_key"] == "router-api-key" + resolved_model, resolved_config = result + assert resolved_model == "openai/text-embedding-3-small" + assert resolved_config["api_key"] == "router-api-key" # DB should NOT have been called since router found the model mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_not_called() @@ -2345,6 +2362,7 @@ async def test_resolve_embedding_config_falls_back_to_db(): # DB has the model mock_db_model = MagicMock() mock_db_model.litellm_params = { + "model": "openai/text-embedding-3-small", "api_key": "db-api-key", "api_base": "https://db-api-base.com", } @@ -2363,7 +2381,9 @@ async def test_resolve_embedding_config_falls_back_to_db(): ) assert result is not None - assert result["api_key"] == "db-api-key" + resolved_model, resolved_config = result + assert resolved_model == "openai/text-embedding-3-small" + assert resolved_config["api_key"] == "db-api-key" # DB should have been called since router didn't find the model mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_called() From 5635811726ed05811abe5a242645dafd48eca9a0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 12:15:03 -0700 Subject: [PATCH 18/69] fix(vector-store): route embeddings through router --- .../base_llm/vector_store/transformation.py | 61 +- litellm/llms/custom_httpx/llm_http_handler.py | 6 + .../valkey/vector_stores/transformation.py | 39 +- .../proxy/vector_store_endpoints/endpoints.py | 32 +- .../management_endpoints.py | 289 +-------- litellm/router.py | 56 +- litellm/vector_stores/main.py | 29 +- .../test_router_embedding_integration.py | 94 ++- .../test_valkey_transformation.py | 34 +- .../test_vector_store_endpoints.py | 576 +++++------------- uv.lock | 22 +- 11 files changed, 469 insertions(+), 769 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 02a51a8bace..772e4f849a0 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -1,10 +1,14 @@ +from __future__ import annotations + from abc import abstractmethod from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, NoReturn +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, NoReturn, Protocol, runtime_checkable import httpx from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import EmbeddingResponse from litellm.types.vector_stores import ( VECTOR_STORE_OPENAI_PARAMS, BaseVectorStoreAuthCredentials, @@ -17,6 +21,7 @@ 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 from ..chat.transformation import BaseLLMException as _BaseLLMException @@ -27,6 +32,58 @@ else: BaseLLMException = Any +@runtime_checkable +class VectorStoreEmbeddingExecutor(Protocol): + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: ... + + async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: ... + + +@dataclass(frozen=True, slots=True) +class LiteLLMVectorStoreEmbeddingExecutor: + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + import litellm + + return litellm.embedding( # pyright: ignore[reportCallIssue, reportUnknownMemberType, reportUnknownVariableType] # provider kwargs are intentionally dynamic + model=model, + input=[query], # mutable-ok: LiteLLM embedding requires a mutable input list + **dict(configuration), # pyright: ignore[reportArgumentType] # provider-specific embedding config is validated downstream # mutable-ok: kwargs require a concrete dict + ) + + async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + import litellm + + return await litellm.aembedding( # pyright: ignore[reportUnknownMemberType] # provider kwargs are intentionally dynamic + model=model, + input=[query], # mutable-ok: LiteLLM embedding requires a mutable input list + **dict(configuration), # pyright: ignore[reportArgumentType] # provider-specific embedding config is validated downstream # mutable-ok: kwargs require a concrete dict + ) + + +@dataclass(frozen=True, slots=True) +class RouterVectorStoreEmbeddingExecutor: + router: Router + metadata: Mapping[str, object] + + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + if configuration: + return LiteLLMVectorStoreEmbeddingExecutor().embed(model, query, configuration) + return self.router.embedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list + model=model, + input=[query], # mutable-ok: Router embedding requires a mutable input list + metadata=dict(self.metadata), # mutable-ok: Router metadata requires a concrete dict + ) + + async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + if configuration: + return await LiteLLMVectorStoreEmbeddingExecutor().aembed(model, query, configuration) + return await self.router.aembedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list + model=model, + input=[query], # mutable-ok: Router embedding requires a mutable input list + metadata=dict(self.metadata), # mutable-ok: Router metadata requires a concrete dict + ) + + class BaseVectorStoreConfig: def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]: return [] @@ -172,6 +229,7 @@ class BaseDirectVectorStoreConfig(BaseVectorStoreConfig): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None = None, timeout: float | httpx.Timeout | None = None, ) -> VectorStoreSearchResponse: pass @@ -184,6 +242,7 @@ class BaseDirectVectorStoreConfig(BaseVectorStoreConfig): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None = None, timeout: float | httpx.Timeout | None = None, ) -> VectorStoreSearchResponse: pass diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 834f7d564a2..118656b81a2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -70,6 +70,7 @@ from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeech from litellm.llms.base_llm.vector_store.transformation import ( BaseDirectVectorStoreConfig, BaseVectorStoreConfig, + VectorStoreEmbeddingExecutor, ) from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, @@ -9683,6 +9684,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, extra_headers: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, @@ -9702,6 +9704,7 @@ class BaseLLMHTTPHandler: vector_store_search_optional_params=vector_store_search_optional_params, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), # mutable-ok: snapshot GenericLiteLLMParams into the Mapping shape + embedding_executor=embedding_executor, timeout=timeout, ) @@ -9797,6 +9800,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, extra_headers: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, @@ -9812,6 +9816,7 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, + embedding_executor=embedding_executor, extra_headers=extra_headers, extra_body=extra_body, timeout=timeout, @@ -9831,6 +9836,7 @@ class BaseLLMHTTPHandler: vector_store_search_optional_params=vector_store_search_optional_params, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), # mutable-ok: snapshot GenericLiteLLMParams into the Mapping shape + embedding_executor=embedding_executor, timeout=timeout, ) diff --git a/litellm/llms/valkey/vector_stores/transformation.py b/litellm/llms/valkey/vector_stores/transformation.py index 3cbfca0f1a9..b250f71cf3f 100644 --- a/litellm/llms/valkey/vector_stores/transformation.py +++ b/litellm/llms/valkey/vector_stores/transformation.py @@ -15,7 +15,10 @@ import httpx from pydantic import BaseModel, ConfigDict import litellm -from litellm.llms.base_llm.vector_store.transformation import BaseDirectVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import ( + BaseDirectVectorStoreConfig, + VectorStoreEmbeddingExecutor, +) from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector from litellm.types.utils import EmbeddingResponse from litellm.types.vector_stores import ( @@ -213,6 +216,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, litellm_logging_obj: "LiteLLMLoggingObj", litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None = None, timeout: float | httpx.Timeout | None = None, ) -> VectorStoreSearchResponse: params: Final = _ValkeySearchParams.model_validate(litellm_params) @@ -222,10 +226,18 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): embedding_field=params.embedding_field, text_field=params.text_field, ) - embedding_response: Final = self.embedding_fn( - model=params.require_embedding_model(), - input=[query_text], # mutable-ok: litellm.embedding's input contract is a list - **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG), + embedding_response: Final = ( + embedding_executor.embed( + params.require_embedding_model(), + query_text, + params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG, + ) + if embedding_executor is not None + else self.embedding_fn( + model=params.require_embedding_model(), + input=[query_text], # mutable-ok: the injected embedding callable requires list input + **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG), + ) ) vec_params: Final = {"vec": pack_vector(embedding_response.data[0]["embedding"])} # mutable-ok: redis-py API @@ -252,6 +264,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, litellm_logging_obj: "LiteLLMLoggingObj", litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None = None, timeout: float | httpx.Timeout | None = None, ) -> VectorStoreSearchResponse: params: Final = _ValkeySearchParams.model_validate(litellm_params) @@ -261,10 +274,18 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): embedding_field=params.embedding_field, text_field=params.text_field, ) - embedding_response: Final = await self.aembedding_fn( - model=params.require_embedding_model(), - input=[query_text], # mutable-ok: litellm.embedding's input contract is a list - **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG), + embedding_response: Final = ( + await embedding_executor.aembed( + params.require_embedding_model(), + query_text, + params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG, + ) + if embedding_executor is not None + else await self.aembedding_fn( + model=params.require_embedding_model(), + input=[query_text], # mutable-ok: the injected embedding callable requires list input + **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG), + ) ) vec_params: Final = {"vec": pack_vector(embedding_response.data[0]["embedding"])} # mutable-ok: redis-py API diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index e0b6cf8817a..3fc67181d5b 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -14,9 +14,6 @@ from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import jsonify_object -from litellm.proxy.vector_store_endpoints.management_endpoints import ( - _resolve_embedding_config, -) from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, assert_user_can_access_vector_store, @@ -65,32 +62,9 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( data["litellm_credential_name"] = vector_store_to_run.get("litellm_credential_name") if "litellm_params" in vector_store_to_run: - litellm_params = vector_store_to_run.get("litellm_params", {}) or {} - # Resolve ``litellm_embedding_config`` here, at request-handling - # time, instead of at row-creation time. The resolved - # ``api_key`` / ``api_base`` / ``api_version`` lives only in - # this per-request ``data`` dict and is never persisted. - # Legacy rows that carry a resolved config are refreshed when the - # embedding model is an alias so the provider-qualified model is used. - embedding_model: Final = litellm_params.get("litellm_embedding_model") - if embedding_model: - from litellm.proxy.proxy_server import prisma_client - - embedding_resolution: Final = await _resolve_embedding_config( - embedding_model=embedding_model, prisma_client=prisma_client - ) - if embedding_resolution: - resolved_model, resolved_config = embedding_resolution - # Build a fresh dict via spread instead of mutating - # ``litellm_params`` in place — the registry hands back - # a reference to its cached object, so an in-place - # update would persist the resolved cleartext into the - # in-memory cache for the lifetime of the process. - litellm_params = { - **litellm_params, - "litellm_embedding_model": resolved_model, - "litellm_embedding_config": resolved_config, - } + litellm_params: Final = ( + vector_store_to_run.get("litellm_params", {}) or {} + ) # mutable-ok: request execution merges persisted params into a mutable body data.update(litellm_params) return data diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 1930ec4aaba..8ca45f736ae 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -10,7 +10,7 @@ All /vector_store management endpoints import copy import json -from typing import TYPE_CHECKING, Any, Final, TypeAlias +from typing import TYPE_CHECKING, Any, Final from fastapi import APIRouter, Depends, HTTPException @@ -18,11 +18,8 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_ManagedVectorStoresTable as _VectorStoreRow from litellm.proxy.utils import PrismaClient - from litellm.router import Router - import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker @@ -32,13 +29,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store -from litellm.repositories.model_repository import ModelRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ManagedVectorStoresRepository -from litellm.secret_managers.main import get_secret from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, LiteLLM_ManagedVectorStoreListResponse, @@ -49,7 +43,6 @@ from litellm.types.vector_stores import ( from litellm.vector_stores.vector_store_registry import VectorStoreRegistry router: Final = APIRouter() -EmbeddingResolution: TypeAlias = tuple[str, dict[str, object]] def _vector_store_table(prisma_client: "PrismaClient") -> "TableActions[_VectorStoreRow]": @@ -65,28 +58,6 @@ _LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker() _REDACT_LITELLM_PARAMS_MAX_DEPTH: Final = 10 -# Use-time embedding-config resolution runs on every vector-store request -# whose persisted row carries only a model reference (the post-fix shape). -# Without a cache, that's one ``litellm_proxymodeltable.find_first`` per -# request — the no-DB-in-critical-path rule. Hold the resolved config in -# memory for a short TTL so a hot model name pays the DB lookup at most -# once per ``_EMBEDDING_CONFIG_CACHE_TTL`` seconds. Cleartext credentials -# only ever live in process memory (never persisted, never echoed in -# management responses), so the cache doesn't widen the disclosure surface. -_EMBEDDING_CONFIG_CACHE_TTL: Final = 60 -_EMBEDDING_CONFIG_CACHE_MAX_SIZE: Final = 256 -_embedding_config_cache: InMemoryCache | None = None - - -def _get_embedding_config_cache() -> InMemoryCache: - global _embedding_config_cache - if _embedding_config_cache is None: - _embedding_config_cache = InMemoryCache( - max_size_in_memory=_EMBEDDING_CONFIG_CACHE_MAX_SIZE, - default_ttl=_EMBEDDING_CONFIG_CACHE_TTL, - ) - return _embedding_config_cache - def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> Any: """ @@ -156,264 +127,6 @@ async def _fetch_and_authorize_vector_store( return typed -def _provider_qualified_embedding_model( - fallback: str, - model: object, - custom_llm_provider: object, -) -> str: - if not isinstance(model, str) or not model: - return fallback - if "/" in model or not isinstance(custom_llm_provider, str) or not custom_llm_provider: - return model - return f"{custom_llm_provider}/{model}" - - -def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> EmbeddingResolution | None: - """ - Resolve embedding config from router's config-defined models. - - Config-defined models (from proxy_config.yaml) are stored in the router's model_list, - not in the database. This function looks up the model in the router and extracts - api_key, api_base, and api_version from the deployment's litellm_params. - - Args: - embedding_model: The embedding model string (e.g., "text-embedding-ada-002" or "azure/text-embedding-3-large") - llm_router: The LiteLLM router instance - - Returns: - Provider-qualified model and its connection config if found, otherwise None - """ - if not embedding_model or llm_router is None: - return None - - # Extract model name candidates - could be "text-embedding-ada-002" or "azure/text-embedding-3-large" - # Try exact match first, then try without provider prefix - model_name_candidates: Final = [embedding_model] - if "/" in embedding_model: - # If it has a provider prefix, also try without it - _, model_name = embedding_model.split("/", 1) - model_name_candidates.append(model_name) - - # Try to find model in router - for model_name in model_name_candidates: - try: - # Try to get deployment by model group name (model_name in config) - deployment = llm_router.get_deployment_by_model_group_name(model_group_name=model_name) - - if deployment is not None and deployment.litellm_params is not None: - litellm_params = deployment.litellm_params - - # Build embedding config from model params - embedding_config: dict[str, object] = {} - - # Extract api_key - api_key = getattr(litellm_params, "api_key", None) - if api_key: - # Handle os.environ/ prefix - if isinstance(api_key, str) and api_key.startswith("os.environ/"): - api_key = get_secret(api_key) - embedding_config["api_key"] = api_key - - # Extract api_base - api_base = getattr(litellm_params, "api_base", None) - if api_base: - # Handle os.environ/ prefix - if isinstance(api_base, str) and api_base.startswith("os.environ/"): - api_base = get_secret(api_base) - embedding_config["api_base"] = api_base - - # Extract api_version - api_version = getattr(litellm_params, "api_version", None) - if api_version: - embedding_config["api_version"] = api_version - - project_id = getattr(litellm_params, "project_id", None) - if project_id: - embedding_config["project_id"] = project_id - - resolved_model: Final = _provider_qualified_embedding_model( - fallback=embedding_model, - model=getattr(litellm_params, "model", None), - custom_llm_provider=getattr(litellm_params, "custom_llm_provider", None), - ) - - # Only return config if we have at least api_key or api_base - if embedding_config: - verbose_proxy_logger.debug( - "Resolved embedding config from router model %s: %s", model_name, list(embedding_config.keys()) - ) - return ( - resolved_model, - embedding_config, - ) - except Exception as e: - verbose_proxy_logger.debug("Error resolving embedding config from router for model %s: %s", model_name, e) - continue - - return None - - -async def _resolve_embedding_config_from_db( - embedding_model: str, prisma_client: "PrismaClient" -) -> EmbeddingResolution | None: - """ - Resolve embedding config from database model configuration. - - If litellm_embedding_model is provided but litellm_embedding_config is not, - this function looks up the model in the database and extracts api_key, api_base, - and api_version from the model's litellm_params to build the embedding config. - - Args: - embedding_model: The embedding model string (e.g., "text-embedding-ada-002" or "azure/text-embedding-3-large") - prisma_client: The Prisma client instance - - Returns: - Provider-qualified model and its connection config if found, otherwise None - """ - if not embedding_model: - return None - - # Extract model name - could be "text-embedding-ada-002" or "azure/text-embedding-3-large" - # Try to find model by exact match first, then try without provider prefix - model_name_candidates: Final = [embedding_model] - if "/" in embedding_model: - # If it has a provider prefix, also try without it - _, model_name = embedding_model.split("/", 1) - model_name_candidates.append(model_name) - - # Try to find model in database - for model_name in model_name_candidates: - try: - db_model = await ModelRepository(prisma_client).table.find_first(where={"model_name": model_name}) - - if db_model and db_model.litellm_params: - # Extract litellm_params (could be dict or JSON string) - model_params = db_model.litellm_params - if isinstance(model_params, str): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json is str - model_params = json.loads(model_params) - - # Decrypt values from database (similar to how proxy_server.py does it) - # Values stored in DB are encrypted, so we need to decrypt them first - decrypted_params = {} - if isinstance(model_params, dict): - for k, v in model_params.items(): - if isinstance(v, str): - # Decrypt value - returns original value if decryption fails or no key is set - decrypted_value = decrypt_value_helper(value=v, key=k, return_original_value=True) - decrypted_params[k] = decrypted_value - else: - decrypted_params[k] = v - else: - decrypted_params = model_params - - # Build embedding config from model params - embedding_config = {} - - # Extract api_key - api_key = decrypted_params.get("api_key") - if api_key: - # Handle os.environ/ prefix (after decryption, values may be os.environ/ prefixed) - if isinstance(api_key, str) and api_key.startswith("os.environ/"): - api_key = get_secret(api_key) - embedding_config["api_key"] = api_key - - # Extract api_base - api_base = decrypted_params.get("api_base") - if api_base: - # Handle os.environ/ prefix (after decryption, values may be os.environ/ prefixed) - if isinstance(api_base, str) and api_base.startswith("os.environ/"): - api_base = get_secret(api_base) - embedding_config["api_base"] = api_base - - # Extract api_version - api_version = decrypted_params.get("api_version") - if api_version: - embedding_config["api_version"] = api_version - - # Only return config if we have at least api_key or api_base - if embedding_config: - verbose_proxy_logger.debug( - "Resolved embedding config from database model %s: %s", - model_name, - list(embedding_config.keys()), - ) - resolved_model: Final = _provider_qualified_embedding_model( - fallback=embedding_model, - model=decrypted_params.get("model"), - custom_llm_provider=decrypted_params.get("custom_llm_provider"), - ) - return ( - resolved_model, - embedding_config, - ) - except Exception as e: - verbose_proxy_logger.debug("Error resolving embedding config for model %s: %s", model_name, e) - continue - - return None - - -async def _resolve_embedding_config( - embedding_model: str, prisma_client: "PrismaClient | None", llm_router: "Router | None" = None -) -> EmbeddingResolution | None: - """ - Resolve embedding config from either router (config-defined) or database models. - - This function first checks the router for config-defined models, then falls back - to the database. This allows users to use models defined in either location. - - Results are cached in process memory for ``_EMBEDDING_CONFIG_CACHE_TTL`` - seconds so the request-handling path doesn't hit the database on every - vector-store call. Negative results (model not found) are intentionally - not cached to avoid blocking a freshly-added model behind the TTL. - - Args: - embedding_model: The embedding model string (e.g., "text-embedding-ada-002" or "azure/text-embedding-3-large") - prisma_client: The Prisma client instance - llm_router: The LiteLLM router instance (optional, will be imported if not provided) - - Returns: - Provider-qualified model and its connection config if found, otherwise None - """ - if not embedding_model: - return None - - cache: Final = _get_embedding_config_cache() - cached: Final = cache.get_cache(embedding_model) - if cached is not None: - return cached - - # Import llm_router if not provided - if llm_router is None: - try: - from litellm.proxy.proxy_server import llm_router - except ImportError: - llm_router = None - - # First try to resolve from router (config-defined models) - if llm_router is not None: - router_config = _resolve_embedding_config_from_router(embedding_model=embedding_model, llm_router=llm_router) - if router_config: - verbose_proxy_logger.debug("Resolved embedding config from router for model %s", embedding_model) - cache.set_cache(embedding_model, router_config) - return router_config - - # Fall back to database - if prisma_client is not None: - db_config: Final = await _resolve_embedding_config_from_db( - embedding_model=embedding_model, prisma_client=prisma_client - ) - if db_config: - verbose_proxy_logger.debug("Resolved embedding config from database for model %s", embedding_model) - cache.set_cache(embedding_model, db_config) - return db_config - - verbose_proxy_logger.debug( - "Could not resolve embedding config for model %s from router or database", embedding_model - ) - return None - - ######################################################## # Helper Functions ######################################################## diff --git a/litellm/router.py b/litellm/router.py index 471a1116f44..9e0e267f21c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -84,6 +84,9 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_credentials_in_payload, mask_sensitive_structure, ) +from litellm.llms.base_llm.vector_store.transformation import ( + RouterVectorStoreEmbeddingExecutor, +) from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.least_busy import LeastBusyLoggingHandler @@ -6319,6 +6322,34 @@ class Router: client: object | None = None, **kwargs, ): + if call_type == "vector_store_search": + metadata: Final = self._vector_store_request_metadata(kwargs) + provider_kwargs: Final = ( + { + "custom_llm_provider": custom_llm_provider + } # mutable-ok: provider kwargs are expanded into the request + if custom_llm_provider is not None + else MappingProxyType({}) + ) + search_kwargs: Final = { # mutable-ok: the routed request requires dynamic keyword arguments + **kwargs, + **provider_kwargs, + "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( + router=self, + metadata=metadata, + ), + } + model: Final = search_kwargs.get("model") + if isinstance(model, str) and model: + routed_kwargs: Final = { # mutable-ok: model must be removed before expanding routed kwargs + key: value for key, value in search_kwargs.items() if key != "model" + } + return self._generic_api_call_with_fallbacks( + model=model, + original_function=original_function, + **routed_kwargs, + ) + return original_function(**search_kwargs) return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs) return sync_wrapper @@ -6512,10 +6543,21 @@ class Router: "avector_store_update", "avector_store_delete", ): + vector_store_kwargs: Final = ( + { # mutable-ok: the async routed request requires dynamic keyword arguments + **kwargs, + "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( + router=self, + metadata=self._vector_store_request_metadata(kwargs), + ), + } + if call_type == "avector_store_search" + else kwargs + ) return await self._init_vector_store_api_endpoints( original_function=original_function, custom_llm_provider=custom_llm_provider, - **kwargs, + **vector_store_kwargs, ) elif call_type in ("afile_delete", "afile_content"): return await self._ageneric_api_call_with_fallbacks( @@ -6551,6 +6593,18 @@ class Router: return async_wrapper + @staticmethod + def _vector_store_request_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: + litellm_metadata: Final = kwargs.get("litellm_metadata") + if isinstance(litellm_metadata, dict): + return cast( # cast-ok: isinstance validates the runtime dict boundary + "dict[str, object]", litellm_metadata + ) + metadata: Final = kwargs.get("metadata") + if isinstance(metadata, dict): + return cast("dict[str, object]", metadata) # cast-ok: isinstance validates the runtime dict boundary + return MappingProxyType({}) + async def _init_vector_store_api_endpoints( self, original_function: Callable, diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 9b0ff71730a..89c3319ca5a 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -15,6 +15,10 @@ import litellm from litellm.constants import request_timeout from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.vector_store.transformation import ( + LiteLLMVectorStoreEmbeddingExecutor, + VectorStoreEmbeddingExecutor, +) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -35,6 +39,14 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# +def _direct_vector_store_embedding_executor(value: object) -> VectorStoreEmbeddingExecutor: + if value is None: + return LiteLLMVectorStoreEmbeddingExecutor() + if isinstance(value, VectorStoreEmbeddingExecutor): + return value + raise TypeError("Invalid direct vector store embedding executor") + + def mock_vector_store_search_response( mock_results: list[VectorStoreSearchResult] | None = None, ): @@ -285,7 +297,12 @@ async def asearch( """ Async: Search a vector store for relevant chunks based on a query and file attributes filter. """ - local_vars: Final = locals() + embedding_executor: Final = _direct_vector_store_embedding_executor( + kwargs.pop("_direct_vector_store_embedding_executor", None) + ) + local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot + key: value for key, value in locals().items() if key != "embedding_executor" + } try: loop: Final = asyncio.get_event_loop() @@ -308,6 +325,7 @@ async def asearch( extra_body=extra_body, timeout=timeout, custom_llm_provider=custom_llm_provider, + _direct_vector_store_embedding_executor=embedding_executor, **kwargs, ) @@ -363,12 +381,16 @@ def search( Returns: VectorStoreSearchResponse containing the search results. """ - local_vars: Final = locals() + embedding_executor: Final = _direct_vector_store_embedding_executor( + kwargs.pop("_direct_vector_store_embedding_executor", None) + ) + local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot + key: value for key, value in locals().items() if key != "embedding_executor" + } try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) _is_async: Final = kwargs.pop("asearch", False) is True - # pull credentials from registry if available if litellm.vector_store_registry is not None and vector_store_id is not None: try: @@ -445,6 +467,7 @@ def search( custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, logging_obj=litellm_logging_obj, + embedding_executor=embedding_executor, extra_headers=extra_headers, extra_body=extra_body, timeout=timeout or request_timeout, diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index 75dacbaf08e..5c01587a6fe 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -5,17 +5,107 @@ These tests simulate real-world scenarios where headers and configuration need to be properly propagated through the router to the LLM API. """ -from unittest.mock import MagicMock, patch, AsyncMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest - from litellm import Router +from litellm.llms.base_llm.vector_store.transformation import ( + LiteLLMVectorStoreEmbeddingExecutor, + RouterVectorStoreEmbeddingExecutor, +) +from litellm.types.utils import EmbeddingResponse class TestRouterEmbeddingIntegration: """Integration tests for embedding with router configuration.""" + def test_vector_store_request_metadata_prefers_litellm_metadata(self): + assert Router._vector_store_request_metadata( + { + "litellm_metadata": {"user_api_key_team_id": "team-a"}, + "metadata": {"user_api_key_team_id": "team-b"}, + } + ) == {"user_api_key_team_id": "team-a"} + + assert Router._vector_store_request_metadata({"metadata": {"user_api_key_team_id": "team-b"}}) == { + "user_api_key_team_id": "team-b" + } + assert Router._vector_store_request_metadata({}) == {} + + def test_sync_vector_store_wrapper_injects_router_embedding_executor(self): + router = Router(model_list=[]) + original = MagicMock(return_value="searched") + wrapped = router.factory_function(original, call_type="vector_store_search") + + assert ( + wrapped( + vector_store_id="store", + query="query", + custom_llm_provider="valkey", + metadata={"user_api_key_team_id": "team-a"}, + ) + == "searched" + ) + + call_kwargs = original.call_args.kwargs + assert call_kwargs["custom_llm_provider"] == "valkey" + executor = call_kwargs["_direct_vector_store_embedding_executor"] + assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) + assert executor.metadata == {"user_api_key_team_id": "team-a"} + + def test_sync_vector_store_wrapper_preserves_model_routing(self): + router = Router(model_list=[]) + original = MagicMock() + wrapped = router.factory_function(original, call_type="vector_store_search") + + with patch.object(router, "_generic_api_call_with_fallbacks", return_value="routed") as fallback: + assert wrapped(model="vector-alias", vector_store_id="store", query="query") == "routed" + + assert fallback.call_args.kwargs["model"] == "vector-alias" + assert fallback.call_args.kwargs["original_function"] is original + assert isinstance( + fallback.call_args.kwargs["_direct_vector_store_embedding_executor"], + RouterVectorStoreEmbeddingExecutor, + ) + + @pytest.mark.asyncio + async def test_vector_store_embedding_executors_cover_sdk_and_router_paths(self): + response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + sdk_executor = LiteLLMVectorStoreEmbeddingExecutor() + + with ( + patch("litellm.embedding", return_value=response) as embedding, + patch("litellm.aembedding", new=AsyncMock(return_value=response)) as aembedding, + ): + assert sdk_executor.embed("openai/model", "sync", {"api_key": "explicit"}) is response + assert await sdk_executor.aembed("openai/model", "async", {"api_key": "explicit"}) is response + + embedding.assert_called_once_with(model="openai/model", input=["sync"], api_key="explicit") + aembedding.assert_awaited_once_with(model="openai/model", input=["async"], api_key="explicit") + + mock_router = MagicMock() + mock_router.embedding.return_value = response + router_executor = RouterVectorStoreEmbeddingExecutor( + router=mock_router, + metadata={"user_api_key_team_id": "team-a"}, + ) + assert router_executor.embed("team-alias", "query", {}) is response + mock_router.embedding.assert_called_once_with( + model="team-alias", + input=["query"], + metadata={"user_api_key_team_id": "team-a"}, + ) + + with patch("litellm.embedding", return_value=response) as explicit_embedding: + assert router_executor.embed("openai/model", "query", {"api_key": "store-key"}) is response + explicit_embedding.assert_called_once_with(model="openai/model", input=["query"], api_key="store-key") + mock_router.embedding.assert_called_once() + + with patch("litellm.aembedding", new=AsyncMock(return_value=response)) as explicit_aembedding: + assert await router_executor.aembed("openai/model", "query", {"api_key": "store-key"}) is response + explicit_aembedding.assert_awaited_once_with(model="openai/model", input=["query"], api_key="store-key") + def test_embedding_with_deployment_specific_headers(self): """ Test that deployment-specific headers are propagated. diff --git a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py b/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py index a2ee2c2bdb1..aa114f128c5 100644 --- a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py +++ b/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py @@ -67,20 +67,52 @@ class FakeAsyncEmbeddingFn(FakeEmbeddingFn): return SimpleNamespace(data=[{"embedding": self.embedding}]) +class FakeEmbeddingExecutor: + def __init__(self, embedding): + self.embedding = embedding + self.captured = None + + def embed(self, model, query, configuration): + self.captured = (model, query, configuration) + return SimpleNamespace(data=[{"embedding": self.embedding}]) + + async def aembed(self, model, query, configuration): + self.captured = (model, query, configuration) + return SimpleNamespace(data=[{"embedding": self.embedding}]) + + def _doc(doc_id, distance, **fields): return SimpleNamespace(id=doc_id, vector_distance=str(distance), **fields) -def _search(config, client=None, query="what is litellm", optional_params=None, litellm_params=None): +def _search(config, client=None, query="what is litellm", optional_params=None, litellm_params=None, executor=None): return config.execute_search_vector_store_request( vector_store_id="my_index", query=query, vector_store_search_optional_params=optional_params or {}, litellm_logging_obj=MagicMock(), litellm_params={"litellm_embedding_model": "openai/text-embedding-3-small", **(litellm_params or {})}, + embedding_executor=executor, ) +def test_sync_search_uses_request_embedding_executor_without_overwriting_explicit_config(): + executor = FakeEmbeddingExecutor([0.1, 0.2]) + config = ValkeyVectorStoreConfig(sync_client=FakeRedis()) + embedding_config = {"api_key": "store-specific-key", "aws_region_name": "us-west-2"} + + _search( + config, + litellm_params={ + "litellm_embedding_model": "team-embedding-alias", + "litellm_embedding_config": embedding_config, + }, + executor=executor, + ) + + assert executor.captured == ("team-embedding-alias", "what is litellm", embedding_config) + + def test_sync_search_builds_knn_query_with_packed_vector(): embedding_fn = FakeEmbeddingFn([0.1, 0.2, 0.3]) client = FakeRedis() diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 1484adb258f..ad411e874ca 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -2,29 +2,24 @@ from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import Request - - -from fastapi import HTTPException +from fastapi import HTTPException, Request import litellm from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) +from litellm.llms.base_llm.vector_store.transformation import ( + LiteLLMVectorStoreEmbeddingExecutor, + RouterVectorStoreEmbeddingExecutor, +) from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.vector_store_endpoints.endpoints import ( _update_request_data_with_litellm_managed_vector_store_registry, index_create, index_list, ) -from litellm.proxy.vector_store_files_endpoints.endpoints import ( - _update_request_data_with_model_routing_hint, -) from litellm.proxy.vector_store_endpoints.management_endpoints import ( _check_vector_store_access, - _resolve_embedding_config, - _resolve_embedding_config_from_db, - _resolve_embedding_config_from_router, create_vector_store_in_db, new_vector_store, ) @@ -33,8 +28,12 @@ from litellm.proxy.vector_store_endpoints.utils import ( is_allowed_to_call_vector_store_endpoint, is_allowed_to_call_vector_store_files_endpoint, ) +from litellm.proxy.vector_store_files_endpoints.endpoints import ( + _update_request_data_with_model_routing_hint, +) +from litellm.types.utils import EmbeddingResponse, LlmProviders from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse -from litellm.types.utils import LlmProviders +from litellm.vector_stores.main import _direct_vector_store_embedding_executor def _serialize_litellm_params(litellm_params): @@ -51,17 +50,98 @@ def _serialize_litellm_params(litellm_params): return json.dumps(litellm_params or {}) -@pytest.fixture(autouse=True) -def _reset_embedding_config_cache(): - """The use-time embedding-config resolver caches results in process - memory across calls. Reset it before every test so the resolver - actually exercises the router/DB path under test instead of returning - a value cached by an earlier test.""" - from litellm.proxy.vector_store_endpoints import management_endpoints +def test_direct_vector_store_embedding_executor_rejects_invalid_value(): + with pytest.raises(TypeError, match="Invalid direct vector store embedding executor"): + _direct_vector_store_embedding_executor(object()) - management_endpoints._embedding_config_cache = None - yield - management_endpoints._embedding_config_cache = None + +def test_router_vector_store_search_injects_executor_and_request_metadata(): + router = litellm.Router(model_list=[]) + original = MagicMock(return_value="searched") + wrapped = router.factory_function(original, call_type="vector_store_search") + + assert ( + wrapped( + vector_store_id="store", + query="query", + custom_llm_provider="valkey", + litellm_metadata={"user_api_key_team_id": "team-a"}, + ) + == "searched" + ) + + call_kwargs = original.call_args.kwargs + assert call_kwargs["custom_llm_provider"] == "valkey" + executor = call_kwargs["_direct_vector_store_embedding_executor"] + assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) + assert executor.metadata == {"user_api_key_team_id": "team-a"} + assert litellm.Router._vector_store_request_metadata({"metadata": {"user_api_key_team_id": "team-b"}}) == { + "user_api_key_team_id": "team-b" + } + assert litellm.Router._vector_store_request_metadata({}) == {} + + with patch.object( # test-quality-ok: fallback dispatch is the boundary this wrapper delegates to + router, "_generic_api_call_with_fallbacks", return_value="routed" + ) as fallback: + assert wrapped(model="vector-alias", vector_store_id="store", query="query") == "routed" + assert fallback.call_args.kwargs["model"] == "vector-alias" + assert fallback.call_args.kwargs["original_function"] is original + + create_original = MagicMock() + wrapped_create = router.factory_function(create_original, call_type="vector_store_create") + with patch.object( # test-quality-ok: fallback dispatch is the boundary this wrapper delegates to + router, "_generic_api_call_with_fallbacks", return_value="created" + ) as fallback: + assert wrapped_create(name="store") == "created" + fallback.assert_called_once_with(original_function=create_original, name="store") + + +@pytest.mark.asyncio +async def test_vector_store_embedding_executors_preserve_explicit_configuration(): + response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + sdk_executor = LiteLLMVectorStoreEmbeddingExecutor() + + with ( + patch( # test-quality-ok: isolates SDK dispatch from external embedding providers + "litellm.embedding", return_value=response + ) as embedding, + patch( # test-quality-ok: isolates async SDK dispatch from external embedding providers + "litellm.aembedding", new=AsyncMock(return_value=response) + ) as aembedding, + ): + assert sdk_executor.embed("openai/model", "sync", {"api_key": "explicit"}) is response + assert await sdk_executor.aembed("openai/model", "async", {"api_key": "explicit"}) is response + + embedding.assert_called_once_with(model="openai/model", input=["sync"], api_key="explicit") + aembedding.assert_awaited_once_with(model="openai/model", input=["async"], api_key="explicit") + + mock_router = MagicMock() + mock_router.embedding.return_value = response + router_executor = RouterVectorStoreEmbeddingExecutor( + router=mock_router, + metadata={"user_api_key_team_id": "team-a"}, + ) + + assert router_executor.embed("team-alias", "query", {}) is response + mock_router.embedding.assert_called_once_with( + model="team-alias", + input=["query"], + metadata={"user_api_key_team_id": "team-a"}, + ) + + with ( + patch( # test-quality-ok: verifies explicit store configuration at the SDK boundary + "litellm.embedding", return_value=response + ) as explicit_embedding, + patch( # test-quality-ok: verifies async explicit store configuration at the SDK boundary + "litellm.aembedding", new=AsyncMock(return_value=response) + ) as explicit_aembedding, + ): + assert router_executor.embed("openai/model", "query", {"api_key": "store-key"}) is response + assert await router_executor.aembed("openai/model", "query", {"api_key": "store-key"}) is response + + explicit_embedding.assert_called_once_with(model="openai/model", input=["query"], api_key="store-key") + explicit_aembedding.assert_awaited_once_with(model="openai/model", input=["query"], api_key="store-key") @pytest.mark.asyncio @@ -82,10 +162,11 @@ async def test_router_avector_store_search_passes_correct_args(): } # Call router's avector_store_search - result = await router.avector_store_search( + await router.avector_store_search( vector_store_id="test_store_id", query="test query", custom_llm_provider="bedrock", + metadata={"user_api_key_team_id": "team-a"}, ) # Verify the internal method was called with correct args @@ -96,6 +177,38 @@ async def test_router_avector_store_search_passes_correct_args(): assert call_args[1]["vector_store_id"] == "test_store_id" assert call_args[1]["query"] == "test query" assert call_args[1]["custom_llm_provider"] == "bedrock" + executor = call_args[1]["_direct_vector_store_embedding_executor"] + assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) + assert executor.metadata["user_api_key_team_id"] == "team-a" + + +@pytest.mark.asyncio +async def test_vector_store_embedding_executor_uses_team_scoped_router_deployment(): + router = litellm.Router( + model_list=[ + { + "model_name": "shared-embedding", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "team-a-key"}, + "model_info": {"team_id": "team-a", "team_public_model_name": "shared-embedding"}, + }, + { + "model_name": "shared-embedding", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "team-b-key"}, + "model_info": {"team_id": "team-b", "team_public_model_name": "shared-embedding"}, + }, + ] + ) + executor = RouterVectorStoreEmbeddingExecutor( + router=router, + metadata={"user_api_key_team_id": "team-b"}, + ) + response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + + with patch("litellm.aembedding", new=AsyncMock(return_value=response)) as mock_aembedding: + result = await executor.aembed("shared-embedding", "query", {}) + + assert result is response + assert mock_aembedding.await_args.kwargs["api_key"] == "team-b-key" @pytest.mark.asyncio @@ -502,89 +615,30 @@ async def test_update_request_data_with_litellm_managed_vector_store_registry(): @pytest.mark.asyncio -async def test_update_request_data_resolves_embedding_config_at_use_time(): - """When the persisted vector store row carries only a - ``litellm_embedding_model`` reference (the new behaviour after - moving the auto-resolve out of write time), the request-handling - layer must resolve the embedding config so the downstream embed - call still has ``api_key`` / ``api_base`` / ``api_version``. The - resolved config lives in this per-request data dict only — never - persisted.""" - mock_vector_store: LiteLLM_ManagedVectorStore = { +async def test_managed_vector_store_keeps_embedding_reference_and_explicit_config(): + explicit_config = {"api_key": "store-specific-key", "api_base": "https://embedding.example"} + managed_vector_store: LiteLLM_ManagedVectorStore = { "vector_store_id": "test_store", - "custom_llm_provider": "azure_ai", + "custom_llm_provider": "valkey", "litellm_params": { - "litellm_embedding_model": "multilingual-e5-large", - # Note: no litellm_embedding_config persisted + "litellm_embedding_model": "team-embedding-alias", + "litellm_embedding_config": explicit_config, }, } - mock_registry = MagicMock() - mock_registry.get_litellm_managed_vector_store_from_registry.return_value = ( - mock_vector_store - ) + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = managed_vector_store - resolved = { - "api_key": "use-time-resolved-key", - "api_base": "https://my-azure.example", - "api_version": "2024-09-01", - } - - with ( - patch.object(litellm, "vector_store_registry", mock_registry), - patch( - "litellm.proxy.vector_store_endpoints.endpoints._resolve_embedding_config", - new=AsyncMock(return_value=("azure/multilingual-e5-large", resolved)), - ), - ): + with patch.object(litellm, "vector_store_registry", mock_registry): result = await _update_request_data_with_litellm_managed_vector_store_registry( - data={}, vector_store_id="test_store" + data={}, + vector_store_id="test_store", ) - assert result["litellm_embedding_model"] == "azure/multilingual-e5-large" - assert result["litellm_embedding_config"] == resolved + assert result["litellm_embedding_model"] == "team-embedding-alias" + assert result["litellm_embedding_config"] == explicit_config + assert managed_vector_store["litellm_params"]["litellm_embedding_config"] == explicit_config -@pytest.mark.asyncio -async def test_update_request_data_preserves_legacy_embedding_config_when_model_not_resolved(): - """A vector store row created by an older proxy version may already - carry a fully-resolved ``litellm_embedding_config`` in its persisted - ``litellm_params``. Preserve it when the model cannot be resolved.""" - legacy_config = { - "api_key": "legacy-cleartext-key", - "api_base": "https://legacy-azure.example", - "api_version": "2024-01-01", - } - mock_vector_store: LiteLLM_ManagedVectorStore = { - "vector_store_id": "legacy_store", - "custom_llm_provider": "azure_ai", - "litellm_params": { - "litellm_embedding_model": "azure/text-embedding-3-large", - "litellm_embedding_config": legacy_config, - }, - } - - mock_registry = MagicMock() - mock_registry.get_litellm_managed_vector_store_from_registry.return_value = ( - mock_vector_store - ) - - resolve_mock = AsyncMock(return_value=None) - - with ( - patch.object(litellm, "vector_store_registry", mock_registry), - patch( - "litellm.proxy.vector_store_endpoints.endpoints._resolve_embedding_config", - new=resolve_mock, - ), - ): - result = await _update_request_data_with_litellm_managed_vector_store_registry( - data={}, vector_store_id="legacy_store" - ) - - assert result["litellm_embedding_config"] == legacy_config - resolve_mock.assert_awaited_once() - class TestCheckVectorStorePermission: """Test suite for check_vector_store_permission function.""" @@ -2001,60 +2055,7 @@ async def test_vector_store_update_and_list_synchronization(): @pytest.mark.asyncio -async def test_resolve_embedding_config_from_db(): - """Test that _resolve_embedding_config_from_db correctly resolves embedding config from database.""" - mock_prisma_client = MagicMock() - - # Mock database model with litellm_params - mock_db_model = MagicMock() - mock_db_model.litellm_params = { - "model": "openai/text-embedding-3-small", - "api_key": "test-api-key", - "api_base": "https://api.openai.com", - "api_version": "2024-01-01", - } - - mock_prisma_client.db.litellm_proxymodeltable.find_first = AsyncMock( - return_value=mock_db_model - ) - - with patch( - "litellm.proxy.vector_store_endpoints.management_endpoints.decrypt_value_helper", - side_effect=lambda value, key, return_original_value: value, - ): - result = await _resolve_embedding_config_from_db( - embedding_model="text-embedding-ada-002", prisma_client=mock_prisma_client - ) - - assert result is not None - resolved_model, resolved_config = result - assert resolved_model == "openai/text-embedding-3-small" - assert resolved_config["api_key"] == "test-api-key" - assert resolved_config["api_base"] == "https://api.openai.com" - assert resolved_config["api_version"] == "2024-01-01" - mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_called_once_with( - where={"model_name": "text-embedding-ada-002"} - ) - - # Test with empty embedding_model - result_empty = await _resolve_embedding_config_from_db( - embedding_model="", prisma_client=mock_prisma_client - ) - assert result_empty is None - - # Test with model not found - mock_prisma_client.db.litellm_proxymodeltable.find_first = AsyncMock( - return_value=None - ) - result_not_found = await _resolve_embedding_config_from_db( - embedding_model="non-existent-model", prisma_client=mock_prisma_client - ) - assert result_not_found is None - - -@pytest.mark.asyncio -async def test_new_vector_store_auto_resolves_embedding_config(): - """Test that new_vector_store auto-resolves embedding config when embedding_model is provided but config is not.""" +async def test_new_vector_store_persists_embedding_reference_without_credentials(): import json from litellm.types.vector_stores import LiteLLM_ManagedVectorStore @@ -2071,14 +2072,6 @@ async def test_new_vector_store_auto_resolves_embedding_config(): }, } - # Mock database model lookup for embedding config resolution - mock_db_model = MagicMock() - mock_db_model.litellm_params = { - "api_key": "resolved-api-key", - "api_base": "https://api.openai.com", - "api_version": "2024-01-01", - } - # Mock user API key mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.user_role = None @@ -2089,10 +2082,6 @@ async def test_new_vector_store_auto_resolves_embedding_config(): mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( return_value=None # Vector store doesn't exist yet ) - mock_prisma_client.db.litellm_proxymodeltable.find_first = AsyncMock( - return_value=mock_db_model - ) - # Track what was passed to create captured_create_data = {} @@ -2113,280 +2102,21 @@ async def test_new_vector_store_auto_resolves_embedding_config(): mock_registry = MagicMock() mock_registry.add_vector_store_to_registry = MagicMock() - # Mock router to return None (so it falls back to DB resolution) - mock_router = MagicMock() - mock_router.get_deployment_by_model_group_name.return_value = None - with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), - patch("litellm.proxy.proxy_server.llm_router", mock_router), - patch( - "litellm.proxy.vector_store_endpoints.management_endpoints.decrypt_value_helper", - side_effect=lambda value, key, return_original_value: value, - ), patch.object(litellm, "vector_store_registry", mock_registry), ): - result = await new_vector_store( - vector_store=vector_store_data, user_api_key_dict=mock_user_api_key - ) + result = await new_vector_store(vector_store=vector_store_data, user_api_key_dict=mock_user_api_key) assert result["status"] == "success" - # Auto-resolve no longer happens at create time — the persisted row - # carries only the model reference, never the resolved cleartext - # credential. Resolution now happens at request-handling time inside - # ``_update_request_data_with_litellm_managed_vector_store_registry``, - # where the resolved config lives in per-request memory and is never - # written to the database. litellm_params_json = captured_create_data.get("litellm_params") assert litellm_params_json is not None litellm_params_dict = json.loads(litellm_params_json) assert "litellm_embedding_config" not in litellm_params_dict assert litellm_params_dict["litellm_embedding_model"] == "text-embedding-ada-002" - # The response must also not echo a cleartext credential — even on - # the create response, where redaction guards against caller-supplied - # cleartext or pre-existing rows that were created by an earlier - # proxy version. response_vs = result["vector_store"] - assert "resolved-api-key" not in _serialize_litellm_params( - response_vs.get("litellm_params") - ) - - -def test_resolve_embedding_config_from_router(): - """Test that _resolve_embedding_config_from_router correctly extracts credentials from config-defined models.""" - from litellm.types.router import Deployment, LiteLLM_Params - - # Create a mock router with a model - mock_router = MagicMock() - - # Create a mock deployment with litellm_params - mock_litellm_params = MagicMock(spec=LiteLLM_Params) - mock_litellm_params.api_key = "config-api-key" - mock_litellm_params.api_base = "https://config-api-base.com" - mock_litellm_params.api_version = "2024-02-01" - mock_litellm_params.model = "text-embedding-3-small" - mock_litellm_params.custom_llm_provider = "openai" - - mock_deployment = MagicMock(spec=Deployment) - mock_deployment.litellm_params = mock_litellm_params - - mock_router.get_deployment_by_model_group_name.return_value = mock_deployment - - # Test resolution - result = _resolve_embedding_config_from_router( - embedding_model="text-embedding-ada-002", llm_router=mock_router - ) - - assert result is not None - resolved_model, resolved_config = result - assert resolved_model == "openai/text-embedding-3-small" - assert resolved_config["api_key"] == "config-api-key" - assert resolved_config["api_base"] == "https://config-api-base.com" - assert resolved_config["api_version"] == "2024-02-01" - - mock_router.get_deployment_by_model_group_name.assert_called_once_with( - model_group_name="text-embedding-ada-002" - ) - - -def test_resolve_embedding_config_from_router_with_provider_prefix(): - """Test that _resolve_embedding_config_from_router handles provider prefixes like 'azure/model-name'.""" - from litellm.types.router import Deployment, LiteLLM_Params - - # Create a mock router - mock_router = MagicMock() - - # Create a mock deployment - mock_litellm_params = MagicMock(spec=LiteLLM_Params) - mock_litellm_params.api_key = "azure-api-key" - mock_litellm_params.api_base = "https://azure-endpoint.openai.azure.com" - mock_litellm_params.api_version = "2024-02-15" - mock_litellm_params.model = "text-embedding-3-large" - mock_litellm_params.custom_llm_provider = "azure" - - mock_deployment = MagicMock(spec=Deployment) - mock_deployment.litellm_params = mock_litellm_params - - # First call with full name returns None, second call with stripped name returns deployment - mock_router.get_deployment_by_model_group_name.side_effect = [None, mock_deployment] - - result = _resolve_embedding_config_from_router( - embedding_model="azure/text-embedding-3-large", llm_router=mock_router - ) - - assert result is not None - resolved_model, resolved_config = result - assert resolved_model == "azure/text-embedding-3-large" - assert resolved_config["api_key"] == "azure-api-key" - assert resolved_config["api_base"] == "https://azure-endpoint.openai.azure.com" - assert resolved_config["api_version"] == "2024-02-15" - - # Should have tried both the full name and stripped name - assert mock_router.get_deployment_by_model_group_name.call_count == 2 - - -def test_resolve_embedding_config_from_router_returns_none_when_not_found(): - """Test that _resolve_embedding_config_from_router returns None when model is not in router.""" - mock_router = MagicMock() - mock_router.get_deployment_by_model_group_name.return_value = None - - result = _resolve_embedding_config_from_router( - embedding_model="nonexistent-model", llm_router=mock_router - ) - - assert result is None - - -def test_resolve_embedding_config_from_router_handles_os_environ(): - """Test that _resolve_embedding_config_from_router handles os.environ/ prefixed values.""" - from litellm.types.router import Deployment, LiteLLM_Params - - mock_router = MagicMock() - - mock_litellm_params = MagicMock(spec=LiteLLM_Params) - mock_litellm_params.api_key = "os.environ/OPENAI_API_KEY" - mock_litellm_params.api_base = "https://direct-url.com" - mock_litellm_params.api_version = None - mock_litellm_params.model = "text-embedding-3-small" - mock_litellm_params.custom_llm_provider = "openai" - - mock_deployment = MagicMock(spec=Deployment) - mock_deployment.litellm_params = mock_litellm_params - - mock_router.get_deployment_by_model_group_name.return_value = mock_deployment - - with patch( - "litellm.proxy.vector_store_endpoints.management_endpoints.get_secret", - return_value="resolved-from-env", - ) as mock_get_secret: - result = _resolve_embedding_config_from_router( - embedding_model="text-embedding-ada-002", llm_router=mock_router - ) - - assert result is not None - resolved_model, resolved_config = result - assert resolved_model == "openai/text-embedding-3-small" - assert resolved_config["api_key"] == "resolved-from-env" - assert resolved_config["api_base"] == "https://direct-url.com" - assert "api_version" not in resolved_config - - mock_get_secret.assert_called_once_with("os.environ/OPENAI_API_KEY") - - -@pytest.mark.asyncio -async def test_resolve_embedding_config_tries_router_then_db(): - """Test that _resolve_embedding_config tries router first, then falls back to DB.""" - from litellm.types.router import Deployment, LiteLLM_Params - - mock_prisma_client = MagicMock() - mock_router = MagicMock() - - # Router has the model - mock_litellm_params = MagicMock(spec=LiteLLM_Params) - mock_litellm_params.api_key = "router-api-key" - mock_litellm_params.api_base = "https://router-api-base.com" - mock_litellm_params.api_version = None - mock_litellm_params.model = "text-embedding-3-small" - mock_litellm_params.custom_llm_provider = "openai" - - mock_deployment = MagicMock(spec=Deployment) - mock_deployment.litellm_params = mock_litellm_params - - mock_router.get_deployment_by_model_group_name.return_value = mock_deployment - - # DB should NOT be called since router has the model - mock_prisma_client.db.litellm_proxymodeltable.find_first = AsyncMock() - - result = await _resolve_embedding_config( - embedding_model="text-embedding-ada-002", - prisma_client=mock_prisma_client, - llm_router=mock_router, - ) - - assert result is not None - resolved_model, resolved_config = result - assert resolved_model == "openai/text-embedding-3-small" - assert resolved_config["api_key"] == "router-api-key" - - # DB should NOT have been called since router found the model - mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_not_called() - - -@pytest.mark.asyncio -async def test_resolve_embedding_config_caches_result(): - """The first lookup should hit the router/DB; subsequent lookups for - the same model name should return the cached value without touching - the router or the database.""" - from litellm.types.router import Deployment, LiteLLM_Params - - mock_prisma_client = MagicMock() - mock_router = MagicMock() - - mock_litellm_params = MagicMock(spec=LiteLLM_Params) - mock_litellm_params.api_key = "router-api-key" - mock_litellm_params.api_base = "https://router-api-base.com" - mock_litellm_params.api_version = None - - mock_deployment = MagicMock(spec=Deployment) - mock_deployment.litellm_params = mock_litellm_params - mock_router.get_deployment_by_model_group_name.return_value = mock_deployment - - first = await _resolve_embedding_config( - embedding_model="cached-model", - prisma_client=mock_prisma_client, - llm_router=mock_router, - ) - assert first is not None - assert mock_router.get_deployment_by_model_group_name.call_count == 1 - - second = await _resolve_embedding_config( - embedding_model="cached-model", - prisma_client=mock_prisma_client, - llm_router=mock_router, - ) - assert second == first - # Router (and by extension the DB) was not consulted again. - assert mock_router.get_deployment_by_model_group_name.call_count == 1 - - -@pytest.mark.asyncio -async def test_resolve_embedding_config_falls_back_to_db(): - """Test that _resolve_embedding_config falls back to DB when router doesn't have the model.""" - mock_prisma_client = MagicMock() - mock_router = MagicMock() - - # Router doesn't have the model - mock_router.get_deployment_by_model_group_name.return_value = None - - # DB has the model - mock_db_model = MagicMock() - mock_db_model.litellm_params = { - "model": "openai/text-embedding-3-small", - "api_key": "db-api-key", - "api_base": "https://db-api-base.com", - } - mock_prisma_client.db.litellm_proxymodeltable.find_first = AsyncMock( - return_value=mock_db_model - ) - - with patch( - "litellm.proxy.vector_store_endpoints.management_endpoints.decrypt_value_helper", - side_effect=lambda value, key, return_original_value: value, - ): - result = await _resolve_embedding_config( - embedding_model="text-embedding-ada-002", - prisma_client=mock_prisma_client, - llm_router=mock_router, - ) - - assert result is not None - resolved_model, resolved_config = result - assert resolved_model == "openai/text-embedding-3-small" - assert resolved_config["api_key"] == "db-api-key" - - # DB should have been called since router didn't find the model - mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_called() + assert "api_key" not in _serialize_litellm_params(response_vs.get("litellm_params")) @pytest.mark.asyncio @@ -2445,9 +2175,7 @@ async def test_new_vector_store_auto_resolves_from_router(): } return mock_created_vector_store - mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock( - side_effect=mock_create - ) + mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock(side_effect=mock_create) mock_registry = MagicMock() mock_registry.add_vector_store_to_registry = MagicMock() diff --git a/uv.lock b/uv.lock index 27be919eea1..8d886044083 100644 --- a/uv.lock +++ b/uv.lock @@ -9441,19 +9441,19 @@ wheels = [ [[package]] name = "tornado" -version = "6.5.7" +version = "6.5.8" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/64/24/95ec527ad67b76d59299e5465b3935d05e4294b7e0290a3924b7487df30b/tornado-6.5.7.tar.gz", hash = "sha256:66c513a76cda70d53907bc27cf1447557699c2e95aa48ba27a442ff61c3ddfc2", size = 519252, upload-time = "2026-06-08T17:34:51.232Z" } +sdist = { url = "https://files.pythonhosted.org/packages/10/d3/343e5bb989d6515b1646cf3d40135d73f3d5e45339bded401b56cdac24dd/tornado-6.5.8.tar.gz", hash = "sha256:9452e1b208a8bd771e2cb1f2ff564985b9b214bdebbe622793e1799e0a6bd23f", size = 520493, upload-time = "2026-08-07T02:12:42.971Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/02/dc/c7043cab6fed8ae159fc1923ce829ada35c4dbd797d408a43858ffaf9639/tornado-6.5.7-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:148b2eb15c2c765a50796172c1e499649b35f30d2e3c3d3e15913cfa56bfb163", size = 448543, upload-time = "2026-06-08T17:34:38.052Z" }, - { url = "https://files.pythonhosted.org/packages/92/4f/090b1431e5a43df696feceffc268c5383cc079ecb5f08ce58f917109aafe/tornado-6.5.7-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9da38de27f1da3b78a966f0dae12b5a1ea9afe72ca805d84ff06508272ddf100", size = 446707, upload-time = "2026-06-08T17:34:39.594Z" }, - { url = "https://files.pythonhosted.org/packages/37/d8/ef374952fd5da67d4463122c2b8e5a96536ec10b4b339254c6dcde81d01c/tornado-6.5.7-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8d759e71906ee783f8867b93bf26a265743da4c1e2f4a018464c1ba019862972", size = 449774, upload-time = "2026-06-08T17:34:41.204Z" }, - { url = "https://files.pythonhosted.org/packages/35/37/d434c73f4c6e014b745b9b37085f34f40c022f007efff3d7fe65991899f3/tornado-6.5.7-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a46347a18f23fb92b396beebe0fb78f61dda0cc302445202c16203d8a18848b", size = 450745, upload-time = "2026-06-08T17:34:42.531Z" }, - { url = "https://files.pythonhosted.org/packages/b6/2b/56b9aff361d7f1ab728a805ec7d7ea835f8807afa9f5cc690ea0e630efb9/tornado-6.5.7-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7778b30bef919231265e91c69963ce0f49a1e9c07ac900bbe75b19ce2575ba92", size = 450578, upload-time = "2026-06-08T17:34:43.787Z" }, - { url = "https://files.pythonhosted.org/packages/02/30/a7444fb23aa76860a14198fab96ac79f1866b0a6e19e26c4381b0938e50f/tornado-6.5.7-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e726f0c75da7726eec023aa62751ff8878bd2737e34fbdd33b1ae5897d2200f5", size = 449985, upload-time = "2026-06-08T17:34:45.326Z" }, - { url = "https://files.pythonhosted.org/packages/5c/42/5f0e56c01e8d9d36f4e23f367b85ae6cae0c1ecddd5e6977d8388ad27488/tornado-6.5.7-cp39-abi3-win32.whl", hash = "sha256:f8de3bf12d3efdd0cbe7c8887868198f8a91415e3f29fcf258d9b8eb7b1d9ae4", size = 451047, upload-time = "2026-06-08T17:34:46.784Z" }, - { url = "https://files.pythonhosted.org/packages/c9/a4/b393076ffb21b469eec5b328a0534cf03a3b90bfc6b1f09507cdd075d938/tornado-6.5.7-cp39-abi3-win_amd64.whl", hash = "sha256:de942f843533a039ef9fa3d9c88c7cd8a7c94553fb5ad0154270989b3d99a2c4", size = 451485, upload-time = "2026-06-08T17:34:48.248Z" }, - { url = "https://files.pythonhosted.org/packages/71/2e/7b1c769803121b809112cf9a00681c472eae1d80e32d7ec0e0bd61d0d0e1/tornado-6.5.7-cp39-abi3-win_arm64.whl", hash = "sha256:ff934fce95643af5f11efdae618eaa73d469dc588641e5c8d19295a0c65c4796", size = 450506, upload-time = "2026-06-08T17:34:49.702Z" }, + { url = "https://files.pythonhosted.org/packages/f2/d5/007086fd8df5489338e204f65adce33fd4f21a4999dbb2b9cff2f897b5f4/tornado-6.5.8-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:cc6aa787d7cfab7c3d35189dc7a56fbd2399a569624c730c6b55b3d6531d0403", size = 449487, upload-time = "2026-08-07T02:12:28.682Z" }, + { url = "https://files.pythonhosted.org/packages/70/c8/5a24a99495903f594f6a199dd7beead1cbc0a13e2cb9102727bcaaf2a997/tornado-6.5.8-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9715b5eb79735b2bcd454ce216a9275b7c0470e64ea1bf5742f78b2f72b26eeb", size = 447649, upload-time = "2026-08-07T02:12:30.306Z" }, + { url = "https://files.pythonhosted.org/packages/6e/de/f2e733f386b85962d1b1dc82cd63d169b5b4580062b35397eac9244a41fe/tornado-6.5.8-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:547d63f450d570c14fe0e8db2cfb14c9bbd1c2503b4a6612586267955aa47b58", size = 450707, upload-time = "2026-08-07T02:12:31.95Z" }, + { url = "https://files.pythonhosted.org/packages/0b/94/20efeee9a01c141e9ac47c397f81679dfda24b32768fc4fff24e76d36c2c/tornado-6.5.8-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e2360a0ffbe145eca8af0b19cb7203d79b1a98dd4cccdd6b368f6f49c2e3808", size = 451677, upload-time = "2026-08-07T02:12:33.512Z" }, + { url = "https://files.pythonhosted.org/packages/42/ec/a96ccb8ccf0de2b7bc2c5fa1608a4803735018242e90c4882365a9fd418f/tornado-6.5.8-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:5d242290bdf7ab3151bc1065fdd75c0dcc21cbc7b49f22a4c56329c2d6566d22", size = 451510, upload-time = "2026-08-07T02:12:35.346Z" }, + { url = "https://files.pythonhosted.org/packages/29/b5/93185859245ad3f00e62175f29607346788b696369347f0146e0421286bb/tornado-6.5.8-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7b94ff0e128fe0542f3bd331fb44d06260fc4ac16881545159f34ef08aad4195", size = 450917, upload-time = "2026-08-07T02:12:36.963Z" }, + { url = "https://files.pythonhosted.org/packages/97/cf/fe33cf062834487d34d1559746a4a12521033c22645b6d74d4bca702e018/tornado-6.5.8-cp39-abi3-win32.whl", hash = "sha256:67832909c4779c64942380cb5f044a5c6163d00831472d80e25e115de9917836", size = 451952, upload-time = "2026-08-07T02:12:38.512Z" }, + { url = "https://files.pythonhosted.org/packages/cb/e1/468ad54333e92ccb62627e62cb88e5fc14a2171daa67ed47b1b8542d5b86/tornado-6.5.8-cp39-abi3-win_amd64.whl", hash = "sha256:11881db6b7c168494be2c2d12e65931451bdf7ee718535418ae1d8855dd5a0ee", size = 452391, upload-time = "2026-08-07T02:12:39.971Z" }, + { url = "https://files.pythonhosted.org/packages/ad/3e/cd5e4f06e34cde33b8ef66cf36aa2b5ad46354cc1af7d2136bbe365fee1d/tornado-6.5.8-cp39-abi3-win_arm64.whl", hash = "sha256:68a7468c7e289f8514d7d664101753903217eff1bb6822c6b5994a0b5f5bcb26", size = 451411, upload-time = "2026-08-07T02:12:41.469Z" }, ] [[package]] From 6805d01709f9401f36bba163dc7f9c3192643c51 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 14:05:12 -0700 Subject: [PATCH 19/69] fix(vector-store): preserve aliases with embedding config --- .../base_llm/vector_store/transformation.py | 21 +++++--- .../test_router_embedding_integration.py | 53 ++++++++++++++++--- .../test_vector_store_endpoints.py | 17 +++++- 3 files changed, 75 insertions(+), 16 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 772e4f849a0..e9c925448a8 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -3,7 +3,7 @@ from __future__ import annotations from abc import abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, NoReturn, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, runtime_checkable import httpx @@ -65,22 +65,29 @@ class RouterVectorStoreEmbeddingExecutor: router: Router metadata: Mapping[str, object] + def _embedding_kwargs(self, configuration: Mapping[str, object]) -> dict[str, object]: + configured_metadata: Final = configuration.get("metadata") + metadata: Final = { + **(configured_metadata if isinstance(configured_metadata, Mapping) else {}), + **self.metadata, + } + return { + **{key: value for key, value in configuration.items() if key not in ("input", "metadata", "model")}, + "metadata": metadata, + } + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: - if configuration: - return LiteLLMVectorStoreEmbeddingExecutor().embed(model, query, configuration) return self.router.embedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, input=[query], # mutable-ok: Router embedding requires a mutable input list - metadata=dict(self.metadata), # mutable-ok: Router metadata requires a concrete dict + **self._embedding_kwargs(configuration), # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic ) async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: - if configuration: - return await LiteLLMVectorStoreEmbeddingExecutor().aembed(model, query, configuration) return await self.router.aembedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, input=[query], # mutable-ok: Router embedding requires a mutable input list - metadata=dict(self.metadata), # mutable-ok: Router metadata requires a concrete dict + **self._embedding_kwargs(configuration), # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic ) diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index 5c01587a6fe..d5e0c750d88 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -97,14 +97,53 @@ class TestRouterEmbeddingIntegration: metadata={"user_api_key_team_id": "team-a"}, ) - with patch("litellm.embedding", return_value=response) as explicit_embedding: - assert router_executor.embed("openai/model", "query", {"api_key": "store-key"}) is response - explicit_embedding.assert_called_once_with(model="openai/model", input=["query"], api_key="store-key") - mock_router.embedding.assert_called_once() + alias_router = Router( + model_list=[ + { + "model_name": "team-alias", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + } + ] + ) + alias_executor = RouterVectorStoreEmbeddingExecutor( + router=alias_router, + metadata={"user_api_key_team_id": "team-a"}, + ) + explicit_config = { + "api_base": "https://embedding.example/v1", + "api_key": "store-key", + "metadata": { + "configured": True, + "user_api_key_team_id": "untrusted-team", + }, + "model": "untrusted-model", + } - with patch("litellm.aembedding", new=AsyncMock(return_value=response)) as explicit_aembedding: - assert await router_executor.aembed("openai/model", "query", {"api_key": "store-key"}) is response - explicit_aembedding.assert_awaited_once_with(model="openai/model", input=["query"], api_key="store-key") + with ( + patch("litellm.embedding", return_value=response) as explicit_embedding, + patch("litellm.aembedding", new=AsyncMock(return_value=response)) as explicit_aembedding, + ): + assert alias_executor.embed("team-alias", "sync query", explicit_config) is response + assert await alias_executor.aembed("team-alias", "async query", explicit_config) is response + + sync_kwargs = explicit_embedding.call_args.kwargs + assert sync_kwargs["model"] == "openai/text-embedding-3-small" + assert sync_kwargs["input"] == ["sync query"] + assert sync_kwargs["api_base"] == "https://embedding.example/v1" + assert sync_kwargs["api_key"] == "store-key" + assert sync_kwargs["metadata"]["configured"] is True + assert sync_kwargs["metadata"]["user_api_key_team_id"] == "team-a" + + async_kwargs = explicit_aembedding.await_args.kwargs + assert async_kwargs["model"] == "openai/text-embedding-3-small" + assert async_kwargs["input"] == ["async query"] + assert async_kwargs["api_base"] == "https://embedding.example/v1" + assert async_kwargs["api_key"] == "store-key" + assert async_kwargs["metadata"]["configured"] is True + assert async_kwargs["metadata"]["user_api_key_team_id"] == "team-a" def test_embedding_with_deployment_specific_headers(self): """ diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index ad411e874ca..903d5cb55f3 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -117,6 +117,7 @@ async def test_vector_store_embedding_executors_preserve_explicit_configuration( mock_router = MagicMock() mock_router.embedding.return_value = response + mock_router.aembedding = AsyncMock(return_value=response) router_executor = RouterVectorStoreEmbeddingExecutor( router=mock_router, metadata={"user_api_key_team_id": "team-a"}, @@ -140,8 +141,20 @@ async def test_vector_store_embedding_executors_preserve_explicit_configuration( assert router_executor.embed("openai/model", "query", {"api_key": "store-key"}) is response assert await router_executor.aembed("openai/model", "query", {"api_key": "store-key"}) is response - explicit_embedding.assert_called_once_with(model="openai/model", input=["query"], api_key="store-key") - explicit_aembedding.assert_awaited_once_with(model="openai/model", input=["query"], api_key="store-key") + explicit_embedding.assert_not_called() + explicit_aembedding.assert_not_awaited() + assert mock_router.embedding.call_args.kwargs == { + "model": "openai/model", + "input": ["query"], + "api_key": "store-key", + "metadata": {"user_api_key_team_id": "team-a"}, + } + mock_router.aembedding.assert_awaited_once_with( + model="openai/model", + input=["query"], + api_key="store-key", + metadata={"user_api_key_team_id": "team-a"}, + ) @pytest.mark.asyncio From 5799a32cdda6647d2f16460d79b8610dbae49d34 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 15:37:50 -0700 Subject: [PATCH 20/69] fix(vector-store): route pre-call searches through router --- .../vector_store_pre_call_hook.py | 22 ++++++++-- .../test_bedrock_knowledgebase_hook.py | 42 +++++++++++++++++++ 2 files changed, 61 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 07d4f959489..aaf5cb080dc 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -80,10 +80,15 @@ class VectorStorePreCallHook(CustomLogger): # Get prisma_client for database fallback prisma_client = None + llm_router = None try: - from litellm.proxy.proxy_server import prisma_client as _prisma_client + from litellm.proxy.proxy_server import ( + llm_router as _llm_router, + prisma_client as _prisma_client, + ) prisma_client = _prisma_client + llm_router = _llm_router except ImportError: pass @@ -114,12 +119,23 @@ class VectorStorePreCallHook(CustomLogger): vector_store_id = vector_store_to_run.get("vector_store_id", "") custom_llm_provider = vector_store_to_run.get("custom_llm_provider") litellm_params_for_vector_store = vector_store_to_run.get("litellm_params", {}) or {} - # Call litellm.vector_stores.search() with the required parameters - search_response = await litellm.vector_stores.asearch( + request_litellm_params: Final = ( + litellm_logging_obj.model_call_details.get("litellm_params", {}) + if litellm_logging_obj is not None + else {} + ) + request_metadata: Final = ( + request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {} + ) + search_function: Final = ( + llm_router.avector_store_search if llm_router is not None else litellm.vector_stores.asearch + ) + search_response = await search_function( **{ "vector_store_id": vector_store_id, "query": query, "custom_llm_provider": custom_llm_provider, + "metadata": request_metadata, **litellm_params_for_vector_store, }, ) diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 3f9f2bacdd3..06083b77e84 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -71,6 +71,48 @@ def setup_vector_store_registry(): ) +@pytest.mark.asyncio +async def test_vector_store_hook_routes_search_through_proxy_router( + setup_vector_store_registry, +): + proxy_router = Mock() + proxy_router.avector_store_search = AsyncMock( + return_value=VectorStoreSearchResponse( + object="vector_store.search_results.page", + search_query="what is litellm?", + data=[ + VectorStoreSearchResult( + score=1.0, + content=[VectorStoreResultContent(text="routed context", type="text")], + ) + ], + ) + ) + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_params": {"metadata": {"user_api_key_team_id": "team-a"}} + } + + with patch("litellm.proxy.proxy_server.llm_router", proxy_router): + _, messages, _ = await VectorStorePreCallHook().async_get_chat_completion_prompt( + model="chat-model", + messages=[{"role": "user", "content": "what is litellm?"}], + non_default_params={"vector_store_ids": ["T37J8R4WTM"]}, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + litellm_logging_obj=logging_obj, + ) + + proxy_router.avector_store_search.assert_awaited_once_with( + vector_store_id="T37J8R4WTM", + query="what is litellm?", + custom_llm_provider="bedrock", + metadata={"user_api_key_team_id": "team-a"}, + ) + assert messages[0]["content"] == "Context:\n\nrouted context\n\n" + + @pytest.mark.asyncio async def test_e2e_bedrock_knowledgebase_retrieval_with_completion( setup_vector_store_registry, From 0cc0c47f8b2a4196beb4e766bcc65fe84743e9ed Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 15:52:57 -0700 Subject: [PATCH 21/69] style(vector-store): satisfy import lint --- .../vector_store_pre_call_hook.py | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index aaf5cb080dc..e012d35b8f3 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -5,6 +5,7 @@ This hook is called before making an LLM request when a vector store is configur It searches the vector store for relevant context and appends it to the messages. """ +from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any, Final, cast import litellm @@ -84,6 +85,8 @@ class VectorStorePreCallHook(CustomLogger): try: from litellm.proxy.proxy_server import ( llm_router as _llm_router, + ) + from litellm.proxy.proxy_server import ( prisma_client as _prisma_client, ) @@ -119,17 +122,20 @@ class VectorStorePreCallHook(CustomLogger): vector_store_id = vector_store_to_run.get("vector_store_id", "") custom_llm_provider = vector_store_to_run.get("custom_llm_provider") litellm_params_for_vector_store = vector_store_to_run.get("litellm_params", {}) or {} - request_litellm_params: Final = ( - litellm_logging_obj.model_call_details.get("litellm_params", {}) - if litellm_logging_obj is not None - else {} - ) - request_metadata: Final = ( + request_litellm_params = litellm_logging_obj.model_call_details.get("litellm_params", {}) + request_metadata = ( request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {} ) - search_function: Final = ( - llm_router.avector_store_search if llm_router is not None else litellm.vector_stores.asearch - ) + if llm_router is not None: + search_function = cast( # cast-ok: normalize router search callable + Callable[..., Awaitable[VectorStoreSearchResponse]], + llm_router.avector_store_search, + ) + else: + search_function = cast( # cast-ok: normalize SDK search callable + Callable[..., Awaitable[VectorStoreSearchResponse]], + litellm.vector_stores.asearch, + ) search_response = await search_function( **{ "vector_store_id": vector_store_id, From 1cd99a036e0538bb61b280e17639545c75374d81 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 17:29:58 -0700 Subject: [PATCH 22/69] fix(router): route Claude Code subagents through session router --- litellm/router.py | 80 +++++++++++++++++- tests/test_litellm/test_router.py | 129 ++++++++++++++++++++++++++++++ 2 files changed, 208 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 6e4405ebfef..b49269f2457 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -353,6 +353,8 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset({"model", "api_base", "api_key", "api_version"}) _ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" +_CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") +_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) -> bool: @@ -12546,6 +12548,77 @@ class Router: return None return candidates[0] + @staticmethod + def _request_header(request_kwargs: Mapping[str, object], header_name: str) -> str | None: + proxy_server_request: Final = request_kwargs.get("proxy_server_request") + if not isinstance(proxy_server_request, Mapping): + return None + headers: Final = proxy_server_request.get("headers") + if not isinstance(headers, Mapping): + return None + return next( + ( + value + for key, value in headers.items() + if isinstance(key, str) and key.lower() == header_name and isinstance(value, str) + ), + None, + ) + + def _claude_code_session_router_cache_key(self, request_kwargs: Mapping[str, object]) -> str | None: + session_id: Final = self._request_header(request_kwargs, "x-claude-code-session-id") + if session_id is None or _CLAUDE_CODE_SESSION_ID_RE.fullmatch(session_id) is None: + return None + metadata_name: Final = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata" + metadata: Final = request_kwargs.get(metadata_name) + if not isinstance(metadata, Mapping): + return None + caller_scope: Final = metadata.get("user_api_key_hash") + if not isinstance(caller_scope, str) or not caller_scope: + return None + return f"claude_code_session_router:v1:{caller_scope}:{session_id}" + + async def _resolve_claude_code_session_router( + self, + model: str, + registered_model_name: str, + request_kwargs: Mapping[str, object], + ) -> str: + cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) + if cache_key is None or not isinstance(request_kwargs, dict): + return registered_model_name + + agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") + if agent_id is not None: + bound_model: Final = await self.cache.async_get_cache(key=cache_key) + if not isinstance(bound_model, str): + return registered_model_name + bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model + if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: + await self.cache.async_delete_cache(key=cache_key) + return registered_model_name + await self.cache.async_set_cache( + key=cache_key, + value=bound_model, + ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, + ) + self._stamp_or_clear_metadata_key(request_kwargs, "model_group", bound_model) + return bound_registered_model + + if self._request_header(request_kwargs, "x-app") != "cli": + return registered_model_name + if request_kwargs.get("fallback_depth") not in (None, 0): + return registered_model_name + if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: + await self.cache.async_delete_cache(key=cache_key) + return registered_model_name + await self.cache.async_set_cache( + key=cache_key, + value=model, + ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, + ) + return registered_model_name + async def async_pre_routing_hook( self, model: str, @@ -12565,7 +12638,12 @@ class Router: the alias, since spend metadata is stamped before routing and the response carries the tier group the strategy picked. """ - registered_model_name: Final = self._get_model_from_alias(model=model) or model + requested_registered_model_name: Final = self._get_model_from_alias(model=model) or model + registered_model_name: Final = await self._resolve_claude_code_session_router( + model=model, + registered_model_name=requested_registered_model_name, + request_kwargs=request_kwargs, + ) ######################################################### # Run the routing-plugin pipeline, if any plugins are configured. diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 84f6344be35..d22a1cc04ca 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8316,6 +8316,135 @@ class TestConsumedRequestTagsStamp: assert CONSUMED_REQUEST_TAGS_METADATA_KEY not in request_kwargs["metadata"] +class TestClaudeCodeSubagentSessionRouterBinding: + class _RewriteStrategy: + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + return PreRoutingHookResponse( + model="cheap-model", + messages=messages, + routing_decision={ + "router_model_name": "smart-router", + "router_type": "complexity", + "routed_model": "cheap-model", + "cause": "heuristic_scorer", + }, + ) + + @classmethod + def _router(cls) -> "litellm.Router": + from litellm.types.router import TaggedPreRoutingStrategy + + router = litellm.Router( + model_list=[ + { + "model_name": "cheap-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "cheap response"}, + }, + { + "model_name": "expensive-model", + "litellm_params": {"model": "openai/gpt-4o", "mock_response": "expensive response"}, + }, + ] + ) + router.complexity_routers = { + "smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy())] + } + return router + + @staticmethod + def _request_kwargs( + *, + key_hash: str = "key-hash-a", + app: str = "cli", + agent_id: str | None = None, + fallback_depth: int | None = None, + ) -> dict: + headers = { + "X-Claude-Code-Session-Id": "session-1234", + "x-app": app, + **({"x-claude-code-agent-id": agent_id} if agent_id is not None else {}), + } + return { + "metadata": {"user_api_key_hash": key_hash}, + "proxy_server_request": {"headers": headers}, + **({"fallback_depth": fallback_depth} if fallback_depth is not None else {}), + } + + @pytest.mark.asyncio + async def test_subagent_concrete_model_uses_the_main_sessions_router(self): + router = self._router() + + await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "main turn"}], + **self._request_kwargs(), + ) + subagent_kwargs = self._request_kwargs(agent_id="agent-1234") + + response = await router.acompletion( + model="expensive-model", + messages=[{"role": "user", "content": "subagent turn"}], + **subagent_kwargs, + ) + + assert response.choices[0].message.content == "cheap response" + assert subagent_kwargs["metadata"]["model_group"] == "smart-router" + assert subagent_kwargs["metadata"]["routing_decision"]["router_model_name"] == "smart-router" + + @pytest.mark.asyncio + async def test_main_direct_model_clears_the_session_router(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + await router.async_pre_routing_hook(model="expensive-model", request_kwargs=self._request_kwargs()) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert response is None + + @pytest.mark.asyncio + async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(app="cli-bg"), + ) + await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(fallback_depth=1), + ) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert response is not None + assert response.model == "cheap-model" + + @pytest.mark.asyncio + async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(key_hash="key-hash-b", agent_id="agent-1234"), + ) + + assert response is None + + class TestAutoRouterMaxInputCharsWiring: """`auto_router_max_input_chars` on the deployment has to reach the AutoRouter that embeds prompts. From 82d046c42821ca3b24036dd57e362252564c0596 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 17:43:48 -0700 Subject: [PATCH 23/69] fix(router): make Claude session cleanup best effort --- litellm/router.py | 14 ++++++++++++-- tests/test_litellm/test_router.py | 19 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b49269f2457..2bb725b0880 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12578,6 +12578,16 @@ class Router: return None return f"claude_code_session_router:v1:{caller_scope}:{session_id}" + async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: + try: + await self.cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request + verbose_router_logger.warning( + "Failed to delete Claude Code session router binding; " + "the binding may remain until its TTL expires: %s", + e, + ) + async def _resolve_claude_code_session_router( self, model: str, @@ -12595,7 +12605,7 @@ class Router: return registered_model_name bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: - await self.cache.async_delete_cache(key=cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, @@ -12610,7 +12620,7 @@ class Router: if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self.cache.async_delete_cache(key=cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d22a1cc04ca..77b73a2d12a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8409,6 +8409,25 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None + @pytest.mark.asyncio + async def test_redis_cleanup_failure_does_not_reject_a_direct_model_request(self): + from litellm.caching.caching import RedisCache + + router = self._router() + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + router._update_redis_cache(cache=redis_cache) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(), + ) + + assert response is None + redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From e3a61c82da9f8dbe42fdfcc4907af7b3a2901392 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 17:46:41 -0700 Subject: [PATCH 24/69] test(router): register indirect session routing coverage --- tests/code_coverage_tests/router_code_coverage.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index a5e00799519..60b56b7fac6 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -82,6 +82,10 @@ ignored_function_names = [ "_invalidate_access_groups_cache", # Tested indirectly via set_model_list, upsert_model etc. (test files lack "router" in name) "has_buffered_provider_output", # Property, so its reads in test_router.py are never an ast.Call "_resolved_provider", # Tested via get_pattern in test_pattern_match_deployments.py (file lacks "router" in name) + "_request_header", # Tested through Claude Code session routing in test_router.py + "_claude_code_session_router_cache_key", # Tested through Claude Code session routing in test_router.py + "_delete_claude_code_session_router_binding", # Tested through Redis cleanup failure in test_router.py + "_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py ] From 2d4301589c1e741489f68e6eef3c2d112da91e2a Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Wed, 2 Sep 2026 01:10:31 +0000 Subject: [PATCH 25/69] fix(router): keep serving when Claude Code session router cleanup fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 12 ++++-------- tests/test_litellm/test_router.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 2bb725b0880..d6d9f20085f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12578,15 +12578,11 @@ class Router: return None return f"claude_code_session_router:v1:{caller_scope}:{session_id}" - async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: + async def _clear_claude_code_session_router(self, cache_key: str) -> None: try: await self.cache.async_delete_cache(key=cache_key) except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request - verbose_router_logger.warning( - "Failed to delete Claude Code session router binding; " - "the binding may remain until its TTL expires: %s", - e, - ) + verbose_router_logger.debug("Claude Code session router cleanup skipped for %s: %s", cache_key, e) async def _resolve_claude_code_session_router( self, @@ -12605,7 +12601,7 @@ class Router: return registered_model_name bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) + await self._clear_claude_code_session_router(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, @@ -12620,7 +12616,7 @@ class Router: if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) + await self._clear_claude_code_session_router(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 77b73a2d12a..b66dbf6aa7d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8428,6 +8428,24 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_main_direct_model_still_served_when_cache_delete_fails(self): + router = self._router() + await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "main turn"}], **self._request_kwargs() + ) + + async def failing_delete(key: str) -> None: + raise Exception("Redis circuit breaker is open — skipping async_delete_cache") + + router.cache.async_delete_cache = failing_delete + + response = await router.acompletion( + model="expensive-model", messages=[{"role": "user", "content": "direct turn"}], **self._request_kwargs() + ) + + assert response.choices[0].message.content == "expensive response" + @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From 46502f58042a41619648be7acc67108079f737e2 Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Wed, 2 Sep 2026 01:11:23 +0000 Subject: [PATCH 26/69] Revert "fix(router): keep serving when Claude Code session router cleanup fails" This reverts commit 2d4301589c1e741489f68e6eef3c2d112da91e2a. --- litellm/router.py | 12 ++++++++---- tests/test_litellm/test_router.py | 18 ------------------ 2 files changed, 8 insertions(+), 22 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index d6d9f20085f..2bb725b0880 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12578,11 +12578,15 @@ class Router: return None return f"claude_code_session_router:v1:{caller_scope}:{session_id}" - async def _clear_claude_code_session_router(self, cache_key: str) -> None: + async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: try: await self.cache.async_delete_cache(key=cache_key) except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request - verbose_router_logger.debug("Claude Code session router cleanup skipped for %s: %s", cache_key, e) + verbose_router_logger.warning( + "Failed to delete Claude Code session router binding; " + "the binding may remain until its TTL expires: %s", + e, + ) async def _resolve_claude_code_session_router( self, @@ -12601,7 +12605,7 @@ class Router: return registered_model_name bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: - await self._clear_claude_code_session_router(cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, @@ -12616,7 +12620,7 @@ class Router: if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self._clear_claude_code_session_router(cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b66dbf6aa7d..77b73a2d12a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8428,24 +8428,6 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() - @pytest.mark.asyncio - async def test_main_direct_model_still_served_when_cache_delete_fails(self): - router = self._router() - await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "main turn"}], **self._request_kwargs() - ) - - async def failing_delete(key: str) -> None: - raise Exception("Redis circuit breaker is open — skipping async_delete_cache") - - router.cache.async_delete_cache = failing_delete - - response = await router.acompletion( - model="expensive-model", messages=[{"role": "user", "content": "direct turn"}], **self._request_kwargs() - ) - - assert response.choices[0].message.content == "expensive response" - @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From 8a0967443d84ecc02c10caf6ca55385b907b11b2 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 20:38:13 -0700 Subject: [PATCH 27/69] fix(router): isolate Claude session binding cache --- litellm/router.py | 19 +++++++++++++------ tests/test_litellm/test_router.py | 14 ++++++++++++++ 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 2bb725b0880..fb8625bc6ba 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -782,6 +782,10 @@ class Router: self.cache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache() ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. + self._claude_code_session_router_cache: DualCache = DualCache( + redis_cache=redis_cache, + in_memory_cache=InMemoryCache(), + ) ### SCHEDULER ### self.scheduler = Scheduler(polling_interval=polling_interval, redis_cache=redis_cache) @@ -1102,8 +1106,8 @@ class Router: ``` and caching to just work. """ - if self.cache.redis_cache is None: - self.cache.redis_cache = cache + self.cache.attach_redis_cache(cache) + self._claude_code_session_router_cache.attach_redis_cache(cache) # Maps a routing strategy string to the attribute on `self` that holds # the default group's strategy selector for that strategy. (The selectors @@ -12580,7 +12584,7 @@ class Router: async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: try: - await self.cache.async_delete_cache(key=cache_key) + await self._claude_code_session_router_cache.async_delete_cache(key=cache_key) except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request verbose_router_logger.warning( "Failed to delete Claude Code session router binding; " @@ -12600,14 +12604,14 @@ class Router: agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") if agent_id is not None: - bound_model: Final = await self.cache.async_get_cache(key=cache_key) + bound_model: Final = await self._claude_code_session_router_cache.async_get_cache(key=cache_key) if not isinstance(bound_model, str): return registered_model_name bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name - await self.cache.async_set_cache( + await self._claude_code_session_router_cache.async_set_cache( key=cache_key, value=bound_model, ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, @@ -12622,7 +12626,7 @@ class Router: if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name - await self.cache.async_set_cache( + await self._claude_code_session_router_cache.async_set_cache( key=cache_key, value=model, ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, @@ -13450,6 +13454,9 @@ class Router: def flush_cache(self): litellm.cache = None self.cache.flush_cache() + session_in_memory_cache: Final = self._claude_code_session_router_cache.in_memory_cache + if session_in_memory_cache is not None: + session_in_memory_cache.flush_cache() def reset(self): ## clean up on close diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 77b73a2d12a..a528266d738 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8428,6 +8428,20 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_session_bindings_do_not_evict_router_rate_limit_state(self): + router = self._router() + assert router._update_usage(deployment_id="deployment-id", parent_otel_span=None) == 1 + + for session_index in range(201): + request_kwargs = self._request_kwargs() + request_kwargs["proxy_server_request"]["headers"]["X-Claude-Code-Session-Id"] = ( + f"session-{session_index:04d}" + ) + await router.async_pre_routing_hook(model="smart-router", request_kwargs=request_kwargs) + + assert router._update_usage(deployment_id="deployment-id", parent_otel_span=None) == 2 + @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From 6adc14b4b12250ddc927682c8153469976e87888 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 21:13:45 -0700 Subject: [PATCH 28/69] fix(router): preserve Claude subagent fallbacks --- litellm/router.py | 4 ++-- tests/test_litellm/test_router.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index fb8625bc6ba..6f1f1bc700b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12601,6 +12601,8 @@ class Router: cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) if cache_key is None or not isinstance(request_kwargs, dict): return registered_model_name + if request_kwargs.get("fallback_depth") not in (None, 0): + return registered_model_name agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") if agent_id is not None: @@ -12621,8 +12623,6 @@ class Router: if self._request_header(request_kwargs, "x-app") != "cli": return registered_model_name - if request_kwargs.get("fallback_depth") not in (None, 0): - return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index a528266d738..b413bb18f04 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8464,6 +8464,19 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is not None assert response.model == "cheap-model" + @pytest.mark.asyncio + async def test_subagent_fallback_does_not_reapply_the_session_router(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234", fallback_depth=1), + ) + + assert response is None + @pytest.mark.asyncio async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From 1400070d711f645290fb382564e1f21200c5e610 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 07:58:09 +0000 Subject: [PATCH 29/69] chore(techdebt): clear fresh debt from the 2026-09-01 window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 6 +++--- .../integrations/SlackAlerting/slack_alerting.py | 9 ++++----- .../websearch_interception/handler.py | 1 - .../_experimental/mcp_server/rest_endpoints.py | 13 +++++++++---- .../guardrails/guardrail_hooks/alice/alice.py | 16 ++++++++-------- type-discipline-budget.json | 6 +++--- 6 files changed, 27 insertions(+), 24 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 2d4adc02234..f14f8e002dd 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15300 + "limit": 15298 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38347 + "limit": 38344 }, "reportUnknownParameterType": { "limit": 19626 }, "reportUnknownVariableType": { - "limit": 29884 + "limit": 29880 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 748ef938cea..dc41c7dadc8 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -1955,11 +1955,10 @@ Model Info: if not thresholds_enabled and not anomalies_enabled: return - if prisma_client is None: - from litellm.proxy.proxy_server import prisma_client as global_prisma_client + from litellm.proxy.proxy_server import prisma_client as global_prisma_client - prisma_client = global_prisma_client # rebind-ok: fall back to the proxy's global client - if prisma_client is None: + client: Final = prisma_client if prisma_client is not None else global_prisma_client + if client is None: return from litellm.integrations.SlackAlerting.user_spend_alerts import ( @@ -1970,7 +1969,7 @@ Model Info: try: today: Final = datetime.datetime.now(datetime.timezone.utc).date() rows: Final = await fetch_user_spend_rows( - prisma_client=prisma_client, + prisma_client=client, today=today, baseline_days=self.alerting_args.spend_anomaly_baseline_days, ) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index dc61ee38a8c..2d737bc34e7 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -419,7 +419,6 @@ class WebSearchInterceptionLogger(CustomLogger): if call_type in (CallTypes.responses, CallTypes.aresponses): return self._convert_responses_tools(kwargs=kwargs, tools=tools) - # Check if any tool is a web search tool (native or already LiteLLM standard) has_websearch: Final = any(is_web_search_tool(t) for t in tools) if not has_websearch: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d1ef73a15cd..90474bfc5e6 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -92,6 +92,9 @@ def _connection_error_message(exc: BaseException) -> str: if MCP_AVAILABLE: + from mcp.types import Tool as MCPTool + + from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, global_mcp_server_manager, @@ -876,7 +879,6 @@ if MCP_AVAILABLE: return (), classify_list_exception(e) return tools_result, ServerListOk(tool_count=len(tools_result)) - # Query all servers the user has access to queried_servers: Final = tuple( server for server in map(global_mcp_server_manager.get_mcp_server_by_id, allowed_server_ids) @@ -1141,6 +1143,11 @@ if MCP_AVAILABLE: scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None return client_id, client_secret, scopes + async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None: + with anyio.move_on_after(deadline): + return await client.list_tools(raise_on_error=True) + return None + async def _execute_with_mcp_client( request: NewMCPServerRequest, operation: Callable[..., Awaitable[Mapping[str, object]]], @@ -1422,9 +1429,7 @@ if MCP_AVAILABLE: getattr(client, "timeout", MCP_CLIENT_TIMEOUT) or MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT, ) - list_tools_result = None # rebind-ok: set inside the timeout scope below - with anyio.move_on_after(listing_deadline): - list_tools_result = await client.list_tools(raise_on_error=True) # rebind-ok: fills the init above + list_tools_result: Final = await _list_tools_within(client, listing_deadline) if list_tools_result is None: verbose_logger.warning( "MCP tools/list preview timed out after %s seconds while paginating upstream tools", diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py b/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py index 27018769909..9cabac2d0fa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py +++ b/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py @@ -8,6 +8,7 @@ import json import os from collections.abc import Mapping +from itertools import islice from typing import ( TYPE_CHECKING, Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__; see ruff-strict.toml @@ -341,19 +342,18 @@ def _json_safe( if depth >= _MAX_DEPTH or id(value) in seen: return None - nested: Final = seen | {id(value)} # mutable-ok: one-shot set literal, unioned into a frozenset immediately + nested: Final = seen | frozenset((id(value),)) if isinstance(value, dict): - out: dict[str, object] = {} # mutable-ok: bounded accumulator local to this call, never escapes as-is - for key, item in list(value.items())[:_MAX_ITEMS]: # mutable-ok: list() only to slice an unordered view - if isinstance(key, str) and key not in strip_keys: - out[key] = _json_safe(item, depth + 1, nested, strip_keys) - return out + return { + key: _json_safe(item, depth + 1, nested, strip_keys) + for key, item in islice(value.items(), _MAX_ITEMS) + if isinstance(key, str) and key not in strip_keys + } if isinstance(value, (list, tuple, set, frozenset)): return [ # mutable-ok: return value is a one-shot list, discarded by the caller after use - _json_safe(item, depth + 1, nested, strip_keys) - for item in list(value)[:_MAX_ITEMS] # mutable-ok: list() only to slice an unordered view + _json_safe(item, depth + 1, nested, strip_keys) for item in islice(value, _MAX_ITEMS) ] dump: Final = getattr(value, "model_dump", None) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index fbf998533f8..0f39b32670a 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -6,7 +6,7 @@ "limit": 26777 }, "LIT003": { - "limit": 266 + "limit": 265 }, "LIT004": { "limit": 40 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16504 + "limit": 16502 }, "LIT011": { - "limit": 5531 + "limit": 5529 }, "LIT012": { "limit": 4495 From 7b919f89a85ea7671fd0c3ac2fb27d31fb74201b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 10:00:47 -0700 Subject: [PATCH 30/69] fix(router): track routed model in fallback attempts --- .../router_utils/fallback_event_handlers.py | 5 +-- .../test_fallback_event_handlers.py | 21 +++++++++++ tests/test_litellm/test_router.py | 35 +++++++++++++++++-- 3 files changed, 56 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 3d37ca216a7..0167721f9fe 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -470,10 +470,11 @@ async def run_async_fallback( attempted: Final = ( carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets() ) - attempted.record(original_model_group) + failed_model_group: Final = get_pre_routing_selection(kwargs) or original_model_group + attempted.record(failed_model_group) for mg in fallback_model_group: - if mg == original_model_group: + if mg == failed_model_group: continue if same_model_group_only and _get_fallback_target_model_group(mg) != original_model_group: verbose_router_logger.info( diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 894b2d9e74f..9e51a60364b 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -614,6 +614,27 @@ async def test_run_async_fallback_forwards_attempted_model_groups_to_nested_call ) +@pytest.mark.asyncio +async def test_run_async_fallback_can_target_the_requested_group_when_a_pre_router_replaced_it(): + """The requested group was never called when a pre-router selected a tier, so a + tier fallback may legitimately target that originally requested group.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["requested-model"], + original_model_group="requested-model", + original_exception=RuntimeError("selected tier failed"), + max_fallbacks=3, + fallback_depth=0, + model="requested-model", + metadata={"pre_routing_selected_model": "selected-tier"}, + ) + + assert router.received_kwargs["model"] == "requested-model" + assert router.received_kwargs["attempted_targets"].keys == frozenset({"selected-tier", "requested-model"}) + + @pytest.mark.asyncio @pytest.mark.parametrize( "entry", diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b413bb18f04..d255135bf64 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8335,20 +8335,26 @@ class TestClaudeCodeSubagentSessionRouterBinding: ) @classmethod - def _router(cls) -> "litellm.Router": + def _router( + cls, + cheap_response: str = "cheap response", + fallbacks: list[dict[str, list[str]]] | None = None, + ) -> "litellm.Router": from litellm.types.router import TaggedPreRoutingStrategy router = litellm.Router( model_list=[ { "model_name": "cheap-model", - "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "cheap response"}, + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": cheap_response}, }, { "model_name": "expensive-model", "litellm_params": {"model": "openai/gpt-4o", "mock_response": "expensive response"}, }, - ] + ], + fallbacks=fallbacks, + num_retries=0, ) router.complexity_routers = { "smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy())] @@ -8477,6 +8483,29 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None + @pytest.mark.asyncio + async def test_subagent_can_fallback_to_its_original_requested_model(self): + router = self._router( + cheap_response="litellm.RateLimitError", + fallbacks=[{"cheap-model": ["expensive-model"]}], + ) + + await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "main turn"}], + **self._request_kwargs(), + ) + subagent_kwargs = self._request_kwargs(agent_id="agent-1234") + + response = await router.acompletion( + model="expensive-model", + messages=[{"role": "user", "content": "subagent turn"}], + **subagent_kwargs, + ) + + assert response.choices[0].message.content == "expensive response" + assert subagent_kwargs["metadata"]["routing_decision"]["routed_model"] == "cheap-model" + @pytest.mark.asyncio async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From dba190842cea106ab4b03880861038ddbe6aae42 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:26:59 -0700 Subject: [PATCH 31/69] fix(anthropic): keep the cache_control normalizer inside the type-discipline budget --- litellm/llms/anthropic/common_utils.py | 38 +++++++++++++++++--------- 1 file changed, 25 insertions(+), 13 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index b5bfb32c0c6..19d3d6d7043 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1404,13 +1404,33 @@ def _with_portable_cache_control_in_message(message: object) -> object: return message return { # mutable-ok: JSON wire format **message, - "content": [_with_portable_cache_control_in_content_block(block) for block in content], + "content": [ # mutable-ok: JSON wire format + _with_portable_cache_control_in_content_block(block) for block in content + ], } -def normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire format +def _with_portable_cache_control_in_messages(messages: object) -> object: + if isinstance(messages, str) or not isinstance(messages, Sequence): + return messages + return [ # mutable-ok: JSON wire format + _with_portable_cache_control_in_message(message) for message in messages + ] + + +def _with_portable_cache_control_in_scoped_value(key: str, value: object) -> object: + match key: + case "system" | "tools": + return _with_portable_cache_control_in_blocks(value) + case "messages": + return _with_portable_cache_control_in_messages(value) + case _: + return value + + +def normalize_cache_control_in_anthropic_payload( payload: Mapping[str, object], -) -> dict[str, object]: +) -> dict[str, object]: # mutable-ok: JSON wire format """ Return a copy of an Anthropic /v1/messages payload with every ``cache_control`` entry reduced to ``{"type": }`` @@ -1427,17 +1447,9 @@ def normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire forma dropped entirely. The caller's payload is never mutated. """ portable: Final = _with_portable_cache_control(payload) - scoped: Final = { # mutable-ok: JSON wire format - key: ( - _with_portable_cache_control_in_blocks(value) - if key in ("system", "tools") - else [_with_portable_cache_control_in_message(message) for message in value] - if key == "messages" and isinstance(value, Sequence) and not isinstance(value, str) - else value - ) - for key, value in portable.items() + return { # mutable-ok: JSON wire format + key: _with_portable_cache_control_in_scoped_value(key, value) for key, value in portable.items() } - return scoped def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: From 53da9bca8e45af86507cb6b5c83736290913ba71 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:29:44 -0700 Subject: [PATCH 32/69] fix(bedrock): drop client_metadata for every converse model --- .../bedrock/chat/converse_transformation.py | 10 +-- litellm/llms/bedrock/common_utils.py | 9 -- .../chat/test_converse_transformation.py | 85 +++---------------- 3 files changed, 15 insertions(+), 89 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 38b9569856a..df52b78f6b5 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -86,7 +86,6 @@ from litellm.utils import ( from ..common_utils import ( BedrockError, BedrockModelInfo, - bedrock_arn_hides_model_family, bedrock_converse_supports_parallel_tool_use_config, bedrock_model_accepts_cache_points, get_anthropic_beta_from_headers, @@ -1335,14 +1334,7 @@ class AmazonConverseConfig(BaseConfig): ) additional_request_params.pop("parallel_tool_calls", None) - - drops_client_metadata: Final = base_model.startswith("anthropic") or bedrock_arn_hides_model_family(model) - if drops_client_metadata and additional_request_params.pop("client_metadata", None) is not None: - litellm.verbose_logger.debug( - "Bedrock Converse: dropping `client_metadata` for model=%s, Anthropic rejects it with " - "'client_metadata: Extra inputs are not permitted'", - model, - ) + additional_request_params.pop("client_metadata", None) # Only set the topK value in for models that support it additional_request_params.update(self._handle_top_k_value(model, inference_params, drop_params)) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index cf18e3e3ec8..66ee5f10679 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -809,15 +809,6 @@ def get_bedrock_base_model(model: str) -> str: return model -def bedrock_arn_hides_model_family(model: str) -> bool: - """ - True for an ARN-addressed model whose base name carries no ``provider.model`` - id, such as an application inference profile or a provisioned throughput ARN. - Callers that gate behavior on the model family cannot resolve one here. - """ - return "arn:" in model.lower() and "." not in get_bedrock_base_model(model) - - def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool: return any( (litellm.model_cost.get(candidate) or {}).get("supports_parallel_tool_use_config") is True diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 7556e1624be..fb165c38ef9 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -979,16 +979,19 @@ def test_config_blocks_do_not_leak_into_inference_config(): assert data["serviceTier"] == {"type": "priority"} -@pytest.mark.parametrize("model", ["anthropic.claude-opus-4-8", "us.anthropic.claude-opus-4-8"]) -def test_client_metadata_stripped_for_anthropic_converse_request(model): - """``client_metadata`` sent by codex must not reach Anthropic as a passthrough model field. - - Converse forwards ``additionalModelRequestFields`` verbatim to the model, and Anthropic - rejects the request with "client_metadata: Extra inputs are not permitted". - """ - config = AmazonConverseConfig() - - data = config._transform_request_helper( +@pytest.mark.parametrize( + "model", + [ + "anthropic.claude-opus-4-8", + "us.anthropic.claude-opus-4-8", + "amazon.nova-pro-v1:0", + "us.meta.llama4-maverick-17b-instruct-v1:0", + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456", + "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.amazon.nova-pro-v1:0", + ], +) +def test_client_metadata_stripped_from_converse_request(model): + data = AmazonConverseConfig()._transform_request_helper( model=model, system_content_blocks=[], optional_params={ @@ -999,71 +1002,11 @@ def test_client_metadata_stripped_for_anthropic_converse_request(model): messages=None, ) - fields = data.get("additionalModelRequestFields", {}) + fields = data["additionalModelRequestFields"] assert "client_metadata" not in fields assert fields["anthropic_beta"] == ["computer-use-2025-01-24"] -def test_client_metadata_kept_for_non_anthropic_converse_request(): - """Only Anthropic is known to reject ``client_metadata``, so other families keep the passthrough.""" - config = AmazonConverseConfig() - - data = config._transform_request_helper( - model="amazon.nova-pro-v1:0", - system_content_blocks=[], - optional_params={ - "maxTokens": 16, - "client_metadata": {"originator": "codex_cli_rs"}, - }, - messages=None, - ) - - assert data["additionalModelRequestFields"]["client_metadata"] == {"originator": "codex_cli_rs"} - - -@pytest.mark.parametrize( - "model", - [ - "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456", - "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/abcdef123456", - ], -) -def test_client_metadata_stripped_for_arn_models_converse(model): - """An ARN hides which family serves the request, and pointing one at Claude is how - teams route codex traffic, so the field has to go there too or the 400 comes back.""" - config = AmazonConverseConfig() - - data = config._transform_request_helper( - model=model, - system_content_blocks=[], - optional_params={ - "maxTokens": 16, - "client_metadata": {"originator": "codex_cli_rs"}, - }, - messages=None, - ) - - assert "client_metadata" not in data.get("additionalModelRequestFields", {}) - - -def test_client_metadata_kept_for_arn_naming_a_non_anthropic_family(): - """An inference profile ARN that still spells out the family is resolvable, so a - non-Anthropic one keeps its passthrough.""" - config = AmazonConverseConfig() - - data = config._transform_request_helper( - model="arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.amazon.nova-pro-v1:0", - system_content_blocks=[], - optional_params={ - "maxTokens": 16, - "client_metadata": {"originator": "codex_cli_rs"}, - }, - messages=None, - ) - - assert data["additionalModelRequestFields"]["client_metadata"] == {"originator": "codex_cli_rs"} - - def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost From dbc126cfc97734e47066ab988c3ac1090e7f830a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:07:30 -0700 Subject: [PATCH 33/69] fix(hosted_vllm): forward truncate_prompt_tokens on rerank requests --- .../llms/hosted_vllm/rerank/transformation.py | 21 ++- litellm/types/rerank.py | 21 ++- .../test_hosted_vllm_rerank_transformation.py | 122 +++++++++++++++++- 3 files changed, 154 insertions(+), 10 deletions(-) diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 0e8fa294f5d..265eb350fc6 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -3,6 +3,7 @@ Transformation logic for Hosted VLLM rerank """ from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final import httpx @@ -13,6 +14,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( + HostedVLLMRerankTruncationParams, OptionalRerankParams, RerankBilledUnits, RerankRequest, @@ -62,7 +64,11 @@ class HostedVLLMRerankConfig(BaseRerankConfig): "top_n", "rank_fields", "return_documents", + "max_tokens_per_doc", "instruction", + "truncate_prompt_tokens", + "truncation_side", + "max_tokens_per_query", ] def map_cohere_rerank_params( @@ -100,7 +106,15 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if instruction is not None: mapped_params["instruction"] = instruction - return dict(mapped_params) + truncation: Final = HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({})) + forwarded: Final[OptionalRerankParams] = { + **mapped_params, + "max_tokens_per_doc": max_tokens_per_doc, + "truncate_prompt_tokens": truncation.truncate_prompt_tokens, + "truncation_side": truncation.truncation_side, + "max_tokens_per_query": truncation.max_tokens_per_query, + } + return dict(forwarded) def validate_environment( self, @@ -138,6 +152,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if "documents" not in optional_rerank_params: raise ValueError("documents is required for Hosted VLLM rerank") + truncation: Final = HostedVLLMRerankTruncationParams.model_validate(optional_rerank_params) rerank_request: Final = RerankRequest( model=model, query=optional_rerank_params["query"], @@ -146,6 +161,10 @@ class HostedVLLMRerankConfig(BaseRerankConfig): rank_fields=optional_rerank_params.get("rank_fields", None), return_documents=optional_rerank_params.get("return_documents", None), instruction=optional_rerank_params.get("instruction", None), + max_tokens_per_doc=truncation.max_tokens_per_doc, + truncate_prompt_tokens=truncation.truncate_prompt_tokens, + truncation_side=truncation.truncation_side, + max_tokens_per_query=truncation.max_tokens_per_query, ) return rerank_request.model_dump(exclude_none=True) diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 903781b2ccd..a76e6cf1187 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -4,8 +4,10 @@ https://docs.cohere.com/reference/rerank """ -from pydantic import BaseModel, PrivateAttr -from typing_extensions import Required, TypedDict +from typing import Literal + +from pydantic import BaseModel, ConfigDict, PrivateAttr +from typing_extensions import ReadOnly, Required, TypedDict class RerankRequest(BaseModel): @@ -21,6 +23,18 @@ class RerankRequest(BaseModel): # (e.g. hosted vLLM / Qwen3-Reranker, DeepInfra). Omitted from the outgoing # request when None, so this is fully backward-compatible. instruction: str | None = None + truncate_prompt_tokens: int | None = None + truncation_side: Literal["left", "right"] | None = None + max_tokens_per_query: int | None = None + + +class HostedVLLMRerankTruncationParams(BaseModel): + model_config = ConfigDict(frozen=True) + + truncate_prompt_tokens: int | None = None + truncation_side: Literal["left", "right"] | None = None + max_tokens_per_query: int | None = None + max_tokens_per_doc: int | None = None class OptionalRerankParams(TypedDict, total=False): @@ -32,6 +46,9 @@ class OptionalRerankParams(TypedDict, total=False): max_chunks_per_doc: int | None max_tokens_per_doc: int | None instruction: str | None + truncate_prompt_tokens: ReadOnly[int | None] + truncation_side: ReadOnly[Literal["left", "right"] | None] + max_tokens_per_query: ReadOnly[int | None] class RerankBilledUnits(TypedDict, total=False): diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index e6e6aa946d5..da27ea1ac58 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -1,8 +1,13 @@ +import json import os import sys +from unittest.mock import MagicMock, patch +import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig from litellm.rerank_api.rerank_utils import get_optional_rerank_params from litellm.types.rerank import ( @@ -87,9 +92,7 @@ class TestHostedVLLMRerankTransform: assert "instruction" not in body def test_map_cohere_rerank_params_raises_on_max_chunks_per_doc(self): - with pytest.raises( - ValueError, match="Hosted VLLM does not support max_chunks_per_doc" - ): + with pytest.raises(ValueError, match="Hosted VLLM does not support max_chunks_per_doc"): self.config.map_cohere_rerank_params( non_default_params=None, model=self.model, @@ -104,12 +107,10 @@ class TestHostedVLLMRerankTransform: url = self.config.get_complete_url(base, self.model) assert url == "https://api.example.com/rerank" # Already ends with /rerank - url2 = self.config.get_complete_url( - "https://api.example.com/rerank", self.model - ) + url2 = self.config.get_complete_url("https://api.example.com/rerank", self.model) assert url2 == "https://api.example.com/rerank" # Raises if api_base is None - with pytest.raises(ValueError, match='api_base must be provided for Hosted VLLM rerank'): + with pytest.raises(ValueError, match="api_base must be provided for Hosted VLLM rerank"): self.config.get_complete_url(None, self.model) def test_transform_response(self): @@ -173,3 +174,110 @@ class TestGetOptionalRerankParamsInstruction: documents=["doc1", "doc2"], ) assert "instruction" not in params + + +class TestHostedVLLMRerankTruncationParams: + def setup_method(self): + self.config = HostedVLLMRerankConfig() + self.model = "hosted-vllm-model" + + def test_map_cohere_rerank_params_forwards_vllm_truncation_params(self): + params = self.config.map_cohere_rerank_params( + non_default_params={ + "truncate_prompt_tokens": 512, + "truncation_side": "left", + "max_tokens_per_query": 64, + "metadata": {"user_api_key": "sk-test"}, + }, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + max_tokens_per_doc=128, + ) + assert params["truncate_prompt_tokens"] == 512 + assert params["truncation_side"] == "left" + assert params["max_tokens_per_query"] == 64 + assert params["max_tokens_per_doc"] == 128 + assert "metadata" not in params + + def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): + params = self.config.map_cohere_rerank_params( + non_default_params={"metadata": {"user_api_key": "sk-test"}}, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + body = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) + truncation_keys = {"truncate_prompt_tokens", "truncation_side", "max_tokens_per_query", "max_tokens_per_doc"} + assert not truncation_keys & body.keys() + assert body == { + "model": self.model, + "query": "test query", + "documents": ["doc1", "doc2"], + "return_documents": True, + } + + def test_map_cohere_rerank_params_rejects_invalid_truncation_side(self): + with pytest.raises(ValueError, match="truncation_side"): + self.config.map_cohere_rerank_params( + non_default_params={"truncation_side": "middle"}, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + + def test_transform_request_forwards_truncation_params(self): + body = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={ + "query": "test query", + "documents": ["doc1", "doc2"], + "truncate_prompt_tokens": 512, + "truncation_side": "left", + "max_tokens_per_query": 64, + "max_tokens_per_doc": 128, + }, + headers={}, + ) + assert body["truncate_prompt_tokens"] == 512 + assert body["truncation_side"] == "left" + assert body["max_tokens_per_query"] == 64 + assert body["max_tokens_per_doc"] == 128 + + def test_transform_request_omits_truncation_params_when_absent(self): + body = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={"query": "test query", "documents": ["doc1", "doc2"]}, + headers={}, + ) + assert "truncate_prompt_tokens" not in body + assert "truncation_side" not in body + assert "max_tokens_per_query" not in body + assert "max_tokens_per_doc" not in body + + def test_rerank_sends_truncate_prompt_tokens_to_vllm(self): + client = HTTPHandler() + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "score-1", + "results": [{"index": 0, "relevance_score": 0.5}], + "usage": {"total_tokens": 512}, + } + with patch.object(client, "post", return_value=mock_response) as mock_post: + litellm.rerank( + model="hosted_vllm/BAAI/bge-reranker-base", + api_base="http://vllm.local:8000", + query="List all the unique case ids", + documents=["a document longer than the reranker context window"], + truncate_prompt_tokens=512, + truncation_side="left", + client=client, + ) + sent_body = json.loads(mock_post.call_args.kwargs["data"]) + assert mock_post.call_args.kwargs["url"] == "http://vllm.local:8000/rerank" + assert sent_body["truncate_prompt_tokens"] == 512 + assert sent_body["truncation_side"] == "left" From dfaf23523453de12278e3d30801075c4da6ee903 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:08:45 -0700 Subject: [PATCH 34/69] fix(bedrock): honor BEDROCK_MANTLE_API_BASE on bedrock/mantle messages and chat URLs --- litellm/llms/bedrock/common_utils.py | 7 ++-- .../test_litellm/llms/bedrock/test_mantle.py | 42 +++++++++++++++++++ 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 66ee5f10679..048d023a1bd 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -758,12 +758,13 @@ def build_mantle_messages_url( """Build the bedrock-mantle Anthropic /messages URL. Honors an explicit endpoint override (``api_base``, then - ``aws_bedrock_runtime_endpoint``) so private VPC / VPCE / GovCloud Mantle - endpoints are reachable; otherwise falls back to the public regional host. + ``aws_bedrock_runtime_endpoint``, then ``BEDROCK_MANTLE_API_BASE``) so + private VPC / VPCE / GovCloud Mantle endpoints are reachable; otherwise + falls back to the public regional host. The mantle messages path is appended unless the override already carries it, so callers can pass either the host or the full messages URL. """ - override: Final = api_base or aws_bedrock_runtime_endpoint + override: Final = api_base or aws_bedrock_runtime_endpoint or get_secret_str("BEDROCK_MANTLE_API_BASE") if override: base: Final = override.rstrip("/") if base.endswith(MANTLE_MESSAGES_PATH): diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index d34517f61f6..d1d1ba447fb 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -128,6 +128,12 @@ def test_mantle_messages_url_construction(): _VPC_ENDPOINT = "https://vpce-0a1b2c3d.bedrock-mantle.us-gov-west-1.vpce.amazonaws.com" +@pytest.fixture(autouse=True) +def no_ambient_mantle_api_base(monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + + + def test_mantle_chat_url_honors_api_base_host(): config = AmazonMantleConfig() url = config.get_complete_url( @@ -193,6 +199,42 @@ def test_mantle_messages_url_honors_aws_bedrock_runtime_endpoint(): assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" +_ENV_ENDPOINT = "https://bedrock-mantle.us-east-1.api.aws.internal.example.com" + + +@pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig]) +def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT) + url = config_cls().get_complete_url( + api_base=None, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-east-1"}, + litellm_params={}, + ) + assert url == f"{_ENV_ENDPOINT}/anthropic/v1/messages" + + +@pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig]) +@pytest.mark.parametrize( + ("api_base", "optional_params"), + [ + (_VPC_ENDPOINT, {"aws_region_name": "us-gov-west-1"}), + (None, {"aws_region_name": "us-gov-west-1", "aws_bedrock_runtime_endpoint": _VPC_ENDPOINT}), + ], +) +def test_mantle_url_explicit_endpoint_beats_bedrock_mantle_api_base_env(monkeypatch, config_cls, api_base, optional_params): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT) + url = config_cls().get_complete_url( + api_base=api_base, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params=optional_params, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + + def test_mantle_transform_request_strips_prefix_and_adds_model(): config = AmazonMantleConfig() request = config.transform_request( From 8d00220acef335234797b00cf3b1f1ee73db2b3a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:19:20 -0700 Subject: [PATCH 35/69] test(hosted_vllm): annotate rerank truncation test locals as Final --- .../test_hosted_vllm_rerank_transformation.py | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index da27ea1ac58..49d58cc28c4 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -1,6 +1,7 @@ import json import os import sys +from typing import Final from unittest.mock import MagicMock, patch import httpx @@ -182,7 +183,7 @@ class TestHostedVLLMRerankTruncationParams: self.model = "hosted-vllm-model" def test_map_cohere_rerank_params_forwards_vllm_truncation_params(self): - params = self.config.map_cohere_rerank_params( + params: Final = self.config.map_cohere_rerank_params( non_default_params={ "truncate_prompt_tokens": 512, "truncation_side": "left", @@ -202,15 +203,20 @@ class TestHostedVLLMRerankTruncationParams: assert "metadata" not in params def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): - params = self.config.map_cohere_rerank_params( + params: Final = self.config.map_cohere_rerank_params( non_default_params={"metadata": {"user_api_key": "sk-test"}}, model=self.model, drop_params=False, query="test query", documents=["doc1", "doc2"], ) - body = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) - truncation_keys = {"truncate_prompt_tokens", "truncation_side", "max_tokens_per_query", "max_tokens_per_doc"} + body: Final = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) + truncation_keys: Final = { + "truncate_prompt_tokens", + "truncation_side", + "max_tokens_per_query", + "max_tokens_per_doc", + } assert not truncation_keys & body.keys() assert body == { "model": self.model, @@ -230,7 +236,7 @@ class TestHostedVLLMRerankTruncationParams: ) def test_transform_request_forwards_truncation_params(self): - body = self.config.transform_rerank_request( + body: Final = self.config.transform_rerank_request( model=self.model, optional_rerank_params={ "query": "test query", @@ -248,7 +254,7 @@ class TestHostedVLLMRerankTruncationParams: assert body["max_tokens_per_doc"] == 128 def test_transform_request_omits_truncation_params_when_absent(self): - body = self.config.transform_rerank_request( + body: Final = self.config.transform_rerank_request( model=self.model, optional_rerank_params={"query": "test query", "documents": ["doc1", "doc2"]}, headers={}, @@ -259,8 +265,8 @@ class TestHostedVLLMRerankTruncationParams: assert "max_tokens_per_doc" not in body def test_rerank_sends_truncate_prompt_tokens_to_vllm(self): - client = HTTPHandler() - mock_response = MagicMock(spec=httpx.Response) + client: Final = HTTPHandler() + mock_response: Final = MagicMock(spec=httpx.Response) mock_response.status_code = 200 mock_response.json.return_value = { "id": "score-1", @@ -277,7 +283,7 @@ class TestHostedVLLMRerankTruncationParams: truncation_side="left", client=client, ) - sent_body = json.loads(mock_post.call_args.kwargs["data"]) + sent_body: Final = json.loads(mock_post.call_args.kwargs["data"]) assert mock_post.call_args.kwargs["url"] == "http://vllm.local:8000/rerank" assert sent_body["truncate_prompt_tokens"] == 512 assert sent_body["truncation_side"] == "left" From ef14bed0296b4a6c48777b8f9df8018b11179c4b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:31:19 -0700 Subject: [PATCH 36/69] fix(hosted_vllm): reject invalid rerank truncation params with a 400 --- .../llms/hosted_vllm/rerank/transformation.py | 11 +++++++- .../test_hosted_vllm_rerank_transformation.py | 26 ++++++++++++------- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 265eb350fc6..764d80c6f82 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -7,8 +7,10 @@ from types import MappingProxyType from typing import Any, Final import httpx +from pydantic import ValidationError from litellm._uuid import uuid +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig @@ -36,6 +38,13 @@ class HostedVLLMRerankError(BaseLLMException): super().__init__(status_code=status_code, message=message, headers=headers) +def validated_truncation_params(non_default_params: Mapping[str, object] | None) -> HostedVLLMRerankTruncationParams: + try: + return HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({})) + except ValidationError as error: + raise UnsupportedParamsError(status_code=400, message=f"hosted_vllm rerank: {error}") from error + + class HostedVLLMRerankConfig(BaseRerankConfig): def __init__(self) -> None: pass @@ -106,7 +115,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if instruction is not None: mapped_params["instruction"] = instruction - truncation: Final = HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({})) + truncation: Final = validated_truncation_params(non_default_params) forwarded: Final[OptionalRerankParams] = { **mapped_params, "max_tokens_per_doc": max_tokens_per_doc, diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index 49d58cc28c4..9a62fcf6f0f 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -202,6 +202,22 @@ class TestHostedVLLMRerankTruncationParams: assert params["max_tokens_per_doc"] == 128 assert "metadata" not in params + @pytest.mark.parametrize( + "bad_params", + [{"truncation_side": "middle"}, {"truncate_prompt_tokens": "lots"}, {"max_tokens_per_query": -1.5}], + ) + def test_map_cohere_rerank_params_rejects_invalid_truncation_params_as_400(self, bad_params: dict[str, object]): + with pytest.raises(litellm.UnsupportedParamsError) as raised: + self.config.map_cohere_rerank_params( + non_default_params=dict(bad_params), + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + assert raised.value.status_code == 400 + assert next(iter(bad_params)) in str(raised.value) + def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): params: Final = self.config.map_cohere_rerank_params( non_default_params={"metadata": {"user_api_key": "sk-test"}}, @@ -225,16 +241,6 @@ class TestHostedVLLMRerankTruncationParams: "return_documents": True, } - def test_map_cohere_rerank_params_rejects_invalid_truncation_side(self): - with pytest.raises(ValueError, match="truncation_side"): - self.config.map_cohere_rerank_params( - non_default_params={"truncation_side": "middle"}, - model=self.model, - drop_params=False, - query="test query", - documents=["doc1", "doc2"], - ) - def test_transform_request_forwards_truncation_params(self): body: Final = self.config.transform_rerank_request( model=self.model, From d7ee215c57af44219d5a19f5042343ce98e9de2a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:32:04 -0700 Subject: [PATCH 37/69] fix(responses): keep namespace tools intact when a guardrail returns them unchanged Any pre_call guardrail on /v1/responses flattened Codex namespace tools into ns__member functions and wrote the flattened list back to the request, so the model called mcp__server__tool with no namespace and Codex rejected the call as unsupported. The handler now keeps the client's original tools, hands the guardrail a deep copy of the flattened ones, and rebuilds data["tools"] by matching the guardrail's output to the originals by type and name. Unchanged tools go back as the original objects, a dropped or edited namespace member changes only that member, and tools the guardrail injects are still appended. Fixes #39183 --- basedpyright-code-budget.json | 12 +- .../guardrail_translation/handler.py | 114 +++------- .../guardrail_translation/tool_merge.py | 177 +++++++++++++++ .../transformation.py | 161 ++++++++------ ruff-strict-budget.json | 2 +- ...test_openai_responses_guardrail_handler.py | 206 +++++++++++++++++- ...t_openai_responses_guardrail_tool_merge.py | 144 ++++++++++++ .../test_litellm_completion_responses.py | 13 ++ type-discipline-budget.json | 10 +- 9 files changed, 669 insertions(+), 170 deletions(-) create mode 100644 litellm/llms/openai/responses/guardrail_translation/tool_merge.py create mode 100644 tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index da788bf1ce3..e7a069de29a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 14076 }, "reportArgumentType": { - "limit": 2216 + "limit": 2215 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4128 + "limit": 4127 }, "reportFunctionMemberAccess": { "limit": 7 @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44364 + "limit": 44362 }, "reportUnknownLambdaType": { "limit": 109 @@ -117,13 +117,13 @@ "limit": 111 }, "reportUnnecessaryComparison": { - "limit": 692 + "limit": 687 }, "reportUnnecessaryContains": { - "limit": 5 + "limit": 4 }, "reportUnnecessaryIsInstance": { - "limit": 826 + "limit": 823 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1530c154e93..5a5970fb867 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -28,6 +28,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has: - text: str """ +import copy import time import uuid from collections.abc import Mapping, Sequence @@ -36,7 +37,6 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall -from openai.types.responses.tool_param import FunctionToolParam from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -49,6 +49,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, stream_item_field, ) +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -62,7 +63,6 @@ from litellm.types.llms.openai import ( ContentPartDonePartOutputText, ErrorEvent, ErrorEventError, - OpenAIMcpServerTool, OutputItemAddedEvent, OutputItemDoneEvent, OutputTextDeltaEvent, @@ -157,23 +157,31 @@ class OpenAIResponsesHandler(BaseTranslation): Handles both string input and list of message objects. """ input_data: Final[str | ResponseInputParam | None] = data.get("input") - tools_to_check: Final[list[ChatCompletionToolParam]] = [] if input_data is None: return data structured_messages: Final = self.get_structured_messages(data) + raw_tools: Final = data.get("tools") + original_tools: Final[tuple[Mapping[str, object], ...]] = ( + tuple(raw_tools) if isinstance(raw_tools, list) else () + ) + flattened_tool_groups: Final = tuple( + form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) + ) + flattened_tools: Final = tuple( + cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list + for group in flattened_tool_groups + for tool in group + ) + tools_to_check: Final[list[ChatCompletionToolParam]] = list( # mutable-ok: guardrail inputs want a list + copy.deepcopy(flattened_tools) + ) # Handle simple string input if isinstance(input_data, str): inputs = GenericGuardrailAPIInputs(texts=[input_data]) - original_tools: list[dict[str, object]] = [] - - # Extract and transform tools if present - if "tools" in data and data["tools"]: - original_tools = list(data["tools"]) - self._extract_and_transform_tools(data["tools"], tools_to_check) - if tools_to_check: - inputs["tools"] = tools_to_check + if tools_to_check: + inputs["tools"] = tools_to_check if structured_messages: inputs["structured_messages"] = structured_messages # Include model information if available @@ -189,7 +197,9 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts = guardrailed_inputs.get("texts", []) data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data - self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools")) + self._apply_guardrailed_tools_to_data( + data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") + ) verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") return data @@ -200,7 +210,6 @@ class OpenAIResponsesHandler(BaseTranslation): texts_to_check: Final[list[str]] = [] images_to_check: Final[list[str]] = [] task_mappings: Final[list[tuple[int, int | None]]] = [] - original_tools_list: Final[list[dict[str, object]]] = list(data.get("tools") or []) # Step 1: Extract all text content, images, and tools for msg_idx, message in enumerate(input_data): @@ -212,10 +221,6 @@ class OpenAIResponsesHandler(BaseTranslation): task_mappings=task_mappings, ) - # Extract and transform tools if present - if "tools" in data and data["tools"]: - self._extract_and_transform_tools(data["tools"], tools_to_check) - # Step 2: Apply guardrail to all texts in batch if texts_to_check: inputs = GenericGuardrailAPIInputs(texts=texts_to_check) @@ -238,9 +243,7 @@ class OpenAIResponsesHandler(BaseTranslation): guardrailed_texts = guardrailed_inputs.get("texts", []) self._apply_guardrailed_tools_to_data( - data, - original_tools_list, - guardrailed_inputs.get("tools"), + data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") ) # Step 3: Map guardrail responses back to original input structure @@ -267,73 +270,18 @@ class OpenAIResponsesHandler(BaseTranslation): names.append(str(tool["server_label"])) return names - def _extract_and_transform_tools( - self, - tools: list[FunctionToolParam | OpenAIMcpServerTool], - tools_to_check: list[ChatCompletionToolParam], - ) -> None: - """ - Extract and transform tools from Responses API format to Chat Completion format. - - Uses the LiteLLM transformation function to convert Responses API tools - to Chat Completion tools that can be passed to guardrails. - """ - if tools is not None and isinstance(tools, list): - # Transform Responses API tools to Chat Completion tools - ( - transformed_tools, - _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools) - tools_to_check.extend(cast(list[ChatCompletionToolParam], transformed_tools)) - - def _remap_tools_to_responses_api_format(self, guardrailed_tools: list[Any]) -> list[dict[str, object]]: - """ - Remap guardrail-returned tools (Chat Completion format) back to - Responses API request tool format. - """ - return LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools( - guardrailed_tools - ) - - def _merge_tools_after_guardrail( - self, - original_tools: list[dict[str, object]], - remapped: list[dict[str, object]], - ) -> list[dict[str, object]]: - """ - Merge remapped guardrailed tools with original tools that were not sent - to the guardrail (e.g. web_search, web_search_preview), preserving order. - Tools a guardrail appended (``remapped`` longer than ``original_tools``) - have no original slot and are kept so an injected tool is not dropped. - """ - if not original_tools: - return remapped - result: Final[list[dict[str, object]]] = [] - j = 0 - for tool in original_tools: - if isinstance(tool, dict) and tool.get("type") in ( - "web_search", - "web_search_preview", - ): - result.append(tool) - else: - if j < len(remapped): - result.append(remapped[j]) - j += 1 - # Keep guardrail-appended tools that matched no original slot above. - result.extend(remapped[j:]) - return result - def _apply_guardrailed_tools_to_data( self, data: dict, - original_tools: list[dict[str, object]], - guardrailed_tools: list[ChatCompletionToolParam] | None, + original_tools: Sequence[Mapping[str, object]], + flattened_tool_groups: Sequence[Sequence[Mapping[str, object]]], + guardrailed_tools: Sequence[ChatCompletionToolParam] | None, ) -> None: - """Remap guardrailed tools to Responses API format and merge with original, then set data['tools'].""" - if guardrailed_tools is not None: - remapped: Final = self._remap_tools_to_responses_api_format(guardrailed_tools) - data["tools"] = self._merge_tools_after_guardrail(original_tools, remapped) + if guardrailed_tools is None: + return + data["tools"] = list( # mutable-ok: downstream wants a list # rebind-ok: in-place request rewrite + merge_guardrailed_tools(original_tools, flattened_tool_groups, guardrailed_tools) + ) def _extract_input_text_and_images( self, diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py new file mode 100644 index 00000000000..3ae951d3f61 --- /dev/null +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -0,0 +1,177 @@ +from collections.abc import Iterable, Mapping, Sequence +from itertools import accumulate, chain +from types import MappingProxyType +from typing import Final, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger +from litellm.responses.litellm_completion_transformation.transformation import ( + NAMESPACE_DESCRIPTION_SEPARATOR, + LiteLLMCompletionResponsesConfig, +) + +Tool: TypeAlias = Mapping[str, object] +IndexedKey: TypeAlias = tuple[str, int] + +_TOOL_ADAPTER: Final = TypeAdapter(dict[str, object]) +_CHAT_TOOL_TOP_LEVEL_KEYS: Final = frozenset({"type", "function"}) + + +def _as_tool(value: object) -> Tool | None: + try: + return _TOOL_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _validated_tools(values: Iterable[object]) -> tuple[Tool, ...]: + validated: Final = tuple(map(_as_tool, values)) + dropped: Final = sum(tool is None for tool in validated) + if dropped: + verbose_logger.warning("Dropping %d guardrail-returned tools that are not objects", dropped) + return tuple(tool for tool in validated if tool is not None) + + +def _is_function(tool: Tool) -> bool: + return tool.get("type") == "function" + + +def _chat_tool_key(tool: Tool) -> str: + tool_type: Final = str(tool.get("type") or "") + function: Final = _as_tool(tool.get("function")) + if function is not None: + return f"{tool_type}:{function.get('name') or ''}" + return f"{tool_type}:{tool.get('server_label') or tool.get('name') or ''}" + + +def _indexed_keys(tools: Sequence[Tool]) -> tuple[IndexedKey, ...]: + keys: Final = tuple(_chat_tool_key(tool) for tool in tools) + return tuple((key, keys[:position].count(key)) for position, key in enumerate(keys)) + + +def _namespace_members(namespace: Tool) -> tuple[Tool, ...]: + members: Final = namespace.get("tools") + if not isinstance(members, Sequence) or isinstance(members, (str, bytes)): + return () + return tuple(member for member in map(_as_tool, members) if member is not None) + + +def _function_fields(tool: Tool) -> Tool: + function: Final = _as_tool(tool.get("function")) + return function if function is not None else MappingProxyType({}) + + +def _without_namespace_prefix(key: str, value: object, prefix: str) -> object: + if key != "description" or not isinstance(value, str) or not value.startswith(prefix): + return value + return value[len(prefix) :] + + +def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool: + flattened_function: Final = _function_fields(flattened) + prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else "" + changed_function: Final = MappingProxyType( + { + key: _without_namespace_prefix(key, value, prefix) + for key, value in _function_fields(guardrailed).items() + if flattened_function.get(key) != value + } + ) + changed_extras: Final = MappingProxyType( + { + key: value + for key, value in guardrailed.items() + if key not in _CHAT_TOOL_TOP_LEVEL_KEYS and flattened.get(key) != value + } + ) + return {**member, **changed_extras, **changed_function} # mutable-ok: json.dumps rejects MappingProxyType + + +def _rebuilt_function_members( + function_members: Sequence[Tool], + flattened_group: Sequence[Tool], + group_keys: Sequence[IndexedKey], + guardrailed_by_key: Mapping[IndexedKey, Tool], + namespace_description: str, +) -> tuple[Tool | None, ...]: + return tuple( + None + if key not in guardrailed_by_key + else member + if guardrailed_by_key[key] == flattened + else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description) + for member, flattened, key in zip(function_members, flattened_group, group_keys) + ) + + +def _rebuilt_namespace( + original: Tool, + members: Sequence[Tool], + flattened_group: Sequence[Tool], + group_keys: Sequence[IndexedKey], + guardrailed_by_key: Mapping[IndexedKey, Tool], +) -> tuple[Tool, ...]: + namespace_description: Final = str(original.get("description") or "") + rebuilt_functions: Final = iter( + _rebuilt_function_members( + tuple(member for member in members if _is_function(member)), + flattened_group, + group_keys, + guardrailed_by_key, + namespace_description, + ) + ) + rebuilt_members: Final = tuple( + rebuilt + for rebuilt in (next(rebuilt_functions) if _is_function(member) else member for member in members) + if rebuilt is not None + ) + if not rebuilt_members: + return () + return ({**original, "tools": list(rebuilt_members)},) # mutable-ok: json.dumps needs a plain dict and list + + +def _merged_original( + original: Tool, + flattened_group: Sequence[Tool], + group_keys: Sequence[IndexedKey], + guardrailed_by_key: Mapping[IndexedKey, Tool], +) -> tuple[Tool, ...]: + if not group_keys: + return (original,) + guardrailed_group: Final = tuple(guardrailed_by_key[key] for key in group_keys if key in guardrailed_by_key) + if guardrailed_group == tuple(flattened_group): + return (original,) + if not guardrailed_group: + return () + members: Final = _namespace_members(original) if original.get("type") == "namespace" else () + if members and sum(map(_is_function, members)) == len(flattened_group): + return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key) + return tuple( + LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools(guardrailed_group) + ) + + +def merge_guardrailed_tools( + original_tools: Sequence[Tool], + flattened_groups: Sequence[Sequence[Tool]], + guardrailed_tools: Iterable[object], +) -> tuple[Tool, ...]: + guardrailed: Final = _validated_tools(guardrailed_tools) + flattened_keys: Final = _indexed_keys(tuple(chain.from_iterable(flattened_groups))) + guardrailed_keys: Final = _indexed_keys(guardrailed) + guardrailed_by_key: Final = MappingProxyType(dict(zip(guardrailed_keys, guardrailed))) + group_ends: Final = tuple(accumulate(len(group) for group in flattened_groups)) + group_key_slices: Final = tuple( + flattened_keys[end - len(group) : end] for group, end in zip(flattened_groups, group_ends) + ) + merged_originals: Final = chain.from_iterable( + _merged_original(original, group, group_keys, guardrailed_by_key) + for original, group, group_keys in zip(original_tools, flattened_groups, group_key_slices) + ) + owned_keys: Final = frozenset(flattened_keys) + appended: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools( + tuple(tool for key, tool in zip(guardrailed_keys, guardrailed) if key not in owned_keys) + ) + return tuple(chain(merged_originals, appended)) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 5f3e88bb12f..9f91d7527cb 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -6,6 +6,7 @@ import json import re import uuid from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -102,6 +103,15 @@ from .custom_tools import ( NamespaceNameMap: TypeAlias = Mapping[str, tuple[str, str]] NamespaceTool: TypeAlias = Mapping[str, object] ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None +ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool +NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" + + +@dataclass(frozen=True, slots=True) +class ResponsesToolChatForm: + chat_tools: tuple[ChatToolParam, ...] + web_search_options: OpenAIWebSearchOptions | None + if TYPE_CHECKING: from openai.types.responses.response_apply_patch_tool_call import ( @@ -1771,7 +1781,7 @@ class LiteLLMCompletionResponsesConfig: tool_name: Final = str(namespace_tool.get("name") or "") raw_description: Final = str(namespace_tool.get("description") or "") description: Final = ( - f"{namespace_description}\n\n{raw_description}" + f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}" if nested and namespace_description and raw_description else namespace_description if nested and namespace_description @@ -1837,9 +1847,82 @@ class LiteLLMCompletionResponsesConfig: + ", ".join(sorted(conflicting_tool_names)) ) + @staticmethod + def _responses_tool_to_chat_form(tool: Mapping[str, object]) -> ResponsesToolChatForm: + tool_type: Final = tool.get("type") + if tool_type == "mcp": + return ResponsesToolChatForm(chat_tools=(cast(OpenAIMcpServerTool, tool),), web_search_options=None) + if tool_type == "web_search_preview" or tool_type == "web_search": + _search_context_size: Final[Literal["low", "medium", "high"]] = cast( + Literal["low", "medium", "high"], tool.get("search_context_size") + ) + _user_location: Final[OpenAIWebSearchUserLocation | None] = cast( + OpenAIWebSearchUserLocation | None, + tool.get("user_location") or None, + ) + return ResponsesToolChatForm( + chat_tools=(), + web_search_options=OpenAIWebSearchOptions( + search_context_size=_search_context_size, + user_location=_user_location, + ), + ) + if tool_type == "function": + typed_tool: Final = cast(FunctionToolParam, tool) + raw_parameters: Final = typed_tool.get("parameters", {}) or {} + parameters: Final = ( + {**raw_parameters} # mutable-ok: json.dumps rejects MappingProxyType + if "type" in raw_parameters + else {**raw_parameters, "type": "object"} # mutable-ok: json.dumps rejects MappingProxyType + ) + chat_completion_tool: Final[dict[str, object]] = { + "type": "function", + "function": { + "name": typed_tool.get("name") or "", + "description": typed_tool.get("description") or "", + "parameters": parameters, + "strict": typed_tool.get("strict", False) or False, + }, + } + if tool.get("cache_control"): + chat_completion_tool["cache_control"] = tool.get("cache_control") + if tool.get("defer_loading"): + chat_completion_tool["defer_loading"] = tool.get("defer_loading") + if tool.get("allowed_callers"): + chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") + if tool.get("input_examples"): + chat_completion_tool["input_examples"] = tool.get("input_examples") + return ResponsesToolChatForm( + chat_tools=(cast(ChatCompletionToolParam, chat_completion_tool),), web_search_options=None + ) + if tool_type == "namespace": + return ResponsesToolChatForm( + chat_tools=LiteLLMCompletionResponsesConfig._namespace_chat_tools(tool), web_search_options=None + ) + if tool_type == "custom": + converted: Final = convert_custom_tool_to_function_tool(tool) + return ResponsesToolChatForm(chat_tools=() if converted is None else (converted,), web_search_options=None) + if tool_type in ("computer_use", "image_generation", "shell"): + # Drop unsupported Responses-API-only tool types that have no + # Chat Completions equivalent. Passing them through verbatim + # causes providers to reject the request with "'function' is a + # required property". + verbose_logger.warning( + "Dropping Responses API tool of type '%s': it has no Chat Completions " + "equivalent and the target provider would reject the request.", + tool_type, + ) + return ResponsesToolChatForm(chat_tools=(), web_search_options=None) + return ResponsesToolChatForm(chat_tools=(cast(ChatToolParam, tool),), web_search_options=None) + + @staticmethod + def responses_tools_to_chat_forms(tools: ResponseTools) -> tuple[ResponsesToolChatForm, ...]: + LiteLLMCompletionResponsesConfig._validate_namespace_name_collisions(tools) + return tuple(LiteLLMCompletionResponsesConfig._responses_tool_to_chat_form(tool) for tool in tools or ()) + @staticmethod def transform_responses_api_tools_to_chat_completion_tools( - tools: list[FunctionToolParam | OpenAIMcpServerTool] | None, + tools: ResponseTools, ) -> tuple[ list[ChatCompletionToolParam | OpenAIMcpServerTool], OpenAIWebSearchOptions | None, @@ -1849,73 +1932,16 @@ class LiteLLMCompletionResponsesConfig: """ if tools is None: return [], None - LiteLLMCompletionResponsesConfig._validate_namespace_name_collisions(tools) - chat_completion_tools: Final[list[ChatCompletionToolParam | OpenAIMcpServerTool]] = [] - web_search_options: OpenAIWebSearchOptions | None = None - for tool in tools: - if tool.get("type") == "mcp": - chat_completion_tools.append(cast(OpenAIMcpServerTool, tool)) - elif tool.get("type") == "web_search_preview" or tool.get("type") == "web_search": - _search_context_size: Literal["low", "medium", "high"] = cast( - Literal["low", "medium", "high"], tool.get("search_context_size") - ) - _user_location: OpenAIWebSearchUserLocation | None = cast( - OpenAIWebSearchUserLocation | None, - tool.get("user_location") or None, - ) - web_search_options = OpenAIWebSearchOptions( - search_context_size=_search_context_size, - user_location=_user_location, - ) - elif tool.get("type") == "function": - typed_tool = cast(FunctionToolParam, tool) - # Ensure parameters has "type": "object" as required by providers like Anthropic - parameters = dict(typed_tool.get("parameters", {}) or {}) - if not parameters or "type" not in parameters: - parameters["type"] = "object" - chat_completion_tool: dict[str, object] = { - "type": "function", - "function": { - "name": typed_tool.get("name") or "", - "description": typed_tool.get("description") or "", - "parameters": parameters, - "strict": typed_tool.get("strict", False) or False, - }, - } - if tool.get("cache_control"): - chat_completion_tool["cache_control"] = tool.get("cache_control") - if tool.get("defer_loading"): - chat_completion_tool["defer_loading"] = tool.get("defer_loading") - if tool.get("allowed_callers"): - chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") - if tool.get("input_examples"): - chat_completion_tool["input_examples"] = tool.get("input_examples") - chat_completion_tools.append(cast(ChatCompletionToolParam, chat_completion_tool)) - elif tool.get("type") == "namespace": - chat_completion_tools.extend(LiteLLMCompletionResponsesConfig._namespace_chat_tools(tool)) - elif tool.get("type") == "custom": - converted = convert_custom_tool_to_function_tool(tool) - if converted is not None: - chat_completion_tools.append(converted) - else: - _tool_type = tool.get("type") - if _tool_type in ("computer_use", "image_generation", "shell"): - # Drop unsupported Responses-API-only tool types that have no - # Chat Completions equivalent. Passing them through verbatim - # causes providers to reject the request with "'function' is a - # required property". - verbose_logger.warning( - "Dropping Responses API tool of type '%s': it has no Chat Completions " - "equivalent and the target provider would reject the request.", - _tool_type, - ) - continue - chat_completion_tools.append(cast(ChatCompletionToolParam | OpenAIMcpServerTool, tool)) - return chat_completion_tools, web_search_options + forms: Final = LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(tools) + web_search_options: Final = next( + (form.web_search_options for form in reversed(forms) if form.web_search_options is not None), + None, + ) + return [chat_tool for form in forms for chat_tool in form.chat_tools], web_search_options @staticmethod def transform_chat_completion_tool_params_to_responses_api_tools( - chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None, + chat_completion_tools: Sequence[Mapping[str, object]] | None, ) -> list[dict[str, object]]: """ Transform Chat Completion tool params (e.g. from guardrail output) back to @@ -1926,9 +1952,6 @@ class LiteLLMCompletionResponsesConfig: return [] result: Final[list[dict[str, object]]] = [] for tool in chat_completion_tools: - if not isinstance(tool, dict): - result.append(tool) - continue if tool.get("type") == "function": fn = cast(_ToolFunctionDefinition, tool.get("function") or {}) parameters = dict(fn.get("parameters", {}) or {}) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 9b1cc977a64..935fed18a79 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -177,7 +177,7 @@ "limit": 8 }, "RUF019": { - "limit": 31 + "limit": 29 }, "RUF046": { "limit": 4 diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 315b6948bd8..d071ef78c2d 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -5,6 +5,8 @@ Tests the handler's ability to process input/output for the Responses API with guardrail transformations. """ +import copy +from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple from unittest.mock import AsyncMock, MagicMock @@ -19,6 +21,10 @@ from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.responses.main import GenericResponseOutputItem, OutputText from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs @@ -1287,14 +1293,14 @@ class TestOpenAIResponsesHandlerToolInjection: """A tool a guardrail injects must survive the write-back to Responses format.""" def test_merge_keeps_guardrail_appended_tool(self): - """_merge_tools_after_guardrail must not drop the extra appended tool.""" - handler = OpenAIResponsesHandler() + """merge_guardrailed_tools must not drop the extra appended tool.""" original = [{"type": "function", "name": "a"}] - remapped = [ - {"type": "function", "name": "a"}, - {"type": "function", "name": "b"}, + groups = [form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original)] + guardrailed = [ + *groups[0], + {"type": "function", "function": {"name": "b", "description": "", "parameters": {"type": "object"}}}, ] - merged = handler._merge_tools_after_guardrail(original, remapped) + merged = merge_guardrailed_tools(original, groups, guardrailed) assert [t["name"] for t in merged] == ["a", "b"] @pytest.mark.asyncio @@ -1323,6 +1329,194 @@ class TestOpenAIResponsesHandlerToolInjection: assert "injected_tool" in names +class ToolEditingGuardrail(CustomGuardrail): + """Guardrail that rewrites the flattened chat tools it was handed through ``edit``""" + + def __init__(self, edit: Callable[[list[dict]], list[dict]], **kwargs): + super().__init__(**kwargs) + self.edit = edit + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Any | None = None, + ) -> GenericGuardrailAPIInputs: + inputs["tools"] = self.edit(list(inputs.get("tools") or [])) + return inputs + + +def _codex_request(input_value): + """A Responses API request shaped like what the Codex CLI sends when an MCP server is configured""" + return { + "model": "gpt-5.3-codex", + "input": input_value, + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Weather lookup", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + "strict": False, + }, + { + "type": "namespace", + "name": "mcp__confluence", + "description": "Confluence tools", + "tools": [ + { + "type": "function", + "name": "confluence_get_page", + "description": "Get a page", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + "strict": False, + }, + { + "type": "function", + "name": "confluence_search", + "description": "Search pages", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + "strict": False, + }, + ], + }, + { + "type": "custom", + "name": "apply_patch", + "description": "Apply a patch", + "format": {"type": "grammar", "syntax": "lark", "definition": 'start: "x"'}, + }, + {"type": "web_search"}, + ], + } + + +def _tool_named(tools, name): + return next(tool for tool in tools if tool.get("name") == name) + + +class TestOpenAIResponsesHandlerNamespaceTools: + """Codex sends MCP tools as ``namespace`` tools; a guardrail must never flatten them (GH #39183)""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "input_value", + ["hi", [{"role": "user", "content": "hi", "type": "message"}]], + ids=["string_input", "list_input"], + ) + async def test_pass_through_guardrail_leaves_tools_untouched(self, input_value): + data = _codex_request(input_value) + expected_tools = copy.deepcopy(data["tools"]) + + result = await OpenAIResponsesHandler().process_input_messages( + data, MockPassThroughGuardrail(guardrail_name="test") + ) + + assert result["tools"] == expected_tools + + @pytest.mark.asyncio + async def test_appending_guardrail_keeps_namespace_and_adds_tool(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolAppendingGuardrail(guardrail_name="test") + ) + + assert result["tools"][:-1] == expected_tools + assert result["tools"][-1]["type"] == "function" + assert result["tools"][-1]["name"] == "injected_tool" + + @pytest.mark.asyncio + async def test_dropping_one_member_prunes_only_that_member(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + guardrail = ToolEditingGuardrail( + edit=lambda tools: [t for t in tools if t["function"]["name"] != "mcp__confluence__confluence_search"], + guardrail_name="test", + ) + + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + namespace = _tool_named(result["tools"], "mcp__confluence") + assert [member["name"] for member in namespace["tools"]] == ["confluence_get_page"] + assert namespace["tools"][0] == expected_tools[1]["tools"][0] + assert [t for t in result["tools"] if t is not namespace] == [expected_tools[0], *expected_tools[2:]] + + @pytest.mark.asyncio + async def test_editing_a_member_lands_on_that_member_without_the_namespace_prefix(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + def redact_search(tools): + for tool in tools: + if tool["function"]["name"] == "mcp__confluence__confluence_search": + tool["function"]["description"] = "Confluence tools\n\nREDACTED" + return tools + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolEditingGuardrail(edit=redact_search, guardrail_name="test") + ) + + namespace = _tool_named(result["tools"], "mcp__confluence") + assert namespace["tools"][0] == expected_tools[1]["tools"][0] + assert namespace["tools"][1] == {**expected_tools[1]["tools"][1], "description": "REDACTED"} + assert {k: v for k, v in namespace.items() if k != "tools"} == { + k: v for k, v in expected_tools[1].items() if k != "tools" + } + + @pytest.mark.asyncio + async def test_dropping_every_member_drops_the_namespace(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + guardrail = ToolEditingGuardrail( + edit=lambda tools: [t for t in tools if not t["function"]["name"].startswith("mcp__confluence__")], + guardrail_name="test", + ) + + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert result["tools"] == [expected_tools[0], *expected_tools[2:]] + + @pytest.mark.asyncio + async def test_edited_top_level_function_is_rewritten_in_place(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + def rename_weather(tools): + for tool in tools: + if tool["function"]["name"] == "get_weather": + tool["function"]["description"] = "Weather lookup (guarded)" + return tools + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolEditingGuardrail(edit=rename_weather, guardrail_name="test") + ) + + assert result["tools"][0] == {**expected_tools[0], "description": "Weather lookup (guarded)"} + assert result["tools"][1:] == expected_tools[1:] + + +class TestOpenAIResponsesHandlerMalformedTools: + @pytest.mark.asyncio + async def test_request_tools_that_are_not_a_list_never_reach_the_guardrail(self): + handler = OpenAIResponsesHandler() + seen: list[list[dict]] = [] + + def record(tools): + seen.append(tools) + return tools + + guardrail = ToolEditingGuardrail(edit=record, guardrail_name="test") + data = {"input": "hi", "tools": {"type": "function", "name": "get_weather"}} + + result = await handler.process_input_messages(data, guardrail) + + assert seen == [[]] + assert result["input"] == "hi" + + class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE events""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py new file mode 100644 index 00000000000..b80dd0b36aa --- /dev/null +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -0,0 +1,144 @@ +""" +Unit tests for merge_guardrailed_tools, which writes guardrail-returned chat tools back onto the +Responses API request tools they were flattened from +""" + +import copy + +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) + + +def _groups(tools): + return [form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(tools)] + + +def _flat(groups): + return [chat_tool for group in groups for chat_tool in group] + + +def _function(name, description=""): + return {"type": "function", "name": name, "description": description, "parameters": {"type": "object"}} + + +def test_unchanged_tools_come_back_as_the_original_objects(): + original = [ + _function("a"), + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("x"), _function("y")]}, + {"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}, + {"type": "web_search"}, + ] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, _flat(groups)) + + assert list(merged) == original + assert all(merged_tool is original_tool for merged_tool, original_tool in zip(merged, original)) + + +def test_guardrail_reordering_unchanged_tools_keeps_request_order(): + original = [_function("a"), {"type": "namespace", "name": "ns", "tools": [_function("x")]}, {"type": "web_search"}] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, list(reversed(_flat(groups)))) + + assert list(merged) == original + + +def test_duplicate_function_names_are_matched_by_ordinal(): + original = [_function("dup", "first"), _function("dup", "second")] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, _flat(groups)[:1]) + + assert list(merged) == [original[0]] + + +def test_edited_mcp_tool_is_rewritten(): + original = [{"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}] + groups = _groups(original) + edited = [{**groups[0][0], "allowed_tools": ["read_wiki_structure"]}] + + merged = merge_guardrailed_tools(original, groups, edited) + + assert list(merged) == edited + + +def test_injected_tool_lands_after_the_request_tools_when_request_had_none(): + injected = {"type": "function", "function": {"name": "b", "description": "d", "parameters": {"type": "object"}}} + + merged = merge_guardrailed_tools([], [], [injected]) + + assert list(merged) == [ + {"type": "function", "name": "b", "description": "d", "parameters": {"type": "object"}, "strict": False} + ] + + +def test_empty_guardrail_output_keeps_only_tools_never_sent_to_the_guardrail(): + original = [_function("a"), {"type": "web_search"}, {"type": "namespace", "name": "ns", "tools": [_function("x")]}] + + merged = merge_guardrailed_tools(original, _groups(original), []) + + assert list(merged) == [{"type": "web_search"}] + + +def test_member_edit_strips_only_the_namespace_description_prefix(): + original = [{"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("x", "X doc")]}] + groups = _groups(original) + assert groups[0][0]["function"]["description"] == "NS\n\nX doc" + edited = [{**groups[0][0], "function": {**groups[0][0]["function"], "description": "NS\n\nX doc (guarded)"}}] + + merged = merge_guardrailed_tools(original, groups, edited) + + assert list(merged) == [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("x", "X doc (guarded)")]} + ] + + +def test_namespace_keeps_a_non_function_member_when_a_function_member_is_edited(): + custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}} + original = [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read", "Read"), custom_member]} + ] + groups = _groups(original) + edited = copy.deepcopy(_flat(groups)) + edited[0]["function"]["description"] = "NS\n\nEDITED" + + merged = merge_guardrailed_tools(original, groups, edited) + + assert len(merged) == 1 + assert [member["name"] for member in merged[0]["tools"]] == ["read", "grep"] + assert merged[0]["tools"][0]["description"] == "EDITED" + assert merged[0]["tools"][1] == custom_member + + +def test_member_extras_edited_by_the_guardrail_land_on_that_member(): + original = [{"type": "namespace", "name": "ns", "tools": [_function("read")]}] + groups = _groups(original) + edited = copy.deepcopy(_flat(groups)) + edited[0]["cache_control"] = {"type": "ephemeral"} + + merged = merge_guardrailed_tools(original, groups, edited) + + assert merged[0]["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert merged[0]["tools"][0]["name"] == "read" + + +def test_guardrail_output_is_read_once(): + original = [_function("a"), {"type": "namespace", "name": "ns", "tools": [_function("x")]}] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, (chat_tool for chat_tool in _flat(groups))) + + assert list(merged) == original + + +def test_non_object_guardrail_items_are_dropped(): + original = [_function("a")] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, [*_flat(groups), "junk", None]) + + assert list(merged) == original diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b2b8eb5da80..2068f10ea2d 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1928,6 +1928,19 @@ class TestToolTransformation: assert result_tool["function"]["parameters"]["type"] == "object" assert "properties" in result_tool["function"]["parameters"] + def test_transform_function_tools_parameters_keep_client_key_order(self): + tools = [ + {"type": "function", "name": "a", "parameters": {"properties": {"arg": {"type": "string"}}, "required": ["arg"]}}, + {"type": "function", "name": "b", "parameters": {"type": "object", "properties": {}}}, + ] + + result_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=tools + ) + + assert list(result_tools[0]["function"]["parameters"]) == ["properties", "required", "type"] + assert list(result_tools[1]["function"]["parameters"]) == ["type", "properties"] + def test_transform_function_tools_empty_parameters(self): """Test that empty parameters get 'type': 'object' added""" function_tool = { diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 52cb9628252..970a44cd4fa 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22364 + "limit": 22340 }, "LIT002": { - "limit": 26777 + "limit": 26770 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1039 + "limit": 1038 }, "LIT007": { "limit": 0 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16507 + "limit": 16503 }, "LIT011": { - "limit": 5535 + "limit": 5534 }, "LIT012": { "limit": 4495 From 49c69c46b25af2dd962b322bd6c8e6c5668546c6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:51:27 -0700 Subject: [PATCH 38/69] fix(bedrock): drop the OpenAI base suffix from BEDROCK_MANTLE_API_BASE before the mantle messages path --- litellm/llms/bedrock/common_utils.py | 15 +++++++++++++-- tests/test_litellm/llms/bedrock/test_mantle.py | 12 +++++++++--- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 048d023a1bd..1e5329c90dd 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -748,6 +748,15 @@ def strip_bedrock_throughput_suffix(model: str) -> str: MANTLE_MESSAGES_PATH: Final = "/anthropic/v1/messages" +_MANTLE_OPENAI_BASE_SUFFIXES: Final = ("/openai/v1", "/v1") + + +def _mantle_api_base_from_env() -> str | None: + env_base: Final = get_secret_str("BEDROCK_MANTLE_API_BASE") + if env_base is None: + return None + base: Final = env_base.rstrip("/") + return next((base[: -len(suffix)] for suffix in _MANTLE_OPENAI_BASE_SUFFIXES if base.endswith(suffix)), base) def build_mantle_messages_url( @@ -762,9 +771,11 @@ def build_mantle_messages_url( private VPC / VPCE / GovCloud Mantle endpoints are reachable; otherwise falls back to the public regional host. The mantle messages path is appended unless the override already carries it, - so callers can pass either the host or the full messages URL. + so callers can pass either the host or the full messages URL. The env var is + shared with the OpenAI-surface ``bedrock_mantle/*`` routes, which need it to + carry their ``/v1`` or ``/openai/v1`` base, so that suffix is dropped first. """ - override: Final = api_base or aws_bedrock_runtime_endpoint or get_secret_str("BEDROCK_MANTLE_API_BASE") + override: Final = api_base or aws_bedrock_runtime_endpoint or _mantle_api_base_from_env() if override: base: Final = override.rstrip("/") if base.endswith(MANTLE_MESSAGES_PATH): diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index d1d1ba447fb..09be2118001 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -203,8 +203,12 @@ _ENV_ENDPOINT = "https://bedrock-mantle.us-east-1.api.aws.internal.example.com" @pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig]) -def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls): - monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT) +@pytest.mark.parametrize( + "env_value", + [_ENV_ENDPOINT, f"{_ENV_ENDPOINT}/", f"{_ENV_ENDPOINT}/v1", f"{_ENV_ENDPOINT}/openai/v1"], +) +def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls, env_value): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", env_value) url = config_cls().get_complete_url( api_base=None, api_key=None, @@ -223,7 +227,9 @@ def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls): (None, {"aws_region_name": "us-gov-west-1", "aws_bedrock_runtime_endpoint": _VPC_ENDPOINT}), ], ) -def test_mantle_url_explicit_endpoint_beats_bedrock_mantle_api_base_env(monkeypatch, config_cls, api_base, optional_params): +def test_mantle_url_explicit_endpoint_beats_bedrock_mantle_api_base_env( + monkeypatch, config_cls, api_base, optional_params +): monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT) url = config_cls().get_complete_url( api_base=api_base, From 9bd870d47a700b183e0ee0d9bf7647aa7c739561 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:36:28 -0700 Subject: [PATCH 39/69] fix(databricks): upgrade legacy thinking to adaptive on adaptive-only Claude models --- .../llms/databricks/chat/transformation.py | 4 ++++ .../test_databricks_chat_transformation.py | 21 +++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index c587146005f..65622d62af2 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -330,6 +330,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): ) -> dict: is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) mapped_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params) + if "claude" in model: + AnthropicConfig.translate_legacy_thinking_for_adaptive_model( + model=model, optional_params=mapped_params, custom_llm_provider="databricks" + ) if "tools" in mapped_params: mapped_params["tools"] = self._map_openai_to_dbrx_tool(model=model, tools=mapped_params["tools"]) if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens: diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 41fb2589655..71661cc532b 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -422,6 +422,27 @@ def test_databricks_config_probes_capabilities_under_databricks_namespace(): assert DatabricksConfig().custom_llm_provider == "databricks" +@pytest.mark.parametrize( + "model, expected_thinking, expected_output_config", + [ + ("databricks-claude-opus-4-8", {"type": "adaptive"}, {"effort": "high"}), + ("databricks-claude-opus-4-6", {"type": "enabled", "budget_tokens": 4096}, None), + ], + ids=["adaptive_only_upgrades_to_adaptive", "legacy_capable_forwards_verbatim"], +) +def test_map_openai_params_upgrades_legacy_thinking_on_adaptive_only_claude( + model, expected_thinking, expected_output_config +): + mapped = DatabricksConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 4096}}, + optional_params={}, + model=model, + drop_params=False, + ) + assert mapped["thinking"] == expected_thinking + assert mapped.get("output_config") == expected_output_config + + def _streaming_chunk(usage=None, choices=None): base = { "id": "chatcmpl-test", From dc12e4c2b4ad31b1eda1544ecd2e424aabc78151 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:57:01 -0700 Subject: [PATCH 40/69] fix(responses): match guardrail tools by ordinal in one pass Sort the chat-tool keys once and number duplicates with groupby instead of rescanning every preceding key per position, so the guardrail merge stays O(n log n) on client-supplied tool lists. Drop the comment that restated the unsupported-tool warning in the Responses-to-chat transformation. --- .../guardrail_translation/tool_merge.py | 8 ++++++-- .../transformation.py | 4 ---- ...st_openai_responses_guardrail_tool_merge.py | 18 ++++++++++++++++++ 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index 3ae951d3f61..9fbf9f31a0f 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -1,5 +1,5 @@ from collections.abc import Iterable, Mapping, Sequence -from itertools import accumulate, chain +from itertools import accumulate, chain, groupby from types import MappingProxyType from typing import Final, TypeAlias @@ -47,7 +47,11 @@ def _chat_tool_key(tool: Tool) -> str: def _indexed_keys(tools: Sequence[Tool]) -> tuple[IndexedKey, ...]: keys: Final = tuple(_chat_tool_key(tool) for tool in tools) - return tuple((key, keys[:position].count(key)) for position, key in enumerate(keys)) + positions_by_key: Final = groupby(sorted(range(len(keys)), key=keys.__getitem__), key=keys.__getitem__) + ordinal_by_position: Final = MappingProxyType( + {position: ordinal for _, positions in positions_by_key for ordinal, position in enumerate(positions)} + ) + return tuple((key, ordinal_by_position[position]) for position, key in enumerate(keys)) def _namespace_members(namespace: Tool) -> tuple[Tool, ...]: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 9f91d7527cb..b2d1a69e0d8 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1903,10 +1903,6 @@ class LiteLLMCompletionResponsesConfig: converted: Final = convert_custom_tool_to_function_tool(tool) return ResponsesToolChatForm(chat_tools=() if converted is None else (converted,), web_search_options=None) if tool_type in ("computer_use", "image_generation", "shell"): - # Drop unsupported Responses-API-only tool types that have no - # Chat Completions equivalent. Passing them through verbatim - # causes providers to reject the request with "'function' is a - # required property". verbose_logger.warning( "Dropping Responses API tool of type '%s': it has no Chat Completions " "equivalent and the target provider would reject the request.", diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py index b80dd0b36aa..4075c209606 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -56,6 +56,24 @@ def test_duplicate_function_names_are_matched_by_ordinal(): assert list(merged) == [original[0]] +def test_interleaved_duplicate_names_keep_their_own_ordinals(): + original = [ + _function("dup", "a"), + _function("other", "x"), + _function("dup", "b"), + _function("dup", "c"), + _function("other", "y"), + ] + groups = _groups(original) + flat = _flat(groups) + edited = {**flat[3], "function": {**flat[3]["function"], "description": "changed"}} + + merged = merge_guardrailed_tools(original, groups, [*flat[:3], edited, flat[4]]) + + assert list(merged) == [*original[:3], {**_function("dup", "changed"), "strict": False}, original[4]] + assert all(merged[position] is original[position] for position in (0, 1, 2, 4)) + + def test_edited_mcp_tool_is_rewritten(): original = [{"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}] groups = _groups(original) From 0346bb265934a09bd5d8eab336facba1cc5bc01b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:58:05 +0000 Subject: [PATCH 41/69] fix(bedrock): upgrade legacy thinking after the invoke response_format stub model swap Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../anthropic_claude3_transformation.py | 5 +++++ ...ations_anthropic_claude3_transformation.py | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 2a4c38e71ea..07ddf6570f2 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -107,6 +107,11 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): # Restore original model name model = original_model + # The stub model hides the original model from the parent's legacy thinking upgrade + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=original_model, optional_params=optional_params, custom_llm_provider="bedrock" + ) + # The stub model hides the original model from the parent's forced-tool-use backstop response_format_tool_choice: Final = optional_params.get("tool_choice") if ( diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 41d82e4f960..d136cb6450b 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -671,3 +671,22 @@ def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice assert "output_format" not in result assert "tools" in result assert "tool_choice" not in result + + +def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking(local_model_cost_map): + """Regression: the tool-based ``response_format`` path swaps in a Claude 3 stub + model before the shared Anthropic mapping, which hid the adaptive-only model + from the legacy ``thinking`` upgrade and left ``type=enabled`` on the wire.""" + result = AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params={ + "response_format": {"type": "json_object"}, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "max_tokens": 8192, + }, + optional_params={}, + model="us.anthropic.claude-fable-5-1", + drop_params=False, + ) + + assert result["thinking"] == {"type": "adaptive"} + assert result["output_config"] == {"effort": "high"} From f4eca10f1d7f832f6300b588434bcd99003f4bb7 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 20:07:57 +0000 Subject: [PATCH 42/69] ci: retrigger checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> From 0f6d983c7057faf13716639869d828d309bcba5b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:08:14 -0700 Subject: [PATCH 43/69] fix(router): skip Claude Code session binding without pre-routing strategies --- litellm/router.py | 2 ++ tests/test_litellm/test_router.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index fc4f2d227f5..b1038ca6002 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12636,6 +12636,8 @@ class Router: registered_model_name: str, request_kwargs: Mapping[str, object], ) -> str: + if not any((self.auto_routers, self.complexity_routers, self.adaptive_routers, self.quality_routers)): + return registered_model_name cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) if cache_key is None or not isinstance(request_kwargs, dict): return registered_model_name diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 05dcb664ed5..ef25502a4f9 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8499,6 +8499,26 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self): + from litellm.caching.caching import RedisCache + + router = self._router() + router.complexity_routers = {} + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(return_value=None) + redis_cache.async_set_cache = AsyncMock() + redis_cache.async_delete_cache = AsyncMock() + router._update_redis_cache(cache=redis_cache) + + for request_kwargs in (self._request_kwargs(), self._request_kwargs(agent_id="agent-1234")): + response = await router.async_pre_routing_hook(model="expensive-model", request_kwargs=request_kwargs) + assert response is None + + redis_cache.async_get_cache.assert_not_awaited() + redis_cache.async_set_cache.assert_not_awaited() + redis_cache.async_delete_cache.assert_not_awaited() + @pytest.mark.asyncio async def test_session_bindings_do_not_evict_router_rate_limit_state(self): router = self._router() From 4705dc6325fc8b2803e59ad19c85f9d07230a9a8 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 20:34:54 +0000 Subject: [PATCH 44/69] fix: retry P3009 when the deadlocked ledger row was already rolled back Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_proxy_extras/utils.py | 46 +++++++---------- .../tests/test_setup_database_fail_fast.py | 51 +++++++++++++++++++ 2 files changed, 70 insertions(+), 27 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 1183757c827..9cc33e6ca56 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -277,22 +277,16 @@ class ProxyExtrasDBManager: pass @staticmethod - def _failed_migration_logs(migration_name: str) -> str: - """Logs recorded on the migration's failed _prisma_migrations row. - - P3009 stderr does not carry the original failure, so this is the only - way to tell a migration that lost a deadlock race against a concurrent - migrate deploy from one whose SQL is genuinely broken. Returns "" when - psycopg is missing, the DB is unreachable, or no failed row exists. - """ + def _failed_migration_logs(migration_name: str) -> Optional[str]: + """Return failed migration logs, or None if the ledger is unavailable.""" database_url = os.getenv("DATABASE_URL") if not database_url: - return "" + return None try: import psycopg except ImportError: - return "" + return None cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) ledger_table = psycopg.sql.SQL("{}.{}").format( @@ -314,7 +308,7 @@ class ProxyExtrasDBManager: (migration_name,), ).fetchone() except (psycopg.OperationalError, psycopg.DatabaseError): - return "" + return None return (row[0] or "") if row else "" @staticmethod @@ -825,22 +819,20 @@ class ProxyExtrasDBManager: f"Detail: {resolve_err}" ) from resolve_err continue - if migration_match and _MIGRATION_DEADLOCK_MARKER in ( - ProxyExtrasDBManager._failed_migration_logs( - migration_match.group(1) - ) - ): - logger.info( - "Migration %s lost a deadlock race against a " - "concurrent migrate deploy, rolling its ledger " - "row back and retrying", - migration_match.group(1), - ) - ProxyExtrasDBManager._roll_back_migration_best_effort( - migration_match.group(1) - ) - time.sleep(random.randrange(5, 15)) - continue + if migration_match: + migration_name = migration_match.group(1) + ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name) + if ledger_logs is not None and ( + ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs + ): + logger.info( + "Migration %s failed in a concurrent migrate deploy " + "deadlock race, rolling its ledger row back and retrying", + migration_name, + ) + ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name) + time.sleep(random.randrange(5, 15)) + continue raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index c4347a91dce..70e0f216338 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -350,6 +350,57 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] +def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path): + """v2: empty failed ledger logs mean a concurrent deploy moved it on.""" + _stub_v2_env(monkeypatch, tmp_path) + + stderr = ( + "Error: P3009\n" + "migrate found failed migrations in the target database\n" + "The `20260415120000_health_check_latest_per_model_index` migration " + "started at 2026-09-01 18:46:13 UTC failed" + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "") + rolled_back = [] + monkeypatch.setattr( + ProxyExtrasDBManager, + "_roll_back_migration", + lambda name: rolled_back.append(name), + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_resolve_specific_migration", + lambda name: pytest.fail("a deadlocked migration must never be marked applied"), + ) + monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] + + +def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): + """v2: an unreadable ledger cannot establish that P3009 was a deadlock.""" + _stub_v2_env(monkeypatch, tmp_path) + + stderr = ( + "Error: P3009\n" + "migrate found failed migrations in the target database\n" + "The `20260415120000_health_check_latest_per_model_index` migration " + "started at 2026-09-01 18:46:13 UTC failed" + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_roll_back_migration", + lambda name: pytest.fail("an unreadable ledger must not trigger a retry"), + ) + monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): """v2: a failed ledger row whose logs show a real SQL error stays fatal.""" _stub_v2_env(monkeypatch, tmp_path) From fd4b15fae6d8592aacd1ee2a60cdb9738f5dfc4e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:37:14 -0700 Subject: [PATCH 45/69] fix(anthropic): upgrade legacy thinking after the Bedrock Invoke and Vertex structured-output stub swap --- .../anthropic_claude3_transformation.py | 4 ++++ .../anthropic/transformation.py | 4 ++++ ...ations_anthropic_claude3_transformation.py | 24 +++++++++++++++++++ ...partner_models_anthropic_transformation.py | 23 ++++++++++++++++++ 4 files changed, 55 insertions(+) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 2a4c38e71ea..67720451c00 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -107,6 +107,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): # Restore original model name model = original_model + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=original_model, optional_params=optional_params, custom_llm_provider="bedrock" + ) + # The stub model hides the original model from the parent's forced-tool-use backstop response_format_tool_choice: Final = optional_params.get("tool_choice") if ( diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index ef03e61a858..7579bc8c02e 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -177,6 +177,10 @@ class VertexAIAnthropicConfig(AnthropicConfig): # Restore original model name for any other processing model = original_model + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=original_model, optional_params=optional_params, custom_llm_provider="vertex_ai" + ) + return optional_params def transform_response( diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 41d82e4f960..4a1e97ca170 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -671,3 +671,27 @@ def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice assert "output_format" not in result assert "tools" in result assert "tool_choice" not in result + + +@pytest.mark.parametrize("model", ["us.anthropic.claude-sonnet-5", "us.anthropic.claude-fable-5-1"]) +def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_thinking(local_model_cost_map, model): + result = AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + }, + }, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "max_tokens": 8192, + }, + optional_params={}, + model=model, + drop_params=False, + ) + + assert "tools" in result + assert result["thinking"] == {"type": "adaptive"} + assert result["output_config"] == {"effort": "high"} diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 9419f88a981..a57672cfbfb 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -752,3 +752,26 @@ def test_vertex_ai_fable_5_1_response_format_uses_native_output_format(local_mod assert "output_format" in result_params assert "tool_choice" not in result_params assert "tools" not in result_params + + +def test_vertex_ai_anthropic_tool_based_response_format_still_upgrades_legacy_thinking(local_model_cost_map): + result_params = VertexAIAnthropicConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + }, + }, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "max_tokens": 8192, + }, + optional_params={}, + model="claude-opus-4-8", + drop_params=False, + ) + + assert "tools" in result_params + assert result_params["thinking"] == {"type": "adaptive"} + assert result_params["output_config"] == {"effort": "high"} From f6eff1bde0f2464ce06c560b55feb7dab4192ab7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:55:16 -0700 Subject: [PATCH 46/69] fix(router): keep Claude Code session bindings across side calls and workers --- litellm/router.py | 9 ++- .../router_code_coverage.py | 1 + tests/test_litellm/test_router.py | 58 +++++++++++++++---- 3 files changed, 56 insertions(+), 12 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b1038ca6002..f4d912eb5be 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12630,6 +12630,12 @@ class Router: e, ) + async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: + session_cache: Final = self._claude_code_session_router_cache + if session_cache.redis_cache is None: + return await session_cache.async_get_cache(key=cache_key) + return await session_cache.redis_cache.async_get_cache(key=cache_key) + async def _resolve_claude_code_session_router( self, model: str, @@ -12646,7 +12652,7 @@ class Router: agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") if agent_id is not None: - bound_model: Final = await self._claude_code_session_router_cache.async_get_cache(key=cache_key) + bound_model: Final = await self._get_claude_code_session_router_binding(cache_key) if not isinstance(bound_model, str): return registered_model_name bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model @@ -12664,7 +12670,6 @@ class Router: if self._request_header(request_kwargs, "x-app") != "cli": return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self._claude_code_session_router_cache.async_set_cache( key=cache_key, diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 60b56b7fac6..0af29f069c6 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -86,6 +86,7 @@ ignored_function_names = [ "_claude_code_session_router_cache_key", # Tested through Claude Code session routing in test_router.py "_delete_claude_code_session_router_binding", # Tested through Redis cleanup failure in test_router.py "_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py + "_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py ] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index ef25502a4f9..1e5767aa645 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8383,18 +8383,21 @@ class TestConsumedRequestTagsStamp: class TestClaudeCodeSubagentSessionRouterBinding: class _RewriteStrategy: + def __init__(self, routed_model: str = "cheap-model") -> None: + self.routed_model = routed_model + async def async_pre_routing_hook( self, model, request_kwargs, messages=None, input=None, specific_deployment=False ): from litellm.types.router import PreRoutingHookResponse return PreRoutingHookResponse( - model="cheap-model", + model=self.routed_model, messages=messages, routing_decision={ "router_model_name": "smart-router", "router_type": "complexity", - "routed_model": "cheap-model", + "routed_model": self.routed_model, "cause": "heuristic_scorer", }, ) @@ -8422,7 +8425,8 @@ class TestClaudeCodeSubagentSessionRouterBinding: num_retries=0, ) router.complexity_routers = { - "smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy())] + "smart-router": (TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy()),), + "premium-router": (TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy("expensive-model")),), } return router @@ -8467,7 +8471,7 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert subagent_kwargs["metadata"]["routing_decision"]["router_model_name"] == "smart-router" @pytest.mark.asyncio - async def test_main_direct_model_clears_the_session_router(self): + async def test_main_thread_side_calls_to_a_plain_model_keep_the_session_router(self): router = self._router() await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) @@ -8478,33 +8482,67 @@ class TestClaudeCodeSubagentSessionRouterBinding: request_kwargs=self._request_kwargs(agent_id="agent-1234"), ) - assert response is None + assert response is not None + assert response.model == "cheap-model" @pytest.mark.asyncio - async def test_redis_cleanup_failure_does_not_reject_a_direct_model_request(self): + async def test_redis_cleanup_failure_does_not_reject_a_subagent_request(self): from litellm.caching.caching import RedisCache router = self._router() + del router.complexity_routers["smart-router"] redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(return_value="smart-router") redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) - - await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) router._update_redis_cache(cache=redis_cache) response = await router.async_pre_routing_hook( model="expensive-model", - request_kwargs=self._request_kwargs(), + request_kwargs=self._request_kwargs(agent_id="agent-1234"), ) assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_subagents_follow_the_main_threads_latest_router_across_workers(self): + from types import SimpleNamespace + + from litellm.caching.caching import RedisCache + + shared_binding = SimpleNamespace(value=None) + shared_redis = MagicMock(spec=RedisCache) + shared_redis.async_get_cache = AsyncMock(side_effect=lambda key, **_: shared_binding.value) + shared_redis.async_set_cache = AsyncMock( + side_effect=lambda key, value, **_: setattr(shared_binding, "value", value) + ) + main_worker, subagent_worker = self._router(), self._router() + main_worker._update_redis_cache(cache=shared_redis) + subagent_worker._update_redis_cache(cache=shared_redis) + + await main_worker.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + first = await subagent_worker.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + await main_worker.async_pre_routing_hook(model="premium-router", request_kwargs=self._request_kwargs()) + second = await subagent_worker.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert first is not None + assert first.model == "cheap-model" + assert second is not None + assert second.model == "expensive-model" + assert shared_binding.value == "premium-router" + @pytest.mark.asyncio async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self): from litellm.caching.caching import RedisCache router = self._router() - router.complexity_routers = {} + router.complexity_routers.clear() redis_cache = MagicMock(spec=RedisCache) redis_cache.async_get_cache = AsyncMock(return_value=None) redis_cache.async_set_cache = AsyncMock() From 86c8b93bf748ca5a0d5a8e59d366a45f04c7ba34 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:14:41 -0700 Subject: [PATCH 47/69] fix(vector-store): embed Milvus and Azure AI Search queries through the request executor Milvus REST and Azure AI Search still embedded the query through the SDK, so a bare Router alias as litellm_embedding_model kept failing after the executor landed for Valkey. Both now share BaseQueryEmbeddingVectorStoreConfig, which embeds through the injected executor, drops the empty litellm_embedding_config requirement, and awaits aembedding on the async path. The Router executor falls back to the SDK for models the Router does not serve, so inline provider configs such as azure/text-embedding-3-large with their own credentials keep working through the proxy. Tests fake OpenAI and Milvus at the HTTP boundary with respx instead of patching litellm.embedding. --- .../azure_ai/vector_stores/transformation.py | 130 +++++----- .../base_llm/vector_store/transformation.py | 119 +++++++++- litellm/llms/custom_httpx/llm_http_handler.py | 46 ++-- .../milvus/vector_stores/transformation.py | 129 +++++----- .../test_router_embedding_integration.py | 183 ++++++++------ .../test_azure_ai_vector_store.py | 119 +++++++++- .../test_milvus_vector_store.py | 223 +++++++++++++++--- uv.lock | 22 +- 8 files changed, 700 insertions(+), 271 deletions(-) diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index 5e61d0a1dd9..044b8f5243c 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -1,10 +1,13 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx -import litellm from litellm.llms.azure.common_utils import BaseAzureLLM -from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import ( + BaseQueryEmbeddingVectorStoreConfig, + VectorStoreEmbeddingExecutor, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( BaseVectorStoreAuthCredentials, @@ -25,7 +28,7 @@ else: LiteLLMLoggingObj = Any -class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): +class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM): """ Configuration for Azure AI Search Vector Store @@ -109,82 +112,71 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): 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, - ) -> tuple[str, dict[str, Any]]: - """ - Transform search request for Azure AI Search API + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | 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) + return self._search_request( + vector_store_id, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, + litellm_params, + ) - Generates embeddings using litellm.embeddings and constructs Azure AI Search request - """ - # Convert query to string if it's a list - if isinstance(query, list): - query = " ".join(query) + async def atransform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | 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) + return self._search_request( + vector_store_id, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, + litellm_params, + ) - # Get embedding model from litellm_params (required) - embedding_model: Final = litellm_params.get("litellm_embedding_model") - if not embedding_model: - raise ValueError( - "embedding_model is required in litellm_params for Azure AI Search. " - "Example: litellm_params['embedding_model'] = 'azure/text-embedding-3-large'" - ) - - embedding_config: Final = litellm_params.get("litellm_embedding_config", {}) - if not embedding_config: - raise ValueError( - "embedding_config is required in litellm_params for Azure AI Search. " - "Example: litellm_params['embedding_config'] = {'api_base': 'https://krris-mh44uf7y-eastus2.cognitiveservices.azure.com/', 'api_key': 'os.environ/AZURE_API_KEY', 'api_version': '2025-09-01'}" - ) - - # Get vector field name (defaults to contentVector) + @staticmethod + def _search_request( + vector_store_id: str, + query_text: str, + query_vector: Sequence[float], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: vector_field: Final = litellm_params.get("azure_search_vector_field", "contentVector") - - # Get top_k (number of results to return) top_k: Final = vector_store_search_optional_params.get("top_k", 10) - - # Generate embedding for the query using litellm.embeddings - try: - embedding_response: Final = litellm.embedding( - model=embedding_model, - input=[query], - **embedding_config, - ) - query_vector: Final = embedding_response.data[0]["embedding"] - except Exception as e: - raise Exception(f"Failed to generate embedding for query: {e}") - - # Azure AI Search endpoint for search - index_name: Final = vector_store_id # vector_store_id is the index name - url: Final = f"{api_base}/indexes/{index_name}/docs/search?api-version=2024-07-01" - - # Build the request body for Azure AI Search with vector search - request_body: Final = { - "search": "*", # Get all documents (filtered by vector similarity) - "vectorQueries": [ - { - "vector": query_vector, - "fields": vector_field, - "kind": "vector", - "k": top_k, # Number of nearest neighbors to return - } - ], - "select": "id,content", # Fields to return (customize based on schema) + litellm_logging_obj.model_call_details["input"] = query_text + litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model") + litellm_logging_obj.model_call_details["top_k"] = top_k + return f"{api_base}/indexes/{vector_store_id}/docs/search?api-version=2024-07-01", { + "search": "*", + "vectorQueries": [{"vector": query_vector, "fields": vector_field, "kind": "vector", "k": top_k}], + "select": "id,content", "top": top_k, } - ######################################################### - # Update logging object with details of the request - ######################################################### - litellm_logging_obj.model_call_details["input"] = query - litellm_logging_obj.model_call_details["embedding_model"] = embedding_model - litellm_logging_obj.model_call_details["top_k"] = top_k - - return url, request_body - def transform_search_vector_store_response( self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj ) -> VectorStoreSearchResponse: diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index e9c925448a8..95863266bf7 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -3,9 +3,11 @@ from __future__ import annotations from abc import abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, runtime_checkable import httpx +from pydantic import TypeAdapter from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import EmbeddingResponse @@ -65,7 +67,7 @@ class RouterVectorStoreEmbeddingExecutor: router: Router metadata: Mapping[str, object] - def _embedding_kwargs(self, configuration: Mapping[str, object]) -> dict[str, object]: + def _embedding_kwargs(self, configuration: Mapping[str, object]) -> Mapping[str, object]: configured_metadata: Final = configuration.get("metadata") metadata: Final = { **(configured_metadata if isinstance(configured_metadata, Mapping) else {}), @@ -76,18 +78,32 @@ class RouterVectorStoreEmbeddingExecutor: "metadata": metadata, } + def _router_serves(self, model: str) -> bool: + team_id: Final = self.metadata.get("user_api_key_team_id") + resolved: Final = self.router.resolved_litellm_models(model, team_id if isinstance(team_id, str) else None) + deployment_models: Final = ( + deployment.get("litellm_params", {}).get("model") for deployment in self.router.get_model_list() or () + ) + return bool(resolved) or model in deployment_models + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + embedding_kwargs: Final = self._embedding_kwargs(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, input=[query], # mutable-ok: Router embedding requires a mutable input list - **self._embedding_kwargs(configuration), # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic + **embedding_kwargs, # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic ) async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + embedding_kwargs: Final = self._embedding_kwargs(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, input=[query], # mutable-ok: Router embedding requires a mutable input list - **self._embedding_kwargs(configuration), # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic + **embedding_kwargs, # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic ) @@ -221,6 +237,103 @@ class BaseVectorStoreConfig: return 0.0, 0.0 +_EMPTY_EMBEDDING_CONFIGURATION: Final[Mapping[str, object]] = MappingProxyType({}) +_QUERY_VECTOR: Final = TypeAdapter(list[float]) + + +class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): + @abstractmethod + def transform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + pass + + async def atransform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + return self.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=litellm_logging_obj, + litellm_params=litellm_params, + extra_body=extra_body, + embedding_executor=embedding_executor, + ) + + @staticmethod + def query_text(query: str | Sequence[str]) -> str: + return query if isinstance(query, str) else " ".join(query) + + @staticmethod + def query_embedding_model(litellm_params: Mapping[str, object]) -> str: + embedding_model: Final = litellm_params.get("litellm_embedding_model") + if isinstance(embedding_model, str) and embedding_model: + return embedding_model + raise ValueError( + "litellm_embedding_model is required in litellm_params for this vector store. " + "Example: litellm_params['litellm_embedding_model'] = 'openai/text-embedding-3-small'" + ) + + @staticmethod + def query_embedding_configuration(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + configuration: Final = litellm_params.get("litellm_embedding_config") + if isinstance(configuration, Mapping): + return {str(key): value for key, value in configuration.items()} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # litellm_params is an untyped dict, keys are re-validated as str here + return _EMPTY_EMBEDDING_CONFIGURATION + + def embed_query( + self, + query_text: str, + litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None, + ) -> Sequence[float]: + model: Final = self.query_embedding_model(litellm_params) + configuration: Final = self.query_embedding_configuration(litellm_params) + executor: Final = ( + embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor() + ) + try: + response: Final = executor.embed(model, query_text, configuration) + except Exception as e: + raise Exception(f"Failed to generate embedding for query: {e}") + return _QUERY_VECTOR.validate_python(response.data[0]["embedding"]) # pyright: ignore[reportUnknownMemberType] # EmbeddingResponse.data is an untyped list, the vector is validated here + + async def aembed_query( + self, + query_text: str, + litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None, + ) -> Sequence[float]: + model: Final = self.query_embedding_model(litellm_params) + configuration: Final = self.query_embedding_configuration(litellm_params) + executor: Final = ( + embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor() + ) + try: + response: Final = await executor.aembed(model, query_text, configuration) + except Exception as e: + raise Exception(f"Failed to generate embedding for query: {e}") + return _QUERY_VECTOR.validate_python(response.data[0]["embedding"]) # pyright: ignore[reportUnknownMemberType] # EmbeddingResponse.data is an untyped list, the vector is validated here + + class BaseDirectVectorStoreConfig(BaseVectorStoreConfig): """ Base config for vector store providers whose datastore has no HTTP API diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 118656b81a2..c0ef7456680 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -69,6 +69,7 @@ from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig from litellm.llms.base_llm.vector_store.transformation import ( BaseDirectVectorStoreConfig, + BaseQueryEmbeddingVectorStoreConfig, BaseVectorStoreConfig, VectorStoreEmbeddingExecutor, ) @@ -9728,8 +9729,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - # Check if provider has async transform method - if hasattr(vector_store_provider_config, "atransform_search_vector_store_request"): + if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig): ( url, request_body, @@ -9741,12 +9741,13 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, + embedding_executor=embedding_executor, ) else: ( url, request_body, - ) = vector_store_provider_config.transform_search_vector_store_request( + ) = await vector_store_provider_config.atransform_search_vector_store_request( vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, @@ -9857,18 +9858,33 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - ( - url, - request_body, - ) = vector_store_provider_config.transform_search_vector_store_request( - vector_store_id=vector_store_id, - query=query, - vector_store_search_optional_params=vector_store_search_optional_params, - api_base=api_base, - litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), - extra_body=extra_body, - ) + if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig): + ( + url, + request_body, + ) = vector_store_provider_config.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + extra_body=extra_body, + embedding_executor=embedding_executor, + ) + else: + ( + url, + request_body, + ) = vector_store_provider_config.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + extra_body=extra_body, + ) all_optional_params: Final[dict[str, object]] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 34f0cd854c4..b0291c692d5 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -1,9 +1,12 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx -import litellm -from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import ( + BaseQueryEmbeddingVectorStoreConfig, + VectorStoreEmbeddingExecutor, +) from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -36,7 +39,7 @@ MILVUS_OPTIONAL_PARAMS: Final = { } -class MilvusVectorStoreConfig(BaseVectorStoreConfig): +class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): """ Configuration for Milvus Vector Store @@ -117,77 +120,77 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): 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, - ) -> tuple[str, dict[str, Any]]: - """ - Transform search request for Azure AI Search API + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | 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) + return self._search_request( + vector_store_id, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, + litellm_params, + ) - Generates embeddings using litellm.embeddings and constructs Azure AI Search request - """ - # Convert query to string if it's a list - if isinstance(query, list): - query = " ".join(query) + async def atransform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | 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) + return self._search_request( + vector_store_id, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, + litellm_params, + ) - # Get embedding model from litellm_params (required) - embedding_model: Final = litellm_params.get("litellm_embedding_model") - if not embedding_model: - raise ValueError( - "embedding_model is required in litellm_params for Milvus. You can call any litellm embedding model." - "Example: litellm_params['embedding_model'] = 'azure/text-embedding-3-large'" + @staticmethod + def _search_request( + vector_store_id: str, + query_text: str, + query_vector: Sequence[float], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: + scope: Final = { + key: value + for key, value in ( + ("dbName", litellm_params.get("milvus_db_name")), + ("partitionNames", litellm_params.get("milvus_partition_names")), ) - - embedding_config: Final = litellm_params.get("litellm_embedding_config", {}) - if not embedding_config: - raise ValueError( - "embedding_config is required in litellm_params for Milvus. You can call any litellm embedding model." - "Example: litellm_params['embedding_config'] = {'api_base': 'https://krris-mh44uf7y-eastus2.cognitiveservices.azure.com/', 'api_key': 'os.environ/AZURE_API_KEY', 'api_version': '2025-09-01'}" - ) - - # Get top_k (number of results to return) - # Generate embedding for the query using litellm.embeddings - try: - embedding_response: Final = litellm.embedding( - model=embedding_model, - input=[query], - **embedding_config, - ) - query_vector: Final = embedding_response.data[0]["embedding"] - except Exception as e: - raise Exception(f"Failed to generate embedding for query: {e}") - - # Azure AI Search endpoint for search - index_name: Final = vector_store_id # vector_store_id is the index name - url: Final = f"{api_base}/v2/vectordb/entities/search" - - # Build the request body for Azure AI Search with vector search - request_body: Final[dict[str, Any]] = { - "collectionName": index_name, + if value + } + litellm_logging_obj.model_call_details["input"] = query_text + litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model") + return f"{api_base}/v2/vectordb/entities/search", { + "collectionName": vector_store_id, "data": [query_vector], "annsField": "book_intro_vector", **vector_store_search_optional_params, + **scope, } - db_name: Final = litellm_params.get("milvus_db_name") - if db_name: - request_body["dbName"] = db_name - - partition_names: Final = litellm_params.get("milvus_partition_names") - if partition_names: - request_body["partitionNames"] = partition_names - - ######################################################### - # Update logging object with details of the request - ######################################################### - litellm_logging_obj.model_call_details["input"] = query - litellm_logging_obj.model_call_details["embedding_model"] = embedding_model - - return url, request_body - def transform_search_vector_store_response( self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj ) -> VectorStoreSearchResponse: diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index d5e0c750d88..96c6ce3708e 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -5,16 +5,57 @@ These tests simulate real-world scenarios where headers and configuration need to be properly propagated through the router to the LLM API. """ +import json from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx +import litellm from litellm import Router from litellm.llms.base_llm.vector_store.transformation import ( LiteLLMVectorStoreEmbeddingExecutor, RouterVectorStoreEmbeddingExecutor, ) -from litellm.types.utils import EmbeddingResponse + +QUERY_VECTOR = [0.5, -0.25, 0.125] +OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" +STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + + +def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route: + return respx_mock.post(url).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": QUERY_VECTOR}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + + +def _sent(route: respx.Route, index: int) -> tuple[str, str, list[str]]: + request = route.calls[index].request + body = json.loads(request.read()) + return request.headers["authorization"], body["model"], body["input"] + + +def _alias_router() -> Router: + return Router( + model_list=[ + { + "model_name": "team-alias", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + } + ] + ) class TestRouterEmbeddingIntegration: @@ -70,48 +111,23 @@ class TestRouterEmbeddingIntegration: ) @pytest.mark.asyncio - async def test_vector_store_embedding_executors_cover_sdk_and_router_paths(self): - response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + async def test_vector_store_embedding_executors_cover_sdk_and_router_paths( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + openai_route = _mock_embedding_route(respx_mock, OPENAI_EMBEDDINGS_URL) + store_route = _mock_embedding_route(respx_mock, STORE_EMBEDDINGS_URL) sdk_executor = LiteLLMVectorStoreEmbeddingExecutor() - with ( - patch("litellm.embedding", return_value=response) as embedding, - patch("litellm.aembedding", new=AsyncMock(return_value=response)) as aembedding, - ): - assert sdk_executor.embed("openai/model", "sync", {"api_key": "explicit"}) is response - assert await sdk_executor.aembed("openai/model", "async", {"api_key": "explicit"}) is response + sync_response = sdk_executor.embed("openai/text-embedding-3-small", "sync", {"api_key": "explicit"}) + async_response = await sdk_executor.aembed("openai/text-embedding-3-small", "async", {"api_key": "explicit"}) - embedding.assert_called_once_with(model="openai/model", input=["sync"], api_key="explicit") - aembedding.assert_awaited_once_with(model="openai/model", input=["async"], api_key="explicit") + assert sync_response.data[0]["embedding"] == QUERY_VECTOR + assert async_response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(openai_route, 0) == ("Bearer explicit", "text-embedding-3-small", ["sync"]) + assert _sent(openai_route, 1) == ("Bearer explicit", "text-embedding-3-small", ["async"]) - mock_router = MagicMock() - mock_router.embedding.return_value = response - router_executor = RouterVectorStoreEmbeddingExecutor( - router=mock_router, - metadata={"user_api_key_team_id": "team-a"}, - ) - assert router_executor.embed("team-alias", "query", {}) is response - mock_router.embedding.assert_called_once_with( - model="team-alias", - input=["query"], - metadata={"user_api_key_team_id": "team-a"}, - ) - - alias_router = Router( - model_list=[ - { - "model_name": "team-alias", - "litellm_params": { - "model": "openai/text-embedding-3-small", - "api_key": "deployment-key", - }, - } - ] - ) - alias_executor = RouterVectorStoreEmbeddingExecutor( - router=alias_router, - metadata={"user_api_key_team_id": "team-a"}, - ) explicit_config = { "api_base": "https://embedding.example/v1", "api_key": "store-key", @@ -121,29 +137,66 @@ class TestRouterEmbeddingIntegration: }, "model": "untrusted-model", } + mock_router = MagicMock() + mock_router.embedding.return_value = sync_response + router_executor = RouterVectorStoreEmbeddingExecutor( + router=mock_router, + metadata={"user_api_key_team_id": "team-a"}, + ) + assert router_executor.embed("team-alias", "query", explicit_config) is sync_response + mock_router.embedding.assert_called_once_with( + model="team-alias", + input=["query"], + api_base="https://embedding.example/v1", + api_key="store-key", + metadata={"configured": True, "user_api_key_team_id": "team-a"}, + ) - with ( - patch("litellm.embedding", return_value=response) as explicit_embedding, - patch("litellm.aembedding", new=AsyncMock(return_value=response)) as explicit_aembedding, - ): - assert alias_executor.embed("team-alias", "sync query", explicit_config) is response - assert await alias_executor.aembed("team-alias", "async query", explicit_config) is response + alias_executor = RouterVectorStoreEmbeddingExecutor( + router=_alias_router(), + metadata={"user_api_key_team_id": "team-a"}, + ) + sync_alias = alias_executor.embed("team-alias", "sync query", explicit_config) + async_alias = await alias_executor.aembed("team-alias", "async query", explicit_config) - sync_kwargs = explicit_embedding.call_args.kwargs - assert sync_kwargs["model"] == "openai/text-embedding-3-small" - assert sync_kwargs["input"] == ["sync query"] - assert sync_kwargs["api_base"] == "https://embedding.example/v1" - assert sync_kwargs["api_key"] == "store-key" - assert sync_kwargs["metadata"]["configured"] is True - assert sync_kwargs["metadata"]["user_api_key_team_id"] == "team-a" + assert sync_alias.data[0]["embedding"] == QUERY_VECTOR + assert async_alias.data[0]["embedding"] == QUERY_VECTOR + assert openai_route.call_count == 2 + assert _sent(store_route, 0) == ("Bearer store-key", "text-embedding-3-small", ["sync query"]) + assert _sent(store_route, 1) == ("Bearer store-key", "text-embedding-3-small", ["async query"]) - async_kwargs = explicit_aembedding.await_args.kwargs - assert async_kwargs["model"] == "openai/text-embedding-3-small" - assert async_kwargs["input"] == ["async query"] - assert async_kwargs["api_base"] == "https://embedding.example/v1" - assert async_kwargs["api_key"] == "store-key" - assert async_kwargs["metadata"]["configured"] is True - assert async_kwargs["metadata"]["user_api_key_team_id"] == "team-a" + @pytest.mark.asyncio + async def test_router_executor_falls_back_to_sdk_for_models_the_router_does_not_serve( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + store_route = _mock_embedding_route(respx_mock, STORE_EMBEDDINGS_URL) + executor = RouterVectorStoreEmbeddingExecutor( + router=_alias_router(), + metadata={"user_api_key_team_id": "team-a"}, + ) + inline_config = {"api_base": "https://embedding.example/v1", "api_key": "store-key"} + + sync_response = executor.embed("openai/text-embedding-3-large", "sync query", inline_config) + async_response = await executor.aembed("openai/text-embedding-3-large", "async query", inline_config) + + assert sync_response.data[0]["embedding"] == QUERY_VECTOR + assert async_response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(store_route, 0) == ("Bearer store-key", "text-embedding-3-large", ["sync query"]) + assert _sent(store_route, 1) == ("Bearer store-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 + ): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + openai_route = _mock_embedding_route(respx_mock, OPENAI_EMBEDDINGS_URL) + executor = RouterVectorStoreEmbeddingExecutor(router=_alias_router(), metadata={}) + + response = executor.embed("openai/text-embedding-3-small", "query", {}) + + assert response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(openai_route, 0) == ("Bearer deployment-key", "text-embedding-3-small", ["query"]) def test_embedding_with_deployment_specific_headers(self): """ @@ -251,9 +304,7 @@ class TestRouterEmbeddingIntegration: router = Router( model_list=model_list, - default_litellm_params={ - "metadata": {"environment": "test", "service": "embedding-service"} - }, + default_litellm_params={"metadata": {"environment": "test", "service": "embedding-service"}}, ) with patch("litellm.embedding") as mock_embedding: @@ -369,9 +420,7 @@ class TestRouterEmbeddingIntegration: # Make multiple calls and verify headers are always present for i in range(5): with patch("litellm.embedding") as mock_embedding: - mock_embedding.return_value = MagicMock( - data=[{"embedding": [0.1, 0.2]}] - ) + mock_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2]}]) router.embedding(model="shared-embedding-model", input=[f"test {i}"]) @@ -456,9 +505,7 @@ class TestRouterEmbeddingIntegration: router = Router( model_list=model_list, - default_litellm_params={ - "headers": {"X-Custom-Azure-Header": "azure-value"} - }, + default_litellm_params={"headers": {"X-Custom-Azure-Header": "azure-value"}}, ) with patch("litellm.embedding") as mock_embedding: diff --git a/tests/vector_store_tests/test_azure_ai_vector_store.py b/tests/vector_store_tests/test_azure_ai_vector_store.py index 58e45f259ab..d1fc8436fc9 100644 --- a/tests/vector_store_tests/test_azure_ai_vector_store.py +++ b/tests/vector_store_tests/test_azure_ai_vector_store.py @@ -1,10 +1,19 @@ -import pytest -import litellm import json import os +from unittest.mock import MagicMock + +import httpx +import pytest +import respx + +import litellm +from litellm.llms.azure_ai.vector_stores.transformation import AzureAIVectorStoreConfig +from litellm.types.utils import EmbeddingResponse +from litellm.vector_stores import ( + asearch as vector_store_asearch, +) from litellm.vector_stores import ( search as vector_store_search, - asearch as vector_store_asearch, ) @@ -30,10 +39,108 @@ async def test_basic_search_vector_store(sync_mode): if sync_mode: response = vector_store_search(query=default_query, **base_request_args) else: - response = await vector_store_asearch( - query=default_query, **base_request_args - ) + response = await vector_store_asearch(query=default_query, **base_request_args) except litellm.InternalServerError: pytest.skip("Skipping test due to litellm.InternalServerError") print("litellm response=", json.dumps(response, indent=4, default=str)) + + +class RecordingEmbeddingExecutor: + def __init__(self, response): + self.response = response + self.calls = [] + + def embed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + async def aembed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + +ALIAS_QUERY_VECTOR = [0.5, -0.25, 0.125] +ALIAS_EMBEDDING_RESPONSE = EmbeddingResponse( + data=[{"embedding": ALIAS_QUERY_VECTOR, "index": 0, "object": "embedding"}] +) +STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + + +def _transform_kwargs(executor): + logging_obj = MagicMock() + logging_obj.model_call_details = {} + return { + "vector_store_id": "my-vector-index", + "query": "what is azure search?", + "vector_store_search_optional_params": {"top_k": 2}, + "api_base": "https://azure-kb-search.search.windows.net", + "litellm_logging_obj": logging_obj, + "litellm_params": { + "litellm_embedding_model": "multilingual-e5-large", + "azure_search_vector_field": "embedding", + }, + "embedding_executor": executor, + } + + +@pytest.mark.asyncio +async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): + executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE) + config = AzureAIVectorStoreConfig() + transform_kwargs = _transform_kwargs(executor) + + url, sync_body = config.transform_search_vector_store_request(**transform_kwargs) + _, async_body = await config.atransform_search_vector_store_request(**transform_kwargs) + + assert respx_mock.calls.call_count == 0 + assert executor.calls == [("multilingual-e5-large", "what is azure search?", {})] * 2 + assert ( + url == "https://azure-kb-search.search.windows.net/indexes/my-vector-index/docs/search?api-version=2024-07-01" + ) + assert sync_body == async_body + assert sync_body["vectorQueries"] == [ + {"vector": ALIAS_QUERY_VECTOR, "fields": "embedding", "kind": "vector", "k": 2} + ] + assert sync_body["top"] == 2 + logging_details = transform_kwargs["litellm_logging_obj"].model_call_details + assert logging_details["embedding_model"] == "multilingual-e5-large" + assert logging_details["top_k"] == 2 + + +def test_transform_falls_back_to_sdk_embedding_without_executor( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = respx_mock.post(STORE_EMBEDDINGS_URL).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": ALIAS_QUERY_VECTOR}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + transform_kwargs = _transform_kwargs(None) + transform_kwargs["litellm_params"] = { + "litellm_embedding_model": "openai/text-embedding-3-small", + "litellm_embedding_config": {"api_base": "https://embedding.example/v1", "api_key": "store-key"}, + } + + _, body = AzureAIVectorStoreConfig().transform_search_vector_store_request(**transform_kwargs) + + embedding_request = embedding_route.calls.last.request + assert embedding_request.headers["authorization"] == "Bearer store-key" + assert json.loads(embedding_request.read())["input"] == ["what is azure search?"] + assert body["vectorQueries"][0]["vector"] == ALIAS_QUERY_VECTOR + assert body["vectorQueries"][0]["fields"] == "contentVector" + + +def test_transform_requires_embedding_model(): + transform_kwargs = _transform_kwargs(RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE)) + transform_kwargs["litellm_params"] = {"litellm_embedding_config": {"api_key": "store-key"}} + + with pytest.raises(ValueError, match="litellm_embedding_model is required"): + AzureAIVectorStoreConfig().transform_search_vector_store_request(**transform_kwargs) diff --git a/tests/vector_store_tests/test_milvus_vector_store.py b/tests/vector_store_tests/test_milvus_vector_store.py index 6627f6006d1..ea3c1883e46 100644 --- a/tests/vector_store_tests/test_milvus_vector_store.py +++ b/tests/vector_store_tests/test_milvus_vector_store.py @@ -3,16 +3,19 @@ Tests for Milvus Vector Store """ import json -import os from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm +from litellm import Router +from litellm.llms.milvus.vector_stores.transformation import MilvusVectorStoreConfig +from litellm.types.utils import EmbeddingResponse from litellm.vector_stores import asearch as vector_store_asearch from litellm.vector_stores import search as vector_store_search - # Mock response from actual Milvus API MOCK_MILVUS_SEARCH_RESPONSE = { "code": 0, @@ -98,7 +101,7 @@ class TestMilvusVectorStore: mock_response.json.return_value = MOCK_MILVUS_SEARCH_RESPONSE mock_response.text = json.dumps(MOCK_MILVUS_SEARCH_RESPONSE) - with patch("litellm.embedding") as mock_embedding: + with patch("litellm.aembedding", new_callable=AsyncMock) as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE with patch( @@ -147,16 +150,10 @@ class TestMilvusVectorStore: else: # Fallback: check for json kwarg or in args request_data = call_args.kwargs.get("json") - if ( - request_data is None - and len(call_args.args) > 0 - and isinstance(call_args.args[0], dict) - ): + if request_data is None and len(call_args.args) > 0 and isinstance(call_args.args[0], dict): request_data = call_args.args[0] - assert ( - request_data is not None - ), f"Could not extract request data. Call args: {call_args}" + assert request_data is not None, f"Could not extract request data. Call args: {call_args}" print("Request data:", json.dumps(request_data, indent=2, default=str)) # Validate request structure @@ -213,9 +210,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response # Make the search request @@ -252,16 +247,10 @@ class TestMilvusVectorStore: else: # Fallback: check for json kwarg or in args request_data = call_args.kwargs.get("json") - if ( - request_data is None - and len(call_args.args) > 0 - and isinstance(call_args.args[0], dict) - ): + if request_data is None and len(call_args.args) > 0 and isinstance(call_args.args[0], dict): request_data = call_args.args[0] - assert ( - request_data is not None - ), f"Could not extract request data. Call args: {call_args}" + assert request_data is not None, f"Could not extract request data. Call args: {call_args}" # Validate request structure assert "collectionName" in request_data @@ -316,11 +305,7 @@ class TestMilvusVectorStore: if request_data_str: return json.loads(request_data_str) request_data = call_args.kwargs.get("json") - if ( - request_data is None - and len(call_args.args) > 0 - and isinstance(call_args.args[0], dict) - ): + if request_data is None and len(call_args.args) > 0 and isinstance(call_args.args[0], dict): request_data = call_args.args[0] return request_data @@ -334,9 +319,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response vector_store_search( @@ -375,9 +358,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response vector_store_search( @@ -413,9 +394,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response vector_store_search( @@ -492,3 +471,175 @@ if __name__ == "__main__": test.test_basic_search_with_mock_sync() print("\n✅ All mock tests passed!") + + +class RecordingEmbeddingExecutor: + def __init__(self, response): + self.response = response + self.calls = [] + + def embed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + async def aembed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + +ALIAS_QUERY_VECTOR = [0.5, -0.25, 0.125] +ALIAS_EMBEDDING_RESPONSE = EmbeddingResponse( + data=[{"embedding": ALIAS_QUERY_VECTOR, "index": 0, "object": "embedding"}] +) +OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" +MILVUS_SEARCH_URL = "https://milvus.example/v2/vectordb/entities/search" +ALIAS_SEARCH_KWARGS = { + "query": "what is machine learning?", + "vector_store_id": "book_2", + "custom_llm_provider": "milvus", + "api_base": "https://milvus.example", + "api_key": "mock_milvus_api_key", + "litellm_embedding_model": "multilingual-e5-large", + "milvus_text_field": "book_intro_text", +} + + +def _alias_router(): + return Router( + model_list=[ + { + "model_name": "multilingual-e5-large", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + } + ] + ) + + +def _mock_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(OPENAI_EMBEDDINGS_URL).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": ALIAS_QUERY_VECTOR}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + + +def _mock_search_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(MILVUS_SEARCH_URL).mock(return_value=httpx.Response(200, json=MOCK_MILVUS_SEARCH_RESPONSE)) + + +def _assert_alias_resolved(embedding_route: respx.Route, search_route: respx.Route, response): + embedding_request = embedding_route.calls.last.request + assert embedding_request.headers["authorization"] == "Bearer deployment-key" + embedding_body = json.loads(embedding_request.read()) + assert embedding_body["model"] == "text-embedding-3-small" + assert embedding_body["input"] == ["what is machine learning?"] + search_request = search_route.calls.last.request + assert search_request.headers["authorization"] == "Bearer mock_milvus_api_key" + assert json.loads(search_request.read())["data"] == [ALIAS_QUERY_VECTOR] + assert len(response["data"]) == len(MOCK_MILVUS_SEARCH_RESPONSE["data"]) + assert response["data"][0]["content"][0]["text"] == MOCK_MILVUS_SEARCH_RESPONSE["data"][0]["book_intro_text"] + + +def test_router_search_resolves_bare_embedding_alias_sync( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = _alias_router().vector_store_search(**ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_router_search_resolves_bare_embedding_alias_async( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = await _alias_router().avector_store_search(**ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): + executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE) + config = MilvusVectorStoreConfig() + logging_obj = MagicMock() + logging_obj.model_call_details = {} + transform_kwargs = { + "vector_store_id": "book_2", + "query": ["what is", "milvus?"], + "vector_store_search_optional_params": {"limit": 3}, + "api_base": "https://milvus.example", + "litellm_logging_obj": logging_obj, + "litellm_params": {"litellm_embedding_model": "multilingual-e5-large", "milvus_db_name": "docs"}, + "embedding_executor": executor, + } + + url, sync_body = config.transform_search_vector_store_request(**transform_kwargs) + _, async_body = await config.atransform_search_vector_store_request(**transform_kwargs) + + assert respx_mock.calls.call_count == 0 + assert executor.calls == [("multilingual-e5-large", "what is milvus?", {})] * 2 + assert url == MILVUS_SEARCH_URL + assert sync_body == async_body + assert sync_body == { + "collectionName": "book_2", + "data": [ALIAS_QUERY_VECTOR], + "annsField": "book_intro_vector", + "limit": 3, + "dbName": "docs", + } + assert logging_obj.model_call_details["input"] == "what is milvus?" + assert logging_obj.model_call_details["embedding_model"] == "multilingual-e5-large" + + +def test_transform_falls_back_to_sdk_embedding_without_executor_or_config( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + embedding_route = _mock_embedding_route(respx_mock) + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + _, body = MilvusVectorStoreConfig().transform_search_vector_store_request( + vector_store_id="book_2", + query="q", + vector_store_search_optional_params={}, + api_base="https://milvus.example", + litellm_logging_obj=logging_obj, + litellm_params={"litellm_embedding_model": "openai/text-embedding-3-small"}, + ) + + embedding_request = embedding_route.calls.last.request + assert embedding_request.headers["authorization"] == "Bearer env-key" + assert json.loads(embedding_request.read())["input"] == ["q"] + assert body["data"] == [ALIAS_QUERY_VECTOR] + + +def test_transform_requires_embedding_model(): + with pytest.raises(ValueError, match="litellm_embedding_model is required"): + MilvusVectorStoreConfig().transform_search_vector_store_request( + vector_store_id="book_2", + query="q", + vector_store_search_optional_params={}, + api_base="https://milvus.example", + litellm_logging_obj=MagicMock(), + litellm_params={"litellm_embedding_config": {"api_key": "store-key"}}, + embedding_executor=RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE), + ) diff --git a/uv.lock b/uv.lock index 8d886044083..27be919eea1 100644 --- a/uv.lock +++ b/uv.lock @@ -9441,19 +9441,19 @@ wheels = [ [[package]] name = "tornado" -version = "6.5.8" +version = "6.5.7" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/10/d3/343e5bb989d6515b1646cf3d40135d73f3d5e45339bded401b56cdac24dd/tornado-6.5.8.tar.gz", hash = "sha256:9452e1b208a8bd771e2cb1f2ff564985b9b214bdebbe622793e1799e0a6bd23f", size = 520493, upload-time = "2026-08-07T02:12:42.971Z" } +sdist = { url = "https://files.pythonhosted.org/packages/64/24/95ec527ad67b76d59299e5465b3935d05e4294b7e0290a3924b7487df30b/tornado-6.5.7.tar.gz", hash = "sha256:66c513a76cda70d53907bc27cf1447557699c2e95aa48ba27a442ff61c3ddfc2", size = 519252, upload-time = "2026-06-08T17:34:51.232Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f2/d5/007086fd8df5489338e204f65adce33fd4f21a4999dbb2b9cff2f897b5f4/tornado-6.5.8-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:cc6aa787d7cfab7c3d35189dc7a56fbd2399a569624c730c6b55b3d6531d0403", size = 449487, upload-time = "2026-08-07T02:12:28.682Z" }, - { url = "https://files.pythonhosted.org/packages/70/c8/5a24a99495903f594f6a199dd7beead1cbc0a13e2cb9102727bcaaf2a997/tornado-6.5.8-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9715b5eb79735b2bcd454ce216a9275b7c0470e64ea1bf5742f78b2f72b26eeb", size = 447649, upload-time = "2026-08-07T02:12:30.306Z" }, - { url = "https://files.pythonhosted.org/packages/6e/de/f2e733f386b85962d1b1dc82cd63d169b5b4580062b35397eac9244a41fe/tornado-6.5.8-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:547d63f450d570c14fe0e8db2cfb14c9bbd1c2503b4a6612586267955aa47b58", size = 450707, upload-time = "2026-08-07T02:12:31.95Z" }, - { url = "https://files.pythonhosted.org/packages/0b/94/20efeee9a01c141e9ac47c397f81679dfda24b32768fc4fff24e76d36c2c/tornado-6.5.8-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e2360a0ffbe145eca8af0b19cb7203d79b1a98dd4cccdd6b368f6f49c2e3808", size = 451677, upload-time = "2026-08-07T02:12:33.512Z" }, - { url = "https://files.pythonhosted.org/packages/42/ec/a96ccb8ccf0de2b7bc2c5fa1608a4803735018242e90c4882365a9fd418f/tornado-6.5.8-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:5d242290bdf7ab3151bc1065fdd75c0dcc21cbc7b49f22a4c56329c2d6566d22", size = 451510, upload-time = "2026-08-07T02:12:35.346Z" }, - { url = "https://files.pythonhosted.org/packages/29/b5/93185859245ad3f00e62175f29607346788b696369347f0146e0421286bb/tornado-6.5.8-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7b94ff0e128fe0542f3bd331fb44d06260fc4ac16881545159f34ef08aad4195", size = 450917, upload-time = "2026-08-07T02:12:36.963Z" }, - { url = "https://files.pythonhosted.org/packages/97/cf/fe33cf062834487d34d1559746a4a12521033c22645b6d74d4bca702e018/tornado-6.5.8-cp39-abi3-win32.whl", hash = "sha256:67832909c4779c64942380cb5f044a5c6163d00831472d80e25e115de9917836", size = 451952, upload-time = "2026-08-07T02:12:38.512Z" }, - { url = "https://files.pythonhosted.org/packages/cb/e1/468ad54333e92ccb62627e62cb88e5fc14a2171daa67ed47b1b8542d5b86/tornado-6.5.8-cp39-abi3-win_amd64.whl", hash = "sha256:11881db6b7c168494be2c2d12e65931451bdf7ee718535418ae1d8855dd5a0ee", size = 452391, upload-time = "2026-08-07T02:12:39.971Z" }, - { url = "https://files.pythonhosted.org/packages/ad/3e/cd5e4f06e34cde33b8ef66cf36aa2b5ad46354cc1af7d2136bbe365fee1d/tornado-6.5.8-cp39-abi3-win_arm64.whl", hash = "sha256:68a7468c7e289f8514d7d664101753903217eff1bb6822c6b5994a0b5f5bcb26", size = 451411, upload-time = "2026-08-07T02:12:41.469Z" }, + { url = "https://files.pythonhosted.org/packages/02/dc/c7043cab6fed8ae159fc1923ce829ada35c4dbd797d408a43858ffaf9639/tornado-6.5.7-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:148b2eb15c2c765a50796172c1e499649b35f30d2e3c3d3e15913cfa56bfb163", size = 448543, upload-time = "2026-06-08T17:34:38.052Z" }, + { url = "https://files.pythonhosted.org/packages/92/4f/090b1431e5a43df696feceffc268c5383cc079ecb5f08ce58f917109aafe/tornado-6.5.7-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9da38de27f1da3b78a966f0dae12b5a1ea9afe72ca805d84ff06508272ddf100", size = 446707, upload-time = "2026-06-08T17:34:39.594Z" }, + { url = "https://files.pythonhosted.org/packages/37/d8/ef374952fd5da67d4463122c2b8e5a96536ec10b4b339254c6dcde81d01c/tornado-6.5.7-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8d759e71906ee783f8867b93bf26a265743da4c1e2f4a018464c1ba019862972", size = 449774, upload-time = "2026-06-08T17:34:41.204Z" }, + { url = "https://files.pythonhosted.org/packages/35/37/d434c73f4c6e014b745b9b37085f34f40c022f007efff3d7fe65991899f3/tornado-6.5.7-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a46347a18f23fb92b396beebe0fb78f61dda0cc302445202c16203d8a18848b", size = 450745, upload-time = "2026-06-08T17:34:42.531Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2b/56b9aff361d7f1ab728a805ec7d7ea835f8807afa9f5cc690ea0e630efb9/tornado-6.5.7-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7778b30bef919231265e91c69963ce0f49a1e9c07ac900bbe75b19ce2575ba92", size = 450578, upload-time = "2026-06-08T17:34:43.787Z" }, + { url = "https://files.pythonhosted.org/packages/02/30/a7444fb23aa76860a14198fab96ac79f1866b0a6e19e26c4381b0938e50f/tornado-6.5.7-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e726f0c75da7726eec023aa62751ff8878bd2737e34fbdd33b1ae5897d2200f5", size = 449985, upload-time = "2026-06-08T17:34:45.326Z" }, + { url = "https://files.pythonhosted.org/packages/5c/42/5f0e56c01e8d9d36f4e23f367b85ae6cae0c1ecddd5e6977d8388ad27488/tornado-6.5.7-cp39-abi3-win32.whl", hash = "sha256:f8de3bf12d3efdd0cbe7c8887868198f8a91415e3f29fcf258d9b8eb7b1d9ae4", size = 451047, upload-time = "2026-06-08T17:34:46.784Z" }, + { url = "https://files.pythonhosted.org/packages/c9/a4/b393076ffb21b469eec5b328a0534cf03a3b90bfc6b1f09507cdd075d938/tornado-6.5.7-cp39-abi3-win_amd64.whl", hash = "sha256:de942f843533a039ef9fa3d9c88c7cd8a7c94553fb5ad0154270989b3d99a2c4", size = 451485, upload-time = "2026-06-08T17:34:48.248Z" }, + { url = "https://files.pythonhosted.org/packages/71/2e/7b1c769803121b809112cf9a00681c472eae1d80e32d7ec0e0bd61d0d0e1/tornado-6.5.7-cp39-abi3-win_arm64.whl", hash = "sha256:ff934fce95643af5f11efdae618eaa73d469dc588641e5c8d19295a0c65c4796", size = 450506, upload-time = "2026-06-08T17:34:49.702Z" }, ] [[package]] From 5561b8438c34054b3fd1475e011ed5cfd6f6ec97 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:18:16 -0700 Subject: [PATCH 48/69] fix(responses): keep a namespace's non-function members when every function member is dropped --- .../responses/guardrail_translation/tool_merge.py | 4 ++-- .../test_openai_responses_guardrail_tool_merge.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index 9fbf9f31a0f..9a34946838a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -147,11 +147,11 @@ def _merged_original( guardrailed_group: Final = tuple(guardrailed_by_key[key] for key in group_keys if key in guardrailed_by_key) if guardrailed_group == tuple(flattened_group): return (original,) - if not guardrailed_group: - return () members: Final = _namespace_members(original) if original.get("type") == "namespace" else () if members and sum(map(_is_function, members)) == len(flattened_group): return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key) + if not guardrailed_group: + return () return tuple( LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools(guardrailed_group) ) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py index 4075c209606..cbcddb98235 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -132,6 +132,19 @@ def test_namespace_keeps_a_non_function_member_when_a_function_member_is_edited( assert merged[0]["tools"][1] == custom_member +def test_namespace_keeps_its_non_function_members_when_every_function_member_is_dropped(): + custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}} + original = [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read"), custom_member]}, + _function("a"), + ] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, [groups[1][0]]) + + assert list(merged) == [{"type": "namespace", "name": "ns", "description": "NS", "tools": [custom_member]}, _function("a")] + + def test_member_extras_edited_by_the_guardrail_land_on_that_member(): original = [{"type": "namespace", "name": "ns", "tools": [_function("read")]}] groups = _groups(original) From e3f61ce1347286774bc7c5da301c2bda1d0b0a96 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 22:13:23 +0000 Subject: [PATCH 49/69] chore(techdebt): refresh lint budgets after staging merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 10 +++++----- type-discipline-budget.json | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 64a22209d86..a53d7976263 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15298 + "limit": 15296 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,14 +105,14 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38344 + "limit": 38341 }, "reportUnknownParameterType": { "limit": 19625 }, - "reportUnknownVariableType": { - "limit": 29877 - }, + "reportUnknownVariableType": { + "limit": 29873 + }, "reportUnnecessaryCast": { "limit": 111 }, diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 5823d83d0bd..5475fc848fe 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -6,7 +6,7 @@ "limit": 26777 }, "LIT003": { - "limit": 265 + "limit": 264 }, "LIT004": { "limit": 40 @@ -26,11 +26,11 @@ "LIT009": { "limit": 0 }, - "LIT010": { - "limit": 16494 - }, + "LIT010": { + "limit": 16492 + }, "LIT011": { - "limit": 5529 + "limit": 5527 }, "LIT012": { "limit": 4495 From 7cbc74399d683b1113a4236768ac8414e7e9f07d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:24:43 -0700 Subject: [PATCH 50/69] fix(responses): accept pydantic tool objects returned by guardrails --- .../guardrail_translation/tool_merge.py | 5 ++-- ...t_openai_responses_guardrail_tool_merge.py | 23 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index 9a34946838a..b596adfad6f 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -3,7 +3,7 @@ from itertools import accumulate, chain, groupby from types import MappingProxyType from typing import Final, TypeAlias -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.responses.litellm_completion_transformation.transformation import ( @@ -19,8 +19,9 @@ _CHAT_TOOL_TOP_LEVEL_KEYS: Final = frozenset({"type", "function"}) def _as_tool(value: object) -> Tool | None: + candidate: Final = value.model_dump(exclude_unset=True) if isinstance(value, BaseModel) else value try: - return _TOOL_ADAPTER.validate_python(value) + return _TOOL_ADAPTER.validate_python(candidate) except ValidationError: return None diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py index cbcddb98235..9c236d81f51 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -9,6 +9,7 @@ from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) +from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GuardrailToolParam def _groups(tools): @@ -166,6 +167,28 @@ def test_guardrail_output_is_read_once(): assert list(merged) == original +def test_pydantic_guardrail_tools_round_trip_like_dicts(): + original = [_function("a"), {"type": "namespace", "name": "ns", "tools": [_function("x", "X doc")]}] + groups = _groups(original) + models = [GuardrailToolParam.model_validate(chat_tool) for chat_tool in _flat(groups)] + + merged = merge_guardrailed_tools(original, groups, models) + + assert list(merged) == original + assert all(merged_tool is original_tool for merged_tool, original_tool in zip(merged, original)) + + +def test_pydantic_guardrail_edit_lands_on_the_member(): + original = [{"type": "namespace", "name": "ns", "tools": [_function("x", "X doc")]}] + groups = _groups(original) + edited = copy.deepcopy(_flat(groups)) + edited[0]["function"]["description"] = "EDITED" + + merged = merge_guardrailed_tools(original, groups, [GuardrailToolParam.model_validate(edited[0])]) + + assert list(merged) == [{"type": "namespace", "name": "ns", "tools": [_function("x", "EDITED")]}] + + def test_non_object_guardrail_items_are_dropped(): original = [_function("a")] groups = _groups(original) From 979c3cb1aa0dc579c38d10f44839b93ecc0d9823 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:29:10 -0700 Subject: [PATCH 51/69] fix(proxy-extras): retry when the advisory lock wait times out behind a concurrent migrate deploy --- .../litellm_proxy_extras/utils.py | 9 ++++++ .../tests/test_setup_database_fail_fast.py | 28 +++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 9cc33e6ca56..d22484bc0e8 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -905,6 +905,15 @@ class ProxyExtrasDBManager: time.sleep(random.randrange(5, 15)) continue + if "P1002" in stderr and "advisory lock" in stderr: + logger.info( + "prisma migrate deploy attempt %s timed out waiting for " + "the advisory lock a concurrent migrate deploy holds, retrying", + attempt + 1, + ) + time.sleep(random.randrange(5, 15)) + continue + raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 70e0f216338..406f07eb792 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -432,3 +432,31 @@ def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path): ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True + + +_P1002_ADVISORY_LOCK_STDERR = ( + "Error: P1002\n\n" + "The database server at `127.0.0.1`:`45743` was reached but timed out.\n\n" + "Context: Timed out trying to acquire a postgres advisory lock " + "(SELECT pg_advisory_lock(72707369)). Elapsed: 10000ms." +) + + +def test_v2_advisory_lock_timeout_retries(monkeypatch, tmp_path): + """v2: the advisory-lock waiter that times out while a peer's retry holds + the lock retries instead of dying.""" + _stub_v2_env(monkeypatch, tmp_path) + monkeypatch.setattr("subprocess.run", _succeed_after(2, _P1002_ADVISORY_LOCK_STDERR)) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + + +def test_v2_p1002_without_advisory_lock_context_still_raises(monkeypatch, tmp_path): + """v2: a plain P1002 (database unreachable) stays fatal.""" + _stub_v2_env(monkeypatch, tmp_path) + stderr = "Error: P1002\n\nThe database server at `db`:`5432` was reached but timed out." + monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) From 6265437592e9d2401cae1149b9747e2eb2082572 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 22:35:02 +0000 Subject: [PATCH 52/69] chore(techdebt): avoid new mutable literal in gigachat tool conversion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 6 +++--- litellm/llms/gigachat/chat/transformation.py | 8 +++++++- type-discipline-budget.json | 6 +++--- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 449ed5bfa1e..e03ed8abdca 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15296 + "limit": 15292 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38341 + "limit": 38335 }, "reportUnknownParameterType": { "limit": 19625 }, "reportUnknownVariableType": { - "limit": 29873 + "limit": 29865 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 991a93ccb21..89920ebd27b 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -10,6 +10,7 @@ import json import time import uuid from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx @@ -34,6 +35,9 @@ else: LiteLLMLoggingObj = Any +_EMPTY_FUNCTION: Final[Mapping[str, object]] = MappingProxyType({}) + + def is_valid_json(value: str) -> bool: """Checks whether the value passed is a valid serialized JSON string""" try: @@ -213,7 +217,9 @@ class GigaChatConfig(BaseConfig): "parameters": function.get("parameters", {}), } for function in ( - tool.get("function", {}) for tool in tools if isinstance(tool, dict) and tool.get("type") == "function" + tool.get("function", _EMPTY_FUNCTION) + for tool in tools + if isinstance(tool, dict) and tool.get("type") == "function" ) ] diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 5a4cae71d41..45d00a71d39 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -6,7 +6,7 @@ "limit": 26774 }, "LIT003": { - "limit": 264 + "limit": 262 }, "LIT004": { "limit": 40 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16492 + "limit": 16488 }, "LIT011": { - "limit": 5527 + "limit": 5523 }, "LIT012": { "limit": 4495 From ed0d3f94422e32c498781fe63722219989987285 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 15:21:37 -0700 Subject: [PATCH 53/69] fix(router): fail open on session binding reads --- litellm/router.py | 13 ++++++-- tests/test_litellm/test_router.py | 50 +++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index f4d912eb5be..2d8b4a43759 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12632,9 +12632,16 @@ class Router: async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: session_cache: Final = self._claude_code_session_router_cache - if session_cache.redis_cache is None: - return await session_cache.async_get_cache(key=cache_key) - return await session_cache.redis_cache.async_get_cache(key=cache_key) + try: + if session_cache.redis_cache is None: + return await session_cache.async_get_cache(key=cache_key) + return await session_cache.redis_cache.async_get_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # an optional binding must not make routing depend on Redis + verbose_router_logger.warning( + "Failed to read Claude Code session router binding; using the requested model: %s", + e, + ) + return None async def _resolve_claude_code_session_router( self, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 1e5767aa645..16c4d997ad7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8504,6 +8504,56 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_redis_read_failure_does_not_reject_a_subagent_request(self): + from litellm.caching.caching import RedisCache + + router = self._router() + request_kwargs = self._request_kwargs(agent_id="agent-1234") + cache_key = router._claude_code_session_router_cache_key(request_kwargs) + assert cache_key is not None + await router._claude_code_session_router_cache.in_memory_cache.async_set_cache( + cache_key, + "smart-router", + ) + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(side_effect=Exception("Redis circuit breaker is open")) + router._update_redis_cache(cache=redis_cache) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=request_kwargs, + ) + + assert response is None + assert "model_group" not in request_kwargs["metadata"] + redis_cache.async_get_cache.assert_awaited_once() + + @pytest.mark.asyncio + async def test_redis_write_failures_do_not_reject_main_or_subagent_requests(self): + from litellm.caching.caching import RedisCache + + router = self._router() + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(return_value="smart-router") + redis_cache.async_set_cache = AsyncMock(side_effect=Exception("redis unavailable")) + router._update_redis_cache(cache=redis_cache) + + main_response = await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=self._request_kwargs(), + ) + subagent_response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert main_response is not None + assert main_response.model == "cheap-model" + assert subagent_response is not None + assert subagent_response.model == "cheap-model" + assert redis_cache.async_set_cache.await_count == 2 + @pytest.mark.asyncio async def test_subagents_follow_the_main_threads_latest_router_across_workers(self): from types import SimpleNamespace From 1306a4505a03bdced3ccdf5ac6b26a719867bbbb Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 16:01:15 -0700 Subject: [PATCH 54/69] fix(router): preserve bound router fallbacks for subagents --- .../router_utils/fallback_event_handlers.py | 12 +++++----- .../test_fallback_event_handlers.py | 22 +++++++++++++++++++ tests/test_litellm/test_router.py | 21 ++++++++++++++++++ 3 files changed, 50 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 0167721f9fe..f7855cb38ff 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, @@ -231,8 +232,6 @@ def record_pre_routing_selection(request_kwargs: Mapping[str, Any] | None, selec on /v1/messages the top-level ``metadata`` dict is the provider's own request field, so a blanket write would forward the tier stamp upstream. """ - from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs - if request_kwargs is None: return bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)) @@ -267,10 +266,13 @@ def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None: def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]: """ Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins, - and the requested group still resolves when no tier-keyed chain exists, so configs keyed - on the router name (the documented contract) keep working behind auto-routers. + then the routed group, then the requested group. The routed group differs when Claude Code + session affinity remaps a subagent's concrete model to its bound router. """ - ordered: Final = (get_pre_routing_selection(kwargs), model_group) + metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) + routed_group_value: Final = metadata.get("model_group") if isinstance(metadata, Mapping) else None + routed_group: Final = routed_group_value if isinstance(routed_group_value, str) else None + ordered: Final = (get_pre_routing_selection(kwargs), routed_group, model_group) return tuple(dict.fromkeys(group for group in ordered if group)) diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 9e51a60364b..a4965c49f07 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -1220,6 +1220,28 @@ class TestOrderedFallbackLookupGroups: assert fallback_lookup_groups({}, "smart-router") == ("smart-router",) assert fallback_lookup_groups({}, None) == () + def test_session_remap_keeps_the_bound_router_between_tier_and_requested_group(self): + from litellm.router_utils.fallback_event_handlers import ( + PRE_ROUTING_SELECTED_MODEL_KEY, + fallback_lookup_groups, + ) + + kwargs = { + "litellm_metadata": { + PRE_ROUTING_SELECTED_MODEL_KEY: "tier1", + "model_group": "smart-router", + } + } + + assert fallback_lookup_groups(kwargs, "requested-model") == ( + "tier1", + "smart-router", + "requested-model", + ) + assert fallback_lookup_groups({"metadata": {"model_group": []}}, "requested-model") == ( + "requested-model", + ) + def test_first_resolving_group_wins_and_generic_idx_survives_a_miss(self): from litellm.router_utils.fallback_event_handlers import ( get_fallback_model_group_for_lookup_groups, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 16c4d997ad7..0a433d0b0a1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8679,6 +8679,27 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response.choices[0].message.content == "expensive response" assert subagent_kwargs["metadata"]["routing_decision"]["routed_model"] == "cheap-model" + @pytest.mark.asyncio + async def test_subagent_can_use_the_bound_router_name_fallback(self): + router = self._router( + cheap_response="litellm.RateLimitError", + fallbacks=[{"smart-router": ["expensive-model"]}], + ) + + await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "main turn"}], + **self._request_kwargs(), + ) + + response = await router.acompletion( + model="expensive-model", + messages=[{"role": "user", "content": "subagent turn"}], + **self._request_kwargs(agent_id="agent-1234"), + ) + + assert response.choices[0].message.content == "expensive response" + @pytest.mark.asyncio async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From ddc25823740e8ed36a8d5ccddb93c4dee9ff95bc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:15:20 -0700 Subject: [PATCH 55/69] fix(vector-store): build the Router executor from the router kwarg and gate the SDK fallback on explicit config --- .../base_llm/vector_store/transformation.py | 7 +++-- litellm/vector_stores/main.py | 16 +++++------- .../test_router_embedding_integration.py | 19 ++++++++++++++ .../test_vector_store_endpoints.py | 2 +- .../test_milvus_vector_store.py | 26 +++++++++++++++++++ 5 files changed, 58 insertions(+), 12 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index e7c549b7358..3547bf81448 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -86,9 +86,12 @@ 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 not self._router_serves(model): + if self._embeds_through_sdk(model, configuration): return LiteLLMVectorStoreEmbeddingExecutor().embed(model, query, embedding_kwargs) return self.router.embedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, @@ -98,7 +101,7 @@ class RouterVectorStoreEmbeddingExecutor: async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: embedding_kwargs: Final = self._embedding_kwargs(configuration) - if not self._router_serves(model): + if self._embeds_through_sdk(model, configuration): 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/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index ade4d19815b..d530a1c9d52 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -16,7 +16,7 @@ from litellm.constants import request_timeout from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.vector_store.transformation import ( - LiteLLMVectorStoreEmbeddingExecutor, + BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler @@ -42,12 +42,10 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# -def _direct_vector_store_embedding_executor(value: object) -> VectorStoreEmbeddingExecutor: - if value is None: - return LiteLLMVectorStoreEmbeddingExecutor() - if isinstance(value, VectorStoreEmbeddingExecutor): - return value - raise TypeError("Invalid direct vector store embedding executor") +def _direct_vector_store_embedding_executor(value: object, router: "Router | None") -> VectorStoreEmbeddingExecutor: + if value is not None and not isinstance(value, VectorStoreEmbeddingExecutor): + raise TypeError("Invalid direct vector store embedding executor") + return BaseQueryEmbeddingVectorStoreConfig.query_embedding_executor(value, router) def mock_vector_store_search_response( @@ -302,7 +300,7 @@ async def asearch( Async: Search a vector store for relevant chunks based on a query and file attributes filter. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None) + kwargs.pop("_direct_vector_store_embedding_executor", None), router ) local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot key: value for key, value in locals().items() if key != "embedding_executor" @@ -388,7 +386,7 @@ def search( VectorStoreSearchResponse containing the search results. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None) + kwargs.pop("_direct_vector_store_embedding_executor", None), router ) local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot key: value for key, value in locals().items() if key != "embedding_executor" diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index 96c6ce3708e..2cc9914c9b3 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -186,6 +186,25 @@ class TestRouterEmbeddingIntegration: assert _sent(store_route, 0) == ("Bearer store-key", "text-embedding-3-large", ["sync query"]) 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( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + openai_route = _mock_embedding_route(respx_mock, OPENAI_EMBEDDINGS_URL) + executor = RouterVectorStoreEmbeddingExecutor( + router=_alias_router(), + 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", {}) + + assert openai_route.call_count == 0 + 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/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 854ce351f2d..a956aef51d6 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -52,7 +52,7 @@ def _serialize_litellm_params(litellm_params): def test_direct_vector_store_embedding_executor_rejects_invalid_value(): with pytest.raises(TypeError, match="Invalid direct vector store embedding executor"): - _direct_vector_store_embedding_executor(object()) + _direct_vector_store_embedding_executor(object(), None) def test_router_vector_store_search_injects_executor_and_request_metadata(): diff --git a/tests/vector_store_tests/test_milvus_vector_store.py b/tests/vector_store_tests/test_milvus_vector_store.py index ea3c1883e46..3b3bd444b20 100644 --- a/tests/vector_store_tests/test_milvus_vector_store.py +++ b/tests/vector_store_tests/test_milvus_vector_store.py @@ -575,6 +575,32 @@ async def test_router_search_resolves_bare_embedding_alias_async( _assert_alias_resolved(embedding_route, search_route, response) +def test_sdk_search_with_router_kwarg_resolves_bare_embedding_alias_sync( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = litellm.vector_stores.search(router=_alias_router(), **ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_sdk_search_with_router_kwarg_resolves_bare_embedding_alias_async( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = await litellm.vector_stores.asearch(router=_alias_router(), **ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + @pytest.mark.asyncio async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE) From f5721ff6f063cd6cfed4285a4bee80e6cf1fc81b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 16:17:35 -0700 Subject: [PATCH 56/69] test(router): cover Claude session multi-hop fallbacks --- tests/test_litellm/test_router.py | 59 +++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 0a433d0b0a1..db73f0da3c2 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8700,6 +8700,65 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response.choices[0].message.content == "expensive response" + @pytest.mark.asyncio + async def test_anthropic_subagent_four_fallback_hops_use_each_current_model_chain(self): + from litellm.types.router import TaggedPreRoutingStrategy + + failing_groups = ("cheap-model", "fallback-1", "fallback-2", "fallback-3") + router = litellm.Router( + model_list=[ + *( + { + "model_name": group, + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "mock_response": "litellm.RateLimitError", + }, + } + for group in failing_groups + ), + { + "model_name": "requested-model", + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "mock_response": "requested response", + }, + }, + { + "model_name": "fallback-4", + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "mock_response": "fourth fallback response", + }, + }, + ], + fallbacks=[ + {"smart-router": ["fallback-1"]}, + {"fallback-1": ["fallback-2"]}, + {"fallback-2": ["fallback-3"]}, + {"fallback-3": ["fallback-4"]}, + ], + num_retries=0, + max_fallbacks=4, + ) + router.complexity_routers = { + "smart-router": (TaggedPreRoutingStrategy(tags=(), strategy=self._RewriteStrategy()),) + } + main_kwargs = self._request_kwargs() + main_kwargs["litellm_metadata"] = main_kwargs.pop("metadata") + await router.async_pre_routing_hook(model="smart-router", request_kwargs=main_kwargs) + subagent_kwargs = self._request_kwargs(agent_id="agent-1234") + subagent_kwargs["litellm_metadata"] = subagent_kwargs.pop("metadata") + + response = await router.aanthropic_messages( + model="requested-model", + messages=[{"role": "user", "content": "subagent turn"}], + max_tokens=64, + **subagent_kwargs, + ) + + assert response["content"][0]["text"] == "fourth fallback response" + @pytest.mark.asyncio async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From 683fc340440f4c6003075ec7fe42451e0875e734 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:23:52 -0700 Subject: [PATCH 57/69] fix(azure_ai): let the caller's output_config from extra_body win over the legacy thinking upgrade --- .../llms/azure_ai/anthropic/transformation.py | 7 +++-- .../test_azure_anthropic_transformation.py | 30 +++++++++++++++++-- 2 files changed, 31 insertions(+), 6 deletions(-) diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index c5053448627..864d2134a84 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/litellm/llms/azure_ai/anthropic/transformation.py @@ -17,13 +17,14 @@ def _promote_extra_body_to_optional_params(optional_params: dict) -> None: ``output_config`` get auto-routed into ``extra_body`` by ``add_provider_specific_params_to_optional_params``. For the Azure→Anthropic route those keys must reach the request body and be validated, so promote - them. ``setdefault`` keeps explicit top-level values authoritative. + them. The caller's values overwrite mapped top-level duplicates, matching + the native ``anthropic`` provider, where the same passthrough lands on + top-level ``optional_params`` after mapping. """ extra_body: Final = optional_params.get("extra_body") if not isinstance(extra_body, dict) or not extra_body: return - for k, v in extra_body.items(): - optional_params.setdefault(k, v) + optional_params.update(extra_body) optional_params.pop("extra_body", None) diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py index 9dac914ca4d..9e2bfb08852 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py @@ -362,8 +362,8 @@ class TestAzureAnthropicConfig: ) assert "xhigh" in str(exc_info.value) - def test_extra_body_promotion_does_not_clobber_top_level(self): - """Top-level ``optional_params`` wins over duplicates in ``extra_body``.""" + def test_extra_body_promotion_overrides_mapped_top_level(self): + """The caller's ``extra_body`` wins over a mapped top-level duplicate, like the native ``anthropic`` passthrough.""" config = AzureAnthropicConfig() messages = [{"role": "user", "content": "Hello"}] @@ -383,7 +383,31 @@ class TestAzureAnthropicConfig: headers=headers, ) - assert result["output_config"] == {"effort": "low"} + assert result["output_config"] == {"effort": "high"} + + def test_legacy_thinking_upgrade_keeps_caller_effort_from_extra_body(self, local_model_cost_map): + config = AzureAnthropicConfig() + + mapped = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 1024}, "max_tokens": 100}, + optional_params={}, + model="claude-opus-4-8", + drop_params=False, + ) + assert mapped["thinking"] == {"type": "adaptive"} + assert mapped["output_config"] == {"effort": "low"} + + result = config.transform_request( + model="claude-opus-4-8", + messages=[{"role": "user", "content": "Hello"}], + optional_params={**mapped, "extra_body": {"output_config": {"effort": "high"}}}, + litellm_params={"api_key": "test-key"}, + headers={"api-key": "test-key", "anthropic-version": "2023-06-01"}, + ) + + assert result["thinking"] == {"type": "adaptive"} + assert result["output_config"] == {"effort": "high"} + assert "extra_body" not in result def test_context_management_mixed_edits_beta_headers(self): """Test that context_management with both compact and other edits adds both beta headers""" From 63f5de12ae4e21f4cc49bf2ac23a89dc077b1854 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:25:58 -0700 Subject: [PATCH 58/69] chore(techdebt): ratchet lint budgets to the merged tree --- basedpyright-code-budget.json | 6 +++--- type-discipline-budget.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index e03ed8abdca..7d3c83a567f 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15292 + "limit": 15290 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38335 + "limit": 38332 }, "reportUnknownParameterType": { "limit": 19625 }, "reportUnknownVariableType": { - "limit": 29865 + "limit": 29861 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 45d00a71d39..091c1f46e02 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -6,7 +6,7 @@ "limit": 26774 }, "LIT003": { - "limit": 262 + "limit": 261 }, "LIT004": { "limit": 40 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16488 + "limit": 16486 }, "LIT011": { - "limit": 5523 + "limit": 5521 }, "LIT012": { "limit": 4495 From 198906495f7de129e9eaa5aae55c4304b7dca99a Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 2 Sep 2026 16:26:35 -0700 Subject: [PATCH 59/69] refactor(python-bridge): split routes and add shared function tracing (#39031) * refactor(python-bridge): split non-streaming bridge modules * refactor(python-bridge): bring shared function tracing into route layer * feat(dev): list Python route functions and call sites * feat(dev): list Rust route functions and call sites * docs(dev): record OCR parity gaps across Python and Rust * feat(dev): list executed SDK calls with runtime tracing * feat(dev): report Python vs Rust SDK pipeline steps in one CLI * feat(dev): side-by-side pipeline step report in compare CLI * fix(dev): drop invalid Final annotations in compare cell loop * feat(dev): blue python-only and yellow rust-only steps in compare CLI * feat(dev): vertical layout with section spacing in compare CLI * fix(dev): validate SDK trace stages across sync and async routes * refactor(rust): align SDK route call structure with Python * refactor(python-bridge): share sync and async route call wrappers * refactor(dev): split compare CLI into fixtures, runtime, and report modules * fix(ci): run SDK trace tests and satisfy test lint --- .github/workflows/test-unit.yml | 1 + litellm-rust/Cargo.lock | 40 ++ litellm-rust/Cargo.toml | 2 + litellm-rust/crates/ai-gateway/Cargo.toml | 1 + .../crates/ai-gateway/src/ocr/common_utils.rs | 7 +- .../crates/ai-gateway/src/ocr/handler.rs | 14 +- .../crates/ai-gateway/src/ocr/hooks.rs | 49 +- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 5 +- .../crates/ai-gateway/src/ocr/prepare.rs | 20 +- .../crates/ai-gateway/src/ocr/tests.rs | 3 +- .../crates/ai-gateway/src/ocr/types.rs | 1 + litellm-rust/crates/core/Cargo.toml | 1 + .../core/src/audio_transcription/handler.rs | 6 +- .../core/src/audio_transcription/mod.rs | 1 + .../core/src/audio_transcription/prepare.rs | 2 + .../src/audio_transcription/transformation.rs | 1 + .../core/src/chat_completions/common_utils.rs | 1 + .../core/src/chat_completions/handler.rs | 10 +- .../crates/core/src/chat_completions/mod.rs | 5 +- .../core/src/chat_completions/prepare.rs | 44 +- .../crates/core/src/chat_completions/tests.rs | 10 +- .../src/chat_completions/transformation.rs | 13 +- .../crates/core/src/chat_completions/types.rs | 11 + litellm-rust/crates/core/src/http_utils.rs | 7 + .../crates/core/src/messages/common_utils.rs | 1 + .../crates/core/src/messages/handler.rs | 17 +- litellm-rust/crates/core/src/messages/mod.rs | 6 +- .../crates/core/src/messages/prepare.rs | 63 +- .../core/src/messages/transformation.rs | 2 + .../crates/core/src/ocr/transformation.rs | 20 + .../chat_completions/transformation.rs | 14 +- .../anthropic/messages/transformation.rs | 1 + .../azure_ai/messages/transformation.rs | 1 + .../providers/bedrock/audio_transcription.rs | 3 + .../chat_completions/transformation.rs | 72 +- .../providers/mistral/ocr/transformation.rs | 7 + litellm-rust/crates/python-bridge/Cargo.toml | 3 + .../crates/python-bridge/src/constants.rs | 1 + .../crates/python-bridge/src/diagnostics.rs | 23 + .../crates/python-bridge/src/errors.rs | 61 ++ .../python-bridge/src/function_trace.rs | 216 ++++++ litellm-rust/crates/python-bridge/src/lib.rs | 623 ++---------------- .../crates/python-bridge/src/marshal.rs | 63 ++ .../src/routes/audio_transcription.rs | 126 ++++ .../src/routes/chat_completions.rs | 156 +++++ .../python-bridge/src/routes/messages.rs | 122 ++++ .../crates/python-bridge/src/routes/mod.rs | 53 ++ .../crates/python-bridge/src/routes/ocr.rs | 128 ++++ tests/sdk_function_trace/README.md | 30 + tests/sdk_function_trace/__init__.py | 13 + tests/sdk_function_trace/compare.py | 45 ++ tests/sdk_function_trace/fixtures.py | 200 ++++++ tests/sdk_function_trace/harness.py | 39 ++ tests/sdk_function_trace/mock_provider.py | 67 ++ tests/sdk_function_trace/ocr-comparison.md | 59 ++ tests/sdk_function_trace/profiler.py | 76 +++ tests/sdk_function_trace/report.py | 175 +++++ tests/sdk_function_trace/route-comparison.md | 26 + tests/sdk_function_trace/runtime.py | 128 ++++ tests/sdk_function_trace/steps.py | 181 +++++ tests/sdk_function_trace/table.py | 72 ++ .../sdk_function_trace/test_mock_provider.py | 33 + tests/sdk_function_trace/test_profiler.py | 167 +++++ tests/sdk_function_trace/test_runtime.py | 47 ++ tests/sdk_function_trace/test_steps.py | 244 +++++++ tests/sdk_function_trace/test_table.py | 67 ++ 66 files changed, 2993 insertions(+), 713 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/constants.rs create mode 100644 litellm-rust/crates/python-bridge/src/diagnostics.rs create mode 100644 litellm-rust/crates/python-bridge/src/errors.rs create mode 100644 litellm-rust/crates/python-bridge/src/function_trace.rs create mode 100644 litellm-rust/crates/python-bridge/src/marshal.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/chat_completions.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/ocr.rs create mode 100644 tests/sdk_function_trace/README.md create mode 100644 tests/sdk_function_trace/__init__.py create mode 100644 tests/sdk_function_trace/compare.py create mode 100644 tests/sdk_function_trace/fixtures.py create mode 100644 tests/sdk_function_trace/harness.py create mode 100644 tests/sdk_function_trace/mock_provider.py create mode 100644 tests/sdk_function_trace/ocr-comparison.md create mode 100644 tests/sdk_function_trace/profiler.py create mode 100644 tests/sdk_function_trace/report.py create mode 100644 tests/sdk_function_trace/route-comparison.md create mode 100644 tests/sdk_function_trace/runtime.py create mode 100644 tests/sdk_function_trace/steps.py create mode 100644 tests/sdk_function_trace/table.py create mode 100644 tests/sdk_function_trace/test_mock_provider.py create mode 100644 tests/sdk_function_trace/test_profiler.py create mode 100644 tests/sdk_function_trace/test_runtime.py create mode 100644 tests/sdk_function_trace/test_steps.py create mode 100644 tests/sdk_function_trace/test_table.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index c2dff805772..6da5fc07e80 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -96,6 +96,7 @@ jobs: - shard: misc artifact-name: misc test-path: >- + tests/sdk_function_trace tests/test_litellm/batches tests/test_litellm/secret_managers tests/test_litellm/a2a_protocol diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index dd41cf0e84b..71c9bb95ec0 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1392,6 +1392,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.186" @@ -1416,6 +1422,7 @@ dependencies = [ "tokio", "tokio-tungstenite", "tower", + "tracing", ] [[package]] @@ -1435,6 +1442,7 @@ dependencies = [ "sha2 0.10.9", "thiserror 2.0.19", "tokio", + "tracing", ] [[package]] @@ -1447,8 +1455,11 @@ dependencies = [ "litellm-python-interop", "pyo3", "pyo3-async-runtimes", + "serde", "serde_json", "tokio", + "tracing", + "tracing-subscriber", ] [[package]] @@ -2276,6 +2287,15 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "2.0.1" @@ -2414,6 +2434,15 @@ dependencies = [ "syn 3.0.0", ] +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + [[package]] name = "time" version = "0.3.53" @@ -2662,6 +2691,17 @@ dependencies = [ "once_cell", ] +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "sharded-slab", + "thread_local", + "tracing-core", +] + [[package]] name = "try-lock" version = "0.2.5" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index c447d915abe..a096643d0d3 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -14,6 +14,8 @@ license = "MIT" repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] +tracing = "0.1" +tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } litellm-core = { path = "crates/core" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } litellm-python-interop = { path = "crates/python-interop" } diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 541beabe170..e3dbdf24ce6 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -14,6 +14,7 @@ path = "src/main.rs" required-features = ["server"] [dependencies] +tracing.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the # Python proxy callbacks API. diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index e0ce165dc93..c1fb328893b 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -32,6 +32,7 @@ pub(super) fn truncate_error_body(body: &str) -> String { format!("{truncated}... (truncated)") } +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) fn ocr_provider_config( provider: &str, model: &str, @@ -73,12 +74,6 @@ pub(super) fn string_headers( .collect() } -pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool { - headers - .iter() - .any(|(key, _)| key.eq_ignore_ascii_case(name)) -} - fn document_url_field(document: &Value) -> Result, Error> { let Some(object) = document.as_object() else { return Ok(None); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 815bc84363a..856d9571201 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -1,12 +1,19 @@ use litellm_core::error::Error; +use litellm_core::http_utils::http_request; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::Value; use super::common_utils::{poll_document_intelligence, truncate_error_body}; -use super::types::ProviderOcrRequest; +use super::hooks::OcrLifecycleHooks; +use super::types::PreparedOcrRequest; use crate::client::http_client; -pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Result { +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub(crate) async fn execute_ocr_provider_call( + request: PreparedOcrRequest, + hooks: &OcrLifecycleHooks, +) -> Result { + let request = hooks.prepare_provider_request(request).await?; let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -15,8 +22,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Re request_builder = request_builder.timeout(duration); } - let response = request_builder - .send() + let response = http_request(request_builder) .await .map_err(|err| Error::Network(err.to_string()))?; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 401e26d3b29..f8c4f8fe8c5 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -1,13 +1,10 @@ use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use litellm_core::error::Error; -use litellm_core::ocr::transformation::OcrAuthStrategy; use serde_json::{Map, Value, json}; use std::future::Future; use std::pin::Pin; -use super::common_utils::{ - convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers, -}; +use super::common_utils::{convert_document_url_to_data_uri, string_headers}; use super::types::{PreparedOcrRequest, ProviderOcrRequest}; use crate::integrations::custom_guardrail::{ CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, @@ -62,6 +59,10 @@ impl OcrLifecycleHooks { .await .map_err(guardrail_error_to_core_error)?; let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?; + let optional_params = match &request.config { + Ok(config) => config.map_ocr_params(&optional_params), + Err(_) => optional_params, + }; Ok(PreparedOcrRequest { document, optional_params, @@ -69,25 +70,23 @@ impl OcrLifecycleHooks { }) } - async fn prepare_provider_request( + pub(crate) async fn prepare_provider_request( &self, request: PreparedOcrRequest, ) -> Result { - let config = ocr_provider_config(&request.custom_llm_provider, &request.model) - .ok_or_else(|| Error::InvalidProvider(request.custom_llm_provider.clone()))?; + let config = request.config?; let env_lookup = |key: &str| std::env::var(key).ok(); - let headers = string_headers(request.extra_headers)?; - let auth_strategy = config.auth_strategy(); - let api_key = (!has_header(&headers, auth_strategy.header_name())) - .then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup)) - .transpose()?; + let upstream_headers = config.validate_environment( + string_headers(request.extra_headers)?, + request.api_key.as_deref(), + &env_lookup, + )?; let url = config.complete_url( request.api_base.as_deref(), &request.model, &request.optional_params, &env_lookup, )?; - let filtered_params = config.map_ocr_params(&request.optional_params); let model = request.model.clone(); let custom_llm_provider = request.custom_llm_provider.clone(); let document = if config.requires_data_uri_document() { @@ -96,9 +95,8 @@ impl OcrLifecycleHooks { request.document }; let body = config - .transform_ocr_request(&request.model, document, filtered_params)? + .transform_ocr_request(&request.model, document, request.optional_params)? .data; - let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref()); let body = self .run_during_call_guardrails(&model, &custom_llm_provider, &url, body) .await?; @@ -167,9 +165,9 @@ impl OcrLifecycleHooks { } } -impl CallLifecycleHooks for OcrLifecycleHooks { +impl CallLifecycleHooks for OcrLifecycleHooks { type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; - type DuringCallFuture<'a> = OcrFuture<'a, ProviderOcrRequest>; + type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; type SuccessFuture<'a> = OcrLogFuture<'a>; type FailureFuture<'a> = OcrLogFuture<'a>; @@ -186,7 +184,7 @@ impl CallLifecycleHooks for OcrLi _context: &'a CallLifecycleContext, request: PreparedOcrRequest, ) -> Self::DuringCallFuture<'a> { - Box::pin(async move { self.prepare_provider_request(request).await }) + Box::pin(async move { Ok(request) }) } fn async_log_success_event<'a>( @@ -247,21 +245,6 @@ impl CallLifecycleHooks for OcrLi } } -fn upstream_headers( - headers: &[(String, String)], - auth_strategy: OcrAuthStrategy, - api_key: Option<&str>, -) -> Vec<(String, String)> { - api_key - .map(|api_key| match auth_strategy { - OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), - OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()), - }) - .into_iter() - .chain(headers.iter().cloned()) - .collect() -} - fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { GuardrailContext { call_type: CallType::Ocr, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index b59ab626fd3..d9230af1c59 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -13,10 +13,13 @@ pub use types::OcrRequest; use handler::execute_ocr_provider_call; use prepare::{PreparedOcrCall, prepare_ocr_call}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn ocr(request: OcrRequest<'_>) -> Result { let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); CallLifecycle::default() - .run_request(request, &hooks, execute_ocr_provider_call) + .run_request(request, &hooks, |request| { + execute_ocr_provider_call(request, &hooks) + }) .await } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs index 6231393c889..fedacc62760 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -3,6 +3,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; +use super::common_utils::ocr_provider_config; use super::hooks::OcrLifecycleHooks; use super::types::{OcrRequest, PreparedOcrRequest}; use crate::integrations::custom_guardrail::CustomGuardrailRunner; @@ -13,6 +14,7 @@ pub(crate) struct PreparedOcrCall { pub(crate) hooks: OcrLifecycleHooks, } +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { let call_id = request .litellm_call_id @@ -25,9 +27,25 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { }); let model = provider_info.model.to_string(); let custom_llm_provider = provider_info.custom_llm_provider.to_string(); + let config = ocr_provider_config(&custom_llm_provider, &model) + .ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone())); + let optional_params = match &config { + Ok(config) => { + let supported = config.supported_ocr_params(); + config.map_ocr_params( + &request + .optional_params + .into_iter() + .filter(|(name, _)| supported.contains(&name.as_str())) + .collect(), + ) + } + Err(_) => request.optional_params, + }; PreparedOcrCall { request: PreparedOcrRequest { + config, model, custom_llm_provider, litellm_call_id: call_id, @@ -35,7 +53,7 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { api_key: request.api_key.map(str::to_string), api_base: request.api_base.map(str::to_string), extra_headers: request.extra_headers, - optional_params: request.optional_params, + optional_params, timeout: request.timeout, }, hooks: OcrLifecycleHooks::new( diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index 8c3f0425149..85e4c408045 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -2,12 +2,13 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use litellm_core::error::Error; +use litellm_core::http_utils::has_header; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body}; +use super::common_utils::{ocr_provider_config, string_headers, truncate_error_body}; use super::{OcrRequest, ocr}; use crate::integrations::custom_guardrail::{ CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs index bde734a4dd1..95e551d79ca 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/types.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -25,6 +25,7 @@ pub struct OcrRequest<'a> { } pub(crate) struct PreparedOcrRequest { + pub(crate) config: Result<&'static dyn OcrProviderConfig, litellm_core::Error>, pub(crate) model: String, pub(crate) custom_llm_provider: String, pub(crate) litellm_call_id: String, diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index ab8050734f2..389dbd49505 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -11,6 +11,7 @@ reqwest.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true +tracing.workspace = true sha2.workspace = true aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true } diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 30ba0da5e68..9a96b9d1140 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,11 +1,12 @@ use serde_json::Value; use crate::error::Error; -use crate::http_utils::truncate_error_body; +use crate::http_utils::{http_request, truncate_error_body}; use super::client::http_client; use super::types::ProviderAudioTranscriptionRequest; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn execute_audio_transcription_provider_call( request: ProviderAudioTranscriptionRequest, ) -> Result { @@ -19,8 +20,7 @@ pub async fn execute_audio_transcription_provider_call( if let Some(duration) = request.timeout { request_builder = request_builder.timeout(duration); } - let response = request_builder - .send() + let response = http_request(request_builder) .await .map_err(|error| Error::Network(error.to_string()))?; let status = response.status(); diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index b71748082bf..31b6de4b3e4 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -11,6 +11,7 @@ pub use handler::execute_audio_transcription_provider_call; pub use prepare::prepare_audio_transcription_provider_call; pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?) .await diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 6288e96b380..bbef97341a9 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -7,6 +7,7 @@ use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider} use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig}; use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> { #[cfg(feature = "bedrock-auth")] if provider == "bedrock" { @@ -16,6 +17,7 @@ fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProv None } +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub fn prepare_audio_transcription_provider_call( request: AudioTranscriptionRequest<'_>, ) -> Result { diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/audio_transcription/transformation.rs index 16a28fbcac0..aa9846427dc 100644 --- a/litellm-rust/crates/core/src/audio_transcription/transformation.rs +++ b/litellm-rust/crates/core/src/audio_transcription/transformation.rs @@ -15,6 +15,7 @@ pub enum AudioTranscriptionAuth { pub trait AudioTranscriptionProviderConfig: Sync { fn supported_transcription_params(&self) -> &'static [&'static str]; + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn map_transcription_params(&self, params: &Map) -> Map { params .iter() diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs index ca51471eb7c..69e5f175ad5 100644 --- a/litellm-rust/crates/core/src/chat_completions/common_utils.rs +++ b/litellm-rust/crates/core/src/chat_completions/common_utils.rs @@ -7,6 +7,7 @@ use super::transformation::ChatCompletionsProviderConfig; const HEADER_CONTEXT: &str = "chat completions"; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) fn chat_completions_provider_config( provider: &str, ) -> Option<&'static dyn ChatCompletionsProviderConfig> { diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 7e2731442cc..96d001e2892 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,17 +1,21 @@ use serde_json::Value; use crate::error::Error; -use crate::http_utils::truncate_error_body; +use crate::http_utils::{http_request, truncate_error_body}; use super::client::http_client; +use super::prepare::prepare_provider_request; use super::transformation::ChatCompletionsAuth; use super::types::{ ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData, + ResolvedChatCompletionsRequest, }; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) async fn execute_chat_completions_provider_call( - request: ProviderChatCompletionsRequest, + request: ResolvedChatCompletionsRequest<'_>, ) -> Result { + let request = prepare_provider_request(request)?; let body = serde_json::to_vec(&request.body).map_err(|err| { Error::InvalidRequest(format!( "failed to serialize chat completions request: {err}" @@ -27,7 +31,7 @@ pub(super) async fn execute_chat_completions_provider_call( request_builder = request_builder.timeout(duration); } - let response = request_builder.send().await.map_err(|err| { + let response = http_request(request_builder).await.map_err(|err| { // Failing to establish the connection means the request never went out, // so the host can still serve it. Everything else here, a timeout // above all, may have reached the provider and been answered. diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 0d009d36d16..32dea17d202 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -19,13 +19,14 @@ pub mod types; use serde_json::{Map, Value}; use handler::execute_chat_completions_provider_call; -use prepare::{parse_messages, prepare_chat_completions_call, resolve_provider_config}; +use prepare::{parse_messages, resolve_provider_config, resolve_request}; use types::{ChatCompletionsRequest, ChatCompletionsResponse}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn chat_completions( request: ChatCompletionsRequest<'_>, ) -> Result { - execute_chat_completions_provider_call(prepare_chat_completions_call(request)?).await + execute_chat_completions_provider_call(resolve_request(request)?).await } /// Whether the core would accept this request, without resolving credentials or diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index 142b2f2aaed..3be2ba21de4 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -6,7 +6,10 @@ use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider} use super::common_utils::{chat_completions_provider_config, string_headers}; use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig}; -use super::types::{ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest}; +use super::types::{ + ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest, + ResolvedChatCompletionsRequest, +}; pub(super) fn resolve_provider_config<'a>( model: &'a str, @@ -34,12 +37,10 @@ pub(super) fn parse_messages(messages: Value) -> Result, Error> .map_err(|err| Error::InvalidRequest(format!("invalid chat completions messages: {err}"))) } -pub(super) fn prepare_chat_completions_call( +pub(super) fn resolve_request( request: ChatCompletionsRequest<'_>, -) -> Result { +) -> Result, Error> { let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?; - let env_lookup = |key: &str| std::env::var(key).ok(); - let messages = parse_messages(request.messages)?; if messages.is_empty() { return Err(Error::InvalidRequest( @@ -49,11 +50,29 @@ pub(super) fn prepare_chat_completions_call( if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) { return Err(Error::Unsupported(reason.0)); } + Ok(ResolvedChatCompletionsRequest { + model, + config, + messages, + optional_params: request.optional_params, + api_key: request.api_key, + api_base: request.api_base, + extra_headers: request.extra_headers, + timeout: request.timeout, + }) +} - let mut headers = string_headers(request.extra_headers)?; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +fn validate_environment( + request: &ResolvedChatCompletionsRequest<'_>, + model: &str, + config: &dyn ChatCompletionsProviderConfig, +) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> { + let env_lookup = |key: &str| std::env::var(key).ok(); + let mut headers = string_headers(request.extra_headers.clone())?; let auth = config.auth( request.api_key, - &model, + model, &request.optional_params, &env_lookup, )?; @@ -94,7 +113,16 @@ pub(super) fn prepare_chat_completions_call( headers.push(((*name).to_string(), (*value).to_string())); } } + Ok((headers, auth)) +} +pub(super) fn prepare_provider_request( + request: ResolvedChatCompletionsRequest<'_>, +) -> Result { + let (headers, auth) = validate_environment(&request, &request.model, request.config)?; + let model = request.model; + let config = request.config; + let env_lookup = |key: &str| std::env::var(key).ok(); let url = config.complete_url( request.api_base, &model, @@ -102,7 +130,7 @@ pub(super) fn prepare_chat_completions_call( &env_lookup, )?; let transformed = - config.transform_request(&model, messages, request.optional_params.clone())?; + config.transform_request(&model, request.messages, request.optional_params.clone())?; Ok(ProviderChatCompletionsRequest { model, diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index 2858d180e27..f8594dee447 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -2,9 +2,15 @@ use serde_json::{Map, Value, json}; use crate::error::Error; -use super::prepare::prepare_chat_completions_call; +use super::prepare::{prepare_provider_request, resolve_request}; use super::transformation::ChatCompletionsAuth; -use super::types::ChatCompletionsRequest; +use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest}; + +fn prepare_chat_completions_call( + request: ChatCompletionsRequest<'_>, +) -> Result { + prepare_provider_request(resolve_request(request)?) +} fn request<'a>( model: &'a str, diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index a0868209305..d7b9704c46c 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -62,9 +62,8 @@ pub trait ChatCompletionsProviderConfig: Sync { false } - /// Provider parameter names (post-mapping) the Rust path knows how to place - /// in the upstream body. Anything outside this set declines the request. - fn supported_params(&self) -> &'static [&'static str]; + /// Supported OpenAI parameter names paired with their provider names. + fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)]; /// Parameters consumed as call configuration (credentials, endpoints) /// rather than placed in the body. Accepted, never serialized. @@ -78,7 +77,7 @@ pub trait ChatCompletionsProviderConfig: Sync { optional_params: &Map, ) -> Option { unsupported_param( - self.supported_params(), + self.supported_openai_params(), self.config_params(), optional_params, ) @@ -100,7 +99,7 @@ pub trait ChatCompletionsProviderConfig: Sync { } pub fn unsupported_param( - supported: &'static [&'static str], + supported: &'static [(&'static str, &'static str)], config: &'static [&'static str], optional_params: &Map, ) -> Option { @@ -115,7 +114,9 @@ pub fn unsupported_param( .keys() .any(|key| { key != STREAM_PARAM - && !supported.contains(&key.as_str()) + && !supported + .iter() + .any(|(_, provider_name)| *provider_name == key) && !config.contains(&key.as_str()) }) .then_some(Unsupported("unrecognized request parameter")) diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 35dd543a986..3238d09b6b5 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -22,6 +22,17 @@ pub struct ChatCompletionsRequest<'a> { pub timeout: Option, } +pub(super) struct ResolvedChatCompletionsRequest<'a> { + pub(super) model: String, + pub(super) config: &'static dyn ChatCompletionsProviderConfig, + pub(super) messages: Vec, + pub(super) optional_params: Map, + pub(super) api_key: Option<&'a str>, + pub(super) api_base: Option<&'a str>, + pub(super) extra_headers: Option>, + pub(super) timeout: Option, +} + pub(super) struct ProviderChatCompletionsRequest { pub(super) model: String, pub(super) config: &'static dyn ChatCompletionsProviderConfig, diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index 10661fadf96..3633130528d 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -5,6 +5,13 @@ use serde_json::{Map, Value}; use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; use crate::error::{Error, json_type_name}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub async fn http_request( + request: reqwest::RequestBuilder, +) -> Result { + request.send().await +} + /// Bound an upstream error body before it crosses a host boundary, so provider /// bodies stay data-minimized. pub fn truncate_error_body(body: &str) -> String { diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 8dfdb2e361a..8f0f6652fa4 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -10,6 +10,7 @@ pub(super) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_b const HEADER_CONTEXT: &str = "messages"; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) fn messages_provider_config( provider: &str, ) -> Option<&'static dyn AnthropicMessagesProviderConfig> { diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 13a65d86131..61ff81bcdc8 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,13 +1,17 @@ use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; use crate::error::Error; +use crate::http_utils::http_request; use super::client::http_client; use super::common_utils::truncate_error_body; -use super::types::{AnthropicMessagesResponse, ProviderMessagesRequest}; +use super::prepare::prepare_provider_request; +use super::types::{AnthropicMessagesResponse, MessagesRequest}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) async fn execute_messages_provider_call( - request: ProviderMessagesRequest, + request: MessagesRequest<'_>, ) -> Result { + let request = prepare_provider_request(request)?; let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -16,8 +20,7 @@ pub(super) async fn execute_messages_provider_call( request_builder = request_builder.timeout(duration); } - let response = request_builder - .send() + let response = http_request(request_builder) .await .map_err(|err| Error::Network(err.to_string()))?; @@ -40,8 +43,9 @@ pub(super) async fn execute_messages_provider_call( } pub(super) async fn execute_messages_provider_stream( - request: ProviderMessagesRequest, + request: MessagesRequest<'_>, ) -> Result { + let request = prepare_provider_request(request)?; if request.provider != ANTHROPIC_MESSAGES_PROVIDER { return Err(Error::InvalidRequest( "streaming messages is not supported for this provider".to_string(), @@ -56,8 +60,7 @@ pub(super) async fn execute_messages_provider_stream( request_builder = request_builder.timeout(duration); } - let response = request_builder - .send() + let response = http_request(request_builder) .await .map_err(|err| Error::Network(err.to_string()))?; let status = response.status(); diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index ee2877e61fc..cfa8bda1104 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -16,15 +16,15 @@ pub mod transformation; pub mod types; use handler::{execute_messages_provider_call, execute_messages_provider_stream}; -use prepare::prepare_messages_call; use types::{AnthropicMessagesResponse, MessagesRequest}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn messages(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_call(prepare_messages_call(request)?).await + execute_messages_provider_call(request).await } pub async fn messages_stream(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_stream(prepare_messages_call(request)?).await + execute_messages_provider_stream(request).await } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 3b253ac3766..ec83d03f535 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -2,10 +2,11 @@ use crate::error::Error; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; -use super::transformation::MessagesAuthStrategy; +use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; use super::types::{MessagesRequest, ProviderMessagesRequest}; +use serde_json::{Map, Value}; -pub(super) fn prepare_messages_call( +pub(super) fn prepare_provider_request( request: MessagesRequest<'_>, ) -> Result { let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) @@ -29,13 +30,46 @@ pub(super) fn prepare_messages_call( .ok_or_else(|| Error::InvalidProvider(provider.to_string()))?; let env_lookup = |key: &str| std::env::var(key).ok(); - let mut headers = string_headers(request.extra_headers)?; + let headers = + validate_environment(config, request.extra_headers, request.api_key, &env_lookup)?; + + let typed_request = serde_json::from_value(request.body).map_err(|err| { + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + })?; + let transformed = config.transform_request(typed_request)?; + let body = serde_json::to_value(transformed).map_err(|err| { + Error::InvalidRequest(format!( + "failed to serialize Anthropic messages request: {err}" + )) + })?; + + let url = config.complete_url(request.api_base, &model, &env_lookup)?; + + Ok(ProviderMessagesRequest { + provider: provider.to_string(), + model, + config, + url, + body, + upstream_headers: headers, + timeout: request.timeout, + }) +} + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +fn validate_environment( + config: &dyn AnthropicMessagesProviderConfig, + extra_headers: Option>, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result, Error> { + let mut headers = string_headers(extra_headers)?; let auth_strategy = config.auth_strategy(); let already_authorized = has_header(&headers, auth_strategy.header_name()) || (config.accepts_bearer_auth() && has_bearer_auth(&headers)); if !already_authorized { - let api_key = config.resolve_api_key(request.api_key, &env_lookup)?; + let api_key = config.resolve_api_key(api_key, env_lookup)?; let auth_header = match auth_strategy { MessagesAuthStrategy::Bearer => { ("authorization".to_string(), format!("Bearer {api_key}")) @@ -51,24 +85,5 @@ pub(super) fn prepare_messages_call( } } - let url = config.complete_url(request.api_base, &model, &env_lookup)?; - let typed_request = serde_json::from_value(request.body).map_err(|err| { - Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) - })?; - let transformed = config.transform_request(typed_request)?; - let body = serde_json::to_value(transformed).map_err(|err| { - Error::InvalidRequest(format!( - "failed to serialize Anthropic messages request: {err}" - )) - })?; - - Ok(ProviderMessagesRequest { - provider: provider.to_string(), - model, - config, - url, - body, - upstream_headers: headers, - timeout: request.timeout, - }) + Ok(headers) } diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index 673a5728aca..a5904c085a0 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -45,6 +45,7 @@ pub trait AnthropicMessagesProviderConfig: Sync { ] } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_request( &self, request: AnthropicMessagesRequest, @@ -52,6 +53,7 @@ pub trait AnthropicMessagesProviderConfig: Sync { Ok(request) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_response( &self, _model: &str, diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 3d3c16c8cb6..ad484c8f968 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -27,6 +27,7 @@ pub enum OcrResponseHandling { pub trait OcrProviderConfig: Sync { fn supported_ocr_params(&self) -> &'static [&'static str]; + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn map_ocr_params(&self, non_default_params: &Map) -> Map { let mut mapped_params = Map::new(); for (param, value) in non_default_params { @@ -64,6 +65,25 @@ pub trait OcrProviderConfig: Sync { env_lookup: &dyn Fn(&str) -> Option, ) -> Result; + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + fn validate_environment( + &self, + headers: Vec<(String, String)>, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result, Error> { + let strategy = self.auth_strategy(); + if crate::http_utils::has_header(&headers, strategy.header_name()) { + return Ok(headers); + } + let api_key = self.resolve_api_key(api_key, env_lookup)?; + let auth_header = match strategy { + OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), + OcrAuthStrategy::Header(name) => (name.to_string(), api_key), + }; + Ok(std::iter::once(auth_header).chain(headers).collect()) + } + fn auth_strategy(&self) -> OcrAuthStrategy { OcrAuthStrategy::Bearer } diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs index 97cc48aa6f2..a7d5a8ad0cf 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs @@ -27,7 +27,12 @@ use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage /// per-model gate inside `transform_request`, the function this route replaces. /// Forwarding it would send `top_k` to a model that removed sampling params and /// take a 400 after the call, where Python drops it and succeeds. -const SUPPORTED_PARAMS: &[&str] = &["max_tokens", "temperature", "top_p", "stop_sequences"]; +const SUPPORTED_PARAMS: &[(&str, &str)] = &[ + ("max_tokens", "max_tokens"), + ("temperature", "temperature"), + ("top_p", "top_p"), + ("stop", "stop_sequences"), +]; pub struct AnthropicChatCompletionsConfig; @@ -112,7 +117,8 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { }) } - fn supported_params(&self) -> &'static [&'static str] { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] { SUPPORTED_PARAMS } @@ -121,7 +127,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { messages: &[ChatMessage], optional_params: &Map, ) -> Option { - unsupported_param(SUPPORTED_PARAMS, &[], optional_params) + unsupported_param(self.supported_openai_params(), &[], optional_params) .or_else(|| messages.iter().find_map(unsupported_message)) // Anthropic rejects a request whose first turn is not a user turn. // Python only repairs that under `litellm.modify_params`, which the @@ -132,6 +138,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_request( &self, model: &str, @@ -143,6 +150,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_response( &self, _model: &str, diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs index 8fcc0f36c7d..f31b961e78a 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -47,6 +47,7 @@ pub fn complete_anthropic_url( } impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( &self, api_base: Option<&str>, diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 70dad0300f1..b8ca10461fb 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -142,6 +142,7 @@ fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMess } impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( &self, api_base: Option<&str>, diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs index bb4f6afe5f9..9bf1f73a74d 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -46,10 +46,12 @@ fn optional_string<'a>(params: &'a Map, key: &str) -> Option<&'a } impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_transcription_params(&self) -> &'static [&'static str] { SUPPORTED_PARAMS } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_transcription_request( &self, _model: &str, @@ -83,6 +85,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_transcription_response( &self, _model: &str, diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs index ef5f44b4a14..7be3d108d44 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs @@ -23,11 +23,12 @@ use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT /// `additionalModelRequestFields` for Anthropic base models and to /// `inferenceConfig` otherwise, and that branch reads the model catalog the /// core cannot see. -const SUPPORTED_PARAMS: &[&str] = &["maxTokens", "temperature", "topP", "stopSequences"]; - -/// Params that belong in `inferenceConfig`, in the order Python's -/// `AmazonConverseConfig` declares them, so bodies compare cleanly. -const INFERENCE_CONFIG_PARAMS: &[&str] = SUPPORTED_PARAMS; +const SUPPORTED_PARAMS: &[(&str, &str)] = &[ + ("max_tokens", "maxTokens"), + ("temperature", "temperature"), + ("top_p", "topP"), + ("stop", "stopSequences"), +]; const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "aws_bedrock_runtime_endpoint"; @@ -66,7 +67,7 @@ fn converse_body(conversation: &Conversation, params: &Map) -> Va }) .collect(); - let inference_config = Map::from_iter(INFERENCE_CONFIG_PARAMS.iter().filter_map(|name| { + let inference_config = Map::from_iter(SUPPORTED_PARAMS.iter().filter_map(|(_, name)| { params .get(*name) .map(|value| ((*name).to_string(), value.clone())) @@ -162,7 +163,8 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { &[("Content-Type", "application/json")] } - fn supported_params(&self) -> &'static [&'static str] { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] { SUPPORTED_PARAMS } @@ -175,32 +177,36 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { messages: &[ChatMessage], optional_params: &Map, ) -> Option { - unsupported_param(SUPPORTED_PARAMS, CONFIG_PARAMS, optional_params) - .or_else(|| messages.iter().find_map(unsupported_message)) - // Python's Converse translation drops blank text blocks instead of - // substituting the placeholder the shared conversation builder - // applies, so decline blank text rather than diverge. - .or_else(|| { - messages - .iter() - .any(has_blank_text) - .then_some(Unsupported("blank message text")) - }) - // Converse has no assistant prefill: Python inserts a continue turn - // when a conversation opens or closes on an assistant message, and - // only under `litellm.modify_params`, which the core cannot see. - // Declining both ends also keeps the shared builder's final - // assistant right-strip (an Anthropic rule) unreachable here. - .or_else(|| { - let conversation = build_conversation(messages); - let ends_on_assistant = conversation - .turns - .last() - .is_some_and(|turn| turn.role == TurnRole::Assistant); - (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported( - "conversation does not run user turn to user turn", - )) - }) + unsupported_param( + self.supported_openai_params(), + CONFIG_PARAMS, + optional_params, + ) + .or_else(|| messages.iter().find_map(unsupported_message)) + // Python's Converse translation drops blank text blocks instead of + // substituting the placeholder the shared conversation builder + // applies, so decline blank text rather than diverge. + .or_else(|| { + messages + .iter() + .any(has_blank_text) + .then_some(Unsupported("blank message text")) + }) + // Converse has no assistant prefill: Python inserts a continue turn + // when a conversation opens or closes on an assistant message, and + // only under `litellm.modify_params`, which the core cannot see. + // Declining both ends also keeps the shared builder's final + // assistant right-strip (an Anthropic rule) unreachable here. + .or_else(|| { + let conversation = build_conversation(messages); + let ends_on_assistant = conversation + .turns + .last() + .is_some_and(|turn| turn.role == TurnRole::Assistant); + (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported( + "conversation does not run user turn to user turn", + )) + }) } fn transform_request( 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 6a8a38204a9..9648321d7ff 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -70,10 +70,12 @@ pub struct MistralOcrConfig; pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig; impl OcrProviderConfig for MistralOcrConfig { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_ocr_params(&self) -> &'static [&'static str] { SUPPORTED_OCR_PARAMS } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_ocr_request( &self, model: &str, @@ -100,6 +102,7 @@ impl OcrProviderConfig for MistralOcrConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_ocr_response( &self, model: &str, @@ -134,6 +137,7 @@ impl OcrProviderConfig for MistralOcrConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( &self, api_base: Option<&str>, @@ -153,6 +157,7 @@ impl OcrProviderConfig for MistralOcrConfig { } } +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub fn supported_ocr_params() -> &'static [&'static str] { MISTRAL_OCR_CONFIG.supported_ocr_params() } @@ -161,6 +166,7 @@ pub fn map_ocr_params(non_default_params: &Map) -> Map Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 498003de149..275e00300d4 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -16,6 +16,9 @@ extension-module = ["pyo3/extension-module"] panic-test = [] [dependencies] +serde.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } litellm-python-interop.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/constants.rs b/litellm-rust/crates/python-bridge/src/constants.rs new file mode 100644 index 00000000000..07b2836b838 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/constants.rs @@ -0,0 +1 @@ +pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace"; diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs new file mode 100644 index 00000000000..cc153a89b8f --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/diagnostics.rs @@ -0,0 +1,23 @@ +use litellm_python_interop::release_count; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +#[pyfunction] +fn gil_stats(py: Python<'_>) -> PyResult> { + let stats = PyDict::new(py); + stats.set_item("releases", release_count())?; + Ok(stats.into_any().unbind()) +} + +#[cfg(feature = "panic-test")] +#[pyfunction] +fn _panic_for_test() { + panic!("intentional PyO3 panic smoke test"); +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(gil_stats, module)?)?; + #[cfg(feature = "panic-test")] + module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; + Ok(()) +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs new file mode 100644 index 00000000000..914e2e1e033 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -0,0 +1,61 @@ +use litellm_core::error::Error; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; + +pyo3::create_exception!( + _native, + RustBridgeDeclined, + pyo3::exceptions::PyException, + "The route declined before calling the provider, so the host may retry on its own path." +); + +pyo3::create_exception!( + _native, + RustUpstreamError, + pyo3::exceptions::PyException, + "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." +); + +pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { + match err { + Error::Auth(message) => PyValueError::new_err(message), + Error::InvalidProvider(_) + | Error::InvalidRequest(_) + | Error::InvalidType { .. } + | Error::MissingField(_) => PyValueError::new_err(err.to_string()), + other => PyRuntimeError::new_err(other.to_string()), + } +} + +/// Map a core error for a route whose host keeps a Python implementation. +/// +/// The distinction the host needs is whether the provider was already called. +/// Everything raised before the request goes out is safe for the host to retry +/// on its own path; anything after it is not, because the provider has already +/// done the work and billed for it. +pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { + match err { + Error::Unsupported(_) + | Error::Auth(_) + | Error::InvalidProvider(_) + | Error::InvalidRequest(_) + | Error::InvalidType { .. } + | Error::MissingField(_) + | Error::Routing(_) + // Nothing reached the provider, so serving it on Python cannot double + // bill and is the only way the caller gets an answer at all. + | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), + Error::Http { status, body } => { + RustUpstreamError::new_err((status, format!("{status}: {body}"))) + } + Error::Network(message) | Error::InvalidResponse(message) => { + RustUpstreamError::new_err((0u16, message)) + } + } +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + let py = module.py(); + module.add("RustBridgeDeclined", py.get_type::())?; + module.add("RustUpstreamError", py.get_type::()) +} diff --git a/litellm-rust/crates/python-bridge/src/function_trace.rs b/litellm-rust/crates/python-bridge/src/function_trace.rs new file mode 100644 index 00000000000..420d237c79d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/function_trace.rs @@ -0,0 +1,216 @@ +use std::future::Future; +use std::sync::{Arc, Mutex}; + +use serde::Serialize; +use tracing::instrument::WithSubscriber; +use tracing::span::{Attributes, Id}; +use tracing::{Dispatch, Level, Subscriber}; +use tracing_subscriber::filter::{LevelFilter, filter_fn}; +use tracing_subscriber::layer::Context; +use tracing_subscriber::prelude::*; +use tracing_subscriber::registry::LookupSpan; +use tracing_subscriber::{Layer, Registry}; + +use crate::constants::FUNCTION_TRACE_TARGET; + +#[derive(Serialize)] +#[serde(untagged)] +pub(crate) enum TraceResponse { + Plain(T), + Traced { + response: T, + trace: Vec, + }, +} + +pub(crate) async fn trace_call( + future: impl Future>, + enabled: bool, +) -> Result, E> { + if !enabled { + return future.await.map(TraceResponse::Plain); + } + let trace = FunctionTrace::default(); + let response = future.with_subscriber(trace.dispatcher()).await?; + Ok(TraceResponse::Traced { + response, + trace: trace.events(), + }) +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct FunctionTraceEvent { + pub function: &'static str, + pub depth: usize, +} + +#[derive(Clone, Default)] +pub struct FunctionTrace { + events: Arc>>, +} + +impl FunctionTrace { + pub fn dispatcher(&self) -> Dispatch { + let filter = filter_fn(|metadata| { + metadata.is_span() + && metadata.target() == FUNCTION_TRACE_TARGET + && *metadata.level() == Level::TRACE + }) + .with_max_level_hint(LevelFilter::TRACE); + Dispatch::new( + Registry::default().with( + FunctionTraceLayer { + trace: self.clone(), + } + .with_filter(filter), + ), + ) + } + + pub fn events(&self) -> Vec { + self.events + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() + } +} + +struct FunctionTraceLayer { + trace: FunctionTrace, +} + +impl Layer for FunctionTraceLayer +where + S: Subscriber + for<'lookup> LookupSpan<'lookup>, +{ + fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) { + let depth = context + .span(id) + .map(|span| span.scope().skip(1).count()) + .unwrap_or_default(); + self.trace + .events + .lock() + .unwrap_or_else(|error| error.into_inner()) + .push(FunctionTraceEvent { + function: attributes.metadata().name(), + depth, + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + async fn outer() { + tokio::task::yield_now().await; + inner().await; + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + async fn inner() { + tokio::task::yield_now().await; + } + + #[tokio::test] + async fn concurrent_futures_keep_separate_traces_across_yields() { + use tracing::instrument::WithSubscriber; + + let first = FunctionTrace::default(); + let second = FunctionTrace::default(); + let outside = FunctionTrace::default(); + + async { + tokio::join!( + outer().with_subscriber(first.dispatcher()), + inner().with_subscriber(second.dispatcher()), + ); + inner().await; + } + .with_subscriber(outside.dispatcher()) + .await; + + assert_eq!( + first.events(), + vec![ + FunctionTraceEvent { + function: "outer", + depth: 0 + }, + FunctionTraceEvent { + function: "inner", + depth: 1 + }, + ], + ); + assert_eq!( + second.events(), + vec![FunctionTraceEvent { + function: "inner", + depth: 0 + }], + ); + assert_eq!( + outside.events(), + vec![FunctionTraceEvent { + function: "inner", + depth: 0 + }], + ); + } + + #[test] + fn records_matching_spans_in_creation_order() { + let trace = FunctionTrace::default(); + let dispatch = trace.dispatcher(); + + tracing::dispatcher::with_default(&dispatch, || { + let _ignored = tracing::trace_span!(target: "other", "ignored"); + let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name"); + let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level"); + let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name"); + }); + + assert_eq!( + trace.events(), + vec![ + FunctionTraceEvent { + function: "same_name", + depth: 0, + }, + FunctionTraceEvent { + function: "same_name", + depth: 0, + }, + ] + ); + } + + #[test] + fn records_matching_span_nesting_depth() { + let trace = FunctionTrace::default(); + let dispatch = trace.dispatcher(); + + tracing::dispatcher::with_default(&dispatch, || { + let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer"); + let _outer_guard = outer.enter(); + let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner"); + }); + + assert_eq!( + trace.events(), + vec![ + FunctionTraceEvent { + function: "outer", + depth: 0, + }, + FunctionTraceEvent { + function: "inner", + depth: 1, + }, + ] + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 2e2624acbe1..1516fcd7d19 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,142 +1,16 @@ -use std::collections::HashMap; -use std::time::Duration; +mod constants; +mod diagnostics; +mod errors; +pub mod function_trace; +mod marshal; +mod routes; -use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; -use litellm_core::audio_transcription::{ - AudioTranscriptionRequest, audio_transcription as run_audio_transcription, -}; -use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; -use litellm_core::chat_completions::{ - chat_completions as run_chat_completions, chat_completions_decline_reason, -}; -use litellm_core::error::Error; -use litellm_core::messages::messages as run_messages; -use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; -use litellm_python_interop::{from_py, release_count, release_gil, to_py}; -use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; -use pyo3::types::{PyAny, PyDict}; -use serde_json::{Map, Value}; +use pyo3::types::PyAny; -pyo3::create_exception!( - _native, - RustBridgeDeclined, - pyo3::exceptions::PyException, - "The route declined before calling the provider, so the host may retry on its own path." -); - -pyo3::create_exception!( - _native, - RustUpstreamError, - pyo3::exceptions::PyException, - "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." -); - -type MarshaledOcrInputs = ( - Value, - Option>, - Map, - Option, -); - -fn messages_response_to_py( - py: Python<'_>, - response: AnthropicMessagesResponse, -) -> PyResult> { - to_py(py, &response) -} - -fn chat_completions_response_to_py( - py: Python<'_>, - response: ChatCompletionsResponse, -) -> PyResult> { - to_py(py, &response) -} - -fn core_error_to_pyerr(err: Error) -> PyErr { - match err { - Error::Auth(message) => PyValueError::new_err(message), - Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::InvalidType { .. } - | Error::MissingField(_) => PyValueError::new_err(err.to_string()), - other => PyRuntimeError::new_err(other.to_string()), - } -} - -/// Map a core error for a route whose host keeps a Python implementation. -/// -/// The distinction the host needs is whether the provider was already called. -/// Everything raised before the request goes out is safe for the host to retry -/// on its own path; anything after it is not, because the provider has already -/// done the work and billed for it. -fn chat_completions_error_to_pyerr(err: Error) -> PyErr { - match err { - Error::Unsupported(_) - | Error::Auth(_) - | Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::InvalidType { .. } - | Error::MissingField(_) - | Error::Routing(_) - // Nothing reached the provider, so serving it on Python cannot double - // bill and is the only way the caller gets an answer at all. - | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), - Error::Http { status, body } => { - RustUpstreamError::new_err((status, format!("{status}: {body}"))) - } - Error::Network(message) | Error::InvalidResponse(message) => { - RustUpstreamError::new_err((0u16, message)) - } - } -} - -fn optional_object_to_map( - py: Python<'_>, - name: &'static str, - value: Option>, -) -> PyResult> { - match value { - Some(value) => match from_py(value.bind(py))? { - Value::Object(map) => Ok(map), - _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), - }, - None => Ok(Map::new()), - } -} - -fn optional_timeout(timeout_seconds: Option) -> Option { - timeout_seconds.and_then(|secs| { - if secs.is_finite() && secs > 0.0 { - Some(Duration::from_secs_f64(secs)) - } else { - None - } - }) -} - -fn marshal_headers( - py: Python<'_>, - headers: Option>, -) -> PyResult> { - let value = match headers { - Some(headers) => from_py(headers.bind(py))?, - None => Value::Object(Map::new()), - }; - let Value::Object(headers) = value else { - return Err(PyValueError::new_err("headers must be a dict")); - }; - headers - .into_iter() - .map(|(name, value)| { - value - .as_str() - .map(|value| (name, value.to_string())) - .ok_or_else(|| PyValueError::new_err("header values must be strings")) - }) - .collect() -} +use crate::errors::core_error_to_pyerr; +use crate::marshal::{marshal_headers, optional_timeout}; #[pyclass] struct ResponsesWebSocketConnection { @@ -186,445 +60,48 @@ impl ResponsesWebSocketConnection { } } -fn marshal_inputs( - py: Python<'_>, - document: Py, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult { - let document = from_py(document.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - - Ok((document, extra_headers, optional_params, timeout)) -} - -#[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn ocr( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (document, extra_headers, optional_params, timeout) = marshal_inputs( - py, - document, - extra_headers, - optional_params, - timeout_seconds, - )?; - - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest { - model: &model, - document, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, - })) - }); - - match result { - Ok(value) => to_py(py, &value), - Err(err) => Err(core_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn aocr( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (document, extra_headers, optional_params, timeout) = marshal_inputs( - py, - document, - extra_headers, - optional_params, - timeout_seconds, - )?; - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let value = run_ocr(OcrRequest { - model: &model, - document, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, - }) - .await - .map_err(core_error_to_pyerr)?; - - Python::attach(|py| to_py(py, &value)) - }) -} - -#[pyfunction] -#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn transcription( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let audio = from_py(audio.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription( - AudioTranscriptionRequest { - model: &model, - audio, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - }, - )) - }); - match result { - Ok(value) => to_py(py, &value), - Err(err) => Err(core_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn atranscription( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let audio = from_py(audio.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let value = run_audio_transcription(AudioTranscriptionRequest { - model: &model, - audio, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - }) - .await - .map_err(core_error_to_pyerr)?; - Python::attach(|py| to_py(py, &value)) - }) -} - -type MarshaledMessagesInputs = (Value, Option>, Option); - -fn marshal_messages_inputs( - py: Python<'_>, - body: Py, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - let body: Value = from_py(body.bind(py))?; - if !body.is_object() { - return Err(PyValueError::new_err("body must be a dict")); - } - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - Ok((body, extra_headers, optional_timeout(timeout_seconds))) -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn messages( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (body, extra_headers, timeout) = - marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; - - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest { - model: &model, - body, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - })) - }); - - match result { - Ok(response) => messages_response_to_py(py, response), - Err(err) => Err(core_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn amessages( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (body, extra_headers, timeout) = - marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let response = run_messages(MessagesRequest { - model: &model, - body, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }) - .await - .map_err(core_error_to_pyerr)?; - - Python::attach(|py| messages_response_to_py(py, response)) - }) -} - -type MarshaledChatCompletionsInputs = ( - Value, - Map, - Option>, - Option, -); - -fn marshal_chat_completions_inputs( - py: Python<'_>, - messages: Py, - optional_params: Option>, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - let messages: Value = from_py(messages.bind(py))?; - if !messages.is_array() { - return Err(PyValueError::new_err("messages must be a list")); - } - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - Ok(( - messages, - optional_params, - extra_headers, - optional_timeout(timeout_seconds), - )) -} - -/// The decline reason for this request, or `None` when the Rust path accepts -/// it. Resolves no credentials and performs no I/O, so a host can ask before -/// committing to either path. -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))] -fn chat_completions_decline( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - custom_llm_provider: Option, -) -> PyResult> { - let messages = from_py(messages.bind(py))?; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - Ok(chat_completions_decline_reason( - &model, - custom_llm_provider.as_deref(), - messages, - &optional_params, - ) - .map(str::to_string)) -} - -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn chat_completions( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs( - py, - messages, - optional_params, - extra_headers, - timeout_seconds, - )?; - - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions( - ChatCompletionsRequest { - model: &model, - messages, - optional_params, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }, - )) - }); - - match result { - Ok(response) => chat_completions_response_to_py(py, response), - Err(err) => Err(chat_completions_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn achat_completions( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs( - py, - messages, - optional_params, - extra_headers, - timeout_seconds, - )?; - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let response = run_chat_completions(ChatCompletionsRequest { - model: &model, - messages, - optional_params, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }) - .await - .map_err(chat_completions_error_to_pyerr)?; - - Python::attach(|py| chat_completions_response_to_py(py, response)) - }) -} - -#[pyfunction] -fn gil_stats(py: Python<'_>) -> PyResult> { - let stats = PyDict::new(py); - stats.set_item("releases", release_count())?; - Ok(stats.into_any().unbind()) -} - -#[cfg(feature = "panic-test")] -#[pyfunction] -fn _panic_for_test() { - panic!("intentional PyO3 panic smoke test"); -} - #[pymodule] fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { - let py = module.py(); - module.add_function(wrap_pyfunction!(ocr, module)?)?; - module.add_function(wrap_pyfunction!(aocr, module)?)?; - module.add_function(wrap_pyfunction!(transcription, module)?)?; - module.add_function(wrap_pyfunction!(atranscription, module)?)?; - module.add_function(wrap_pyfunction!(messages, module)?)?; - module.add_function(wrap_pyfunction!(amessages, module)?)?; - module.add("RustBridgeDeclined", py.get_type::())?; - module.add("RustUpstreamError", py.get_type::())?; - module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?; - module.add_function(wrap_pyfunction!(chat_completions, module)?)?; - module.add_function(wrap_pyfunction!(achat_completions, module)?)?; + errors::register(module)?; + routes::register(module)?; module.add_class::()?; - module.add_function(wrap_pyfunction!(gil_stats, module)?)?; - #[cfg(feature = "panic-test")] - module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; - Ok(()) + diagnostics::register(module) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn module_registration_preserves_the_public_surface() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "_native").expect("module should be created"); + _native(&module).expect("module should register"); + + let expected = [ + "RustBridgeDeclined", + "RustUpstreamError", + "ocr", + "aocr", + "transcription", + "atranscription", + "messages", + "amessages", + "chat_completions_decline", + "chat_completions", + "achat_completions", + "ResponsesWebSocketConnection", + "gil_stats", + ]; + + for name in expected { + assert!( + module + .hasattr(name) + .expect("attribute lookup should succeed") + ); + } + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs new file mode 100644 index 00000000000..8d801109dfa --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -0,0 +1,63 @@ +use std::collections::HashMap; +use std::time::Duration; + +use litellm_python_interop::from_py; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +pub(crate) fn optional_object_to_map( + py: Python<'_>, + name: &'static str, + value: Option>, +) -> PyResult> { + match value { + Some(value) => match from_py(value.bind(py))? { + Value::Object(map) => Ok(map), + _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), + }, + None => Ok(Map::new()), + } +} + +pub(crate) fn optional_object( + py: Python<'_>, + name: &'static str, + value: Option>, +) -> PyResult>> { + value + .map(|value| optional_object_to_map(py, name, Some(value))) + .transpose() +} + +pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option { + timeout_seconds.and_then(|secs| { + if secs.is_finite() && secs > 0.0 { + Some(Duration::from_secs_f64(secs)) + } else { + None + } + }) +} + +pub(crate) fn marshal_headers( + py: Python<'_>, + headers: Option>, +) -> PyResult> { + let value = match headers { + Some(headers) => from_py(headers.bind(py))?, + None => Value::Object(Map::new()), + }; + let Value::Object(headers) = value else { + return Err(PyValueError::new_err("headers must be a dict")); + }; + headers + .into_iter() + .map(|(name, value)| { + value + .as_str() + .map(|value| (name, value.to_string())) + .ok_or_else(|| PyValueError::new_err("header values must be strings")) + }) + .collect() +} diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs new file mode 100644 index 00000000000..c3e290310fa --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -0,0 +1,126 @@ +use std::time::Duration; + +use litellm_core::audio_transcription::{ + AudioTranscriptionRequest, audio_transcription as run_audio_transcription, +}; +use litellm_core::error::Error; +use litellm_python_interop::from_py; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::core_error_to_pyerr; +use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; + +use super::{block_on, into_py_future}; + +struct TranscriptionInputs { + model: String, + audio: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Map, + timeout: Option, +} + +#[allow(clippy::too_many_arguments)] +fn marshal_inputs( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult { + Ok(TranscriptionInputs { + model, + audio: from_py(audio.bind(py))?, + api_key, + api_base, + custom_llm_provider, + extra_headers: optional_object(py, "extra_headers", extra_headers)?, + optional_params: optional_object_to_map(py, "optional_params", optional_params)?, + timeout: optional_timeout(timeout_seconds), + }) +} + +async fn call(inputs: TranscriptionInputs) -> Result { + run_audio_transcription(AudioTranscriptionRequest { + model: &inputs.model, + audio: inputs.audio, + api_key: inputs.api_key.as_deref(), + api_base: inputs.api_base.as_deref(), + custom_llm_provider: inputs.custom_llm_provider.as_deref(), + extra_headers: inputs.extra_headers, + optional_params: inputs.optional_params, + timeout: inputs.timeout, + }) + .await +} + +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn transcription( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + audio, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + )?; + block_on(py, call(inputs), trace, core_error_to_pyerr) +} + +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn atranscription( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + audio, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + )?; + into_py_future(py, call(inputs), trace, core_error_to_pyerr) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(transcription, module)?)?; + module.add_function(wrap_pyfunction!(atranscription, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs new file mode 100644 index 00000000000..2ef3810015e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -0,0 +1,156 @@ +use std::time::Duration; + +use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; +use litellm_core::chat_completions::{ + chat_completions as run_chat_completions, chat_completions_decline_reason, +}; +use litellm_core::error::Error; +use litellm_python_interop::from_py; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::chat_completions_error_to_pyerr; +use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; + +use super::{block_on, into_py_future}; + +struct ChatCompletionsInputs { + model: String, + messages: Value, + optional_params: Map, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout: Option, +} + +#[allow(clippy::too_many_arguments)] +fn marshal_inputs( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult { + let messages: Value = from_py(messages.bind(py))?; + if !messages.is_array() { + return Err(PyValueError::new_err("messages must be a list")); + } + Ok(ChatCompletionsInputs { + model, + messages, + optional_params: optional_object_to_map(py, "optional_params", optional_params)?, + api_key, + api_base, + custom_llm_provider, + extra_headers: optional_object(py, "extra_headers", extra_headers)?, + timeout: optional_timeout(timeout_seconds), + }) +} + +async fn call(inputs: ChatCompletionsInputs) -> Result { + run_chat_completions(ChatCompletionsRequest { + model: &inputs.model, + messages: inputs.messages, + optional_params: inputs.optional_params, + api_key: inputs.api_key.as_deref(), + api_base: inputs.api_base.as_deref(), + custom_llm_provider: inputs.custom_llm_provider.as_deref(), + extra_headers: inputs.extra_headers, + timeout: inputs.timeout, + }) + .await +} + +/// The decline reason for this request, or `None` when the Rust path accepts +/// it. Resolves no credentials and performs no I/O, so a host can ask before +/// committing to either path. +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))] +fn chat_completions_decline( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + custom_llm_provider: Option, +) -> PyResult> { + let messages = from_py(messages.bind(py))?; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + Ok(chat_completions_decline_reason( + &model, + custom_llm_provider.as_deref(), + messages, + &optional_params, + ) + .map(str::to_string)) +} + +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn chat_completions( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + messages, + optional_params, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout_seconds, + )?; + block_on(py, call(inputs), trace, chat_completions_error_to_pyerr) +} + +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn achat_completions( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + messages, + optional_params, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout_seconds, + )?; + into_py_future(py, call(inputs), trace, chat_completions_error_to_pyerr) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?; + module.add_function(wrap_pyfunction!(chat_completions, module)?)?; + module.add_function(wrap_pyfunction!(achat_completions, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs new file mode 100644 index 00000000000..a29fbdd9fdc --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -0,0 +1,122 @@ +use std::time::Duration; + +use litellm_core::error::Error; +use litellm_core::messages::messages as run_messages; +use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; +use litellm_python_interop::from_py; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::core_error_to_pyerr; +use crate::marshal::{optional_object, optional_timeout}; + +use super::{block_on, into_py_future}; + +struct MessagesInputs { + model: String, + body: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout: Option, +} + +#[allow(clippy::too_many_arguments)] +fn marshal_inputs( + py: Python<'_>, + model: String, + body: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult { + let body: Value = from_py(body.bind(py))?; + if !body.is_object() { + return Err(PyValueError::new_err("body must be a dict")); + } + Ok(MessagesInputs { + model, + body, + api_key, + api_base, + custom_llm_provider, + extra_headers: optional_object(py, "extra_headers", extra_headers)?, + timeout: optional_timeout(timeout_seconds), + }) +} + +async fn call(inputs: MessagesInputs) -> Result { + run_messages(MessagesRequest { + model: &inputs.model, + body: inputs.body, + api_key: inputs.api_key.as_deref(), + api_base: inputs.api_base.as_deref(), + custom_llm_provider: inputs.custom_llm_provider.as_deref(), + extra_headers: inputs.extra_headers, + timeout: inputs.timeout, + }) + .await +} + +#[pyfunction] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn messages( + py: Python<'_>, + model: String, + body: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + body, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout_seconds, + )?; + block_on(py, call(inputs), trace, core_error_to_pyerr) +} + +#[pyfunction] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn amessages( + py: Python<'_>, + model: String, + body: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + body, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout_seconds, + )?; + into_py_future(py, call(inputs), trace, core_error_to_pyerr) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(messages, module)?)?; + module.add_function(wrap_pyfunction!(amessages, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs new file mode 100644 index 00000000000..ab31702d70f --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -0,0 +1,53 @@ +use std::future::Future; + +use litellm_core::error::Error; +use litellm_python_interop::{release_gil, to_py}; +use pyo3::prelude::*; +use serde::Serialize; + +use crate::function_trace::trace_call; + +mod audio_transcription; +mod chat_completions; +mod messages; +mod ocr; + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + ocr::register(module)?; + audio_transcription::register(module)?; + messages::register(module)?; + chat_completions::register(module) +} + +fn block_on( + py: Python<'_>, + call: impl Future> + Send, + trace: bool, + map_err: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send, +{ + let result = release_gil(py, || { + pyo3_async_runtimes::tokio::get_runtime().block_on(trace_call(call, trace)) + }); + match result { + Ok(response) => to_py(py, &response), + Err(err) => Err(map_err(err)), + } +} + +fn into_py_future<'py, T>( + py: Python<'py>, + call: impl Future> + Send + 'static, + trace: bool, + map_err: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, +{ + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let response = trace_call(call, trace).await.map_err(map_err)?; + Python::attach(|py| to_py(py, &response)) + }) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs new file mode 100644 index 00000000000..ad2a839de07 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -0,0 +1,128 @@ +use std::time::Duration; + +use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; +use litellm_core::error::Error; +use litellm_python_interop::from_py; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::core_error_to_pyerr; +use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; + +use super::{block_on, into_py_future}; + +struct OcrInputs { + model: String, + document: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Map, + timeout: Option, +} + +#[allow(clippy::too_many_arguments)] +fn marshal_inputs( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult { + Ok(OcrInputs { + model, + document: from_py(document.bind(py))?, + api_key, + api_base, + custom_llm_provider, + extra_headers: optional_object(py, "extra_headers", extra_headers)?, + optional_params: optional_object_to_map(py, "optional_params", optional_params)?, + timeout: optional_timeout(timeout_seconds), + }) +} + +async fn call(inputs: OcrInputs) -> Result { + run_ocr(OcrRequest { + model: &inputs.model, + document: inputs.document, + api_key: inputs.api_key.as_deref(), + api_base: inputs.api_base.as_deref(), + custom_llm_provider: inputs.custom_llm_provider.as_deref(), + extra_headers: inputs.extra_headers, + optional_params: inputs.optional_params, + timeout: inputs.timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }) + .await +} + +#[pyfunction] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn ocr( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + document, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + )?; + block_on(py, call(inputs), trace, core_error_to_pyerr) +} + +#[pyfunction] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn aocr( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + document, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + )?; + into_py_future(py, call(inputs), trace, core_error_to_pyerr) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(ocr, module)?)?; + module.add_function(wrap_pyfunction!(aocr, module)?) +} diff --git a/tests/sdk_function_trace/README.md b/tests/sdk_function_trace/README.md new file mode 100644 index 00000000000..d3a3b654aea --- /dev/null +++ b/tests/sdk_function_trace/README.md @@ -0,0 +1,30 @@ +# SDK function tracing + +The compare runner executes the same SDK calls through the Python engine and the Rust native bridge against a local HTTP provider fixture, then prints their pipeline trees side by side. Matching calls align on the same row in green; Python-only calls are blue, Rust-only calls yellow, and reordered calls red. Gaps preserve execution order and each column retains its own nesting. A comparison column labels every row even without color. Colors are enabled in terminals unless `NO_COLOR` is set. A difference summary follows (shared step order, python-only steps, rust-only steps). Each invocation must issue exactly one HTTP request. It requires the LiteLLM Python dependencies and the native extension built with tracing support + +From the repository root, using the project's Python environment: + +```bash +uv run python -m tests.sdk_function_trace.compare +uv run python -m tests.sdk_function_trace.compare --route ocr +uv run python -m tests.sdk_function_trace.compare --route ocr --sync +uv run python -m tests.sdk_function_trace.compare --route all --both --check +``` + +Calls default to async; use `--sync` for synchronous calls or `--both` for the complete matrix. Python sync Messages raises `not implemented for sync calls`; only that exact failure is marked `SKIP`, and the runner still executes Rust sync Messages and subsequent routes. Bedrock transcription has no independent Python provider implementation: its Python trace covers SDK dispatch into Rust + +Both engines are projected onto a shared per-route step table (`steps.py`): canonical names such as `transform_ocr_request` map Python functions (`MistralOCRConfig.transform_ocr_request`) and Rust spans (`transform_ocr_request`) to the same label. Only the first occurrence of each step is kept. Python indentation uses each event's actual frame ancestors and the nearest already displayed ancestor, so returned helpers and coroutine resumptions do not create false parents. Rust indentation uses instrumented span ancestry. Unmatched Rust span names pass through unchanged. `--full` prints every captured runtime event; validation still uses projected steps + +Every report checks required stage presence and dependency order. Provider lookup must precede request transformation, which must precede HTTP, followed by response transformation. The handler must precede HTTP; parameter mapping and supported-parameter checks must precede request transformation. Environment validation and URL construction, where mapped, must precede HTTP. Python transcription is checked only through native dispatch. `--check` also requires identical canonical step sequences for comparable routes and exits nonzero for missing, extra, or reordered steps, or an unexpected call failure, after finishing all selected cases + +Individual stage checks are separate from cross-language `step parity`. Passing stage checks cannot override a failing step comparison. Bedrock transcription and Python sync Messages report `UNAVAILABLE` for cross-language parity because they lack an independent Python execution to compare. Absolute nesting depth is not a cross-language gate: async Python Messages dispatches its handler onto another thread. See `route-comparison.md` for the audited matrix and remaining contract limitations + +The Python runner uses the existing `profile_python` / `sys.setprofile` collector, selecting executed code under the installed `litellm` source directory instead of maintaining a function-name allowlist. It prints source locations and qualified function names, including repeated calls. Coroutine resumptions are counted once per invocation. It profiles the current thread and threads created during the call, including the fresh async executor. Existing worker threads are not retroactively profiled; background Python calls may appear, and indentation follows selected Python stack ancestors within each thread + +The Rust runner calls the compiled PyO3 SDK entrypoints with `trace=True`. The existing `FunctionTrace` subscriber collects `#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]` spans for the route entrypoint, preparation, provider lookup, HTTP handler, and selected provider transformations. The shared `http_request` helper instruments the existing Rust send operation without changing clients, timeouts, signing, or error mapping. Function names come from the actual functions. `WithSubscriber` attaches the collector to each future across async polls. Arguments and provider payloads are not recorded in trace events. Uninstrumented functions do not appear; this is scoped instrumentation, not an exhaustive native call graph + +Tracing is opt-in: native calls without `trace=True` keep their original response shape. Traced calls return `{"response": ..., "trace": [{"function": ..., "depth": ...}]}`. The runners print only trace events. Missing native support or empty traces fail instead of falling back to source searching. The old `--repo`, `--signatures`, and `--calls` options are removed + +`profile_python(functions)` still supports direct function references for focused parity checks. `assert_function_trace_parity` compares selected Python events with Rust events supplied by an executable scenario. Successful stage checks prove the declared pipeline ran in a valid dependency order for this fixture; they do not assert identical function contracts, request bodies, responses, streaming behavior, or live-provider correctness + +Build the extension with `maturin develop` in the project's virtual environment. Then run either command above to get the executed function order diff --git a/tests/sdk_function_trace/__init__.py b/tests/sdk_function_trace/__init__.py new file mode 100644 index 00000000000..da62b8041f6 --- /dev/null +++ b/tests/sdk_function_trace/__init__.py @@ -0,0 +1,13 @@ +from tests.sdk_function_trace.harness import ( + TraceScenario, + TraceStep, + assert_function_trace_parity, +) +from tests.sdk_function_trace.profiler import FunctionTraceEvent + +__all__ = [ + "FunctionTraceEvent", + "TraceScenario", + "TraceStep", + "assert_function_trace_parity", +] diff --git a/tests/sdk_function_trace/compare.py b/tests/sdk_function_trace/compare.py new file mode 100644 index 00000000000..941c1b6e067 --- /dev/null +++ b/tests/sdk_function_trace/compare.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import argparse +import os +import sys +from typing import Final + +from tests.sdk_function_trace.fixtures import ROUTES +from tests.sdk_function_trace.report import compare, render + + +def _run(route: str, asynchronous: bool, *, full: bool, colorize: bool) -> bool: + comparison: Final = compare(route, asynchronous=asynchronous) + sys.stdout.write(render(comparison, full=full, colorize=colorize)) + return comparison.passed + + +def main() -> None: + parser: Final = argparse.ArgumentParser(description="Compare Python and Rust SDK pipeline steps per route") + parser.add_argument("--route", choices=("all", *ROUTES), default="all") + mode: Final = parser.add_mutually_exclusive_group() + mode.add_argument("--async", dest="asynchronous", action="store_true", default=True) + mode.add_argument("--sync", dest="asynchronous", action="store_false") + mode.add_argument("--both", action="store_true", help="run async and sync for every selected route") + parser.add_argument( + "--check", action="store_true", help="exit nonzero for missing, extra, or reordered comparable steps" + ) + parser.add_argument( + "--full", action="store_true", help="print every captured runtime event instead of pipeline steps" + ) + args: Final = parser.parse_args() + os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") + colorize: Final = sys.stdout.isatty() and "NO_COLOR" not in os.environ + results: Final = tuple( + _run(selected, selected_mode, full=args.full, colorize=colorize) + for selected in ROUTES + if args.route in ("all", selected) + for selected_mode in ((True, False) if args.both else (args.asynchronous,)) + ) + if args.check and not all(results): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/sdk_function_trace/fixtures.py b/tests/sdk_function_trace/fixtures.py new file mode 100644 index 00000000000..47bbe839627 --- /dev/null +++ b/tests/sdk_function_trace/fixtures.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +import base64 +import io +import json +import wave +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final, Protocol, cast + +from tests.sdk_function_trace.mock_provider import MockProviderResponse +from tests.sdk_function_trace.steps import Engine + +ANTHROPIC_MODEL: Final = "claude-sonnet-5" +OCR_MODEL: Final = "mistral-ocr-latest" +AUDIO_MODEL: Final = "mistral.voxtral-mini-3b-2507" + + +class SdkCall(Protocol): + def __call__(self, **kwargs: object) -> object: ... + + +@dataclass(frozen=True, slots=True) +class Fixture: + kwargs: dict[str, object] + provider_response: MockProviderResponse + + +@dataclass(frozen=True, slots=True) +class RouteSpec: + label: str + python_entrypoints: tuple[str, str] + rust_entrypoints: tuple[str, str] + fixture: Callable[[Engine], Fixture] + + +@dataclass(frozen=True, slots=True) +class Invocation: + function: SdkCall + kwargs: dict[str, object] + provider_response: MockProviderResponse + label: str + + +def audio_bytes() -> bytes: + with io.BytesIO() as buffer: + with wave.open(buffer, "wb") as audio: + audio.setnchannels(1) + audio.setsampwidth(2) + audio.setframerate(16000) + audio.writeframes(b"\x00\x00" * 1600) + return buffer.getvalue() + + +def _anthropic_message_response() -> MockProviderResponse: + body: Final = { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": ANTHROPIC_MODEL, + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 3}, + } + return MockProviderResponse(200, (("content-type", "application/json"),), json.dumps(body).encode()) + + +def _conversation() -> dict[str, object]: + return {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} + + +def _ocr_fixture(engine: Engine) -> Fixture: + return Fixture( + kwargs={ + "model": f"mistral/{OCR_MODEL}", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + **({"optional_params": {"pages": [0]}} if engine == "rust" else {"pages": [0]}), + }, + provider_response=MockProviderResponse( + 200, + (("content-type", "application/json"),), + json.dumps( + { + "pages": [{"index": 0, "markdown": "hello"}], + "model": OCR_MODEL, + "usage_info": {"pages_processed": 1}, + } + ).encode(), + ), + ) + + +def _chat_completions_fixture(engine: Engine) -> Fixture: + conversation: Final = _conversation() + payload: Final = ( + {"messages": conversation["messages"], "optional_params": {"max_tokens": 16}} + if engine == "rust" + else conversation + ) + return Fixture( + kwargs={"model": f"anthropic/{ANTHROPIC_MODEL}", **payload}, + provider_response=_anthropic_message_response(), + ) + + +def _messages_fixture(engine: Engine) -> Fixture: + conversation: Final = _conversation() + payload: Final = {"body": {**conversation, "model": ANTHROPIC_MODEL}} if engine == "rust" else conversation + return Fixture( + kwargs={"model": f"anthropic/{ANTHROPIC_MODEL}", **payload}, + provider_response=_anthropic_message_response(), + ) + + +def _transcription_fixture(engine: Engine) -> Fixture: + credentials: Final = { + "aws_access_key_id": "test-access", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + } + payload: Final = ( + { + "audio": {"data": base64.b64encode(audio_bytes()).decode(), "format": "wav"}, + "optional_params": credentials, + } + if engine == "rust" + else {"file": ("sample.wav", audio_bytes(), "audio/wav"), **credentials} + ) + return Fixture( + kwargs={"model": f"bedrock/{AUDIO_MODEL}", **payload}, + provider_response=MockProviderResponse( + 200, + (("content-type", "application/json"),), + json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, + } + ).encode(), + ), + ) + + +ROUTE_SPECS: Final[dict[str, RouteSpec]] = { + "chat_completions": RouteSpec( + label="anthropic", + python_entrypoints=("completion", "acompletion"), + rust_entrypoints=("chat_completions", "achat_completions"), + fixture=_chat_completions_fixture, + ), + "audio_transcription": RouteSpec( + label="bedrock (Rust-only provider; Python trace covers SDK dispatch)", + python_entrypoints=("transcription", "atranscription"), + rust_entrypoints=("transcription", "atranscription"), + fixture=_transcription_fixture, + ), + "messages": RouteSpec( + label="anthropic", + python_entrypoints=("create", "acreate"), + rust_entrypoints=("messages", "amessages"), + fixture=_messages_fixture, + ), + "ocr": RouteSpec( + label="mistral", + python_entrypoints=("ocr", "aocr"), + rust_entrypoints=("ocr", "aocr"), + fixture=_ocr_fixture, + ), +} + +ROUTES: Final = tuple(ROUTE_SPECS) + + +def sdk_invocation(route: str, *, engine: Engine, asynchronous: bool) -> Invocation: + import litellm + from litellm.anthropic_interface import messages as sdk_messages + from litellm.rust_bridge import get_native_bridge + + rust: Final = engine == "rust" + bridge: Final = get_native_bridge() if rust else None + if rust and bridge is None: + raise RuntimeError("Build the native extension first: maturin develop") + spec: Final = ROUTE_SPECS.get(route) + if spec is None: + raise ValueError(f"Unknown route: {route}") + fixture: Final = spec.fixture(engine) + owner: Final = bridge if rust else (sdk_messages if route == "messages" else litellm) + entrypoint: Final = (spec.rust_entrypoints if rust else spec.python_entrypoints)[int(asynchronous)] + return Invocation( + function=cast(SdkCall, getattr(owner, entrypoint)), + kwargs={ + **fixture.kwargs, + "api_key": "test-key", + **({"trace": True, "timeout_seconds": 5} if rust else {"timeout": 5}), + }, + provider_response=fixture.provider_response, + label=spec.label, + ) diff --git a/tests/sdk_function_trace/harness.py b/tests/sdk_function_trace/harness.py new file mode 100644 index 00000000000..8f707402449 --- /dev/null +++ b/tests/sdk_function_trace/harness.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from types import FunctionType +from typing import Final, cast + +from tests.sdk_function_trace.profiler import FunctionTraceEvent, profile_python + + +@dataclass(frozen=True, slots=True) +class TraceStep: + function: FunctionType + depth: int + + +@dataclass(frozen=True, slots=True) +class TraceScenario: + steps: tuple[TraceStep, ...] + invoke_python: Callable[[], object] + invoke_rust: Callable[[], Sequence[FunctionTraceEvent]] + + +def assert_function_trace_parity(scenario: TraceScenario) -> None: + expected: Final = tuple( + FunctionTraceEvent(function=step.function.__name__, depth=step.depth) for step in scenario.steps + ) + functions: Final = cast(tuple[FunctionType, ...], tuple(step.function for step in scenario.steps)) + with profile_python(functions) as profiler: + scenario.invoke_python() + python_trace: Final = tuple(profiler.events) + rust_trace: Final = tuple(scenario.invoke_rust()) + + if python_trace != expected: + raise AssertionError(f"Python function trace differs: {python_trace!r} != {expected!r}") + if rust_trace != expected: + raise AssertionError(f"Rust function trace differs: {rust_trace!r} != {expected!r}") + if python_trace != rust_trace: + raise AssertionError(f"Python and Rust function traces differ: {python_trace!r} != {rust_trace!r}") diff --git a/tests/sdk_function_trace/mock_provider.py b/tests/sdk_function_trace/mock_provider.py new file mode 100644 index 00000000000..37eca665586 --- /dev/null +++ b/tests/sdk_function_trace/mock_provider.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Lock, Thread +from typing import Final, cast + + +@dataclass(frozen=True, slots=True) +class MockProviderResponse: + status_code: int + headers: tuple[tuple[str, str], ...] + body: bytes + + +class _MockProviderServer(ThreadingHTTPServer): + def __init__(self, response: MockProviderResponse) -> None: + super().__init__(("127.0.0.1", 0), _MockProviderHandler) + self.response: Final = response + self._request_count = 0 + self._request_count_lock: Final = Lock() + + def record_request(self) -> None: + with self._request_count_lock: + self._request_count += 1 + + @property + def request_count(self) -> int: + with self._request_count_lock: + return self._request_count + + +class _MockProviderHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + content_length: Final = int(self.headers.get("content-length", "0")) + self.rfile.read(content_length) + server: Final = cast(_MockProviderServer, self.server) + server.record_request() + self.send_response(server.response.status_code) + for name, value in server.response.headers: + self.send_header(name, value) + self.send_header("content-length", str(len(server.response.body))) + self.end_headers() + self.wfile.write(server.response.body) + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 # matches BaseHTTPRequestHandler + pass + + +@contextmanager +def mock_provider(response: MockProviderResponse) -> Generator[str]: + server: Final = _MockProviderServer(response) + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + host, port = cast(tuple[str, int], server.server_address) + try: + yield f"http://{host}:{port}" + finally: + server.shutdown() + server.server_close() + thread.join() + if server.request_count != 1: + raise AssertionError(f"expected one provider request, received {server.request_count}") diff --git a/tests/sdk_function_trace/ocr-comparison.md b/tests/sdk_function_trace/ocr-comparison.md new file mode 100644 index 00000000000..d252480e218 --- /dev/null +++ b/tests/sdk_function_trace/ocr-comparison.md @@ -0,0 +1,59 @@ +# OCR Python and Rust comparison + +Audited implementation revision: `edcba483b2`. The implementations do not match in function contracts, call structure, or all tested response behavior. This audit changes the source listing coverage, not OCR runtime behavior + +Run both source listings from the repository root: + +```bash +python3 tests/sdk_function_trace/list_python_steps.py --route ocr --signatures --calls +uv run tests/sdk_function_trace/list_rust_steps.py --route ocr --signatures --calls +``` + +Both cover Mistral, Azure AI Mistral, Azure Document Intelligence, Vertex Mistral, and Vertex DeepSeek. Listings show declarations and source call sites, not executed traces + +## Function contracts + +Comparing Python `BaseOCRConfig` with Rust `OcrProviderConfig`, omitting `self` and language-specific ownership details: + +| Python | Rust | Difference | +| --- | --- | --- | +| `get_supported_ocr_params(model)` | `supported_ocr_params()` | Name and model argument | +| `get_api_key_env_var()` | No corresponding method | Missing contract | +| `map_ocr_params(non_default_params, optional_params, model)` | `map_ocr_params(non_default_params)` | Missing accumulator and model | +| `validate_environment(headers, model, api_key, api_base, litellm_params, **kwargs)` | Separate auth/key/header helpers | Different contract | +| `get_complete_url(api_base, model, optional_params, litellm_params, **kwargs)` | `complete_url(api_base, model, optional_params, env_lookup)` | Name and context | +| `transform_ocr_request(model, document, optional_params, headers, **kwargs)` | `transform_ocr_request(model, document, optional_params)` | Missing headers and extra context | +| `async_transform_ocr_request(...)` | No corresponding method | Missing async override | +| `transform_ocr_response(model, raw_response, logging_obj, **kwargs)` | `transform_ocr_response(model, response_json)` | Missing HTTP metadata, logging and extra context | +| `async_transform_ocr_response(...)` | No corresponding method | Missing async override | +| `get_error_class(error_message, status_code, headers)` | Central Rust error mapping | Different contract | + +Python's default mapper returns the supplied `optional_params`; Rust's filters `non_default_params`. Provider overrides must also be compared + +Python maps parameters during SDK preparation, before HTTP-handler environment validation and URL construction. Rust resolves auth and URL before mapping parameters in `prepare_provider_request`. Python has async provider transforms; both native entrypoints execute the same Rust async route using synchronous transform hooks, with polling and document downloading in gateway helpers + +The native bindings also accept `optional_params` and `timeout_seconds`, while the Python SDK accepts `**kwargs` and `timeout`. Public SDK calls with Rust enabled still execute Python preparation before entering Rust, so matching SDK responses would not prove matching standalone Rust steps + +## Runtime results + +Built the native extension from the audited source using `cargo build -p litellm-python-bridge --features extension-module --offline`. Supplied that build's functions through `use_litellm_rust` dependency injection. Ran public `litellm.ocr` and `litellm.aocr` with Rust disabled and enabled against identical local HTTP response fixtures, requiring one request per invocation + +Successful `model_dump()` results and failure exception classes were compared. These checks cover Mistral response outcomes only, not request equality, error messages, live providers, or every execution branch + +| Mistral response fixture | Sync | Async | Observation | +| --- | --- | --- | --- | +| Valid page/model/usage | Match | Match | Same normalized response | +| Model omitted | Match | Match | Both use the requested model | +| `model: null` | Different | Different | Python rejects; Rust uses the requested model | +| `pages: null` | Different | Different | Python rejects; Rust returns an empty array | +| Invalid page element | Match | Match | Both reject during response validation | + +Six of ten fixture/mode comparisons match, four differ. Rust's Mistral response transform conflates missing values with explicit nulls through `as_array`/`as_str` fallbacks. Python preserves explicit nulls into response validation, which rejects them + +## Other provider gaps found in source + +Azure Document Intelligence's Python configuration supports `pages`, `features`, and `req_format`; Rust lists only `pages`. Python normalizes parameters before URL construction; Rust normalizes pages during URL construction + +Python preserves Azure `content`, `tables`, and `keyValuePairs`, and supports retaining the native operation payload. Rust's `OcrResponseData` has no corresponding fields, and its Azure transform does not preserve those values + +Azure and Vertex async document transforms and Azure polling also use different helper contracts. Their runtime equivalence was not tested in this audit diff --git a/tests/sdk_function_trace/profiler.py b/tests/sdk_function_trace/profiler.py new file mode 100644 index 00000000000..c71c74ab0d3 --- /dev/null +++ b/tests/sdk_function_trace/profiler.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import sys +import threading +from collections.abc import Generator, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import CodeType, FrameType, FunctionType +from typing import Final + + +@dataclass(frozen=True, slots=True) +class FunctionTraceEvent: + function: str + depth: int + ancestors: tuple[str, ...] | None = None + + +class PythonProfiler: + def __init__(self, functions: Sequence[FunctionType], source_root: Path | None = None) -> None: + self._source_root: Final = str(source_root.resolve()) + "/" if source_root is not None else None + self._names_by_code: Final = {function.__code__: function.__name__ for function in functions} + self._seen_frames: Final[set[FrameType]] = set() + self.events: Final[list[FunctionTraceEvent]] = [] + + def __call__(self, frame: FrameType, event: str, _arg: object) -> None: + if event != "call" or frame in self._seen_frames: + return + function_name: Final = self.function_name(frame.f_code) + if function_name is None: + return + ancestors: Final = tuple( + name for ancestor in _frame_ancestors(frame) if (name := self.function_name(ancestor.f_code)) is not None + ) + self._seen_frames.add(frame) + self.events.append( + FunctionTraceEvent( + function=function_name, + depth=len(ancestors), + ancestors=ancestors if self._source_root is not None else None, + ) + ) + + def function_name(self, code: CodeType) -> str | None: + if self._source_root is None: + return self._names_by_code.get(code) + if not code.co_filename.startswith(self._source_root): + return None + relative: Final = code.co_filename.removeprefix(self._source_root) + return f"{relative}:{code.co_firstlineno} {getattr(code, 'co_qualname', code.co_name)}" + + +def _frame_ancestors(frame: FrameType) -> Generator[FrameType]: + ancestor: Final = frame.f_back + if ancestor is not None: + yield ancestor + yield from _frame_ancestors(ancestor) + + +@contextmanager +def profile_python( + functions: Sequence[FunctionType] = (), *, source_root: Path | None = None, threads: bool = False +) -> Generator[PythonProfiler]: + profiler: Final = PythonProfiler(functions, source_root) + previous_thread: Final = threading.getprofile() + if threads: + threading.setprofile(profiler) + previous: Final = sys.getprofile() + sys.setprofile(profiler) + try: + yield profiler + finally: + sys.setprofile(previous) + if threads: + threading.setprofile(previous_thread) diff --git a/tests/sdk_function_trace/report.py b/tests/sdk_function_trace/report.py new file mode 100644 index 00000000000..9b654e571f8 --- /dev/null +++ b/tests/sdk_function_trace/report.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +from tests.sdk_function_trace.fixtures import ROUTE_SPECS +from tests.sdk_function_trace.profiler import FunctionTraceEvent +from tests.sdk_function_trace.runtime import ( + TraceDiff, + TraceFailed, + TraceOk, + TraceRun, + TraceSkipped, + attempt_trace, + trace_diff, +) +from tests.sdk_function_trace.steps import Engine, pipeline_issues, pipeline_steps +from tests.sdk_function_trace.table import format_trace_table + +_PYTHON_ONLY_COLOR: Final = "\033[34m" +_RUST_ONLY_COLOR: Final = "\033[33m" +_RESET: Final = "\033[0m" + +_ENGINE_COLOR: Final[dict[Engine, str]] = {"python": _PYTHON_ONLY_COLOR, "rust": _RUST_ONLY_COLOR} + + +@dataclass(frozen=True, slots=True) +class EngineReport: + engine: Engine + run: TraceRun + events: tuple[FunctionTraceEvent, ...] + steps: tuple[FunctionTraceEvent, ...] + issues: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class Comparison: + route: str + label: str + asynchronous: bool + engines: tuple[EngineReport, ...] + diff: TraceDiff + + @property + def comparable(self) -> bool: + return self.route != "audio_transcription" and all(isinstance(report.run, TraceOk) for report in self.engines) + + @property + def passed(self) -> bool: + return ( + (not self.comparable or self.diff.matches) + and not any(report.issues for report in self.engines) + and all(not isinstance(report.run, TraceFailed) for report in self.engines) + ) + + +def _events(run: TraceRun) -> tuple[FunctionTraceEvent, ...]: + match run: + case TraceOk(events=events): + return events + case TraceSkipped() | TraceFailed(): + return () + + +def _engine_report(route: str, engine: Engine, run: TraceRun) -> EngineReport: + events: Final = _events(run) + steps: Final = pipeline_steps(route, engine, events) + issues: Final = pipeline_issues(route, engine, steps) if isinstance(run, TraceOk) else () + return EngineReport(engine=engine, run=run, events=events, steps=steps, issues=issues) + + +def compare(route: str, *, asynchronous: bool) -> Comparison: + runs: Final = { + engine: attempt_trace(route, engine=engine, asynchronous=asynchronous) for engine in ("python", "rust") + } + engines: Final = tuple(_engine_report(route, engine, run) for engine, run in runs.items()) + return Comparison( + route=route, + label=ROUTE_SPECS[route].label, + asynchronous=asynchronous, + engines=engines, + diff=trace_diff(engines[0].steps, engines[1].steps), + ) + + +def _tree_line(event: FunctionTraceEvent, only: frozenset[str], marker: str, color: str, *, colorize: bool) -> str: + line: Final = f"{' ' * event.depth}{event.function}" + (f" {marker}" if event.function in only else "") + return f"{color}{line}{_RESET}\n" if colorize and event.function in only else f"{line}\n" + + +def _tree_lines( + events: tuple[FunctionTraceEvent, ...], + only: frozenset[str], + marker: str, + color: str, + *, + colorize: bool, +) -> tuple[str, ...]: + return tuple(_tree_line(event, only, marker, color, colorize=colorize) for event in events) + + +def _engine_lines( + report: EngineReport, diff: TraceDiff, *, comparable: bool, full: bool, colorize: bool +) -> tuple[str, ...]: + match report.run: + case TraceSkipped(reason=reason): + return (f"{report.engine}: SKIP ({reason})\n\n",) + case TraceFailed(reason=reason): + return (f"{report.engine}: FAIL ({reason})\n\n",) + case TraceOk(): + shown: Final = report.events if full else report.steps + only: Final = ( + () if full or not comparable else (diff.python_only if report.engine == "python" else diff.rust_only) + ) + return ( + f"{report.engine} ({len(shown)} steps)\n\n", + *_tree_lines( + shown, + frozenset(only), + f"<- {report.engine} only", + _ENGINE_COLOR[report.engine], + colorize=colorize, + ), + "\n", + ) + + +def _parity_lines(comparison: Comparison) -> tuple[str, ...]: + if not comparison.comparable: + if comparison.route == "audio_transcription": + return ("step parity: UNAVAILABLE (Bedrock transcription has no independent Python implementation)\n",) + return ("step parity: UNAVAILABLE (both engines must complete)\n",) + diff: Final = comparison.diff + order: Final = "the same" if diff.shared_order_matches else "a different" + return ( + "diff\n\n", + f"shared steps appear in {order} order\n", + f"python-only: {', '.join(diff.python_only) or 'none'}\n", + f"rust-only: {', '.join(diff.rust_only) or 'none'}\n\n", + f"step parity: {'PASS' if diff.matches else 'FAIL'}\n", + ) + + +def _stage_lines(comparison: Comparison) -> tuple[str, ...]: + return tuple( + f"{report.engine} " + f"{'SDK dispatch only' if comparison.route == 'audio_transcription' and report.engine == 'python' else 'pipeline'}: " + f"{'FAIL: ' + '; '.join(report.issues) if report.issues else 'PASS'}\n" + for report in comparison.engines + if isinstance(report.run, TraceOk) + ) + + +def render(comparison: Comparison, *, full: bool, colorize: bool) -> str: + mode: Final = "async" if comparison.asynchronous else "sync" + traces: Final = ( + (format_trace_table(comparison.engines[0].steps, comparison.engines[1].steps, colorize=colorize) + "\n\n",) + if not full and all(isinstance(report.run, TraceOk) for report in comparison.engines) + else tuple( + line + for report in comparison.engines + for line in _engine_lines( + report, comparison.diff, comparable=comparison.comparable, full=full, colorize=colorize + ) + ) + ) + return "".join( + ( + f"route: {comparison.route} provider: {comparison.label} mode: {mode}\n\n", + *traces, + *_parity_lines(comparison), + *_stage_lines(comparison), + "Each successful invocation issued exactly one local provider request\n\n", + ) + ) diff --git a/tests/sdk_function_trace/route-comparison.md b/tests/sdk_function_trace/route-comparison.md new file mode 100644 index 00000000000..009d3544d05 --- /dev/null +++ b/tests/sdk_function_trace/route-comparison.md @@ -0,0 +1,26 @@ +# SDK route trace audit + +Run the four native HTTP route families in both modes from the repository root: + +```bash +uv run python -m tests.sdk_function_trace.compare --route all --both --check +``` + +The local fixture matrix on 2026-09-02 completed 15 successful engine invocations and one expected skip. Every successful invocation issued exactly one local HTTP request. All five comparable route/mode pairs have identical canonical steps in the same order, with no Python-only or Rust-only steps + +| Route | Python async | Python sync | Rust async | Rust sync | +| --- | --- | --- | --- | --- | +| Chat completions, Anthropic | Pass | Pass | Pass | Pass | +| Messages, Anthropic | Pass | Unsupported, skipped | Pass | Pass | +| OCR, Mistral | Pass | Pass | Pass | Pass | +| Audio transcription, Bedrock | Dispatch only | Dispatch only | Pass | Pass | + +The same canonical step sequence ran in sync and async for each engine with both modes available. Bedrock transcription's Python SDK delegates to Rust, so its two successful calls do not establish independent provider parity. Realtime and Responses WebSockets are outside this HTTP fixture runner + +Chat and OCR also have identical projected nesting in both modes. Async Messages has the same helper nesting beneath its handler, but Python starts that handler on a worker thread, so it appears as a second root. The comparison preserves this physical thread boundary and checks step order independently of absolute depth + +Rust now resolves chat providers and supported parameters before entering its handler. Chat and Messages validate the environment and transform requests inside their handlers. Messages builds the final URL after transformation. OCR resolves its config and maps supported parameters during preparation, then validates credentials, builds the URL, and transforms the request inside its handler. Its during-call guardrails still run before HTTP, within the provider-call lifecycle phase + +The environment hooks execute credential and header validation. Chat's supported-parameter hooks return OpenAI names paired with provider names and feed the existing request acceptance checks. The direct Rust API still accepts provider-mapped parameters, and its supported subset is smaller than Python's. Matching the pipeline does not establish identical parameter contracts + +`--check` now fails if either comparable engine has missing, extra, or reordered canonical steps, even if its individual stage checks pass. Bedrock transcription and sync Messages report `UNAVAILABLE` for cross-language parity; native execution is still checked. Passing establishes step coverage and order for one non-streaming fixture per route, not complete request, response, error, or provider parity. The previously recorded OCR response gaps remain in `ocr-comparison.md` diff --git a/tests/sdk_function_trace/runtime.py b/tests/sdk_function_trace/runtime.py new file mode 100644 index 00000000000..d5bf15694bc --- /dev/null +++ b/tests/sdk_function_trace/runtime.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import asyncio +import os +from collections.abc import Awaitable, Generator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final, cast +from unittest.mock import patch + +from pydantic import BaseModel, ConfigDict + +from tests.sdk_function_trace.fixtures import Invocation, sdk_invocation +from tests.sdk_function_trace.mock_provider import mock_provider +from tests.sdk_function_trace.profiler import FunctionTraceEvent, profile_python +from tests.sdk_function_trace.steps import Engine + + +class TraceEventPayload(BaseModel): + model_config = ConfigDict(strict=True, extra="forbid") + function: str + depth: int + + +class TraceResponsePayload(BaseModel): + model_config = ConfigDict(strict=True, extra="forbid") + response: object + trace: tuple[TraceEventPayload, ...] | list[TraceEventPayload] + + +@contextmanager +def _python_engine() -> Generator[None]: + from litellm.rust_bridge import ocr as ocr_bridge + + previous_ocr: Final = ocr_bridge.rust_ocr_enabled() + with patch.dict(os.environ, {"LITELLM_RUST": "false"}): + ocr_bridge.use_litellm_rust(False) + try: + yield + finally: + ocr_bridge.use_litellm_rust(previous_ocr) + + +def _invoke(case: Invocation, api_base: str, *, asynchronous: bool) -> object: + async def invoke_async() -> object: + return await cast("Awaitable[object]", case.function(**case.kwargs, api_base=api_base)) + + if asynchronous: + return asyncio.run(invoke_async()) + return case.function(**case.kwargs, api_base=api_base) + + +def collect(case: Invocation, api_base: str, *, engine: Engine, asynchronous: bool) -> tuple[FunctionTraceEvent, ...]: + import litellm + + if engine == "rust": + payload: Final = TraceResponsePayload.model_validate(_invoke(case, api_base, asynchronous=asynchronous)) + return tuple(FunctionTraceEvent(event.function, event.depth) for event in payload.trace) + with profile_python(source_root=Path(litellm.__file__).parent, threads=True) as profiler: + _invoke(case, api_base, asynchronous=asynchronous) + return tuple(profiler.events) + + +def run_trace(route: str, *, engine: Engine, asynchronous: bool = False) -> tuple[FunctionTraceEvent, ...]: + case: Final = sdk_invocation(route, engine=engine, asynchronous=asynchronous) + with _python_engine(), mock_provider(case.provider_response) as api_base: + events: Final = collect(case, api_base, engine=engine, asynchronous=asynchronous) + if not events: + raise RuntimeError(f"No runtime events for {route}; rebuild the native extension with tracing support") + return events + + +@dataclass(frozen=True, slots=True) +class TraceOk: + events: tuple[FunctionTraceEvent, ...] + + +@dataclass(frozen=True, slots=True) +class TraceSkipped: + reason: str + + +@dataclass(frozen=True, slots=True) +class TraceFailed: + reason: str + + +TraceRun = TraceOk | TraceSkipped | TraceFailed + + +def attempt_trace(route: str, *, engine: Engine, asynchronous: bool) -> TraceRun: + try: + return TraceOk(run_trace(route, engine=engine, asynchronous=asynchronous)) + except Exception as error: + reason: Final = f"{type(error).__name__}: {error}" + if ( + route == "messages" + and engine == "python" + and not asynchronous + and isinstance(error, ValueError) + and str(error) == "anthropic_messages_handler is not implemented for sync calls" + ): + return TraceSkipped(reason) + return TraceFailed(reason) + + +@dataclass(frozen=True, slots=True) +class TraceDiff: + python_only: tuple[str, ...] + rust_only: tuple[str, ...] + shared_order_matches: bool + + @property + def matches(self) -> bool: + return not self.python_only and not self.rust_only and self.shared_order_matches + + +def trace_diff(python: tuple[FunctionTraceEvent, ...], rust: tuple[FunctionTraceEvent, ...]) -> TraceDiff: + python_names: Final = {event.function for event in python} + rust_names: Final = {event.function for event in rust} + shared_python: Final = tuple(event.function for event in python if event.function in rust_names) + shared_rust: Final = tuple(event.function for event in rust if event.function in python_names) + return TraceDiff( + python_only=tuple(event.function for event in python if event.function not in rust_names), + rust_only=tuple(event.function for event in rust if event.function not in python_names), + shared_order_matches=bool(shared_python) and shared_python == shared_rust, + ) diff --git a/tests/sdk_function_trace/steps.py b/tests/sdk_function_trace/steps.py new file mode 100644 index 00000000000..bb50d4ebe57 --- /dev/null +++ b/tests/sdk_function_trace/steps.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import re +from collections.abc import Sequence +from dataclasses import dataclass +from functools import reduce +from typing import Final, Literal + +from tests.sdk_function_trace.profiler import FunctionTraceEvent + +Engine = Literal["python", "rust"] + + +@dataclass(frozen=True, slots=True) +class Step: + name: str + python: re.Pattern[str] | None + rust: str | None + + +def _step(name: str, python: str | None = None, rust: str | None = None) -> Step: + return Step(name, re.compile(python) if python is not None else None, rust) + + +_POST: Final = r"AsyncHTTPHandler\.post$|HTTPHandler\.post$" + +STEPS: Final[dict[str, tuple[Step, ...]]] = { + "ocr": ( + _step("ocr", r"ocr/main\.py:\d+ a?ocr$", "ocr"), + _step("prepare_ocr_call", r"ocr/main\.py:\d+ _prepare_ocr_request$", "prepare_ocr_call"), + _step("get_provider_ocr_config", r"ProviderConfigManager\.get_provider_ocr_config$", "ocr_provider_config"), + _step("supported_ocr_params", r"get_supported_ocr_params$", "supported_ocr_params"), + _step("map_ocr_params", r"(? tuple[str, ...]: + names: Final = tuple(event.function for event in events) + required: Final = tuple(step.name for step in STEPS[route] if getattr(step, engine) is not None) + missing: Final = tuple(f"missing {name}" for name in required if name not in names) + provider: Final = next(name for name in required if name.startswith("get_provider_")) + handler: Final = next(name for name in required if name.startswith("execute_")) + dispatch_only: Final = route == "audio_transcription" and engine == "python" + request: Final = next( + (name for name in required if name.startswith("transform_") and name.endswith("request")), handler + ) + response: Final = next( + (name for name in required if name.startswith("transform_") and name.endswith("response")), handler + ) + phases: Final = ( + (route, "map_transcription_params", provider, handler) + if dispatch_only + else (route, provider, request, "http_request", response) + ) + extra_edges: Final = ( + () + if dispatch_only + else ( + (handler, "http_request"), + *((name, request) for name in required if name.startswith(("map_", "supported_"))), + *((name, "http_request") for name in ("validate_environment", "complete_url") if name in required), + ) + ) + edges: Final = (*zip(phases, phases[1:]), *extra_edges) + return missing + tuple( + f"{before} must precede {after}" + for before, after in edges + if before in names and after in names and names.index(before) >= names.index(after) + ) + + +def _canonical_name(route: str, engine: Engine, function: str) -> str | None: + for step in STEPS[route]: + if engine == "python": + if step.python is not None and step.python.search(function): + return step.name + elif step.rust is not None and function == step.rust: + return step.name + return function if engine == "rust" else None + + +@dataclass(frozen=True, slots=True) +class _Projection: + shown: tuple[FunctionTraceEvent, ...] = () + stack: tuple[tuple[int, int], ...] = () + seen: frozenset[str] = frozenset() + + +def _project(route: str, engine: Engine, state: _Projection, event: FunctionTraceEvent) -> _Projection: + stack: Final = tuple(pair for pair in state.stack if event.depth > pair[0]) + name: Final = _canonical_name(route, engine, event.function) + if name is None or name in state.seen: + return _Projection(state.shown, stack, state.seen) + depth: Final = ( + next( + ( + kept.depth + 1 + for ancestor in event.ancestors + for kept in state.shown + if kept.function == _canonical_name(route, engine, ancestor) + ), + 0, + ) + if event.ancestors is not None + else stack[-1][1] + 1 + if stack + else 0 + ) + return _Projection( + state.shown + (FunctionTraceEvent(function=name, depth=depth),), + stack + ((event.depth, depth),), + state.seen | {name}, + ) + + +def pipeline_steps(route: str, engine: Engine, events: Sequence[FunctionTraceEvent]) -> tuple[FunctionTraceEvent, ...]: + projection: Final = reduce(lambda state, event: _project(route, engine, state, event), events, _Projection()) + return projection.shown diff --git a/tests/sdk_function_trace/table.py b/tests/sdk_function_trace/table.py new file mode 100644 index 00000000000..2124d7e3faf --- /dev/null +++ b/tests/sdk_function_trace/table.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from collections.abc import Iterator +from difflib import SequenceMatcher +from typing import Final + +from tests.sdk_function_trace.profiler import FunctionTraceEvent + + +def _aligned_rows( + python: tuple[FunctionTraceEvent, ...], rust: tuple[FunctionTraceEvent, ...] +) -> Iterator[tuple[FunctionTraceEvent | None, FunctionTraceEvent | None]]: + matcher: Final = SequenceMatcher( + a=tuple(event.function for event in python), + b=tuple(event.function for event in rust), + autojunk=False, + ) + for tag, python_start, python_end, rust_start, rust_end in matcher.get_opcodes(): + if tag == "equal": + yield from zip(python[python_start:python_end], rust[rust_start:rust_end]) + else: + yield from ((event, None) for event in python[python_start:python_end]) + yield from ((None, event) for event in rust[rust_start:rust_end]) + + +def _label(event: FunctionTraceEvent | None) -> str: + return f"{' ' * event.depth}{event.function}" if event is not None else "" + + +def _status( + python: FunctionTraceEvent | None, + rust: FunctionTraceEvent | None, + python_names: frozenset[str], + rust_names: frozenset[str], +) -> tuple[str, str]: + if python is not None and rust is not None: + return "match", "\033[32m" + if python is not None: + return ("reordered", "\033[31m") if python.function in rust_names else ("python only", "\033[34m") + if rust is not None: + return ("reordered", "\033[31m") if rust.function in python_names else ("rust only", "\033[33m") + return "", "" + + +def format_trace_table( + python: tuple[FunctionTraceEvent, ...], + rust: tuple[FunctionTraceEvent, ...], + *, + colorize: bool, +) -> str: + python_header: Final = f"python ({len(python)} steps)" + rust_header: Final = f"rust ({len(rust)} steps)" + python_width: Final = max(len(python_header), *(len(_label(event)) for event in python), 0) + rust_width: Final = max(len(rust_header), *(len(_label(event)) for event in rust), 0) + python_names: Final = frozenset(event.function for event in python) + rust_names: Final = frozenset(event.function for event in rust) + border: Final = f"+-{'-' * python_width}-+-{'-' * rust_width}-+-------------+" + rows: Final = tuple( + f"{color}{line}\033[0m" if colorize else line + for left, right in _aligned_rows(python, rust) + for status, color in (_status(left, right, python_names, rust_names),) + for line in (f"| {_label(left):<{python_width}} | {_label(right):<{rust_width}} | {status:<11} |",) + ) + return "\n".join( + ( + border, + f"| {python_header:<{python_width}} | {rust_header:<{rust_width}} | {'comparison':<11} |", + border, + *rows, + border, + ) + ) diff --git a/tests/sdk_function_trace/test_mock_provider.py b/tests/sdk_function_trace/test_mock_provider.py new file mode 100644 index 00000000000..88d7d5392d0 --- /dev/null +++ b/tests/sdk_function_trace/test_mock_provider.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from contextlib import ExitStack +from typing import Final +from urllib.error import HTTPError +from urllib.request import Request, urlopen + +import pytest + +from tests.sdk_function_trace.mock_provider import MockProviderResponse, mock_provider + + +def test_mock_provider_preserves_error_response() -> None: + response: Final = MockProviderResponse(429, (("retry-after", "2"),), b'{"error":"rate limited"}') + with mock_provider(response) as api_base: + with pytest.raises(HTTPError) as error: + urlopen(Request(api_base, data=b"{}"), timeout=5) + with error.value as received: + assert received.code == 429 + assert received.headers["retry-after"] == "2" + assert received.read() == response.body + + +@pytest.mark.parametrize("request_count", [0, 2]) +def test_mock_provider_rejects_missing_or_duplicate_requests(request_count: int) -> None: + response: Final = MockProviderResponse(200, (), b"{}") + with ExitStack() as stack: + api_base: Final = stack.enter_context(mock_provider(response)) + for _ in range(request_count): + with urlopen(Request(api_base, data=b"{}"), timeout=5) as received: + assert received.read() == response.body + with pytest.raises(AssertionError, match=f"expected one provider request, received {request_count}"): + stack.close() diff --git a/tests/sdk_function_trace/test_profiler.py b/tests/sdk_function_trace/test_profiler.py new file mode 100644 index 00000000000..10a266fb1e8 --- /dev/null +++ b/tests/sdk_function_trace/test_profiler.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path +from types import FunctionType +from typing import Final, cast + +import pytest + +from tests.sdk_function_trace import ( + FunctionTraceEvent, + TraceScenario, + TraceStep, + assert_function_trace_parity, +) +from tests.sdk_function_trace.profiler import profile_python + + +class First: + @staticmethod + def run() -> None: + return None + + +class Second: + @staticmethod + def run() -> None: + return None + + +def test_profiler_matches_code_objects_and_keeps_repeated_calls() -> None: + with profile_python((First.run,)) as profiler: + Second.run() + First.run() + First.run() + + assert profiler.events == [ + FunctionTraceEvent(function="run", depth=0), + FunctionTraceEvent(function="run", depth=0), + ] + + +def test_profiler_records_selected_function_nesting_depth() -> None: + class Nested: + @staticmethod + def run() -> None: + First.run() + + with profile_python((Nested.run, First.run)) as profiler: + Nested.run() + + assert profiler.events == [ + FunctionTraceEvent(function="run", depth=0), + FunctionTraceEvent(function="run", depth=1), + ] + + +def test_profiler_restores_previous_profiler_after_failure() -> None: + previous: Final = sys.getprofile() + + with profile_python((First.run,)) as outer: + with pytest.raises(RuntimeError, match="stop"): + with profile_python((Second.run,)): + raise RuntimeError("stop") + assert sys.getprofile() is outer + First.run() + + assert sys.getprofile() is previous + assert outer.events == [FunctionTraceEvent(function="run", depth=0)] + + +def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None: + async def suspended() -> None: + await asyncio.sleep(0) + First.run() + await asyncio.sleep(0) + + with profile_python((suspended, First.run)) as profiler: + asyncio.run(suspended()) + + assert profiler.events == [ + FunctionTraceEvent(function="suspended", depth=0), + FunctionTraceEvent(function="run", depth=1), + ] + + +def test_source_profiler_records_real_frame_ancestry() -> None: + def outer() -> None: + First.run() + + with profile_python(source_root=Path(__file__).parent) as profiler: + outer() + Second.run() + + outer_event, first_event, second_event = ( + event for event in profiler.events if event.function.startswith("test_profiler.py:") + ) + assert first_event.ancestors is not None + assert outer_event.function in first_event.ancestors + assert second_event.ancestors is not None + assert outer_event.function not in second_event.ancestors + + +@pytest.mark.parametrize( + "rust_trace", + [ + (), + (FunctionTraceEvent(function="renamed", depth=0),), + (FunctionTraceEvent(function="run", depth=1),), + (FunctionTraceEvent(function="run", depth=0),) * 2, + ], + ids=["missing", "renamed", "wrong-depth", "extra-call"], +) +def test_harness_rejects_rust_function_trace_drift(rust_trace: tuple[FunctionTraceEvent, ...]) -> None: + with pytest.raises(AssertionError, match="Rust function trace differs"): + assert_function_trace_parity( + TraceScenario( + steps=(TraceStep(cast(FunctionType, First.run), depth=0),), + invoke_python=First.run, + invoke_rust=lambda: rust_trace, + ) + ) + + +def test_harness_rejects_python_function_trace_drift() -> None: + with pytest.raises(AssertionError, match="Python function trace differs"): + assert_function_trace_parity( + TraceScenario( + steps=(TraceStep(cast(FunctionType, First.run), depth=0),), + invoke_python=Second.run, + invoke_rust=lambda: (FunctionTraceEvent(function="run", depth=0),), + ) + ) + + +def test_harness_accepts_matching_traces() -> None: + assert_function_trace_parity( + TraceScenario( + steps=(TraceStep(cast(FunctionType, First.run), depth=0),), + invoke_python=First.run, + invoke_rust=lambda: (FunctionTraceEvent(function="run", depth=0),), + ) + ) + + +def test_harness_rejects_reordered_calls() -> None: + def begin() -> None: + return None + + def finish() -> None: + return None + + with pytest.raises(AssertionError, match="Rust function trace differs"): + assert_function_trace_parity( + TraceScenario( + steps=( + TraceStep(cast(FunctionType, begin), depth=0), + TraceStep(cast(FunctionType, finish), depth=0), + ), + invoke_python=lambda: (begin(), finish()), + invoke_rust=lambda: ( + FunctionTraceEvent(function="finish", depth=0), + FunctionTraceEvent(function="begin", depth=0), + ), + ) + ) diff --git a/tests/sdk_function_trace/test_runtime.py b/tests/sdk_function_trace/test_runtime.py new file mode 100644 index 00000000000..015cba55083 --- /dev/null +++ b/tests/sdk_function_trace/test_runtime.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from typing import Final + +import pytest + +from tests.sdk_function_trace.runtime import ( + TraceFailed, + TraceSkipped, + attempt_trace, + run_trace, + trace_diff, +) +from tests.sdk_function_trace.steps import pipeline_issues, pipeline_steps + + +def test_sync_messages_records_the_known_python_limitation() -> None: + result: Final = attempt_trace("messages", engine="python", asynchronous=False) + + assert isinstance(result, TraceSkipped) + assert result.reason == "ValueError: anthropic_messages_handler is not implemented for sync calls" + + +def test_unexpected_call_failure_is_not_skipped() -> None: + result: Final = attempt_trace("unknown", engine="python", asynchronous=False) + + assert isinstance(result, TraceFailed) + assert result.reason == "ValueError: Unknown route: unknown" + + +@pytest.mark.parametrize( + ("route", "asynchronous"), + (("chat_completions", False), ("chat_completions", True), ("messages", True), ("ocr", False), ("ocr", True)), +) +def test_compiled_routes_match_python_steps(route: str, asynchronous: bool) -> None: + from litellm.rust_bridge import get_native_bridge + + if get_native_bridge() is None: + pytest.skip("build the native bridge to run executed route parity") + python: Final = pipeline_steps(route, "python", run_trace(route, engine="python", asynchronous=asynchronous)) + rust: Final = pipeline_steps(route, "rust", run_trace(route, engine="rust", asynchronous=asynchronous)) + + assert pipeline_issues(route, "python", python) == () + assert pipeline_issues(route, "rust", rust) == () + assert trace_diff(python, rust).matches + if route != "messages": + assert python == rust diff --git a/tests/sdk_function_trace/test_steps.py b/tests/sdk_function_trace/test_steps.py new file mode 100644 index 00000000000..b5432951187 --- /dev/null +++ b/tests/sdk_function_trace/test_steps.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +from typing import Final + +import pytest + +from tests.sdk_function_trace.profiler import FunctionTraceEvent +from tests.sdk_function_trace.runtime import trace_diff +from tests.sdk_function_trace.steps import pipeline_issues, pipeline_steps + + +def test_python_ocr_projection_keeps_pipeline_and_drops_noise() -> None: + events: Final = ( + FunctionTraceEvent("utils.py:1747 client..wrapper_async", 0), + FunctionTraceEvent("ocr/main.py:331 aocr", 1), + FunctionTraceEvent("ocr/main.py:70 _prepare_ocr_request", 2), + FunctionTraceEvent("litellm_core_utils/get_llm_provider_logic.py:142 get_llm_provider", 3), + FunctionTraceEvent("utils.py:9303 ProviderConfigManager.get_provider_ocr_config", 3), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:34 MistralOCRConfig.get_supported_ocr_params", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:72 MistralOCRConfig.map_ocr_params", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:34 MistralOCRConfig.get_supported_ocr_params", 5), + FunctionTraceEvent("llms/custom_httpx/llm_http_handler.py:1705 BaseLLMHTTPHandler.async_ocr", 2), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:94 MistralOCRConfig.validate_environment", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:124 MistralOCRConfig.get_complete_url", 4), + FunctionTraceEvent("llms/base_llm/ocr/transformation.py:209 BaseOCRConfig.async_transform_ocr_request", 5), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:149 MistralOCRConfig.transform_ocr_request", 6), + FunctionTraceEvent("llms/custom_httpx/http_handler.py:654 AsyncHTTPHandler.post", 6), + FunctionTraceEvent("llms/base_llm/ocr/transformation.py:255 BaseOCRConfig.async_transform_ocr_response", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:200 MistralOCRConfig.transform_ocr_response", 5), + FunctionTraceEvent("cost_calculator.py:1874 ocr_cost", 6), + ) + + assert pipeline_steps("ocr", "python", events) == ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("prepare_ocr_call", 1), + FunctionTraceEvent("get_provider_ocr_config", 2), + FunctionTraceEvent("supported_ocr_params", 3), + FunctionTraceEvent("map_ocr_params", 3), + FunctionTraceEvent("execute_ocr_provider_call", 1), + FunctionTraceEvent("validate_environment", 2), + FunctionTraceEvent("complete_url", 2), + FunctionTraceEvent("transform_ocr_request", 3), + FunctionTraceEvent("http_request", 3), + FunctionTraceEvent("transform_ocr_response", 2), + ) + + +def test_rust_ocr_projection_reuses_step_names_and_keeps_unknown_spans() -> None: + events: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("prepare_ocr_call", 1), + FunctionTraceEvent("map_ocr_params", 2), + FunctionTraceEvent("supported_ocr_params", 3), + FunctionTraceEvent("map_ocr_params", 2), + FunctionTraceEvent("transform_ocr_request", 2), + FunctionTraceEvent("execute_ocr_provider_call", 2), + FunctionTraceEvent("transform_ocr_response", 3), + FunctionTraceEvent("new_uninstrumented_span", 3), + ) + + assert pipeline_steps("ocr", "rust", events) == ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("prepare_ocr_call", 1), + FunctionTraceEvent("map_ocr_params", 2), + FunctionTraceEvent("supported_ocr_params", 3), + FunctionTraceEvent("transform_ocr_request", 2), + FunctionTraceEvent("execute_ocr_provider_call", 2), + FunctionTraceEvent("transform_ocr_response", 3), + FunctionTraceEvent("new_uninstrumented_span", 3), + ) + + +def test_projection_resets_depth_on_thread_root() -> None: + events: Final = ( + FunctionTraceEvent("main.py:387 acompletion", 1), + FunctionTraceEvent("llms/anthropic/chat/handler.py:255 AnthropicChatCompletion.acompletion_function", 2), + FunctionTraceEvent( + "llms/anthropic/experimental_pass_through/messages/handler.py:416 anthropic_messages_handler", 0 + ), + FunctionTraceEvent( + "llms/anthropic/experimental_pass_through/messages/transformation.py:575" + " AnthropicMessagesConfig.transform_anthropic_messages_request", + 4, + ), + ) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("chat_completions", 0), + FunctionTraceEvent("execute_chat_completions_provider_call", 1), + ) + assert pipeline_steps("messages", "python", events) == ( + FunctionTraceEvent("execute_messages_provider_call", 0), + FunctionTraceEvent("transform_request", 1), + ) + + +@pytest.mark.parametrize("function", ("completion", "completion_function", "acompletion_function")) +def test_chat_projection_includes_sync_and_async_handlers(function: str) -> None: + events: Final = (FunctionTraceEvent(f"llms/anthropic/chat/handler.py:100 AnthropicChatCompletion.{function}", 0),) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("execute_chat_completions_provider_call", 0), + ) + + +def test_trace_diff_reports_no_difference_for_identical_steps() -> None: + steps: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("transform_ocr_request", 1), + ) + + diff: Final = trace_diff(steps, steps) + + assert diff.python_only == () + assert diff.rust_only == () + assert diff.shared_order_matches + assert diff.matches + + +def test_trace_diff_reports_exclusive_steps_and_reordered_shared_steps() -> None: + python: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("supported_ocr_params", 1), + FunctionTraceEvent("map_ocr_params", 1), + FunctionTraceEvent("http_request", 2), + ) + rust: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("map_ocr_params", 1), + FunctionTraceEvent("supported_ocr_params", 2), + FunctionTraceEvent("transform_ocr_response", 2), + ) + + diff: Final = trace_diff(python, rust) + + assert diff.python_only == ("http_request",) + assert diff.rust_only == ("transform_ocr_response",) + assert not diff.shared_order_matches + assert not diff.matches + + +def test_trace_diff_does_not_claim_empty_or_disjoint_traces_match() -> None: + assert not trace_diff((), ()).shared_order_matches + assert not trace_diff((FunctionTraceEvent("ocr", 0),), (FunctionTraceEvent("messages", 0),)).shared_order_matches + + +def test_projection_uses_actual_ancestors_after_coroutine_resumption() -> None: + entrypoint: Final = "main.py:387 acompletion" + handler: Final = "llms/anthropic/chat/handler.py:255 AnthropicChatCompletion.acompletion_function" + events: Final = ( + FunctionTraceEvent(entrypoint, 0, ()), + FunctionTraceEvent(handler, 1, (entrypoint,)), + FunctionTraceEvent("utils.py:100 unrelated_worker", 0, ()), + FunctionTraceEvent("llms/anthropic/chat/transformation.py:100 transform_response", 1, (handler,)), + ) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("chat_completions", 0), + FunctionTraceEvent("execute_chat_completions_provider_call", 1), + FunctionTraceEvent("transform_response", 2), + ) + + +def test_projection_does_not_nest_siblings_under_a_returned_config_lookup() -> None: + events: Final = ( + FunctionTraceEvent("main.py:387 completion", 0), + FunctionTraceEvent("utils.py:100 ProviderConfigManager.get_provider_chat_config", 1), + FunctionTraceEvent("utils.py:200 unrelated_helper", 1), + FunctionTraceEvent("llms/anthropic/chat/transformation.py:100 transform_request", 2), + ) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("chat_completions", 0), + FunctionTraceEvent("get_provider_chat_config", 1), + FunctionTraceEvent("transform_request", 1), + ) + + +CHAT_RUST_STEPS: Final = ( + "chat_completions", + "get_provider_chat_config", + "supported_openai_params", + "execute_chat_completions_provider_call", + "validate_environment", + "transform_request", + "http_request", + "transform_response", +) + + +@pytest.mark.parametrize("missing", CHAT_RUST_STEPS) +def test_pipeline_check_rejects_missing_stages(missing: str) -> None: + steps: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS if name != missing) + + assert f"missing {missing}" in pipeline_issues("chat_completions", "rust", steps) + + +def test_pipeline_check_rejects_http_before_request_transformation() -> None: + steps: Final = tuple( + FunctionTraceEvent(name, 0) + for name in ( + "chat_completions", + "get_provider_chat_config", + "supported_openai_params", + "execute_chat_completions_provider_call", + "validate_environment", + "http_request", + "transform_request", + "transform_response", + ) + ) + + assert "transform_request must precede http_request" in pipeline_issues("chat_completions", "rust", steps) + + +def test_step_parity_rejects_different_handler_boundaries_even_with_valid_stages() -> None: + rust: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS) + python: Final = tuple( + FunctionTraceEvent(name, 0) + for name in ( + "chat_completions", + "get_provider_chat_config", + "supported_openai_params", + "validate_environment", + "transform_request", + "execute_chat_completions_provider_call", + "http_request", + "transform_response", + ) + ) + + assert not trace_diff(python, rust).shared_order_matches + assert not trace_diff(python, rust).matches + assert pipeline_issues("chat_completions", "python", python) == () + assert pipeline_issues("chat_completions", "rust", rust) == () + + +def test_step_parity_rejects_an_exclusive_helper_with_matching_shared_order() -> None: + rust: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS) + python: Final = (*rust, FunctionTraceEvent("unmatched_helper", 0)) + diff: Final = trace_diff(python, rust) + + assert diff.shared_order_matches + assert not diff.matches diff --git a/tests/sdk_function_trace/test_table.py b/tests/sdk_function_trace/test_table.py new file mode 100644 index 00000000000..c2341a391a9 --- /dev/null +++ b/tests/sdk_function_trace/test_table.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import re +from typing import Final + +from tests.sdk_function_trace.profiler import FunctionTraceEvent +from tests.sdk_function_trace.table import format_trace_table + + +def test_table_aligns_matches_after_missing_steps_and_preserves_indentation() -> None: + python: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("python_helper", 1), + FunctionTraceEvent("http_request", 2), + ) + rust: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("rust_helper", 1), + FunctionTraceEvent("http_request", 1), + ) + output: Final = format_trace_table(python, rust, colorize=False) + rows: Final = tuple(line.split("|")[1:-1] for line in output.splitlines() if line.startswith("|")) + + assert tuple(tuple(cell.strip() for cell in row) for row in rows) == ( + ("python (3 steps)", "rust (3 steps)", "comparison"), + ("ocr", "ocr", "match"), + ("python_helper", "", "python only"), + ("", "rust_helper", "rust only"), + ("http_request", "http_request", "match"), + ) + assert rows[-1][0].startswith(" http_request") + assert rows[-1][1].startswith(" http_request") + assert len({len(line) for line in output.splitlines()}) == 1 + assert "\033[" not in output + + +def test_table_marks_reordered_calls_and_keeps_both_execution_orders() -> None: + python: Final = tuple(FunctionTraceEvent(name, 0) for name in ("ocr", "map", "validate", "http")) + rust: Final = tuple(FunctionTraceEvent(name, 0) for name in ("ocr", "validate", "map", "http")) + output: Final = format_trace_table(python, rust, colorize=True) + plain: Final = re.sub(r"\033\[[0-9;]*m", "", output) + rows: Final = tuple(line.split("|")[1:-1] for line in plain.splitlines() if line.startswith("|"))[1:] + + assert tuple(row[0].strip() for row in rows if row[0].strip()) == tuple(event.function for event in python) + assert tuple(row[1].strip() for row in rows if row[1].strip()) == tuple(event.function for event in rust) + assert plain.count("reordered") == 2 + assert output.count("\033[31m") == 2 + assert "only" not in output + + +def test_table_colors_match_and_exclusive_rows_without_changing_alignment() -> None: + python: Final = (FunctionTraceEvent("ocr", 0), FunctionTraceEvent("python_helper", 1)) + rust: Final = (FunctionTraceEvent("ocr", 0), FunctionTraceEvent("rust_helper", 1)) + colored: Final = format_trace_table(python, rust, colorize=True) + + assert re.sub(r"\033\[[0-9;]*m", "", colored) == format_trace_table(python, rust, colorize=False) + assert next(line for line in colored.splitlines() if "match" in line).startswith("\033[32m") + assert next(line for line in colored.splitlines() if "python only" in line).startswith("\033[34m") + assert next(line for line in colored.splitlines() if "rust only" in line).startswith("\033[33m") + + +def test_table_handles_empty_traces() -> None: + output: Final = format_trace_table((), (), colorize=False) + + assert "python (0 steps)" in output + assert "rust (0 steps)" in output + assert "match" not in output From 62e318de8eef9ef90af744e7c31e732cea66d758 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 2 Sep 2026 16:26:35 -0700 Subject: [PATCH 60/69] fix(python-bridge): harden sync and async route boundaries (#39332) --- .github/workflows/test-rust.yml | 5 + litellm-rust/Cargo.lock | 22 +- litellm-rust/Cargo.toml | 2 +- litellm-rust/crates/python-bridge/Cargo.toml | 4 +- .../crates/python-bridge/src/execution.rs | 423 ++++++++++++++++++ litellm-rust/crates/python-bridge/src/lib.rs | 116 ++++- .../crates/python-bridge/src/marshal.rs | 7 +- .../src/routes/audio_transcription.rs | 8 +- .../src/routes/chat_completions.rs | 16 +- .../python-bridge/src/routes/messages.rs | 8 +- .../crates/python-bridge/src/routes/mod.rs | 40 -- .../crates/python-bridge/src/routes/ocr.rs | 8 +- litellm-rust/crates/python-interop/src/lib.rs | 2 +- .../crates/python-interop/src/marshal.rs | 72 +++ .../rust_bridge/native_route_wheel_test.py | 349 +++++++++++++++ 15 files changed, 990 insertions(+), 92 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/execution.rs create mode 100644 tests/test_litellm/rust_bridge/native_route_wheel_test.py diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index aada0fcf239..1b71232bc2e 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -9,6 +9,7 @@ on: - "rust-toolchain.toml" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" + - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" - ".github/workflows/test-rust.yml" pull_request: branches: @@ -23,6 +24,7 @@ on: - "rust-toolchain.toml" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" + - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" - ".github/workflows/test-rust.yml" permissions: @@ -121,3 +123,6 @@ jobs: env: RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl + + - name: Test native route wheel + run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 71c9bb95ec0..b3dac5ca935 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1450,6 +1450,7 @@ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ "criterion", + "futures-util", "litellm-ai-gateway", "litellm-core", "litellm-python-interop", @@ -1458,6 +1459,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "tokio-tungstenite", "tracing", "tracing-subscriber", ] @@ -1680,9 +1682,9 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.29.0" +version = "0.29.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" +checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b" dependencies = [ "libc", "once_cell", @@ -1708,18 +1710,18 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.29.0" +version = "0.29.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" +checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9" dependencies = [ "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.29.0" +version = "0.29.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" +checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327" dependencies = [ "libc", "pyo3-build-config", @@ -1727,9 +1729,9 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.29.0" +version = "0.29.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" +checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -1739,9 +1741,9 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.29.0" +version = "0.29.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" +checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952" dependencies = [ "heck", "proc-macro2", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a096643d0d3..a13dd4c04b0 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -20,7 +20,7 @@ litellm-core = { path = "crates/core" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } litellm-python-interop = { path = "crates/python-interop" } axum = "0.7" -pyo3 = "0.29.0" +pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 275e00300d4..637e5580170 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -16,7 +16,7 @@ extension-module = ["pyo3/extension-module"] panic-test = [] [dependencies] -serde.workspace = true +futures-util.workspace = true tracing.workspace = true tracing-subscriber.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } @@ -24,11 +24,13 @@ litellm-ai-gateway = { workspace = true, default-features = false } litellm-python-interop.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true +serde.workspace = true serde_json.workspace = true tokio.workspace = true [dev-dependencies] criterion = "0.8.2" +tokio-tungstenite.workspace = true [[bench]] name = "serialization" diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs new file mode 100644 index 00000000000..f3648158cf6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -0,0 +1,423 @@ +use std::future::Future; +use std::panic::AssertUnwindSafe; +use std::time::Duration; + +use futures_util::FutureExt; +use litellm_core::error::Error; +use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use serde::Serialize; +use tokio::runtime::{Handle, Runtime}; +use tokio::time::{self, MissedTickBehavior}; + +pub(crate) fn run_sync( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + Send + 'static, +{ + run_sync_on( + py, + pyo3_async_runtimes::tokio::get_runtime(), + future, + map_error, + ) +} + +fn run_sync_on( + py: Python<'_>, + runtime: &Runtime, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + Send + 'static, +{ + if Handle::try_current().is_ok() { + return Err(PyRuntimeError::new_err( + "synchronous native routes cannot run from a Tokio context; use the async route", + )); + } + + let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?; + let result = map_core_result(result, map_error)?; + Pythonized(result).into_pyobject(py).map(Bound::unbind) +} + +pub(crate) fn run_async( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + Send + 'static, +{ + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let result = catch_future_panic(future).await?; + let result = map_core_result(result, map_error)?; + Ok(Pythonized(result)) + }) +} + +fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { + match result { + Ok(value) => Ok(value), + Err(error) => Err( + std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error))) + .map_err(panic_to_pyerr)?, + ), + } +} + +async fn catch_future_panic(future: F) -> PyResult> +where + F: Future>, +{ + AssertUnwindSafe(future) + .catch_unwind() + .await + .map_err(panic_to_pyerr) +} + +async fn wait_for_sync_result(future: F) -> PyResult> +where + F: Future>, +{ + let future = catch_future_panic(future); + tokio::pin!(future); + + let signal_interval = Duration::from_millis(50); + let mut signal_checks = + time::interval_at(time::Instant::now() + signal_interval, signal_interval); + signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay); + loop { + tokio::select! { + result = &mut future => return result, + _ = signal_checks.tick() => Python::attach(|py| py.check_signals())?, + } + } +} + +#[cfg(test)] +mod tests { + use std::ffi::CString; + use std::future::poll_fn; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, mpsc}; + use std::task::Poll; + use std::thread; + use std::time::Instant; + + use pyo3::panic::PanicException; + use pyo3::types::{PyDict, PyModule}; + use serde::Serializer; + use tokio::runtime::Builder; + + use super::*; + + fn runtime_error(error: Error) -> PyErr { + PyRuntimeError::new_err(error.to_string()) + } + + fn panicking_error_mapper(_error: Error) -> PyErr { + panic!("error mapper panicked") + } + + struct PanickingOutput; + + static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0); + + impl Serialize for PanickingOutput { + fn serialize(&self, _serializer: S) -> Result + where + S: Serializer, + { + panic!("serializer panicked") + } + } + + #[pyfunction] + fn async_serialization_panic(py: Python<'_>) -> PyResult> { + run_async(py, async { Ok(PanickingOutput) }, runtime_error) + } + + #[pyfunction] + fn async_runtime_probe(py: Python<'_>) -> PyResult> { + run_async( + py, + async { + ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst); + Ok(true) + }, + runtime_error, + ) + } + + #[pyfunction] + fn runtime_worker_count() -> usize { + pyo3_async_runtimes::tokio::get_runtime() + .metrics() + .num_workers() + } + + #[pyfunction] + fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool { + let completion_deadline = Instant::now() + Duration::from_secs(2); + while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions { + if Instant::now() >= completion_deadline { + return false; + } + thread::sleep(Duration::from_millis(1)); + } + + let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1); + pyo3_async_runtimes::tokio::get_runtime().spawn(async move { + let _ = heartbeat_tx.send(()); + }); + heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok() + } + + fn extract_bool(py: Python<'_>, result: PyResult>) -> bool { + result + .expect("route should complete") + .bind(py) + .extract() + .expect("result should convert") + } + + #[test] + fn sync_runner_polls_future_on_the_caller_thread() { + Python::initialize(); + Python::attach(|py| { + let caller_thread = std::thread::current().id(); + let result = run_sync( + py, + async move { Ok(std::thread::current().id() == caller_thread) }, + runtime_error, + ); + + assert!(extract_bool(py, result)); + }); + } + + #[test] + fn sync_runner_releases_gil_while_waiting() { + Python::initialize(); + Python::attach(|py| { + let result = run_sync( + py, + async { + let gil_acquired = tokio::time::timeout( + Duration::from_secs(2), + tokio::task::spawn_blocking(|| Python::attach(|_| true)), + ) + .await; + Ok(matches!(gil_acquired, Ok(Ok(true)))) + }, + runtime_error, + ); + + assert!(extract_bool(py, result)); + }); + } + + #[test] + fn sync_runner_rejects_calls_from_a_tokio_context() { + Python::initialize(); + let runtime = Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime should build"); + + let error = runtime.block_on(async { + Python::attach(|py| { + run_sync::(py, async { Ok(true) }, runtime_error) + .expect_err("sync route should reject a nested Tokio runtime") + }) + }); + + assert_eq!( + error.to_string(), + "RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route" + ); + } + + #[test] + fn sync_runner_can_drive_a_current_thread_runtime() { + Python::initialize(); + let runtime = Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime should build"); + Python::attach(|py| { + let result = run_sync_on( + py, + &runtime, + async { + tokio::task::yield_now().await; + Ok(true) + }, + runtime_error, + ); + assert!(extract_bool(py, result)); + }); + } + + #[test] + fn sync_runner_maps_a_panicked_future() { + Python::initialize(); + Python::attach(|py| { + let error = run_sync::( + py, + poll_fn(|_| -> Poll> { panic!("route future panicked") }), + runtime_error, + ) + .expect_err("panicked route should become a Python exception"); + + assert!(error.is_instance_of::(py)); + assert_eq!(error.to_string(), "PanicException: route future panicked"); + }); + } + + #[test] + fn sync_runner_maps_a_panicked_error_mapper() { + Python::initialize(); + Python::attach(|py| { + let error = run_sync::( + py, + async { Err(Error::InvalidRequest("invalid".to_string())) }, + panicking_error_mapper, + ) + .expect_err("panicked mapper should become a Python exception"); + + assert!(error.is_instance_of::(py)); + assert_eq!(error.to_string(), "PanicException: error mapper panicked"); + }); + } + + #[test] + fn sync_runner_surfaces_serializer_panics() { + Python::initialize(); + Python::attach(|py| { + let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error) + .expect_err("serializer panic should become a Python exception"); + + assert!(error.is_instance_of::(py)); + assert_eq!(error.to_string(), "PanicException: serializer panicked"); + }); + } + + #[test] + fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() { + Python::initialize(); + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let callers: Vec<_> = (0..2) + .map(|_| { + let barrier = Arc::clone(&barrier); + thread::spawn(move || { + Python::attach(|py| { + extract_bool( + py, + run_sync( + py, + async move { + Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait()) + .await + .is_ok()) + }, + runtime_error, + ), + ) + }) + }) + }) + .collect(); + let results: Vec<_> = callers + .into_iter() + .map(|caller| caller.join().expect("caller should not panic")) + .collect(); + + assert_eq!(results, vec![true, true]); + } + + #[test] + fn async_runner_surfaces_serializer_panics() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "runtime").expect("module should be created"); + module + .add_function( + wrap_pyfunction!(async_serialization_panic, &module) + .expect("function should wrap"), + ) + .expect("function should register"); + let locals = PyDict::new(py); + locals + .set_item("runtime", &module) + .expect("module should enter Python locals"); + let code = CString::new( + r#" +import asyncio + +async def exercise(): + try: + await runtime.async_serialization_panic() + except BaseException as error: + assert type(error).__name__ == "PanicException" + assert str(error) == "serializer panicked" + else: + raise AssertionError("serializer panic was not raised") + +asyncio.run(exercise()) +"#, + ) + .expect("Python source should not contain null bytes"); + py.run(&code, Some(&locals), Some(&locals)) + .expect("serializer panic should reach the Python awaiter"); + }); + } + + #[test] + fn async_result_delivery_does_not_stall_tokio_workers() { + Python::initialize(); + ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst); + Python::attach(|py| { + let module = PyModule::new(py, "runtime").expect("module should be created"); + for function in [ + wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"), + wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"), + wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"), + ] { + module + .add_function(function) + .expect("function should register"); + } + let locals = PyDict::new(py); + locals + .set_item("runtime", &module) + .expect("module should enter Python locals"); + let code = CString::new( + r#" +import asyncio + +async def exercise(): + worker_count = runtime.runtime_worker_count() + awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)] + assert runtime.runtime_is_responsive(worker_count) + assert await asyncio.gather(*awaitables) == [True] * worker_count + +asyncio.run(exercise()) +"#, + ) + .expect("Python source should not contain null bytes"); + py.run(&code, Some(&locals), Some(&locals)) + .expect("result delivery should leave Tokio workers responsive"); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 1516fcd7d19..5f36a22370a 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,6 +1,7 @@ mod constants; mod diagnostics; mod errors; +mod execution; pub mod function_trace; mod marshal; mod routes; @@ -8,6 +9,7 @@ mod routes; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use pyo3::prelude::*; use pyo3::types::PyAny; +use serde_json::Value; use crate::errors::core_error_to_pyerr; use crate::marshal::{marshal_headers, optional_timeout}; @@ -25,16 +27,16 @@ impl ResponsesWebSocketConnection { _cls: &Bound<'py, pyo3::types::PyType>, py: Python<'py>, url: String, - headers: Option>, + #[pyo3(from_py_with = litellm_python_interop::from_py)] headers: Option, timeout_seconds: Option, ) -> PyResult> { - let headers = marshal_headers(py, headers)?; + let headers = marshal_headers(headers)?; let timeout = optional_timeout(timeout_seconds); pyo3_async_runtimes::tokio::future_into_py(py, async move { let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) .await .map_err(core_error_to_pyerr)?; - Python::attach(|py| Py::new(py, ResponsesWebSocketConnection { inner })) + Ok(ResponsesWebSocketConnection { inner }) }) } @@ -60,24 +62,36 @@ impl ResponsesWebSocketConnection { } } -#[pymodule] -fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { - errors::register(module)?; - routes::register(module)?; - module.add_class::()?; - diagnostics::register(module) +#[pymodule(gil_used = false)] +mod _native { + use pyo3::prelude::*; + + #[pymodule_init] + fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { + super::errors::register(module)?; + super::routes::register(module)?; + module.add_class::()?; + super::diagnostics::register(module) + } } #[cfg(test)] mod tests { + use std::ffi::CString; + use std::time::Duration; + + use futures_util::{SinkExt, StreamExt}; + use pyo3::types::PyDict; + use tokio::net::TcpListener; + use tokio_tungstenite::{accept_async, tungstenite::Message}; + use super::*; #[test] fn module_registration_preserves_the_public_surface() { Python::initialize(); Python::attach(|py| { - let module = PyModule::new(py, "_native").expect("module should be created"); - _native(&module).expect("module should register"); + let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py); let expected = [ "RustBridgeDeclined", @@ -95,13 +109,79 @@ mod tests { "gil_stats", ]; - for name in expected { - assert!( - module - .hasattr(name) - .expect("attribute lookup should succeed") - ); - } + let public_names: Vec = module + .dict() + .keys() + .extract::>() + .expect("module names should be strings") + .into_iter() + .filter(|name| !name.starts_with("__")) + .collect(); + assert_eq!(public_names, expected); }); } + + #[test] + fn responses_websocket_connection_round_trips_through_python() { + Python::initialize(); + let runtime = pyo3_async_runtimes::tokio::get_runtime(); + let listener = runtime + .block_on(TcpListener::bind("127.0.0.1:0")) + .expect("listener should bind"); + let address = listener + .local_addr() + .expect("listener should have an address"); + let server = runtime.spawn(async move { + let (stream, _) = listener.accept().await.expect("server should accept"); + let mut socket = accept_async(stream) + .await + .expect("handshake should succeed"); + + let message = socket + .next() + .await + .expect("client should send a frame") + .expect("client frame should be valid"); + assert_eq!(message, Message::Text("from-python".into())); + socket + .send(Message::Text("from-server".into())) + .await + .expect("server should reply"); + assert!(matches!(socket.next().await, Some(Ok(Message::Close(_))))); + }); + + Python::attach(|py| { + let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py); + let locals = PyDict::new(py); + locals + .set_item("native", &module) + .expect("module should enter Python locals"); + locals + .set_item("url", format!("ws://{address}")) + .expect("URL should enter Python locals"); + let code = CString::new( + r#" +import asyncio + +async def exercise(): + connection = await native.ResponsesWebSocketConnection.connect(url) + assert type(connection) is native.ResponsesWebSocketConnection + await connection.send_text("from-python") + assert await connection.recv_text() == "from-server" + await connection.close() + assert await connection.recv_text() is None + +asyncio.run(asyncio.wait_for(exercise(), timeout=5)) +"#, + ) + .expect("Python source should not contain null bytes"); + py.run(&code, Some(&locals), Some(&locals)) + .expect("Python WebSocket methods should round trip"); + }); + + runtime + .block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await }) + .expect("server should finish") + .expect("server task should not panic"); + } } diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 8d801109dfa..1c780b2baf9 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -40,12 +40,9 @@ pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option }) } -pub(crate) fn marshal_headers( - py: Python<'_>, - headers: Option>, -) -> PyResult> { +pub(crate) fn marshal_headers(headers: Option) -> PyResult> { let value = match headers { - Some(headers) => from_py(headers.bind(py))?, + Some(headers) => headers, None => Value::Object(Map::new()), }; let Value::Object(headers) = value else { diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index c3e290310fa..1c925ed8592 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -9,10 +9,10 @@ use pyo3::prelude::*; use serde_json::{Map, Value}; use crate::errors::core_error_to_pyerr; +use crate::execution::{run_async, run_sync}; +use crate::function_trace::trace_call; use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; -use super::{block_on, into_py_future}; - struct TranscriptionInputs { model: String, audio: Value, @@ -88,7 +88,7 @@ fn transcription( optional_params, timeout_seconds, )?; - block_on(py, call(inputs), trace, core_error_to_pyerr) + run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr) } #[pyfunction] @@ -117,7 +117,7 @@ fn atranscription( optional_params, timeout_seconds, )?; - into_py_future(py, call(inputs), trace, core_error_to_pyerr) + run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr) } pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 2ef3810015e..4dc446f9985 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -11,10 +11,10 @@ use pyo3::prelude::*; use serde_json::{Map, Value}; use crate::errors::chat_completions_error_to_pyerr; +use crate::execution::{run_async, run_sync}; +use crate::function_trace::trace_call; use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; -use super::{block_on, into_py_future}; - struct ChatCompletionsInputs { model: String, messages: Value, @@ -117,7 +117,11 @@ fn chat_completions( extra_headers, timeout_seconds, )?; - block_on(py, call(inputs), trace, chat_completions_error_to_pyerr) + run_sync( + py, + trace_call(call(inputs), trace), + chat_completions_error_to_pyerr, + ) } #[pyfunction] @@ -146,7 +150,11 @@ fn achat_completions( extra_headers, timeout_seconds, )?; - into_py_future(py, call(inputs), trace, chat_completions_error_to_pyerr) + run_async( + py, + trace_call(call(inputs), trace), + chat_completions_error_to_pyerr, + ) } pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index a29fbdd9fdc..b91f1567e06 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -9,10 +9,10 @@ use pyo3::prelude::*; use serde_json::{Map, Value}; use crate::errors::core_error_to_pyerr; +use crate::execution::{run_async, run_sync}; +use crate::function_trace::trace_call; use crate::marshal::{optional_object, optional_timeout}; -use super::{block_on, into_py_future}; - struct MessagesInputs { model: String, body: Value, @@ -86,7 +86,7 @@ fn messages( extra_headers, timeout_seconds, )?; - block_on(py, call(inputs), trace, core_error_to_pyerr) + run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr) } #[pyfunction] @@ -113,7 +113,7 @@ fn amessages( extra_headers, timeout_seconds, )?; - into_py_future(py, call(inputs), trace, core_error_to_pyerr) + run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr) } pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index ab31702d70f..a2eb8355767 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -1,11 +1,4 @@ -use std::future::Future; - -use litellm_core::error::Error; -use litellm_python_interop::{release_gil, to_py}; use pyo3::prelude::*; -use serde::Serialize; - -use crate::function_trace::trace_call; mod audio_transcription; mod chat_completions; @@ -18,36 +11,3 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { messages::register(module)?; chat_completions::register(module) } - -fn block_on( - py: Python<'_>, - call: impl Future> + Send, - trace: bool, - map_err: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send, -{ - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(trace_call(call, trace)) - }); - match result { - Ok(response) => to_py(py, &response), - Err(err) => Err(map_err(err)), - } -} - -fn into_py_future<'py, T>( - py: Python<'py>, - call: impl Future> + Send + 'static, - trace: bool, - map_err: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send + 'static, -{ - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let response = trace_call(call, trace).await.map_err(map_err)?; - Python::attach(|py| to_py(py, &response)) - }) -} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index ad2a839de07..bed63a871b5 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -7,10 +7,10 @@ use pyo3::prelude::*; use serde_json::{Map, Value}; use crate::errors::core_error_to_pyerr; +use crate::execution::{run_async, run_sync}; +use crate::function_trace::trace_call; use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; -use super::{block_on, into_py_future}; - struct OcrInputs { model: String, document: Value, @@ -90,7 +90,7 @@ fn ocr( optional_params, timeout_seconds, )?; - block_on(py, call(inputs), trace, core_error_to_pyerr) + run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr) } #[pyfunction] @@ -119,7 +119,7 @@ fn aocr( optional_params, timeout_seconds, )?; - into_py_future(py, call(inputs), trace, core_error_to_pyerr) + run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr) } pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/litellm-rust/crates/python-interop/src/lib.rs b/litellm-rust/crates/python-interop/src/lib.rs index df2bd260fdb..2e562bdae70 100644 --- a/litellm-rust/crates/python-interop/src/lib.rs +++ b/litellm-rust/crates/python-interop/src/lib.rs @@ -2,4 +2,4 @@ mod gil; mod marshal; pub use gil::{release_count, release_gil}; -pub use marshal::{from_py, to_py}; +pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py}; diff --git a/litellm-rust/crates/python-interop/src/marshal.rs b/litellm-rust/crates/python-interop/src/marshal.rs index c3d0638427c..a16d1e0ae13 100644 --- a/litellm-rust/crates/python-interop/src/marshal.rs +++ b/litellm-rust/crates/python-interop/src/marshal.rs @@ -1,4 +1,8 @@ +use std::any::Any; +use std::panic::{AssertUnwindSafe, catch_unwind}; + use pyo3::exceptions::PyValueError; +use pyo3::panic::PanicException; use pyo3::prelude::*; use serde::Serialize; use serde::de::DeserializeOwned; @@ -18,3 +22,71 @@ where .map(Bound::unbind) .map_err(|error| PyValueError::new_err(error.to_string())) } + +pub struct Pythonized(pub T); + +impl<'py, T> IntoPyObject<'py> for Pythonized +where + T: Serialize, +{ + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, &self.0))) + .map_err(panic_to_pyerr)? + .map_err(|error| PyValueError::new_err(error.to_string())) + } +} + +pub fn panic_to_pyerr(payload: Box) -> PyErr { + let message = payload + .downcast_ref::() + .map(String::as_str) + .or_else(|| payload.downcast_ref::<&str>().copied()) + .unwrap_or("panic from Rust code"); + PanicException::new_err(message.to_string()) +} + +#[cfg(test)] +mod tests { + use serde::Serializer; + + use super::*; + + struct PanickingSerializer; + + impl Serialize for PanickingSerializer { + fn serialize(&self, _serializer: S) -> Result + where + S: Serializer, + { + panic!("serializer panicked") + } + } + + #[test] + fn pythonized_converts_on_the_attached_thread() { + Python::initialize(); + Python::attach(|py| { + let value: Vec = Pythonized(vec![1, 2, 3]) + .into_pyobject(py) + .and_then(|value| value.extract()) + .expect("value should convert"); + assert_eq!(value, vec![1, 2, 3]); + }); + } + + #[test] + fn pythonized_maps_serializer_panics_to_a_base_exception() { + Python::initialize(); + Python::attach(|py| { + let error = Pythonized(PanickingSerializer) + .into_pyobject(py) + .expect_err("serializer panic should become a Python exception"); + assert!(error.is_instance_of::(py)); + assert_eq!(error.to_string(), "PanicException: serializer panicked"); + }); + } +} diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py new file mode 100644 index 00000000000..0ce49d51aed --- /dev/null +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -0,0 +1,349 @@ +from __future__ import annotations + +import asyncio +import importlib.util +import json +import os +import signal +import subprocess +import sys +import tempfile +import threading +import zipfile +from http.client import HTTPMessage +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from socket import socket as Socket +from typing import Final + +REQUEST_STARTED: Final = threading.Event() +REQUEST_CANCELLED: Final = threading.Event() + +ANTHROPIC_RESPONSE: Final = ( + b'{"id":"msg_native","type":"message","role":"assistant",' + b'"model":"claude-sonnet-4-5","content":[{"type":"text","text":"native-message"}],' + b'"stop_reason":"end_turn","stop_sequence":null,' + b'"usage":{"input_tokens":2,"output_tokens":3}}' +) + + +class NativeRouteHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + content_length: Final = int(self.headers.get("content-length", "0")) + body: Final = json.loads(self.rfile.read(content_length)) + route: Final = self.headers.get("x-test-route") + outcome: Final = self.headers.get("x-test-outcome") + assert_native_request(route, outcome, self.path, self.headers, body) + if outcome == "hang": + REQUEST_STARTED.set() + self.connection.settimeout(5) + if connection_was_cancelled(self.connection): + REQUEST_CANCELLED.set() + return + + status: Final = 429 if outcome == "429" else 200 + response_body: Final = native_response(status, route) + + self.send_response(status) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(response_body))) + self.send_header("connection", "close") + self.end_headers() + self.wfile.write(response_body) + + def log_message(self, _message_format: str, *_args: object) -> None: + pass + + +def connection_was_cancelled(connection: Socket) -> bool: + try: + return connection.recv(1) == b"" + except TimeoutError: + return False + except OSError: + return True + + +def assert_native_request( + route: str | None, + outcome: str | None, + path: str, + headers: HTTPMessage, + body: object, +) -> None: + if route not in {"ocr", "transcription", "messages", "chat_completions"}: + raise AssertionError(f"unexpected route marker: {route!r}") + if outcome not in {"success", "429", "hang"}: + raise AssertionError(f"unexpected outcome marker: {outcome!r}") + if not isinstance(body, dict): + raise TypeError(f"{route} sent {type(body).__name__}, expected a JSON object") + if route == "ocr": + assert path == "/v1/ocr" + assert headers.get("authorization") == "Bearer sk-native" + assert body["model"] == "mistral-ocr-latest" + assert body["document"]["document_url"] == "https://example.com/document.pdf" + assert body["include_image_base64"] is True + return + if route == "transcription": + assert path == "/model/mistral.voxtral-mini-3b-2507/converse" + assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ") + assert headers.get("x-amz-date") + assert body["messages"][0]["content"][0]["audio"]["source"]["bytes"] == "AQI=" + assert "The audio language is en" in body["messages"][0]["content"][1]["text"] + return + assert path == "/v1/messages" + assert headers.get("x-api-key") == "sk-native" + assert body["model"] == "claude-sonnet-4-5" + if route == "messages": + assert body["max_tokens"] == 16 + assert body["messages"][0]["content"] == "hello-from-messages" + return + assert body["max_tokens"] == 17 + assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}] + + +def native_response(status: int, route: str | None) -> bytes: + if status == 429: + return b'{"error":"native-rate-limit"}' + if route == "ocr": + return b'{"pages":[{"index":0,"markdown":"native-ocr"}]}' + if route == "transcription": + return b'{"output":{"message":{"content":[{"text":"native-transcription"}]}}}' + return ANTHROPIC_RESPONSE + + +def load_native(native_path: Path) -> object: + module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path) + if module_spec is None or module_spec.loader is None: + raise RuntimeError("cannot create native extension import specification") + native_module: Final = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(native_module) + return native_module + + +def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: + common: Final = { + "api_base": api_base, + "extra_headers": {"x-test-outcome": outcome, "x-test-route": route}, + "timeout_seconds": 3.0, + } + if route == "ocr": + return common | { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + "api_key": "sk-native", + "custom_llm_provider": "mistral", + "optional_params": {"include_image_base64": True}, + } + if route == "transcription": + return common | { + "model": "mistral.voxtral-mini-3b-2507", + "audio": {"data": "AQI=", "format": "wav", "filename": "audio.wav"}, + "custom_llm_provider": "bedrock", + "optional_params": { + "aws_access_key_id": "native-access-key", + "aws_secret_access_key": "native-secret-key", + "aws_region_name": "us-east-1", + "language": "en", + }, + } + if route == "messages": + return common | { + "model": "claude-sonnet-4-5", + "body": { + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hello-from-messages"}], + }, + "api_key": "sk-native", + "custom_llm_provider": "anthropic", + } + if route == "chat_completions": + return common | { + "model": "anthropic/claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hello-from-chat"}], + "optional_params": {"max_tokens": 17}, + "api_key": "sk-native", + } + raise AssertionError(f"unknown route: {route}") + + +def assert_success(route: str, response: object) -> None: + if not isinstance(response, dict): + raise TypeError(f"{route} returned {type(response).__name__}, expected dict") + actual: Final = success_value(route, response) + expected: Final = ( + "native-ocr" if route == "ocr" else "native-transcription" if route == "transcription" else "native-message" + ) + if actual != expected: + raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}") + + +def assert_traced_success(route: str, response: object) -> None: + if not isinstance(response, dict): + raise TypeError(f"{route} returned {type(response).__name__}, expected a traced dict") + assert_success(route, response["response"]) + expected_function: Final = "audio_transcription" if route == "transcription" else route + assert response["trace"][0] == {"function": expected_function, "depth": 0} + + +def success_value(route: str, response: dict[object, object]) -> object: + if route == "ocr": + return response["pages"][0]["markdown"] + if route == "transcription": + return response["text"] + if route == "messages": + return response["content"][0]["text"] + return response["choices"][0]["message"]["content"] + + +def assert_rate_limit(native: object, route: str, error: BaseException) -> None: + if route == "chat_completions": + upstream_error: Final = native.RustUpstreamError + if not isinstance(error, upstream_error) or error.args[0] != 429: + raise AssertionError(f"{route} returned the wrong 429 error: {error!r}") + return + if not isinstance(error, RuntimeError) or "429" not in str(error): + raise AssertionError(f"{route} returned the wrong 429 error: {error!r}") + + +def exercise_sync(native: object, api_base: str) -> None: + for route in ("ocr", "transcription", "messages", "chat_completions"): + function: Final = getattr(native, route) + assert_success(route, function(**route_kwargs(route, api_base, "success"))) + assert_traced_success(route, function(**route_kwargs(route, api_base, "success"), trace=True)) + try: + function(**route_kwargs(route, api_base, "429")) + except (RuntimeError, native.RustUpstreamError) as error: + assert_rate_limit(native, route, error) + else: + raise AssertionError(f"{route} accepted a 429 response") + + +async def exercise_async(native: object, api_base: str) -> None: + for route in ("ocr", "transcription", "messages", "chat_completions"): + function: Final = getattr(native, f"a{route}") + assert_success(route, await function(**route_kwargs(route, api_base, "success"))) + assert_traced_success(route, await function(**route_kwargs(route, api_base, "success"), trace=True)) + try: + await function(**route_kwargs(route, api_base, "429")) + except (RuntimeError, native.RustUpstreamError) as error: + assert_rate_limit(native, route, error) + else: + raise AssertionError(f"a{route} accepted a 429 response") + + +async def exercise_async_concurrency(native: object, api_base: str) -> None: + responses: Final = await asyncio.wait_for( + asyncio.gather( + *( + native.amessages(**route_kwargs("messages", api_base, "success")) + for _ in range(32) + ) + ), + timeout=15, + ) + for response in responses: + assert_success("messages", response) + + +def exercise_routes(native_path: Path, api_base: str) -> object: + native: Final = load_native(native_path) + exercise_sync(native, api_base) + asyncio.run(exercise_async(native, api_base)) + asyncio.run(exercise_async_concurrency(native, api_base)) + return native + + +def exercise_signal(native: object, api_base: str) -> int: + try: + native.messages( + **route_kwargs("messages", api_base, "hang"), + ) + except KeyboardInterrupt: + sys.stdout.write("KeyboardInterrupt\n") + sys.stdout.flush() + sys.stdin.read(1) + return 0 + raise AssertionError("sync native route ignored SIGINT") + + +def verify_sigint(native_path: Path, api_base: str) -> None: + REQUEST_STARTED.clear() + REQUEST_CANCELLED.clear() + process: Final = subprocess.Popen( + (sys.executable, __file__, "child", str(native_path), api_base), + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + if not REQUEST_STARTED.wait(30): + process.kill() + stdout, stderr = process.communicate(timeout=5) + raise AssertionError( + f"native route matrix did not reach the hanging upstream\nstdout:\n{stdout}\nstderr:\n{stderr}" + ) + os.kill(process.pid, signal.SIGINT) + if not REQUEST_CANCELLED.wait(5): + raise AssertionError("interrupted native route did not cancel its upstream future") + if process.poll() is not None: + raise AssertionError("signal child exited before cancellation was observed") + stdout, stderr = process.communicate(input="\n", timeout=5) + if process.returncode != 0 or stdout != "KeyboardInterrupt\n": + raise AssertionError( + f"signal child exited with status {process.returncode}\nstdout:\n{stdout}\nstderr:\n{stderr}" + ) + finally: + if process.poll() is None: + process.kill() + process.wait(timeout=5) + + +def verify_wheel(wheel: Path) -> int: + with tempfile.TemporaryDirectory() as temporary_directory, zipfile.ZipFile(wheel) as archive: + wheel_root: Final = Path(temporary_directory) + for member in archive.infolist(): + target: Final = wheel_root / member.filename + if member.is_dir(): + target.mkdir(parents=True, exist_ok=True) + else: + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(archive.read(member)) + native_members: Final = tuple( + member + for member in archive.infolist() + if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so") + ) + if len(native_members) != 1: + raise AssertionError(f"expected one native extension, found {len(native_members)}") + native_path: Final = wheel_root / native_members[0].filename + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), NativeRouteHandler) + server_thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + api_base: Final = f"http://127.0.0.1:{server.server_address[1]}" + try: + verify_sigint(native_path, api_base) + finally: + server.shutdown() + server.server_close() + server_thread.join(timeout=5) + return 0 + + +def main() -> int: + if len(sys.argv) == 2: + return verify_wheel(Path(sys.argv[1])) + if len(sys.argv) == 4 and sys.argv[1] == "child": + native: Final = exercise_routes(Path(sys.argv[2]), sys.argv[3]) + return exercise_signal(native, sys.argv[3]) + sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") + return 2 + + +if __name__ == "__main__": + sys.exit(main()) From 34b45f3d7996d654968ee9ba501767c2f8e34f40 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 2 Sep 2026 16:26:36 -0700 Subject: [PATCH 61/69] refactor(python-bridge): declare sync and async routes once (#39333) * refactor(rust): move audio transcription into core * refactor(python-bridge): split non-streaming bridge modules * fix(python-bridge): harden sync and async route boundaries * refactor(python-bridge): declare sync and async routes once --- .../crates/python-bridge/src/marshal.rs | 72 ++- .../src/routes/audio_transcription.rs | 171 +++---- .../src/routes/chat_completions.rs | 197 +++----- .../python-bridge/src/routes/definition.rs | 429 ++++++++++++++++++ .../python-bridge/src/routes/messages.rs | 165 +++---- .../crates/python-bridge/src/routes/mod.rs | 3 + .../crates/python-bridge/src/routes/ocr.rs | 179 +++----- .../python-bridge/src/routes/runtime.rs | 423 +++++++++++++++++ 8 files changed, 1149 insertions(+), 490 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/routes/definition.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/runtime.rs diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 1c780b2baf9..a14e4b55d82 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -1,33 +1,77 @@ use std::collections::HashMap; use std::time::Duration; -use litellm_python_interop::from_py; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use serde_json::{Map, Value}; -pub(crate) fn optional_object_to_map( - py: Python<'_>, +pub(crate) struct RouteOptions { + pub(crate) model: String, + pub(crate) api_key: Option, + pub(crate) api_base: Option, + pub(crate) custom_llm_provider: Option, + pub(crate) extra_headers: Option>, + pub(crate) timeout: Option, +} + +pub(crate) struct RouteOptionsInputs { + pub(crate) model: String, + pub(crate) api_key: Option, + pub(crate) api_base: Option, + pub(crate) custom_llm_provider: Option, + pub(crate) extra_headers: Option, + pub(crate) timeout_seconds: Option, +} + +impl RouteOptions { + pub(crate) fn from_python(inputs: RouteOptionsInputs) -> PyResult { + Ok(Self { + model: inputs.model, + api_key: inputs.api_key, + api_base: inputs.api_base, + custom_llm_provider: inputs.custom_llm_provider, + extra_headers: optional_object("extra_headers", inputs.extra_headers)?, + timeout: optional_timeout(inputs.timeout_seconds), + }) + } +} + +pub(crate) fn required_value( name: &'static str, - value: Option>, + value: Value, + expected: fn(&Value) -> bool, + expected_name: &'static str, +) -> PyResult { + if expected(&value) { + return Ok(value); + } + Err(PyValueError::new_err(format!( + "{name} must be a {expected_name}" + ))) +} + +pub(crate) fn object_or_empty( + name: &'static str, + value: Option, ) -> PyResult> { match value { - Some(value) => match from_py(value.bind(py))? { - Value::Object(map) => Ok(map), - _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), - }, + Some(value) => object(name, value), None => Ok(Map::new()), } } -pub(crate) fn optional_object( - py: Python<'_>, +fn optional_object( name: &'static str, - value: Option>, + value: Option, ) -> PyResult>> { - value - .map(|value| optional_object_to_map(py, name, Some(value))) - .transpose() + value.map(|value| object(name, value)).transpose() +} + +fn object(name: &'static str, value: Value) -> PyResult> { + match value { + Value::Object(map) => Ok(map), + _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), + } } pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option { diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index 1c925ed8592..10b86132be7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -1,126 +1,71 @@ -use std::time::Duration; +use litellm_core::Error; +use std::future::Future; use litellm_core::audio_transcription::{ AudioTranscriptionRequest, audio_transcription as run_audio_transcription, }; -use litellm_core::error::Error; -use litellm_python_interop::from_py; use pyo3::prelude::*; -use serde_json::{Map, Value}; +use serde_json::Value; use crate::errors::core_error_to_pyerr; -use crate::execution::{run_async, run_sync}; -use crate::function_trace::trace_call; -use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; -struct TranscriptionInputs { - model: String, - audio: Value, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Map, - timeout: Option, -} - -#[allow(clippy::too_many_arguments)] -fn marshal_inputs( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult { - Ok(TranscriptionInputs { - model, - audio: from_py(audio.bind(py))?, - api_key, - api_base, - custom_llm_provider, - extra_headers: optional_object(py, "extra_headers", extra_headers)?, - optional_params: optional_object_to_map(py, "optional_params", optional_params)?, - timeout: optional_timeout(timeout_seconds), - }) -} - -async fn call(inputs: TranscriptionInputs) -> Result { - run_audio_transcription(AudioTranscriptionRequest { - model: &inputs.model, - audio: inputs.audio, - api_key: inputs.api_key.as_deref(), - api_base: inputs.api_base.as_deref(), - custom_llm_provider: inputs.custom_llm_provider.as_deref(), +fn prepare_transcription( + inputs: AudioTranscriptionInputs, +) -> PyResult> + Send + 'static> { + let audio = inputs.audio; + let options = RouteOptions::from_python(RouteOptionsInputs { + model: inputs.model, + api_key: inputs.api_key, + api_base: inputs.api_base, + custom_llm_provider: inputs.custom_llm_provider, extra_headers: inputs.extra_headers, - optional_params: inputs.optional_params, - timeout: inputs.timeout, + timeout_seconds: inputs.timeout_seconds, + })?; + let optional_params = object_or_empty("optional_params", inputs.optional_params)?; + + Ok(async move { + let RouteOptions { + model, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout, + } = options; + run_audio_transcription(AudioTranscriptionRequest { + model: &model, + audio, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + }) + .await }) - .await } -#[pyfunction] -#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn transcription( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - audio, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - timeout_seconds, - )?; - run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -#[pyfunction] -#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn atranscription( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - audio, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - timeout_seconds, - )?; - run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(transcription, module)?)?; - module.add_function(wrap_pyfunction!(atranscription, module)?) +bridge_route! { + sync = transcription, + asynchronous = atranscription, + inputs = AudioTranscriptionInputs, + required = { + model: String, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + audio: Value, + }, + optional = { + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + optional_params: Option, + timeout_seconds: Option, + }, + prepare = prepare_transcription, + errors = core_error_to_pyerr, } diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 4dc446f9985..68b7762cb10 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -1,87 +1,62 @@ -use std::time::Duration; +use litellm_core::Error; +use std::future::Future; use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; use litellm_core::chat_completions::{ chat_completions as run_chat_completions, chat_completions_decline_reason, }; -use litellm_core::error::Error; -use litellm_python_interop::from_py; -use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use serde_json::{Map, Value}; +use serde_json::Value; use crate::errors::chat_completions_error_to_pyerr; -use crate::execution::{run_async, run_sync}; -use crate::function_trace::trace_call; -use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value}; -struct ChatCompletionsInputs { - model: String, - messages: Value, - optional_params: Map, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout: Option, -} - -#[allow(clippy::too_many_arguments)] -fn marshal_inputs( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - let messages: Value = from_py(messages.bind(py))?; - if !messages.is_array() { - return Err(PyValueError::new_err("messages must be a list")); - } - Ok(ChatCompletionsInputs { - model, - messages, - optional_params: optional_object_to_map(py, "optional_params", optional_params)?, - api_key, - api_base, - custom_llm_provider, - extra_headers: optional_object(py, "extra_headers", extra_headers)?, - timeout: optional_timeout(timeout_seconds), - }) -} - -async fn call(inputs: ChatCompletionsInputs) -> Result { - run_chat_completions(ChatCompletionsRequest { - model: &inputs.model, - messages: inputs.messages, - optional_params: inputs.optional_params, - api_key: inputs.api_key.as_deref(), - api_base: inputs.api_base.as_deref(), - custom_llm_provider: inputs.custom_llm_provider.as_deref(), +fn prepare_chat_completions( + inputs: ChatCompletionsInputs, +) -> PyResult> + Send + 'static> { + let messages = required_value("messages", inputs.messages, Value::is_array, "list")?; + let optional_params = object_or_empty("optional_params", inputs.optional_params)?; + let options = RouteOptions::from_python(RouteOptionsInputs { + model: inputs.model, + api_key: inputs.api_key, + api_base: inputs.api_base, + custom_llm_provider: inputs.custom_llm_provider, extra_headers: inputs.extra_headers, - timeout: inputs.timeout, + timeout_seconds: inputs.timeout_seconds, + })?; + + Ok(async move { + let RouteOptions { + model, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout, + } = options; + run_chat_completions(ChatCompletionsRequest { + model: &model, + messages, + optional_params, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + timeout, + }) + .await }) - .await } -/// The decline reason for this request, or `None` when the Rust path accepts -/// it. Resolves no credentials and performs no I/O, so a host can ask before -/// committing to either path. #[pyfunction] #[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))] fn chat_completions_decline( - py: Python<'_>, model: String, - messages: Py, - optional_params: Option>, + #[pyo3(from_py_with = litellm_python_interop::from_py)] messages: Value, + #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option, custom_llm_provider: Option, ) -> PyResult> { - let messages = from_py(messages.bind(py))?; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let optional_params = object_or_empty("optional_params", optional_params)?; Ok(chat_completions_decline_reason( &model, custom_llm_provider.as_deref(), @@ -91,74 +66,26 @@ fn chat_completions_decline( .map(str::to_string)) } -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn chat_completions( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - messages, - optional_params, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout_seconds, - )?; - run_sync( - py, - trace_call(call(inputs), trace), - chat_completions_error_to_pyerr, - ) -} - -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn achat_completions( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - messages, - optional_params, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout_seconds, - )?; - run_async( - py, - trace_call(call(inputs), trace), - chat_completions_error_to_pyerr, - ) -} - -pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?; - module.add_function(wrap_pyfunction!(chat_completions, module)?)?; - module.add_function(wrap_pyfunction!(achat_completions, module)?) +bridge_route! { + sync = chat_completions, + asynchronous = achat_completions, + inputs = ChatCompletionsInputs, + required = { + model: String, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + messages: Value, + }, + optional = { + #[pyo3(from_py_with = litellm_python_interop::from_py)] + optional_params: Option, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + timeout_seconds: Option, + }, + prepare = prepare_chat_completions, + errors = chat_completions_error_to_pyerr, + extra = [chat_completions_decline], } diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs new file mode 100644 index 00000000000..21a7fd5a766 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -0,0 +1,429 @@ +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyCFunction; + +macro_rules! bridge_route { + ( + sync = $sync_name:ident, + asynchronous = $async_name:ident, + inputs = $inputs:ident, + required = { $($(#[$required_attr:meta])* $required_name:ident: $required_type:ty),+ $(,)? }, + optional = { $($(#[$optional_attr:meta])* $optional_name:ident: $optional_type:ty),* $(,)? }, + prepare = $prepare:path, + errors = $map_error:path + $(, extra = [$($extra:ident),* $(,)?])? + $(,)? + ) => { + struct $inputs { + $($required_name: $required_type,)* + $($optional_name: $optional_type),* + } + + #[pyfunction] + #[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))] + #[allow(clippy::too_many_arguments)] + fn $sync_name( + py: pyo3::Python<'_>, + $($(#[$required_attr])* $required_name: $required_type,)* + $($(#[$optional_attr])* $optional_name: $optional_type,)* + trace: bool, + ) -> pyo3::PyResult> { + let future = $prepare($inputs { + $($required_name,)* + $($optional_name),* + })?; + $crate::execution::run_sync( + py, + $crate::function_trace::trace_call(future, trace), + $map_error, + ) + } + + #[pyfunction] + #[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))] + #[allow(clippy::too_many_arguments)] + fn $async_name( + py: pyo3::Python<'_>, + $($(#[$required_attr])* $required_name: $required_type,)* + $($(#[$optional_attr])* $optional_name: $optional_type,)* + trace: bool, + ) -> pyo3::PyResult> { + let future = $prepare($inputs { + $($required_name,)* + $($optional_name),* + })?; + $crate::execution::run_async( + py, + $crate::function_trace::trace_call(future, trace), + $map_error, + ) + } + + pub(super) fn register( + module: &pyo3::Bound<'_, pyo3::types::PyModule>, + ) -> pyo3::PyResult<()> { + $($($crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($extra, module)?)?;)*)? + $crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($sync_name, module)?)?; + $crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($async_name, module)?)?; + Ok(()) + } + }; +} + +pub(super) fn add_function( + module: &Bound<'_, PyModule>, + function: Bound<'_, PyCFunction>, +) -> PyResult<()> { + let name: String = function.getattr("__name__")?.extract()?; + if module.hasattr(&name)? { + return Err(PyRuntimeError::new_err(format!( + "duplicate native route: {name}" + ))); + } + module.add_function(function) +} + +#[cfg(test)] +mod tests { + use std::ffi::CString; + use std::sync::atomic::{AtomicBool, Ordering}; + + use litellm_core::error::Error; + use pyo3::exceptions::PyLookupError; + use pyo3::types::{PyDict, PyList}; + + use super::*; + + mod synthetic { + use std::future::{Future, pending}; + + use super::*; + + static FUTURE_DROPPED: AtomicBool = AtomicBool::new(false); + + struct DropGuard; + + impl Drop for DropGuard { + fn drop(&mut self) { + FUTURE_DROPPED.store(true, Ordering::SeqCst); + } + } + + #[pyfunction] + fn future_dropped() -> bool { + FUTURE_DROPPED.load(Ordering::SeqCst) + } + + bridge_route! { + sync = echo, + asynchronous = aecho, + inputs = EchoInputs, + required = { value: String }, + optional = {}, + prepare = prepare_echo, + errors = map_error, + extra = [future_dropped], + } + + fn prepare_echo( + inputs: EchoInputs, + ) -> PyResult> + Send + 'static> { + FUTURE_DROPPED.store(false, Ordering::SeqCst); + let drop_guard = (inputs.value == "pending").then_some(DropGuard); + Ok(async move { + let _drop_guard = drop_guard; + tokio::task::yield_now().await; + match inputs.value.as_str() { + "error" => Err(Error::InvalidRequest("synthetic error".to_string())), + "map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())), + "panic" => panic!("synthetic panic"), + "pending" => { + pending::<()>().await; + unreachable!() + } + _ => Ok(inputs.value), + } + }) + } + + fn map_error(error: Error) -> PyErr { + if matches!(&error, Error::InvalidRequest(message) if message == "panic in mapper") { + panic!("synthetic mapper panic") + } + PyLookupError::new_err(error.to_string()) + } + } + + #[test] + fn sync_and_async_route_signatures_match_the_python_contract() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "routes").expect("module should be created"); + crate::routes::register(&module).expect("routes should register"); + let routes = [ + ( + "ocr", + "aocr", + "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)", + ), + ( + "transcription", + "atranscription", + "(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)", + ), + ( + "messages", + "amessages", + "(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)", + ), + ( + "chat_completions", + "achat_completions", + "(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)", + ), + ]; + + for (sync_name, async_name, expected) in routes { + let sync_signature: String = module + .getattr(sync_name) + .and_then(|function| function.getattr("__text_signature__")) + .and_then(|signature| signature.extract()) + .expect("sync signature should be available"); + let async_signature: String = module + .getattr(async_name) + .and_then(|function| function.getattr("__text_signature__")) + .and_then(|signature| signature.extract()) + .expect("async signature should be available"); + + assert_eq!(sync_signature, expected); + assert_eq!(async_signature, expected); + } + }); + } + + #[test] + fn sync_and_async_routes_apply_the_same_input_validation() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "routes").expect("module should be created"); + crate::routes::register(&module).expect("routes should register"); + + let invalid_messages = PyDict::new(py); + let sync_chat_error = module + .getattr("chat_completions") + .and_then(|function| function.call1(("model", &invalid_messages))) + .expect_err("sync chat should reject a non-list messages value"); + let async_chat_error = module + .getattr("achat_completions") + .and_then(|function| function.call1(("model", &invalid_messages))) + .expect_err("async chat should reject a non-list messages value"); + + assert_eq!( + sync_chat_error.to_string(), + "ValueError: messages must be a list" + ); + assert_eq!(async_chat_error.to_string(), sync_chat_error.to_string()); + + let invalid_body = PyList::empty(py); + let sync_messages_error = module + .getattr("messages") + .and_then(|function| function.call1(("model", &invalid_body))) + .expect_err("sync Messages should reject a non-dict body"); + let async_messages_error = module + .getattr("amessages") + .and_then(|function| function.call1(("model", &invalid_body))) + .expect_err("async Messages should reject a non-dict body"); + + assert_eq!( + sync_messages_error.to_string(), + "ValueError: body must be a dict" + ); + assert_eq!( + async_messages_error.to_string(), + sync_messages_error.to_string() + ); + + let invalid_headers = PyList::empty(py); + let kwargs = PyDict::new(py); + kwargs + .set_item("extra_headers", &invalid_headers) + .expect("kwargs should accept extra_headers"); + let document = PyDict::new(py); + + for (sync_name, async_name) in [("ocr", "aocr"), ("transcription", "atranscription")] { + let sync_error = module + .getattr(sync_name) + .and_then(|function| function.call(("model", &document), Some(&kwargs))) + .expect_err("sync route should reject non-dict extra_headers"); + let async_error = module + .getattr(async_name) + .and_then(|function| function.call(("model", &document), Some(&kwargs))) + .expect_err("async route should reject non-dict extra_headers"); + + assert_eq!( + sync_error.to_string(), + "ValueError: extra_headers must be a dict" + ); + assert_eq!(async_error.to_string(), sync_error.to_string()); + } + }); + } + + #[test] + fn route_input_validation_preserves_left_to_right_order() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "routes").expect("module should be created"); + crate::routes::register(&module).expect("routes should register"); + let invalid = PyList::empty(py); + + let chat_kwargs = PyDict::new(py); + chat_kwargs + .set_item("optional_params", &invalid) + .expect("kwargs should accept optional_params"); + chat_kwargs + .set_item("extra_headers", &invalid) + .expect("kwargs should accept extra_headers"); + let invalid_messages = PyDict::new(py); + let error = module + .getattr("chat_completions") + .and_then(|function| { + function.call(("model", &invalid_messages), Some(&chat_kwargs)) + }) + .expect_err("messages should be validated first"); + assert_eq!(error.to_string(), "ValueError: messages must be a list"); + + let valid_messages = PyList::empty(py); + let error = module + .getattr("chat_completions") + .and_then(|function| function.call(("model", &valid_messages), Some(&chat_kwargs))) + .expect_err("optional_params should be validated before headers"); + assert_eq!( + error.to_string(), + "ValueError: optional_params must be a dict" + ); + + let headers_kwargs = PyDict::new(py); + headers_kwargs + .set_item("extra_headers", &invalid) + .expect("kwargs should accept extra_headers"); + let invalid_body = PyList::empty(py); + let error = module + .getattr("messages") + .and_then(|function| function.call(("model", &invalid_body), Some(&headers_kwargs))) + .expect_err("body should be validated before headers"); + assert_eq!(error.to_string(), "ValueError: body must be a dict"); + + let invalid_payload = + PyModule::new(py, "invalid_payload").expect("invalid payload should be created"); + for name in ["ocr", "transcription"] { + let error = module + .getattr(name) + .and_then(|function| { + function.call(("model", &invalid_payload), Some(&headers_kwargs)) + }) + .expect_err("payload should be validated before headers"); + assert!(!error.to_string().contains("extra_headers")); + } + }); + } + + #[test] + fn generated_routes_execute_sync_and_async_contracts() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "synthetic").expect("module should be created"); + synthetic::register(&module).expect("routes should register"); + + let sync_value: String = module + .getattr("echo") + .and_then(|function| function.call1(("sync",))) + .and_then(|value| value.extract()) + .expect("sync route should return its value"); + assert_eq!(sync_value, "sync"); + + let sync_error = module + .getattr("echo") + .and_then(|function| function.call1(("error",))) + .expect_err("sync route should map its error"); + assert!(sync_error.is_instance_of::(py)); + assert_eq!( + sync_error.to_string(), + "LookupError: invalid request: synthetic error" + ); + + let locals = PyDict::new(py); + locals + .set_item("routes", &module) + .expect("module should enter Python locals"); + let code = CString::new( + r#" +import asyncio + +async def exercise(): + assert await routes.aecho("async") == "async" + + try: + await routes.aecho("error") + except LookupError as error: + assert str(error) == "invalid request: synthetic error" + else: + raise AssertionError("mapped error was not raised") + + try: + await routes.aecho("panic") + except BaseException as error: + assert type(error).__name__ == "PanicException" + assert str(error) == "synthetic panic" + else: + raise AssertionError("panic was not raised") + + try: + await routes.aecho("map_panic") + except BaseException as error: + assert type(error).__name__ == "PanicException" + assert str(error) == "synthetic mapper panic" + else: + raise AssertionError("mapper panic was not raised") + + task = asyncio.ensure_future(routes.aecho("pending")) + await asyncio.sleep(0) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + else: + raise AssertionError("cancelled route completed") + + for _ in range(100): + if routes.future_dropped(): + break + await asyncio.sleep(0.001) + assert routes.future_dropped() + +asyncio.run(exercise()) +"#, + ) + .expect("Python source should not contain null bytes"); + py.run(&code, Some(&locals), Some(&locals)) + .expect("async route contract should hold"); + }); + } + + #[test] + fn route_registration_rejects_duplicate_python_names() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "synthetic").expect("module should be created"); + synthetic::register(&module).expect("first registration should succeed"); + let error = synthetic::register(&module) + .expect_err("duplicate registration should be rejected"); + + assert_eq!( + error.to_string(), + "RuntimeError: duplicate native route: future_dropped" + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index b91f1567e06..2bb64a7a763 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -1,122 +1,65 @@ -use std::time::Duration; - -use litellm_core::error::Error; +use litellm_core::Error; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; -use litellm_python_interop::from_py; -use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use serde_json::{Map, Value}; +use serde_json::Value; +use std::future::Future; use crate::errors::core_error_to_pyerr; -use crate::execution::{run_async, run_sync}; -use crate::function_trace::trace_call; -use crate::marshal::{optional_object, optional_timeout}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value}; -struct MessagesInputs { - model: String, - body: Value, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout: Option, -} - -#[allow(clippy::too_many_arguments)] -fn marshal_inputs( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - let body: Value = from_py(body.bind(py))?; - if !body.is_object() { - return Err(PyValueError::new_err("body must be a dict")); - } - Ok(MessagesInputs { - model, - body, - api_key, - api_base, - custom_llm_provider, - extra_headers: optional_object(py, "extra_headers", extra_headers)?, - timeout: optional_timeout(timeout_seconds), - }) -} - -async fn call(inputs: MessagesInputs) -> Result { - run_messages(MessagesRequest { - model: &inputs.model, - body: inputs.body, - api_key: inputs.api_key.as_deref(), - api_base: inputs.api_base.as_deref(), - custom_llm_provider: inputs.custom_llm_provider.as_deref(), +fn prepare_messages( + inputs: MessagesInputs, +) -> PyResult> + Send + 'static> { + let body = required_value("body", inputs.body, Value::is_object, "dict")?; + let options = RouteOptions::from_python(RouteOptionsInputs { + model: inputs.model, + api_key: inputs.api_key, + api_base: inputs.api_base, + custom_llm_provider: inputs.custom_llm_provider, extra_headers: inputs.extra_headers, - timeout: inputs.timeout, + timeout_seconds: inputs.timeout_seconds, + })?; + + Ok(async move { + let RouteOptions { + model, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout, + } = options; + run_messages(MessagesRequest { + model: &model, + body, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + timeout, + }) + .await }) - .await } -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn messages( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - body, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout_seconds, - )?; - run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn amessages( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - body, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout_seconds, - )?; - run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(messages, module)?)?; - module.add_function(wrap_pyfunction!(amessages, module)?) +bridge_route! { + sync = messages, + asynchronous = amessages, + inputs = MessagesInputs, + required = { + model: String, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + body: Value, + }, + optional = { + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + timeout_seconds: Option, + }, + prepare = prepare_messages, + errors = core_error_to_pyerr, } diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index a2eb8355767..bf611c26d44 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -1,5 +1,8 @@ use pyo3::prelude::*; +#[macro_use] +mod definition; + mod audio_transcription; mod chat_completions; mod messages; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index bed63a871b5..5588c400972 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -1,128 +1,73 @@ -use std::time::Duration; +use litellm_core::Error; +use std::future::Future; use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; -use litellm_core::error::Error; -use litellm_python_interop::from_py; use pyo3::prelude::*; -use serde_json::{Map, Value}; +use serde_json::Value; use crate::errors::core_error_to_pyerr; -use crate::execution::{run_async, run_sync}; -use crate::function_trace::trace_call; -use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; -struct OcrInputs { - model: String, - document: Value, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Map, - timeout: Option, -} - -#[allow(clippy::too_many_arguments)] -fn marshal_inputs( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult { - Ok(OcrInputs { - model, - document: from_py(document.bind(py))?, - api_key, - api_base, - custom_llm_provider, - extra_headers: optional_object(py, "extra_headers", extra_headers)?, - optional_params: optional_object_to_map(py, "optional_params", optional_params)?, - timeout: optional_timeout(timeout_seconds), - }) -} - -async fn call(inputs: OcrInputs) -> Result { - run_ocr(OcrRequest { - model: &inputs.model, - document: inputs.document, - api_key: inputs.api_key.as_deref(), - api_base: inputs.api_base.as_deref(), - custom_llm_provider: inputs.custom_llm_provider.as_deref(), +fn prepare_ocr( + inputs: OcrInputs, +) -> PyResult> + Send + 'static> { + let document = inputs.document; + let options = RouteOptions::from_python(RouteOptionsInputs { + model: inputs.model, + api_key: inputs.api_key, + api_base: inputs.api_base, + custom_llm_provider: inputs.custom_llm_provider, extra_headers: inputs.extra_headers, - optional_params: inputs.optional_params, - timeout: inputs.timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, + timeout_seconds: inputs.timeout_seconds, + })?; + let optional_params = object_or_empty("optional_params", inputs.optional_params)?; + + Ok(async move { + let RouteOptions { + model, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout, + } = options; + run_ocr(OcrRequest { + model: &model, + document, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }) + .await }) - .await } -#[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn ocr( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - document, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - timeout_seconds, - )?; - run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -#[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn aocr( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - document, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - timeout_seconds, - )?; - run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(ocr, module)?)?; - module.add_function(wrap_pyfunction!(aocr, module)?) +bridge_route! { + sync = ocr, + asynchronous = aocr, + inputs = OcrInputs, + required = { + model: String, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + document: Value, + }, + optional = { + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + optional_params: Option, + timeout_seconds: Option, + }, + prepare = prepare_ocr, + errors = core_error_to_pyerr, } diff --git a/litellm-rust/crates/python-bridge/src/routes/runtime.rs b/litellm-rust/crates/python-bridge/src/routes/runtime.rs new file mode 100644 index 00000000000..87a0c3e0104 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/runtime.rs @@ -0,0 +1,423 @@ +use std::future::Future; +use std::panic::AssertUnwindSafe; +use std::time::Duration; + +use futures_util::FutureExt; +use litellm_core::error::Error; +use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use serde::Serialize; +use tokio::runtime::{Handle, Runtime}; +use tokio::time::{self, MissedTickBehavior}; + +pub(super) fn run_sync( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + Send + 'static, +{ + run_sync_on( + py, + pyo3_async_runtimes::tokio::get_runtime(), + future, + map_error, + ) +} + +fn run_sync_on( + py: Python<'_>, + runtime: &Runtime, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + Send + 'static, +{ + if Handle::try_current().is_ok() { + return Err(PyRuntimeError::new_err( + "synchronous native routes cannot run from a Tokio context; use the async route", + )); + } + + let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?; + let result = map_core_result(result, map_error)?; + Pythonized(result).into_pyobject(py).map(Bound::unbind) +} + +pub(super) fn run_async( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + Send + 'static, +{ + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let result = catch_route_panic(future).await?; + let result = map_core_result(result, map_error)?; + Ok(Pythonized(result)) + }) +} + +fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { + match result { + Ok(value) => Ok(value), + Err(error) => Err( + std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error))) + .map_err(panic_to_pyerr)?, + ), + } +} + +async fn catch_route_panic(future: F) -> PyResult> +where + F: Future>, +{ + AssertUnwindSafe(future) + .catch_unwind() + .await + .map_err(panic_to_pyerr) +} + +async fn wait_for_sync_result(future: F) -> PyResult> +where + F: Future>, +{ + let future = catch_route_panic(future); + tokio::pin!(future); + + let signal_interval = Duration::from_millis(50); + let mut signal_checks = + time::interval_at(time::Instant::now() + signal_interval, signal_interval); + signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay); + loop { + tokio::select! { + result = &mut future => return result, + _ = signal_checks.tick() => Python::attach(|py| py.check_signals())?, + } + } +} + +#[cfg(test)] +mod tests { + use std::ffi::CString; + use std::future::poll_fn; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, mpsc}; + use std::task::Poll; + use std::thread; + use std::time::Instant; + + use pyo3::panic::PanicException; + use pyo3::types::{PyDict, PyModule}; + use serde::Serializer; + use tokio::runtime::Builder; + + use super::*; + + fn runtime_error(error: Error) -> PyErr { + PyRuntimeError::new_err(error.to_string()) + } + + fn panicking_error_mapper(_error: Error) -> PyErr { + panic!("error mapper panicked") + } + + struct PanickingOutput; + + static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0); + + impl Serialize for PanickingOutput { + fn serialize(&self, _serializer: S) -> Result + where + S: Serializer, + { + panic!("serializer panicked") + } + } + + #[pyfunction] + fn async_serialization_panic(py: Python<'_>) -> PyResult> { + run_async(py, async { Ok(PanickingOutput) }, runtime_error) + } + + #[pyfunction] + fn async_runtime_probe(py: Python<'_>) -> PyResult> { + run_async( + py, + async { + ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst); + Ok(true) + }, + runtime_error, + ) + } + + #[pyfunction] + fn runtime_worker_count() -> usize { + pyo3_async_runtimes::tokio::get_runtime() + .metrics() + .num_workers() + } + + #[pyfunction] + fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool { + let completion_deadline = Instant::now() + Duration::from_secs(2); + while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions { + if Instant::now() >= completion_deadline { + return false; + } + thread::sleep(Duration::from_millis(1)); + } + + let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1); + pyo3_async_runtimes::tokio::get_runtime().spawn(async move { + let _ = heartbeat_tx.send(()); + }); + heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok() + } + + fn extract_bool(py: Python<'_>, result: PyResult>) -> bool { + result + .expect("route should complete") + .bind(py) + .extract() + .expect("result should convert") + } + + #[test] + fn sync_runner_polls_future_on_the_caller_thread() { + Python::initialize(); + Python::attach(|py| { + let caller_thread = std::thread::current().id(); + let result = run_sync( + py, + async move { Ok(std::thread::current().id() == caller_thread) }, + runtime_error, + ); + + assert!(extract_bool(py, result)); + }); + } + + #[test] + fn sync_runner_releases_gil_while_waiting() { + Python::initialize(); + Python::attach(|py| { + let result = run_sync( + py, + async { + let gil_acquired = tokio::time::timeout( + Duration::from_secs(2), + tokio::task::spawn_blocking(|| Python::attach(|_| true)), + ) + .await; + Ok(matches!(gil_acquired, Ok(Ok(true)))) + }, + runtime_error, + ); + + assert!(extract_bool(py, result)); + }); + } + + #[test] + fn sync_runner_rejects_calls_from_a_tokio_context() { + Python::initialize(); + let runtime = Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime should build"); + + let error = runtime.block_on(async { + Python::attach(|py| { + run_sync::(py, async { Ok(true) }, runtime_error) + .expect_err("sync route should reject a nested Tokio runtime") + }) + }); + + assert_eq!( + error.to_string(), + "RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route" + ); + } + + #[test] + fn sync_runner_can_drive_a_current_thread_runtime() { + Python::initialize(); + let runtime = Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime should build"); + Python::attach(|py| { + let result = run_sync_on( + py, + &runtime, + async { + tokio::task::yield_now().await; + Ok(true) + }, + runtime_error, + ); + assert!(extract_bool(py, result)); + }); + } + + #[test] + fn sync_runner_maps_a_panicked_future() { + Python::initialize(); + Python::attach(|py| { + let error = run_sync::( + py, + poll_fn(|_| -> Poll> { panic!("route future panicked") }), + runtime_error, + ) + .expect_err("panicked route should become a Python exception"); + + assert!(error.is_instance_of::(py)); + assert_eq!(error.to_string(), "PanicException: route future panicked"); + }); + } + + #[test] + fn sync_runner_maps_a_panicked_error_mapper() { + Python::initialize(); + Python::attach(|py| { + let error = run_sync::( + py, + async { Err(Error::InvalidRequest("invalid".to_string())) }, + panicking_error_mapper, + ) + .expect_err("panicked mapper should become a Python exception"); + + assert!(error.is_instance_of::(py)); + assert_eq!(error.to_string(), "PanicException: error mapper panicked"); + }); + } + + #[test] + fn sync_runner_surfaces_serializer_panics() { + Python::initialize(); + Python::attach(|py| { + let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error) + .expect_err("serializer panic should become a Python exception"); + + assert!(error.is_instance_of::(py)); + assert_eq!(error.to_string(), "PanicException: serializer panicked"); + }); + } + + #[test] + fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() { + Python::initialize(); + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let callers: Vec<_> = (0..2) + .map(|_| { + let barrier = Arc::clone(&barrier); + thread::spawn(move || { + Python::attach(|py| { + extract_bool( + py, + run_sync( + py, + async move { + Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait()) + .await + .is_ok()) + }, + runtime_error, + ), + ) + }) + }) + }) + .collect(); + let results: Vec<_> = callers + .into_iter() + .map(|caller| caller.join().expect("caller should not panic")) + .collect(); + + assert_eq!(results, vec![true, true]); + } + + #[test] + fn async_runner_surfaces_serializer_panics() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "runtime").expect("module should be created"); + module + .add_function( + wrap_pyfunction!(async_serialization_panic, &module) + .expect("function should wrap"), + ) + .expect("function should register"); + let locals = PyDict::new(py); + locals + .set_item("runtime", &module) + .expect("module should enter Python locals"); + let code = CString::new( + r#" +import asyncio + +async def exercise(): + try: + await runtime.async_serialization_panic() + except BaseException as error: + assert type(error).__name__ == "PanicException" + assert str(error) == "serializer panicked" + else: + raise AssertionError("serializer panic was not raised") + +asyncio.run(exercise()) +"#, + ) + .expect("Python source should not contain null bytes"); + py.run(&code, Some(&locals), Some(&locals)) + .expect("serializer panic should reach the Python awaiter"); + }); + } + + #[test] + fn async_result_delivery_does_not_stall_tokio_workers() { + Python::initialize(); + ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst); + Python::attach(|py| { + let module = PyModule::new(py, "runtime").expect("module should be created"); + for function in [ + wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"), + wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"), + wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"), + ] { + module + .add_function(function) + .expect("function should register"); + } + let locals = PyDict::new(py); + locals + .set_item("runtime", &module) + .expect("module should enter Python locals"); + let code = CString::new( + r#" +import asyncio + +async def exercise(): + worker_count = runtime.runtime_worker_count() + awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)] + assert runtime.runtime_is_responsive(worker_count) + assert await asyncio.gather(*awaitables) == [True] * worker_count + +asyncio.run(exercise()) +"#, + ) + .expect("Python source should not contain null bytes"); + py.run(&code, Some(&locals), Some(&locals)) + .expect("result delivery should leave Tokio workers responsive"); + }); + } +} From 082bea851e43db2f6b4447d9c08f6a6e39105e34 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 2 Sep 2026 16:26:36 -0700 Subject: [PATCH 62/69] Merge pull request #39334 from BerriAI/litellm_rust_opt_in_configuration feat(python): unify Rust opt-in and bridge policy --- basedpyright-code-budget.json | 2 +- litellm/__init__.py | 2 +- litellm/llms/custom_httpx/llm_http_handler.py | 17 +- litellm/ocr/main.py | 10 +- litellm/rust_bridge/__init__.py | 2 +- litellm/rust_bridge/bindings.py | 49 +++++ litellm/rust_bridge/chat_completions.py | 12 +- litellm/rust_bridge/configuration.py | 148 ++++++++++++++ litellm/rust_bridge/ocr.py | 61 +----- litellm/rust_bridge/runtime.py | 180 ++++++++++++++++++ litellm/types/router.py | 1 + ruff-strict-budget.json | 6 +- .../test_rust_bridge_messages.py | 36 ++-- tests/test_litellm/ocr/test_rust_bridge.py | 35 ++-- .../responses/test_rust_bridge_websocket.py | 23 ++- .../test_litellm/rust_bridge/test_bindings.py | 35 ++++ .../rust_bridge/test_chat_completions.py | 83 ++++---- .../rust_bridge/test_configuration.py | 175 +++++++++++++++++ .../test_litellm/rust_bridge/test_runtime.py | 95 +++++++++ type-discipline-budget.json | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 21 files changed, 813 insertions(+), 165 deletions(-) create mode 100644 litellm/rust_bridge/bindings.py create mode 100644 litellm/rust_bridge/configuration.py create mode 100644 litellm/rust_bridge/runtime.py create mode 100644 tests/test_litellm/rust_bridge/test_bindings.py create mode 100644 tests/test_litellm/rust_bridge/test_configuration.py create mode 100644 tests/test_litellm/rust_bridge/test_runtime.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index d094c98f5ec..60e691e0ef5 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 14074 }, "reportArgumentType": { - "limit": 2216 + "limit": 2215 }, "reportAssignmentType": { "limit": 319 diff --git a/litellm/__init__.py b/litellm/__init__.py index 61794dabddc..44f2e7c1f02 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1417,7 +1417,7 @@ from .skills.main import ( ) from .containers.main import * from .ocr.main import * -from .rust_bridge.ocr import use_litellm_rust +from .rust_bridge import use_litellm_rust from .rag.main import * from .sandbox.main import * from .search.main import * diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 6f42d42de00..3cd6ee54069 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,6 +1,5 @@ import asyncio import json -import os import ssl from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager @@ -159,7 +158,11 @@ def _rust_responses_websocket_enabled( custom_llm_provider: str | None, litellm_params: GenericLiteLLMParams, ) -> bool: - return custom_llm_provider == "openai" and litellm_params.get("rust") is True + from litellm.rust_bridge.configuration import rust_enabled + + raw_request_override: Final = litellm_params.get("rust") + request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None + return custom_llm_provider == "openai" and rust_enabled(request_override=request_override) from .http_handler import get_shared_realtime_ssl_context @@ -2364,10 +2367,6 @@ class BaseLLMHTTPHandler: "anthropic_messages", ) - @staticmethod - def _rust_env_enabled() -> bool: - return os.getenv("LITELLM_RUST", "").strip().lower() in {"1", "true", "yes", "on"} - @staticmethod async def _maybe_rust_anthropic_messages( *, @@ -2383,7 +2382,11 @@ class BaseLLMHTTPHandler: ) -> AnthropicMessagesResponse | None: if custom_llm_provider not in ("azure_ai", "anthropic"): return None - if litellm_params.get("rust") is not True and not BaseLLMHTTPHandler._rust_env_enabled(): + from litellm.rust_bridge.configuration import rust_enabled + + raw_request_override: Final = litellm_params.get("rust") + request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None + if not rust_enabled(request_override=request_override): return None if has_agentic_hook: return None diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index b918f013700..b260ec6e06f 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -194,6 +194,12 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS +def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool: + raw_request_override: Final = prepared_request.litellm_params.get("rust") + request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None + return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override) + + def _rust_bridge_optional_params( prepared_request: _PreparedOCRRequest, resolve_secret: Callable[[str], str | None], @@ -422,7 +428,7 @@ async def aocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): + if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): from litellm.secret_managers.main import get_secret_str rust_response: Final = await _run_rust_aocr( @@ -694,7 +700,7 @@ def ocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): + if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): from litellm.secret_managers.main import get_secret_str rust_response: Final = _run_rust_ocr( diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py index 3da5b98449b..e6d8ffef48c 100644 --- a/litellm/rust_bridge/__init__.py +++ b/litellm/rust_bridge/__init__.py @@ -1,9 +1,9 @@ """LiteLLM Rust bridge package.""" +from litellm.rust_bridge.configuration import use_litellm_rust from litellm.rust_bridge.loader import ( get_native_bridge, native_bridge_available, ) -from litellm.rust_bridge.ocr import use_litellm_rust __all__ = ["get_native_bridge", "native_bridge_available", "use_litellm_rust"] diff --git a/litellm/rust_bridge/bindings.py b/litellm/rust_bridge/bindings.py new file mode 100644 index 00000000000..d16f150a2aa --- /dev/null +++ b/litellm/rust_bridge/bindings.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import Final, Generic, TypeVar + +from litellm.rust_bridge.loader import get_native_bridge + +BindingT = TypeVar("BindingT") + + +class _Unset: + pass + + +_UNSET: Final = _Unset() + + +class NativeBinding(Generic[BindingT]): + """Resolve one native attribute with an explicit, resettable test override.""" + + def __init__(self, attribute: str, *, validate: Callable[[object], BindingT | None]) -> None: + self._attribute: Final = attribute + self._validate: Final = validate + self._override: BindingT | None | _Unset = _UNSET + + def load(self) -> BindingT | None: + if not isinstance(self._override, _Unset): + return self._override + native: Final = get_native_bridge() + if native is None: + return None + return self._validate(getattr(native, self._attribute, None)) + + def override(self, value: BindingT | None) -> None: + self._override = value + + def reset(self) -> None: + self._override = _UNSET + + +def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None: + native: Final = get_native_bridge() + if native is None: + return None + declined: Final = getattr(native, "RustBridgeDeclined", None) + upstream: Final = getattr(native, "RustUpstreamError", None) + if not isinstance(declined, type) or not isinstance(upstream, type): + return None + return declined, upstream diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index acda3086051..c599667ab17 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -13,7 +13,6 @@ retrying it there would bill the customer for the same work twice. from __future__ import annotations import json -import os from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Protocol @@ -27,6 +26,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo convert_to_model_response_object, ) from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned +from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.loader import get_native_bridge from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.types.utils import ModelResponse @@ -44,8 +44,6 @@ _LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) RUST_RESPONSE_HEADER: Final = "x-litellm-rust" -_TRUTHY_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"}) - class RustChatCompletions(Protocol): def __call__( @@ -181,10 +179,6 @@ def load_rust_achat_completions() -> RustAchatCompletions | None: return loaded -def _env_enables_rust() -> bool: - return os.getenv("LITELLM_RUST", "").strip().lower() in _TRUTHY_ENV_VALUES - - def _load_rust_decline() -> RustChatCompletionsDecline | None: if _STATE.decline is not None: return _STATE.decline @@ -253,8 +247,8 @@ def rust_chat_completions_accepts( return False if stream: return False - opted_in: Final = litellm_params is not None and litellm_params.get("rust") is True - if not opted_in and not _env_enables_rust(): + request_override: Final = litellm_params.get("rust") if litellm_params is not None else None + if not rust_enabled(request_override=request_override if isinstance(request_override, bool) else None): return False if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params): verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path") diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py new file mode 100644 index 00000000000..d54b15f060c --- /dev/null +++ b/litellm/rust_bridge/configuration.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import os +import warnings +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + from litellm.rust_bridge.messages import RustAmessages, RustMessages + from litellm.rust_bridge.ocr import RustAocr, RustOcr + from litellm.rust_bridge.responses_websocket import RustResponsesWebSocketConnection + from litellm.rust_bridge.transcription import RustAtranscription, RustTranscription + +DEFAULT_RUST_ENABLED: Final = False +_TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"}) +_GLOBAL_ENV_NAME: Final = "LITELLM_RUST" +_LEGACY_OCR_ENV_NAME: Final = "LITELLM_USE_RUST_OCR" + + +class _Unset: + pass + + +_UNSET: Final = _Unset() + + +class _RustConfiguration: + def __init__(self) -> None: + self.override: bool | None = None + + +_CONFIGURATION: Final = _RustConfiguration() + + +def _parse_env_bool(value: str | None) -> bool | None: + if value is None: + return None + return value.strip().lower() in _TRUE_ENV_VALUES + + +def resolve_rust_enabled( + *, + request_override: bool | None, + process_override: bool | None, + environment_override: bool | None, + legacy_ocr_override: bool | None = None, + release_default: bool = DEFAULT_RUST_ENABLED, +) -> bool: + if request_override is not None: + return request_override + if process_override is not None: + return process_override + if environment_override is not None: + return environment_override + if legacy_ocr_override is not None: + return legacy_ocr_override + return release_default + + +def rust_enabled(*, request_override: bool | None = None) -> bool: + if request_override is not None: + return request_override + process_override: Final = _CONFIGURATION.override + if process_override is not None: + return process_override + return resolve_rust_enabled( + request_override=None, + process_override=None, + environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)), + ) + + +def rust_ocr_enabled(*, request_override: bool | None = None) -> bool: + if request_override is not None: + return request_override + process_override: Final = _CONFIGURATION.override + if process_override is not None: + return process_override + global_override: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)) + legacy_override: Final = None if global_override is not None else _parse_env_bool(os.getenv(_LEGACY_OCR_ENV_NAME)) + if legacy_override is not None: + warnings.warn( + f"{_LEGACY_OCR_ENV_NAME} is deprecated; use {_GLOBAL_ENV_NAME} instead", + DeprecationWarning, + stacklevel=2, + ) + return resolve_rust_enabled( + request_override=None, + process_override=None, + environment_override=global_override, + legacy_ocr_override=legacy_override, + ) + + +def reset_rust_configuration() -> None: + _CONFIGURATION.override = None + + +def use_litellm_rust( + enabled: bool = True, + *, + ocr: RustOcr | None | _Unset = _UNSET, + aocr: RustAocr | None | _Unset = _UNSET, + messages: RustMessages | None | _Unset = _UNSET, + amessages: RustAmessages | None | _Unset = _UNSET, + responses_websocket: type[RustResponsesWebSocketConnection] | None | _Unset = _UNSET, + transcription: RustTranscription | None | _Unset = _UNSET, + atranscription: RustAtranscription | None | _Unset = _UNSET, +) -> None: + """Set the process override for optional Rust paths. + + Rust-only paths, including Bedrock transcription, are not controlled by this switch. + """ + _CONFIGURATION.override = enabled + bindings: Final = (ocr, aocr, messages, amessages, responses_websocket, transcription, atranscription) + if all(isinstance(binding, _Unset) for binding in bindings): + return + warnings.warn( + "Injecting Rust bridge implementations through use_litellm_rust() is deprecated; " + "use the internal bridge setters in tests", + DeprecationWarning, + stacklevel=2, + ) + + if not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset): + from litellm.rust_bridge.ocr import set_rust_ocr + + if not isinstance(ocr, _Unset): + set_rust_ocr(ocr=ocr) + if not isinstance(aocr, _Unset): + set_rust_ocr(aocr=aocr) + if not isinstance(messages, _Unset) or not isinstance(amessages, _Unset): + from litellm.rust_bridge.messages import set_rust_messages + + if not isinstance(messages, _Unset): + set_rust_messages(messages=messages) + if not isinstance(amessages, _Unset): + set_rust_messages(amessages=amessages) + if not isinstance(responses_websocket, _Unset): + from litellm.rust_bridge.responses_websocket import set_rust_responses_websocket + + set_rust_responses_websocket(connection=responses_websocket) + if not isinstance(transcription, _Unset) or not isinstance(atranscription, _Unset): + from litellm.rust_bridge.transcription import configure_rust_transcription + + if not isinstance(transcription, _Unset): + configure_rust_transcription(transcription=transcription) + if not isinstance(atranscription, _Unset): + configure_rust_transcription(atranscription=atranscription) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 82297d35170..b5b0a35a498 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,16 +2,16 @@ from __future__ import annotations -import os from collections.abc import Awaitable -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables import httpx +from litellm.rust_bridge import configuration as _configuration from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds -if TYPE_CHECKING: - from litellm.rust_bridge.messages import RustAmessages, RustMessages +rust_ocr_enabled = _configuration.rust_ocr_enabled +use_litellm_rust = _configuration.use_litellm_rust class RustOcr(Protocol): @@ -51,69 +51,20 @@ class _Unset: _UNSET: Final[_Unset] = _Unset() -def _env_enables_rust_ocr() -> bool: - return os.getenv("LITELLM_USE_RUST_OCR", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - - -_rust_ocr_enabled = _env_enables_rust_ocr() _rust_ocr_impl: RustOcr | None = None _rust_aocr_impl: RustAocr | None = None -def use_litellm_rust( - enabled: bool = True, +def set_rust_ocr( *, ocr: RustOcr | None | _Unset = _UNSET, aocr: RustAocr | None | _Unset = _UNSET, - messages: RustMessages | None | _Unset = _UNSET, - amessages: RustAmessages | None | _Unset = _UNSET, - responses_websocket: Any | None | _Unset = _UNSET, - transcription: Any | None | _Unset = _UNSET, - atranscription: Any | None | _Unset = _UNSET, ) -> None: - global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl - configuring_ocr: Final = not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset) - configuring_messages: Final = not isinstance(messages, _Unset) or not isinstance(amessages, _Unset) - configuring_responses_websocket: Final = not isinstance(responses_websocket, _Unset) - configuring_transcription: Final = not isinstance(transcription, _Unset) or not isinstance(atranscription, _Unset) - if configuring_ocr or (not configuring_messages and not configuring_responses_websocket): - _rust_ocr_enabled = enabled + global _rust_ocr_impl, _rust_aocr_impl if not isinstance(ocr, _Unset): _rust_ocr_impl = ocr if not isinstance(aocr, _Unset): _rust_aocr_impl = aocr - if configuring_transcription: - from litellm.rust_bridge.transcription import configure_rust_transcription - - configure_rust_transcription( - enabled=enabled, - transcription=transcription, - atranscription=atranscription, - ) - if not configuring_messages and not configuring_responses_websocket: - return - if configuring_messages: - from litellm.rust_bridge.messages import set_rust_messages - - if not isinstance(messages, _Unset) and not isinstance(amessages, _Unset): - set_rust_messages(messages=messages, amessages=amessages) - elif not isinstance(messages, _Unset): - set_rust_messages(messages=messages) - else: - set_rust_messages(amessages=amessages) - if configuring_responses_websocket: - from litellm.rust_bridge.responses_websocket import set_rust_responses_websocket - - set_rust_responses_websocket(connection=responses_websocket) - - -def rust_ocr_enabled() -> bool: - return _rust_ocr_enabled def load_rust_ocr() -> RustOcr | None: diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py new file mode 100644 index 00000000000..00f06c046a2 --- /dev/null +++ b/litellm/rust_bridge/runtime.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from enum import Enum +from typing import Final, Generic, NoReturn, TypeAlias, TypeVar + +from litellm.exceptions import APIError +from litellm.rust_bridge.bindings import native_exception_types + +NativeT = TypeVar("NativeT") +ResultT = TypeVar("ResultT") + + +class FallbackMode(Enum): + PYTHON = "python" + RUST_REQUIRED = "rust_required" + + +@dataclass(frozen=True, slots=True) +class RustHandled(Generic[ResultT]): + value: ResultT + + +@dataclass(frozen=True, slots=True) +class RustDeclined: + reason: str + + +@dataclass(frozen=True, slots=True) +class RustUnavailable: + pass + + +RustAttempt: TypeAlias = RustHandled[ResultT] | RustDeclined | RustUnavailable + + +@dataclass(frozen=True, slots=True) +class BridgeErrorContext: + route: str + provider: str + model: str + + +def invoke( + *, + native_call: Callable[[], NativeT] | None, + fallback: Callable[[], ResultT], + adapt: Callable[[NativeT], ResultT], + mode: FallbackMode, + context: BridgeErrorContext, +) -> ResultT: + result: Final = attempt(native_call=native_call, adapt=adapt, context=context) + if isinstance(result, RustHandled): + return result.value + if mode is FallbackMode.PYTHON: + return fallback() + _raise_required(result, context) + + +async def ainvoke( + *, + native_call: Callable[[], Awaitable[NativeT]] | None, + fallback: Callable[[], Awaitable[ResultT]], + adapt: Callable[[NativeT], ResultT], + mode: FallbackMode, + context: BridgeErrorContext, +) -> ResultT: + result: Final = await aattempt(native_call=native_call, adapt=adapt, context=context) + if isinstance(result, RustHandled): + return result.value + if mode is FallbackMode.PYTHON: + return await fallback() + _raise_required(result, context) + + +def attempt( + *, + native_call: Callable[[], NativeT] | None, + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, +) -> RustAttempt[ResultT]: + if native_call is None: + return RustUnavailable() + exceptions: Final = native_exception_types() + if exceptions is None: + return RustHandled(adapt(native_call())) + declined, upstream = exceptions + try: + value: Final = native_call() + except declined as error: + return RustDeclined(reason=_decline_reason(error)) + except upstream as error: + _raise_upstream(error, context) + return RustHandled(adapt(value)) + + +async def aattempt( + *, + native_call: Callable[[], Awaitable[NativeT]] | None, + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, +) -> RustAttempt[ResultT]: + if native_call is None: + return RustUnavailable() + exceptions: Final = native_exception_types() + if exceptions is None: + return RustHandled(adapt(await native_call())) + declined, upstream = exceptions + try: + value: Final = await native_call() + except declined as error: + return RustDeclined(reason=_decline_reason(error)) + except upstream as error: + _raise_upstream(error, context) + return RustHandled(adapt(value)) + + +def call(operation: Callable[[], ResultT], context: BridgeErrorContext) -> ResultT: + exceptions: Final = native_exception_types() + if exceptions is None: + return operation() + upstream: Final = exceptions[1] + try: + return operation() + except upstream as error: + _raise_upstream(error, context) + + +async def acall(operation: Callable[[], Awaitable[ResultT]], context: BridgeErrorContext) -> ResultT: + exceptions: Final = native_exception_types() + if exceptions is None: + return await operation() + upstream: Final = exceptions[1] + try: + return await operation() + except upstream as error: + _raise_upstream(error, context) + + +def _decline_reason(error: BaseException) -> str: + reason: Final[object] = error.args[0] if error.args else str(error) + return reason if isinstance(reason, str) else str(reason) + + +def _raise_required( + result: RustDeclined | RustUnavailable, + context: BridgeErrorContext, +) -> NoReturn: + raise RuntimeError(f"Rust {context.route} bridge {_required_reason(result)}") + + +def _required_reason(result: RustDeclined | RustUnavailable) -> str: + match result: + case RustUnavailable(): + return "is unavailable" + case RustDeclined(reason=reason): + return f"declined the request: {reason}" + + +def _raise_upstream(error: BaseException, context: BridgeErrorContext) -> NoReturn: + args: Final[tuple[object, ...]] = error.args + status_value: Final = args[0] if args else 0 + message_value: Final = args[1] if len(args) > 1 else str(error) + status: Final = status_value if isinstance(status_value, int) else 0 + message: Final = message_value if isinstance(message_value, str) else str(message_value) + raise APIError( + status_code=status or 500, + message=f"litellm rust {context.route}: {message}", + llm_provider=context.provider, + model=context.model, + ) from error + + +def identity(value: ResultT) -> ResultT: + return value + + +async def async_none() -> None: + return None diff --git a/litellm/types/router.py b/litellm/types/router.py index 2a5f264cee3..4f4df1a8d2e 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -305,6 +305,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): """ custom_llm_provider: str | None = None + rust: bool | None = None tpm: int | None = None rpm: int | None = None itpm: int | None = None diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index ae91b711e13..a390c61dcf2 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 307 + "limit": 304 }, "ASYNC230": { "limit": 11 @@ -156,7 +156,7 @@ "limit": 215 }, "PLW0603": { - "limit": 191 + "limit": 190 }, "PLW1508": { "limit": 190 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1073 + "limit": 1071 }, "TRY002": { "limit": 524 diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index fbd7e36e298..293f75b7592 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -8,6 +8,7 @@ import pytest import litellm from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.rust_bridge import configuration from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -109,10 +110,12 @@ class RaisingAsyncMessages: @pytest.fixture(autouse=True) def _reset_rust_flag(): - litellm.use_litellm_rust(False, messages=None, amessages=None) + rust_messages.set_rust_messages(messages=None, amessages=None) + configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - litellm.use_litellm_rust(False, messages=None, amessages=None) + rust_messages.set_rust_messages(messages=None, amessages=None) + configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -122,17 +125,6 @@ def test_load_rust_messages_returns_injected_impl(): assert rust_messages.load_rust_messages() is bridge -def test_configuring_messages_does_not_enable_ocr(): - from litellm.rust_bridge.ocr import rust_ocr_enabled - - litellm.use_litellm_rust(False) - assert rust_ocr_enabled() is False - - litellm.use_litellm_rust(True, messages=RecordingMessages()) - - assert rust_ocr_enabled() is False - - def test_bare_use_litellm_rust_still_toggles_ocr(): from litellm.rust_bridge.ocr import rust_ocr_enabled @@ -264,7 +256,7 @@ async def test_gate_falls_back_to_python_when_bridge_raises(): @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_absent(): bridge = ExplodingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + rust_messages.set_rust_messages(amessages=bridge) response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure")) @@ -272,6 +264,18 @@ async def test_gate_skips_rust_when_flag_absent(): assert bridge.calls == 0 +@pytest.mark.asyncio +async def test_gate_uses_process_enable_without_request_override(): + bridge = RecordingAsyncMessages() + rust_messages.set_rust_messages(amessages=bridge) + litellm.use_litellm_rust(True) + + response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure")) + + assert response is not None + assert bridge.calls[0]["custom_llm_provider"] == "azure_ai" + + @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_false(): bridge = ExplodingAsyncMessages() @@ -305,7 +309,7 @@ async def test_gate_invokes_rust_for_native_anthropic_provider(): @pytest.mark.asyncio async def test_gate_invokes_rust_when_env_var_set(monkeypatch): bridge = RecordingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + rust_messages.set_rust_messages(amessages=bridge) monkeypatch.setenv("LITELLM_RUST", "1") response = await _gate( @@ -320,7 +324,7 @@ async def test_gate_invokes_rust_when_env_var_set(monkeypatch): @pytest.mark.asyncio async def test_gate_env_var_falsey_does_not_enable(monkeypatch): bridge = ExplodingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + rust_messages.set_rust_messages(amessages=bridge) monkeypatch.setenv("LITELLM_RUST", "0") response = await _gate( diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index acad249a2bb..0764aec7185 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,7 +1,7 @@ """Tests for the optional Rust-backed OCR path.""" -import importlib import builtins +import importlib import types from typing import Any @@ -10,6 +10,7 @@ import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge import configuration # `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` # function onto `litellm.ocr` and shadows the submodule, so import the modules @@ -214,10 +215,12 @@ def build_prepared_request( @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) + rust_bridge.set_rust_ocr(ocr=None, aocr=None) + configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) + rust_bridge.set_rust_ocr(ocr=None, aocr=None) + configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -247,7 +250,14 @@ def test_use_litellm_rust_toggles_flag(): def test_env_var_enables_rust_ocr(monkeypatch): monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - assert rust_bridge._env_enables_rust_ocr() is True + with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): + assert rust_bridge.rust_ocr_enabled() is True + + +def test_explicit_false_overrides_process_enable(): + litellm.use_litellm_rust(True) + + assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False def test_load_rust_ocr_returns_injected_impl(): @@ -471,9 +481,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): ocr_main._run_rust_ocr( prepared_request=build_prepared_request(api_key=None, timeout=None), - resolve_api_key=lambda name: ( - "sk-from-vault" if name == "MISTRAL_API_KEY" else None - ), + resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None, ) assert bridge.calls[0]["api_key"] == "sk-from-vault" @@ -580,9 +588,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): api_base=None, timeout=None, ), - resolve_api_key=lambda name: ( - "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None - ), + resolve_api_key=lambda name: "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None, ) assert bridge.calls[0]["api_base"] == "https://azure.example.com" @@ -600,9 +606,7 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): timeout=None, ), resolve_api_key=lambda name: ( - "https://document-intelligence.example.com" - if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" - else None + "https://document-intelligence.example.com" if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" else None ), ) @@ -815,9 +819,6 @@ def test_ocr_provider_configs_expose_api_key_env_vars(): assert BaseOCRConfig().get_api_key_env_var() is None assert MistralOCRConfig().get_api_key_env_var() == "MISTRAL_API_KEY" assert AzureAIOCRConfig().get_api_key_env_var() == "AZURE_AI_API_KEY" - assert ( - AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() - == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" - ) + assert AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index c9a5b988be6..1233ddf1785 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -3,7 +3,7 @@ from __future__ import annotations import pytest from litellm.llms.custom_httpx.llm_http_handler import _rust_responses_websocket_enabled -from litellm.rust_bridge import responses_websocket +from litellm.rust_bridge import configuration, responses_websocket from litellm.types.router import GenericLiteLLMParams @@ -39,12 +39,33 @@ class _FakeNativeBridge: return _FakeNativeConnection() +@pytest.fixture(autouse=True) +def reset_responses_websocket(): + responses_websocket.set_rust_responses_websocket(connection=None) + configuration.reset_rust_configuration() + yield + responses_websocket.set_rust_responses_websocket(connection=None) + configuration.reset_rust_configuration() + + def test_rust_websocket_bridge_is_disabled_without_flag() -> None: assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) assert not _rust_responses_websocket_enabled("anthropic", GenericLiteLLMParams(rust=True)) assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=True)) +def test_explicit_false_overrides_process_enable() -> None: + configuration.use_litellm_rust(True) + + assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=False)) + + +def test_process_enable_applies_without_request_override() -> None: + configuration.use_litellm_rust(True) + + assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) + + @pytest.mark.asyncio async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: adapter = responses_websocket._ConnectionAdapter(_ClosedNativeConnection()) diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/test_litellm/rust_bridge/test_bindings.py new file mode 100644 index 00000000000..88036a5a556 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_bindings.py @@ -0,0 +1,35 @@ +from types import SimpleNamespace +from typing import Final + +import pytest + +from litellm.rust_bridge import bindings + + +def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None: + native = SimpleNamespace(route=lambda: "native") + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + binding: bindings.NativeBinding[object] = bindings.NativeBinding("route", validate=lambda value: value) + + assert binding.load() is native.route + + binding.override(None) + assert binding.load() is None + + replacement = object() + binding.override(replacement) + assert binding.load() is replacement + + binding.reset() + assert binding.load() is native.route + + +@pytest.mark.parametrize(("value", "expected"), ((3, 3), ("invalid", None), (None, None))) +def test_binding_validates_native_attribute( + monkeypatch: pytest.MonkeyPatch, value: object, expected: int | None +) -> None: + native: Final = SimpleNamespace(route=value) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None) + + assert binding.load() == expected diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 47cb66932b7..03921133c77 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -11,6 +11,7 @@ import pytest import litellm from litellm.rust_bridge import chat_completions as bridge +from litellm.rust_bridge import configuration from litellm.types.utils import ModelResponse RUST_RESPONSE = { @@ -68,13 +69,11 @@ def _hide_native_bridge(monkeypatch): @pytest.fixture(autouse=True) def reset_bridge(): """Every test starts with no injected callables, and leaves none behind.""" - bridge.set_rust_chat_completions( - chat_completions=None, achat_completions=None, decline=None - ) + bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) + configuration.reset_rust_configuration() yield - bridge.set_rust_chat_completions( - chat_completions=None, achat_completions=None, decline=None - ) + bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) + configuration.reset_rust_configuration() class _RecordingDecline: @@ -138,6 +137,18 @@ class TestGate: assert gate.calls[0]["model"] == "claude-sonnet-4-5" assert gate.calls[0]["custom_llm_provider"] == "anthropic" + def test_explicit_false_overrides_process_enable(self): + bridge.set_rust_chat_completions(decline=_RecordingDecline()) + configuration.use_litellm_rust(True) + + assert _accepts(litellm_params={"rust": False}) is False + + def test_process_enable_applies_without_request_override(self): + bridge.set_rust_chat_completions(decline=_RecordingDecline()) + configuration.use_litellm_rust(True) + + assert _accepts(litellm_params={}) is True + def test_the_env_var_opts_in_without_a_per_model_flag(self, monkeypatch): monkeypatch.setenv("LITELLM_RUST", "true") bridge.set_rust_chat_completions(decline=_RecordingDecline()) @@ -253,9 +264,7 @@ class TestSyncCall: assert result.usage.completion_tokens == 4 assert result.usage.total_tokens == 15 assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} - assert result.id == original_id, ( - "the rust path must keep the chatcmpl id litellm already minted" - ) + assert result.id == original_id, "the rust path must keep the chatcmpl id litellm already minted" def test_passes_the_timeout_through_as_seconds(self): native = _RecordingCall() @@ -269,9 +278,7 @@ class TestSyncCall: def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch): _fake_native_bridge(monkeypatch) - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=_FakeDeclined("streaming")) - ) + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming"))) assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None @@ -290,13 +297,9 @@ class TestAsyncCall: assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None @pytest.mark.asyncio - async def test_falls_back_when_the_core_declines_before_calling_the_provider( - self, monkeypatch - ): + async def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch): _fake_native_bridge(monkeypatch) - bridge.set_rust_chat_completions( - achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")) - ) + bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming"))) assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None @@ -310,25 +313,19 @@ class TestAsyncFallbackWrapper: ran.append(True) return "python" - result = await bridge.achat_completions_or_fallback( - **_call_kwargs(ModelResponse()), python_fallback=fallback - ) + result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert result.choices[0].message.content == "hello from rust" assert ran == [] @pytest.mark.asyncio async def test_runs_the_fallback_when_the_core_declines(self, monkeypatch): _fake_native_bridge(monkeypatch) - bridge.set_rust_chat_completions( - achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")) - ) + bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming"))) async def fallback(): return "python" - result = await bridge.achat_completions_or_fallback( - **_call_kwargs(ModelResponse()), python_fallback=fallback - ) + result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert result == "python" @pytest.mark.asyncio @@ -338,9 +335,7 @@ class TestAsyncFallbackWrapper: async def fallback(): return "python" - result = await bridge.achat_completions_or_fallback( - **_call_kwargs(ModelResponse()), python_fallback=fallback - ) + result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert result == "python" @@ -353,17 +348,13 @@ class TestFailureClassification: _fake_native_bridge(monkeypatch) def test_a_decline_falls_back_because_nothing_was_sent(self): - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=_FakeDeclined("streaming")) - ) + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming"))) assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None def test_an_upstream_failure_is_surfaced_with_its_status(self): from litellm.exceptions import APIError - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited")) - ) + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited"))) with pytest.raises(APIError) as raised: bridge.chat_completions(**_call_kwargs(ModelResponse())) assert raised.value.status_code == 429 @@ -372,17 +363,13 @@ class TestFailureClassification: def test_a_transport_failure_with_no_response_surfaces_as_a_500(self): from litellm.exceptions import APIError - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=_FakeUpstream(0, "connection reset")) - ) + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(0, "connection reset"))) with pytest.raises(APIError) as raised: bridge.chat_completions(**_call_kwargs(ModelResponse())) assert raised.value.status_code == 500 def test_an_unrecognized_error_is_not_swallowed(self): - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=RuntimeError("something else")) - ) + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=RuntimeError("something else"))) with pytest.raises(RuntimeError): bridge.chat_completions(**_call_kwargs(ModelResponse())) @@ -390,9 +377,7 @@ class TestFailureClassification: async def test_the_async_wrapper_does_not_fall_back_on_an_upstream_failure(self): from litellm.exceptions import APIError - bridge.set_rust_chat_completions( - achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom")) - ) + bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom"))) ran = [] async def fallback(): @@ -400,9 +385,7 @@ class TestFailureClassification: return "python" with pytest.raises(APIError): - await bridge.achat_completions_or_fallback( - **_call_kwargs(ModelResponse()), python_fallback=fallback - ) + await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert ran == [], "a request the provider already served must not be re-issued" @pytest.mark.asyncio @@ -414,7 +397,5 @@ class TestFailureClassification: async def fallback(): return "python" - result = await bridge.achat_completions_or_fallback( - **_call_kwargs(ModelResponse()), python_fallback=fallback - ) + result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert result == "python" diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/test_litellm/rust_bridge/test_configuration.py new file mode 100644 index 00000000000..1c81c1fb624 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_configuration.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +import os +import subprocess +import sys +from collections.abc import Generator +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +import pytest + +from litellm.rust_bridge import configuration +from litellm.rust_bridge import ocr as rust_ocr + + +class _OcrBridge: + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + return {} + + +@pytest.fixture(autouse=True) +def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest discovers fixtures dynamically + monkeypatch: pytest.MonkeyPatch, +) -> Generator[None]: + configuration.reset_rust_configuration() + monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.delenv("LITELLM_USE_RUST_OCR", raising=False) + rust_ocr.set_rust_ocr(ocr=None, aocr=None) + yield + configuration.reset_rust_configuration() + rust_ocr.set_rust_ocr(ocr=None, aocr=None) + + +@pytest.mark.parametrize( + ("request_override", "process", "environment", "legacy_ocr", "release_default", "expected"), + ( + (False, True, True, True, True, False), + (True, False, False, False, False, True), + (None, False, True, True, True, False), + (None, True, False, False, False, True), + (None, None, False, True, True, False), + (None, None, True, False, False, True), + (None, None, None, False, True, False), + (None, None, None, True, False, True), + (None, None, None, None, False, False), + (None, None, None, None, True, True), + ), +) +def test_resolution_precedence( + request_override: bool | None, + process: bool | None, + environment: bool | None, + legacy_ocr: bool | None, + release_default: bool, + expected: bool, +) -> None: + assert ( + configuration.resolve_rust_enabled( + request_override=request_override, + process_override=process, + environment_override=environment, + legacy_ocr_override=legacy_ocr, + release_default=release_default, + ) + is expected + ) + + +def test_release_default_remains_disabled() -> None: + assert configuration.DEFAULT_RUST_ENABLED is False + assert configuration.rust_enabled() is False + + +def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "0") + configuration.use_litellm_rust(True) + + assert configuration.rust_enabled() is True + assert configuration.rust_enabled(request_override=False) is False + + +def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "off") + + assert configuration.rust_enabled() is False + + +@pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) +def test_invalid_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv("LITELLM_RUST", value) + monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") + + assert configuration.rust_enabled() is False + assert configuration.rust_ocr_enabled() is False + + +@pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) +def test_invalid_legacy_environment_value_disables_ocr(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv("LITELLM_USE_RUST_OCR", value) + + with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): + assert configuration.rust_ocr_enabled() is False + + +def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "1") + + with ThreadPoolExecutor(max_workers=1) as executor: + assert executor.submit(configuration.rust_enabled).result() is True + configuration.use_litellm_rust(False) + assert executor.submit(configuration.rust_enabled).result() is False + assert executor.submit(configuration.rust_ocr_enabled).result() is False + configuration.reset_rust_configuration() + assert executor.submit(configuration.rust_enabled).result() is True + assert executor.submit(configuration.rust_ocr_enabled).result() is True + + +def test_explicit_override_precedes_invalid_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "sometimes") + + assert configuration.rust_enabled(request_override=False) is False + configuration.use_litellm_rust(True) + assert configuration.rust_enabled() is True + + +def test_legacy_ocr_environment_is_deprecated_and_ocr_only(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") + + with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): + assert configuration.rust_ocr_enabled() is True + assert configuration.rust_enabled() is False + + +def test_global_environment_precedes_legacy_ocr_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "0") + monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") + + assert configuration.rust_ocr_enabled() is False + + +def test_deprecated_public_injection_delegates_to_internal_binding() -> None: + bridge: Final = _OcrBridge() + + with pytest.warns(DeprecationWarning, match="Injecting Rust bridge implementations"): + configuration.use_litellm_rust(True, ocr=bridge) + + assert rust_ocr.load_rust_ocr() is bridge + + +@pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False"))) +def test_environment_controls_startup(value: str, expected: str) -> None: + environment: Final = {**os.environ, "LITELLM_RUST": value} + result: Final = subprocess.run( + ( + sys.executable, + "-c", + "from litellm.rust_bridge.configuration import rust_enabled; print(rust_enabled())", + ), + check=True, + capture_output=True, + text=True, + env=environment, + ) + + assert result.stdout.strip() == expected diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py new file mode 100644 index 00000000000..b0fa510069b --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from litellm.exceptions import APIError +from litellm.rust_bridge import bindings, runtime + + +class RustBridgeDeclined(Exception): + pass + + +class RustUpstreamError(Exception): + pass + + +@pytest.fixture(autouse=True) +def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None: + native = SimpleNamespace( + RustBridgeDeclined=RustBridgeDeclined, + RustUpstreamError=RustUpstreamError, + ) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + + +def context() -> runtime.BridgeErrorContext: + return runtime.BridgeErrorContext(route="messages", provider="anthropic", model="model") + + +def test_invoke_tags_native_decline_before_running_fallback() -> None: + calls: list[str] = [] + + def decline() -> object: + calls.append("rust") + raise RustBridgeDeclined("unsupported") + + value = runtime.invoke( + native_call=decline, + fallback=lambda: calls.append("python") or "fallback", + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), + ) + + assert value == "fallback" + assert calls == ["rust", "python"] + + +def test_invoke_translates_upstream_without_fallback() -> None: + def fail() -> object: + raise RustUpstreamError(429, "rate limited") + + with pytest.raises(APIError, match="rate limited") as caught: + runtime.invoke( + native_call=fail, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), + ) + + assert caught.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_ainvoke_handles_native_success() -> None: + async def native() -> int: + return 3 + + async def fallback() -> str: + pytest.fail("fallback must not run") + + assert ( + await runtime.ainvoke( + native_call=native, + fallback=fallback, + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), + ) + == "3" + ) + + +def test_required_mode_rejects_unavailable_bridge() -> None: + with pytest.raises(RuntimeError, match="is unavailable"): + runtime.invoke( + native_call=None, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + mode=runtime.FallbackMode.RUST_REQUIRED, + context=context(), + ) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 6273fbce595..d2f70a5b730 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22358 }, "LIT002": { - "limit": 26774 + "limit": 26772 }, "LIT003": { "limit": 269 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 522a2de574b..746a933cdfd 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29322,6 +29322,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; + /** Rust */ + rust?: boolean | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Encryption Key Id */ @@ -39196,6 +39198,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; + /** Rust */ + rust?: boolean | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Encryption Key Id */ From 9aeeca4ce3b8c9d6a9545ac65523c44bcd1f41db Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 2 Sep 2026 16:33:04 -0700 Subject: [PATCH 63/69] feat(router): add heuristic v2 complexity routing (#39276) * feat(router): add trained heuristic complexity routing * feat(router): expose heuristic v2 classifier * style(router): format heuristic v2 predictor --- .../complexity_router/README.md | 30 + .../artifacts/ultrafeedback_tiers.json | 4069 +++++++++++++++++ .../complexity_router/complexity_router.py | 31 + .../complexity_router/config.py | 21 +- .../complexity_router/tier_predictor.py | 156 + litellm/types/utils.py | 1 + pyproject.toml | 5 +- .../router_strategy/test_complexity_router.py | 70 +- .../test_complexity_tier_predictor.py | 91 + .../add_model/ClassificationMethodConfig.tsx | 16 + .../add_model/ComplexityRouterConfig.test.tsx | 25 + .../add_model/ComplexityRouterConfig.tsx | 19 +- .../add_model/HeuristicScoringConfig.test.tsx | 1 + .../build_complexity_router_config.test.ts | 21 +- .../LogDetailsDrawer/RoutingDecisionCard.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 115 +- 16 files changed, 4653 insertions(+), 19 deletions(-) create mode 100644 litellm/router_strategy/complexity_router/artifacts/ultrafeedback_tiers.json create mode 100644 litellm/router_strategy/complexity_router/tier_predictor.py create mode 100644 tests/test_litellm/router_strategy/test_complexity_tier_predictor.py diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index bc8df67cc28..8e0cad39561 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -68,6 +68,36 @@ still resolve to a deployment in `model_list`; this configuration does not creat - abc ``` +### Heuristic v2 + +Set `classifier_type: heuristic_v2` to classify with the bundled calibrated +success-probability model instead of the hand-written weighted scorer + +```yaml +model_list: + - model_name: smart-router + litellm_params: + model: auto_router/complexity_router + complexity_router_config: + classifier_type: heuristic_v2 + tiers: + SIMPLE: luna + MEDIUM: terra + COMPLEX: sol + REASONING: sol-ultra +``` + +No classifier model call or per-model training data is required. The classifier +uses global tier quality, request-type quality, and similar-request cohorts from +the bundled UltraFeedback artifact. It estimates success at every tier, enforces +monotonic probabilities, and returns the first tier meeting the trained 0.75 +threshold. The existing complexity-router tier pool then selects and dispatches +a model from that tier + +Spend logs record `routing_decision.cause: heuristic_v2`, the detected request +type, and all four predicted probabilities. Existing `classifier_type: heuristic` +configurations keep the original weighted scorer unchanged + ### Renaming the tiers `tier_labels` puts your own vocabulary on the four tiers: diff --git a/litellm/router_strategy/complexity_router/artifacts/ultrafeedback_tiers.json b/litellm/router_strategy/complexity_router/artifacts/ultrafeedback_tiers.json new file mode 100644 index 00000000000..4fcb599907c --- /dev/null +++ b/litellm/router_strategy/complexity_router/artifacts/ultrafeedback_tiers.json @@ -0,0 +1,4069 @@ +{ + "schema_version": 1, + "global_statistics": [ + { + "tier": 1, + "successes": 36619.0, + "observations": 45504.0 + }, + { + "tier": 2, + "successes": 59797.0, + "observations": 70062.0 + }, + { + "tier": 3, + "successes": 48604.0, + "observations": 52245.0 + }, + { + "tier": 4, + "successes": 11393.0, + "observations": 11561.0 + } + ], + "domain_statistics": [ + { + "tier": 1, + "successes": 1592.0, + "observations": 2211.0, + "request_type": "analytical_reasoning" + }, + { + "tier": 2, + "successes": 2654.0, + "observations": 3374.0, + "request_type": "analytical_reasoning" + }, + { + "tier": 3, + "successes": 2243.0, + "observations": 2481.0, + "request_type": "analytical_reasoning" + }, + { + "tier": 4, + "successes": 538.0, + "observations": 546.0, + "request_type": "analytical_reasoning" + }, + { + "tier": 1, + "successes": 750.0, + "observations": 1015.0, + "request_type": "code_generation" + }, + { + "tier": 2, + "successes": 1271.0, + "observations": 1511.0, + "request_type": "code_generation" + }, + { + "tier": 3, + "successes": 1030.0, + "observations": 1111.0, + "request_type": "code_generation" + }, + { + "tier": 4, + "successes": 233.0, + "observations": 235.0, + "request_type": "code_generation" + }, + { + "tier": 1, + "successes": 243.0, + "observations": 277.0, + "request_type": "code_understanding" + }, + { + "tier": 2, + "successes": 385.0, + "observations": 425.0, + "request_type": "code_understanding" + }, + { + "tier": 3, + "successes": 322.0, + "observations": 334.0, + "request_type": "code_understanding" + }, + { + "tier": 4, + "successes": 74.0, + "observations": 76.0, + "request_type": "code_understanding" + }, + { + "tier": 1, + "successes": 2014.0, + "observations": 2170.0, + "request_type": "factual_lookup" + }, + { + "tier": 2, + "successes": 3120.0, + "observations": 3266.0, + "request_type": "factual_lookup" + }, + { + "tier": 3, + "successes": 2612.0, + "observations": 2670.0, + "request_type": "factual_lookup" + }, + { + "tier": 4, + "successes": 540.0, + "observations": 542.0, + "request_type": "factual_lookup" + }, + { + "tier": 1, + "successes": 30571.0, + "observations": 38161.0, + "request_type": "general" + }, + { + "tier": 2, + "successes": 50037.0, + "observations": 58821.0, + "request_type": "general" + }, + { + "tier": 3, + "successes": 40460.0, + "observations": 43618.0, + "request_type": "general" + }, + { + "tier": 4, + "successes": 9565.0, + "observations": 9716.0, + "request_type": "general" + }, + { + "tier": 1, + "successes": 159.0, + "observations": 170.0, + "request_type": "technical_design" + }, + { + "tier": 2, + "successes": 282.0, + "observations": 303.0, + "request_type": "technical_design" + }, + { + "tier": 3, + "successes": 231.0, + "observations": 236.0, + "request_type": "technical_design" + }, + { + "tier": 4, + "successes": 47.0, + "observations": 47.0, + "request_type": "technical_design" + }, + { + "tier": 1, + "successes": 1290.0, + "observations": 1500.0, + "request_type": "writing" + }, + { + "tier": 2, + "successes": 2048.0, + "observations": 2362.0, + "request_type": "writing" + }, + { + "tier": 3, + "successes": 1706.0, + "observations": 1795.0, + "request_type": "writing" + }, + { + "tier": 4, + "successes": 396.0, + "observations": 399.0, + "request_type": "writing" + } + ], + "cohort_statistics": [ + { + "tier": 1, + "successes": 272.0, + "observations": 372.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 420.0, + "observations": 519.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 341.0, + "observations": 381.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 82.0, + "observations": 84.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 1.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 2.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 11.0, + "observations": 18.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 33.0, + "observations": 45.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 28.0, + "observations": 34.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 131.0, + "observations": 176.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 210.0, + "observations": 274.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 187.0, + "observations": 209.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 39.0, + "observations": 41.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 9.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 11.0, + "observations": 13.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 6.0, + "observations": 7.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 10.0, + "observations": 15.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 14.0, + "observations": 16.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 15.0, + "observations": 16.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 15.0, + "observations": 20.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 33.0, + "observations": 38.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 28.0, + "observations": 29.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 2.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 26.0, + "observations": 34.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 39.0, + "observations": 50.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 42.0, + "observations": 45.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 11.0, + "observations": 11.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 3.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 4.0, + "observations": 4.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 452.0, + "observations": 634.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 745.0, + "observations": 962.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 634.0, + "observations": 691.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 151.0, + "observations": 153.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 3.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 4.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 10.0, + "observations": 20.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 30.0, + "observations": 44.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 12.0, + "observations": 15.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 178.0, + "observations": 270.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 304.0, + "observations": 402.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 266.0, + "observations": 295.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 73.0, + "observations": 73.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 7.0, + "observations": 18.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 17.0, + "observations": 32.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 17.0, + "observations": 27.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 7.0, + "observations": 7.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 11.0, + "observations": 17.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 17.0, + "observations": 25.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 10.0, + "observations": 16.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 6.0, + "observations": 6.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 58.0, + "observations": 72.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 90.0, + "observations": 103.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 72.0, + "observations": 78.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 19.0, + "observations": 19.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 37.0, + "observations": 56.0, + "cohort": "analytical_reasoning|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 95.0, + "observations": 110.0, + "cohort": "analytical_reasoning|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 73.0, + "observations": 82.0, + "cohort": "analytical_reasoning|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 20.0, + "observations": 20.0, + "cohort": "analytical_reasoning|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 69.0, + "observations": 83.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 102.0, + "observations": 110.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 93.0, + "observations": 98.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 17.0, + "observations": 17.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 6.0, + "observations": 6.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 21.0, + "observations": 32.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 45.0, + "observations": 59.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 45.0, + "observations": 48.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 9.0, + "observations": 9.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 7.0, + "observations": 9.0, + "cohort": "analytical_reasoning|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 8.0, + "observations": 12.0, + "cohort": "analytical_reasoning|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 11.0, + "observations": 11.0, + "cohort": "analytical_reasoning|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 5.0, + "observations": 9.0, + "cohort": "analytical_reasoning|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 9.0, + "observations": 11.0, + "cohort": "analytical_reasoning|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 8.0, + "observations": 8.0, + "cohort": "analytical_reasoning|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 136.0, + "observations": 176.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 221.0, + "observations": 267.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 172.0, + "observations": 194.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 37.0, + "observations": 39.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 6.0, + "observations": 11.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 10.0, + "observations": 15.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 8.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 68.0, + "observations": 84.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 103.0, + "observations": 136.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 113.0, + "observations": 119.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 21.0, + "observations": 21.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 3.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 4.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 13.0, + "observations": 16.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 27.0, + "observations": 34.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 15.0, + "observations": 16.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 6.0, + "observations": 6.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 5.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 4.0, + "observations": 7.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 32.0, + "observations": 39.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 48.0, + "observations": 63.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 35.0, + "observations": 40.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 14.0, + "observations": 14.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 3.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 2.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 2.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 28.0, + "observations": 31.0, + "cohort": "code_generation|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 41.0, + "observations": 46.0, + "cohort": "code_generation|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 25.0, + "observations": 26.0, + "cohort": "code_generation|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "code_generation|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 6.0, + "observations": 9.0, + "cohort": "code_generation|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 20.0, + "observations": 23.0, + "cohort": "code_generation|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 16.0, + "observations": 16.0, + "cohort": "code_generation|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "code_generation|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 46.0, + "observations": 60.0, + "cohort": "code_generation|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 72.0, + "observations": 91.0, + "cohort": "code_generation|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 59.0, + "observations": 64.0, + "cohort": "code_generation|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 17.0, + "observations": 17.0, + "cohort": "code_generation|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 32.0, + "observations": 49.0, + "cohort": "code_generation|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 60.0, + "observations": 76.0, + "cohort": "code_generation|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 48.0, + "observations": 52.0, + "cohort": "code_generation|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 15.0, + "observations": 15.0, + "cohort": "code_generation|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 93.0, + "observations": 121.0, + "cohort": "code_generation|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 149.0, + "observations": 170.0, + "cohort": "code_generation|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 146.0, + "observations": 151.0, + "cohort": "code_generation|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 26.0, + "observations": 26.0, + "cohort": "code_generation|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 12.0, + "observations": 16.0, + "cohort": "code_generation|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 22.0, + "observations": 25.0, + "cohort": "code_generation|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 13.0, + "observations": 15.0, + "cohort": "code_generation|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "code_generation|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 196.0, + "observations": 268.0, + "cohort": "code_generation|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 312.0, + "observations": 370.0, + "cohort": "code_generation|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 234.0, + "observations": 253.0, + "cohort": "code_generation|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 56.0, + "observations": 57.0, + "cohort": "code_generation|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 1.0, + "cohort": "code_generation|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_generation|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_generation|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 40.0, + "observations": 59.0, + "cohort": "code_generation|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 73.0, + "observations": 92.0, + "cohort": "code_generation|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 78.0, + "observations": 83.0, + "cohort": "code_generation|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 14.0, + "observations": 14.0, + "cohort": "code_generation|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 3.0, + "cohort": "code_generation|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_generation|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 106.0, + "observations": 140.0, + "cohort": "code_generation|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 220.0, + "observations": 247.0, + "cohort": "code_generation|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 167.0, + "observations": 178.0, + "cohort": "code_generation|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 43.0, + "observations": 43.0, + "cohort": "code_generation|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 4.0, + "cohort": "code_generation|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 9.0, + "cohort": "code_generation|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 7.0, + "cohort": "code_generation|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 147.0, + "observations": 199.0, + "cohort": "code_generation|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 233.0, + "observations": 269.0, + "cohort": "code_generation|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 189.0, + "observations": 209.0, + "cohort": "code_generation|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 34.0, + "observations": 35.0, + "cohort": "code_generation|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 7.0, + "observations": 10.0, + "cohort": "code_generation|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 18.0, + "observations": 22.0, + "cohort": "code_generation|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 10.0, + "observations": 10.0, + "cohort": "code_generation|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_generation|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 4.0, + "cohort": "code_generation|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 5.0, + "observations": 8.0, + "cohort": "code_generation|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 8.0, + "cohort": "code_generation|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_generation|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 5.0, + "cohort": "code_generation|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "code_generation|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_generation|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 7.0, + "observations": 9.0, + "cohort": "code_generation|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 8.0, + "cohort": "code_generation|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 9.0, + "observations": 9.0, + "cohort": "code_generation|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_generation|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 24.0, + "observations": 33.0, + "cohort": "code_generation|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 25.0, + "observations": 45.0, + "cohort": "code_generation|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 19.0, + "observations": 27.0, + "cohort": "code_generation|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 7.0, + "observations": 7.0, + "cohort": "code_generation|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_understanding|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_understanding|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 6.0, + "observations": 6.0, + "cohort": "code_understanding|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_understanding|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 5.0, + "cohort": "code_understanding|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 3.0, + "cohort": "code_understanding|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 3.0, + "cohort": "code_understanding|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 17.0, + "observations": 23.0, + "cohort": "code_understanding|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 27.0, + "observations": 37.0, + "cohort": "code_understanding|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 32.0, + "observations": 33.0, + "cohort": "code_understanding|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "code_understanding|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 24.0, + "observations": 28.0, + "cohort": "code_understanding|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 37.0, + "observations": 44.0, + "cohort": "code_understanding|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 33.0, + "observations": 33.0, + "cohort": "code_understanding|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 7.0, + "observations": 7.0, + "cohort": "code_understanding|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 8.0, + "observations": 9.0, + "cohort": "code_understanding|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 19.0, + "observations": 21.0, + "cohort": "code_understanding|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 10.0, + "observations": 10.0, + "cohort": "code_understanding|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "code_understanding|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 3.0, + "cohort": "code_understanding|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 3.0, + "cohort": "code_understanding|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 73.0, + "observations": 76.0, + "cohort": "code_understanding|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 107.0, + "observations": 111.0, + "cohort": "code_understanding|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 87.0, + "observations": 87.0, + "cohort": "code_understanding|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 22.0, + "observations": 22.0, + "cohort": "code_understanding|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 32.0, + "observations": 33.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 44.0, + "observations": 44.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 44.0, + "observations": 47.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 8.0, + "observations": 8.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 1.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 2.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 15.0, + "observations": 16.0, + "cohort": "code_understanding|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 23.0, + "observations": 23.0, + "cohort": "code_understanding|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 12.0, + "observations": 14.0, + "cohort": "code_understanding|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "code_understanding|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 3.0, + "cohort": "code_understanding|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 36.0, + "observations": 43.0, + "cohort": "code_understanding|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 71.0, + "observations": 74.0, + "cohort": "code_understanding|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 48.0, + "observations": 50.0, + "cohort": "code_understanding|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 12.0, + "observations": 13.0, + "cohort": "code_understanding|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 5.0, + "observations": 6.0, + "cohort": "code_understanding|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 10.0, + "observations": 10.0, + "cohort": "code_understanding|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 7.0, + "cohort": "code_understanding|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 1.0, + "cohort": "code_understanding|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_understanding|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 9.0, + "observations": 11.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 18.0, + "observations": 23.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 11.0, + "observations": 13.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 3.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 13.0, + "observations": 17.0, + "cohort": "code_understanding|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 19.0, + "observations": 23.0, + "cohort": "code_understanding|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 27.0, + "observations": 28.0, + "cohort": "code_understanding|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 7.0, + "observations": 8.0, + "cohort": "code_understanding|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 49.0, + "observations": 54.0, + "cohort": "factual_lookup|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 64.0, + "observations": 75.0, + "cohort": "factual_lookup|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 80.0, + "observations": 80.0, + "cohort": "factual_lookup|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 15.0, + "observations": 15.0, + "cohort": "factual_lookup|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 9.0, + "observations": 9.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 12.0, + "observations": 14.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 9.0, + "observations": 10.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 3.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 19.0, + "observations": 22.0, + "cohort": "factual_lookup|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 35.0, + "observations": 38.0, + "cohort": "factual_lookup|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 37.0, + "observations": 41.0, + "cohort": "factual_lookup|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 7.0, + "observations": 7.0, + "cohort": "factual_lookup|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 66.0, + "observations": 75.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 73.0, + "observations": 86.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 69.0, + "observations": 74.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 21.0, + "observations": 21.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 174.0, + "observations": 181.0, + "cohort": "factual_lookup|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 204.0, + "observations": 215.0, + "cohort": "factual_lookup|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 206.0, + "observations": 210.0, + "cohort": "factual_lookup|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 42.0, + "observations": 42.0, + "cohort": "factual_lookup|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 31.0, + "observations": 37.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 41.0, + "observations": 48.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 44.0, + "observations": 45.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 10.0, + "observations": 10.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 1.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 131.0, + "observations": 142.0, + "cohort": "factual_lookup|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 162.0, + "observations": 171.0, + "cohort": "factual_lookup|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 146.0, + "observations": 151.0, + "cohort": "factual_lookup|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 24.0, + "observations": 24.0, + "cohort": "factual_lookup|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 41.0, + "observations": 46.0, + "cohort": "factual_lookup|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 55.0, + "observations": 59.0, + "cohort": "factual_lookup|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 45.0, + "observations": 46.0, + "cohort": "factual_lookup|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 9.0, + "observations": 9.0, + "cohort": "factual_lookup|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1382.0, + "observations": 1473.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2290.0, + "observations": 2348.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1824.0, + "observations": 1853.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 368.0, + "observations": 370.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 5.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 12.0, + "observations": 13.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 7.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 11.0, + "observations": 16.0, + "cohort": "factual_lookup|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 28.0, + "observations": 31.0, + "cohort": "factual_lookup|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 15.0, + "observations": 17.0, + "cohort": "factual_lookup|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "factual_lookup|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 36.0, + "observations": 39.0, + "cohort": "factual_lookup|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 63.0, + "observations": 66.0, + "cohort": "factual_lookup|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 64.0, + "observations": 66.0, + "cohort": "factual_lookup|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 13.0, + "observations": 13.0, + "cohort": "factual_lookup|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 4.0, + "cohort": "factual_lookup|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 5.0, + "observations": 6.0, + "cohort": "factual_lookup|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 5.0, + "cohort": "factual_lookup|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 27.0, + "observations": 31.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 29.0, + "observations": 37.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 26.0, + "observations": 29.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 11.0, + "observations": 11.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 2.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 7.0, + "observations": 8.0, + "cohort": "factual_lookup|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 17.0, + "observations": 20.0, + "cohort": "factual_lookup|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 14.0, + "observations": 14.0, + "cohort": "factual_lookup|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "factual_lookup|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 4.0, + "cohort": "factual_lookup|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 6.0, + "observations": 6.0, + "cohort": "factual_lookup|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 5.0, + "observations": 5.0, + "cohort": "factual_lookup|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 17.0, + "observations": 21.0, + "cohort": "factual_lookup|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 19.0, + "observations": 25.0, + "cohort": "factual_lookup|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 13.0, + "observations": 13.0, + "cohort": "factual_lookup|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "factual_lookup|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2828.0, + "observations": 3552.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 4621.0, + "observations": 5551.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 3628.0, + "observations": 3931.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 911.0, + "observations": 922.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 127.0, + "observations": 406.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 250.0, + "observations": 583.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 219.0, + "observations": 396.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 98.0, + "observations": 107.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 94.0, + "observations": 127.0, + "cohort": "general|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 168.0, + "observations": 196.0, + "cohort": "general|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 135.0, + "observations": 146.0, + "cohort": "general|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 34.0, + "observations": 35.0, + "cohort": "general|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 912.0, + "observations": 1219.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1417.0, + "observations": 1736.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1163.0, + "observations": 1258.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 278.0, + "observations": 283.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 33.0, + "observations": 122.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 41.0, + "observations": 150.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 61.0, + "observations": 96.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 24.0, + "observations": 24.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 16.0, + "observations": 24.0, + "cohort": "general|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 24.0, + "observations": 34.0, + "cohort": "general|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 18.0, + "observations": 21.0, + "cohort": "general|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "general|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 271.0, + "observations": 348.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 488.0, + "observations": 555.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 369.0, + "observations": 386.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 67.0, + "observations": 71.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 3.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 3.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 9.0, + "observations": 9.0, + "cohort": "general|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 14.0, + "observations": 17.0, + "cohort": "general|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 9.0, + "observations": 10.0, + "cohort": "general|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "general|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 340.0, + "observations": 426.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 544.0, + "observations": 635.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 472.0, + "observations": 491.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 120.0, + "observations": 120.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 2.0, + "cohort": "general|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 0.0, + "observations": 1.0, + "cohort": "general|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 8332.0, + "observations": 10133.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 13225.0, + "observations": 15509.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 10614.0, + "observations": 11347.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2499.0, + "observations": 2531.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 437.0, + "observations": 1294.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 863.0, + "observations": 1932.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 738.0, + "observations": 1269.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 296.0, + "observations": 325.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 117.0, + "observations": 173.0, + "cohort": "general|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 183.0, + "observations": 252.0, + "cohort": "general|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 171.0, + "observations": 191.0, + "cohort": "general|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 47.0, + "observations": 52.0, + "cohort": "general|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 927.0, + "observations": 1353.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1550.0, + "observations": 2023.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1273.0, + "observations": 1430.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 348.0, + "observations": 354.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 60.0, + "observations": 255.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 143.0, + "observations": 377.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 145.0, + "observations": 243.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 58.0, + "observations": 61.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 27.0, + "observations": 47.0, + "cohort": "general|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 30.0, + "observations": 48.0, + "cohort": "general|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 30.0, + "observations": 36.0, + "cohort": "general|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "general|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 1119.0, + "observations": 1348.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1816.0, + "observations": 2024.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1472.0, + "observations": 1552.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 301.0, + "observations": 304.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 8.0, + "observations": 10.0, + "cohort": "general|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 9.0, + "observations": 12.0, + "cohort": "general|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 8.0, + "observations": 9.0, + "cohort": "general|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 349.0, + "observations": 439.0, + "cohort": "general|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 549.0, + "observations": 622.0, + "cohort": "general|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 495.0, + "observations": 526.0, + "cohort": "general|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 111.0, + "observations": 113.0, + "cohort": "general|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 4.0, + "observations": 6.0, + "cohort": "general|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 4.0, + "cohort": "general|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 11635.0, + "observations": 12910.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 19249.0, + "observations": 20591.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 15415.0, + "observations": 15867.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3362.0, + "observations": 3392.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 68.0, + "observations": 155.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 113.0, + "observations": 202.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 97.0, + "observations": 124.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 31.0, + "observations": 31.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 6.0, + "observations": 11.0, + "cohort": "general|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 20.0, + "observations": 21.0, + "cohort": "general|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 23.0, + "observations": 24.0, + "cohort": "general|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "general|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 199.0, + "observations": 270.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 334.0, + "observations": 414.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 309.0, + "observations": 337.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 75.0, + "observations": 75.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 5.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 5.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 629.0, + "observations": 778.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1131.0, + "observations": 1286.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 919.0, + "observations": 989.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 227.0, + "observations": 227.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 4.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 3.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 5.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|short|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|short|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 28.0, + "observations": 37.0, + "cohort": "general|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 54.0, + "observations": 64.0, + "cohort": "general|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 41.0, + "observations": 46.0, + "cohort": "general|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 9.0, + "observations": 9.0, + "cohort": "general|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1066.0, + "observations": 1437.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1639.0, + "observations": 2007.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1397.0, + "observations": 1507.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 333.0, + "observations": 341.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 6.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 4.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 65.0, + "observations": 89.0, + "cohort": "general|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 119.0, + "observations": 139.0, + "cohort": "general|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 98.0, + "observations": 103.0, + "cohort": "general|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 21.0, + "observations": 21.0, + "cohort": "general|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 403.0, + "observations": 546.0, + "cohort": "general|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 700.0, + "observations": 875.0, + "cohort": "general|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 556.0, + "observations": 612.0, + "cohort": "general|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 133.0, + "observations": 135.0, + "cohort": "general|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 17.0, + "observations": 27.0, + "cohort": "general|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 19.0, + "observations": 33.0, + "cohort": "general|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 25.0, + "observations": 27.0, + "cohort": "general|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 9.0, + "observations": 9.0, + "cohort": "general|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 210.0, + "observations": 280.0, + "cohort": "general|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 307.0, + "observations": 378.0, + "cohort": "general|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 258.0, + "observations": 289.0, + "cohort": "general|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 64.0, + "observations": 65.0, + "cohort": "general|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 17.0, + "observations": 23.0, + "cohort": "general|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 31.0, + "observations": 39.0, + "cohort": "general|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 16.0, + "observations": 19.0, + "cohort": "general|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "general|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 202.0, + "observations": 282.0, + "cohort": "general|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 361.0, + "observations": 478.0, + "cohort": "general|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 262.0, + "observations": 310.0, + "cohort": "general|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 82.0, + "observations": 82.0, + "cohort": "general|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 5.0, + "cohort": "general|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 9.0, + "observations": 11.0, + "cohort": "general|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 4.0, + "cohort": "general|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 19.0, + "observations": 21.0, + "cohort": "technical_design|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 23.0, + "observations": 25.0, + "cohort": "technical_design|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 28.0, + "observations": 29.0, + "cohort": "technical_design|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "technical_design|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 8.0, + "observations": 9.0, + "cohort": "technical_design|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 9.0, + "observations": 9.0, + "cohort": "technical_design|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 14.0, + "observations": 15.0, + "cohort": "technical_design|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 24.0, + "observations": 25.0, + "cohort": "technical_design|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 18.0, + "observations": 18.0, + "cohort": "technical_design|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 44.0, + "observations": 45.0, + "cohort": "technical_design|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 84.0, + "observations": 86.0, + "cohort": "technical_design|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 74.0, + "observations": 74.0, + "cohort": "technical_design|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 15.0, + "observations": 15.0, + "cohort": "technical_design|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 6.0, + "observations": 6.0, + "cohort": "technical_design|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 23.0, + "observations": 24.0, + "cohort": "technical_design|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 47.0, + "observations": 53.0, + "cohort": "technical_design|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 28.0, + "observations": 29.0, + "cohort": "technical_design|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 10.0, + "observations": 10.0, + "cohort": "technical_design|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 4.0, + "cohort": "technical_design|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 7.0, + "cohort": "technical_design|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 8.0, + "cohort": "technical_design|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 13.0, + "observations": 13.0, + "cohort": "technical_design|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 31.0, + "observations": 31.0, + "cohort": "technical_design|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 23.0, + "observations": 23.0, + "cohort": "technical_design|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "technical_design|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 4.0, + "cohort": "technical_design|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 10.0, + "observations": 15.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 26.0, + "observations": 30.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 19.0, + "observations": 21.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 5.0, + "observations": 6.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 6.0, + "observations": 7.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 6.0, + "observations": 6.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 4.0, + "observations": 6.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 8.0, + "observations": 8.0, + "cohort": "technical_design|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 5.0, + "observations": 6.0, + "cohort": "technical_design|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 5.0, + "observations": 5.0, + "cohort": "technical_design|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 176.0, + "observations": 226.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 267.0, + "observations": 329.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 229.0, + "observations": 245.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 51.0, + "observations": 52.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 2.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 2.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 5.0, + "cohort": "writing|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 4.0, + "observations": 5.0, + "cohort": "writing|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 36.0, + "observations": 47.0, + "cohort": "writing|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 53.0, + "observations": 59.0, + "cohort": "writing|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 47.0, + "observations": 51.0, + "cohort": "writing|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 14.0, + "observations": 15.0, + "cohort": "writing|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 17.0, + "observations": 21.0, + "cohort": "writing|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 27.0, + "observations": 32.0, + "cohort": "writing|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 24.0, + "observations": 24.0, + "cohort": "writing|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "writing|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 2.0, + "cohort": "writing|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 2.0, + "cohort": "writing|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 6.0, + "observations": 6.0, + "cohort": "writing|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 541.0, + "observations": 598.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 892.0, + "observations": 997.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 728.0, + "observations": 756.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 164.0, + "observations": 165.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 7.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 15.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 10.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 25.0, + "observations": 35.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 41.0, + "observations": 63.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 39.0, + "observations": 46.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 12.0, + "observations": 12.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 2.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 4.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 24.0, + "observations": 27.0, + "cohort": "writing|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 49.0, + "observations": 56.0, + "cohort": "writing|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 46.0, + "observations": 49.0, + "cohort": "writing|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 8.0, + "observations": 8.0, + "cohort": "writing|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 11.0, + "observations": 13.0, + "cohort": "writing|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 13.0, + "observations": 13.0, + "cohort": "writing|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 10.0, + "observations": 10.0, + "cohort": "writing|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 327.0, + "observations": 353.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 508.0, + "observations": 555.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 415.0, + "observations": 432.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 92.0, + "observations": 92.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 4.0, + "cohort": "writing|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 7.0, + "cohort": "writing|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "writing|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 3.0, + "cohort": "writing|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 8.0, + "observations": 9.0, + "cohort": "writing|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 4.0, + "cohort": "writing|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 63.0, + "observations": 85.0, + "cohort": "writing|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 111.0, + "observations": 139.0, + "cohort": "writing|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 95.0, + "observations": 99.0, + "cohort": "writing|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 29.0, + "observations": 29.0, + "cohort": "writing|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 5.0, + "cohort": "writing|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 2.0, + "cohort": "writing|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 4.0, + "cohort": "writing|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 26.0, + "observations": 36.0, + "cohort": "writing|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 30.0, + "observations": 44.0, + "cohort": "writing|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 23.0, + "observations": 28.0, + "cohort": "writing|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 8.0, + "observations": 8.0, + "cohort": "writing|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 14.0, + "observations": 14.0, + "cohort": "writing|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 14.0, + "observations": 16.0, + "cohort": "writing|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 6.0, + "observations": 7.0, + "cohort": "writing|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "writing|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 5.0, + "observations": 5.0, + "cohort": "writing|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 8.0, + "cohort": "writing|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 7.0, + "cohort": "writing|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 1.0, + "cohort": "writing|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 0.0, + "observations": 1.0, + "cohort": "writing|very_long|code=1|math=1|mc=1|intl=0" + } + ], + "domain_prior_mass": 200.0, + "cohort_prior_mass": 20.0, + "routing_threshold": 0.75, + "datasets": [ + { + "name": "openbmb/UltraFeedback", + "url": "https://huggingface.co/datasets/openbmb/UltraFeedback", + "license": "MIT", + "rows": 255864, + "success_definition": "UltraFeedback overall_score >= 4" + } + ], + "success_definition": "UltraFeedback overall_score >= 4", + "split_method": "sha256(prompt): 70% train, 15% validation, 15% test" +} diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 577cee0920d..a123a75dd81 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -33,6 +33,11 @@ from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload from litellm.llms.base_llm.base_utils import type_to_response_format_param +from litellm.router_strategy.adaptive_router.classifier import classify_prompt +from litellm.router_strategy.complexity_router.tier_predictor import ( + TierSuccessPredictor, + resolve_tier_artifact, +) from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, ModelResponse, @@ -790,6 +795,7 @@ class ClassificationOutcome(NamedTuple): signals: tuple[str, ...] cause: Literal[ "heuristic_scorer", + "heuristic_v2", "reasoning_override", "llm_classifier", "heuristic_first_short_circuit", @@ -978,6 +984,11 @@ class ComplexityRouter(CustomLogger): if llm_classifier_configured else None ) + self._tier_success_predictor: TierSuccessPredictor | None = ( + TierSuccessPredictor(resolve_tier_artifact(self.config.heuristic_v2_artifact)) + if self.config.classifier_type == "heuristic_v2" + else None + ) verbose_router_logger.debug("ComplexityRouter initialized for %s with tiers: %s", model_name, self.config.tiers) @@ -1350,6 +1361,8 @@ class ComplexityRouter(CustomLogger): custom tier set, and classifier_fallback otherwise decides between the heuristic scorer and default_model. The outcome's `cause` reports which path actually ran. """ + if self.config.classifier_type == "heuristic_v2": + return self._classify_with_heuristic_v2(prompt) if self.config.classifier_type == "custom": return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages) if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None: @@ -1359,6 +1372,24 @@ class ComplexityRouter(CustomLogger): return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause) return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages) + def _classify_with_heuristic_v2(self, prompt: str) -> ClassificationOutcome: + predictor: Final = self._tier_success_predictor + if predictor is None: + raise ValueError("heuristic v2 predictor is not configured") + request_type: Final = classify_prompt(prompt) + prediction: Final = predictor.predict(prompt, request_type) + tier: Final = TIER_SEVERITY_ORDER[prediction.required_tier - 1] + probability_signals: Final = tuple( + f"tier-probability:{candidate.value.lower()}={prediction.probabilities[index]:.6f}" + for index, candidate in enumerate(TIER_SEVERITY_ORDER, start=1) + ) + return ClassificationOutcome( + tier=tier, + score=None, + signals=(f"request-type:{request_type.value}", *probability_signals), + cause="heuristic_v2", + ) + async def _classify_heuristic_first( self, prompt: str, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 70aeecb31c6..fdf2a3a0b39 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -14,6 +14,8 @@ from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_seriali from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin +from .tier_predictor import TrainedTierArtifact + class ComplexityTier(str, Enum): """Complexity tiers for routing decisions.""" @@ -625,12 +627,19 @@ class ComplexityRouterConfig(BaseModel): ) # Classifier strategy - classifier_type: Literal["heuristic", "llm", "custom", "heuristic_first"] = Field( + classifier_type: Literal["heuristic", "heuristic_v2", "llm", "custom", "heuristic_first"] = Field( default="heuristic", description=( - "Classification strategy: local regex/keyword scoring, an LLM call, a custom classifier " - "plugin, or 'heuristic_first', which scores locally and only pays for the LLM classifier " - "when the local scorer does not confidently land a cheap tier" + "Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, " + "an LLM call, a custom classifier plugin, or 'heuristic_first', which scores locally and only pays " + "for the LLM classifier when the local scorer does not confidently land a cheap tier" + ), + ) + heuristic_v2_artifact: TrainedTierArtifact | Literal["ultrafeedback"] = Field( + default="ultrafeedback", + description=( + "Success-probability artifact used by classifier_type 'heuristic_v2'. The bundled " + "UltraFeedback artifact is selected by default; an inline trained artifact may replace it" ), ) classifier_llm_config: ClassifierLLMConfig | None = Field( @@ -1248,10 +1257,10 @@ class ComplexityRouterConfig(BaseModel): ) if duplicated: raise ValueError(f"tier_definitions names must be unique (case-insensitive): {', '.join(duplicated)}") - if self.classifier_type in ("heuristic", "heuristic_first"): + if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first"): raise ValueError( "tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only " - "produces the built-in tiers" + "produces the four built-in tiers, as does heuristic_v2" ) conflicts: Final = self._tier_definition_conflicts() if conflicts: diff --git a/litellm/router_strategy/complexity_router/tier_predictor.py b/litellm/router_strategy/complexity_router/tier_predictor.py new file mode 100644 index 00000000000..764f6e6ad56 --- /dev/null +++ b/litellm/router_strategy/complexity_router/tier_predictor.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import re +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal + +from pydantic import BaseModel, Field, model_validator + +from litellm.types.router import RequestType + + +class TierGlobalStatistic(BaseModel): + tier: int = Field(ge=1, le=4) + successes: float = Field(ge=0.0) + observations: float = Field(gt=0.0) + + @model_validator(mode="after") + def _successes_do_not_exceed_observations(self) -> TierGlobalStatistic: + if self.successes > self.observations: + raise ValueError("successes cannot exceed observations") + return self + + +class TierDomainStatistic(TierGlobalStatistic): + request_type: RequestType + + +class TierCohortStatistic(TierGlobalStatistic): + cohort: str = Field(min_length=1) + + +class TierDataset(BaseModel): + name: str = Field(min_length=1) + url: str = Field(min_length=1) + license: str = Field(min_length=1) + rows: int = Field(gt=0) + success_definition: str = Field(default="quality score meets the dataset success threshold", min_length=1) + + +class TrainedTierArtifact(BaseModel): + schema_version: Literal[1] = 1 + global_statistics: tuple[TierGlobalStatistic, ...] + domain_statistics: tuple[TierDomainStatistic, ...] = () + cohort_statistics: tuple[TierCohortStatistic, ...] = () + domain_prior_mass: float = Field(default=200.0, gt=0.0) + cohort_prior_mass: float = Field(default=20.0, gt=0.0) + routing_threshold: float = Field(default=0.75, ge=0.0, le=1.0) + datasets: tuple[TierDataset, ...] = () + success_definition: str = Field(default="quality score meets the dataset success threshold", min_length=1) + split_method: str = Field(default="sha256(prompt): 70% train, 15% validation, 15% test", min_length=1) + + @model_validator(mode="after") + def _statistics_are_unique(self) -> TrainedTierArtifact: + global_tiers: Final = tuple(stat.tier for stat in self.global_statistics) + if frozenset(global_tiers) != frozenset((1, 2, 3, 4)) or len(global_tiers) != 4: + raise ValueError("global statistics must contain each tier exactly once") + domain_keys: Final = tuple((stat.request_type, stat.tier) for stat in self.domain_statistics) + if len(domain_keys) != len(frozenset(domain_keys)): + raise ValueError("domain statistics must contain unique request_type and tier pairs") + cohort_keys: Final = tuple((stat.cohort, stat.tier) for stat in self.cohort_statistics) + if len(cohort_keys) != len(frozenset(cohort_keys)): + raise ValueError("cohort statistics must contain unique cohort and tier pairs") + return self + + +_CODE_PATTERN: Final = re.compile( + r"```|\b(def|class|function|python|javascript|typescript|sql|code)\b", + re.IGNORECASE, +) +_MATH_PATTERN: Final = re.compile( + r"\b(solve|calculate|equation|probability|theorem|proof|integral)\b|[$=]", + re.IGNORECASE, +) +_MULTIPLE_CHOICE_PATTERN: Final = re.compile(r"(?:^|\s)[A-D][.)]\s") +_TIERS: Final = (1, 2, 3, 4) +_BUILTIN_ARTIFACTS: Final = MappingProxyType({"ultrafeedback": "ultrafeedback_tiers.json"}) + + +def resolve_tier_artifact(artifact: TrainedTierArtifact | str) -> TrainedTierArtifact: + if isinstance(artifact, TrainedTierArtifact): + return artifact + filename: Final = _BUILTIN_ARTIFACTS.get(artifact) + if filename is None: + raise ValueError(f"unknown complexity router tier artifact: {artifact}") + path: Final = Path(__file__).with_name("artifacts") / filename + return TrainedTierArtifact.model_validate_json(path.read_text()) + + +def similarity_cohort(prompt: str, request_type: RequestType) -> str: + length: Final = len(prompt) + length_bucket: Final = ( + "short" if length < 200 else "medium" if length < 800 else "long" if length < 2000 else "very_long" + ) + code: Final = int(bool(_CODE_PATTERN.search(prompt))) + math: Final = int(bool(_MATH_PATTERN.search(prompt))) + multiple_choice: Final = int(bool(_MULTIPLE_CHOICE_PATTERN.search(prompt))) + non_ascii: Final = int(sum(ord(character) > 127 for character in prompt) / max(1, length) > 0.1) + return f"{request_type.value}|{length_bucket}|code={code}|math={math}|mc={multiple_choice}|intl={non_ascii}" + + +@dataclass(frozen=True, slots=True) +class TierPrediction: + probabilities: Mapping[int, float] + required_tier: int + + +class TierSuccessPredictor: + def __init__(self, artifact: TrainedTierArtifact) -> None: + self._artifact = artifact + self._global: Mapping[int, TierGlobalStatistic] = MappingProxyType( + {stat.tier: stat for stat in artifact.global_statistics} + ) + self._domain: Mapping[tuple[RequestType, int], TierDomainStatistic] = MappingProxyType( + {(stat.request_type, stat.tier): stat for stat in artifact.domain_statistics} + ) + self._cohort: Mapping[tuple[str, int], TierCohortStatistic] = MappingProxyType( + {(stat.cohort, stat.tier): stat for stat in artifact.cohort_statistics} + ) + + @property + def routing_threshold(self) -> float: + return self._artifact.routing_threshold + + def predict(self, prompt: str, request_type: RequestType) -> TierPrediction: + cohort: Final = similarity_cohort(prompt, request_type) + raw: Final = tuple(self._probability(tier, request_type, cohort) for tier in _TIERS) + monotonic: Final = tuple(max(raw[:index]) for index in range(1, len(raw) + 1)) + probabilities: Final[Mapping[int, float]] = MappingProxyType( + {int(tier): probability for tier, probability in zip(_TIERS, monotonic)} + ) + required_tier: Final = next( + (tier for tier in _TIERS if probabilities[tier] >= self._artifact.routing_threshold), + 4, + ) + return TierPrediction(probabilities=probabilities, required_tier=required_tier) + + def _probability(self, tier: int, request_type: RequestType, cohort: str) -> float: + global_stat: Final = self._global[tier] + global_mean: Final = (global_stat.successes + 1.0) / (global_stat.observations + 2.0) + domain_stat: Final = self._domain.get((request_type, tier)) + domain_mean: Final = self._posterior_mean(domain_stat, self._artifact.domain_prior_mass, global_mean) + cohort_stat: Final = self._cohort.get((cohort, tier)) + return self._posterior_mean(cohort_stat, self._artifact.cohort_prior_mass, domain_mean) + + @staticmethod + def _posterior_mean( + statistic: TierGlobalStatistic | None, + prior_mass: float, + prior_mean: float, + ) -> float: + if statistic is None: + return prior_mean + return (statistic.successes + prior_mass * prior_mean) / (statistic.observations + prior_mass) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5783a39b30c..09c01873a9d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2839,6 +2839,7 @@ class StandardLoggingRoutingDecisionTierBoundaries(TypedDict): RoutingDecisionCause = Literal[ "heuristic_scorer", + "heuristic_v2", # The scorer found 2+ reasoning markers and forced REASONING regardless of score. # A distinct cause rather than a marker inside `signals`, because it is the fact # that tells a reader the score did NOT choose the tier; encoding it as free text diff --git a/pyproject.toml b/pyproject.toml index 60162544612..d0e5723d1cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -278,7 +278,10 @@ bindings = "pyo3" features = ["extension-module"] profile = "release" editable-profile = "dev" -include = ["litellm/proxy/_experimental/out/**"] +include = [ + "litellm/proxy/_experimental/out/**", + "litellm/router_strategy/complexity_router/artifacts/*.json", +] exclude = [ "litellm/proxy/enterprise", "litellm/proxy/enterprise/**", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 1ec8be88c9b..82659571de4 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -12,7 +12,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from pydantic import ValidationError - import litellm from litellm import Router from litellm._logging import verbose_router_logger @@ -34,10 +33,14 @@ from litellm.router_strategy.complexity_router.config import ( DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, DEFAULT_COMPLEXITY_CONFIG, DEFAULT_TECHNICAL_KEYWORDS, + ClassificationRubric, ClassifierLLMConfig, ComplexityRouterConfig, ComplexityTier, - ClassificationRubric, +) +from litellm.router_strategy.complexity_router.tier_predictor import ( + TierGlobalStatistic, + TrainedTierArtifact, ) from litellm.types.router import ( Deployment, @@ -46,6 +49,16 @@ from litellm.types.router import ( ) +def _heuristic_v2_artifact() -> TrainedTierArtifact: + return TrainedTierArtifact( + global_statistics=tuple( + TierGlobalStatistic(tier=tier, successes=successes, observations=100) + for tier, successes in enumerate((10, 20, 90, 99), start=1) + ), + routing_threshold=0.8, + ) + + @pytest.fixture def mock_router_instance(): """Create a mock LiteLLM Router instance.""" @@ -1696,6 +1709,59 @@ class TestLLMClassifier: assert outcome.cause == "heuristic_scorer" assert outcome.score is not None + @pytest.mark.asyncio + async def test_heuristic_v2_routes_directly_to_predicted_builtin_tier(self, mock_router_instance): + router = ComplexityRouter( + model_name="tier-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "classifier_type": "heuristic_v2", + "heuristic_v2_artifact": _heuristic_v2_artifact(), + "tiers": { + "SIMPLE": "simple-model", + "MEDIUM": "medium-model", + "COMPLEX": "complex-model", + "REASONING": "reasoning-model", + }, + }, + ) + + response = await router.async_pre_routing_hook( + model="tier-router", + request_kwargs={}, + messages=[{"role": "user", "content": "Handle this new request"}], + ) + + assert response is not None + assert response.model == "complex-model" + assert response.routing_decision["tier"] == "COMPLEX" + assert response.routing_decision["cause"] == "heuristic_v2" + assert response.routing_decision["signals"] == [ + "request-type:general", + "tier-probability:simple=0.107843", + "tier-probability:medium=0.205882", + "tier-probability:complex=0.892157", + "tier-probability:reasoning=0.980392", + ] + + def test_heuristic_v2_needs_no_classifier_model(self): + config = ComplexityRouterConfig(classifier_type="heuristic_v2") + + assert config.classifier_llm_config is None + assert config.heuristic_v2_artifact == "ultrafeedback" + + def test_heuristic_v2_rejects_custom_tier_definitions(self): + with pytest.raises(ValidationError, match="as does heuristic_v2"): + ComplexityRouterConfig( + classifier_type="heuristic_v2", + tier_definitions=( + {"name": "low", "description": "easy work"}, + {"name": "high", "description": "hard work"}, + ), + tiers={"low": "cheap", "high": "expensive"}, + fallback_tier="high", + ) + @pytest.mark.asyncio async def test_aclassify_llm_success_routes_by_llm_verdict(self, llm_complexity_router, mock_router_instance): """A well-formed structured LLM response should decide the tier directly. diff --git a/tests/test_litellm/router_strategy/test_complexity_tier_predictor.py b/tests/test_litellm/router_strategy/test_complexity_tier_predictor.py new file mode 100644 index 00000000000..5bbe0fb5669 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_complexity_tier_predictor.py @@ -0,0 +1,91 @@ +from typing import Final + +import pytest + +from litellm.router_strategy.complexity_router.tier_predictor import ( + TierCohortStatistic, + TierDomainStatistic, + TierGlobalStatistic, + TierSuccessPredictor, + TrainedTierArtifact, + resolve_tier_artifact, + similarity_cohort, +) +from litellm.types.router import RequestType + + +def _artifact( + global_successes: tuple[float, float, float, float] = (4.0, 5.0, 6.0, 7.0), + threshold: float = 0.75, + domain_statistics: tuple[TierDomainStatistic, ...] = (), + cohort_statistics: tuple[TierCohortStatistic, ...] = (), +) -> TrainedTierArtifact: + return TrainedTierArtifact( + global_statistics=tuple( + TierGlobalStatistic(tier=tier, successes=successes, observations=10.0) + for tier, successes in enumerate(global_successes, start=1) + ), + domain_statistics=domain_statistics, + cohort_statistics=cohort_statistics, + domain_prior_mass=10.0, + cohort_prior_mass=10.0, + routing_threshold=threshold, + ) + + +def test_predictions_are_monotonic_across_tiers() -> None: + predictor: Final = TierSuccessPredictor(_artifact(global_successes=(9.0, 2.0, 7.0, 6.0))) + + prediction: Final = predictor.predict("hello", RequestType.GENERAL) + + probabilities: Final = tuple(prediction.probabilities.values()) + assert probabilities == tuple(sorted(probabilities)) + + +def test_domain_and_cohort_statistics_back_off_hierarchically() -> None: + matching_cohort: Final = similarity_cohort("hello", RequestType.GENERAL) + artifact: Final = _artifact( + global_successes=(1.0, 5.0, 6.0, 7.0), + domain_statistics=( + TierDomainStatistic( + tier=1, + request_type=RequestType.GENERAL, + successes=10.0, + observations=10.0, + ), + ), + cohort_statistics=( + TierCohortStatistic( + tier=1, + cohort=matching_cohort, + successes=0.0, + observations=10.0, + ), + ), + ) + predictor: Final = TierSuccessPredictor(artifact) + + cohort_probability: Final = predictor.predict("hello", RequestType.GENERAL).probabilities[1] + domain_probability: Final = predictor.predict("hello " * 100, RequestType.GENERAL).probabilities[1] + global_probability: Final = predictor.predict("hello", RequestType.WRITING).probabilities[1] + + assert cohort_probability == pytest.approx(7.0 / 24.0) + assert domain_probability == pytest.approx(7.0 / 12.0) + assert global_probability == pytest.approx(1.0 / 6.0) + + +def test_selects_first_tier_above_probability_threshold() -> None: + predictor: Final = TierSuccessPredictor(_artifact(global_successes=(4.0, 6.0, 8.0, 9.0), threshold=0.7)) + + prediction: Final = predictor.predict("hello", RequestType.GENERAL) + + assert prediction.required_tier == 3 + + +def test_builtin_ultrafeedback_artifact_is_loadable() -> None: + artifact: Final = resolve_tier_artifact("ultrafeedback") + + assert artifact.routing_threshold == 0.75 + assert artifact.domain_prior_mass == 200.0 + assert artifact.cohort_prior_mass == 20.0 + assert artifact.datasets[0].license == "MIT" diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 96c93306611..c2ed4b1f55a 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -43,6 +43,10 @@ const DEFAULT_SCORING_EXPLANATION = "The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical " + "terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the tier:"; +const HEURISTIC_V2_EXPLANATION = + "The router estimates success probability for all four tiers with the bundled calibrated model, then selects " + + "the first tier that meets its trained threshold. It runs locally with no classifier API call."; + const CLASSIFIER_TIMEOUT_ID = "classifier-timeout-ms"; const CLASSIFIER_CONTEXT_WINDOW_SIZE_ID = "classifier-context-window-size"; const CLASSIFIER_CONTEXT_BUDGET_CHARS_ID = "classifier-context-budget-chars"; @@ -62,6 +66,7 @@ const CUSTOM_PROMPT_WITH_DEFAULT_MODEL_FALLBACK = * at all, so the panel must not keep implying a score is involved on either router. */ const scoringExplanation = (value: ComplexityRouterConfigValue): string => { + if (value.classifier_type === "heuristic_v2") return HEURISTIC_V2_EXPLANATION; const usesCustomPrompt = usesLlmClassifier(value.classifier_type) && Boolean(value.classifier_llm_config?.system_prompt?.trim()); if (!usesCustomPrompt) return DEFAULT_SCORING_EXPLANATION; @@ -179,6 +184,17 @@ const ClassifierTypeRadios: React.FC<{ + + + + + + ); @@ -263,6 +277,8 @@ const ClassificationMethodConfig: React.FC = ({ classifierType === "heuristic_first" ? value.heuristic_first_max_tier ?? DEFAULT_HEURISTIC_FIRST_MAX_TIER : undefined, + hybrid_boundary_margin: + classifierType === "hybrid" ? value.hybrid_boundary_margin ?? DEFAULT_HYBRID_BOUNDARY_MARGIN : undefined, }; onChange(nextValue); }; @@ -271,6 +287,13 @@ const ClassificationMethodConfig: React.FC = ({ onChange({ ...value, heuristic_first_max_tier: tier }); }; + const handleHybridBoundaryMarginChange = (raw: string) => { + setDraft({ id: HYBRID_BOUNDARY_MARGIN_ID, raw }); + const parsed = Number(raw); + if (raw.trim() === "" || !Number.isFinite(parsed)) return; + onChange({ ...value, hybrid_boundary_margin: Math.min(1, Math.max(0, parsed)) }); + }; + const handleClassificationPromptChange = (classificationPrompt: string | undefined) => { onChange({ ...value, classification_prompt: classificationPrompt }); }; @@ -391,6 +414,30 @@ const ClassificationMethodConfig: React.FC = ({ )} + {classifierType === "hybrid" && ( +
+ Boundary margin + handleHybridBoundaryMarginChange(event.target.value)} + onBlur={() => setDraft(null)} + className="w-full" + /> +

+ A score further than this from every tier boundary routes on the scorer's own tier, however expensive + that tier is. A score closer than this, and anything the scorer found no signal for at all, goes to the + classifier to break the tie +

+
+ )} +
How often to classify - classifierType === "llm" || classifierType === "heuristic_first"; + classifierType === "llm" || classifierType === "heuristic_first" || classifierType === "hybrid"; export type ClassifierFallback = "heuristic" | "default_model"; @@ -162,7 +162,8 @@ export const heuristicScoringRoleFor = ( classifierFallback: ClassifierFallback | undefined, ): HeuristicScoringRole => { if (classifierType === "heuristic_v2") return "never"; - if (classifierType === "heuristic" || classifierType === "heuristic_first") return "decides"; + if (classifierType === "heuristic" || classifierType === "heuristic_first" || classifierType === "hybrid") + return "decides"; return (classifierFallback ?? DEFAULT_CLASSIFIER_FALLBACK) === "heuristic" ? "fallback_only" : "never"; }; @@ -404,6 +405,8 @@ export interface ComplexityRouterConfigValue { classification_prompt?: string; /** Highest tier the scorer may decide alone under heuristic_first. Required by that type, rejected by the others. */ heuristic_first_max_tier?: string; + /** How near a tier boundary a score may land before hybrid defers to the classifier. Required by that type, rejected by the others. */ + hybrid_boundary_margin?: number; classification_mode?: ClassificationMode; session_affinity?: boolean; modality_routing?: boolean; @@ -516,6 +519,9 @@ export const effectiveTierLabel = (tier: keyof ComplexityTiers, tierLabels: Comp export const DEFAULT_HEURISTIC_FIRST_MAX_TIER = "SIMPLE"; +/** What the Hybrid radio starts at. Required by that type, so the form always has a value to send. */ +export const DEFAULT_HYBRID_BOUNDARY_MARGIN = 0.03; + /** * Tiers the heuristic_first threshold may name. The top tier is excluded because it would short * circuit every request and leave the classifier unreachable, which the backend rejects. diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 03b7432e4ec..81acb7b7e50 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -342,6 +342,7 @@ const AddAutoRouterTab: React.FC = ({ planModeMinTier: complexityRouterConfig.plan_mode_min_tier, classificationPrompt: complexityRouterConfig.classification_prompt, heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier, + hybridBoundaryMargin: complexityRouterConfig.hybrid_boundary_margin, classificationMode: complexityRouterConfig.classification_mode, tierLabels: complexityRouterConfig.tier_labels, classifierType: complexityRouterConfig.classifier_type, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 2d8f55b2b64..c6370541a45 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -817,6 +817,39 @@ describe("heuristic_first", () => { }); }); +describe("hybrid", () => { + const hybridParams: BuildComplexityRouterConfigParams = { + ...baseParams, + classifierType: "hybrid", + hybridBoundaryMargin: 0.03, + classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 }, + classifierFallback: "default_model", + }; + + it("emits hybrid_boundary_margin, zero included since exactly-on-a-boundary is a real setting", () => { + expect(buildComplexityRouterConfig(hybridParams).hybrid_boundary_margin).toBe(0.03); + expect(buildComplexityRouterConfig({ ...hybridParams, hybridBoundaryMargin: 0 }).hybrid_boundary_margin).toBe(0); + }); + + it("keeps every classifier key the operator set, since hybrid still calls the classifier", () => { + const config = buildComplexityRouterConfig(hybridParams); + expect(config.classifier_type).toBe("hybrid"); + expect(config.classifier_llm_config).toEqual({ model: "gpt-4o-mini", timeout_ms: 3000 }); + expect(config.classifier_fallback).toBe("default_model"); + }); + + it("omits hybrid_boundary_margin on every other classifier type, which the backend rejects it on", () => { + for (const classifierType of ["heuristic", "llm", "heuristic_first"] as const) { + const config = buildComplexityRouterConfig({ + ...hybridParams, + classifierType, + ...(classifierType === "heuristic_first" && { heuristicFirstMaxTier: "SIMPLE" }), + }); + expect(config.hybrid_boundary_margin).toBeUndefined(); + } + }); +}); + describe("classification_mode", () => { it("emits user_turn", () => { const config = buildComplexityRouterConfig({ ...baseParams, classificationMode: "user_turn" }); @@ -924,10 +957,12 @@ describe("buildComplexityRouterConfig with an edited tier set", () => { dimensionWeights: { length: 1 }, reasoningOverrideMinScore: 0.5, heuristicFirstMaxTier: "SIMPLE", + hybridBoundaryMargin: 0.03, customTechnicalKeywords: ["kubernetes"], }; const emittingType = key === "heuristic_first_max_tier" ? "heuristic_first" : "llm"; - expect(buildComplexityRouterConfig({ ...baseParams, ...loaded, classifierType: emittingType })).toHaveProperty(key); + const typeForKey = key === "hybrid_boundary_margin" ? "hybrid" : emittingType; + expect(buildComplexityRouterConfig({ ...baseParams, ...loaded, classifierType: typeForKey })).toHaveProperty(key); expect(build(loaded)).not.toHaveProperty(key); }); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index af34dc92c0f..94c5badf6d1 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -107,6 +107,7 @@ export interface BuildComplexityRouterConfigParams { classifierFallback: ClassifierFallback | undefined; classificationPrompt: string | undefined; heuristicFirstMaxTier: string | undefined; + hybridBoundaryMargin?: number; classificationMode: ClassificationMode | undefined; sessionAffinity: boolean; modalityRouting?: boolean; @@ -163,6 +164,7 @@ export interface ComplexityRouterConfigPayload { classifier_fallback?: ClassifierFallback; classification_prompt?: string; heuristic_first_max_tier?: string; + hybrid_boundary_margin?: number; classification_mode: ClassificationMode; session_affinity: boolean; deployment_affinity: boolean; @@ -352,6 +354,7 @@ const classifierWireFields = ( classifierLlmConfig, classifierFallback, heuristicFirstMaxTier, + hybridBoundaryMargin, classifierContextWindowSize, classifierContextBudgetChars, classifierContextIncludeAssistantTurns, @@ -360,6 +363,7 @@ const classifierWireFields = ( | "classifierLlmConfig" | "classifierFallback" | "heuristicFirstMaxTier" + | "hybridBoundaryMargin" | "classifierContextWindowSize" | "classifierContextBudgetChars" | "classifierContextIncludeAssistantTurns" @@ -371,6 +375,8 @@ const classifierWireFields = ( classifierFallback !== undefined && { classifier_fallback: classifierFallback }), ...(effectiveType === "heuristic_first" && heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }), + ...(effectiveType === "hybrid" && + hybridBoundaryMargin !== undefined && { hybrid_boundary_margin: hybridBoundaryMargin }), ...(usesLlmClassifier(effectiveType) && classifierContextWindowSize !== undefined && { classifier_context_window_size: classifierContextWindowSize, @@ -399,6 +405,7 @@ export const buildComplexityRouterConfig = ({ classifierFallback, classificationPrompt, heuristicFirstMaxTier, + hybridBoundaryMargin, classificationMode, sessionAffinity, modalityRouting, @@ -444,6 +451,7 @@ export const buildComplexityRouterConfig = ({ classifierLlmConfig, classifierFallback, heuristicFirstMaxTier, + hybridBoundaryMargin, classifierContextWindowSize, classifierContextBudgetChars, classifierContextIncludeAssistantTurns, diff --git a/ui/litellm-dashboard/src/components/add_model/tier_rows.ts b/ui/litellm-dashboard/src/components/add_model/tier_rows.ts index d2706a82037..3c5b149f4da 100644 --- a/ui/litellm-dashboard/src/components/add_model/tier_rows.ts +++ b/ui/litellm-dashboard/src/components/add_model/tier_rows.ts @@ -122,10 +122,10 @@ export const CUSTOM_TIER_RESTRICTIONS = { reason: "Session pinning escalates along the built-in tier ladder, which your tier set replaces", }, heuristicClassifier: { - omit: ["heuristic_first_max_tier"], + omit: ["heuristic_first_max_tier", "hybrid_boundary_margin"], reason: "The heuristic scorer only produces the built-in tiers, so an edited set needs the LLM classifier. " + - "Heuristic first is out for the same reason: its local scorer decides the cheap traffic", + "Heuristic first and hybrid are out for the same reason: their local scorer decides the traffic it is sure of", }, heuristicScoring: { omit: [ diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 3e63cbd3b32..fae6744d3ac 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -520,16 +520,21 @@ describe("managed keys survive an untouched open-and-save", () => { }; // tier_definitions, fallback_tier and classification_prompt cannot sit beside heuristic_first, which - // this fixture uses, so no single stored config can hold every managed key. They get their own round - // trip below. - const CUSTOM_TIER_ONLY_KEYS = new Set(["tier_definitions", "fallback_tier", "classification_prompt"]); + // this fixture uses, and hybrid_boundary_margin belongs to the sibling hybrid type, so no single + // stored config can hold every managed key. Each gets its own round trip below. + const KEYS_ANOTHER_CLASSIFIER_TYPE_OWNS = new Set([ + "tier_definitions", + "fallback_tier", + "classification_prompt", + "hybrid_boundary_margin", + ]); it("carries every managed key a built-in router can hold through hydrate then save", () => { const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined); const saved = buildUpdatedComplexityRouterConfig(STORED_ALL_MANAGED, hydrated); const dropped = [...MANAGED_COMPLEXITY_ROUTER_KEYS] - .filter((key) => !CUSTOM_TIER_ONLY_KEYS.has(key)) + .filter((key) => !KEYS_ANOTHER_CLASSIFIER_TYPE_OWNS.has(key)) .filter((key) => saved[key] === undefined); expect(dropped).toEqual([]); }); @@ -583,6 +588,18 @@ describe("managed keys survive an untouched open-and-save", () => { ); }); + it("round-trips a hybrid router's margin, which save requires and the backend rejects without", () => { + const storedHybrid: Record = { + ...STORED_ALL_MANAGED, + classifier_type: "hybrid", + hybrid_boundary_margin: 0.05, + }; + delete storedHybrid.heuristic_first_max_tier; + const hydrated = hydrateComplexityRouterConfig(storedHybrid, undefined); + expect(hydrated.hybrid_boundary_margin).toBe(0.05); + expect(buildUpdatedComplexityRouterConfig(storedHybrid, hydrated).hybrid_boundary_margin).toBe(0.05); + }); + it("round-trips the heuristic_first threshold, which save requires and the backend rejects without", () => { const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined); expect(hydrated.heuristic_first_max_tier).toBe("SIMPLE"); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index e18582f77a0..7aa5f4beb1f 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -89,6 +89,7 @@ export interface StoredComplexityRouterConfig { plan_mode_min_tier?: unknown; classification_prompt?: unknown; heuristic_first_max_tier?: unknown; + hybrid_boundary_margin?: unknown; tier_labels?: unknown; classifier_type?: ClassifierType; classifier_llm_config?: ClassifierLLMConfig; @@ -167,6 +168,8 @@ export const hydrateComplexityRouterConfig = ( typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" ? parsedConfig.heuristic_first_max_tier : undefined, + hybrid_boundary_margin: + typeof parsedConfig.hybrid_boundary_margin === "number" ? parsedConfig.hybrid_boundary_margin : undefined, classification_mode: parsedConfig.classification_mode === "user_turn" || parsedConfig.classification_mode === "every_request" ? parsedConfig.classification_mode @@ -214,6 +217,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classifier_fallback", "classification_prompt", "heuristic_first_max_tier", + "hybrid_boundary_margin", "classification_mode", "session_affinity", "modality_routing", @@ -303,6 +307,7 @@ export const buildUpdatedComplexityRouterConfig = ( planModeMinTier: value.plan_mode_min_tier, classificationPrompt: value.classification_prompt, heuristicFirstMaxTier: value.heuristic_first_max_tier, + hybridBoundaryMargin: value.hybrid_boundary_margin, classificationMode: value.classification_mode, tierLabels: value.tier_labels, classifierType: value.classifier_type, diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx index 66e80dcc655..ffef01010b7 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -86,6 +86,7 @@ const CONSTANT_CAUSE_LABELS: Record = { heuristic_scorer: "Heuristic scorer", heuristic_v2: "Heuristic v2", heuristic_first_short_circuit: "Heuristic scorer, classifier skipped", + hybrid_short_circuit: "Heuristic scorer, score clear of every boundary", classifier_plugin: "Custom classifier plugin", semantic_keyword_match: "Semantic keyword match", session_affinity_pin: "Pinned to session", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 79e44de4f1f..be81450b941 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34560,7 +34560,7 @@ export interface components { * @enum {string} */ classifier_fallback: "heuristic" | "default_model"; - /** @description Configuration for the LLM classifier; required when classifier_type is 'llm' or 'heuristic_first' */ + /** @description Configuration for the LLM classifier; required when classifier_type is 'llm', 'heuristic_first' or 'hybrid' */ classifier_llm_config?: components["schemas"]["ClassifierLLMConfig"] | null; /** * Classifier Plugin @@ -34575,11 +34575,11 @@ export interface components { classifier_plugin_timeout_ms: number; /** * Classifier Type - * @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM call, a custom classifier plugin, or 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier + * @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM call, a custom classifier plugin, 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', which trusts the local scorer everywhere except when its score lands near a tier boundary * @default heuristic * @enum {string} */ - classifier_type: "heuristic" | "heuristic_v2" | "llm" | "custom" | "heuristic_first"; + classifier_type: "heuristic" | "heuristic_v2" | "llm" | "custom" | "heuristic_first" | "hybrid"; /** * Code Keywords * @description Keywords indicating code-related content @@ -34651,6 +34651,11 @@ export interface components { * @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings. */ housekeeping_patterns?: string[] | null; + /** + * Hybrid Boundary Margin + * @description How close to a tier boundary a heuristic score has to land before the LLM classifier breaks the tie; required when classifier_type is 'hybrid' and rejected otherwise. Everything further than this from every active boundary routes on the scorer's own tier with no classifier call, at any tier, which is what separates 'hybrid' from 'heuristic_first' and its cheap-tier ceiling. A prompt where no dimension fired still goes to the classifier, since the scorer has no opinion to be near a boundary with. 0 escalates only scores sitting exactly on a boundary. + */ + hybrid_boundary_margin?: number | null; /** * Keyword Tier Rules * @description Rules that force a specific tier when their keywords match the prompt @@ -35867,7 +35872,7 @@ export interface components { * Cause * @enum {string} */ - cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; + cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; /** Classifier Cost */ classifier_cost?: number; /** Classifier Model */ From 1df402b80d2d3f4fed8d504a5a73946162c133b6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:53:05 -0700 Subject: [PATCH 66/69] fix(vector-store): carry request metadata into the Router executor built from the router kwarg --- .../base_llm/vector_store/transformation.py | 16 ++++++- litellm/router.py | 11 +---- litellm/vector_stores/main.py | 13 ++++-- .../test_vector_store_endpoints.py | 2 +- .../test_milvus_vector_store.py | 46 +++++++++++++++++++ 5 files changed, 73 insertions(+), 15 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 3547bf81448..9624a721870 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -62,6 +62,19 @@ class LiteLLMVectorStoreEmbeddingExecutor: ) +_REQUEST_METADATA: Final = TypeAdapter(dict[str, object]) + + +def vector_store_request_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: + litellm_metadata: Final = kwargs.get("litellm_metadata") + if isinstance(litellm_metadata, dict): + return _REQUEST_METADATA.validate_python(litellm_metadata) + metadata: Final = kwargs.get("metadata") + if isinstance(metadata, dict): + return _REQUEST_METADATA.validate_python(metadata) + return MappingProxyType({}) + + @dataclass(frozen=True, slots=True) class RouterVectorStoreEmbeddingExecutor: router: Router @@ -312,11 +325,12 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): def query_embedding_executor( embedding_executor: VectorStoreEmbeddingExecutor | None, router: Router | None, + request_metadata: Mapping[str, object] = MappingProxyType({}), ) -> VectorStoreEmbeddingExecutor: if embedding_executor is not None: return embedding_executor if router is not None: - return RouterVectorStoreEmbeddingExecutor(router=router, metadata=MappingProxyType({})) + return RouterVectorStoreEmbeddingExecutor(router=router, metadata=request_metadata) return LiteLLMVectorStoreEmbeddingExecutor() def embed_query( diff --git a/litellm/router.py b/litellm/router.py index 991f1bc2828..8b510e748f5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -87,6 +87,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( ) from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, + vector_store_request_metadata, ) from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.router_strategy.budget_limiter import RouterBudgetLimiting @@ -6697,15 +6698,7 @@ class Router: @staticmethod def _vector_store_request_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: - litellm_metadata: Final = kwargs.get("litellm_metadata") - if isinstance(litellm_metadata, dict): - return cast( # cast-ok: isinstance validates the runtime dict boundary - "dict[str, object]", litellm_metadata - ) - metadata: Final = kwargs.get("metadata") - if isinstance(metadata, dict): - return cast("dict[str, object]", metadata) # cast-ok: isinstance validates the runtime dict boundary - return MappingProxyType({}) + return vector_store_request_metadata(kwargs) async def _init_vector_store_api_endpoints( self, diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index d530a1c9d52..636bdd4b52e 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, + vector_store_request_metadata, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.router import GenericLiteLLMParams @@ -42,10 +43,14 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# -def _direct_vector_store_embedding_executor(value: object, router: "Router | None") -> VectorStoreEmbeddingExecutor: +def _direct_vector_store_embedding_executor( + value: object, router: "Router | None", request_kwargs: Mapping[str, object] +) -> VectorStoreEmbeddingExecutor: if value is not None and not isinstance(value, VectorStoreEmbeddingExecutor): raise TypeError("Invalid direct vector store embedding executor") - return BaseQueryEmbeddingVectorStoreConfig.query_embedding_executor(value, router) + return BaseQueryEmbeddingVectorStoreConfig.query_embedding_executor( + value, router, vector_store_request_metadata(request_kwargs) + ) def mock_vector_store_search_response( @@ -300,7 +305,7 @@ async def asearch( Async: Search a vector store for relevant chunks based on a query and file attributes filter. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None), router + kwargs.pop("_direct_vector_store_embedding_executor", None), router, kwargs ) local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot key: value for key, value in locals().items() if key != "embedding_executor" @@ -386,7 +391,7 @@ def search( VectorStoreSearchResponse containing the search results. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None), router + kwargs.pop("_direct_vector_store_embedding_executor", None), router, kwargs ) local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot key: value for key, value in locals().items() if key != "embedding_executor" diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index a956aef51d6..1abbbe91e97 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -52,7 +52,7 @@ def _serialize_litellm_params(litellm_params): def test_direct_vector_store_embedding_executor_rejects_invalid_value(): with pytest.raises(TypeError, match="Invalid direct vector store embedding executor"): - _direct_vector_store_embedding_executor(object(), None) + _direct_vector_store_embedding_executor(object(), None, {}) def test_router_vector_store_search_injects_executor_and_request_metadata(): diff --git a/tests/vector_store_tests/test_milvus_vector_store.py b/tests/vector_store_tests/test_milvus_vector_store.py index 3b3bd444b20..2ba9168b49f 100644 --- a/tests/vector_store_tests/test_milvus_vector_store.py +++ b/tests/vector_store_tests/test_milvus_vector_store.py @@ -601,6 +601,52 @@ async def test_sdk_search_with_router_kwarg_resolves_bare_embedding_alias_async( _assert_alias_resolved(embedding_route, search_route, response) +def _team_alias_router(): + return Router( + model_list=[ + { + "model_name": "team-a-embedder", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + "model_info": {"team_id": "team-a", "team_public_model_name": "multilingual-e5-large"}, + } + ] + ) + + +@pytest.mark.asyncio +async def test_sdk_search_with_router_kwarg_resolves_team_alias_from_request_metadata( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = await litellm.vector_stores.asearch( + router=_team_alias_router(), metadata={"user_api_key_team_id": "team-a"}, **ALIAS_SEARCH_KWARGS + ) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_sdk_search_with_router_kwarg_rejects_team_alias_without_team_metadata( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + _mock_search_route(respx_mock) + + with pytest.raises(litellm.APIConnectionError): + await litellm.vector_stores.asearch(router=_team_alias_router(), **ALIAS_SEARCH_KWARGS) + + assert embedding_route.call_count == 0 + + @pytest.mark.asyncio async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE) From b0fe71010b89425501587a3e7e66ede9a4e6e775 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 16:56:35 -0700 Subject: [PATCH 67/69] fix(ollama_chat): stamp finish_reason tool_calls when tool calls streamed before the done chunk (#39010) * fix(ollama_chat): stamp finish_reason tool_calls when tool calls streamed before the done chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ollama_chat): trim finish_reason override comment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mubashir1osmani --- litellm/llms/ollama/chat/transformation.py | 7 +- .../test_streaming_iterator_stop_reason.py | 79 +++++++++++++++++++ .../ollama/test_ollama_chat_transformation.py | 40 ++++++++++ 3 files changed, 124 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index de626b468f0..181894646e3 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -423,6 +423,7 @@ class OllamaChatConfig(BaseConfig): class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): started_reasoning_content: bool = False finished_reasoning_content: bool = False + seen_tool_calls: bool = False def _is_function_call_complete(self, function_args: str | dict) -> bool: if isinstance(function_args, dict): @@ -468,6 +469,7 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): # process tool calls - if complete function arg - add id to tool call tool_calls: Final = chunk["message"].get("tool_calls") if tool_calls is not None: + self.seen_tool_calls = True for tool_call in tool_calls: function_args = tool_call.get("function").get("arguments") if function_args is not None and len(function_args) > 0: @@ -508,9 +510,10 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): if chunk["done"] is True: finish_reason = chunk.get("done_reason") or "stop" - # Override finish_reason when tool_calls are present + # Override finish_reason when tool_calls appeared in any chunk # Fixes: https://github.com/BerriAI/litellm/issues/18922 - if tool_calls is not None: + # Fixes: https://github.com/BerriAI/litellm/issues/34692 + if self.seen_tool_calls: finish_reason = "tool_calls" choices = [ StreamingChoices( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py new file mode 100644 index 00000000000..4b95b36fec3 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py @@ -0,0 +1,79 @@ +""" +Regression tests for issue #34692. + +ollama_chat streams tool_calls in a mid-stream chunk while its final +(``done: true``) chunk carries only ``done_reason: "stop"``. The provider +iterator must remember the earlier tool_calls and stamp +``finish_reason="tool_calls"`` on the final chunk, so the Anthropic +``/v1/messages`` bridge emits ``stop_reason: "tool_use"``. Before the fix the +bridge emitted ``stop_reason: "end_turn"`` and Anthropic tool-runners +(Claude Code, ``messages.stream``) silently dropped the tool call. +""" + +import pytest + +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, +) +from litellm.llms.ollama.chat.transformation import ( + OllamaChatCompletionResponseIterator, +) +from litellm.types.utils import ModelResponseStream + +_OLLAMA_TOOL_CHUNK = { + "model": "qwen3:8b", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "get_weather", "arguments": {"city": "San Francisco"}}}], + }, + "done": False, +} +_OLLAMA_DONE_CHUNK = { + "model": "qwen3:8b", + "message": {"role": "assistant", "content": ""}, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 100, + "eval_count": 20, +} + + +def _ollama_streamed_chunks() -> list[ModelResponseStream]: + iterator = OllamaChatCompletionResponseIterator(streaming_response=iter([]), sync_stream=True) + return [iterator.chunk_parser(_OLLAMA_TOOL_CHUNK), iterator.chunk_parser(_OLLAMA_DONE_CHUNK)] + + +class _AsyncStream: + def __init__(self, items: list[ModelResponseStream]): + self._it = iter(items) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._it) + except StopIteration: + raise StopAsyncIteration + + +def _assert_tool_use_stop_reason(events: list[dict]) -> None: + block_types = [e["content_block"]["type"] for e in events if e.get("type") == "content_block_start"] + assert "tool_use" in block_types, f"no tool_use content block opened: {events}" + message_deltas = [e for e in events if e.get("type") == "message_delta"] + assert message_deltas, f"no message_delta emitted: {events}" + assert message_deltas[-1]["delta"]["stop_reason"] == "tool_use", ( + f"expected stop_reason 'tool_use', got: {message_deltas[-1]}" + ) + + +def test_ollama_mid_stream_tool_call_yields_tool_use_stop_reason_sync(): + wrapper = AnthropicStreamWrapper(completion_stream=iter(_ollama_streamed_chunks()), model="qwen3:8b") + _assert_tool_use_stop_reason(list(wrapper)) + + +@pytest.mark.asyncio +async def test_ollama_mid_stream_tool_call_yields_tool_use_stop_reason_async(): + wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(_ollama_streamed_chunks()), model="qwen3:8b") + _assert_tool_use_stop_reason([event async for event in wrapper]) diff --git a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py index 8f3dbf7b0d9..25f9645faa0 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py @@ -615,6 +615,46 @@ class TestOllamaFinishReasonLength: result.choices[0].finish_reason == "stop" ), f"Expected 'stop' for natural finish, got '{result.choices[0].finish_reason}'" + def test_finish_reason_tool_calls_streamed_before_done_chunk(self): + """Streaming: tool_calls arriving mid-stream (not on the done chunk) must + still produce finish_reason='tool_calls' on the final chunk. + + Regression test for https://github.com/BerriAI/litellm/issues/34692: + Ollama emits tool_calls in an earlier chunk and the done chunk carries + none, which left finish_reason at 'stop' and made the Anthropic + /v1/messages bridge emit stop_reason 'end_turn' instead of 'tool_use'. + """ + iterator = OllamaChatCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + ) + + tool_chunk = { + "model": "qwen3:8b", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + {"function": {"name": "get_weather", "arguments": {"city": "San Francisco"}}} + ], + }, + "done": False, + } + done_chunk = { + "model": "qwen3:8b", + "message": {"role": "assistant", "content": ""}, + "done": True, + "done_reason": "stop", + } + + tool_result = iterator.chunk_parser(tool_chunk) + assert tool_result.choices[0].delta.tool_calls is not None + + done_result = iterator.chunk_parser(done_chunk) + assert ( + done_result.choices[0].finish_reason == "tool_calls" + ), f"Expected 'tool_calls' when tool_calls were streamed earlier, got '{done_result.choices[0].finish_reason}'" + class TestOllamaReasoningContentStreaming: """Test that reasoning_content is properly extracted from all thinking chunks.""" From e641864e6f863ef03cd8724b1cb4d69097d447b0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:58:53 -0700 Subject: [PATCH 68/69] test(e2e/ui): give the seeded users passwords that pass the default password policy The default password policy from #39381 (12+ chars with upper, lower, number, and special) makes globalSetup's POST /user/update with password "test" fail with a 400, so every role login in the UI e2e suite has been failing since that merge. Seed the roles, the noteam@test.local login, and the second-admin spec with policy-compliant passwords and update the scrypt hash in seed.sql to match. --- tests/e2e/ui/constants.ts | 1 + tests/e2e/ui/fixtures/seed.sql | 20 +++++++++---------- tests/e2e/ui/fixtures/users.ts | 9 +++++---- .../internal-user/internalUserNoTeam.spec.ts | 3 ++- .../ui/tests/proxy-admin/secondAdmin.spec.ts | 2 +- 5 files changed, 19 insertions(+), 16 deletions(-) diff --git a/tests/e2e/ui/constants.ts b/tests/e2e/ui/constants.ts index bb33c90ddf3..3a62252915c 100644 --- a/tests/e2e/ui/constants.ts +++ b/tests/e2e/ui/constants.ts @@ -30,6 +30,7 @@ export const E2E_PROXY_ADMIN_EMAIL = "admin@test.local"; export const E2E_INTERNAL_USER_ID = "e2e-internal-user"; export const E2E_INTERNAL_USER_EMAIL = "internal@test.local"; export const E2E_TEAM_ADMIN_USER_ID = "e2e-team-admin"; +export const E2E_SEEDED_USER_PASSWORD = "E2e-Test-Pass-2026!"; // Key aliases for seeded test keys (match seed.sql) export const E2E_UPDATE_LIMITS_KEY_ALIAS = "e2eUpdateLimitsKey"; diff --git a/tests/e2e/ui/fixtures/seed.sql b/tests/e2e/ui/fixtures/seed.sql index e77b4a16b3d..00ea668ed8f 100644 --- a/tests/e2e/ui/fixtures/seed.sql +++ b/tests/e2e/ui/fixtures/seed.sql @@ -24,18 +24,18 @@ INSERT INTO "LiteLLM_OrganizationTable" ( 'e2e-proxy-admin', 'e2e-proxy-admin' ); --- 4. Users (password hash is scrypt of "test") +-- 4. Users (password hash is scrypt of E2E_SEEDED_USER_PASSWORD from constants.ts) INSERT INTO "LiteLLM_UserTable" ("user_id", "user_email", "user_role", "teams", "password") VALUES - ('e2e-proxy-admin', 'admin@test.local', 'proxy_admin', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-admin-viewer', 'adminviewer@test.local', 'proxy_admin_viewer', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-internal-user', 'internal@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-org","e2e-team-keygen"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-internal-viewer', 'viewer@test.local', 'internal_user_viewer', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-team-admin', 'teamadmin@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-delete"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-invitable-user', 'invitable@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-internal-noteam', 'noteam@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-invitable-by-team-admin', 'invitable-team@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-removable-member', 'removable@test.local', 'internal_user', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'); + ('e2e-proxy-admin', 'admin@test.local', 'proxy_admin', '{"e2e-team-crud"}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-admin-viewer', 'adminviewer@test.local', 'proxy_admin_viewer', '{}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-internal-user', 'internal@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-org","e2e-team-keygen"}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-internal-viewer', 'viewer@test.local', 'internal_user_viewer', '{"e2e-team-crud"}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-team-admin', 'teamadmin@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-delete"}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-invitable-user', 'invitable@test.local', 'internal_user', '{}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-internal-noteam', 'noteam@test.local', 'internal_user', '{}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-invitable-by-team-admin', 'invitable-team@test.local', 'internal_user', '{}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-removable-member', 'removable@test.local', 'internal_user', '{"e2e-team-crud"}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'); -- 5. Teams (members_with_roles is required JSON) INSERT INTO "LiteLLM_TeamTable" ( diff --git a/tests/e2e/ui/fixtures/users.ts b/tests/e2e/ui/fixtures/users.ts index 79ee237f334..0457361d9d4 100644 --- a/tests/e2e/ui/fixtures/users.ts +++ b/tests/e2e/ui/fixtures/users.ts @@ -1,6 +1,7 @@ import { ADMIN_STORAGE_PATH, ADMIN_VIEWER_STORAGE_PATH, + E2E_SEEDED_USER_PASSWORD, INTERNAL_USER_STORAGE_PATH, INTERNAL_VIEWER_STORAGE_PATH, TEAM_ADMIN_STORAGE_PATH, @@ -23,22 +24,22 @@ export const users: Record { // Log in via the form as the no-team seeded user. await page.goto("/ui/login"); await page.getByPlaceholder("Enter your username").fill("noteam@test.local"); - await page.getByPlaceholder("Enter your password").fill("test"); + await page.getByPlaceholder("Enter your password").fill(E2E_SEEDED_USER_PASSWORD); await page.getByRole("button", { name: "Login", exact: true }).click(); await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 30_000 }); expect(new URL(page.url()).pathname).not.toMatch(/\/connect$/); diff --git a/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts index 5a8bc84cc13..73263c844fa 100644 --- a/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts @@ -10,7 +10,7 @@ test.describe("Second proxy admin", () => { test("an invited admin can log in, mint a key, and call a model with it", async ({ page, browser, request }) => { const suffix = Date.now(); const email = `second-admin-${suffix}@test.local`; - const password = "e2e-second-admin-password"; + const password = "E2e-Second-Admin-Pass-1!"; const auth = { Authorization: `Bearer ${masterKey()}` }; const inviteAdminUser = async (): Promise => { From 5ad330f6203af47ced983a6fc3852288d0a5cdc1 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 17:30:46 -0700 Subject: [PATCH 69/69] fix(http_handler): honor HTTP(S)_PROXY / NO_PROXY when force_ipv4 uses the httpx transport (#39443) * fix(http_handler): honor HTTP(S)_PROXY / NO_PROXY when force_ipv4 uses the httpx transport Passing an explicit transport to httpx.AsyncClient / httpx.Client disables its automatic environment proxy mounts, so force_ipv4 on the httpx path sent every LLM request direct and silently bypassed HTTPS_PROXY. Mount the same env-derived proxy transports next to the IPv4-pinned direct transport in AsyncHTTPHandler, HTTPHandler and the OpenAI async client factory. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(http_handler): carry the client's TLS verify and cert settings onto env proxy mounts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/http_handler.py | 49 +++- litellm/llms/openai/common_utils.py | 12 +- .../llms/custom_httpx/test_http_handler.py | 233 ++++++++++++++++++ 3 files changed, 285 insertions(+), 9 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index b6e93f590ca..e1f0fc9e7d3 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -9,13 +9,15 @@ import threading import time from collections.abc import AsyncIterable, Callable, Iterable, Mapping from http.cookiejar import CookieJar, DefaultCookiePolicy -from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict, TypeVar import certifi import httpx from aiohttp import ClientSession, DummyCookieJar, TCPConnector from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport -from httpx._types import RequestFiles +from httpx._types import CertTypes, RequestFiles +from httpx._utils import get_environment_proxies import litellm from litellm._logging import verbose_logger @@ -66,6 +68,22 @@ _AddrInfo: TypeAlias = tuple[int | socket.AddressFamily, int | socket.SocketKind _RequestContent: TypeAlias = str | bytes | Iterable[bytes] | AsyncIterable[bytes] +_IPV4_LOCAL_ADDRESS: Final = "0.0.0.0" + +_HttpxTransportT = TypeVar("_HttpxTransportT", HTTPTransport, AsyncHTTPTransport) + + +def _environment_proxy_mounts( + build_proxy_transport: Callable[[str], _HttpxTransportT], +) -> Mapping[str, _HttpxTransportT | None]: + """httpx skips its own HTTP(S)_PROXY / NO_PROXY mounts whenever an explicit `transport=` is passed.""" + return MappingProxyType( + { + pattern: None if proxy_url is None else build_proxy_transport(proxy_url) + for pattern, proxy_url in get_environment_proxies().items() + } + ) + class _TCPConnectorKwargs(TypedDict, total=False): local_addr: tuple[str, int] | None @@ -607,6 +625,7 @@ class AsyncHTTPHandler: return httpx.AsyncClient( transport=transport, + mounts=AsyncHTTPHandler._create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert), event_hooks=event_hooks, timeout=timeout, verify=ssl_config, @@ -1191,10 +1210,22 @@ class AsyncHTTPHandler: - [Default] If force_ipv4 is False, it will return None """ if litellm.force_ipv4: - return AsyncHTTPTransport(local_address="0.0.0.0") + return AsyncHTTPTransport(local_address=_IPV4_LOCAL_ADDRESS) else: return None + @staticmethod + def _create_httpx_proxy_mounts( + transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | None, + verify: VerifyTypes, + cert: CertTypes | None, + ) -> Mapping[str, AsyncHTTPTransport | None] | None: + if not isinstance(transport, AsyncHTTPTransport): + return None + return _environment_proxy_mounts( + lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert) + ) + class HTTPHandler: def __init__( @@ -1227,6 +1258,7 @@ class HTTPHandler: # Create a client with a connection pool return httpx.Client( transport=self._create_sync_transport(), + mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert), timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT, verify=ssl_config, cert=cert, @@ -1507,10 +1539,19 @@ class HTTPHandler: Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them """ if litellm.force_ipv4: - return HTTPTransport(local_address="0.0.0.0") + return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS) else: return getattr(litellm, "sync_transport", None) + @staticmethod + def _create_sync_proxy_mounts( + verify: VerifyTypes, + cert: CertTypes | None, + ) -> Mapping[str, HTTPTransport | None] | None: + if not litellm.force_ipv4: + return None + return _environment_proxy_mounts(lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert)) + def get_async_httpx_client( llm_provider: LlmProviders | httpxSpecialProvider, diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 4d774f6f165..bcd4ea43243 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -305,14 +305,16 @@ class BaseOpenAILLM: # Get unified SSL configuration ssl_config: Final = get_ssl_configuration() + transport: Final = AsyncHTTPHandler._create_async_transport( + ssl_context=(ssl_config if isinstance(ssl_config, ssl.SSLContext) else None), + ssl_verify=ssl_config if isinstance(ssl_config, bool) else None, + shared_session=shared_session, + ) return httpx.AsyncClient( verify=ssl_config, - transport=AsyncHTTPHandler._create_async_transport( - ssl_context=(ssl_config if isinstance(ssl_config, ssl.SSLContext) else None), - ssl_verify=ssl_config if isinstance(ssl_config, bool) else None, - shared_session=shared_session, - ), + transport=transport, + mounts=AsyncHTTPHandler._create_httpx_proxy_mounts(transport, verify=ssl_config, cert=None), follow_redirects=True, ) diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 16d57437043..9e64bfafa54 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1314,3 +1314,236 @@ async def test_finalizer_on_live_loop_disposes_foreign_loop_session_without_sche assert AsyncHTTPHandler._finalizer_close_tasks == baseline_tasks assert session.closed + + +@pytest.fixture +def forward_proxy_server(): + """Plain HTTP forward proxy that records the absolute URIs it is asked to fetch.""" + from http.server import BaseHTTPRequestHandler, HTTPServer + from socketserver import ThreadingMixIn + + seen_uris: list[str] = [] + + class RecordingProxyHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): + seen_uris.append(self.path) + self.send_response(200) + self.send_header("Content-Length", "9") + self.end_headers() + self.wfile.write(b"via-proxy") + + def log_message(self, format, *args): + pass + + class ThreadedServer(ThreadingMixIn, HTTPServer): + daemon_threads = True + + server = ThreadedServer(("127.0.0.1", 0), RecordingProxyHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}", seen_uris + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +# `.invalid` never resolves (RFC 6761), so the only way this request can succeed is through the proxy +_PROXY_ONLY_UPSTREAM_URL = "http://upstream.invalid/v1/models" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("disable_aiohttp_transport", [True, False]) +@pytest.mark.parametrize("force_ipv4", [True, False]) +async def test_async_handler_honours_proxy_env_for_every_transport( + forward_proxy_server, monkeypatch: pytest.MonkeyPatch, disable_aiohttp_transport: bool, force_ipv4: bool +): + proxy_url, seen_uris = forward_proxy_server + monkeypatch.setenv("HTTP_PROXY", proxy_url) + monkeypatch.delenv("NO_PROXY", raising=False) + monkeypatch.delenv("no_proxy", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", disable_aiohttp_transport) + monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) + + handler = AsyncHTTPHandler() + try: + response = await handler.get(_PROXY_ONLY_UPSTREAM_URL) + finally: + await handler.close() + + assert response.text == "via-proxy" + assert seen_uris == [_PROXY_ONLY_UPSTREAM_URL] + + +@pytest.mark.parametrize("force_ipv4", [True, False]) +def test_sync_handler_honours_proxy_env(forward_proxy_server, monkeypatch: pytest.MonkeyPatch, force_ipv4: bool): + proxy_url, seen_uris = forward_proxy_server + monkeypatch.setenv("HTTP_PROXY", proxy_url) + monkeypatch.delenv("NO_PROXY", raising=False) + monkeypatch.delenv("no_proxy", raising=False) + monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) + + handler = HTTPHandler() + try: + response = handler.get(_PROXY_ONLY_UPSTREAM_URL) + finally: + handler.close() + + assert response.text == "via-proxy" + assert seen_uris == [_PROXY_ONLY_UPSTREAM_URL] + + +@pytest.mark.asyncio +async def test_force_ipv4_httpx_transport_honours_no_proxy(keepalive_server, monkeypatch: pytest.MonkeyPatch): + """NO_PROXY hosts must still go direct when the proxy mounts are supplied by litellm instead of httpx.""" + monkeypatch.setenv("HTTP_PROXY", "http://proxy.invalid:3128") + monkeypatch.setenv("NO_PROXY", "127.0.0.1") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "force_ipv4", True) + + handler = AsyncHTTPHandler() + try: + response = await handler.get(keepalive_server) + finally: + await handler.close() + + assert response.text == "ok" + + +@pytest.fixture +def private_ca_tls_upstream(tmp_path: pathlib.Path): + """HTTPS server behind a CONNECT proxy, both on localhost; the server's cert is signed by a test-only CA.""" + import datetime + import select + import socket + from http.server import BaseHTTPRequestHandler, HTTPServer + from socketserver import ThreadingMixIn + + from cryptography import x509 + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.x509.oid import NameOID + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "upstream.invalid")]) + now = datetime.datetime.now(datetime.timezone.utc) + cert = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(minutes=1)) + .not_valid_after(now + datetime.timedelta(hours=1)) + .add_extension(x509.SubjectAlternativeName([x509.DNSName("upstream.invalid")]), critical=False) + .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True) + .sign(key, hashes.SHA256()) + ) + ca_pem = tmp_path / "ca.pem" + ca_pem.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) + key_pem = tmp_path / "key.pem" + key_pem.write_bytes( + key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() + ) + ) + + class OkTlsHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): + self.send_response(200) + self.send_header("Content-Length", "6") + self.end_headers() + self.wfile.write(b"ok-tls") + + def log_message(self, format, *args): + pass + + class ThreadedServer(ThreadingMixIn, HTTPServer): + daemon_threads = True + + tls_server = ThreadedServer(("127.0.0.1", 0), OkTlsHandler) + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(str(ca_pem), str(key_pem)) + tls_server.socket = server_ctx.wrap_socket(tls_server.socket, server_side=True) + tls_port = tls_server.server_port + + class ConnectProxyHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_CONNECT(self): + upstream = socket.create_connection(("127.0.0.1", tls_port)) + self.send_response(200, "Connection established") + self.end_headers() + sockets = [self.connection, upstream] + while True: + readable, _, _ = select.select(sockets, [], [], 5) + if not readable: + break + for src in readable: + data = src.recv(65536) + if not data: + upstream.close() + return + (upstream if src is self.connection else self.connection).sendall(data) + + def log_message(self, format, *args): + pass + + proxy_server = ThreadedServer(("127.0.0.1", 0), ConnectProxyHandler) + threads = [ + threading.Thread(target=tls_server.serve_forever, daemon=True), + threading.Thread(target=proxy_server.serve_forever, daemon=True), + ] + for thread in threads: + thread.start() + try: + yield f"http://127.0.0.1:{proxy_server.server_port}", str(ca_pem) + finally: + for server in (proxy_server, tls_server): + server.shutdown() + server.server_close() + for thread in threads: + thread.join(timeout=5) + + +@pytest.mark.asyncio +async def test_force_ipv4_https_proxy_mount_uses_handler_ca_bundle( + private_ca_tls_upstream, monkeypatch: pytest.MonkeyPatch +): + proxy_url, ca_pem = private_ca_tls_upstream + monkeypatch.setenv("HTTPS_PROXY", proxy_url) + monkeypatch.delenv("NO_PROXY", raising=False) + monkeypatch.delenv("no_proxy", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "force_ipv4", True) + + handler = AsyncHTTPHandler(ssl_verify=ca_pem) + try: + response = await handler.get("https://upstream.invalid/v1/models") + finally: + await handler.close() + + assert response.text == "ok-tls" + + +def test_sync_force_ipv4_https_proxy_mount_uses_handler_ca_bundle( + private_ca_tls_upstream, monkeypatch: pytest.MonkeyPatch +): + proxy_url, ca_pem = private_ca_tls_upstream + monkeypatch.setenv("HTTPS_PROXY", proxy_url) + monkeypatch.delenv("NO_PROXY", raising=False) + monkeypatch.delenv("no_proxy", raising=False) + monkeypatch.setattr(litellm, "force_ipv4", True) + + handler = HTTPHandler(ssl_verify=ca_pem) + try: + response = handler.get("https://upstream.invalid/v1/models") + finally: + handler.close() + + assert response.text == "ok-tls"