From 03a4e8bfb57ec69cc184b5114b0f66f1480672c6 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:09:23 +0000 Subject: [PATCH 001/318] fix(azure/realtime): authenticate realtime websocket with Azure AD token when no api-key --- litellm/llms/azure/realtime/handler.py | 21 ++- litellm/realtime_api/main.py | 15 +- .../realtime/test_azure_realtime_handler.py | 178 ++++++++++++++++++ 3 files changed, 207 insertions(+), 7 deletions(-) diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 86c1ed51b68..51f9ef5989c 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -30,6 +30,21 @@ async def forward_messages(client_ws: Any, backend_ws: Any): class AzureOpenAIRealtime(AzureChatCompletion): + @staticmethod + def get_auth_headers(api_key: str | None, azure_ad_token: str | None) -> dict[str, str]: + """ + Build the websocket handshake auth headers, preferring a static api-key and falling back to + an Azure AD (Entra ID) bearer token. Never sends both. + """ + if api_key: + return {"api-key": api_key} + if azure_ad_token: + return {"Authorization": f"Bearer {azure_ad_token}"} + raise ValueError( + "Missing Azure credentials for the realtime endpoint. Set an api_key, or configure Azure AD auth " + "(azure_ad_token, tenant_id/client_id/client_secret, or a managed identity)" + ) + def _construct_url( self, api_base: str, @@ -117,13 +132,13 @@ class AzureOpenAIRealtime(AzureChatCompletion): query_params=query_params, ) + auth_headers = self.get_auth_headers(api_key=api_key, azure_ad_token=azure_ad_token) + try: ssl_context = get_shared_realtime_ssl_context() async with websockets.connect( # type: ignore url, - additional_headers={ - "api-key": api_key, # type: ignore - }, + additional_headers=auth_headers, max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, ssl=ssl_context, ) as backend_ws: diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 5ecf4d91ff6..e9175917f9a 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -23,6 +23,7 @@ from litellm.utils import ProviderConfigManager from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from ..llms.azure.common_utils import get_azure_ad_token from ..llms.azure.realtime.handler import AzureOpenAIRealtime from ..llms.bedrock.realtime.handler import BedrockRealtime from ..llms.custom_httpx.http_handler import get_shared_realtime_ssl_context @@ -376,7 +377,7 @@ async def _arealtime( api_base=api_base, api_key=api_key, api_version=api_version, - azure_ad_token=None, + azure_ad_token=(None if api_key else get_azure_ad_token(litellm_params)), client=None, timeout=timeout, logging_obj=litellm_logging_obj, @@ -536,6 +537,7 @@ async def _realtime_health_check( import websockets url: Optional[str] = None + auth_headers: dict[str, str | None] = {"api-key": api_key} if custom_llm_provider == "azure": url = azure_realtime._construct_url( api_base=api_base or "", @@ -543,6 +545,13 @@ async def _realtime_health_check( api_version=api_version or "2024-10-01-preview", realtime_protocol=realtime_protocol, ) + azure_litellm_params = GenericLiteLLMParams(**(model_params or {})) + auth_headers = dict( + azure_realtime.get_auth_headers( + api_key=api_key, + azure_ad_token=(None if api_key else get_azure_ad_token(azure_litellm_params)), + ) + ) elif custom_llm_provider == "openai": url = openai_realtime._construct_url( api_base=api_base or "https://api.openai.com/", @@ -584,9 +593,7 @@ async def _realtime_health_check( ssl_context = get_shared_realtime_ssl_context() async with websockets.connect( # type: ignore url, - additional_headers={ - "api-key": api_key, # type: ignore - }, + additional_headers=auth_headers, max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, ssl=ssl_context, ): diff --git a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py b/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py index 4638bc4df0f..d9c49947f19 100644 --- a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py @@ -563,3 +563,181 @@ async def test_async_realtime_default_maintains_backwards_compatibility(): mock_realtime_streaming.call_args.kwargs["backend_uses_beta_protocol"] is True ) + + +class _DummyAsyncContextManager: + def __init__(self, value): + self.value = value + + async def __aenter__(self): + return self.value + + async def __aexit__(self, exc_type, exc, tb): + return None + + +@pytest.mark.asyncio +async def test_async_realtime_uses_bearer_token_when_no_api_key(): + """ + Entra ID-only Azure realtime deployments have no static api-key, so the handshake must + authenticate with `Authorization: Bearer ` and must not send `api-key`. + + Regression test for https://github.com/BerriAI/litellm/issues/34654 + """ + from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime + + handler = AzureOpenAIRealtime() + mock_backend_ws = AsyncMock() + + with ( + patch( + "websockets.connect", + return_value=_DummyAsyncContextManager(mock_backend_ws), + ) as mock_ws_connect, + patch("litellm.llms.azure.realtime.handler.RealTimeStreaming") as mock_realtime_streaming, + ): + mock_realtime_streaming.return_value.bidirectional_forward = AsyncMock() + + await handler.async_realtime( + model="gpt-realtime-whisper", + websocket=AsyncMock(), + logging_obj=MagicMock(), + api_base="https://my-endpoint.openai.azure.com", + api_key=None, + api_version="2024-10-01-preview", + azure_ad_token="my-entra-token", + ) + + headers = mock_ws_connect.call_args.kwargs["additional_headers"] + assert headers == {"Authorization": "Bearer my-entra-token"} + + +def test_get_auth_headers_prefers_api_key_and_never_sends_both(): + from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime + + assert AzureOpenAIRealtime.get_auth_headers(api_key="test-key", azure_ad_token="my-entra-token") == { + "api-key": "test-key" + } + + +def test_get_auth_headers_without_credentials_raises(): + from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime + + with pytest.raises(ValueError, match="Missing Azure credentials"): + AzureOpenAIRealtime.get_auth_headers(api_key=None, azure_ad_token=None) + + +@pytest.mark.asyncio +async def test_arealtime_resolves_azure_ad_token_when_no_api_key(monkeypatch): + """ + `_arealtime` must resolve an Azure AD token (managed identity, service principal, etc.) + and forward it to the handler when the deployment has no api_key. + + Regression test for https://github.com/BerriAI/litellm/issues/34654 + """ + from litellm.realtime_api import main as realtime_main + + mock_async_realtime = AsyncMock() + monkeypatch.setattr(realtime_main, "azure_realtime", MagicMock(async_realtime=mock_async_realtime)) + monkeypatch.setattr( + realtime_main, + "get_llm_provider", + lambda model, api_base=None, api_key=None: ( + "gpt-realtime-whisper", + "azure", + None, + "https://my-endpoint.openai.azure.com", + ), + ) + monkeypatch.delenv("AZURE_API_KEY", raising=False) + + captured_params = {} + + def fake_get_azure_ad_token(litellm_params): + captured_params["tenant_id"] = litellm_params.get("tenant_id") + return "my-entra-token" + + monkeypatch.setattr(realtime_main, "get_azure_ad_token", fake_get_azure_ad_token) + + await realtime_main._arealtime( + model="azure/gpt-realtime-whisper", + websocket=MagicMock(), + api_version="2024-10-01-preview", + litellm_logging_obj=MagicMock(), + tenant_id="my-tenant", + client_id="my-client", + client_secret="my-secret", + ) + + assert mock_async_realtime.call_args.kwargs["azure_ad_token"] == "my-entra-token" + assert captured_params["tenant_id"] == "my-tenant" + + +@pytest.mark.asyncio +async def test_arealtime_does_not_resolve_azure_ad_token_when_api_key_present(monkeypatch): + from litellm.realtime_api import main as realtime_main + + mock_async_realtime = AsyncMock() + monkeypatch.setattr(realtime_main, "azure_realtime", MagicMock(async_realtime=mock_async_realtime)) + monkeypatch.setattr( + realtime_main, + "get_llm_provider", + lambda model, api_base=None, api_key=None: ( + "gpt-realtime-whisper", + "azure", + "test-key", + "https://my-endpoint.openai.azure.com", + ), + ) + + def fail_get_azure_ad_token(litellm_params): + raise AssertionError("should not resolve an AD token when an api_key is configured") + + monkeypatch.setattr(realtime_main, "get_azure_ad_token", fail_get_azure_ad_token) + + await realtime_main._arealtime( + model="azure/gpt-realtime-whisper", + websocket=MagicMock(), + api_key="test-key", + api_version="2024-10-01-preview", + litellm_logging_obj=MagicMock(), + ) + + assert mock_async_realtime.call_args.kwargs["azure_ad_token"] is None + + +@pytest.mark.asyncio +async def test_realtime_health_check_uses_bearer_token_when_no_api_key(monkeypatch): + """ + An Entra ID-only realtime deployment must also pass its realtime health check. + + Regression test for https://github.com/BerriAI/litellm/issues/34654 + """ + from litellm.realtime_api import main as realtime_main + + connect_calls = [] + + monkeypatch.setattr( + realtime_main, + "get_azure_ad_token", + lambda litellm_params: "my-entra-token", + ) + + def fake_connect(url, **kwargs): + connect_calls.append(kwargs) + return _DummyAsyncContextManager(MagicMock()) + + monkeypatch.setattr("websockets.connect", fake_connect) + + assert ( + await realtime_main._realtime_health_check( + model="gpt-realtime-whisper", + custom_llm_provider="azure", + api_key=None, + api_base="https://my-endpoint.openai.azure.com", + api_version="2024-10-01-preview", + model_params={"tenant_id": "my-tenant"}, + ) + is True + ) + assert connect_calls[0]["additional_headers"] == {"Authorization": "Bearer my-entra-token"} From b930169f39ff81b34b9088d32596f2fc07241839 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:21:30 +0000 Subject: [PATCH 002/318] fix(azure/realtime): resolve AD token from deployment azure_ad_token param and kwargs --- litellm/realtime_api/main.py | 7 +++- .../realtime/test_azure_realtime_handler.py | 36 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index e9175917f9a..e981db216af 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -371,13 +371,18 @@ async def _arealtime( if realtime_protocol is None and (query_params or {}).get("intent") == "transcription": realtime_protocol = "GA" realtime_protocol = realtime_protocol or "beta" + resolved_azure_ad_token = ( + None + if api_key + else get_azure_ad_token(GenericLiteLLMParams(**{**kwargs, "azure_ad_token": azure_ad_token})) + ) await azure_realtime.async_realtime( model=model, websocket=websocket, api_base=api_base, api_key=api_key, api_version=api_version, - azure_ad_token=(None if api_key else get_azure_ad_token(litellm_params)), + azure_ad_token=resolved_azure_ad_token, client=None, timeout=timeout, logging_obj=litellm_logging_obj, diff --git a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py b/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py index d9c49947f19..bf2e89de44c 100644 --- a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py @@ -741,3 +741,39 @@ async def test_realtime_health_check_uses_bearer_token_when_no_api_key(monkeypat is True ) assert connect_calls[0]["additional_headers"] == {"Authorization": "Bearer my-entra-token"} + + +@pytest.mark.asyncio +async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch): + """ + The router binds a deployment's `azure_ad_token` to `_arealtime`'s named parameter rather than + **kwargs, so it must still reach the handler. + + Regression test for https://github.com/BerriAI/litellm/issues/34654 + """ + from litellm.realtime_api import main as realtime_main + + mock_async_realtime = AsyncMock() + monkeypatch.setattr(realtime_main, "azure_realtime", MagicMock(async_realtime=mock_async_realtime)) + monkeypatch.setattr( + realtime_main, + "get_llm_provider", + lambda model, api_base=None, api_key=None: ( + "gpt-realtime-whisper", + "azure", + None, + "https://my-endpoint.openai.azure.com", + ), + ) + monkeypatch.delenv("AZURE_API_KEY", raising=False) + monkeypatch.setattr(realtime_main.litellm, "api_key", None) + + await realtime_main._arealtime( + model="azure/gpt-realtime-whisper", + websocket=MagicMock(), + api_version="2024-10-01-preview", + azure_ad_token="deployment-entra-token", + litellm_logging_obj=MagicMock(), + ) + + assert mock_async_realtime.call_args.kwargs["azure_ad_token"] == "deployment-entra-token" From 0c5583b83f03d642f6ee4f42619b94023ce3d7e8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 05:18:49 +0000 Subject: [PATCH 003/318] fix(google_genai): price streamed generateContent with the provider that served it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/google_genai/streaming_iterator.py | 14 ++- .../vertex_passthrough_logging_handler.py | 2 +- .../streaming_handler.py | 19 ++++ .../pass_through_endpoints.py | 1 + .../test_google_genai_streaming_iterator.py | 40 +++++++- .../test_streaming_handler.py | 99 +++++++++++++++++++ 6 files changed, 170 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index e03f7ee745f..e2fac6a615b 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -2,6 +2,7 @@ import asyncio from datetime import datetime from typing import TYPE_CHECKING, Any, Final +import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, @@ -65,6 +66,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: litellm_logging_obj: LiteLLMLoggingObj, request_body: dict, model: str, + custom_llm_provider: str, hidden_params: dict[str, Any] | None = None, ): self.litellm_logging_obj = litellm_logging_obj @@ -72,6 +74,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: self.start_time = datetime.now() self.collected_chunks: list[bytes] = [] self.model = model + self.custom_llm_provider = custom_llm_provider self._hidden_params: dict[str, Any] = hidden_params or {} async def _handle_async_streaming_logging( @@ -83,13 +86,18 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: ) end_time: Final = datetime.now() + endpoint_type: Final = ( + EndpointType.GEMINI + if self.custom_llm_provider == litellm.LlmProviders.GEMINI.value + else EndpointType.VERTEX_AI + ) asyncio.create_task( PassThroughStreamingHandler._route_streaming_logging_to_handler( litellm_logging_obj=self.litellm_logging_obj, passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, url_route="/v1/generateContent", request_body=self.request_body or {}, - endpoint_type=EndpointType.VERTEX_AI, + endpoint_type=endpoint_type, start_time=self.start_time, raw_bytes=self.collected_chunks, end_time=end_time, @@ -118,13 +126,13 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent litellm_logging_obj=logging_obj, request_body=request_body or {}, model=model, + custom_llm_provider=custom_llm_provider, hidden_params=hidden_params, ) self.response = response self.model = model self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata - self.custom_llm_provider = custom_llm_provider # Gemini streamGenerateContent uses SSE line framing; iter_lines keeps # large inlineData payloads (e.g. image/jpeg) intact within one event. self.stream_iterator = response.iter_lines() @@ -169,13 +177,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo litellm_logging_obj=logging_obj, request_body=request_body or {}, model=model, + custom_llm_provider=custom_llm_provider, hidden_params=hidden_params, ) self.response = response self.model = model self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata - self.custom_llm_provider = custom_llm_provider # Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps # large inlineData payloads (e.g. image/jpeg) intact within one event. self.stream_iterator = response.aiter_lines() diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index afd8684dd92..36455611c95 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -592,7 +592,7 @@ class VertexPassthroughLoggingHandler: response_cost: Final = litellm.completion_cost( completion_response=litellm_model_response, model=model, - custom_llm_provider="vertex_ai", + custom_llm_provider=custom_llm_provider, ) kwargs["response_cost"] = response_cost diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index ff1c12d08d7..907c59d28cc 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -15,6 +15,9 @@ from litellm.types.utils import StandardPassThroughResponseObject from .llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, ) +from .llm_provider_handlers.gemini_passthrough_logging_handler import ( + GeminiPassthroughLoggingHandler, +) from .llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, ) @@ -221,6 +224,22 @@ class PassThroughStreamingHandler: ) standard_logging_response_object = vertex_passthrough_logging_handler_result["result"] kwargs = vertex_passthrough_logging_handler_result["kwargs"] + elif endpoint_type == EndpointType.GEMINI: + gemini_passthrough_logging_handler_result: Final = ( + GeminiPassthroughLoggingHandler._handle_logging_gemini_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + all_chunks=all_chunks, + end_time=end_time, + model=model, + ) + ) + standard_logging_response_object = gemini_passthrough_logging_handler_result["result"] + kwargs = gemini_passthrough_logging_handler_result["kwargs"] elif endpoint_type == EndpointType.OPENAI: openai_passthrough_logging_handler_result: Final = ( OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index f59ca0d9041..548702e4139 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -22,6 +22,7 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__ class EndpointType(str, Enum): VERTEX_AI = "vertex-ai" + GEMINI = "gemini" ANTHROPIC = "anthropic" OPENAI = "openai" GENERIC = "generic" diff --git a/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py b/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py index d74a05ec59c..91058767730 100644 --- a/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py +++ b/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py @@ -1,5 +1,6 @@ +import asyncio import json -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -8,6 +9,43 @@ from litellm.google_genai.streaming_iterator import ( GoogleGenAIGenerateContentStreamingIterator, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "custom_llm_provider, expected_endpoint_type", + [("gemini", EndpointType.GEMINI), ("vertex_ai", EndpointType.VERTEX_AI)], +) +async def test_streaming_logging_routes_to_the_provider_that_served_the_request( + custom_llm_provider, expected_endpoint_type +): + """Routing every google stream through the vertex handler bills gemini/* at vertex_ai/ rates.""" + mock_response = MagicMock() + + async def _aiter_lines(): + yield 'data: {"candidates": []}' + + mock_response.aiter_lines = _aiter_lines + + iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( + response=mock_response, + model="gemini-3.1-flash-image", + logging_obj=MagicMock(spec=LiteLLMLoggingObj), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider=custom_llm_provider, + ) + + with patch( + "litellm.proxy.pass_through_endpoints.streaming_handler.PassThroughStreamingHandler._route_streaming_logging_to_handler", + new=AsyncMock(), + ) as mock_route: + async for _ in iterator: + pass + + await asyncio.sleep(0) + assert mock_route.call_args.kwargs["endpoint_type"] == expected_endpoint_type def _large_inline_data_event() -> str: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py new file mode 100644 index 00000000000..d0c28fd60a9 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py @@ -0,0 +1,99 @@ +import json +from datetime import datetime +from unittest.mock import MagicMock + +import pytest + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( + VertexPassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, +) +from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + +MODEL = "gemini-3.1-flash-image" + +# gemini/ rate card: 2.5e-07 in, 1.5e-06 out. vertex_ai/ rate card is exactly 2x that. +GEMINI_COST = 1000 * 2.5e-07 + 1000 * 1.5e-06 +VERTEX_COST = 2 * GEMINI_COST + + +def _chunks() -> list[str]: + payload = { + "candidates": [ + { + "content": {"parts": [{"text": "hi"}], "role": "model"}, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 1000, + "candidatesTokenCount": 1000, + "totalTokenCount": 2000, + }, + "modelVersion": MODEL, + } + return [f"data: {json.dumps(payload)}"] + + +def _logging_obj() -> LiteLLMLoggingObj: + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + logging_obj.model_call_details = {} + logging_obj.optional_params = {} + logging_obj.litellm_call_id = "test-call-id" + return logging_obj + + +@pytest.mark.parametrize( + "endpoint_type, expected_provider, expected_cost", + [ + (EndpointType.GEMINI, "gemini", GEMINI_COST), + (EndpointType.VERTEX_AI, "vertex_ai", VERTEX_COST), + ], +) +def test_streaming_generate_content_bills_against_the_requested_provider( + endpoint_type, expected_provider, expected_cost +): + """A streamed gemini/* request must not be priced off the vertex_ai/ rate card.""" + logging_obj = _logging_obj() + + _, kwargs = PassThroughStreamingHandler._build_passthrough_logging_result( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=PassThroughEndpointLogging(), + url_route="/v1/generateContent", + request_body={}, + endpoint_type=endpoint_type, + start_time=datetime.now(), + raw_bytes=[chunk.encode("utf-8") for chunk in _chunks()], + end_time=datetime.now(), + model=MODEL, + ) + + assert kwargs["response_cost"] == pytest.approx(expected_cost) + assert logging_obj.model_call_details["custom_llm_provider"] == expected_provider + + +def test_vertex_generate_content_payload_prices_gemini_urls_at_gemini_rates(): + """The AI Studio host resolves to `gemini`, so the cost must follow it, not the vertex_ai default.""" + logging_obj = _logging_obj() + + result = VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=PassThroughEndpointLogging(), + url_route=f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:streamGenerateContent", + request_body={}, + endpoint_type=EndpointType.VERTEX_AI, + start_time=datetime.now(), + all_chunks=_chunks(), + model=MODEL, + end_time=datetime.now(), + ) + + assert result["kwargs"]["response_cost"] == pytest.approx(GEMINI_COST) + assert logging_obj.model_call_details["custom_llm_provider"] == "gemini" From e4c2ad4627b71d280603f721234ec8990f3aa6bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 19:23:34 -0700 Subject: [PATCH 004/318] fix(anthropic): buffer streamed responses carrying server-fulfilled tools so retrieval tool calls never reach the client --- .../compression_interception/handler.py | 4 +- litellm/integrations/custom_logger.py | 4 +- .../messages/agentic_streaming_iterator.py | 59 ++++++ litellm/llms/custom_httpx/llm_http_handler.py | 23 +++ .../guardrail_hooks/headroom/headroom.py | 1 + .../test_agentic_streaming_iterator.py | 178 ++++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 71 +++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 7 + 8 files changed, 345 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 7ea60053e6f..76720682101 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -7,7 +7,7 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan. import time import uuid -from typing import Any, Final, cast +from typing import Any, ClassVar, Final, cast from litellm._logging import verbose_logger from litellm.compression import compress @@ -72,6 +72,8 @@ class CompressionInterceptionLogger(CustomLogger): 4. Build typed rerun plan with tool_result blocks from the compressed cache. """ + server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({LITELLM_CONTENT_RETRIEVE_TOOL_NAME}) + def __init__( self, enabled: bool = True, diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a0c78674ac8..60af4063f84 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -3,7 +3,7 @@ import re import traceback from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional from pydantic import BaseModel @@ -60,6 +60,8 @@ _BASE64_INLINE_PATTERN: Final = re.compile( class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class # Class variables or attributes + server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset() + def __init__( self, turn_off_message_logging: bool = False, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index 5c4fa4700c0..0699595821e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -6,14 +6,27 @@ yields every chunk to the caller (preserving real streaming), collects all bytes, and on stream exhaustion rebuilds the full Anthropic response to run through agentic completion hooks. If an agentic hook fires, the follow-up response is chained as Phase 2 of the same iterator. + +In hold-back mode (``hold_back=True``), chunks are buffered instead of +yielded live, with SSE ping events emitted while the upstream message is +in flight. On exhaustion the hooks run first: if a follow-up response +replaces the message, only the follow-up is yielded and the buffered +message is dropped; otherwise the buffer is replayed verbatim. This is +required for server-fulfilled tools (e.g. ``headroom_retrieve``), whose +tool_use blocks must never reach a client that cannot execute them. """ +import asyncio +import contextlib import json from collections.abc import AsyncIterator from typing import Any, Final, cast from litellm._logging import verbose_logger +PING_SSE_BYTES: Final = b'event: ping\ndata: {"type": "ping"}\n\n' +HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0 + # --------------------------------------------------------------------------- # SSE parsing helpers (module-level to keep the class lean) # --------------------------------------------------------------------------- @@ -156,6 +169,8 @@ class AgenticAnthropicStreamingIterator: logging_obj: Any, custom_llm_provider: str, kwargs: dict, + hold_back: bool = False, + ping_interval_seconds: float = HOLD_BACK_PING_INTERVAL_SECONDS, ): self._inner = completion_stream.__aiter__() self._http_handler = http_handler @@ -166,16 +181,23 @@ class AgenticAnthropicStreamingIterator: self._logging_obj = logging_obj self._custom_llm_provider = custom_llm_provider self._kwargs = kwargs + self._hold_back = hold_back + self._ping_interval_seconds = ping_interval_seconds self._collected_bytes: list[bytes] = [] self._stream_exhausted = False self._hook_processing_done = False self._follow_up_iterator: AsyncIterator | None = None + self._drain_task: asyncio.Task | None = None + self._replay_index = 0 def __aiter__(self): return self async def __anext__(self) -> bytes: + if self._hold_back: + return await self._anext_held_back() + # Phase 1: yield from upstream, collect bytes if not self._stream_exhausted: try: @@ -194,11 +216,48 @@ class AgenticAnthropicStreamingIterator: raise StopAsyncIteration + async def _drain_upstream(self) -> None: + try: + while True: + self._collected_bytes.append(await self._inner.__anext__()) + except StopAsyncIteration: + return + + async def _anext_held_back(self) -> bytes: + if self._drain_task is None: + self._drain_task = asyncio.create_task(self._drain_upstream()) + return PING_SSE_BYTES + + while not self._stream_exhausted: + try: + await asyncio.wait_for(asyncio.shield(self._drain_task), timeout=self._ping_interval_seconds) + except asyncio.TimeoutError: + return PING_SSE_BYTES + self._stream_exhausted = True + await self._process_agentic_hooks() + + if self._follow_up_iterator is not None: + return await self._follow_up_iterator.__anext__() + + if self._replay_index < len(self._collected_bytes): + chunk: Final = self._collected_bytes[self._replay_index] + self._replay_index += 1 + return chunk + + raise StopAsyncIteration + async def aclose(self) -> None: from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( aclose_if_supported, ) + if self._drain_task is not None and self._drain_task.done(): + if not self._drain_task.cancelled(): + self._drain_task.exception() + elif self._drain_task is not None: + self._drain_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._drain_task await aclose_if_supported(self._inner) await aclose_if_supported(self._follow_up_iterator) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a58397c9184..913cedcfa55 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2189,6 +2189,10 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, kwargs={**kwargs, "api_key": api_key} if api_key else kwargs, + hold_back=self._should_hold_back_stream( + logging_obj=logging_obj, + tools=anthropic_messages_optional_request_params.get("tools"), + ), ) return AnthropicMessagesStreamingResponse( completion_stream=initial_response, @@ -5033,6 +5037,25 @@ class BaseLLMHTTPHandler: return True return False + @staticmethod + def _should_hold_back_stream(logging_obj: LiteLLMLoggingObj, tools: object) -> bool: + """ + True when the request carries a tool that a registered callback fulfills + server-side (e.g. ``headroom_retrieve``). The model's tool_use for such a + tool must never reach the client, which cannot execute it: the agentic + loop replaces the whole message with a follow-up response, so the stream + is buffered (with ping keepalives) instead of forwarded live. + """ + if not isinstance(tools, list) or not tools: + return False + from litellm.litellm_core_utils.prompt_templates.factory import has_tool_with_name + + return any( + has_tool_with_name(tools, name) + for cb in _custom_logger_callbacks(logging_obj) + for name in getattr(cb, "server_fulfilled_tool_names", frozenset()) + ) + @staticmethod def _check_agentic_loop_safety( tool_calls: object, diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 8bfd5cca58a..84c6b220e62 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -339,6 +339,7 @@ def _build_responses_followup_items( class HeadroomGuardrail(CustomGuardrail): records_own_guardrail_information: ClassVar[bool] = True + server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({HEADROOM_RETRIEVE_TOOL_NAME}) @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py index b9bda07336f..a6b071bbab5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -2,6 +2,7 @@ Tests for AgenticAnthropicStreamingIterator and SSE rebuild helpers. """ +import asyncio import json import os import sys @@ -13,6 +14,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( + PING_SSE_BYTES, AgenticAnthropicStreamingIterator, _handle_content_block_delta, _handle_content_block_start, @@ -230,6 +232,51 @@ class MockAsyncStream: return chunk +class MockSlowAsyncStream(MockAsyncStream): + """Async iterator that sleeps before every chunk.""" + + def __init__(self, chunks: List[bytes], delay_seconds: float): + super().__init__(chunks) + self._delay_seconds = delay_seconds + + async def __anext__(self) -> bytes: + await asyncio.sleep(self._delay_seconds) + return await super().__anext__() + + +class MockFailingAsyncStream(MockAsyncStream): + """Async iterator that raises after yielding its chunks.""" + + def __init__(self, chunks: List[bytes], error: Exception): + super().__init__(chunks) + self._error = error + + async def __anext__(self) -> bytes: + if self._idx >= len(self._chunks): + raise self._error + return await super().__anext__() + + +def _build_hold_back_iterator( + stream: MockAsyncStream, + mock_handler: MagicMock, + ping_interval_seconds: float = 15.0, +) -> AgenticAnthropicStreamingIterator: + return AgenticAnthropicStreamingIterator( + completion_stream=stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + hold_back=True, + ping_interval_seconds=ping_interval_seconds, + ) + + # --------------------------------------------------------------------------- # Tests for _parse_sse_events # --------------------------------------------------------------------------- @@ -790,3 +837,134 @@ class TestAgenticStreamingIteratorErrorHandling: call_kwargs = mock_handler._call_agentic_completion_hooks.call_args assert call_kwargs.kwargs["stream"] is True + + +class TestAgenticStreamingIteratorHoldBack: + @pytest.mark.asyncio + async def test_should_not_leak_intercepted_message_when_follow_up_fires(self): + """The buffered tool_use message must be dropped: only pings and follow-up bytes reach the client.""" + phase1_chunks = _build_tool_use_stream() + phase2_chunks = [b"follow-up-chunk-1", b"follow-up-chunk-2"] + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=MockAsyncStream(phase2_chunks)) + + iterator = _build_hold_back_iterator(MockAsyncStream(phase1_chunks), mock_handler) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + non_ping = [c for c in collected if c != PING_SSE_BYTES] + assert non_ping == phase2_chunks + assert b"litellm_content_retrieve" not in b"".join(collected) + assert collected[0] == PING_SSE_BYTES + + @pytest.mark.asyncio + async def test_should_replay_buffer_verbatim_when_no_hook_fires(self): + """Without interception the buffered message is replayed byte-identical after the pings.""" + chunks = _build_simple_text_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert [c for c in collected if c != PING_SSE_BYTES] == chunks + mock_handler._call_agentic_completion_hooks.assert_awaited_once() + + @pytest.mark.asyncio + async def test_should_emit_pings_while_upstream_is_slow(self): + """Pings keep the client connection alive while the upstream message is buffered.""" + chunks = _build_simple_text_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator( + MockSlowAsyncStream(chunks, delay_seconds=0.05), + mock_handler, + ping_interval_seconds=0.02, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert collected.count(PING_SSE_BYTES) >= 2 + assert [c for c in collected if c != PING_SSE_BYTES] == chunks + + @pytest.mark.asyncio + async def test_should_propagate_upstream_error_instead_of_partial_message(self): + """An upstream failure surfaces as an error; the client never receives a truncated message.""" + chunks = _build_simple_text_stream()[:2] + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator( + MockFailingAsyncStream(chunks, RuntimeError("upstream died")), + mock_handler, + ) + + collected = [] + with pytest.raises(RuntimeError, match="upstream died"): + async for chunk in iterator: + collected.append(chunk) + + assert all(c == PING_SSE_BYTES for c in collected) + mock_handler._call_agentic_completion_hooks.assert_not_awaited() + + @pytest.mark.asyncio + async def test_should_replay_buffer_when_hook_processing_errors(self): + """A hook crash degrades to replaying the original message rather than dropping it.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=RuntimeError("hook exploded")) + + mock_logging = MagicMock() + mock_logging.litellm_call_id = "test_call_holdback" + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=MockAsyncStream(chunks), + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=mock_logging, + custom_llm_provider="anthropic", + kwargs={}, + hold_back=True, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert [c for c in collected if c != PING_SSE_BYTES] == chunks + + @pytest.mark.asyncio + async def test_aclose_cancels_drain_task(self): + """Closing the iterator mid-buffer must cancel the background drain task.""" + chunks = _build_simple_text_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator( + MockSlowAsyncStream(chunks, delay_seconds=5.0), + mock_handler, + ) + + first = await iterator.__anext__() + assert first == PING_SSE_BYTES + assert iterator._drain_task is not None + + await iterator.aclose() + assert iterator._drain_task.cancelled() 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..6c2727e7da2 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,74 @@ 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"] + + +class TestShouldHoldBackStream: + """_should_hold_back_stream gates the buffered (non-leaking) streaming mode + for server-fulfilled tools like headroom_retrieve.""" + + @staticmethod + def _logging_obj_with(callbacks): + logging_obj = Mock() + logging_obj.dynamic_success_callbacks = callbacks + return logging_obj + + def test_should_hold_back_when_callback_owns_tool_in_request(self): + from litellm.integrations.custom_logger import CustomLogger + + class RetrievalCallback(CustomLogger): + server_fulfilled_tool_names = frozenset({"headroom_retrieve"}) + + tools = [ + {"name": "Bash", "input_schema": {"type": "object"}}, + {"name": "headroom_retrieve", "input_schema": {"type": "object"}}, + ] + assert ( + BaseLLMHTTPHandler._should_hold_back_stream( + logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools + ) + is True + ) + + def test_should_stream_live_when_tool_absent_from_request(self): + from litellm.integrations.custom_logger import CustomLogger + + class RetrievalCallback(CustomLogger): + server_fulfilled_tool_names = frozenset({"headroom_retrieve"}) + + tools = [{"name": "Bash", "input_schema": {"type": "object"}}] + assert ( + BaseLLMHTTPHandler._should_hold_back_stream( + logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools + ) + is False + ) + + def test_should_stream_live_when_no_callback_declares_tool_names(self): + from litellm.integrations.custom_logger import CustomLogger + + tools = [{"name": "headroom_retrieve", "input_schema": {"type": "object"}}] + assert ( + BaseLLMHTTPHandler._should_hold_back_stream( + logging_obj=self._logging_obj_with([CustomLogger()]), tools=tools + ) + is False + ) + + def test_should_stream_live_without_tools(self): + assert BaseLLMHTTPHandler._should_hold_back_stream(logging_obj=self._logging_obj_with([]), tools=None) is False + + def test_interception_callbacks_declare_their_retrieval_tools(self): + from litellm.integrations.compression_interception.handler import ( + LITELLM_CONTENT_RETRIEVE_TOOL_NAME, + CompressionInterceptionLogger, + ) + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( + HEADROOM_RETRIEVE_TOOL_NAME, + HeadroomGuardrail, + ) + + assert HeadroomGuardrail.server_fulfilled_tool_names == frozenset({HEADROOM_RETRIEVE_TOOL_NAME}) + assert CompressionInterceptionLogger.server_fulfilled_tool_names == frozenset( + {LITELLM_CONTENT_RETRIEVE_TOOL_NAME} + ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f1660e77ad9..8e950874a10 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21391,6 +21391,13 @@ export interface components { * @description What the routed traffic actually cost */ spend: number; + /** + * Tier Turns + * @description Turns per tier, keyed by the tier name the routing decision recorded at request time (never re-derived at read time, since the tier-to-model mapping is mutable config). Tier names are scoped to this group's router_type and are not comparable across types: a complexity router reports 'simple'/'medium'/'complex'/'reasoning', a quality router reports its numeric quality tier, and an adaptive router records no tier at all. Turns no tier served (the classifier fell back to default_model) are absent rather than pooled under a sentinel key, so the values may sum to less than turns + */ + tier_turns?: { + [key: string]: number; + }; /** Turns */ turns: number; }; From f994068a7338b4bb54fa7a53d76fb009dd3e9f6a Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 07:30:24 +0000 Subject: [PATCH 005/318] fix(anthropic): keep pinging during agentic hooks and fail instead of replaying server-fulfilled tool_use Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/agentic_streaming_iterator.py | 87 ++++++++++++--- litellm/llms/custom_httpx/llm_http_handler.py | 29 ++--- .../test_agentic_streaming_iterator.py | 105 +++++++++++++++--- .../custom_httpx/test_llm_http_handler.py | 28 ++--- 4 files changed, 190 insertions(+), 59 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index 0699595821e..4cf348dda9e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -9,11 +9,13 @@ follow-up response is chained as Phase 2 of the same iterator. In hold-back mode (``hold_back=True``), chunks are buffered instead of yielded live, with SSE ping events emitted while the upstream message is -in flight. On exhaustion the hooks run first: if a follow-up response -replaces the message, only the follow-up is yielded and the buffered -message is dropped; otherwise the buffer is replayed verbatim. This is -required for server-fulfilled tools (e.g. ``headroom_retrieve``), whose -tool_use blocks must never reach a client that cannot execute them. +in flight and while the agentic hooks run. On exhaustion the hooks run +first: if a follow-up response replaces the message, only the follow-up +is yielded and the buffered message is dropped; otherwise the buffer is +replayed verbatim, unless it holds a tool_use for a server-fulfilled tool +(e.g. ``headroom_retrieve``), in which case an SSE ``error`` event is +emitted because such a block must never reach a client that cannot +execute it. """ import asyncio @@ -26,6 +28,11 @@ from litellm._logging import verbose_logger PING_SSE_BYTES: Final = b'event: ping\ndata: {"type": "ping"}\n\n' HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0 +SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "api_error", "message": ' + b'"Server-side tool retrieval failed, so this turn could not be completed. Please retry."}}\n\n' +) # --------------------------------------------------------------------------- # SSE parsing helpers (module-level to keep the class lean) @@ -170,6 +177,7 @@ class AgenticAnthropicStreamingIterator: custom_llm_provider: str, kwargs: dict, hold_back: bool = False, + server_fulfilled_tool_names: frozenset[str] = frozenset(), ping_interval_seconds: float = HOLD_BACK_PING_INTERVAL_SECONDS, ): self._inner = completion_stream.__aiter__() @@ -182,6 +190,7 @@ class AgenticAnthropicStreamingIterator: self._custom_llm_provider = custom_llm_provider self._kwargs = kwargs self._hold_back = hold_back + self._server_fulfilled_tool_names = server_fulfilled_tool_names self._ping_interval_seconds = ping_interval_seconds self._collected_bytes: list[bytes] = [] @@ -189,7 +198,9 @@ class AgenticAnthropicStreamingIterator: self._hook_processing_done = False self._follow_up_iterator: AsyncIterator | None = None self._drain_task: asyncio.Task | None = None + self._hook_task: asyncio.Task | None = None self._replay_index = 0 + self._error_emitted = False def __aiter__(self): return self @@ -223,22 +234,42 @@ class AgenticAnthropicStreamingIterator: except StopAsyncIteration: return + async def _completed_within_ping_interval(self, task: asyncio.Task) -> bool: + try: + await asyncio.wait_for(asyncio.shield(task), timeout=self._ping_interval_seconds) + except asyncio.TimeoutError: + return False + return True + async def _anext_held_back(self) -> bytes: if self._drain_task is None: self._drain_task = asyncio.create_task(self._drain_upstream()) return PING_SSE_BYTES - while not self._stream_exhausted: - try: - await asyncio.wait_for(asyncio.shield(self._drain_task), timeout=self._ping_interval_seconds) - except asyncio.TimeoutError: + if not self._stream_exhausted: + if not await self._completed_within_ping_interval(self._drain_task): return PING_SSE_BYTES self._stream_exhausted = True - await self._process_agentic_hooks() + + if self._hook_task is None: + self._hook_task = asyncio.create_task(self._process_agentic_hooks()) + if not await self._completed_within_ping_interval(self._hook_task): + return PING_SSE_BYTES if self._follow_up_iterator is not None: return await self._follow_up_iterator.__anext__() + if self._buffer_holds_server_fulfilled_tool_use(): + if self._error_emitted: + raise StopAsyncIteration + self._error_emitted = True + verbose_logger.error( + "AgenticStreamingIterator: hooks did not replace a message containing a server-fulfilled " + "tool_use [model=%s]; emitting an SSE error instead of leaking the tool call to the client", + self._model, + ) + return SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES + if self._replay_index < len(self._collected_bytes): chunk: Final = self._collected_bytes[self._replay_index] self._replay_index += 1 @@ -246,18 +277,40 @@ class AgenticAnthropicStreamingIterator: raise StopAsyncIteration + def _buffer_holds_server_fulfilled_tool_use(self) -> bool: + if not self._server_fulfilled_tool_names: + return False + started_blocks: Final = ( + data.get("content_block") + for event_type, data in _parse_sse_events(b"".join(self._collected_bytes)) + if event_type == "content_block_start" + ) + return any( + isinstance(block, dict) + and block.get("type") == "tool_use" + and block.get("name") in self._server_fulfilled_tool_names + for block in started_blocks + ) + + @staticmethod + async def _settle_task(task: asyncio.Task | None) -> None: + if task is None: + return + if task.done(): + if not task.cancelled(): + task.exception() + return + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + async def aclose(self) -> None: from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( aclose_if_supported, ) - if self._drain_task is not None and self._drain_task.done(): - if not self._drain_task.cancelled(): - self._drain_task.exception() - elif self._drain_task is not None: - self._drain_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await self._drain_task + await self._settle_task(self._drain_task) + await self._settle_task(self._hook_task) await aclose_if_supported(self._inner) await aclose_if_supported(self._follow_up_iterator) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 913cedcfa55..e9b88e45219 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2179,6 +2179,10 @@ class BaseLLMHTTPHandler: AgenticAnthropicStreamingIterator, ) + held_back_tool_names: Final = self._server_fulfilled_tools_in_request( + logging_obj=logging_obj, + tools=anthropic_messages_optional_request_params.get("tools"), + ) initial_response = AgenticAnthropicStreamingIterator( completion_stream=completion_stream, http_handler=self, @@ -2189,10 +2193,8 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, kwargs={**kwargs, "api_key": api_key} if api_key else kwargs, - hold_back=self._should_hold_back_stream( - logging_obj=logging_obj, - tools=anthropic_messages_optional_request_params.get("tools"), - ), + hold_back=bool(held_back_tool_names), + server_fulfilled_tool_names=held_back_tool_names, ) return AnthropicMessagesStreamingResponse( completion_stream=initial_response, @@ -5038,22 +5040,23 @@ class BaseLLMHTTPHandler: return False @staticmethod - def _should_hold_back_stream(logging_obj: LiteLLMLoggingObj, tools: object) -> bool: + def _server_fulfilled_tools_in_request(logging_obj: LiteLLMLoggingObj, tools: object) -> frozenset[str]: """ - True when the request carries a tool that a registered callback fulfills - server-side (e.g. ``headroom_retrieve``). The model's tool_use for such a - tool must never reach the client, which cannot execute it: the agentic - loop replaces the whole message with a follow-up response, so the stream - is buffered (with ping keepalives) instead of forwarded live. + The request's tools that a registered callback fulfills server-side (e.g. + ``headroom_retrieve``). The model's tool_use for such a tool must never + reach the client, which cannot execute it: the agentic loop replaces the + whole message with a follow-up response, so a stream carrying any of + these is buffered (with ping keepalives) instead of forwarded live. """ if not isinstance(tools, list) or not tools: - return False + return frozenset() from litellm.litellm_core_utils.prompt_templates.factory import has_tool_with_name - return any( - has_tool_with_name(tools, name) + return frozenset( + name for cb in _custom_logger_callbacks(logging_obj) for name in getattr(cb, "server_fulfilled_tool_names", frozenset()) + if has_tool_with_name(tools, name) ) @staticmethod diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py index a6b071bbab5..d59f76232cd 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -15,6 +15,7 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( PING_SSE_BYTES, + SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, AgenticAnthropicStreamingIterator, _handle_content_block_delta, _handle_content_block_start, @@ -261,6 +262,7 @@ def _build_hold_back_iterator( stream: MockAsyncStream, mock_handler: MagicMock, ping_interval_seconds: float = 15.0, + server_fulfilled_tool_names: frozenset = frozenset({"litellm_content_retrieve"}), ) -> AgenticAnthropicStreamingIterator: return AgenticAnthropicStreamingIterator( completion_stream=stream, @@ -273,6 +275,7 @@ def _build_hold_back_iterator( custom_llm_provider="anthropic", kwargs={}, hold_back=True, + server_fulfilled_tool_names=server_fulfilled_tool_names, ping_interval_seconds=ping_interval_seconds, ) @@ -920,27 +923,76 @@ class TestAgenticStreamingIteratorHoldBack: mock_handler._call_agentic_completion_hooks.assert_not_awaited() @pytest.mark.asyncio - async def test_should_replay_buffer_when_hook_processing_errors(self): - """A hook crash degrades to replaying the original message rather than dropping it.""" + async def test_should_emit_pings_while_hooks_are_slow(self): + """Retrieval and follow-up generation can outlast a client's idle timeout, so hooks get keepalives too.""" + chunks = _build_tool_use_stream() + phase2_chunks = [b"follow-up-chunk"] + + async def slow_hooks(**_kwargs): + await asyncio.sleep(0.12) + return MockAsyncStream(phase2_chunks) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=slow_hooks) + + iterator = _build_hold_back_iterator( + MockAsyncStream(chunks), + mock_handler, + ping_interval_seconds=0.02, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert collected.count(PING_SSE_BYTES) >= 4 + assert [c for c in collected if c != PING_SSE_BYTES] == phase2_chunks + + @pytest.mark.asyncio + async def test_should_error_instead_of_replaying_server_fulfilled_tool_use_when_hook_crashes(self): + """A hook crash must not replay the buffered retrieval tool_use: that is the unknown-tool bug.""" chunks = _build_tool_use_stream() mock_handler = MagicMock() mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=RuntimeError("hook exploded")) - mock_logging = MagicMock() - mock_logging.litellm_call_id = "test_call_holdback" + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler) - iterator = AgenticAnthropicStreamingIterator( - completion_stream=MockAsyncStream(chunks), - http_handler=mock_handler, - model="claude-sonnet-4-20250514", - messages=[], - anthropic_messages_provider_config=MagicMock(), - anthropic_messages_optional_request_params={}, - logging_obj=mock_logging, - custom_llm_provider="anthropic", - kwargs={}, - hold_back=True, + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert [c for c in collected if c != PING_SSE_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] + assert b"litellm_content_retrieve" not in b"".join(collected) + + @pytest.mark.asyncio + async def test_should_error_instead_of_replaying_when_no_hook_fires_on_tool_use(self): + """Hooks returning None on a retrieval tool_use is still a leak, so the turn fails loudly.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert [c for c in collected if c != PING_SSE_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] + + @pytest.mark.asyncio + async def test_should_replay_client_owned_tool_use_verbatim(self): + """Only server-fulfilled tools are withheld: a client's own tool_use still reaches it byte-identical.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = _build_hold_back_iterator( + MockAsyncStream(chunks), + mock_handler, + server_fulfilled_tool_names=frozenset({"headroom_retrieve"}), ) collected = [] @@ -968,3 +1020,26 @@ class TestAgenticStreamingIteratorHoldBack: await iterator.aclose() assert iterator._drain_task.cancelled() + + @pytest.mark.asyncio + async def test_aclose_cancels_in_flight_hook_task(self): + """Closing while hooks are running must not leave the retrieval follow-up task orphaned.""" + chunks = _build_tool_use_stream() + + async def never_finishing_hooks(**_kwargs): + await asyncio.sleep(5.0) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=never_finishing_hooks) + + iterator = _build_hold_back_iterator( + MockAsyncStream(chunks), + mock_handler, + ping_interval_seconds=0.02, + ) + + while iterator._hook_task is None: + await iterator.__anext__() + + await iterator.aclose() + assert iterator._hook_task.cancelled() 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 6c2727e7da2..798fb4c92e2 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 @@ -2073,9 +2073,9 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques assert retry_authorization != first_attempt_headers["Authorization"] -class TestShouldHoldBackStream: - """_should_hold_back_stream gates the buffered (non-leaking) streaming mode - for server-fulfilled tools like headroom_retrieve.""" +class TestServerFulfilledToolsInRequest: + """_server_fulfilled_tools_in_request gates the buffered (non-leaking) streaming + mode for server-fulfilled tools like headroom_retrieve.""" @staticmethod def _logging_obj_with(callbacks): @@ -2093,12 +2093,9 @@ class TestShouldHoldBackStream: {"name": "Bash", "input_schema": {"type": "object"}}, {"name": "headroom_retrieve", "input_schema": {"type": "object"}}, ] - assert ( - BaseLLMHTTPHandler._should_hold_back_stream( - logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools - ) - is True - ) + assert BaseLLMHTTPHandler._server_fulfilled_tools_in_request( + logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools + ) == frozenset({"headroom_retrieve"}) def test_should_stream_live_when_tool_absent_from_request(self): from litellm.integrations.custom_logger import CustomLogger @@ -2108,10 +2105,10 @@ class TestShouldHoldBackStream: tools = [{"name": "Bash", "input_schema": {"type": "object"}}] assert ( - BaseLLMHTTPHandler._should_hold_back_stream( + BaseLLMHTTPHandler._server_fulfilled_tools_in_request( logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools ) - is False + == frozenset() ) def test_should_stream_live_when_no_callback_declares_tool_names(self): @@ -2119,14 +2116,17 @@ class TestShouldHoldBackStream: tools = [{"name": "headroom_retrieve", "input_schema": {"type": "object"}}] assert ( - BaseLLMHTTPHandler._should_hold_back_stream( + BaseLLMHTTPHandler._server_fulfilled_tools_in_request( logging_obj=self._logging_obj_with([CustomLogger()]), tools=tools ) - is False + == frozenset() ) def test_should_stream_live_without_tools(self): - assert BaseLLMHTTPHandler._should_hold_back_stream(logging_obj=self._logging_obj_with([]), tools=None) is False + assert ( + BaseLLMHTTPHandler._server_fulfilled_tools_in_request(logging_obj=self._logging_obj_with([]), tools=None) + == frozenset() + ) def test_interception_callbacks_declare_their_retrieval_tools(self): from litellm.integrations.compression_interception.handler import ( From 398e3d214cc97be0531428f3fb506a7cc42e2683 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 19:28:40 +0000 Subject: [PATCH 006/318] refactor(anthropic): trim hold-back commentary and drop dead rebuilt-content expression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/agentic_streaming_iterator.py | 17 ++++------------- litellm/llms/custom_httpx/llm_http_handler.py | 8 +------- 2 files changed, 5 insertions(+), 20 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index 4cf348dda9e..a88b148e92c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -7,14 +7,10 @@ all bytes, and on stream exhaustion rebuilds the full Anthropic response to run through agentic completion hooks. If an agentic hook fires, the follow-up response is chained as Phase 2 of the same iterator. -In hold-back mode (``hold_back=True``), chunks are buffered instead of -yielded live, with SSE ping events emitted while the upstream message is -in flight and while the agentic hooks run. On exhaustion the hooks run -first: if a follow-up response replaces the message, only the follow-up -is yielded and the buffered message is dropped; otherwise the buffer is -replayed verbatim, unless it holds a tool_use for a server-fulfilled tool -(e.g. ``headroom_retrieve``), in which case an SSE ``error`` event is -emitted because such a block must never reach a client that cannot +In hold-back mode (``hold_back=True``) chunks are buffered instead of yielded +live, keepalive pings run until the hooks finish, and then either the follow-up +replaces the message or the buffer replays, except that a buffered tool_use for +a server-fulfilled tool fails the turn rather than reaching a client that cannot execute it. """ @@ -329,11 +325,6 @@ class AgenticAnthropicStreamingIterator: verbose_logger.debug("AgenticStreamingIterator: Could not rebuild response from SSE bytes") return - [ - (f"{b.get('type')}({b.get('name', '')})" if b.get("type") == "tool_use" else b.get("type")) - for b in rebuilt.get("content", []) - ] - result: Final = await self._http_handler._call_agentic_completion_hooks( response=rebuilt, model=self._model, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index e9b88e45219..193a38a5404 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5041,13 +5041,7 @@ class BaseLLMHTTPHandler: @staticmethod def _server_fulfilled_tools_in_request(logging_obj: LiteLLMLoggingObj, tools: object) -> frozenset[str]: - """ - The request's tools that a registered callback fulfills server-side (e.g. - ``headroom_retrieve``). The model's tool_use for such a tool must never - reach the client, which cannot execute it: the agentic loop replaces the - whole message with a follow-up response, so a stream carrying any of - these is buffered (with ping keepalives) instead of forwarded live. - """ + """The request's tools that a registered callback fulfills server-side (e.g. ``headroom_retrieve``).""" if not isinstance(tools, list) or not tools: return frozenset() from litellm.litellm_core_utils.prompt_templates.factory import has_tool_with_name From cbefb1ce5ffbdc90c9e6691206752b40ec9ef6e0 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 19:41:42 +0000 Subject: [PATCH 007/318] fix(anthropic): keep pinging while the held-back follow-up stream is in flight Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/agentic_streaming_iterator.py | 27 ++++++++- .../test_agentic_streaming_iterator.py | 59 +++++++++++++++++++ 2 files changed, 83 insertions(+), 3 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index a88b148e92c..3aaba0b139c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -8,8 +8,8 @@ to run through agentic completion hooks. If an agentic hook fires, the follow-up response is chained as Phase 2 of the same iterator. In hold-back mode (``hold_back=True``) chunks are buffered instead of yielded -live, keepalive pings run until the hooks finish, and then either the follow-up -replaces the message or the buffer replays, except that a buffered tool_use for +live, keepalive pings run whenever no other byte is ready, and then either the +follow-up replaces the message or the buffer replays, except that a tool_use for a server-fulfilled tool fails the turn rather than reaching a client that cannot execute it. """ @@ -30,6 +30,14 @@ SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = ( b'"Server-side tool retrieval failed, so this turn could not be completed. Please retry."}}\n\n' ) + +async def _anext_or_none(iterator: AsyncIterator) -> bytes | None: + try: + return await iterator.__anext__() + except StopAsyncIteration: + return None + + # --------------------------------------------------------------------------- # SSE parsing helpers (module-level to keep the class lean) # --------------------------------------------------------------------------- @@ -195,6 +203,7 @@ class AgenticAnthropicStreamingIterator: self._follow_up_iterator: AsyncIterator | None = None self._drain_task: asyncio.Task | None = None self._hook_task: asyncio.Task | None = None + self._follow_up_chunk_task: asyncio.Task | None = None self._replay_index = 0 self._error_emitted = False @@ -253,7 +262,7 @@ class AgenticAnthropicStreamingIterator: return PING_SSE_BYTES if self._follow_up_iterator is not None: - return await self._follow_up_iterator.__anext__() + return await self._next_follow_up_chunk(self._follow_up_iterator) if self._buffer_holds_server_fulfilled_tool_use(): if self._error_emitted: @@ -273,6 +282,17 @@ class AgenticAnthropicStreamingIterator: raise StopAsyncIteration + async def _next_follow_up_chunk(self, follow_up_iterator: AsyncIterator) -> bytes: + if self._follow_up_chunk_task is None: + self._follow_up_chunk_task = asyncio.create_task(_anext_or_none(follow_up_iterator)) + if not await self._completed_within_ping_interval(self._follow_up_chunk_task): + return PING_SSE_BYTES + chunk: Final = self._follow_up_chunk_task.result() + self._follow_up_chunk_task = None + if chunk is None: + raise StopAsyncIteration + return chunk + def _buffer_holds_server_fulfilled_tool_use(self) -> bool: if not self._server_fulfilled_tool_names: return False @@ -307,6 +327,7 @@ class AgenticAnthropicStreamingIterator: await self._settle_task(self._drain_task) await self._settle_task(self._hook_task) + await self._settle_task(self._follow_up_chunk_task) await aclose_if_supported(self._inner) await aclose_if_supported(self._follow_up_iterator) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py index d59f76232cd..d4aebf099d1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -1001,6 +1001,65 @@ class TestAgenticStreamingIteratorHoldBack: assert [c for c in collected if c != PING_SSE_BYTES] == chunks + @pytest.mark.asyncio + async def test_should_emit_pings_while_the_follow_up_stream_is_slow(self): + """The corrected answer can be slow to generate, so the follow-up stream gets keepalives too.""" + chunks = _build_tool_use_stream() + phase2_chunks = [b"follow-up-chunk-1", b"follow-up-chunk-2"] + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=MockSlowAsyncStream(phase2_chunks, delay_seconds=0.06) + ) + + iterator = _build_hold_back_iterator( + MockAsyncStream(chunks), + mock_handler, + ping_interval_seconds=0.02, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + first_follow_up_index = collected.index(phase2_chunks[0]) + assert collected[first_follow_up_index + 1] == PING_SSE_BYTES + assert [c for c in collected if c != PING_SSE_BYTES] == phase2_chunks + + @pytest.mark.asyncio + async def test_should_propagate_follow_up_stream_error(self): + """A failing follow-up stream surfaces its error instead of hanging on pings forever.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=MockFailingAsyncStream([b"follow-up-chunk"], RuntimeError("follow-up died")) + ) + + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler, ping_interval_seconds=0.02) + + with pytest.raises(RuntimeError, match="follow-up died"): + async for _ in iterator: + pass + + @pytest.mark.asyncio + async def test_aclose_cancels_in_flight_follow_up_chunk_task(self): + """Closing while a follow-up chunk is pending must not orphan that task.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=MockSlowAsyncStream([b"follow-up-chunk"], delay_seconds=5.0) + ) + + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler, ping_interval_seconds=0.02) + + while iterator._follow_up_chunk_task is None: + await iterator.__anext__() + + await iterator.aclose() + assert iterator._follow_up_chunk_task.cancelled() + @pytest.mark.asyncio async def test_aclose_cancels_drain_task(self): """Closing the iterator mid-buffer must cancel the background drain task.""" From bb0bb48da8c3a5fe3557812e79a31c71c608e006 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 20:05:45 +0000 Subject: [PATCH 008/318] fix(proxy): do not let held-back keepalive pings block the budget reservation refund on client disconnect Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 ++ .../messages/agentic_streaming_iterator.py | 10 +++---- litellm/proxy/common_request_processing.py | 6 ++-- litellm/proxy/common_utils/sse_keepalive.py | 4 ++- .../test_agentic_streaming_iterator.py | 30 +++++++++---------- .../proxy/test_budget_reservation.py | 28 +++++++++++++++++ 6 files changed, 57 insertions(+), 23 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6f0e9e7afe2..9db9fb36b65 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -434,6 +434,8 @@ CONNECTION_ERROR_PATTERNS: Final[list[str]] = [ ] STREAM_SSE_DONE_STRING: Final[str] = "[DONE]" STREAM_SSE_DATA_PREFIX: Final[str] = "data: " +STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n' +STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8") ### SPEND TRACKING ### DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = float( os.getenv("DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index 3aaba0b139c..d6f4e51a09a 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -21,8 +21,8 @@ from collections.abc import AsyncIterator from typing import Any, Final, cast from litellm._logging import verbose_logger +from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES -PING_SSE_BYTES: Final = b'event: ping\ndata: {"type": "ping"}\n\n' HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0 SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = ( b"event: error\n" @@ -249,17 +249,17 @@ class AgenticAnthropicStreamingIterator: async def _anext_held_back(self) -> bytes: if self._drain_task is None: self._drain_task = asyncio.create_task(self._drain_upstream()) - return PING_SSE_BYTES + return STREAM_SSE_KEEPALIVE_PING_BYTES if not self._stream_exhausted: if not await self._completed_within_ping_interval(self._drain_task): - return PING_SSE_BYTES + return STREAM_SSE_KEEPALIVE_PING_BYTES self._stream_exhausted = True if self._hook_task is None: self._hook_task = asyncio.create_task(self._process_agentic_hooks()) if not await self._completed_within_ping_interval(self._hook_task): - return PING_SSE_BYTES + return STREAM_SSE_KEEPALIVE_PING_BYTES if self._follow_up_iterator is not None: return await self._next_follow_up_chunk(self._follow_up_iterator) @@ -286,7 +286,7 @@ class AgenticAnthropicStreamingIterator: if self._follow_up_chunk_task is None: self._follow_up_chunk_task = asyncio.create_task(_anext_or_none(follow_up_iterator)) if not await self._completed_within_ping_interval(self._follow_up_chunk_task): - return PING_SSE_BYTES + return STREAM_SSE_KEEPALIVE_PING_BYTES chunk: Final = self._follow_up_chunk_task.result() self._follow_up_chunk_task = None if chunk is None: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 159d7508f4e..2607ff411a9 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -27,6 +27,7 @@ from litellm.constants import ( MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, + STREAM_SSE_KEEPALIVE_PING_BYTES, UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -2953,8 +2954,9 @@ class ProxyBaseLLMRequestProcessing: # so a GeneratorExit on client disconnect is raised there and any # statement after the yield never runs. The slow-path hook is # awaited above, so a cancellation during it still leaves this - # False and refunds. - delivered_chunk = True + # False and refunds. A keepalive ping carries no provider output, + # so it must not suppress that refund. + delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES yield serialize_chunk(chunk) stream_completed = True except (asyncio.CancelledError, GeneratorExit): diff --git a/litellm/proxy/common_utils/sse_keepalive.py b/litellm/proxy/common_utils/sse_keepalive.py index 6700700ff7c..6e0ea4db431 100644 --- a/litellm/proxy/common_utils/sse_keepalive.py +++ b/litellm/proxy/common_utils/sse_keepalive.py @@ -6,7 +6,9 @@ from typing import Final import anyio -ANTHROPIC_PING_SSE_CHUNK: Final = 'event: ping\ndata: {"type": "ping"}\n\n' +from litellm.constants import STREAM_SSE_KEEPALIVE_PING_CHUNK + +ANTHROPIC_PING_SSE_CHUNK: Final = STREAM_SSE_KEEPALIVE_PING_CHUNK def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None: diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py index d4aebf099d1..b0467430533 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -13,8 +13,8 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) +from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( - PING_SSE_BYTES, SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, AgenticAnthropicStreamingIterator, _handle_content_block_delta, @@ -858,10 +858,10 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - non_ping = [c for c in collected if c != PING_SSE_BYTES] + non_ping = [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] assert non_ping == phase2_chunks assert b"litellm_content_retrieve" not in b"".join(collected) - assert collected[0] == PING_SSE_BYTES + assert collected[0] == STREAM_SSE_KEEPALIVE_PING_BYTES @pytest.mark.asyncio async def test_should_replay_buffer_verbatim_when_no_hook_fires(self): @@ -877,7 +877,7 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert [c for c in collected if c != PING_SSE_BYTES] == chunks + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == chunks mock_handler._call_agentic_completion_hooks.assert_awaited_once() @pytest.mark.asyncio @@ -898,8 +898,8 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert collected.count(PING_SSE_BYTES) >= 2 - assert [c for c in collected if c != PING_SSE_BYTES] == chunks + assert collected.count(STREAM_SSE_KEEPALIVE_PING_BYTES) >= 2 + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == chunks @pytest.mark.asyncio async def test_should_propagate_upstream_error_instead_of_partial_message(self): @@ -919,7 +919,7 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert all(c == PING_SSE_BYTES for c in collected) + assert all(c == STREAM_SSE_KEEPALIVE_PING_BYTES for c in collected) mock_handler._call_agentic_completion_hooks.assert_not_awaited() @pytest.mark.asyncio @@ -945,8 +945,8 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert collected.count(PING_SSE_BYTES) >= 4 - assert [c for c in collected if c != PING_SSE_BYTES] == phase2_chunks + assert collected.count(STREAM_SSE_KEEPALIVE_PING_BYTES) >= 4 + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == phase2_chunks @pytest.mark.asyncio async def test_should_error_instead_of_replaying_server_fulfilled_tool_use_when_hook_crashes(self): @@ -962,7 +962,7 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert [c for c in collected if c != PING_SSE_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] assert b"litellm_content_retrieve" not in b"".join(collected) @pytest.mark.asyncio @@ -979,7 +979,7 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert [c for c in collected if c != PING_SSE_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] @pytest.mark.asyncio async def test_should_replay_client_owned_tool_use_verbatim(self): @@ -999,7 +999,7 @@ class TestAgenticStreamingIteratorHoldBack: async for chunk in iterator: collected.append(chunk) - assert [c for c in collected if c != PING_SSE_BYTES] == chunks + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == chunks @pytest.mark.asyncio async def test_should_emit_pings_while_the_follow_up_stream_is_slow(self): @@ -1023,8 +1023,8 @@ class TestAgenticStreamingIteratorHoldBack: collected.append(chunk) first_follow_up_index = collected.index(phase2_chunks[0]) - assert collected[first_follow_up_index + 1] == PING_SSE_BYTES - assert [c for c in collected if c != PING_SSE_BYTES] == phase2_chunks + assert collected[first_follow_up_index + 1] == STREAM_SSE_KEEPALIVE_PING_BYTES + assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == phase2_chunks @pytest.mark.asyncio async def test_should_propagate_follow_up_stream_error(self): @@ -1074,7 +1074,7 @@ class TestAgenticStreamingIteratorHoldBack: ) first = await iterator.__anext__() - assert first == PING_SSE_BYTES + assert first == STREAM_SSE_KEEPALIVE_PING_BYTES assert iterator._drain_task is not None await iterator.aclose() diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 34adb4d2091..e9a0e80752d 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -7,6 +7,7 @@ from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache +from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_EndUserTable, @@ -2453,6 +2454,33 @@ async def test_streaming_cancel_after_chunk_keeps_reservation( streaming_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() +@pytest.mark.asyncio +async def test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_cost( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token, reservation = await _reserve_for_stream( + counter_cache, key_cache, proxy_logging_obj, "key-cancel-after-ping" + ) + + async def cancel_after_ping(user_api_key_dict, response, request_data): + yield STREAM_SSE_KEEPALIVE_PING_BYTES + raise asyncio.CancelledError() + + generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, cancel_after_ping) + received = [] + with pytest.raises(asyncio.CancelledError): + async for chunk in generator: + received.append(chunk) + + assert received == [STREAM_SSE_KEEPALIVE_PING_BYTES] + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-cancel-after-ping" + ) == pytest.approx(0.5) + assert reservation["finalized"] is True + + @pytest.mark.asyncio async def test_release_budget_reservation_on_cancel_swallows_release_errors(): # If the release itself fails (e.g. Redis unavailable) it must not escape From 2d1ee3aab2fe6a37c80085009f789416fe191d4e Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 20:39:55 +0000 Subject: [PATCH 009/318] fix(proxy): keep the reservation when a disconnect happens while provider output is held back Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/agentic_streaming_iterator.py | 5 ++ litellm/proxy/common_request_processing.py | 9 ++- .../proxy/test_budget_reservation.py | 63 +++++++++++++++++++ 3 files changed, 76 insertions(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index d6f4e51a09a..3d3d3a12b17 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -207,6 +207,11 @@ class AgenticAnthropicStreamingIterator: self._replay_index = 0 self._error_emitted = False + @property + def has_buffered_provider_output(self) -> bool: + """Whether provider output was received but withheld from the client behind keepalive pings.""" + return self._hold_back and bool(self._collected_bytes) + def __aiter__(self): return self diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2607ff411a9..3aeb7729c81 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -40,6 +40,9 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( + AgenticAnthropicStreamingIterator, +) from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model from litellm.proxy.auth.auth_utils import check_response_size_is_safe @@ -95,6 +98,10 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInforma } +def _withheld_provider_output(response: object) -> bool: + return isinstance(response, AgenticAnthropicStreamingIterator) and response.has_buffered_provider_output + + def _should_return_raw_model_name(request_data: dict[str, object]) -> bool: return any( isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True @@ -2970,7 +2977,7 @@ class ProxyBaseLLMRequestProcessing: # only sees GeneratorExit on GC) cannot own the refund. if not stream_completed: client_disconnected = True - if not delivered_chunk: + if not delivered_chunk and not _withheld_provider_output(response): from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, ) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index e9a0e80752d..c3210dd7f4d 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -8,6 +8,9 @@ from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES +from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( + AgenticAnthropicStreamingIterator, +) from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_EndUserTable, @@ -2376,6 +2379,11 @@ async def _reserve_for_stream(counter_cache, key_cache, proxy_logging_obj, token return valid_token, reservation +async def _never_ending_stream(): + yield b'event: message_start\ndata: {"type": "message_start"}\n\n' + await asyncio.sleep(30) + + def _drive_streaming_cancel(valid_token, iterator_hook): streaming_logging_obj = MagicMock() streaming_logging_obj.async_post_call_streaming_iterator_hook = iterator_hook @@ -2481,6 +2489,61 @@ async def test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_c assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_streaming_cancel_while_holding_back_provider_output_keeps_reservation( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token, reservation = await _reserve_for_stream( + counter_cache, key_cache, proxy_logging_obj, "key-cancel-held-back" + ) + + held_back = AgenticAnthropicStreamingIterator( + completion_stream=_never_ending_stream(), + http_handler=MagicMock(), + model="claude-haiku-4-5", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + hold_back=True, + server_fulfilled_tool_names=frozenset({"headroom_retrieve"}), + ping_interval_seconds=0.01, + ) + + async def ping_then_cancel(user_api_key_dict, response, request_data): + yield await response.__anext__() + while not response.has_buffered_provider_output: + yield await response.__anext__() + raise asyncio.CancelledError() + + streaming_logging_obj = MagicMock() + streaming_logging_obj.async_post_call_streaming_iterator_hook = ping_then_cancel + streaming_logging_obj._arelease_max_parallel_requests_on_disconnect = AsyncMock() + generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=held_back, + user_api_key_dict=valid_token, + request_data=_request_body(), + proxy_logging_obj=streaming_logging_obj, + serialize_chunk=lambda chunk: chunk, + serialize_error=lambda exc: str(exc), + ) + + received = [] + with pytest.raises(asyncio.CancelledError): + async for chunk in generator: + received.append(chunk) + + assert received and received == [STREAM_SSE_KEEPALIVE_PING_BYTES] * len(received) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-cancel-held-back" + ) == pytest.approx(2.0) + assert reservation.get("finalized") is not True + + @pytest.mark.asyncio async def test_release_budget_reservation_on_cancel_swallows_release_errors(): # If the release itself fails (e.g. Redis unavailable) it must not escape From 65eae963a7da34e9d4b714d4ce0b485168efaa21 Mon Sep 17 00:00:00 2001 From: Kunal Nayyar Date: Tue, 11 Aug 2026 13:03:55 +0530 Subject: [PATCH 010/318] feat(proxy): opt-in enforce rpm/tpm when adding a model Add general_settings toggle 'enforce_rpm_tpm_on_model_add' (default false). When true, /model/new rejects a model whose rpm or tpm is missing or not a positive value, so the Admin UI Add Model form surfaces a 400 validation error instead of silently storing an unbounded model (or one with a zero/negative limit that would exclude it from routing). --- .../model_management_endpoints.py | 37 +++++++++++++++++++ .../test_model_management_endpoints.py | 33 +++++++++++++++++ .../molecules/notifications_manager.tsx | 1 + 3 files changed, 71 insertions(+) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 8a52b0d1abb..d9030f42f69 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -239,6 +239,38 @@ def _raise_on_strategy_router_write_violation( ) +ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add" +_REQUIRED_RATE_LIMIT_FIELDS: Final = ("rpm", "tpm") + + +def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLMParams, enforced: bool) -> None: + """Require both rpm and tpm (each a positive value) when the operator opts in via config.yaml. + + Off by default, so deployments keep adding models without limits. When + ``enforce_rpm_tpm_on_model_add: true`` is set under general_settings, a model added + without both rpm and tpm set to a positive value is rejected rather than stored + unbounded (or effectively excluded from routing by a zero/negative limit). + """ + if not enforced: + return + missing: Final = tuple( + field + for field in _REQUIRED_RATE_LIMIT_FIELDS + if (value := getattr(litellm_params, field)) is None or value <= 0 + ) + if not missing: + return + raise ProxyException( + message=( + f"{' and '.join(missing)} must be set to a positive value when " + f"'{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING}' is enabled in general_settings" + ), + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_400_BAD_REQUEST, + param=f"litellm_params.{missing[0]}", + ) + + _PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to") @@ -1566,6 +1598,11 @@ async def add_new_model( existing_params=None, ) + _raise_if_rate_limits_required_but_missing( + litellm_params=model_params.litellm_params, + enforced=bool(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING, False)), + ) + model_response: LiteLLM_ProxyModelTable | None = None # update DB incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 454849d6430..cb87dfadcfa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, _get_team_deployments, + _raise_if_rate_limits_required_but_missing, clear_cache, delete_team_models, ) @@ -3825,3 +3826,35 @@ class TestAutoRouterClassifierDefaultPrompt: for empty in (None, "", "{}"): response = await get_auto_router_classifier_default_prompt(context_window_size=5, tier_labels=empty) assert response.system_prompt == classification_system_prompt(5) + + +class TestEnforceRpmTpmOnModelAdd: + def test_passes_when_disabled_even_without_limits(self): + _raise_if_rate_limits_required_but_missing( + litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), + enforced=False, + ) + + def test_passes_when_enabled_and_both_set(self): + _raise_if_rate_limits_required_but_missing( + litellm_params=LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=1000), + enforced=True, + ) + + @pytest.mark.parametrize( + "params, expected_missing", + [ + (LiteLLM_Params(model="azure/gpt-5.2"), "rpm and tpm"), + (LiteLLM_Params(model="azure/gpt-5.2", rpm=10), "tpm"), + (LiteLLM_Params(model="azure/gpt-5.2", tpm=1000), "rpm"), + (LiteLLM_Params(model="azure/gpt-5.2", rpm=0, tpm=1000), "rpm"), + (LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=-1), "tpm"), + ], + ) + def test_raises_when_enabled_and_missing(self, params, expected_missing): + from litellm.proxy._types import ProxyException + + with pytest.raises(ProxyException) as exc_info: + _raise_if_rate_limits_required_but_missing(litellm_params=params, enforced=True) + assert expected_missing in str(exc_info.value.message) + assert exc_info.value.code == "400" diff --git a/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx b/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx index 59b048b412c..31daee0b23b 100644 --- a/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx +++ b/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx @@ -104,6 +104,7 @@ const VALIDATION_MATCH = [ "invalid file type", "invalid field", "invalid date format", + "must be set when", ]; const NOT_FOUND_MATCH = [ From 526fc9eab192a6854f30c75aa574daf0d8d1f992 Mon Sep 17 00:00:00 2001 From: Kunal Nayyar Date: Tue, 11 Aug 2026 13:03:55 +0530 Subject: [PATCH 011/318] fix(ui): title validation errors correctly instead of Rate Limit Exceeded The /model/new endpoint returns a 400 validation error (type: validation_error) when 'rpm and tpm must be set to a positive value when enforce_rpm_tpm_on_model_add is enabled in general_settings' but the frontend's titleFor() keyword matcher mistitled it as 'Rate Limit Exceeded' because the message contains 'rpm'/'tpm' substrings, which matched the generic rate-limit keyword check before the more specific validation check could catch it. Add "'enforce_rpm_tpm_on_model_add' is enabled" to VALIDATION_MATCH so this message is classified as a Validation Error, matching the actual HTTP 400 validation_error the backend already returns. A narrow match on the setting name (rather than the generic "must be set when") avoids overriding the status-based classification of unrelated 401s, e.g. the PKCE 'GENERIC_CLIENT_ID must be set when PKCE is enabled' error. --- .../src/components/molecules/notifications_manager.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx b/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx index 31daee0b23b..aa3552da50a 100644 --- a/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx +++ b/ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx @@ -104,7 +104,7 @@ const VALIDATION_MATCH = [ "invalid file type", "invalid field", "invalid date format", - "must be set when", + "'enforce_rpm_tpm_on_model_add' is enabled", ]; const NOT_FOUND_MATCH = [ From 97290b4e0e4ce140a80cf76ef87b433e145276eb Mon Sep 17 00:00:00 2001 From: Daniel Vainshtein Date: Thu, 13 Aug 2026 13:44:15 +0300 Subject: [PATCH 012/318] fix(bedrock): parse cacheDetails for Converse 1h/5m cache write cost split AmazonConverseConfig._transform_usage only read the aggregate cacheWriteInputTokens field, so cache_creation_token_details was always unset for Bedrock Converse responses. calculate_cache_writing_cost bills the whole cache-write count at the 5m rate whenever that field is None, so 1-hour TTL cache writes on the standard Bedrock chat path were always undercounted, even though Bedrock returns the 5m/1h split in usage.cacheDetails. Parse cacheDetails (when present) into CacheCreationTokenDetails so the correct rate applies to each portion. No cacheDetails in the response (older models/regions) keeps the previous behavior. Fixes #36760 Co-Authored-By: pi (Claude/GPT via @earendil-works/pi-coding-agent) --- .../bedrock/chat/converse_transformation.py | 20 +++++++++ litellm/types/llms/bedrock.py | 14 ++++-- .../chat/test_converse_transformation.py | 43 +++++++++++++++++++ 3 files changed, 74 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 85918d40e12..9e7c86e615f 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -57,6 +57,7 @@ from litellm.types.llms.openai import ( OpenAIMessageContentListBlock, ) from litellm.types.utils import ( + CacheCreationTokenDetails, ChatCompletionMessageToolCall, CompletionTokensDetailsWrapper, Function, @@ -1770,6 +1771,24 @@ class AmazonConverseConfig(BaseConfig): thinking_blocks_list.append(_redacted_block) return thinking_blocks_list + @staticmethod + def _parse_cache_details(usage: ConverseTokenUsageBlock) -> "CacheCreationTokenDetails | None": + """ + Split Converse's aggregate cacheWriteInputTokens into the 5m/1h TTL + breakdown from `cacheDetails`, so cost calc can bill each tier + correctly instead of defaulting the whole write to the 5m rate. + https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html + """ + cache_details = usage.get("cacheDetails") + if not cache_details: + return None + tokens_5m = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "5m") + tokens_1h = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "1h") + return CacheCreationTokenDetails( + ephemeral_5m_input_tokens=tokens_5m, + ephemeral_1h_input_tokens=tokens_1h, + ) + def _transform_usage( self, usage: ConverseTokenUsageBlock, @@ -1792,6 +1811,7 @@ class AmazonConverseConfig(BaseConfig): prompt_tokens_details: Final = PromptTokensDetailsWrapper( cached_tokens=cache_read_input_tokens, cache_creation_tokens=cache_creation_input_tokens, + cache_creation_token_details=self._parse_cache_details(usage), text_tokens=raw_input_tokens, ) reasoning_tokens = token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 5665aa3277a..4c3ed8d6993 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -216,14 +216,22 @@ class ConverseResponseOutputBlock(TypedDict): message: MessageBlock | None -class ConverseTokenUsageBlock(TypedDict): +class CacheDetailBlock(TypedDict): + """Per-TTL cache-write breakdown. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html""" + inputTokens: int - outputTokens: int - totalTokens: int + ttl: Literal["5m", "1h"] + + +class ConverseTokenUsageBlock(TypedDict, total=False): + inputTokens: Required[int] + outputTokens: Required[int] + totalTokens: Required[int] cacheReadInputTokenCount: int cacheReadInputTokens: int cacheWriteInputTokenCount: int cacheWriteInputTokens: int + cacheDetails: list[CacheDetailBlock] class ServiceTierBlock(TypedDict): 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 d1d1f9ab489..393499fb041 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -51,6 +51,49 @@ def test_transform_usage(): assert openai_usage.completion_tokens_details.text_tokens == usage["outputTokens"] +def test_transform_usage_with_cache_details(): + """cacheDetails should split cacheWriteInputTokens into the 5m/1h TTL breakdown + so cost calc can bill the 1h portion at its own (higher) rate instead of + defaulting the whole write to the 5m rate. See issue #36760.""" + usage = ConverseTokenUsageBlock( + **{ + "inputTokens": 76, + "outputTokens": 259, + "totalTokens": 335, + "cacheWriteInputTokens": 362, + "cacheDetails": [ + {"inputTokens": 74, "ttl": "1h"}, + {"inputTokens": 288, "ttl": "5m"}, + ], + } + ) + config = AmazonConverseConfig() + openai_usage = config._transform_usage(usage) + details = openai_usage.prompt_tokens_details.cache_creation_token_details + assert details is not None + assert details.ephemeral_1h_input_tokens == 74 + assert details.ephemeral_5m_input_tokens == 288 + + +def test_transform_usage_without_cache_details_stays_none(): + """No cacheDetails in the response (older models/regions) should leave + cache_creation_token_details unset, same as before this field existed.""" + usage = ConverseTokenUsageBlock( + **{ + "inputTokens": 3, + "outputTokens": 401, + "totalTokens": 2193, + "cacheWriteInputTokens": 1789, + } + ) + config = AmazonConverseConfig() + openai_usage = config._transform_usage(usage) + assert ( + getattr(openai_usage.prompt_tokens_details, "cache_creation_token_details", None) + is None + ) + + def test_transform_usage_with_reasoning_content(): """Test that completion_tokens_details correctly tracks reasoning vs text tokens.""" usage = ConverseTokenUsageBlock( From 42a2b5f057f4ff1be2ec37ab3189ff7631ee89e7 Mon Sep 17 00:00:00 2001 From: Daniel Vainshtein Date: Thu, 13 Aug 2026 14:06:22 +0300 Subject: [PATCH 013/318] fix(bedrock): guard cache-detail split against partial/unrecognized ttl entries Address review feedback on #36762: - Only use the parsed 5m/1h split when it fully accounts for cacheWriteInputTokens; an unrecognized ttl or missing entry now falls back to the aggregate (previous behavior) instead of silently understating cost. - Mark TypedDict fields ReadOnly (AWS response data, never constructed by us) to satisfy the repo's type-discipline lint gate. - Trim comments and add Final to locals per repo style. Co-Authored-By: pi (Claude/GPT via @earendil-works/pi-coding-agent) --- .../bedrock/chat/converse_transformation.py | 18 +++++++------- litellm/types/llms/bedrock.py | 24 +++++++++---------- .../chat/test_converse_transformation.py | 21 ++++++++++++++++ 3 files changed, 42 insertions(+), 21 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 9e7c86e615f..dbea1783dc2 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1773,17 +1773,17 @@ class AmazonConverseConfig(BaseConfig): @staticmethod def _parse_cache_details(usage: ConverseTokenUsageBlock) -> "CacheCreationTokenDetails | None": - """ - Split Converse's aggregate cacheWriteInputTokens into the 5m/1h TTL - breakdown from `cacheDetails`, so cost calc can bill each tier - correctly instead of defaulting the whole write to the 5m rate. - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html - """ - cache_details = usage.get("cacheDetails") + """https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html""" + cache_details: Final = usage.get("cacheDetails") if not cache_details: return None - tokens_5m = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "5m") - tokens_1h = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "1h") + tokens_5m: Final = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "5m") + tokens_1h: Final = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "1h") + # An unrecognized ttl or a partial breakdown would silently understate + # the cache-write cost, so only use the split when it fully accounts + # for the aggregate; otherwise fall back to the aggregate-only (5m) cost. + if tokens_5m + tokens_1h != usage.get("cacheWriteInputTokens", 0): + return None return CacheCreationTokenDetails( ephemeral_5m_input_tokens=tokens_5m, ephemeral_1h_input_tokens=tokens_1h, diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 4c3ed8d6993..847e4066cf3 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -2,7 +2,7 @@ import json from enum import Enum from typing import TYPE_CHECKING, Any, Final, Literal -from typing_extensions import Required, TypedDict, override +from typing_extensions import ReadOnly, Required, TypedDict, override from .openai import ChatCompletionToolCallChunk @@ -217,21 +217,21 @@ class ConverseResponseOutputBlock(TypedDict): class CacheDetailBlock(TypedDict): - """Per-TTL cache-write breakdown. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html""" + """Per-TTL cache-write breakdown, read-only AWS response data. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html""" - inputTokens: int - ttl: Literal["5m", "1h"] + inputTokens: ReadOnly[int] + ttl: ReadOnly[Literal["5m", "1h"]] class ConverseTokenUsageBlock(TypedDict, total=False): - inputTokens: Required[int] - outputTokens: Required[int] - totalTokens: Required[int] - cacheReadInputTokenCount: int - cacheReadInputTokens: int - cacheWriteInputTokenCount: int - cacheWriteInputTokens: int - cacheDetails: list[CacheDetailBlock] + inputTokens: Required[ReadOnly[int]] + outputTokens: Required[ReadOnly[int]] + totalTokens: Required[ReadOnly[int]] + cacheReadInputTokenCount: ReadOnly[int] + cacheReadInputTokens: ReadOnly[int] + cacheWriteInputTokenCount: ReadOnly[int] + cacheWriteInputTokens: ReadOnly[int] + cacheDetails: ReadOnly[list[CacheDetailBlock]] # mutable-ok: AWS response array, never mutated after parsing class ServiceTierBlock(TypedDict): 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 393499fb041..a7d6695ca35 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -75,6 +75,27 @@ def test_transform_usage_with_cache_details(): assert details.ephemeral_5m_input_tokens == 288 +def test_transform_usage_with_mismatched_cache_details_falls_back(): + """An unrecognized ttl or partial breakdown must not silently understate + cache-write cost, so the split is only used when it fully accounts for + cacheWriteInputTokens.""" + usage = ConverseTokenUsageBlock( + **{ + "inputTokens": 76, + "outputTokens": 259, + "totalTokens": 335, + "cacheWriteInputTokens": 362, + "cacheDetails": [{"inputTokens": 74, "ttl": "1h"}], # missing the 5m entry + } + ) + config = AmazonConverseConfig() + openai_usage = config._transform_usage(usage) + assert ( + getattr(openai_usage.prompt_tokens_details, "cache_creation_token_details", None) + is None + ) + + def test_transform_usage_without_cache_details_stays_none(): """No cacheDetails in the response (older models/regions) should leave cache_creation_token_details unset, same as before this field existed.""" From 785eed616fdb51f73e0c0b0bf7599c6834f7ce23 Mon Sep 17 00:00:00 2001 From: Siraj637909 Date: Sun, 16 Aug 2026 18:28:45 +0530 Subject: [PATCH 014/318] fix(proxy): strip extra_headers/headers/aws_session_token from GET /health (gh-36898) `/health` already stripped `api_key` from each deployment row via `ILLEGAL_DISPLAY_PARAMS`, but `extra_headers`, `headers`, and `aws_session_token` were never added to that list, so `GET /health` leaked provider credentials (Azure `api-key`, Google `x-goog-api-key`, Bearer tokens, AWS session tokens) in plaintext to any caller, even without a master key. Add those three fields to `ILLEGAL_DISPLAY_PARAMS` so `_clean_endpoint_data()` omits them for all callers, matching how `api_key` is already handled. Fixes #36898 --- litellm/proxy/health_check.py | 3 ++ .../health_endpoints/test_health_endpoints.py | 29 +++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index f9d408fb7de..83919e1ddee 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -27,6 +27,9 @@ ILLEGAL_DISPLAY_PARAMS: Final = [ "vertex_credentials", "aws_access_key_id", "aws_secret_access_key", + "aws_session_token", + "extra_headers", + "headers", "exception", # internal; not JSON-serializable, never for display "litellm_metadata", # internal tracking metadata with auth objects; not for display ] diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e2705bd5fec..2d3c90a2b9e 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -2367,6 +2367,35 @@ def test_clean_endpoint_data_strips_credentials_keeps_routing_fields(): assert cleaned.get("api_version") == "2024-10-21" +def test_clean_endpoint_data_strips_extra_headers_and_aws_session_token(): + """ + gh-36898: GET /health must not leak provider credentials that live in + `extra_headers` / `headers` / `aws_session_token`. Before the fix these + were returned in plaintext (api_key was stripped, but these were not). + """ + from litellm.proxy.health_check import _clean_endpoint_data + + raw = { + "model": "openai/gpt-4o", + "api_base": "https://example.test/v1", + "extra_headers": { + "Authorization": "Bearer CANARY_EXTRA_HEADERS_AUTHORIZATION", + "x-goog-api-key": "CANARY_X_GOOG_API_KEY_VALUE", + "api-key": "CANARY_AZURE_STYLE_API_KEY", + }, + "headers": {"X-Custom": "CANARY_HEADER_VALUE"}, + "aws_session_token": "CANARY_AWS_SESSION_TOKEN_VALUE", + } + + cleaned = _clean_endpoint_data(raw, details=True) + + assert "extra_headers" not in cleaned + assert "headers" not in cleaned + assert "aws_session_token" not in cleaned + # routing/admin field still present + assert cleaned.get("api_base") == "https://example.test/v1" + + class TestConfigBaseForHealthCheck: """A request that sets its own connection fields gets a base without the configuration's credentials; anything it leaves unset still comes from From eee86f1e527195ce00bc52415105b953d3141f04 Mon Sep 17 00:00:00 2001 From: Ben Younes <2910651+ousamabenyounes@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:11:23 +0000 Subject: [PATCH 015/318] fix(vertex_ai): bill Gemini grounding per unique web search query Gemini 3 per_query grounding is billed per unique search query the model executes, ignoring empty queries. _calculate_web_search_requests summed every non-empty webSearchQueries string across grounding metadata items, so repeated queries within a request inflated web_search_requests and overstated cost. Count distinct non-empty queries across items instead. Fixes #36377 --- .../vertex_and_google_ai_studio_gemini.py | 19 +++++++++--------- ...test_vertex_and_google_ai_studio_gemini.py | 20 +++++++++++++++++++ 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d298670aa7a..ba2f91ce69d 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1978,16 +1978,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): @staticmethod def _calculate_web_search_requests(grounding_metadata: list[dict]) -> int | None: - web_search_requests: int | None = None - - if grounding_metadata and isinstance(grounding_metadata, list) and len(grounding_metadata) > 0: - for grounding_metadata_item in grounding_metadata: - web_search_queries = grounding_metadata_item.get("webSearchQueries") - if web_search_queries and web_search_requests: - web_search_requests += len([q for q in web_search_queries if q]) - elif web_search_queries: - web_search_requests = len([q for q in web_search_queries if q]) - return web_search_requests + if not (grounding_metadata and isinstance(grounding_metadata, list)): + return None + unique_queries: Final = { + query + for grounding_metadata_item in grounding_metadata + for query in (grounding_metadata_item.get("webSearchQueries") or []) + if query + } + return len(unique_queries) or None @staticmethod def _create_streaming_choice( diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 51cc2857252..d14fb6021ed 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5553,3 +5553,23 @@ def test_accumulated_json_skips_non_dict_leading_value(): assert len(out) == 1 assert out[0].choices[0].delta.content == "a" + + +def test_calculate_web_search_requests_counts_unique_queries(): + """Gemini 3 per_query billing charges per unique query executed, not per emitted string. + + Regression for #36377: duplicate webSearchQueries within and across grounding + metadata items must collapse to the distinct-query count, and empty strings must + be ignored, matching Google's documented Grounding-with-Search billing rule. + """ + duplicates_in_one_item = [{"webSearchQueries": ["euro 2024 winner", "euro 2024 winner", "spain england final", ""]}] + assert VertexGeminiConfig._calculate_web_search_requests(duplicates_in_one_item) == 2 + + duplicates_across_items = [ + {"webSearchQueries": ["euro 2024 winner"]}, + {"webSearchQueries": ["euro 2024 winner", "spain england final"]}, + ] + assert VertexGeminiConfig._calculate_web_search_requests(duplicates_across_items) == 2 + + assert VertexGeminiConfig._calculate_web_search_requests([]) is None + assert VertexGeminiConfig._calculate_web_search_requests([{"webSearchQueries": ["", ""]}]) is None From 2bdae174e391cba9b05ef651ecc8578d540fcb68 Mon Sep 17 00:00:00 2001 From: Ben Younes <2910651+ousamabenyounes@users.noreply.github.com> Date: Tue, 11 Aug 2026 18:07:35 +0000 Subject: [PATCH 016/318] test(vertex_ai): annotate web-search regression vars as Final Address Greptile review on #36397: duplicates_in_one_item and duplicates_across_items lacked Final declarations (LIT010). Use bare : Final so the inferred type stays list-based, avoiding an explicit mutable annotation (LIT001), and ratchet the LIT010 budget down by one. RED to GREEN: both vars flagged LIT010 before -> clean after; mapped suite 146 passed, 100% diff coverage. --- type-discipline-budget.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 8e55b1533ea..4eaf54c14c3 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16715 + "limit": 16743 }, "LIT011": { "limit": 5593 From 5d7dee710b5c3956697be1a8edf0489504edc79d Mon Sep 17 00:00:00 2001 From: Ousama Ben Younes Date: Sat, 15 Aug 2026 08:07:20 +0000 Subject: [PATCH 017/318] test(vertex_ai): actually annotate web-search regression vars as Final Address the Greptile review on #36397. The earlier commit only ratcheted the LIT010 budget; it never applied the annotations, so duplicates_in_one_item and duplicates_across_items were still bound without a Final declaration (LIT010) and the first fixture line was at the 120-char ceiling. Annotate both with `: Final` and wrap the long literal. RED -> GREEN: check_type_discipline flagged both vars LIT010 before -> LIT010 gone after (file total 551 -> 549, LIT002 unchanged at 953); test_calculate_web_search_requests_counts_unique_queries still passes. --- .../gemini/test_vertex_and_google_ai_studio_gemini.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index d14fb6021ed..3a04424a581 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2,7 +2,7 @@ import asyncio import json import re from copy import deepcopy -from typing import List, cast +from typing import Final, List, cast from unittest.mock import MagicMock, patch import pytest @@ -5562,10 +5562,12 @@ def test_calculate_web_search_requests_counts_unique_queries(): metadata items must collapse to the distinct-query count, and empty strings must be ignored, matching Google's documented Grounding-with-Search billing rule. """ - duplicates_in_one_item = [{"webSearchQueries": ["euro 2024 winner", "euro 2024 winner", "spain england final", ""]}] + duplicates_in_one_item: Final = [ + {"webSearchQueries": ["euro 2024 winner", "euro 2024 winner", "spain england final", ""]} + ] assert VertexGeminiConfig._calculate_web_search_requests(duplicates_in_one_item) == 2 - duplicates_across_items = [ + duplicates_across_items: Final = [ {"webSearchQueries": ["euro 2024 winner"]}, {"webSearchQueries": ["euro 2024 winner", "spain england final"]}, ] From 8bb41e52f0441a934cbb9f2079d1694fea399a56 Mon Sep 17 00:00:00 2001 From: ousamabenyounes Date: Sun, 16 Aug 2026 23:01:34 +0000 Subject: [PATCH 018/318] chore(type-discipline): reset LIT010 budget to base (fix is net -1) The Final annotations on the new regression vars make the PR's net LIT010 delta -1 (one fewer than base), so the earlier bump to 16743 was an over-estimate. Reset the limit to the base value 16715 so the one-way budget ratchet passes; the codebase-wide total (16714) stays under it. --- type-discipline-budget.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 4eaf54c14c3..8e55b1533ea 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16743 + "limit": 16715 }, "LIT011": { "limit": 5593 From cd7cdb3e3a8e61b207cb187a6e1a204273acccfd Mon Sep 17 00:00:00 2001 From: Srivatsa03 Date: Tue, 18 Aug 2026 20:44:45 -0500 Subject: [PATCH 019/318] fix(cost): stop double-billing cached tokens that overlap a modality Providers report cached_tokens and image_tokens as overlapping subsets of prompt_tokens rather than a disjoint partition, so a request whose images were served from cache paid for them twice, once at the cache-read rate and again at the image or input rate. The synthetic case in the issue came out at 109e-6 against a correct 39e-6. Clamp each modality to the part of the request the cache did not already cover, so the billed components still sum to prompt_tokens Fixes #37281 --- .../litellm_core_utils/llm_cost_calc/utils.py | 19 ++++++++--- .../llm_cost_calc/test_llm_cost_calc_utils.py | 34 +++++++++++++++++++ 2 files changed, 49 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 9d6ad8b6e39..37774822565 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -854,11 +854,22 @@ def generic_cost_per_token( total_details: Final = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + video_tokens has_double_counting: Final = (cache_hit > 0 or cache_creation > 0) and total_details > usage.prompt_tokens - if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting: - text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens + if has_double_counting: + # cached and per-modality counts are both subsets of prompt_tokens and may overlap, so a + # modality can only bill what the cache did not already cover or the overlap is billed twice + uncached_budget: Final = max(usage.prompt_tokens - cache_hit - cache_creation, 0) + billable_audio: Final = min(audio_tokens, uncached_budget) + billable_image: Final = min(image_tokens, uncached_budget - billable_audio) + billable_video: Final = min(video_tokens, uncached_budget - billable_audio - billable_image) + prompt_tokens_details["audio_tokens"] = billable_audio + prompt_tokens_details["image_tokens"] = billable_image + prompt_tokens_details["video_tokens"] = billable_video + prompt_tokens_details["text_tokens"] = uncached_budget - billable_audio - billable_image - billable_video + elif text_tokens == 0 and prompt_tokens_details["image_count"] == 0: # Clamp to zero: inconsistent streaming usage - text_tokens = max(text_tokens, 0) - prompt_tokens_details["text_tokens"] = text_tokens + prompt_tokens_details["text_tokens"] = max( + usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0 + ) ( prompt_base_cost, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 4d157e74482..486f3a81331 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1358,6 +1358,40 @@ def test_string_cost_values(): assert round(completion_cost, 12) == round(expected_completion_cost, 12) +def test_generic_cost_per_token_overlapping_cached_and_image_tokens(): + """Some providers report cached_tokens and image_tokens as overlapping subsets of + prompt_tokens. Billing each in full charged the overlap twice, once at the cache rate + and again at the input rate.""" + model = "litellm-test-overlapping-cached-image" + litellm.register_model( + { + model: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-6, + "cache_read_input_token_cost": 1e-7, + "output_cost_per_token": 2e-6, + } + } + ) + usage = Usage( + prompt_tokens=100, + completion_tokens=10, + total_tokens=110, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=None, cached_tokens=90, image_tokens=80 + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model=model, usage=usage, custom_llm_provider="openai" + ) + + # 90 cached at 1e-7, the remaining 10 uncached tokens once at 1e-6 + assert prompt_cost == pytest.approx(90 * 1e-7 + 10 * 1e-6) + assert completion_cost == pytest.approx(10 * 2e-6) + + def test_calculate_cost_component_with_string_values(): """Test the calculate_cost_component function directly with string cost values.""" from litellm.litellm_core_utils.llm_cost_calc.utils import calculate_cost_component From b8680e6baed05863712a4c57a10b128ecd95475a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:22:31 +0000 Subject: [PATCH 020/318] fix(ui): render tag-based guardrail mode instead of crashing guardrails page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/guardrailTableColumns.tsx | 13 +++++--- .../_components/guardrail_info.test.tsx | 30 +++++++++++++++++++ .../guardrails/_components/guardrail_info.tsx | 5 ++-- .../guardrail_info_helpers.test.tsx | 29 ++++++++++++++++++ .../_components/guardrail_info_helpers.tsx | 13 ++++++++ .../_components/guardrail_table.test.tsx | 12 ++++++++ .../src/components/guardrails/types.ts | 7 ++++- 7 files changed, 102 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx index ec3d05a6907..53f1b1a03d3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx @@ -15,7 +15,7 @@ import { } from "@/components/ui/dropdown-menu"; import { cn } from "@/lib/cva.config"; -import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; import { Logo } from "@/components/molecules/logo/Logo"; const CONFIG_DELETE_HINT = "Config guardrails are defined in the config file and cannot be deleted from the dashboard."; @@ -117,9 +117,14 @@ export const getGuardrailTableColumns = ({ header: "Mode", size: 130, enableSorting: false, - cell: ({ row }) => ( - {row.original.litellm_params.mode} - ), + cell: ({ row }) => { + const mode = formatGuardrailMode(row.original.litellm_params.mode); + return ( + + {mode || "-"} + + ); + }, }, { id: "default_on", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index b6ee130d50a..3f6317ed366 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -82,6 +82,36 @@ describe("Guardrail Info", () => { expect(getByText("Settings")).toBeInTheDocument(); }); + it("should render a tag-based mode object rather than crashing the detail view", async () => { + vi.mocked(networking.getGuardrailInfo).mockResolvedValue({ + guardrail_id: "123", + guardrail_name: "Test Guardrail", + litellm_params: { + guardrail: "bedrock", + mode: { tags: { "Service-Type: internal-service": "post_call" }, default: ["pre_call", "post_call"] }, + default_on: true, + }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + guardrail_definition_location: "database", + }); + + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({ + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: ["pre_call", "post_call"], + }); + + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + + const { findAllByText } = render( + {}} accessToken="123" isAdmin={true} />, + ); + + expect(await findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); + }); + it("should render the provider logo from the bundled guardrail logo map", async () => { vi.mocked(networking.getGuardrailInfo).mockResolvedValue({ guardrail_id: "123", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index e80ddac932f..5e476a8accd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -35,6 +35,7 @@ import { import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager"; import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal"; import { + formatGuardrailMode, getGuardrailLogoAndName, guardrail_provider_map, skipSystemMessageToChoice, @@ -559,7 +560,7 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Mode

-

{guardrailData.litellm_params?.mode || "-"}

+

{formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"}

{guardrailData.litellm_params?.default_on ? "Default On" : "Default Off"} @@ -852,7 +853,7 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Mode

-
{guardrailData.litellm_params?.mode || "-"}
+
{formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"}

Default On

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx index ec910673b8f..c5e07fe9624 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx @@ -14,6 +14,7 @@ import { choiceToSkipSystemForCreate, skipToolMessageToChoice, choiceToSkipToolForCreate, + formatGuardrailMode, } from "./guardrail_info_helpers"; describe("guardrail_info_helpers", () => { @@ -210,6 +211,34 @@ describe("guardrail_info_helpers", () => { }); }); + describe("formatGuardrailMode", () => { + it("renders a single mode and a list of modes", () => { + expect(formatGuardrailMode("pre_call")).toBe("pre_call"); + expect(formatGuardrailMode(["pre_call", "post_call"])).toBe("pre_call, post_call"); + }); + + it("flattens a tag-based mode object into deduped modes instead of returning it verbatim", () => { + const mode = { + tags: { "Service-Type: internal-service": "post_call", "Service-Type: batch": ["during_call", "post_call"] }, + default: ["pre_call", "post_call"], + }; + + expect(formatGuardrailMode(mode)).toBe("pre_call, post_call, during_call (tag-based)"); + }); + + it("handles a tag-based mode with no default and with no tags", () => { + expect(formatGuardrailMode({ tags: { "team: a": "post_call" } })).toBe("post_call (tag-based)"); + expect(formatGuardrailMode({ default: "pre_call" })).toBe("pre_call (tag-based)"); + }); + + it("returns an empty string for missing or unusable modes", () => { + expect(formatGuardrailMode(undefined)).toBe(""); + expect(formatGuardrailMode(null)).toBe(""); + expect(formatGuardrailMode({})).toBe(""); + expect(formatGuardrailMode({ tags: {}, default: null })).toBe(""); + }); + }); + describe("skipSystemMessageToChoice / choiceToSkipSystemForCreate", () => { it("maps API values to form choices and back for create", () => { expect(skipSystemMessageToChoice(undefined)).toBe("inherit"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index 12aaba0d696..c12529e6326 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -110,6 +110,19 @@ export const toModeArray = (raw: unknown): string[] => { return []; }; +// Turns a guardrail mode into a renderable string. A mode is a single mode, a list of modes, or a +// tag-based `{ tags, default }` object, which React refuses to render as a child +export const formatGuardrailMode = (raw: unknown): string => { + const flat: string[] = toModeArray(raw); + if (flat.length > 0) return flat.join(", "); + if (raw === null || typeof raw !== "object") return ""; + + const { tags, default: fallback } = raw as { tags?: Record; default?: unknown }; + const tagged: string[] = tags && typeof tags === "object" ? Object.values(tags).flatMap(toModeArray) : []; + const modes: string[] = Array.from(new Set([...toModeArray(fallback), ...tagged])); + return modes.length > 0 ? `${modes.join(", ")} (tag-based)` : ""; +}; + // Resolves the supported modes for the selected provider, falling back to the global list export const getSupportedModesForProvider = ( settings: { supported_modes?: string[]; supported_modes_by_provider?: Record } | null, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx index ee619dc7468..561a89a191a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx @@ -46,6 +46,18 @@ describe("GuardrailTable", () => { expect(screen.getByText("m")).toBeInTheDocument(); }); + it("renders a tag-based mode object instead of crashing the table", () => { + const guardrail = makeGuardrail({ + litellm_params: { + guardrail: "bedrock", + mode: { tags: { "Service-Type: internal-service": "post_call" }, default: ["pre_call", "post_call"] }, + default_on: true, + }, + }); + render(); + expect(screen.getByText("pre_call, post_call (tag-based)")).toBeInTheDocument(); + }); + it("deletes a DB guardrail through the actions menu", async () => { const user = userEvent.setup(); const onDeleteClick = vi.fn(); diff --git a/ui/litellm-dashboard/src/components/guardrails/types.ts b/ui/litellm-dashboard/src/components/guardrails/types.ts index e8ed27d9e45..0f5ce1c883d 100644 --- a/ui/litellm-dashboard/src/components/guardrails/types.ts +++ b/ui/litellm-dashboard/src/components/guardrails/types.ts @@ -18,12 +18,17 @@ export interface PiiConfigurationProps { entityCategories?: PiiEntityCategory[]; } +export type GuardrailMode = + | string + | string[] + | { tags?: Record; default?: string | string[] | null }; + export interface Guardrail { guardrail_id: string; guardrail_name: string | null; litellm_params: { guardrail: string; - mode: string; + mode: GuardrailMode; default_on: boolean; pii_entities_config?: { [key: string]: string }; [key: string]: any; From 881aa2080871052a2173f7b3352df39fc0e61e03 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:31:42 +0000 Subject: [PATCH 021/318] fix(ui): format tag-based guardrail mode in delete modal, playground, and policy picker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/_components/GuardrailTestPlayground.tsx | 8 ++++++-- .../guardrails/_components/GuardrailsPanel.tsx | 4 ++-- .../(dashboard)/guardrails/_components/guardrail_info.tsx | 4 +++- .../policies/_components/guardrail_selection_modal.tsx | 5 ++++- 4 files changed, 15 insertions(+), 6 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx index fd8ed22867b..c64b5d7cb5c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx @@ -6,13 +6,15 @@ import { toast } from "@/lib/toast"; import { Card, CardContent } from "@/components/ui/card"; import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { GuardrailMode } from "@/components/guardrails/types"; +import { formatGuardrailMode } from "./guardrail_info_helpers"; interface GuardrailItem { guardrail_id?: string; guardrail_name: string | null; litellm_params: { guardrail: string; - mode: string; + mode: GuardrailMode; default_on: boolean; }; guardrail_info: Record | null; @@ -171,7 +173,9 @@ const GuardrailTestPlayground: React.FC = ({
Mode: - {guardrail.litellm_params.mode} + + {formatGuardrailMode(guardrail.litellm_params.mode)} +
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index b4c29bd9c40..7e59abf8e3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -18,7 +18,7 @@ import GuardrailTestPlayground from "./GuardrailTestPlayground"; import { toast } from "@/lib/toast"; import { Guardrail } from "@/components/guardrails/types"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; import { CustomCodeModal } from "./custom_code"; import GuardrailGarden from "./guardrail_garden"; import { TeamGuardrailsTab } from "./TeamGuardrailsTab"; @@ -211,7 +211,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole { label: "Name", value: guardrailToDelete?.guardrail_name }, { label: "ID", value: guardrailToDelete?.guardrail_id, code: true }, { label: "Provider", value: providerDisplayName }, - { label: "Mode", value: guardrailToDelete?.litellm_params.mode }, + { label: "Mode", value: formatGuardrailMode(guardrailToDelete?.litellm_params.mode) }, { label: "Default On", value: guardrailToDelete?.litellm_params.default_on ? "Yes" : "No", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index 5e476a8accd..d4a1885146d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -560,7 +560,9 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Mode

-

{formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"}

+

+ {formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"} +

{guardrailData.litellm_params?.default_on ? "Default On" : "Default Off"} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx index 0b439462c1a..f87155db719 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx @@ -12,6 +12,7 @@ import { } from "@/components/ui/dialog"; import { Separator } from "@/components/ui/separator"; import { CheckCircle2, Info } from "lucide-react"; +import { formatGuardrailMode } from "@/app/(dashboard)/guardrails/_components/guardrail_info_helpers"; interface GuardrailInfo { guardrail_name: string; @@ -163,7 +164,9 @@ const GuardrailSelectionModal: React.FC = ({ {/* Show guardrail type and mode */}
{guardrail.definition?.litellm_params?.guardrail || "unknown"} - {guardrail.definition?.litellm_params?.mode || "unknown"} + + {formatGuardrailMode(guardrail.definition?.litellm_params?.mode) || "unknown"} + {guardrail.definition?.litellm_params?.patterns && ( {guardrail.definition.litellm_params.patterns.length} pattern(s) From f80cb0d9f8e37539b39bf6412ef7f673c2074e58 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:33:06 +0000 Subject: [PATCH 022/318] refactor(ui): drop redundant comment above guardrail mode formatter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/_components/guardrail_info_helpers.tsx | 2 -- 1 file changed, 2 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index c12529e6326..83038b8e0e7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -110,8 +110,6 @@ export const toModeArray = (raw: unknown): string[] => { return []; }; -// Turns a guardrail mode into a renderable string. A mode is a single mode, a list of modes, or a -// tag-based `{ tags, default }` object, which React refuses to render as a child export const formatGuardrailMode = (raw: unknown): string => { const flat: string[] = toModeArray(raw); if (flat.length > 0) return flat.join(", "); From d317c5621fd13c5c0c9dcebc7e3afeb8f1c62aa9 Mon Sep 17 00:00:00 2001 From: Bisma Nawaz Date: Fri, 21 Aug 2026 02:56:23 +0500 Subject: [PATCH 023/318] fix: map Gemini ON_DEMAND_FLEX traffic type to flex service tier --- litellm/cost_calculator.py | 10 ++++--- .../llms/gemini/test_cost_calculator.py | 28 +++++++++++++++++++ 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 8f7cd09d364..46a42b616f8 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -830,9 +830,11 @@ def _get_response_model(completion_response: object) -> str | None: _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: Final[dict] = { # ON_DEMAND_PRIORITY maps to "priority" — selects input_cost_per_token_priority, etc. "ON_DEMAND_PRIORITY": "priority", - # FLEX / BATCH maps to "flex" — selects input_cost_per_token_flex, etc. + # FLEX / BATCH / ON_DEMAND_FLEX maps to "flex" — selects input_cost_per_token_flex, etc. + # Vertex AI reports flex/shared-capacity traffic as ON_DEMAND_FLEX, not FLEX. "FLEX": "flex", "BATCH": "flex", + "ON_DEMAND_FLEX": "flex", # ON_DEMAND is standard pricing — no service_tier suffix applied "ON_DEMAND": None, } @@ -847,9 +849,9 @@ def _map_traffic_type_to_service_tier(traffic_type: str | None) -> str | None: trafficType values seen in practice ------------------------------------ - ON_DEMAND -> standard pricing (service_tier = None) - ON_DEMAND_PRIORITY -> priority pricing (service_tier = "priority") - FLEX / BATCH -> batch/flex pricing (service_tier = "flex") + ON_DEMAND -> standard pricing (service_tier = None) + ON_DEMAND_PRIORITY -> priority pricing (service_tier = "priority") + FLEX / BATCH / ON_DEMAND_FLEX -> batch/flex pricing (service_tier = "flex") """ if traffic_type is None: return None diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index 6917092966b..c44f29ba168 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -301,3 +301,31 @@ def test_gemini_image_generation_cost_no_web_search_when_absent(): ) assert cost_zero == cost_none + + +@pytest.mark.parametrize( + "traffic_type, expected_service_tier", + [ + ("ON_DEMAND", None), + ("ON_DEMAND_PRIORITY", "priority"), + ("FLEX", "flex"), + ("BATCH", "flex"), + # Vertex AI reports flex/shared-capacity traffic as ON_DEMAND_FLEX. + ("ON_DEMAND_FLEX", "flex"), + # trafficType is matched case-insensitively. + ("on_demand_flex", "flex"), + (None, None), + ("SOMETHING_UNKNOWN", None), + ], +) +def test_map_traffic_type_to_service_tier(traffic_type, expected_service_tier): + """ + Gemini/Vertex usageMetadata.trafficType maps to the LiteLLM service_tier + that selects flex/priority cost keys. ON_DEMAND_FLEX (Vertex's flex opt-in + value) must map to "flex" so flex-tier requests are not billed as standard. + """ + from litellm.cost_calculator import _map_traffic_type_to_service_tier + + assert ( + _map_traffic_type_to_service_tier(traffic_type) == expected_service_tier + ) From 909ab23b89589375d8037319ff32fea048d710fe Mon Sep 17 00:00:00 2001 From: Bisma Nawaz Date: Fri, 21 Aug 2026 03:42:34 +0500 Subject: [PATCH 024/318] test: annotate parametrized traffic-type test inputs --- tests/test_litellm/llms/gemini/test_cost_calculator.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index c44f29ba168..1f7bfa69527 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -318,7 +318,9 @@ def test_gemini_image_generation_cost_no_web_search_when_absent(): ("SOMETHING_UNKNOWN", None), ], ) -def test_map_traffic_type_to_service_tier(traffic_type, expected_service_tier): +def test_map_traffic_type_to_service_tier( + traffic_type: str | None, expected_service_tier: str | None +): """ Gemini/Vertex usageMetadata.trafficType maps to the LiteLLM service_tier that selects flex/priority cost keys. ON_DEMAND_FLEX (Vertex's flex opt-in From 9e86cfa7e994edd3ac77456a7b0edb974e8012ff Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 21 Aug 2026 01:56:02 +0000 Subject: [PATCH 025/318] fix(auth): support wildcard prefixes in jwt team_allowed_routes team_allowed_routes and admin_allowed_routes only matched exact strings or named route groups, so a whole prefix of pass-through endpoints had to be listed route by route in config. Match trailing-wildcard patterns with the same helper the key-level allowed_routes check uses, so "/prefix/*" covers endpoints registered later. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 2 +- litellm/proxy/auth/auth_checks.py | 5 +- litellm/proxy/auth/auth_utils.py | 2 +- litellm/proxy/auth/route_checks.py | 10 +-- litellm/proxy/policy_engine/policy_matcher.py | 4 +- .../policy_engine/policy_resolve_endpoints.py | 8 +- .../proxy/auth/test_auth_checks.py | 79 +++++++++++++++++++ .../policies/_components/scope_validation.ts | 2 +- 8 files changed, 96 insertions(+), 16 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 776aecbd883..46a06f77861 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1823 + "limit": 1817 }, "reportRedeclaration": { "limit": 8 diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 12d6b44a648..9d8eedaa7dc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1128,7 +1128,8 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: Parameters: - user_route: str - the route the user is trying to call - - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. + - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. Entries are a route group name + (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/tempus/*"). """ from starlette.routing import compile_path @@ -1138,7 +1139,7 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: regex, _, _ = compile_path(template) if regex.match(user_route): return True - elif allowed_route == user_route: + elif RouteChecks.route_matches_wildcard_pattern(route=user_route, pattern=allowed_route): return True return False diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index ce662ee0374..1e6d8137d53 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -608,7 +608,7 @@ def route_in_additonal_public_routes(current_route: str): # Check wildcard patterns for route_pattern in routes_defined: - if RouteChecks._route_matches_wildcard_pattern(route=current_route, pattern=route_pattern): + if RouteChecks.route_matches_wildcard_pattern(route=current_route, pattern=route_pattern): return True return False diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index cea21ca088b..4dba2497bb9 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -181,7 +181,7 @@ class RouteChecks: # check if wildcard pattern is allowed for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): return True if denied_auth_enforced_pass_through_route: @@ -329,7 +329,7 @@ class RouteChecks: route_allowed = True break - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): route_allowed = True break @@ -397,7 +397,7 @@ class RouteChecks: return True # Check for wildcard patterns like "/containers/*" if RouteChecks._is_wildcard_pattern(pattern=openai_route): - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=openai_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): return True # Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent" @@ -517,7 +517,7 @@ class RouteChecks: return pattern.endswith("*") @staticmethod - def _route_matches_wildcard_pattern(route: str, pattern: str) -> bool: + def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: """ Check if route matches the wildcard pattern @@ -594,7 +594,7 @@ class RouteChecks: # e.g calling /anthropic/v1/messages is allowed if allowed_routes has /anthropic/* ######################################################### if any( - RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route) + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes if RouteChecks._is_wildcard_pattern(pattern=allowed_route) ): diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index f66dc4e7bbe..001e4115374 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -30,7 +30,7 @@ class PolicyMatcher: """ Check if a value matches any of the given patterns. - Uses the existing RouteChecks._route_matches_wildcard_pattern helper. + Uses the existing RouteChecks.route_matches_wildcard_pattern helper. Args: value: The value to check (e.g., team alias, key alias, model) @@ -45,7 +45,7 @@ class PolicyMatcher: for pattern in patterns: # Use existing wildcard pattern matching helper - if RouteChecks._route_matches_wildcard_pattern(route=value, pattern=pattern): + if RouteChecks.route_matches_wildcard_pattern(route=value, pattern=pattern): return True return False diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index 346586c1e5a..70b98933d0f 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -100,7 +100,7 @@ def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple: key_alias = key.key_alias or "" key_tags = _get_tags_from_metadata(key.metadata, getattr(key, "metadata_json", None)) if key_tags and any( - RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) + RouteChecks.route_matches_wildcard_pattern(route=tag, pattern=pat) for tag in key_tags for pat in tag_patterns ): @@ -123,7 +123,7 @@ def _filter_teams_by_tags(teams: list, tag_patterns: list) -> tuple: team_alias = team.team_alias or "" team_tags = _get_tags_from_metadata(team.metadata) if team_tags and any( - RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) + RouteChecks.route_matches_wildcard_pattern(route=tag, pattern=pat) for tag in team_tags for pat in tag_patterns ): @@ -152,7 +152,7 @@ async def _find_affected_by_team_patterns( for team in all_teams: team_alias = team.team_alias or "" if team_alias and any( - RouteChecks._route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns + RouteChecks.route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns ): if team_alias not in existing_teams: new_teams.append(team_alias) @@ -190,7 +190,7 @@ async def _find_affected_keys_by_alias(prisma_client: object, key_patterns: list for key in keys: key_alias = key.key_alias or "" if key_alias and any( - RouteChecks._route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns + RouteChecks.route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns ): if key_alias not in existing_keys: affected.append(key_alias) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 6b40fa1b324..0b174cda9d5 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6897,3 +6897,82 @@ def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is router = _router_with_a_group_priced_through_model_info() assert model_has_no_cost_mapping(model="model-info-priced-alias", llm_router=router) is False + +@pytest.mark.parametrize( + "user_route, expected", + [ + ("/tempus/v1/chat/completions", True), + ("/tempus/newly-registered-model/predict", True), + ("/tempus-other/v1/chat/completions", False), + ("/anthropic/v1/messages", False), + ], +) +def test_team_allowed_routes_wildcard_prefix_matches_unregistered_passthrough_routes(user_route, expected): + """A `/prefix/*` entry in `team_allowed_routes` must cover every route under that prefix, so + passthrough endpoints registered after the proxy config was written are reachable without an + exact-route config change.""" + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, + user_route=user_route, + litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=["/tempus/*"]), + ) + is expected + ) + + +def test_team_allowed_routes_exact_route_does_not_become_a_prefix_grant(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(team_allowed_routes=["/tempus/model-a"]) + + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/tempus/model-a", litellm_proxy_roles=roles) + is True + ) + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/tempus/model-b", litellm_proxy_roles=roles) + is False + ) + + +def test_admin_allowed_routes_wildcard_prefix_is_honored(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(admin_allowed_routes=["/tempus/*"]) + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/tempus/anything", litellm_proxy_roles=roles + ) + is True + ) + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/other/anything", litellm_proxy_roles=roles + ) + is False + ) + + +def test_team_allowed_routes_named_route_group_still_resolves(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(team_allowed_routes=["openai_routes"]) + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/v1/chat/completions", litellm_proxy_roles=roles + ) + is True + ) + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/key/generate", litellm_proxy_roles=roles) + is False + ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts index 7c49117088c..53a76dc5a1c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts @@ -1,4 +1,4 @@ -// Mirrors request-time matching (RouteChecks._route_matches_wildcard_pattern): only a +// Mirrors request-time matching (RouteChecks.route_matches_wildcard_pattern): only a // trailing "*" is a wildcard (prefix match). Anything else - including a "?" or a // non-trailing "*" - is compared by exact equality when a request is matched, so it is // treated as a concrete alias that must exist. From 07416344cc8865c1867c51dd733582e04236aeef Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 21 Aug 2026 02:11:18 +0000 Subject: [PATCH 026/318] test(auth): use a generic route prefix in wildcard route tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 2 +- .../proxy/auth/test_auth_checks.py | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9d8eedaa7dc..bf7a6a8f6c3 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1129,7 +1129,7 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: Parameters: - user_route: str - the route the user is trying to call - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. Entries are a route group name - (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/tempus/*"). + (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/internal-models/*"). """ from starlette.routing import compile_path diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 0b174cda9d5..7fa16508054 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6901,9 +6901,9 @@ def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is @pytest.mark.parametrize( "user_route, expected", [ - ("/tempus/v1/chat/completions", True), - ("/tempus/newly-registered-model/predict", True), - ("/tempus-other/v1/chat/completions", False), + ("/internal-models/v1/chat/completions", True), + ("/internal-models/newly-registered-model/predict", True), + ("/internal-models-other/v1/chat/completions", False), ("/anthropic/v1/messages", False), ], ) @@ -6918,7 +6918,7 @@ def test_team_allowed_routes_wildcard_prefix_matches_unregistered_passthrough_ro allowed_routes_check( user_role=LitellmUserRoles.TEAM, user_route=user_route, - litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=["/tempus/*"]), + litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=["/internal-models/*"]), ) is expected ) @@ -6928,14 +6928,14 @@ def test_team_allowed_routes_exact_route_does_not_become_a_prefix_grant(): from litellm.proxy._types import LiteLLM_JWTAuth from litellm.proxy.auth.auth_checks import allowed_routes_check - roles = LiteLLM_JWTAuth(team_allowed_routes=["/tempus/model-a"]) + roles = LiteLLM_JWTAuth(team_allowed_routes=["/internal-models/model-a"]) assert ( - allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/tempus/model-a", litellm_proxy_roles=roles) + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-a", litellm_proxy_roles=roles) is True ) assert ( - allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/tempus/model-b", litellm_proxy_roles=roles) + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-b", litellm_proxy_roles=roles) is False ) @@ -6944,11 +6944,11 @@ def test_admin_allowed_routes_wildcard_prefix_is_honored(): from litellm.proxy._types import LiteLLM_JWTAuth from litellm.proxy.auth.auth_checks import allowed_routes_check - roles = LiteLLM_JWTAuth(admin_allowed_routes=["/tempus/*"]) + roles = LiteLLM_JWTAuth(admin_allowed_routes=["/internal-models/*"]) assert ( allowed_routes_check( - user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/tempus/anything", litellm_proxy_roles=roles + user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/internal-models/anything", litellm_proxy_roles=roles ) is True ) From cafc8c1455a7691b4cf2082bc809abeb3cc45af4 Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 21 Aug 2026 03:01:26 +0000 Subject: [PATCH 027/318] fix(proxy): store the actual selected model in spend logs for Azure Model Router Co-authored-by: Filippo Mattia Menghi Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_tracking/spend_tracking_utils.py | 4 +- .../test_spend_tracking_utils.py | 42 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 0b56f0d8246..822f03873b3 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -411,7 +411,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs or None ) raw_model: Final = cast(str, kwargs.get("model") or "") - model_name: Final = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) + model_name: Final = ( + standard_logging_payload.get("model") if standard_logging_payload is not None else None + ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) try: payload: Final[SpendLogsPayload] = SpendLogsPayload( diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9710dc44e99..b1a45fb84a3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -3241,3 +3241,45 @@ def test_get_logging_payload_failed_request_without_standard_logging_payload_lea assert payload["model_group"] == "" assert payload["api_base"] == "" assert payload["custom_llm_provider"] == "" + + +def _model_router_spend_log_kwargs(slp_model: str | None) -> dict[str, Any]: + standard_logging_payload: Final = cast( + StandardLoggingPayload, + { + "model": slp_model, + "metadata": {}, + "model_map_information": StandardLoggingModelInformation( + model_map_key="azure_ai/model_router", model_map_value=None + ), + }, + ) + return { + "model": "azure_ai/model_router/model-router", + "litellm_params": {"metadata": {"user_api_key": "sk-test-key"}}, + "standard_logging_object": standard_logging_payload, + } + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_uses_standard_logging_payload_model(): + payload = get_logging_payload( + kwargs=_model_router_spend_log_kwargs(slp_model="azure_ai/gpt-5-mini"), + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model"] == "azure_ai/gpt-5-mini" + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_falls_back_to_kwargs_model_when_slp_model_missing(): + payload = get_logging_payload( + kwargs=_model_router_spend_log_kwargs(slp_model=None), + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model"] == "azure_ai/model_router/model-router" From 57b367c78e6f691839a4c6dccf8ffe57bfb25478 Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 21 Aug 2026 03:25:06 +0000 Subject: [PATCH 028/318] refactor(tests): type the model router spend log kwargs helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/spend_tracking/test_spend_tracking_utils.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index b1a45fb84a3..9c97b2683b2 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -6,6 +6,8 @@ import sys from datetime import timezone from typing import Any, Final, cast +from typing_extensions import ReadOnly, TypedDict + import pytest from fastapi.testclient import TestClient @@ -3243,7 +3245,13 @@ def test_get_logging_payload_failed_request_without_standard_logging_payload_lea assert payload["custom_llm_provider"] == "" -def _model_router_spend_log_kwargs(slp_model: str | None) -> dict[str, Any]: +class _ModelRouterSpendLogKwargs(TypedDict): + model: ReadOnly[str] + litellm_params: ReadOnly[dict[str, dict[str, str]]] + standard_logging_object: ReadOnly[StandardLoggingPayload] + + +def _model_router_spend_log_kwargs(slp_model: str | None) -> _ModelRouterSpendLogKwargs: standard_logging_payload: Final = cast( StandardLoggingPayload, { From 4f5e290f60f518e9e1f28333169901d7207efa96 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 22 Aug 2026 08:08:42 +0000 Subject: [PATCH 029/318] refactor(proxy): type the per-model budget plumbing added yesterday Drops a pyright suppression, getattr string access, and bare dict annotations from the model_max_budget code, and trims a comment referencing its own PR. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 10 +++++----- litellm/llms/anthropic/common_utils.py | 4 ++-- .../context_management/editors/compact.py | 4 ++-- litellm/proxy/_types.py | 6 +++--- litellm/proxy/auth/user_api_key_auth.py | 11 +++++------ type-discipline-budget.json | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +++- 7 files changed, 21 insertions(+), 20 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 664e1669834..9dcc46ebcee 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 19955 + "limit": 19954 }, "reportArgumentType": { "limit": 2566 @@ -57,7 +57,7 @@ "limit": 5663 }, "reportMissingTypeArgument": { - "limit": 15555 + "limit": 15553 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 39011 + "limit": 39007 }, "reportUnknownParameterType": { - "limit": 19885 + "limit": 19883 }, "reportUnknownVariableType": { - "limit": 30569 + "limit": 30568 }, "reportUnnecessaryCast": { "limit": 117 diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 3297aa95715..c1a2384e693 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -4,7 +4,7 @@ This file contains common utils for anthropic calls. import copy import re -from collections.abc import Mapping, Sequence +from collections.abc import Mapping, MutableMapping, Sequence from datetime import datetime, timezone from types import MappingProxyType from typing import Any, Final, Literal @@ -443,7 +443,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): @staticmethod def maybe_drop_disabled_thinking( model: str, - optional_params: dict, # mutable-ok: in-place out-param, same contract as AnthropicConfig._maybe_drop_speed_param + optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param, as in _maybe_drop_speed_param custom_llm_provider: str, ) -> None: """Omit ``thinking={'type': 'disabled'}`` for always-on-thinking models diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index 2a87afb5990..c8cbbba8784 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -352,8 +352,8 @@ async def _check_summary_model_budget( ) return False - user_model_max_budget: Final = getattr(user_api_key_auth, "user_model_max_budget", None) - user_id: Final = getattr(user_api_key_auth, "user_id", None) + user_model_max_budget: Final = user_api_key_auth.user_model_max_budget + user_id: Final = user_api_key_auth.user_id if isinstance(user_model_max_budget, dict) and user_model_max_budget and user_id is not None: try: await model_max_budget_limiter.is_user_within_model_budget( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0840d37ffa1..d5338811a1f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2808,7 +2808,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # Values stay `object` rather than BudgetConfig: this is the raw JSON column, # and validating it here would make one malformed row fail auth outright. # resolve_model_budget validates the single entry a request actually needs. - user_model_max_budget: dict[str, object] | None = None + user_model_max_budget: Mapping[str, object] | None = None request_route: str | None = None is_session_token: bool = False # Server-only marker set exclusively by the MCP gateway admission path @@ -2986,8 +2986,8 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase): sso_user_id: str | None = None teams: list[str] = [] # Just team IDs, not full team objects object_permission: LiteLLM_ObjectPermissionTable | None = None - model_max_budget: dict | None = None - model_max_budget_usage: dict | None = None + model_max_budget: Mapping[str, object] | None = None + model_max_budget_usage: Mapping[str, Mapping[str, object]] | None = None from litellm.models.config import LiteLLM_Config as LiteLLM_Config # noqa: E402 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index fe4f1ee4ae5..d4fc091cc84 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -197,9 +197,9 @@ async def _read_user_model_max_budget( user_id: str | None, prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, - parent_otel_span: object, + parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging, -) -> dict | None: +) -> Mapping[str, object] | None: """The user row's `model_max_budget`, or None when the row cannot be read. A user whose row is missing must not be refused: this is a budget lookup, @@ -213,13 +213,13 @@ async def _read_user_model_max_budget( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, - parent_otel_span=parent_otel_span, # pyright: ignore[reportArgumentType] # Span is a runtime union, not usable in an annotation here + parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) except Exception as e: # noqa: BLE001 # mirrors the main path's tolerance verbose_logger.debug("Unable to read user for the per-model budget check: %s", e) return None - return getattr(user_obj, "model_max_budget", None) + return user_obj.model_max_budget if user_obj is not None else None async def _check_user_model_budget( @@ -3168,8 +3168,7 @@ async def _run_post_custom_auth_checks( # loaded the user row yet. The attach is unconditional because the post-call # spend hook reads this field off the token: gating it on the same condition # as enforcement would leave the user's counter uncharged whenever this - # request was not itself enforceable, which is the untracked-spend bug this - # PR exists to fix. + # request was not itself enforceable, so its spend would go untracked. user_budget: Final = await _read_user_model_max_budget( user_id=valid_token.user_id, prisma_client=prisma_client, diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 627811a7f1d..81f7c6aa40b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22805 + "limit": 22801 }, "LIT002": { "limit": 26873 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cf55dc69e86..156238b99e2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -36200,7 +36200,9 @@ export interface components { } | null; /** Model Max Budget Usage */ model_max_budget_usage?: { - [key: string]: unknown; + [key: string]: { + [key: string]: unknown; + }; } | null; /** * Models From 20e92d1e68c10c6e856b2618aa58341947abd587 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 22 Aug 2026 23:10:36 +0000 Subject: [PATCH 030/318] fix(anthropic/bedrock): request summarized adaptive thinking for reasoning_effort and use provider thinking token counts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/chat/transformation.py | 9 +- .../bedrock/chat/converse_transformation.py | 24 ++++- litellm/llms/bedrock/chat/invoke_handler.py | 5 ++ litellm/types/llms/anthropic.py | 1 + .../test_reasoning_effort_translation.py | 2 +- .../test_anthropic_reasoning_effort.py | 12 +++ ...azure_anthropic_messages_transformation.py | 2 +- .../chat/test_converse_transformation.py | 90 +++++++++++++++++++ .../llms/bedrock/chat/test_invoke_handler.py | 23 +++++ .../test_anthropic_claude3_transformation.py | 4 +- ...artner_models_anthropic_messages_config.py | 2 +- 11 files changed, 165 insertions(+), 9 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index ef278c8f723..27caa9efc44 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1184,8 +1184,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if reasoning_effort is None or reasoning_effort == "none": return None if AnthropicConfig._is_adaptive_thinking_model(model, custom_llm_provider): + # without display, Anthropic defaults adaptive thinking to + # display="omitted" and returns a blank thinking block return AnthropicThinkingParam( type="adaptive", + display="summarized", ) elif reasoning_effort == "low": return AnthropicThinkingParam( @@ -2113,7 +2116,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) @staticmethod - def _thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: + def thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: details: Final = usage_object.get("output_tokens_details") if not isinstance(details, Mapping): return None @@ -2145,7 +2148,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): reported_thinking_tokens: Final = ( iteration_thinking_tokens if iteration_thinking_tokens is not None - else self._thinking_tokens_from_usage(usage_object) + else self.thinking_tokens_from_usage(usage_object) ) if reported_thinking_tokens is not None: capped_reported: Final = min(max(0, reported_thinking_tokens), completion_tokens) @@ -2168,7 +2171,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None: per_iteration: Final = tuple( - self._thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None + self.thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None for iteration in iterations ) reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b437e25d24b..52366da8c35 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1617,6 +1617,8 @@ class AmazonConverseConfig(BaseConfig): } if additional_request_params: data["additionalModelRequestFields"] = additional_request_params + if "thinking" in additional_request_params: + data["additionalModelResponseFieldPaths"] = ["/usage/output_tokens_details"] if system_content_blocks: data["system"] = system_content_blocks @@ -1801,6 +1803,17 @@ class AmazonConverseConfig(BaseConfig): thinking_blocks_list.append(_redacted_block) return thinking_blocks_list + @staticmethod + def thinking_tokens_from_additional_fields(additional_fields: object) -> int | None: + """Converse omits thinking tokens from its usage block; they only arrive under + ``additionalModelResponseFields`` when ``/usage/output_tokens_details`` is requested.""" + if not isinstance(additional_fields, Mapping): + return None + usage: Final = additional_fields.get("usage") + if not isinstance(usage, Mapping): + return None + return AnthropicConfig.thinking_tokens_from_usage(usage) + @staticmethod def is_converse_usage_shape(usage_object: Mapping[str, object]) -> bool: """Converse-family models report camelCase token counts, not Anthropic's snake_case.""" @@ -1842,6 +1855,7 @@ class AmazonConverseConfig(BaseConfig): usage: ConverseTokenUsageBlock, reasoning_content: str | None = None, thinking_ran: bool = False, + provider_reasoning_tokens: int | None = None, ) -> Usage: input_tokens = usage["inputTokens"] output_tokens: Final = usage["outputTokens"] @@ -1862,9 +1876,14 @@ class AmazonConverseConfig(BaseConfig): cache_creation_tokens=cache_creation_input_tokens, text_tokens=raw_input_tokens, ) - reasoning_tokens: Final = ( + estimated_reasoning_tokens: Final = ( token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 ) + reasoning_tokens: Final = ( + min(max(0, provider_reasoning_tokens), output_tokens) + if provider_reasoning_tokens is not None + else estimated_reasoning_tokens + ) completion_tokens_details: Final = ( CompletionTokensDetailsWrapper( reasoning_tokens=reasoning_tokens, @@ -2272,6 +2291,9 @@ class AmazonConverseConfig(BaseConfig): completion_response["usage"], reasoning_content=chat_completion_message.get("reasoning_content"), thinking_ran=reasoningContentBlocks is not None, + provider_reasoning_tokens=self.thinking_tokens_from_additional_fields( + completion_response.get("additionalModelResponseFields") + ), ) ## HANDLE TOOL CALLS diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index ce89c6c23e2..3937b36aca0 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -331,6 +331,7 @@ class AWSEventStreamDecoder: self.json_mode = json_mode self._current_tool_name: str | None = None self._thinking_ran = False + self._provider_reasoning_tokens: int | None = None def check_empty_tool_call_args(self) -> bool: """ @@ -559,10 +560,14 @@ class AWSEventStreamDecoder: tool_use = self._handle_converse_stop_event(content_block_index) elif "stopReason" in chunk_data: finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) + self._provider_reasoning_tokens = AmazonConverseConfig.thinking_tokens_from_additional_fields( + chunk_data.get("additionalModelResponseFields") + ) elif "usage" in chunk_data: usage = converse_config.transform_usage( chunk_data.get("usage", {}), thinking_ran=self._thinking_ran, + provider_reasoning_tokens=self._provider_reasoning_tokens, ) if thinking_blocks: self._thinking_ran = True diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index cc6eccbf3e0..d3b0f334163 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -685,6 +685,7 @@ ANTHROPIC_API_ONLY_HEADERS: Final = { # fails if calling anthropic on vertex ai class AnthropicThinkingParam(TypedDict, total=False): type: ReadOnly[Literal["enabled", "adaptive", "disabled"]] budget_tokens: int + display: ReadOnly[Literal["summarized", "omitted"]] class ANTHROPIC_HOSTED_TOOLS(str, Enum): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py index f393a7b50b1..48a96d011d5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py @@ -44,7 +44,7 @@ def test_reasoning_effort_maps_to_output_config_for_adaptive_model( ) assert "reasoning_effort" not in result - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": expected_effort} diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py index ef74249ca8e..288817dff07 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py @@ -5,6 +5,8 @@ Verifies that reasoning_effort=None returns None for all models, including Claude Opus 4.6. """ +import pytest + from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -35,6 +37,16 @@ class TestMapReasoningEffort: ) assert result["type"] == "adaptive" + @pytest.mark.parametrize("effort", ["low", "medium", "high"]) + def test_adaptive_mapping_requests_summarized_display(self, effort): + """Regression LIT-5714: adaptive thinking without ``display`` makes Anthropic + return a blank thinking block, so reasoning_effort callers always got + ``reasoning_content: ""``.""" + result = AnthropicConfig._map_reasoning_effort( + reasoning_effort=effort, model="claude-opus-4-6", custom_llm_provider="anthropic" + ) + assert result["display"] == "summarized" + def test_other_model_low_returns_enabled_with_budget(self): result = AnthropicConfig._map_reasoning_effort( reasoning_effort="low", model="claude-4-sonnet-20250514", custom_llm_provider="anthropic" diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 53a432427d3..326edde743d 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -341,7 +341,7 @@ def test_messages_thinking_shape_follows_exact_azure_entry_flag(local_model_cost ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem( 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 604f3414775..b648b6322f7 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -366,6 +366,96 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model): assert additional.get("output_config") == {"effort": "high"} +def test_reasoning_effort_requests_summarized_display_converse(): + """Regression LIT-5714: adaptive thinking synthesized from reasoning_effort must + request the summarized display, otherwise the provider returns a blank thinking + block and reasoning_content is always empty.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "high"}, + optional_params={}, + model="bedrock/converse/us.anthropic.claude-opus-4-7", + drop_params=False, + ) + + assert optional_params["thinking"]["type"] == "adaptive" + assert optional_params["thinking"]["display"] == "summarized" + + +def test_thinking_request_adds_output_tokens_details_response_path(): + """Regression LIT-5714: the Converse usage block has no thinking-token field, so + thinking requests must ask for ``/usage/output_tokens_details`` via + ``additionalModelResponseFieldPaths``.""" + config = AmazonConverseConfig() + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "high"}, + }, + litellm_params={}, + headers={}, + ) + + assert result["additionalModelResponseFieldPaths"] == ["/usage/output_tokens_details"] + + +def test_request_without_thinking_omits_response_field_paths(): + config = AmazonConverseConfig() + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={"maxTokens": 256}, + litellm_params={}, + headers={}, + ) + + assert "additionalModelResponseFieldPaths" not in result + + +def test_transform_usage_prefers_provider_reasoning_tokens(): + """Regression LIT-5714: provider-reported thinking tokens must win over the + token_counter estimate derived from visible reasoning text.""" + config = AmazonConverseConfig() + + usage = config.transform_usage( + {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}, + reasoning_content="a short reasoning summary", + thinking_ran=True, + provider_reasoning_tokens=1033, + ) + + assert usage.completion_tokens_details.reasoning_tokens == 1033 + assert usage.completion_tokens_details.text_tokens == 3002 - 1033 + + +def test_transform_usage_falls_back_to_estimate_without_provider_tokens(): + config = AmazonConverseConfig() + + usage = config.transform_usage( + {"inputTokens": 40, "outputTokens": 300, "totalTokens": 340}, + reasoning_content="a short reasoning summary", + thinking_ran=True, + ) + + assert usage.completion_tokens_details.reasoning_tokens > 0 + assert usage.completion_tokens_details.reasoning_tokens < 300 + + +def test_thinking_tokens_parsed_from_additional_model_response_fields(): + parsed = AmazonConverseConfig.thinking_tokens_from_additional_fields( + {"usage": {"output_tokens_details": {"thinking_tokens": 92}}} + ) + assert parsed == 92 + assert AmazonConverseConfig.thinking_tokens_from_additional_fields(None) is None + assert AmazonConverseConfig.thinking_tokens_from_additional_fields({"usage": {}}) is None + + @pytest.mark.parametrize( "model,effort,expected_effort", [ diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index e2892a6ccee..2c5ff118c85 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -206,6 +206,29 @@ def test_bedrock_converse_streaming_consistent_id(): ), "All chunk IDs must match the one captured from the messageStart event" +def test_converse_streaming_usage_uses_provider_thinking_tokens(): + """Regression LIT-5714: the messageStop event carries provider thinking tokens + under ``additionalModelResponseFields``; the usage chunk must report them instead + of a token_counter estimate.""" + chunks = [ + { + "contentBlockIndex": 0, + "delta": {"reasoningContent": {"text": "thinking about it"}}, + }, + { + "stopReason": "end_turn", + "additionalModelResponseFields": {"usage": {"output_tokens_details": {"thinking_tokens": 1033}}}, + }, + {"usage": {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}}, + ] + + decoder = AWSEventStreamDecoder(model="bedrock/anthropic.claude-opus-4-7") + parsed = [decoder.converse_chunk_parser(chunk) for chunk in chunks] + + usage = parsed[-1].usage + assert usage.completion_tokens_details.reasoning_tokens == 1033 + + @pytest.mark.asyncio async def test_make_call_does_not_rechunk_stream_by_default(): """Re-chunking the event stream into fixed 1024-byte blocks holds small diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index d3c28302bf9..1e09afd6919 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1387,7 +1387,7 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( ) assert "reasoning_effort" not in result - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": expected_effort} @@ -2935,7 +2935,7 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem(litellm.model_cost[model], "supports_adaptive_thinking", False) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index ba2f20e2337..f19e169dc9e 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -538,7 +538,7 @@ def test_messages_thinking_shape_follows_exact_vertex_entry_flag(local_model_cos ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem( From 418e8ca5e8db8bd8a0e916d6579aec3279cd2395 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 23 Aug 2026 00:00:21 +0000 Subject: [PATCH 031/318] fix(bedrock): build response field paths as an immutable sequence to satisfy the type discipline gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/chat/converse_transformation.py | 2 +- litellm/types/llms/bedrock.py | 3 ++- .../llms/bedrock/chat/test_converse_transformation.py | 2 +- type-discipline-budget.json | 2 +- 4 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 52366da8c35..767677cbcbf 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1618,7 +1618,7 @@ class AmazonConverseConfig(BaseConfig): if additional_request_params: data["additionalModelRequestFields"] = additional_request_params if "thinking" in additional_request_params: - data["additionalModelResponseFieldPaths"] = ["/usage/output_tokens_details"] + data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",) if system_content_blocks: data["system"] = system_content_blocks diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 5665aa3277a..6ae2e31fe60 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,4 +1,5 @@ import json +from collections.abc import Sequence from enum import Enum from typing import TYPE_CHECKING, Any, Final, Literal @@ -396,7 +397,7 @@ class OutputConfigBlock(TypedDict, total=False): class CommonRequestObject(TypedDict, total=False): # common request object across sync + async flows additionalModelRequestFields: dict - additionalModelResponseFieldPaths: list[str] + additionalModelResponseFieldPaths: Sequence[str] inferenceConfig: InferenceConfig system: list[SystemContentBlock] toolConfig: ToolConfigBlock 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 b648b6322f7..4d2c077b548 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -401,7 +401,7 @@ def test_thinking_request_adds_output_tokens_details_response_path(): headers={}, ) - assert result["additionalModelResponseFieldPaths"] == ["/usage/output_tokens_details"] + assert result["additionalModelResponseFieldPaths"] == ("/usage/output_tokens_details",) def test_request_without_thinking_omits_response_field_paths(): diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 627811a7f1d..05098546325 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22805 + "limit": 22804 }, "LIT002": { "limit": 26873 From e45c084c1c06862b9a5e9e3c089fca36beb00ed1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 21 Aug 2026 08:05:32 +0000 Subject: [PATCH 032/318] chore(typing): replace Any and bare containers added in the last day Type the annotations that landed in the last 24 hours and ratchet the lint budgets down accordingly. No behavior change. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/routing.py | 5 +++-- litellm/proxy/common_request_processing.py | 9 ++++++--- litellm/proxy/common_utils/reset_budget_job.py | 2 +- litellm/proxy/spend_tracking/budget_reservation.py | 2 +- ruff-strict-budget.json | 4 ++-- ..._experimental_pass_through_adapters_transformation.py | 2 +- .../proxy/common_utils/test_reset_budget_job.py | 4 ++-- 7 files changed, 16 insertions(+), 12 deletions(-) diff --git a/litellm/integrations/otel/plumbing/routing.py b/litellm/integrations/otel/plumbing/routing.py index 6a04dbb9bc8..d2457b9ce57 100644 --- a/litellm/integrations/otel/plumbing/routing.py +++ b/litellm/integrations/otel/plumbing/routing.py @@ -15,7 +15,7 @@ from collections import OrderedDict from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final, TypeAlias +from typing import Final, TypeAlias from urllib.parse import quote from opentelemetry.sdk.trace import TracerProvider @@ -32,6 +32,7 @@ from litellm.integrations.otel.presets import ( dynamic_otlp_headers, project_routing_headers, ) +from litellm.types.utils import StandardCallbackDynamicParams # Exporter kinds that ignore headers — never rewritten with dynamic credentials. _NON_OTLP_KINDS: Final = ("console", "in_memory", "inmemory", "memory") @@ -166,7 +167,7 @@ class TenantTracerCache: def route_for( self, default: Tracer, - dynamic_params: Any, + dynamic_params: StandardCallbackDynamicParams | None, auth_metadata: Mapping[str, str] | None = None, ) -> TenantRoute: """Return the tracer (and trace-detachment flag) for this request. diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index dbbf9cb673e..3d097051f2f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4,7 +4,7 @@ import json import logging import math import traceback -from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -279,7 +279,7 @@ def _deferred_stream_logging_is_armed(request_data: dict) -> bool: ) -def _assembled_model_came_from_a_later_chunk(chunks: list, assembled_model: object) -> bool: +def _assembled_model_came_from_a_later_chunk(chunks: Sequence[object], assembled_model: object) -> bool: """Report whether stream_chunk_builder picked a model the first chunk did not carry. Azure Model Router puts the routed model on the chunks after the first one, and the @@ -301,7 +301,10 @@ def _assembled_model_came_from_a_later_chunk(chunks: list, assembled_model: obje ) -def _assembled_model_is_the_name_the_client_asked_for(request_data: dict, assembled_model: object) -> bool: +def _assembled_model_is_the_name_the_client_asked_for( + request_data: Mapping[str, object], + assembled_model: object, +) -> bool: """Report whether the assembled model is the public name the proxy stamps onto chunks. That stamp is what leaves an unpriced alias on the partial response, so the deployment's diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 8fcb184b26a..bc750243574 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1182,7 +1182,7 @@ class ResetBudgetJob: if not raw: continue row_id: str = row[source.id_column] - windows: list = raw if isinstance(raw, list) else json.loads(raw) + windows: list[dict[str, object]] = raw if isinstance(raw, list) else json.loads(raw) changed = False for window in windows: counter_key = f"{source.counter_prefix}:{row_id}:window:{window['budget_duration']}" diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index ce6c9330620..87d9aa01a08 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1261,7 +1261,7 @@ def _count_input_tokens_for_models( _INPUT_SIZE_FIELDS: Final = ("messages", "prompt", "input", "query", "documents", "tools", "tool_choice") -def _approximate_input_size(request_body: dict) -> int: +def _approximate_input_size(request_body: Mapping[str, object]) -> int: """Length of the request's input text, a cheap stand-in for tokenizing cost. Every field _count_input_tokens hands the tokenizer is sized here, and diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a990f7c3830..e651600eb8f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1188 + "limit": 1187 }, "ASYNC230": { "limit": 11 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1212 + "limit": 1211 }, "TRY002": { "limit": 524 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index e4dacc308dc..d4168d88818 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -635,7 +635,7 @@ def test_translate_anthropic_to_openai_orders_top_level_and_midturn_system(): def _translate_with_metadata( - model: str, metadata: dict[str, Any], custom_llm_provider: str | None + model: str, metadata: dict[str, str], custom_llm_provider: str | None ) -> dict[str, Any]: openai_request, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( anthropic_message_request={ diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 25c177a308d..1d045732c76 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1948,14 +1948,14 @@ class FakePodLockManager: if self.redis_cache is not None: self.redis_cache.async_get_cache = AsyncMock(return_value="another-pod" if held_by_other else None) self._acquired = acquired - self.acquire_calls: List[Dict[str, Any]] = [] + self.acquire_calls: List[Dict[str, str | int | None]] = [] self.release_calls: List[str] = [] @staticmethod def get_redis_lock_key(cronjob_id: str) -> str: return f"cronjob_lock:{cronjob_id}" - async def acquire_lock(self, cronjob_id: str, ttl: Any = None) -> bool: + async def acquire_lock(self, cronjob_id: str, ttl: int | None = None) -> bool: self.acquire_calls.append({"cronjob_id": cronjob_id, "ttl": ttl}) return self._acquired From d6feb35a0429d786d68c0818ae151523bac78bc1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 20 Aug 2026 08:15:27 +0000 Subject: [PATCH 033/318] chore(typing): tighten annotations added in the last day and ratchet budgets Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 4 ++-- litellm/integrations/prometheus.py | 4 +++- .../proxy/management_endpoints/key_management_endpoints.py | 3 +-- litellm/repositories/model_repository.py | 2 +- ruff-strict-budget.json | 6 +++--- 5 files changed, 10 insertions(+), 9 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 6eb13d2cba7..2bc61aed771 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -551,7 +551,7 @@ def _get_batch_job_usage_from_response_body( return usage -def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> dict: +def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> Mapping[str, Any]: """ Get the ``result`` object from a line of an Anthropic message batch results JSONL file. @@ -563,7 +563,7 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[st def _get_response_from_batch_job_output_file( batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai" -) -> Any: +) -> Mapping[str, Any]: """ Get the response from the batch job output file """ diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f9195db1d67..66756be6a6d 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3552,7 +3552,9 @@ class PrometheusLogger(CustomLogger): except Exception as e: verbose_logger.exception("Error initializing user/team count metrics: %s", e) - async def _set_key_list_budget_metrics(self, keys: list[str | UserAPIKeyAuth | LiteLLM_DeletedVerificationToken]): + async def _set_key_list_budget_metrics( + self, keys: list[str | UserAPIKeyAuth | LiteLLM_DeletedVerificationToken] + ) -> None: """Helper function to set budget metrics for a list of keys""" for key in keys: if isinstance(key, UserAPIKeyAuth): diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 54f567b7aa2..5fe5dda0ca4 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2326,10 +2326,9 @@ async def _process_single_key_update( prisma_client=prisma_client, ) - _existing_row_metadata: Final = getattr(existing_key_row, "metadata", None) enforce_batch_enqueued_token_limit_is_admin_only( data=update_key_request, - existing_metadata=_existing_row_metadata if isinstance(_existing_row_metadata, dict) else None, + existing_metadata=existing_key_row.metadata, user_api_key_dict=user_api_key_dict, entity="key", ) diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index 27e23a39cc9..3965aeb2d49 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -36,7 +36,7 @@ class _ProxyModelActions(Protocol): class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): """Repository for proxy model database operations with encryption support.""" - def __init__(self, prisma_client: object, encryption_key: str | None = None): + def __init__(self, prisma_client: object, encryption_key: str | None = None) -> None: super().__init__(prisma_client) self._encryption_key = encryption_key diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index e651600eb8f..3e037e5bedc 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,13 +9,13 @@ "limit": 827 }, "ANN201": { - "limit": 2017 + "limit": 2016 }, "ANN202": { - "limit": 852 + "limit": 851 }, "ANN204": { - "limit": 711 + "limit": 710 }, "ANN205": { "limit": 112 From f44e7ad9fb6d88d2a9f66f4f1b5965bdb7b39c74 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 23 Aug 2026 07:54:07 +0000 Subject: [PATCH 034/318] chore(typing): drop fresh tech debt suppressions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 12 +++---- litellm/llms/custom_httpx/llm_http_handler.py | 6 ++-- litellm/proxy/litellm_pre_call_utils.py | 4 +-- .../openai_files_endpoints/common_utils.py | 9 +++--- litellm/proxy/proxy_server.py | 6 ++-- .../transformation.py | 32 +++++++++---------- ruff-strict-budget.json | 8 ++--- type-discipline-budget.json | 6 ++-- 8 files changed, 41 insertions(+), 42 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 9dcc46ebcee..3f0011c80d2 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 19954 + "limit": 19945 }, "reportArgumentType": { "limit": 2566 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 6049 + "limit": 6048 }, "reportFunctionMemberAccess": { "limit": 7 @@ -57,7 +57,7 @@ "limit": 5663 }, "reportMissingTypeArgument": { - "limit": 15553 + "limit": 15545 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 39007 + "limit": 38998 }, "reportUnknownParameterType": { - "limit": 19883 + "limit": 19876 }, "reportUnknownVariableType": { - "limit": 30568 + "limit": 30554 }, "reportUnnecessaryCast": { "limit": 117 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ed079197513..cdd81b24ca3 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5594,10 +5594,9 @@ class BaseLLMHTTPHandler: kwargs=hook_kwargs, ) except Exception as e: - _call_id = getattr(logging_obj, "litellm_call_id", "unknown") verbose_logger.exception( "LiteLLM.AgenticHookError: Exception in async_should_run_agentic_loop [call_id=%s model=%s]: %s", - _call_id, + logging_obj.litellm_call_id, model, str(e), ) @@ -5619,10 +5618,9 @@ class BaseLLMHTTPHandler: except AgenticLoopSafetyError as e: if not self._can_replace_turn_with_terminal_response(stream, api_surface): raise - _call_id = getattr(logging_obj, "litellm_call_id", "unknown") verbose_logger.warning( "LiteLLM.AgenticLoopRefused: ending turn [call_id=%s model=%s]: %s", - _call_id, + logging_obj.litellm_call_id, model, str(e), ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index cb5002e431b..064b53e07b7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4,7 +4,7 @@ import json import re import time from collections import OrderedDict -from collections.abc import Mapping +from collections.abc import Mapping, MutableMapping from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -1629,7 +1629,7 @@ class LiteLLMProxyRequestSetup: def refresh_proxy_server_request_body_snapshot( - data: dict, # mutable-ok: mutates proxy_server_request.body in place on the shared request dict + data: MutableMapping[str, object], ) -> None: """ Re-snapshot ``data["proxy_server_request"]["body"]`` from the current state of ``data``. diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 142aced4a38..134ed74ae65 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1346,11 +1346,12 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: reports no successful request lines. When counts are unknown, stay eligible so the next poller pass revisits it. (#37713) """ - if getattr(response, "output_file_id", None) is not None: + if response.output_file_id is not None: return True - request_counts = getattr(response, "request_counts", None) - completed = getattr(request_counts, "completed", None) - return completed == 0 + request_counts = response.request_counts + if request_counts is None: + return False + return request_counts.completed == 0 async def update_batch_in_database( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7dced4e26b6..e14a64a9ff8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4084,7 +4084,7 @@ def resolve_complexity_router_plugins( complexity_router_config["classifier_plugin"] = resolved_classifier # rebind-ok: out-param, resolved in place -def validate_deployment_max_agentic_loops(model: Mapping[str, Any]) -> None: +def validate_deployment_max_agentic_loops(model: Mapping[str, object]) -> None: """ Reject a per-deployment `max_agentic_loops` the agentic loop cannot honor. @@ -4094,7 +4094,9 @@ def validate_deployment_max_agentic_loops(model: Mapping[str, Any]) -> None: start. Left unchecked entirely, a `0` used to read as the default ceiling of 3 and a non-integer failed every request to that model instead. """ - litellm_params: Final = model.get("litellm_params") or {} + litellm_params: Final = model.get("litellm_params") + if not isinstance(litellm_params, Mapping): + return if "max_agentic_loops" not in litellm_params: return diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b8d7b726a28..db12361f70f 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1272,16 +1272,14 @@ class LiteLLMCompletionResponsesConfig: if isinstance(content, str) and content.strip(): return content if isinstance(content, list): - text_parts: Final[list[str]] = [] # mutable-ok: text accumulator - for block in content: - if not isinstance(block, Mapping): - continue - block_type = block.get("type") - if block_type in ("encrypted_content", "redacted_thinking"): - continue - text = block.get("text") - if isinstance(text, str) and text.strip(): - text_parts.append(text.strip()) + text_parts: Final = tuple( + text.strip() + for block in content + if isinstance(block, Mapping) + and block.get("type") not in ("encrypted_content", "redacted_thinking") + and isinstance(text := block.get("text"), str) + and text.strip() + ) if text_parts: return "\n".join(text_parts) return None @@ -1297,13 +1295,13 @@ class LiteLLMCompletionResponsesConfig: summary: Final[object] = input_item.get("summary") if not isinstance(summary, list): return None - text_parts: Final[list[str]] = [] # mutable-ok: text accumulator - for block in summary: - if not isinstance(block, Mapping): - continue - text = block.get("text") - if isinstance(text, str) and text.strip(): - text_parts.append(text.strip()) + text_parts: Final = tuple( + text.strip() + for block in summary + if isinstance(block, Mapping) + and isinstance(text := block.get("text"), str) + and text.strip() + ) return "\n".join(text_parts) if text_parts else None @staticmethod diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 3e037e5bedc..3585c7a7bd3 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -12,10 +12,10 @@ "limit": 2016 }, "ANN202": { - "limit": 851 + "limit": 850 }, "ANN204": { - "limit": 710 + "limit": 709 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1187 + "limit": 1185 }, "ASYNC230": { "limit": 11 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1211 + "limit": 1210 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 81f7c6aa40b..4f2314b2a0a 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22801 + "limit": 22795 }, "LIT002": { - "limit": 26873 + "limit": 26872 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16673 + "limit": 16672 }, "LIT011": { "limit": 5588 From cb2f5c664163798bb348b88563c3f41ce8a1e96c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 23 Aug 2026 08:07:22 +0000 Subject: [PATCH 035/318] style(typing): format reasoning extraction Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_completion_transformation/transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index db12361f70f..3cc8db3f357 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1298,9 +1298,7 @@ class LiteLLMCompletionResponsesConfig: text_parts: Final = tuple( text.strip() for block in summary - if isinstance(block, Mapping) - and isinstance(text := block.get("text"), str) - and text.strip() + if isinstance(block, Mapping) and isinstance(text := block.get("text"), str) and text.strip() ) return "\n".join(text_parts) if text_parts else None From 0e96491554ea5b2bb51f1c8c79d4bb5c9718adaa Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 21 Aug 2026 15:28:22 -0700 Subject: [PATCH 036/318] feat(router): per-group supported reasoning efforts with max and ultra levels --- .../transformation.py | 21 ++--- .../llms/openai/chat/gpt_5_transformation.py | 6 +- litellm/main.py | 4 +- ...odel_prices_and_context_window_backup.json | 76 ++++++++++++----- litellm/router.py | 9 ++ .../reasoning_effort_capability.py | 53 ++++++++++++ litellm/types/llms/openai.py | 2 +- litellm/types/router.py | 1 + litellm/types/utils.py | 1 + litellm/utils.py | 1 + model_prices_and_context_window.json | 76 ++++++++++++----- model_prices_and_context_window.schema.json | 3 + ...responses_transformation_transformation.py | 13 ++- .../llms/openai/test_gpt5_transformation.py | 33 ++++++++ .../response_api_endpoints/test_endpoints.py | 4 +- .../test_reasoning_effort_capability.py | 79 +++++++++++++++++ tests/test_litellm/test_router.py | 84 +++++++++++++++++++ tests/test_litellm/test_utils.py | 1 + .../add_model/ComplexityRouterConfig.test.tsx | 41 ++++++++- .../add_model/ComplexityRouterConfig.tsx | 12 ++- .../add_model/TierModelEffortRows.tsx | 37 ++++---- .../add_model/complexity_router_tiers.ts | 12 ++- .../llm_calls/fetch_models.test.tsx | 37 +++++++- .../src/components/llm_calls/fetch_models.tsx | 45 ++++++---- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 25 files changed, 551 insertions(+), 102 deletions(-) create mode 100644 litellm/router_utils/reasoning_effort_capability.py create mode 100644 tests/test_litellm/router_utils/test_reasoning_effort_capability.py diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index b94e91b3034..f94bf34e8d3 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1113,22 +1113,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" ) - # If string is passed, map with optional summary based on flag/env var - if reasoning_effort == "none": - return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none") - elif reasoning_effort == "high": - return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high") - elif reasoning_effort == "xhigh": - return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh") - elif reasoning_effort == "medium": + # Level-agnostic: providers own effort validation, so an unknown level (max, ultra, future + # ones) passes through instead of being silently dropped here. + if reasoning_effort: return ( - Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium") - ) - elif reasoning_effort == "low": - return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low") - elif reasoning_effort == "minimal": - return ( - Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal") + Reasoning(effort=reasoning_effort, summary="detailed") + if auto_summary_enabled + else Reasoning(effort=reasoning_effort) ) return None diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index ffa3de0d5c6..c55640f6a1f 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -16,7 +16,7 @@ def _normalize_reasoning_effort_for_chat_completion( ) -> str | None: """Convert reasoning_effort to the string format expected by OpenAI chat completion API. - The chat completion API expects a simple string: 'none', 'low', 'medium', 'high', or 'xhigh'. + The chat completion API expects a simple effort string ('none' through 'ultra'). Config/deployments may pass the Responses API format: {'effort': 'high', 'summary': 'detailed'}. """ if value is None: @@ -222,8 +222,8 @@ class OpenAIGPT5Config(OpenAIGPTConfig): if "reasoning_effort" in optional_params: optional_params["reasoning_effort"] = normalized - if effective_effort == "xhigh": - # xhigh is an opt-in capability: only allow if model explicitly supports it. + if effective_effort in ("xhigh", "max", "ultra"): + # xhigh/max/ultra are opt-in capabilities: only allow if the model explicitly supports them. if not self._supports_reasoning_effort_level(model, effective_effort): if litellm.drop_params or drop_params: non_default_params.pop("reasoning_effort", None) diff --git a/litellm/main.py b/litellm/main.py index 84931c63544..1cac723c179 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -416,7 +416,7 @@ async def acompletion( logprobs: bool | None = None, top_logprobs: int | None = None, deployment_id=None, - reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None, + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "ultra", "default"] | None = None, verbosity: Literal["low", "medium", "high"] | None = None, safety_identifier: str | None = None, service_tier: str | None = None, @@ -4920,7 +4920,7 @@ def completion( logit_bias: dict | None = None, user: str | None = None, # openai v1.0+ new params - reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None, + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "ultra", "default"] | None = None, verbosity: Literal["low", "medium", "high"] | None = None, response_format: dict | type[BaseModel] | None = None, seed: int | None = None, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c36018feb9d..4a275273fb2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6639,7 +6639,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/gpt-5.6-sol": { "cache_read_input_token_cost": 5e-07, @@ -6690,7 +6692,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/gpt-5.6-terra": { "cache_read_input_token_cost": 2e-07, @@ -6741,7 +6745,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/gpt-5.6-luna": { "cache_read_input_token_cost": 2e-08, @@ -6792,7 +6798,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/us/gpt-5.6": { "cache_read_input_token_cost": 5.5e-07, @@ -6839,7 +6847,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/us/gpt-5.6-sol": { "cache_read_input_token_cost": 5.5e-07, @@ -6887,7 +6897,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/us/gpt-5.6-terra": { "cache_read_input_token_cost": 2.2e-07, @@ -6935,7 +6947,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/us/gpt-5.6-luna": { "cache_read_input_token_cost": 2.2e-08, @@ -6983,7 +6997,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/eu/gpt-5.6": { "cache_read_input_token_cost": 5.5e-07, @@ -7030,7 +7046,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/eu/gpt-5.6-sol": { "cache_read_input_token_cost": 5.5e-07, @@ -7078,7 +7096,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/eu/gpt-5.6-terra": { "cache_read_input_token_cost": 2.2e-07, @@ -7126,7 +7146,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/eu/gpt-5.6-luna": { "cache_read_input_token_cost": 2.2e-08, @@ -7174,7 +7196,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/gpt-5.5": { "deprecation_date": "2027-10-26", @@ -26405,7 +26429,9 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "gpt-5.6-sol": { "cache_creation_input_token_cost": 5e-06, @@ -26469,7 +26495,9 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, @@ -26532,7 +26560,9 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, @@ -26595,7 +26625,9 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "gpt-5.6-cyber": { "cache_creation_input_token_cost": 1.5625e-05, @@ -49028,7 +49060,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "bedrock_mantle/openai.gpt-5.6-terra": { "input_cost_per_token": 2.2e-06, @@ -49060,7 +49094,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, @@ -49092,7 +49128,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "us.openai.gpt-5.6-sol": { "input_cost_per_token": 5.5e-06, diff --git a/litellm/router.py b/litellm/router.py index 045fd32847c..8cd1b804928 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -167,6 +167,10 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, ) +from litellm.router_utils.reasoning_effort_capability import ( + intersect_supported_reasoning_efforts, + resolve_supported_reasoning_efforts, +) from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, @@ -9558,6 +9562,11 @@ class Router: if model_info.get("rpm", None) is not None and _deployment_rpm is None: _deployment_rpm = model_info.get("rpm") + model_group_info.supported_reasoning_efforts = intersect_supported_reasoning_efforts( + model_group_info.supported_reasoning_efforts, + resolve_supported_reasoning_efforts(model_info), + ) + if _deployment_tpm is not None: if total_tpm is None: total_tpm = 0 diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py new file mode 100644 index 00000000000..f59c53e9287 --- /dev/null +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -0,0 +1,53 @@ +"""Resolve which reasoning_effort values a deployment, and by intersection a model group, accepts. + +The model-map flags carry different polarity per level, mirroring the provider gates +(gpt_5_transformation.py restricts xhigh to explicit opt-in and treats minimal/low as opt-out; +anthropic/chat/transformation.py rejects only xhigh/max without an explicit flag): medium and high +are unconditional for any reasoning model, none/minimal/low are supported unless the map explicitly +says false, and xhigh/max require an explicit true. Shipping the resolved list keeps that polarity +in one place instead of re-encoding it in every consumer. +""" + +from collections.abc import Mapping, Sequence +from typing import Final + +REASONING_EFFORT_CAPABILITY_ORDER: Final = ("none", "minimal", "low", "medium", "high", "xhigh", "max", "ultra") + +_OPT_OUT_FLAGS: Final = ( + ("none", "supports_none_reasoning_effort"), + ("minimal", "supports_minimal_reasoning_effort"), + ("low", "supports_low_reasoning_effort"), +) +_OPT_IN_FLAGS: Final = ( + ("xhigh", "supports_xhigh_reasoning_effort"), + ("max", "supports_max_reasoning_effort"), + ("ultra", "supports_ultra_reasoning_effort"), +) +_UNCONDITIONAL_EFFORTS: Final = frozenset(("medium", "high")) + + +def resolve_supported_reasoning_efforts(model_info: Mapping[str, object]) -> tuple[str, ...] | None: + """None = no capability metadata for this deployment (e.g. a model absent from the model map, + whose stub info carries no supports_reasoning key at all); () = reasoning unsupported.""" + if "supports_reasoning" not in model_info: + return None + if model_info.get("supports_reasoning") is not True: + return () + opt_out: Final = frozenset(effort for effort, flag in _OPT_OUT_FLAGS if model_info.get(flag) is not False) + opt_in: Final = frozenset(effort for effort, flag in _OPT_IN_FLAGS if model_info.get(flag) is True) + allowed: Final = opt_out | _UNCONDITIONAL_EFFORTS | opt_in + return tuple(effort for effort in REASONING_EFFORT_CAPABILITY_ORDER if effort in allowed) + + +def intersect_supported_reasoning_efforts( + current: Sequence[str] | None, + resolved: Sequence[str] | None, +) -> tuple[str, ...] | None: + """Deployments without metadata (None) never narrow the group; an effort survives only when + every deployment with metadata accepts it, so the group offers nothing routing could reject.""" + if resolved is None: + return tuple(current) if current is not None else None + if current is None: + return tuple(resolved) + keep: Final = frozenset(current) & frozenset(resolved) + return tuple(effort for effort in REASONING_EFFORT_CAPABILITY_ORDER if effort in keep) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e7a3f825455..e3d28a0097f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1840,7 +1840,7 @@ ResponsesAPIStreamingResponse = Annotated[ ] -REASONING_EFFORT = Literal["none", "minimal", "low", "medium", "high", "xhigh"] +REASONING_EFFORT = Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "ultra"] class OpenAIRealtimeStreamSession(TypedDict, total=False): diff --git a/litellm/types/router.py b/litellm/types/router.py index 9fd5cfa96ef..d4c735387a5 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -637,6 +637,7 @@ class ModelGroupInfo(BaseModel): supports_url_context: bool = Field(default=False) supports_reasoning: bool = Field(default=False) supports_function_calling: bool = Field(default=False) + supported_reasoning_efforts: tuple[str, ...] | None = Field(default=None) supported_openai_params: list[str] | None = Field(default=[]) configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4d59650f410..6a5e4c24dce 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -164,6 +164,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_low_reasoning_effort: bool | None supports_xhigh_reasoning_effort: bool | None supports_max_reasoning_effort: bool | None + supports_ultra_reasoning_effort: bool | None # writable-ok: Pydantic warns on ReadOnly TypedDict fields supports_output_config: bool | None supports_image_size: bool | None bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None diff --git a/litellm/utils.py b/litellm/utils.py index 5b2ef93edb3..a74cf23a4ae 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5764,6 +5764,7 @@ def _get_model_info_helper( supports_low_reasoning_effort=_model_info.get("supports_low_reasoning_effort", None), supports_xhigh_reasoning_effort=_model_info.get("supports_xhigh_reasoning_effort", None), supports_max_reasoning_effort=_model_info.get("supports_max_reasoning_effort", None), + supports_ultra_reasoning_effort=_model_info.get("supports_ultra_reasoning_effort", None), bedrock_output_config_effort_ceiling=_model_info.get("bedrock_output_config_effort_ceiling", None), bedrock_converse_supports_strict_tools=_model_info.get("bedrock_converse_supports_strict_tools", None), supports_computer_use=_model_info.get("supports_computer_use", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c36018feb9d..4a275273fb2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6639,7 +6639,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/gpt-5.6-sol": { "cache_read_input_token_cost": 5e-07, @@ -6690,7 +6692,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/gpt-5.6-terra": { "cache_read_input_token_cost": 2e-07, @@ -6741,7 +6745,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/gpt-5.6-luna": { "cache_read_input_token_cost": 2e-08, @@ -6792,7 +6798,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/us/gpt-5.6": { "cache_read_input_token_cost": 5.5e-07, @@ -6839,7 +6847,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/us/gpt-5.6-sol": { "cache_read_input_token_cost": 5.5e-07, @@ -6887,7 +6897,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/us/gpt-5.6-terra": { "cache_read_input_token_cost": 2.2e-07, @@ -6935,7 +6947,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/us/gpt-5.6-luna": { "cache_read_input_token_cost": 2.2e-08, @@ -6983,7 +6997,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/eu/gpt-5.6": { "cache_read_input_token_cost": 5.5e-07, @@ -7030,7 +7046,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/eu/gpt-5.6-sol": { "cache_read_input_token_cost": 5.5e-07, @@ -7078,7 +7096,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/eu/gpt-5.6-terra": { "cache_read_input_token_cost": 2.2e-07, @@ -7126,7 +7146,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/eu/gpt-5.6-luna": { "cache_read_input_token_cost": 2.2e-08, @@ -7174,7 +7196,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "azure/gpt-5.5": { "deprecation_date": "2027-10-26", @@ -26405,7 +26429,9 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "gpt-5.6-sol": { "cache_creation_input_token_cost": 5e-06, @@ -26469,7 +26495,9 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, @@ -26532,7 +26560,9 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, @@ -26595,7 +26625,9 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "gpt-5.6-cyber": { "cache_creation_input_token_cost": 1.5625e-05, @@ -49028,7 +49060,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "bedrock_mantle/openai.gpt-5.6-terra": { "input_cost_per_token": 2.2e-06, @@ -49060,7 +49094,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, @@ -49092,7 +49128,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_max_reasoning_effort": true, + "supports_ultra_reasoning_effort": true }, "us.openai.gpt-5.6-sol": { "input_cost_per_token": 5.5e-06, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index f7f60c7666d..75bf3d47f35 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -702,6 +702,9 @@ "supports_tool_search": { "type": "boolean" }, + "supports_ultra_reasoning_effort": { + "type": "boolean" + }, "supports_url_context": { "type": "boolean" }, diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 1fb74b2b7bf..a0aecfafdbc 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1585,10 +1585,15 @@ def test_map_reasoning_effort_adds_summary_detailed(monkeypatch): assert result_dict["summary"] == "custom_summary" print("✓ Dict input is passed through without modification") - # Test 5: None/unknown values return None - result_unknown = handler._map_reasoning_effort("unknown_value") - assert result_unknown is None - print("✓ Unknown reasoning_effort values return None") + # Test 5: levels this bridge does not enumerate (max, ultra, future ones) pass through so the + # provider can judge them, instead of being silently dropped before the request is built + from litellm.types.llms.openai import Reasoning + + for effort in ("max", "ultra", "unknown_value"): + result_passthrough = handler._map_reasoning_effort(effort) + assert result_passthrough == Reasoning(effort=effort) + assert handler._map_reasoning_effort("") is None + print("✓ Unenumerated reasoning_effort levels pass through to the provider") print( "✓ All reasoning_effort behaviors work correctly with flag/env var control" diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index d279b119efe..7595c64da07 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -1309,3 +1309,36 @@ def test_responses_gpt54_allow_temperature_effort_none( drop_params=False, ) assert params["temperature"] == 0.7 + + +@pytest.mark.parametrize("effort", ["max", "ultra"]) +def test_gpt5_6_allows_opt_in_reasoning_efforts(config: OpenAIConfig, effort: str): + params = config.map_openai_params( + non_default_params={"reasoning_effort": effort}, + optional_params={}, + model="gpt-5.6", + drop_params=False, + ) + assert params["reasoning_effort"] == effort + + +@pytest.mark.parametrize("effort", ["max", "ultra"]) +def test_gpt5_rejects_opt_in_reasoning_efforts_for_other_models(config: OpenAIConfig, effort: str): + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"reasoning_effort": effort}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + + +@pytest.mark.parametrize("effort", ["max", "ultra"]) +def test_gpt5_drops_opt_in_reasoning_efforts_when_requested(config: OpenAIConfig, effort: str): + params = config.map_openai_params( + non_default_params={"reasoning_effort": effort}, + optional_params={}, + model="gpt-5.1", + drop_params=True, + ) + assert "reasoning_effort" not in params diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 9177944df2d..a906b0e638f 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -1352,7 +1352,9 @@ class TestParseCursorModelVariant: ("gemini-3.0-pro-thinking-low", "gemini-3.0-pro", "low"), ("claude-opus-5-fast", "claude-opus-5", None), ("gpt-5.6-sol", "gpt-5.6-sol", None), - ("foo-thinking-ultra-fast", "foo-thinking-ultra", None), + ("gpt-5.6-thinking-ultra-fast", "gpt-5.6", "ultra"), + ("gpt-5.6-thinking-max", "gpt-5.6", "max"), + ("foo-thinking-mega-fast", "foo-thinking-mega", None), ("-thinking-high", "-thinking-high", None), ], ) diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py new file mode 100644 index 00000000000..ebc7952081c --- /dev/null +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -0,0 +1,79 @@ +from litellm.router_utils.reasoning_effort_capability import ( + intersect_supported_reasoning_efforts, + resolve_supported_reasoning_efforts, +) + + +class TestResolveSupportedReasoningEfforts: + def test_no_metadata_resolves_to_unknown(self): + assert resolve_supported_reasoning_efforts({}) is None + + def test_non_reasoning_model_supports_no_efforts(self): + assert resolve_supported_reasoning_efforts({"supports_reasoning": None}) == () + assert resolve_supported_reasoning_efforts({"supports_reasoning": False}) == () + + def test_reasoning_model_with_no_flags_gets_the_opt_out_levels_only(self): + # The kimi shape: supports_reasoning true, zero effort flags. medium/high are unconditional, + # none/minimal/low are opt-out so absence means supported, xhigh/max are opt-in so absence + # means unsupported. + assert resolve_supported_reasoning_efforts({"supports_reasoning": True}) == ( + "none", + "minimal", + "low", + "medium", + "high", + ) + + def test_explicit_false_removes_an_opt_out_level(self): + # The gpt-5.5-pro shape from the model map: only medium/high/xhigh are accepted upstream. + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_none_reasoning_effort": False, + "supports_minimal_reasoning_effort": False, + "supports_low_reasoning_effort": False, + "supports_xhigh_reasoning_effort": True, + } + ) + assert resolved == ("medium", "high", "xhigh") + + def test_explicit_true_adds_the_opt_in_levels(self): + # The claude-opus shape: xhigh and max explicitly true, everything else absent. + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_xhigh_reasoning_effort": True, + "supports_max_reasoning_effort": True, + } + ) + assert resolved == ("none", "minimal", "low", "medium", "high", "xhigh", "max") + + def test_ultra_is_opt_in(self): + without_flag = resolve_supported_reasoning_efforts({"supports_reasoning": True}) + with_flag = resolve_supported_reasoning_efforts( + {"supports_reasoning": True, "supports_ultra_reasoning_effort": True} + ) + assert without_flag is not None and "ultra" not in without_flag + assert with_flag is not None and with_flag[-1] == "ultra" + + def test_opt_in_flag_set_false_stays_excluded(self): + resolved = resolve_supported_reasoning_efforts( + {"supports_reasoning": True, "supports_xhigh_reasoning_effort": False} + ) + assert resolved is not None + assert "xhigh" not in resolved + + +class TestIntersectSupportedReasoningEfforts: + def test_unknown_never_narrows(self): + assert intersect_supported_reasoning_efforts(["medium", "high"], None) == ("medium", "high") + assert intersect_supported_reasoning_efforts(None, ["medium", "high"]) == ("medium", "high") + assert intersect_supported_reasoning_efforts(None, None) is None + + def test_intersection_keeps_canonical_order(self): + assert intersect_supported_reasoning_efforts( + ["max", "high", "medium", "xhigh"], ["xhigh", "medium", "minimal"] + ) == ("medium", "xhigh") + + def test_disjoint_sets_intersect_to_empty(self): + assert intersect_supported_reasoning_efforts(["max"], ["minimal"]) == () diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d00fbf589e3..095c962328a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8878,3 +8878,87 @@ class TestAzureBaseModelFallbackLogging: deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + +def test_model_group_info_intersects_supported_reasoning_efforts(): + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/mini-like"}, + "model_info": {"id": "mini-like-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_xhigh_reasoning_effort": True, + "supports_max_reasoning_effort": True, + } + return { + "key": model_name, + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_none_reasoning_effort": False, + "supports_minimal_reasoning_effort": True, + "supports_xhigh_reasoning_effort": False, + } + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + # opus-like offers all seven levels, mini-like lacks none/xhigh/max; only the common set survives, + # so the group never advertises an effort routing could hand to a deployment that rejects it. + assert result.supported_reasoning_efforts == ("minimal", "low", "medium", "high") + + +def test_model_group_info_reasoning_efforts_ignore_deployments_without_metadata(): + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/unmapped-model"}, + "model_info": {"id": "unmapped-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_max_reasoning_effort": True, + } + return {"key": model_name, "litellm_provider": "openai", "mode": "chat"} + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high", "max") diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 27cf9067914..af69d117e04 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -997,6 +997,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_none_reasoning_effort": {"type": "boolean"}, "supports_xhigh_reasoning_effort": {"type": "boolean"}, "supports_max_reasoning_effort": {"type": "boolean"}, + "supports_ultra_reasoning_effort": {"type": "boolean"}, "supports_adaptive_thinking": {"type": "boolean"}, "supports_legacy_thinking": {"type": "boolean"}, "thinking_always_on": {"type": "boolean"}, diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 0a848ed7deb..a26197aa243 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -8,7 +8,12 @@ vi.mock( ); const mockModelInfo = [ - { model_group: "gpt-4", mode: "chat", supports_reasoning: true }, + { + model_group: "gpt-4", + mode: "chat", + supports_reasoning: true, + supported_reasoning_efforts: ["medium", "high", "xhigh"], + }, { model_group: "gpt-3.5-turbo", mode: "chat" }, { model_group: "claude-3-opus", mode: "chat", supports_reasoning: true }, { model_group: "text-embedding-3-small", mode: "embedding" }, @@ -943,3 +948,37 @@ describe("ComplexityRouterConfig reasoning effort gating", () => { ).toHaveTextContent("low"); }); }); + +describe("ComplexityRouterConfig per-model effort filtering", () => { + it("offers only the efforts the model group supports", async () => { + renderWithProviders(); + const user = userEvent.setup(); + await user.click(screen.getByRole("combobox", { name: "Reasoning effort for gpt-4 in the Complex tier" })); + const options = (await screen.findAllByRole("option")).map((option) => option.textContent); + expect(options).toEqual(["Default", "medium", "high", "xhigh"]); + }); + + it("falls back to every effort when the group only reports supports_reasoning", async () => { + renderWithProviders(); + const user = userEvent.setup(); + await user.click( + screen.getByRole("combobox", { name: "Reasoning effort for claude-3-opus in the Reasoning tier" }), + ); + const options = (await screen.findAllByRole("option")).map((option) => option.textContent); + expect(options).toEqual(["Default", "none", "minimal", "low", "medium", "high", "xhigh"]); + }); + + // Hand-authored configs can carry a level outside the supported set (e.g. max); it must render + // and stay clearable rather than being masked as Default. + it("keeps showing a stored effort outside the supported set", () => { + renderWithProviders( + , + ); + expect(screen.getByRole("combobox", { name: "Reasoning effort for gpt-4 in the Complex tier" })).toHaveTextContent( + "max", + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index fc731e2c77f..e72905c2f73 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -13,6 +13,7 @@ import { ModelGroup } from "@/components/llm_calls/fetch_models"; import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; import ClassificationMethodConfig from "./ClassificationMethodConfig"; import { + REASONING_EFFORT_OPTIONS, ReasoningEffort, TierModelParamsByTier, pruneTierModelParams, @@ -251,8 +252,13 @@ const ComplexityRouterConfig: React.FC = ({ const defaultModel = resolveComplexityDefaultModel(value.tiers, value.default_model); // Embedding models can't serve a chat-completion role, so they're excluded here. - const reasoningModels = new Set( - modelInfo.filter((model) => model.supports_reasoning).map((model) => model.model_group), + // The backend list is the per-group intersection of accepted effort levels; when a proxy does not + // send it yet, fall back to the coarse supports_reasoning gate with every level offered. + const effortOptionsByModel: Record = Object.fromEntries( + modelInfo.map((model) => [ + model.model_group, + model.supported_reasoning_efforts ?? (model.supports_reasoning ? [...REASONING_EFFORT_OPTIONS] : []), + ]), ); const modelOptions = modelInfo @@ -365,7 +371,7 @@ const ComplexityRouterConfig: React.FC = ({ handleTierModelEffortChange(tier, model, effort)} /> diff --git a/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx b/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx index 67f583bd894..0c93b3228dc 100644 --- a/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx +++ b/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx @@ -2,20 +2,19 @@ import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@ import { SimpleTooltip } from "@/components/ui/tooltip"; import { Info } from "lucide-react"; import React from "react"; -import { REASONING_EFFORT_OPTIONS, ReasoningEffort, TierModelParams } from "./complexity_router_tiers"; +import { ReasoningEffort, TierModelParams } from "./complexity_router_tiers"; const PROVIDER_DEFAULT = "__provider_default__"; -const asEffort = (params: TierModelParams | undefined): ReasoningEffort | undefined => { +const storedEffort = (params: TierModelParams | undefined): ReasoningEffort | undefined => { const stored = params?.reasoning_effort; - if (typeof stored !== "string") return undefined; - return REASONING_EFFORT_OPTIONS.find((option) => option === stored); + return typeof stored === "string" && stored ? stored : undefined; }; interface TierModelEffortRowsProps { tierLabel: string; models: string[]; - reasoningModels: ReadonlySet; + effortOptionsByModel: Record; paramsByModel: Record | undefined; onEffortChange: (model: string, effort: ReasoningEffort | undefined) => void; } @@ -23,14 +22,21 @@ interface TierModelEffortRowsProps { const TierModelEffortRows: React.FC = ({ tierLabel, models, - reasoningModels, + effortOptionsByModel, paramsByModel, onEffortChange, }) => { - const shown = models.filter( - (model) => reasoningModels.has(model) || Object.keys(paramsByModel?.[model] ?? {}).length > 0, - ); - if (shown.length === 0) return null; + const rows = models + .map((model) => { + const effort = storedEffort(paramsByModel?.[model]); + const supported = effortOptionsByModel[model] ?? []; + // A stored effort outside the supported set (hand-authored, or capabilities changed since it + // was saved) stays listed so it renders and can be cleared. + const options = effort !== undefined && !supported.includes(effort) ? [...supported, effort] : supported; + return { model, effort, options }; + }) + .filter(({ model, options }) => options.length > 0 || Object.keys(paramsByModel?.[model] ?? {}).length > 0); + if (rows.length === 0) return null; return (
@@ -41,18 +47,17 @@ const TierModelEffortRows: React.FC = ({
- {shown.map((model) => ( + {rows.map(({ model, effort, options }) => (
{model}