From 2b9ed94aa38f1f91c3849f3c13d43e38f76b8fc1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:10:11 -0700 Subject: [PATCH 01/15] test(vertex_ai): run auth once per Vertex pass-through request Drive /vertex_ai, /vertex-ai and /vertex_ai/discovery requests through the real proxy app with a counting custom_auth, and assert each request authenticates once and is served as the caller that auth returned. The generate-content routes fail today because the route authenticates a second time. --- .../test_llm_pass_through_endpoints.py | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 0186a33dbc4..1e9573d19a3 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -8139,6 +8139,154 @@ class TestAzureSpeechRawBodyThroughRealAuth: assert "Invalid JSON payload" in response.text +class _CountingCustomAuth: + def __init__(self) -> None: + self.api_keys_seen: list[str] = [] + + async def __call__(self, request: Request, api_key: str) -> UserAPIKeyAuth: + self.api_keys_seen.append(api_key) + return UserAPIKeyAuth(user_id=f"custom-auth-caller-{len(self.api_keys_seen)}") + + +class _ServedCallerRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.served_user_ids: list[str | None] = [] + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: + self.served_user_ids.append(user_api_key_dict.user_id) + return data + + +class TestVertexPassthroughAuthenticatesOncePerRequest: + """Each Vertex pass-through request runs the proxy's auth once and is served as the caller that auth resolved""" + + PROJECT: Final = "vertex-proj" + VIRTUAL_KEY: Final = "sk-vertex-caller" + GENERATE_CONTENT: Final = ( + "v1/projects/vertex-proj/locations/us-central1/publishers/google/models/gemini-3.8-flash:generateContent" + ) + DISCOVERY_SEARCH: Final = ( + "v1/projects/vertex-proj/locations/global/collections/default_collection/engines/support-engine" + "/servingConfigs/default_search:search" + ) + + def _post( + self, + monkeypatch: pytest.MonkeyPatch, + path: str, + headers: Mapping[str, str], + custom_auth: _CountingCustomAuth, + credential_locations: tuple[str, ...], + ) -> tuple[httpx.Response, _ServedCallerRecorder]: + from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import ( + PassthroughEndpointRouter, + ) + from litellm.proxy.proxy_server import app + + router: Final = PassthroughEndpointRouter() + for location in credential_locations: + router.add_vertex_credentials( + project_id=self.PROJECT, location=location, vertex_credentials="vertex-service-account" + ) + recorder: Final = _ServedCallerRecorder() + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router", router + ) + monkeypatch.delitem(app.dependency_overrides, user_api_key_auth, raising=False) + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + with ( + patch.multiple( # test-quality-ok: the real user_api_key_auth reads proxy_server module globals (custom auth, master_key, caches) that have no injection seam + "litellm.proxy.proxy_server", **{**_azure_speech_real_auth_attrs(), "user_custom_auth": custom_auth} + ), + patch( # test-quality-ok: the route mints its Google token through its own VertexBase, nothing injects the credential loader + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth", + return_value=(Mock(token="google-access-token"), self.PROJECT), + ), + ): + response: Final = TestClient(app).post( + path, json={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}, headers=dict(headers) + ) + return response, recorder + + @pytest.mark.parametrize("route_prefix", ["/vertex_ai", "/vertex-ai"]) + def test_generate_content_runs_custom_auth_once_and_is_served_as_its_caller( + self, monkeypatch: pytest.MonkeyPatch, route_prefix: str + ) -> None: + custom_auth: Final = _CountingCustomAuth() + vertex_url: Final = f"https://us-central1-aiplatform.googleapis.com/{self.GENERATE_CONTENT}" + with respx.mock(assert_all_called=True) as upstream: + vertex: Final = upstream.post(vertex_url).mock( + return_value=httpx.Response( + 200, + json={ + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "hello"}]}, "finishReason": "STOP"} + ], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + ) + ) + + response, recorder = self._post( + monkeypatch, + f"{route_prefix}/{self.GENERATE_CONTENT}", + {"Authorization": f"Bearer {self.VIRTUAL_KEY}"}, + custom_auth, + ("us-central1",), + ) + + assert response.status_code == 200, response.text + assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] + assert recorder.served_user_ids == ["custom-auth-caller-1"] + assert vertex.call_count == 1 + + def test_discovery_search_runs_custom_auth_once_and_is_served_as_its_caller( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + custom_auth: Final = _CountingCustomAuth() + with respx.mock(assert_all_called=True) as upstream: + discovery: Final = upstream.post(f"https://discoveryengine.googleapis.com/{self.DISCOVERY_SEARCH}").mock( + return_value=httpx.Response(200, json={"results": []}) + ) + + response, recorder = self._post( + monkeypatch, + f"/vertex_ai/discovery/{self.DISCOVERY_SEARCH}", + {"Authorization": f"Bearer {self.VIRTUAL_KEY}"}, + custom_auth, + ("global",), + ) + + assert (response.status_code, response.json()) == (200, {"results": []}) + assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] + assert recorder.served_user_ids == ["custom-auth-caller-1"] + assert discovery.call_count == 1 + + def test_x_litellm_api_key_authenticates_once_and_is_never_forwarded_without_a_vertex_credential( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + custom_auth: Final = _CountingCustomAuth() + with respx.mock(assert_all_called=False) as upstream: + google: Final = upstream.route().mock(return_value=httpx.Response(200)) + + response, _ = self._post( + monkeypatch, + f"/vertex_ai/{self.GENERATE_CONTENT}", + {"x-litellm-api-key": self.VIRTUAL_KEY}, + custom_auth, + (), + ) + + assert response.status_code == 401, response.text + assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] + assert not google.called + + class TestTypeSafePassthroughRoute: @staticmethod def _request(body: object, query_params: Mapping[str, str] | None = None) -> MagicMock: From 60e16040facbe70c6433c807fc3a8eb2aa962e9f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:59:37 -0700 Subject: [PATCH 02/15] fix(vertex_ai): authenticate Vertex pass-through requests once vertex_proxy_route already resolves the caller through Depends(user_api_key_auth), but _base_vertex_proxy_route ignored it and authenticated the request again. A custom_auth function ran twice per request, and a key with max_budget reserved its estimated cost twice, so the unreleased first reservation exhausted the budget after one or two calls. The base route now uses the caller it is given and authenticates once itself only when none is passed, which is the discovery route's case. --- .../llm_passthrough_endpoints.py | 23 ++++----- .../test_llm_pass_through_endpoints.py | 48 +++++++------------ 2 files changed, 26 insertions(+), 45 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 16fbd1de376..679ae3daa78 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2956,7 +2956,7 @@ async def _base_vertex_proxy_route( fastapi_response: FastAPI response object get_vertex_pass_through_handler: Handler for the specific Vertex AI service general_settings: Proxy general settings, read per request for the anthropic-beta values to drop - user_api_key_dict: User API key authentication dict + user_api_key_dict: The caller the route already authenticated, or None to authenticate here router_credentials: Optional vector store credentials from registry (LiteLLM_ManagedVectorStore) """ from litellm.llms.vertex_ai.common_utils import ( @@ -2970,22 +2970,15 @@ async def _base_vertex_proxy_route( encoded_endpoint = httpx.URL(endpoint).path verbose_proxy_logger.debug("requested endpoint %s", endpoint) headers: Mapping[str, str] = {} - api_key_to_use = get_litellm_virtual_key(request=request) - user_api_key_dict = await user_api_key_auth( - request=request, - api_key=api_key_to_use, + caller: Final = ( + user_api_key_dict + if user_api_key_dict is not None + else await user_api_key_auth(request=request, api_key=get_litellm_virtual_key(request=request)) ) - - if user_api_key_dict is None: - api_key_to_use = get_litellm_virtual_key(request=request) - user_api_key_dict = await user_api_key_auth( - request=request, - api_key=api_key_to_use, - ) if router_credentials is not None: await assert_user_can_access_vector_store( vector_store=router_credentials, - user_api_key_dict=user_api_key_dict, + user_api_key_dict=caller, ) vertex_project: str | None = get_vertex_project_id_from_url(endpoint) @@ -3037,7 +3030,7 @@ async def _base_vertex_proxy_route( router_credentials=router_credentials, vertex_project=vertex_project, vertex_location=vertex_location, - user_api_key_dict=user_api_key_dict, + user_api_key_dict=caller, ) base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location) @@ -3083,7 +3076,7 @@ async def _base_vertex_proxy_route( received_value: Final = await endpoint_func( request, fastapi_response, - user_api_key_dict, + caller, ) except ProxyException as e: if headers_passed_through: diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 1e9573d19a3..376b2f18d3b 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -528,17 +528,12 @@ class TestVertexAIPassThroughHandler: mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler" ) as mock_get_handler, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - ) as mock_auth, ): # Mock credentials object with necessary attributes mock_credentials = Mock() mock_credentials.token = default_credentials mock_load_auth.return_value = (mock_credentials, default_project) - mock_auth.return_value = MagicMock() # Mock the vertex handler mock_handler = Mock() @@ -556,6 +551,7 @@ class TestVertexAIPassThroughHandler: endpoint=endpoint, request=mock_request, fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), ) except Exception as e: traceback.print_exc() @@ -630,20 +626,16 @@ class TestVertexAIPassThroughHandler: mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - ) as mock_auth, ): mock_ensure_token.return_value = ("test-auth-header", test_project) mock_get_token.return_value = (test_token, "") - mock_auth.return_value = UserAPIKeyAuth(api_key="sk-test-creds") with pytest.raises(HTTPException) as exc_info: await vertex_proxy_route( endpoint=endpoint, request=mock_request, fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-creds"), ) assert exc_info.value.status_code == 401 @@ -655,8 +647,9 @@ class TestVertexAIPassThroughHandler: Critical This is how Vertex AI JS SDK will Auth to Litellm Proxy: the virtual key - arrives in x-litellm-api-key and must reach user_api_key_auth. With no Vertex - credential configured, that virtual key must not be forwarded to Google, so + arrives in x-litellm-api-key and authenticates the caller (covered through the + real auth dependency in TestVertexPassthroughAuthenticatesOncePerRequest). With no + Vertex credential configured, that virtual key must not be forwarded to Google, so the request fails with a clean 401 (LIT-5997). """ # Mock dependencies @@ -665,26 +658,21 @@ class TestVertexAIPassThroughHandler: mock_request.method = "POST" mock_response = Mock() - with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth") as mock_auth: - mock_auth.return_value = {"api_key": "test-key-123"} + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_pass_through: + mock_pass_through.return_value = AsyncMock(return_value={"status": "success"}) - with patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" - ) as mock_pass_through: - mock_pass_through.return_value = AsyncMock(return_value={"status": "success"}) + with pytest.raises(HTTPException) as exc_info: + await vertex_proxy_route( + endpoint="v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key-123"), + ) - with pytest.raises(HTTPException) as exc_info: - await vertex_proxy_route( - endpoint="v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", - request=mock_request, - fastapi_response=mock_response, - ) - - assert exc_info.value.status_code == 401 - mock_pass_through.assert_not_called() - mock_auth.assert_called_once() - call_args = mock_auth.call_args[1] - assert call_args["api_key"] == "Bearer test-key-123" + assert exc_info.value.status_code == 401 + mock_pass_through.assert_not_called() def test_vertex_passthrough_handler_multimodal_embedding_response(self): """ From f0f046db12cda4eaae4d1a9cdb03d98259b403d2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:32:47 -0700 Subject: [PATCH 03/15] fix(vertex): read the caller's key only from x-litellm-api-key or Authorization on Vertex routes The single auth check on /vertex_ai and /vertex-ai now reads the same two headers main's effective check did, so a key sent only in another key header gets 401 again instead of authenticating. --- .../llm_passthrough_endpoints.py | 16 ++++++++++++-- .../test_llm_pass_through_endpoints.py | 22 ++++++++++++++++++- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 679ae3daa78..c64f3a9db2a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast import httpx import openai -from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket +from fastapi import APIRouter, Depends, HTTPException, Request, Response, Security, WebSocket from fastapi.responses import StreamingResponse from pydantic import ConfigDict, TypeAdapter, ValidationError from starlette.websockets import WebSocketState @@ -77,6 +77,8 @@ from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports _get_bearer_token, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export + api_key_header, + custom_litellm_key_header, get_bearer_token, is_no_auth_dev_mode, user_api_key_auth, @@ -3137,6 +3139,16 @@ async def vertex_discovery_proxy_route( ) +async def _authenticate_vertex_caller( + request: Request, + authorization: Annotated[str | None, Security(api_key_header)], + litellm_api_key: Annotated[str | None, Security(custom_litellm_key_header)], +) -> UserAPIKeyAuth: + return await user_api_key_auth( + request=request, api_key=authorization or "", custom_litellm_key_header=litellm_api_key + ) + + @router.api_route( "/vertex-ai/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -3152,7 +3164,7 @@ async def vertex_proxy_route( endpoint: str, request: Request, fastapi_response: Response, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_authenticate_vertex_caller)], ): """ Call LiteLLM proxy via Vertex AI SDK. diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 376b2f18d3b..54874b51478 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -8128,11 +8128,14 @@ class TestAzureSpeechRawBodyThroughRealAuth: class _CountingCustomAuth: - def __init__(self) -> None: + def __init__(self, only_accepts: str | None = None) -> None: + self.only_accepts: Final = only_accepts self.api_keys_seen: list[str] = [] async def __call__(self, request: Request, api_key: str) -> UserAPIKeyAuth: self.api_keys_seen.append(api_key) + if self.only_accepts is not None and api_key != self.only_accepts: + raise ValueError("unknown caller") return UserAPIKeyAuth(user_id=f"custom-auth-caller-{len(self.api_keys_seen)}") @@ -8274,6 +8277,23 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] assert not google.called + @pytest.mark.parametrize("header", ["x-goog-api-key", "x-api-key", "API-Key", "Ocp-Apim-Subscription-Key"]) + def test_key_outside_litellm_headers_is_rejected_and_never_reaches_google_without_a_vertex_credential( + self, monkeypatch: pytest.MonkeyPatch, header: str + ) -> None: + custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY) + with respx.mock(assert_all_called=False) as upstream: + google: Final = upstream.route().mock(return_value=httpx.Response(200)) + + response, recorder = self._post( + monkeypatch, f"/vertex_ai/{self.GENERATE_CONTENT}", {header: self.VIRTUAL_KEY}, custom_auth, () + ) + + assert response.status_code == 401, response.text + assert self.VIRTUAL_KEY not in custom_auth.api_keys_seen + assert recorder.served_user_ids == [] + assert not google.called + class TestTypeSafePassthroughRoute: @staticmethod From 7ca83f9807e248c454c7615ae005a2ebe0496ce5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:13:23 -0700 Subject: [PATCH 04/15] fix(vertex): authenticate the discovery route through the same Vertex key check Both Vertex pass-through routes now resolve the caller once through _authenticate_vertex_caller and hand it to _base_vertex_proxy_route, which no longer re-runs auth. x-litellm-api-key accepts the same scheme prefixes on /vertex_ai/discovery as on /vertex_ai and every other proxy route. --- litellm/proxy/_lazy_openapi_snapshot.json | 25 ++++ .../llm_passthrough_endpoints.py | 38 +++--- .../test_llm_pass_through_endpoints.py | 123 +++++++++--------- .../test_vertex_passthrough_load_balancing.py | 24 +--- .../test_vector_store_tenant_guard.py | 24 ++-- 5 files changed, 122 insertions(+), 112 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 545b317055a..85dba9c1469 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -34564,6 +34564,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -34603,6 +34608,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -34642,6 +34652,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -34681,6 +34696,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -34720,6 +34740,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index c64f3a9db2a..3dccd2b68be 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -105,7 +105,6 @@ from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) -from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( stream_usage_relay, without_injected_usage_frames, @@ -2942,7 +2941,7 @@ async def _base_vertex_proxy_route( fastapi_response: Response, get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler, general_settings: Mapping[str, object], - user_api_key_dict: UserAPIKeyAuth | None = None, + user_api_key_dict: UserAPIKeyAuth, router_credentials: LiteLLM_ManagedVectorStore | None = None, ): """ @@ -2958,7 +2957,7 @@ async def _base_vertex_proxy_route( fastapi_response: FastAPI response object get_vertex_pass_through_handler: Handler for the specific Vertex AI service general_settings: Proxy general settings, read per request for the anthropic-beta values to drop - user_api_key_dict: The caller the route already authenticated, or None to authenticate here + user_api_key_dict: The caller the route authenticated router_credentials: Optional vector store credentials from registry (LiteLLM_ManagedVectorStore) """ from litellm.llms.vertex_ai.common_utils import ( @@ -2972,15 +2971,10 @@ async def _base_vertex_proxy_route( encoded_endpoint = httpx.URL(endpoint).path verbose_proxy_logger.debug("requested endpoint %s", endpoint) headers: Mapping[str, str] = {} - caller: Final = ( - user_api_key_dict - if user_api_key_dict is not None - else await user_api_key_auth(request=request, api_key=get_litellm_virtual_key(request=request)) - ) if router_credentials is not None: await assert_user_can_access_vector_store( vector_store=router_credentials, - user_api_key_dict=caller, + user_api_key_dict=user_api_key_dict, ) vertex_project: str | None = get_vertex_project_id_from_url(endpoint) @@ -3032,7 +3026,7 @@ async def _base_vertex_proxy_route( router_credentials=router_credentials, vertex_project=vertex_project, vertex_location=vertex_location, - user_api_key_dict=caller, + user_api_key_dict=user_api_key_dict, ) base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location) @@ -3078,7 +3072,7 @@ async def _base_vertex_proxy_route( received_value: Final = await endpoint_func( request, fastapi_response, - caller, + user_api_key_dict, ) except ProxyException as e: if headers_passed_through: @@ -3088,6 +3082,16 @@ async def _base_vertex_proxy_route( return received_value +async def _authenticate_vertex_caller( + request: Request, + authorization: Annotated[str | None, Security(api_key_header)], + litellm_api_key: Annotated[str | None, Security(custom_litellm_key_header)], +) -> UserAPIKeyAuth: + return await user_api_key_auth( + request=request, api_key=authorization or "", custom_litellm_key_header=litellm_api_key + ) + + @router.api_route( "/vertex_ai/discovery/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -3097,6 +3101,7 @@ async def vertex_discovery_proxy_route( endpoint: str, request: Request, fastapi_response: Response, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_authenticate_vertex_caller)], ): """ Call any vertex discovery endpoint using the proxy. @@ -3135,20 +3140,11 @@ async def vertex_discovery_proxy_route( fastapi_response=fastapi_response, get_vertex_pass_through_handler=discovery_handler, general_settings=_proxy_general_settings(), + user_api_key_dict=user_api_key_dict, router_credentials=vector_store_credentials, ) -async def _authenticate_vertex_caller( - request: Request, - authorization: Annotated[str | None, Security(api_key_header)], - litellm_api_key: Annotated[str | None, Security(custom_litellm_key_header)], -) -> UserAPIKeyAuth: - return await user_api_key_auth( - request=request, api_key=authorization or "", custom_litellm_key_header=litellm_api_key - ) - - @router.api_route( "/vertex-ai/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 54874b51478..aa26145c406 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -324,12 +324,6 @@ class TestVertexAIPassThroughHandler: mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_litellm_virtual_key" - ) as mock_get_virtual_key, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth" - ) as mock_user_auth, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler" ) as mock_get_handler, @@ -340,8 +334,6 @@ class TestVertexAIPassThroughHandler: # Setup mocks mock_load_auth.return_value = (mock_credentials, test_project) - mock_get_virtual_key.return_value = "Bearer test-key" - mock_user_auth.return_value = {"api_key": "test-key"} # Mock the vertex handler mock_handler = Mock() @@ -426,12 +418,6 @@ class TestVertexAIPassThroughHandler: mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_litellm_virtual_key" - ) as mock_get_virtual_key, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth" - ) as mock_user_auth, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler" ) as mock_get_handler, @@ -442,8 +428,6 @@ class TestVertexAIPassThroughHandler: # Setup mocks mock_load_auth.return_value = (mock_credentials, test_project) - mock_get_virtual_key.return_value = "Bearer test-key" - mock_user_auth.return_value = {"api_key": "test-key"} # Mock the vertex handler for global location mock_handler = Mock() @@ -1238,12 +1222,6 @@ class TestVertexAIDiscoveryPassThroughHandler: mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_litellm_virtual_key" - ) as mock_get_virtual_key, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth" - ) as mock_user_auth, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler" ) as mock_get_handler, @@ -1254,8 +1232,6 @@ class TestVertexAIDiscoveryPassThroughHandler: # Setup mocks mock_load_auth.return_value = (mock_credentials, test_project) - mock_get_virtual_key.return_value = "Bearer test-key" - mock_user_auth.return_value = {"api_key": "test-key"} # Mock the discovery handler mock_handler = Mock() @@ -1271,6 +1247,7 @@ class TestVertexAIDiscoveryPassThroughHandler: endpoint=endpoint, request=mock_request, fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), ) # Verify create_pass_through_route was called with correct arguments @@ -1285,36 +1262,29 @@ class TestVertexAIDiscoveryPassThroughHandler: @pytest.mark.asyncio async def test_vertex_discovery_proxy_route_api_key_auth(self): """ - The virtual key arrives in x-litellm-api-key and must reach user_api_key_auth. - With no Vertex credential configured, that virtual key must not be forwarded to + With no Vertex credential configured, the caller's virtual key must not be forwarded to Google, so the request fails with a clean 401 (LIT-5997). """ - # Mock dependencies mock_request = Mock() mock_request.headers = {"x-litellm-api-key": "test-key-123"} mock_request.method = "POST" mock_response = Mock() - with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth") as mock_auth: - mock_auth.return_value = {"api_key": "test-key-123"} + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_pass_through: + mock_pass_through.return_value = AsyncMock(return_value={"status": "success"}) - with patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" - ) as mock_pass_through: - mock_pass_through.return_value = AsyncMock(return_value={"status": "success"}) + with pytest.raises(HTTPException) as exc_info: + await vertex_discovery_proxy_route( + endpoint="v1/projects/test-project/locations/us-central1/dataStores/default/servingConfigs/default:search", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key-123"), + ) - with pytest.raises(HTTPException) as exc_info: - await vertex_discovery_proxy_route( - endpoint="v1/projects/test-project/locations/us-central1/dataStores/default/servingConfigs/default:search", - request=mock_request, - fastapi_response=mock_response, - ) - - assert exc_info.value.status_code == 401 - mock_pass_through.assert_not_called() - mock_auth.assert_called_once() - call_args = mock_auth.call_args[1] - assert call_args["api_key"] == "Bearer test-key-123" + assert exc_info.value.status_code == 401 + mock_pass_through.assert_not_called() @pytest.mark.asyncio @@ -4375,8 +4345,6 @@ class TestVertexRawPredictStreamingClassification: return_value=(mock_credentials, "test-project"), ), mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route), - mock.patch(f"{module}.get_litellm_virtual_key", return_value="Bearer test-key"), - mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=UserAPIKeyAuth(api_key="test-key"))), mock.patch(f"{module}.get_vertex_pass_through_handler", return_value=mock_handler), ): await vertex_proxy_route( @@ -8163,6 +8131,10 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: "v1/projects/vertex-proj/locations/global/collections/default_collection/engines/support-engine" "/servingConfigs/default_search:search" ) + GEMINI_ANSWER: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "hello"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + } def _post( self, @@ -8211,17 +8183,7 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: custom_auth: Final = _CountingCustomAuth() vertex_url: Final = f"https://us-central1-aiplatform.googleapis.com/{self.GENERATE_CONTENT}" with respx.mock(assert_all_called=True) as upstream: - vertex: Final = upstream.post(vertex_url).mock( - return_value=httpx.Response( - 200, - json={ - "candidates": [ - {"content": {"role": "model", "parts": [{"text": "hello"}]}, "finishReason": "STOP"} - ], - "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, - }, - ) - ) + vertex: Final = upstream.post(vertex_url).mock(return_value=httpx.Response(200, json=self.GEMINI_ANSWER)) response, recorder = self._post( monkeypatch, @@ -8258,6 +8220,51 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: assert recorder.served_user_ids == ["custom-auth-caller-1"] assert discovery.call_count == 1 + @pytest.mark.parametrize("scheme", ["bearer", "Basic"]) + @pytest.mark.parametrize( + "route, upstream_url, credential_location, answer", + [ + ( + f"/vertex_ai/{GENERATE_CONTENT}", + f"https://us-central1-aiplatform.googleapis.com/{GENERATE_CONTENT}", + "us-central1", + GEMINI_ANSWER, + ), + ( + f"/vertex_ai/discovery/{DISCOVERY_SEARCH}", + f"https://discoveryengine.googleapis.com/{DISCOVERY_SEARCH}", + "global", + {"results": []}, + ), + ], + ids=["generate_content", "discovery_search"], + ) + def test_x_litellm_api_key_with_any_auth_scheme_authenticates_once_on_every_vertex_route( + self, + monkeypatch: pytest.MonkeyPatch, + scheme: str, + route: str, + upstream_url: str, + credential_location: str, + answer: Mapping[str, object], + ) -> None: + custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY) + with respx.mock(assert_all_called=True) as upstream: + google: Final = upstream.post(upstream_url).mock(return_value=httpx.Response(200, json=answer)) + + response, recorder = self._post( + monkeypatch, + route, + {"x-litellm-api-key": f"{scheme} {self.VIRTUAL_KEY}"}, + custom_auth, + (credential_location,), + ) + + assert response.status_code == 200, response.text + assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] + assert recorder.served_user_ids == ["custom-auth-caller-1"] + assert google.call_count == 1 + def test_x_litellm_api_key_authenticates_once_and_is_never_forwarded_without_a_vertex_credential( self, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index ed14eeda6a2..c5b850ea2a9 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -66,10 +66,6 @@ async def test_vertex_passthrough_load_balancing(): patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - ) as mock_auth, ): # Setup additional mocks to avoid side effects @@ -83,7 +79,6 @@ async def test_vertex_passthrough_load_balancing(): mock_endpoint_func = AsyncMock() mock_create_route.return_value = mock_endpoint_func - mock_auth.return_value = {} # Execute await _base_vertex_proxy_route( @@ -92,6 +87,7 @@ async def test_vertex_passthrough_load_balancing(): fastapi_response=mock_response, get_vertex_pass_through_handler=mock_handler, general_settings={}, + user_api_key_dict=UserAPIKeyAuth(), ) # Verify @@ -376,10 +372,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( patch( # test-quality-ok: the upstream call is captured here, the route offers no injection point "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - patch( # test-quality-ok: the route calls auth directly rather than through Depends - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - ) as mock_auth, patch( # test-quality-ok: the route reads the request body for this, a MagicMock request has none "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn", new_callable=AsyncMock, @@ -398,7 +390,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( "global", ) mock_create_route.return_value = AsyncMock() - mock_auth.return_value = UserAPIKeyAuth(api_key="sk-litellm-secret-key") await _base_vertex_proxy_route( endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}", @@ -406,6 +397,7 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( fastapi_response=MagicMock(), get_vertex_pass_through_handler=VertexAIPassThroughHandler(), general_settings={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"), ) upstream_headers = mock_create_route.call_args.kwargs["custom_headers"] @@ -723,10 +715,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url(): patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - ) as mock_auth, ): mock_pt_router.get_vertex_credentials.return_value = MagicMock() @@ -738,7 +726,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url(): ) mock_endpoint_func = AsyncMock() mock_create_route.return_value = mock_endpoint_func - mock_auth.return_value = {} mock_handler.get_default_base_target_url.return_value = ( "https://global-aiplatform.googleapis.com" @@ -750,6 +737,7 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url(): fastapi_response=mock_response, get_vertex_pass_through_handler=mock_handler, general_settings={}, + user_api_key_dict=UserAPIKeyAuth(), ) # Verify the router was called with the custom model name (extracted from URL) @@ -826,11 +814,6 @@ async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", return_value=relay_returning_logging_kwargs, ), - patch( # test-quality-ok: the route calls auth directly rather than through Depends - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - return_value=UserAPIKeyAuth(api_key="hashed-key"), - ), ): mock_pt_router.get_vertex_credentials.return_value = MagicMock() @@ -840,6 +823,7 @@ async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment fastapi_response=MagicMock(), get_vertex_pass_through_handler=mock_handler, general_settings={}, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), ) assert logging_kwargs["litellm_params"]["metadata"]["model_info"]["id"] == "vertex-gemini-38-flash-dep" diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py index 29efc41b484..5b356c43245 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py @@ -441,19 +441,16 @@ async def test_vertex_discovery_denies_other_team_vector_store_credentials(): "team_id": "team-b", } - with patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new=AsyncMock(return_value=UserAPIKeyAuth(team_id="team-a")), - ): - with pytest.raises(HTTPException) as exc_info: - await _base_vertex_proxy_route( - endpoint="projects/p/locations/us-central1/dataStores/vs_other_team", - request=request, - fastapi_response=Response(), - get_vertex_pass_through_handler=MagicMock(), - general_settings={}, - router_credentials=vector_store_credentials, - ) + with pytest.raises(HTTPException) as exc_info: + await _base_vertex_proxy_route( + endpoint="projects/p/locations/us-central1/dataStores/vs_other_team", + request=request, + fastapi_response=Response(), + get_vertex_pass_through_handler=MagicMock(), + general_settings={}, + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + router_credentials=vector_store_credentials, + ) assert exc_info.value.status_code == 403 @@ -541,6 +538,7 @@ async def test_vertex_discovery_allows_unregistered_provider_native_datastore_id endpoint="projects/p/locations/us-central1/dataStores/vs_unknown", request=request, fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(), ) assert response == {"ok": True} From 13386b95503eb681e19daf894e8b2da28104af9f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:50:10 -0700 Subject: [PATCH 05/15] fix(proxy): size a Gemini request's budget hold from its prompt and maxOutputTokens A Gemini generateContent body's budget hold now counts its contents and systemInstruction and reads generationConfig.maxOutputTokens. Before, it held the model's whole context window, about $0.85 on gemini-3.8-flash, so a key with a small budget got 422s on back-to-back calls it could afford Also removes get_litellm_virtual_key, which nothing calls since the Vertex routes stopped using it --- .../pass_through_endpoints/common_utils.py | 18 -------- .../spend_tracking/budget_reservation.py | 2 + litellm/proxy/spend_tracking/input_tokens.py | 42 +++++++++++++++-- ...test_passthrough_endpoints_common_utils.py | 40 ---------------- .../spend_tracking/test_budget_reservation.py | 46 +++++++++++++++++++ 5 files changed, 85 insertions(+), 63 deletions(-) delete mode 100644 litellm/proxy/pass_through_endpoints/common_utils.py diff --git a/litellm/proxy/pass_through_endpoints/common_utils.py b/litellm/proxy/pass_through_endpoints/common_utils.py deleted file mode 100644 index 3ae83cbe906..00000000000 --- a/litellm/proxy/pass_through_endpoints/common_utils.py +++ /dev/null @@ -1,18 +0,0 @@ -from typing import Final - -from fastapi import Request - - -def get_litellm_virtual_key(request: Request) -> str: - """ - Extract and format API key from request headers. - Prioritizes x-litellm-api-key over Authorization header. - - - Vertex JS SDK uses `Authorization` header, we use `x-litellm-api-key` to pass litellm virtual key - - """ - litellm_api_key: Final = request.headers.get("x-litellm-api-key") - if litellm_api_key: - return f"Bearer {litellm_api_key}" - return request.headers.get("Authorization", "") diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 0b91e57f01e..cfc865a84d1 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1654,9 +1654,11 @@ _OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_outpu def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: inference_config: Final = request_body.get("inferenceConfig") + generation_config: Final = request_body.get("generationConfig") candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None, + generation_config.get("maxOutputTokens") if isinstance(generation_config, Mapping) else None, ) return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 927152296e9..a7e9981ac62 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -15,10 +15,13 @@ import asyncio import json from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final +from typing import Final, TypeAlias + +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger +from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter from litellm.responses.litellm_completion_transformation.compaction_items import ( billable_input_since_newest_compaction, holds_compaction_summary, @@ -35,7 +38,20 @@ from litellm.rust_bridge.token_counter import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 -_INPUT_SIZE_FIELDS: Final = ("messages", "prompt", "input", "query", "documents", "tools", "tool_choice") +_GEMINI_SYSTEM_INSTRUCTION_FIELDS: Final = ("systemInstruction", "system_instruction") +_GeminiContents: TypeAlias = list[dict[str, object]] | dict[str, object] +_GEMINI_CONTENTS: Final = TypeAdapter[_GeminiContents](_GeminiContents) +_INPUT_SIZE_FIELDS: Final = ( + "messages", + "prompt", + "input", + "query", + "documents", + "contents", + *_GEMINI_SYSTEM_INSTRUCTION_FIELDS, + "tools", + "tool_choice", +) def _approximate_input_size(request_body: Mapping[str, object]) -> int: @@ -48,7 +64,7 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int: async def count_input_tokens( - request_body: dict, + request_body: dict[str, object], raw_body: bytes | None, models: Sequence[str], ) -> Mapping[str, int]: @@ -74,7 +90,7 @@ async def count_input_tokens( async def _count_group( - request_body: dict, + request_body: dict[str, object], raw_body: bytes | None, tokenizer: RustTokenizer | None, models: tuple[str, ...], @@ -118,7 +134,7 @@ async def _native_counts( def _count_input_tokens_for_models( - request_body: dict, + request_body: dict[str, object], models: Sequence[str], ) -> Mapping[str, int]: return MappingProxyType( @@ -142,6 +158,8 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: ) except ValueError: return _count_text_tokens(model=model, text=request_body.get("messages")) + if "contents" in request_body: + return _count_gemini_input_tokens(request_body=request_body, model=model) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: @@ -160,6 +178,20 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: return None +def _count_gemini_input_tokens(request_body: Mapping[str, object], model: str) -> int | None: + system_instruction: Final = next( + (request_body[field] for field in _GEMINI_SYSTEM_INSTRUCTION_FIELDS if field in request_body), None + ) + translate: Final = GoogleGenAIAdapter().translate_generate_content_to_completion # pyright: ignore[reportUnknownMemberType] # adapter's **kwargs is untyped + chat_request: Final = translate( + model=model, + contents=_GEMINI_CONTENTS.validate_python(request_body["contents"]), + systemInstruction=system_instruction, + tools=request_body.get("tools"), + ) + return count_input_tokens_for_model(request_body=chat_request, model=model) + + def _count_text_tokens(model: str, text: object) -> int: if text is None: return 0 diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py index 078bd4dd402..54161b4883b 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py @@ -1,46 +1,6 @@ -import json -import traceback -from unittest import mock -from unittest.mock import MagicMock, patch - -import httpx -import pytest -from fastapi import Request, Response -from fastapi.testclient import TestClient - from litellm.passthrough.utils import CommonUtils -from unittest.mock import Mock - -from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key - - -@pytest.mark.asyncio -async def test_get_litellm_virtual_key(): - """ - Test that the get_litellm_virtual_key function correctly handles the API key authentication - """ - # Test with x-litellm-api-key - mock_request = Mock() - mock_request.headers = {"x-litellm-api-key": "test-key-123"} - result = get_litellm_virtual_key(mock_request) - assert result == "Bearer test-key-123" - - # Test with Authorization header - mock_request.headers = {"Authorization": "Bearer auth-key-456"} - result = get_litellm_virtual_key(mock_request) - assert result == "Bearer auth-key-456" - - # Test with both headers (x-litellm-api-key should take precedence) - mock_request.headers = { - "x-litellm-api-key": "test-key-123", - "Authorization": "Bearer auth-key-456", - } - result = get_litellm_virtual_key(mock_request) - assert result == "Bearer test-key-123" - - def test_encode_bedrock_runtime_modelid_arn(): # Test application-inference-profile ARN endpoint = "model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse" diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 024fe5229be..f1ae7bcb5cd 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -176,6 +176,52 @@ def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window(): assert invoke_cost < converse_cost < 2 * invoke_cost +GEMINI_PROMPT: Final = "Say hello in five words" +GEMINI_SYSTEM: Final = "Answer tersely" + + +@pytest.mark.parametrize( + "route", + ( + "/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent", + "/gemini/v1beta/models/gemini-3.8-flash:generateContent", + "/v1beta/models/gemini-3.8-flash:generateContent", + ), +) +@pytest.mark.parametrize( + ("native_output_config", "chat_output_limit"), + (({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}), ({}, {})), +) +def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window( + route: str, native_output_config: dict[str, object], chat_output_limit: dict[str, object] +) -> None: + native_cost: Final = estimate_request_max_cost( + request_body={ + "systemInstruction": {"parts": [{"text": GEMINI_SYSTEM}]}, + "contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT}]}], + **native_output_config, + }, + route=route, + llm_router=None, + input_token_counts={}, + ) + chat_cost: Final = estimate_request_max_cost( + request_body={ + "model": "gemini-3.8-flash", + "messages": [ + {"role": "system", "content": GEMINI_SYSTEM}, + {"role": "user", "content": GEMINI_PROMPT}, + ], + **chat_output_limit, + }, + route="/chat/completions", + llm_router=None, + input_token_counts={}, + ) + assert native_cost is not None + assert native_cost == chat_cost + + def _tiered_deployment(input_cost_per_token: float) -> Deployment: return Deployment( model_name="tiered-group", From 3317b2fda9cb1e5fdd495bea272a7853633877ec Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:14:15 -0700 Subject: [PATCH 06/15] fix(proxy): count every Gemini system part and candidateCount in the budget hold A Gemini generateContent body with several systemInstruction text parts lost every part after the first when LiteLLM translated it to a chat request. That dropped the later system text for non-Gemini models behind /v1beta/models/:generateContent, and made the budget hold count only the first part. The adapter now keeps every text part, as a plain string when there is one part and as text blocks when there are several. The output hold now also multiplies by generationConfig.candidateCount, the same way chat bodies use n, so a request asking for several candidates holds budget for all of them. --- .../google_genai/adapters/transformation.py | 10 ++++--- .../spend_tracking/budget_reservation.py | 26 +++++++++++-------- .../google_genai/test_google_genai_adapter.py | 21 +++++++++++++++ .../spend_tracking/test_budget_reservation.py | 18 +++++++++---- 4 files changed, 56 insertions(+), 19 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index e1085dd7857..0363a6d3d1c 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -628,10 +628,14 @@ class GoogleGenAIAdapter: ) -> list[AllMessageValues]: """Transform Google GenAI contents to OpenAI messages format""" system_parts: Final = system_instruction.get("parts", ()) if system_instruction else () + system_texts: Final = tuple(part["text"] for part in system_parts if "text" in part) + system_content: Final = ( + system_texts[0] + if len(system_texts) == 1 + else [ChatCompletionTextObject(type="text", text=text) for text in system_texts] + ) system_messages: Final = ( - (ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"]),) - if system_parts and "text" in system_parts[0] - else () + (ChatCompletionSystemMessage(role="system", content=system_content),) if system_texts else () ) return [*system_messages, *chain.from_iterable(_genai_content_messages(content) for content in contents)] diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index cfc865a84d1..5cada93bd33 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -16,6 +16,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.proxy._types import ( @@ -1652,24 +1653,27 @@ def _estimate_output_tokens( _OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens") +def _nested_field(request_body: Mapping[str, object], config_field: str, field: str) -> object: + config: Final = as_str_mapping(request_body.get(config_field)) + return config.get(field) if config is not None else None + + def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: - inference_config: Final = request_body.get("inferenceConfig") - generation_config: Final = request_body.get("generationConfig") candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), - inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None, - generation_config.get("maxOutputTokens") if isinstance(generation_config, Mapping) else None, + _nested_field(request_body, "inferenceConfig", "maxTokens"), + _nested_field(request_body, "generationConfig", "maxOutputTokens"), ) return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) -def _get_output_multiplier(request_body: dict) -> int: - output_multiplier = 1 - for key in ("n", "best_of"): - value = _to_int(request_body.get(key)) - if value is not None: - output_multiplier = max(output_multiplier, value) - return output_multiplier +def _get_output_multiplier(request_body: Mapping[str, object]) -> int: + candidates: Final = ( + request_body.get("n"), + request_body.get("best_of"), + _nested_field(request_body, "generationConfig", "candidateCount"), + ) + return max((1, *(count for count in map(_to_int, candidates) if count is not None))) def _is_input_only_route(route: str) -> bool: diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 0292e58c029..7d1e8e2caa4 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -182,6 +182,27 @@ def test_native_function_call_preserves_arguments_and_matching_response(args: ob ] +@pytest.mark.parametrize( + ("system_texts", "expected_content"), + [ + (("Answer tersely",), "Answer tersely"), + ( + ("Answer tersely", "Use British spelling"), + [{"type": "text", "text": "Answer tersely"}, {"type": "text", "text": "Use British spelling"}], + ), + ], +) +def test_every_system_instruction_text_part_reaches_the_system_message( + system_texts: tuple[str, ...], expected_content: str | list[dict[str, str]] +) -> None: + request: Final = GoogleGenAIAdapter().translate_generate_content_to_completion( + model="gemini/gemini-3.8-flash", + contents=[{"role": "user", "parts": [{"text": "hi"}]}], + systemInstruction={"parts": [{"text": text} for text in system_texts]}, + ) + assert request["messages"][0] == {"role": "system", "content": expected_content} + + def test_adapter_import(): """Test that the adapter can be imported successfully""" from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index f1ae7bcb5cd..78431304198 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -177,7 +177,7 @@ def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window(): GEMINI_PROMPT: Final = "Say hello in five words" -GEMINI_SYSTEM: Final = "Answer tersely" +GEMINI_SYSTEM_PARTS: Final = ("Answer tersely", "Use British spelling") @pytest.mark.parametrize( @@ -190,14 +190,22 @@ GEMINI_SYSTEM: Final = "Answer tersely" ) @pytest.mark.parametrize( ("native_output_config", "chat_output_limit"), - (({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}), ({}, {})), + ( + ({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}), + ({"generationConfig": {"maxOutputTokens": 5, "candidateCount": 3}}, {"max_tokens": 5, "n": 3}), + ({}, {}), + ), ) +@pytest.mark.parametrize("system_texts", (GEMINI_SYSTEM_PARTS[:1], GEMINI_SYSTEM_PARTS)) def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window( - route: str, native_output_config: dict[str, object], chat_output_limit: dict[str, object] + route: str, + native_output_config: dict[str, object], + chat_output_limit: dict[str, object], + system_texts: tuple[str, ...], ) -> None: native_cost: Final = estimate_request_max_cost( request_body={ - "systemInstruction": {"parts": [{"text": GEMINI_SYSTEM}]}, + "systemInstruction": {"parts": [{"text": text} for text in system_texts]}, "contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT}]}], **native_output_config, }, @@ -209,7 +217,7 @@ def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window request_body={ "model": "gemini-3.8-flash", "messages": [ - {"role": "system", "content": GEMINI_SYSTEM}, + {"role": "system", "content": [{"type": "text", "text": text} for text in system_texts]}, {"role": "user", "content": GEMINI_PROMPT}, ], **chat_output_limit, From b22e2ca8b6581d479032f1044d8f8111514b8506 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:41:39 -0700 Subject: [PATCH 07/15] fix(proxy): keep the context-window hold for Gemini bodies the prompt count cannot see A Gemini generateContent body with media, cachedContent, other roles, or non-function tools keeps the full context-window hold, since the chat translation drops or undercounts those parts. The hold also reads the snake_case generation_config output fields, and several systemInstruction text parts reach non-Gemini models joined into one system string --- .../google_genai/adapters/transformation.py | 7 +- .../spend_tracking/budget_reservation.py | 13 +++- litellm/proxy/spend_tracking/input_tokens.py | 67 +++++++++++++++++-- .../google_genai/test_google_genai_adapter.py | 7 +- .../spend_tracking/test_budget_reservation.py | 63 ++++++++++++++++- 5 files changed, 137 insertions(+), 20 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 0363a6d3d1c..594b97138ef 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -629,13 +629,8 @@ class GoogleGenAIAdapter: """Transform Google GenAI contents to OpenAI messages format""" system_parts: Final = system_instruction.get("parts", ()) if system_instruction else () system_texts: Final = tuple(part["text"] for part in system_parts if "text" in part) - system_content: Final = ( - system_texts[0] - if len(system_texts) == 1 - else [ChatCompletionTextObject(type="text", text=text) for text in system_texts] - ) system_messages: Final = ( - (ChatCompletionSystemMessage(role="system", content=system_content),) if system_texts else () + (ChatCompletionSystemMessage(role="system", content="\n".join(system_texts)),) if system_texts else () ) return [*system_messages, *chain.from_iterable(_genai_content_messages(content) for content in contents)] diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 5cada93bd33..d20e2371f09 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -7,6 +7,7 @@ import time from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from itertools import product from types import MappingProxyType from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast @@ -1651,6 +1652,7 @@ def _estimate_output_tokens( _OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens") +_GEMINI_CONFIG_FIELDS: Final = ("generationConfig", "generation_config") def _nested_field(request_body: Mapping[str, object], config_field: str, field: str) -> object: @@ -1658,11 +1660,18 @@ def _nested_field(request_body: Mapping[str, object], config_field: str, field: return config.get(field) if config is not None else None +def _gemini_config_values(request_body: Mapping[str, object], *fields: str) -> tuple[object, ...]: + return tuple( + _nested_field(request_body, config_field, field) + for config_field, field in product(_GEMINI_CONFIG_FIELDS, fields) + ) + + def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), _nested_field(request_body, "inferenceConfig", "maxTokens"), - _nested_field(request_body, "generationConfig", "maxOutputTokens"), + *_gemini_config_values(request_body, "maxOutputTokens", "max_output_tokens"), ) return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) @@ -1671,7 +1680,7 @@ def _get_output_multiplier(request_body: Mapping[str, object]) -> int: candidates: Final = ( request_body.get("n"), request_body.get("best_of"), - _nested_field(request_body, "generationConfig", "candidateCount"), + *_gemini_config_values(request_body, "candidateCount", "candidate_count"), ) return max((1, *(count for count in map(_to_int, candidates) if count is not None))) diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index a7e9981ac62..a2559f22892 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -22,6 +22,7 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items from litellm.responses.litellm_completion_transformation.compaction_items import ( billable_input_since_newest_compaction, holds_compaction_summary, @@ -39,6 +40,16 @@ from litellm.rust_bridge.token_counter import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 _GEMINI_SYSTEM_INSTRUCTION_FIELDS: Final = ("systemInstruction", "system_instruction") +_GEMINI_CACHED_CONTENT_FIELDS: Final = ("cachedContent", "cached_content") +_GEMINI_FUNCTION_RESPONSE_FIELDS: Final = ("functionResponse", "function_response") +_GEMINI_TEXT_PART_FIELDS: Final = frozenset({"text", "thought", "thoughtSignature", "thought_signature"}) +_GEMINI_COUNTED_PART_FIELDS: Final = _GEMINI_TEXT_PART_FIELDS | { + "functionCall", + "function_call", + *_GEMINI_FUNCTION_RESPONSE_FIELDS, +} +_GEMINI_COUNTED_ROLES: Final = frozenset({"user", "model"}) +_GEMINI_COUNTED_TOOL_FIELDS: Final = frozenset({"functionDeclarations"}) _GeminiContents: TypeAlias = list[dict[str, object]] | dict[str, object] _GEMINI_CONTENTS: Final = TypeAdapter[_GeminiContents](_GeminiContents) _INPUT_SIZE_FIELDS: Final = ( @@ -182,16 +193,60 @@ def _count_gemini_input_tokens(request_body: Mapping[str, object], model: str) - system_instruction: Final = next( (request_body[field] for field in _GEMINI_SYSTEM_INSTRUCTION_FIELDS if field in request_body), None ) + contents: Final = _GEMINI_CONTENTS.validate_python(request_body["contents"]) + tools: Final = request_body.get("tools") + if not _gemini_body_is_countable(request_body, contents, system_instruction, tools): + return None translate: Final = GoogleGenAIAdapter().translate_generate_content_to_completion # pyright: ignore[reportUnknownMemberType] # adapter's **kwargs is untyped - chat_request: Final = translate( - model=model, - contents=_GEMINI_CONTENTS.validate_python(request_body["contents"]), - systemInstruction=system_instruction, - tools=request_body.get("tools"), - ) + chat_request: Final = translate(model=model, contents=contents, systemInstruction=system_instruction, tools=tools) return count_input_tokens_for_model(request_body=chat_request, model=model) +def _gemini_body_is_countable( + request_body: Mapping[str, object], + contents: _GeminiContents, + system_instruction: object, + tools: object, +) -> bool: + content_items: Final = contents if isinstance(contents, list) else (contents,) + system: Final = as_str_mapping(system_instruction) + return ( + not any(field in request_body for field in _GEMINI_CACHED_CONTENT_FIELDS) + and all(_gemini_content_is_countable(content) for content in content_items) + and ( + system_instruction is None + or (system is not None and _gemini_parts_are_countable(system.get("parts", ()), _GEMINI_TEXT_PART_FIELDS)) + ) + and all(_gemini_tool_is_countable(tool) for tool in iter_items(tools or ())) + ) + + +def _gemini_content_is_countable(content: Mapping[str, object]) -> bool: + return content.get("role", "user") in _GEMINI_COUNTED_ROLES and _gemini_parts_are_countable( + content.get("parts", ()), _GEMINI_COUNTED_PART_FIELDS + ) + + +def _gemini_parts_are_countable(parts: object, fields: frozenset[str]) -> bool: + return all(_gemini_part_is_countable(part, fields) for part in iter_items(parts)) + + +def _gemini_part_is_countable(part: object, fields: frozenset[str]) -> bool: + if isinstance(part, str): + return True + mapping: Final = as_str_mapping(part) + return ( + mapping is not None + and mapping.keys() <= fields + and not any("parts" in (as_str_mapping(mapping.get(field)) or ()) for field in _GEMINI_FUNCTION_RESPONSE_FIELDS) + ) + + +def _gemini_tool_is_countable(tool: object) -> bool: + mapping: Final = as_str_mapping(tool) + return mapping is not None and mapping.keys() <= _GEMINI_COUNTED_TOOL_FIELDS + + def _count_text_tokens(model: str, text: object) -> int: if text is None: return 0 diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 7d1e8e2caa4..e1ffaadb46d 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -186,14 +186,11 @@ def test_native_function_call_preserves_arguments_and_matching_response(args: ob ("system_texts", "expected_content"), [ (("Answer tersely",), "Answer tersely"), - ( - ("Answer tersely", "Use British spelling"), - [{"type": "text", "text": "Answer tersely"}, {"type": "text", "text": "Use British spelling"}], - ), + (("Answer tersely", "Use British spelling"), "Answer tersely\nUse British spelling"), ], ) def test_every_system_instruction_text_part_reaches_the_system_message( - system_texts: tuple[str, ...], expected_content: str | list[dict[str, str]] + system_texts: tuple[str, ...], expected_content: str ) -> None: request: Final = GoogleGenAIAdapter().translate_generate_content_to_completion( model="gemini/gemini-3.8-flash", diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 78431304198..2b654d9085d 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -193,6 +193,7 @@ GEMINI_SYSTEM_PARTS: Final = ("Answer tersely", "Use British spelling") ( ({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}), ({"generationConfig": {"maxOutputTokens": 5, "candidateCount": 3}}, {"max_tokens": 5, "n": 3}), + ({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {"max_tokens": 5, "n": 3}), ({}, {}), ), ) @@ -217,7 +218,7 @@ def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window request_body={ "model": "gemini-3.8-flash", "messages": [ - {"role": "system", "content": [{"type": "text", "text": text} for text in system_texts]}, + {"role": "system", "content": "\n".join(system_texts)}, {"role": "user", "content": GEMINI_PROMPT}, ], **chat_output_limit, @@ -230,6 +231,66 @@ def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window assert native_cost == chat_cost +GEMINI_ROUTE: Final = "/v1beta/models/gemini-3.8-flash:generateContent" +GEMINI_TEXT_CONTENT: Final = {"role": "user", "parts": [{"text": GEMINI_PROMPT}]} +GEMINI_AUDIO_PART: Final = {"inlineData": {"mimeType": "audio/wav", "data": "UklGRg" * 40_000}} + + +def _gemini_hold(body: dict[str, object], input_tokens: int | None = None) -> float | None: + return estimate_request_max_cost( + request_body=body, + route=GEMINI_ROUTE, + llm_router=None, + input_token_counts={} if input_tokens is None else {"gemini-3.8-flash": input_tokens}, + ) + + +@pytest.mark.parametrize( + "body", + ( + {"contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT}, GEMINI_AUDIO_PART]}]}, + {"contents": [{"role": "user", "parts": [{"fileData": {"mimeType": "video/mp4", "fileUri": "gs://b/v.mp4"}}]}]}, + {"contents": [{"role": "user", "parts": [{"inline_data": {"mime_type": "image/png", "data": "iVBORw0K"}}]}]}, + {"contents": [{"role": "model", "parts": [{"executableCode": {"language": "PYTHON", "code": "print(1)"}}]}]}, + {"contents": [GEMINI_TEXT_CONTENT, {"role": "function", "parts": [{"text": GEMINI_PROMPT * 500}]}]}, + { + "contents": [ + { + "role": "user", + "parts": [{"functionResponse": {"name": "listen", "response": {}, "parts": [GEMINI_AUDIO_PART]}}], + } + ] + }, + {"systemInstruction": {"parts": [GEMINI_AUDIO_PART]}, "contents": [GEMINI_TEXT_CONTENT]}, + {"tools": [{"function_declarations": [{"name": "lookup"}]}], "contents": [GEMINI_TEXT_CONTENT]}, + {"tools": [{"googleSearch": {}}], "contents": [GEMINI_TEXT_CONTENT]}, + {"cachedContent": "cachedContents/abc", "contents": [GEMINI_TEXT_CONTENT]}, + ), +) +def test_gemini_body_the_prompt_count_cannot_see_whole_keeps_the_context_window_hold(body: dict[str, object]) -> None: + context_window: Final = litellm.get_model_info("gemini-3.8-flash")["max_input_tokens"] + assert context_window is not None + assert _gemini_hold(body) == _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}, input_tokens=context_window) + + +def test_gemini_tool_turns_and_thought_signatures_count_as_prompt() -> None: + context_window: Final = litellm.get_model_info("gemini-3.8-flash")["max_input_tokens"] + assert context_window is not None + tool_turns: Final = { + "tools": [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "object"}}]}], + "contents": [ + GEMINI_TEXT_CONTENT, + {"role": "model", "parts": [{"functionCall": {"name": "lookup", "args": {}}, "thoughtSignature": "c2ln"}]}, + {"role": "user", "parts": [{"functionResponse": {"name": "lookup", "response": {"answer": "hi"}}}]}, + ], + } + text_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) + tool_turns_hold: Final = _gemini_hold(tool_turns) + context_window_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}, input_tokens=context_window) + assert text_hold is not None and tool_turns_hold is not None and context_window_hold is not None + assert text_hold < tool_turns_hold < context_window_hold + + def _tiered_deployment(input_cost_per_token: float) -> Deployment: return Deployment( model_name="tiered-group", From d793d2917fe87f2c7a45346ce5e79d29258dec06 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:07:38 -0700 Subject: [PATCH 08/15] fix(proxy): read only camelCase Gemini output fields in the budget hold /v1beta forwards only generationConfig.maxOutputTokens to non-Gemini models, so a snake_case cap the model never receives must not shrink the hold. Snake_case bodies now get the default output hold. --- litellm/proxy/spend_tracking/budget_reservation.py | 13 ++----------- .../proxy/spend_tracking/test_budget_reservation.py | 2 +- 2 files changed, 3 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index d20e2371f09..5cada93bd33 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -7,7 +7,6 @@ import time from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from itertools import product from types import MappingProxyType from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast @@ -1652,7 +1651,6 @@ def _estimate_output_tokens( _OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens") -_GEMINI_CONFIG_FIELDS: Final = ("generationConfig", "generation_config") def _nested_field(request_body: Mapping[str, object], config_field: str, field: str) -> object: @@ -1660,18 +1658,11 @@ def _nested_field(request_body: Mapping[str, object], config_field: str, field: return config.get(field) if config is not None else None -def _gemini_config_values(request_body: Mapping[str, object], *fields: str) -> tuple[object, ...]: - return tuple( - _nested_field(request_body, config_field, field) - for config_field, field in product(_GEMINI_CONFIG_FIELDS, fields) - ) - - def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), _nested_field(request_body, "inferenceConfig", "maxTokens"), - *_gemini_config_values(request_body, "maxOutputTokens", "max_output_tokens"), + _nested_field(request_body, "generationConfig", "maxOutputTokens"), ) return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) @@ -1680,7 +1671,7 @@ def _get_output_multiplier(request_body: Mapping[str, object]) -> int: candidates: Final = ( request_body.get("n"), request_body.get("best_of"), - *_gemini_config_values(request_body, "candidateCount", "candidate_count"), + _nested_field(request_body, "generationConfig", "candidateCount"), ) return max((1, *(count for count in map(_to_int, candidates) if count is not None))) diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 2b654d9085d..97f65a4402d 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -193,7 +193,7 @@ GEMINI_SYSTEM_PARTS: Final = ("Answer tersely", "Use British spelling") ( ({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}), ({"generationConfig": {"maxOutputTokens": 5, "candidateCount": 3}}, {"max_tokens": 5, "n": 3}), - ({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {"max_tokens": 5, "n": 3}), + ({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {}), ({}, {}), ), ) From fb75a132e5c903916209b43ed4061ad472712ef5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:17:42 -0700 Subject: [PATCH 09/15] fix(proxy): count snake_case candidate_count in the Gemini budget hold --- litellm/proxy/spend_tracking/budget_reservation.py | 11 ++++++++++- .../proxy/spend_tracking/test_budget_reservation.py | 3 ++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 5cada93bd33..4efcc06fdce 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -7,6 +7,7 @@ import time from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from itertools import product from types import MappingProxyType from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast @@ -1651,6 +1652,7 @@ def _estimate_output_tokens( _OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens") +_GEMINI_CONFIG_FIELDS: Final = ("generationConfig", "generation_config") def _nested_field(request_body: Mapping[str, object], config_field: str, field: str) -> object: @@ -1658,6 +1660,13 @@ def _nested_field(request_body: Mapping[str, object], config_field: str, field: return config.get(field) if config is not None else None +def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object, ...]: + return tuple( + _nested_field(request_body, config_field, field) + for config_field, field in product(_GEMINI_CONFIG_FIELDS, ("candidateCount", "candidate_count")) + ) + + def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), @@ -1671,7 +1680,7 @@ def _get_output_multiplier(request_body: Mapping[str, object]) -> int: candidates: Final = ( request_body.get("n"), request_body.get("best_of"), - _nested_field(request_body, "generationConfig", "candidateCount"), + *_gemini_candidate_counts(request_body), ) return max((1, *(count for count in map(_to_int, candidates) if count is not None))) diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 97f65a4402d..01acfbf9a61 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -193,7 +193,8 @@ GEMINI_SYSTEM_PARTS: Final = ("Answer tersely", "Use British spelling") ( ({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}), ({"generationConfig": {"maxOutputTokens": 5, "candidateCount": 3}}, {"max_tokens": 5, "n": 3}), - ({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {}), + ({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {"n": 3}), + ({"generationConfig": {"candidate_count": 3}}, {"n": 3}), ({}, {}), ), ) From edfc8744dbae84eda68a85333383ac4ae5561dcf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 01:18:47 -0700 Subject: [PATCH 10/15] fix(vertex_ai): skip malformed system parts and read Gemini output fields only on Gemini bodies A systemInstruction part whose text is not a string is skipped, as main did, instead of failing the request. The budget hold reads generationConfig output and candidate fields, and counts contents, only when the body has no messages, prompt, input, query or documents field. An output cap of 0 or below counts as no cap. Test recorders keep tuples instead of mutating lists. --- .../google_genai/adapters/transformation.py | 16 ++++++- .../spend_tracking/budget_reservation.py | 17 ++++++-- litellm/proxy/spend_tracking/input_tokens.py | 7 +++- .../google_genai/test_google_genai_adapter.py | 21 ++++++++++ .../test_llm_pass_through_endpoints.py | 24 +++++------ .../spend_tracking/test_budget_reservation.py | 42 +++++++++++++++++++ .../proxy/spend_tracking/test_input_tokens.py | 8 ++++ 7 files changed, 117 insertions(+), 18 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 594b97138ef..f624f69567e 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -9,6 +9,7 @@ from typing_extensions import ReadOnly, Required, TypedDict from litellm import verbose_logger from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types, normalize_tool_schema @@ -244,6 +245,18 @@ def _genai_tool_call(part: object) -> ChatCompletionAssistantToolCall | None: ) +def _system_part_text(part: object) -> str | None: + fields: Final = as_str_mapping(part) + text: Final = fields.get("text") if fields is not None else None + return text if isinstance(text, str) else None + + +def _system_texts(parts: object) -> tuple[str, ...]: + return tuple( + text for text in map(_system_part_text, parts if isinstance(parts, Sequence) else ()) if text is not None + ) + + def _genai_content_messages(content: object) -> tuple[AllMessageValues, ...]: fields: Final = _validated(_CONFIG_FIELDS, content) if fields is None: @@ -627,8 +640,7 @@ class GoogleGenAIAdapter: system_instruction: _GenAISystemInstruction | None = None, ) -> list[AllMessageValues]: """Transform Google GenAI contents to OpenAI messages format""" - system_parts: Final = system_instruction.get("parts", ()) if system_instruction else () - system_texts: Final = tuple(part["text"] for part in system_parts if "text" in part) + system_texts: Final = _system_texts(system_instruction.get("parts") if system_instruction else None) system_messages: Final = ( (ChatCompletionSystemMessage(role="system", content="\n".join(system_texts)),) if system_texts else () ) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 4efcc06fdce..b1746ca830b 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -41,7 +41,11 @@ from litellm.proxy.common_utils.user_api_key_cache import ( tag_cache_key, team_membership_reservation_cache_key, ) -from litellm.proxy.spend_tracking.input_tokens import count_input_tokens, count_input_tokens_for_model +from litellm.proxy.spend_tracking.input_tokens import ( + count_input_tokens, + count_input_tokens_for_model, + is_gemini_contents_body, +) from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -1661,6 +1665,8 @@ def _nested_field(request_body: Mapping[str, object], config_field: str, field: def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object, ...]: + if not is_gemini_contents_body(request_body): + return () return tuple( _nested_field(request_body, config_field, field) for config_field, field in product(_GEMINI_CONFIG_FIELDS, ("candidateCount", "candidate_count")) @@ -1668,12 +1674,17 @@ def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: + gemini_output_cap: Final = ( + _nested_field(request_body, "generationConfig", "maxOutputTokens") + if is_gemini_contents_body(request_body) + else None + ) candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), _nested_field(request_body, "inferenceConfig", "maxTokens"), - _nested_field(request_body, "generationConfig", "maxOutputTokens"), + gemini_output_cap, ) - return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) + return next((tokens for tokens in map(_to_int, candidates) if tokens is not None and tokens > 0), None) def _get_output_multiplier(request_body: Mapping[str, object]) -> int: diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index a2559f22892..452a12784b4 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -39,6 +39,7 @@ from litellm.rust_bridge.token_counter import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 +_NON_GEMINI_INPUT_FIELDS: Final = ("messages", "prompt", "input", "query", "documents") _GEMINI_SYSTEM_INSTRUCTION_FIELDS: Final = ("systemInstruction", "system_instruction") _GEMINI_CACHED_CONTENT_FIELDS: Final = ("cachedContent", "cached_content") _GEMINI_FUNCTION_RESPONSE_FIELDS: Final = ("functionResponse", "function_response") @@ -169,7 +170,7 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: ) except ValueError: return _count_text_tokens(model=model, text=request_body.get("messages")) - if "contents" in request_body: + if is_gemini_contents_body(request_body): return _count_gemini_input_tokens(request_body=request_body, model=model) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) @@ -189,6 +190,10 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: return None +def is_gemini_contents_body(request_body: Mapping[str, object]) -> bool: + return "contents" in request_body and not any(field in request_body for field in _NON_GEMINI_INPUT_FIELDS) + + def _count_gemini_input_tokens(request_body: Mapping[str, object], model: str) -> int | None: system_instruction: Final = next( (request_body[field] for field in _GEMINI_SYSTEM_INSTRUCTION_FIELDS if field in request_body), None diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index e1ffaadb46d..ae04719f4a0 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -200,6 +200,27 @@ def test_every_system_instruction_text_part_reaches_the_system_message( assert request["messages"][0] == {"role": "system", "content": expected_content} +@pytest.mark.parametrize( + ("system_parts", "expected_messages"), + [ + ( + [{"text": "Answer tersely"}, {"text": None}, {"text": 3}, "text", {"text": "Use British spelling"}], + [{"role": "system", "content": "Answer tersely\nUse British spelling"}, {"role": "user", "content": "hi"}], + ), + (None, [{"role": "user", "content": "hi"}]), + ], +) +def test_a_system_instruction_part_without_string_text_is_skipped( + system_parts: object, expected_messages: list[dict[str, str]] +) -> None: + request: Final = GoogleGenAIAdapter().translate_generate_content_to_completion( + model="gemini/gemini-3.8-flash", + contents=[{"role": "user", "parts": [{"text": "hi"}]}], + systemInstruction={"parts": system_parts}, + ) + assert request["messages"] == expected_messages + + def test_adapter_import(): """Test that the adapter can be imported successfully""" from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index ee069c01c04..45b262b250d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -8282,10 +8282,10 @@ class TestAzureSpeechRawBodyThroughRealAuth: class _CountingCustomAuth: def __init__(self, only_accepts: str | None = None) -> None: self.only_accepts: Final = only_accepts - self.api_keys_seen: list[str] = [] + self.api_keys_seen: tuple[str, ...] = () async def __call__(self, request: Request, api_key: str) -> UserAPIKeyAuth: - self.api_keys_seen.append(api_key) + self.api_keys_seen = (*self.api_keys_seen, api_key) if self.only_accepts is not None and api_key != self.only_accepts: raise ValueError("unknown caller") return UserAPIKeyAuth(user_id=f"custom-auth-caller-{len(self.api_keys_seen)}") @@ -8294,12 +8294,12 @@ class _CountingCustomAuth: class _ServedCallerRecorder(CustomLogger): def __init__(self) -> None: super().__init__() - self.served_user_ids: list[str | None] = [] + self.served_user_ids: tuple[str | None, ...] = () async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral ) -> dict: - self.served_user_ids.append(user_api_key_dict.user_id) + self.served_user_ids = (*self.served_user_ids, user_api_key_dict.user_id) return data @@ -8378,8 +8378,8 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: ) assert response.status_code == 200, response.text - assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] - assert recorder.served_user_ids == ["custom-auth-caller-1"] + assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,) + assert recorder.served_user_ids == ("custom-auth-caller-1",) assert vertex.call_count == 1 def test_discovery_search_runs_custom_auth_once_and_is_served_as_its_caller( @@ -8400,8 +8400,8 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: ) assert (response.status_code, response.json()) == (200, {"results": []}) - assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] - assert recorder.served_user_ids == ["custom-auth-caller-1"] + assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,) + assert recorder.served_user_ids == ("custom-auth-caller-1",) assert discovery.call_count == 1 @pytest.mark.parametrize("scheme", ["bearer", "Basic"]) @@ -8445,8 +8445,8 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: ) assert response.status_code == 200, response.text - assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] - assert recorder.served_user_ids == ["custom-auth-caller-1"] + assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,) + assert recorder.served_user_ids == ("custom-auth-caller-1",) assert google.call_count == 1 def test_x_litellm_api_key_authenticates_once_and_is_never_forwarded_without_a_vertex_credential( @@ -8465,7 +8465,7 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: ) assert response.status_code == 401, response.text - assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] + assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,) assert not google.called @pytest.mark.parametrize("header", ["x-goog-api-key", "x-api-key", "API-Key", "Ocp-Apim-Subscription-Key"]) @@ -8482,7 +8482,7 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: assert response.status_code == 401, response.text assert self.VIRTUAL_KEY not in custom_auth.api_keys_seen - assert recorder.served_user_ids == [] + assert recorder.served_user_ids == () assert not google.called diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 01acfbf9a61..42ae046ec72 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -292,6 +292,48 @@ def test_gemini_tool_turns_and_thought_signatures_count_as_prompt() -> None: assert text_hold < tool_turns_hold < context_window_hold +GEMINI_CHAT_BODY: Final = { + "model": "gemini-3.8-flash", + "messages": [{"role": "user", "content": GEMINI_PROMPT}], +} + + +def _chat_hold(body: dict[str, object]) -> float | None: + return estimate_request_max_cost( + request_body=body, + route="/chat/completions", + llm_router=None, + input_token_counts={}, + ) + + +@pytest.mark.parametrize( + "gemini_output_fields", + ( + {"generationConfig": {"maxOutputTokens": 5}}, + {"generationConfig": {"candidateCount": 3}}, + {"generation_config": {"candidate_count": 3}}, + {"contents": [], "generationConfig": {"maxOutputTokens": 5}}, + ), +) +def test_gemini_output_fields_on_a_chat_body_leave_its_hold_alone(gemini_output_fields: dict[str, object]) -> None: + plain_hold: Final = _chat_hold(GEMINI_CHAT_BODY) + assert plain_hold is not None + assert _chat_hold({**GEMINI_CHAT_BODY, **gemini_output_fields}) == plain_hold + + +@pytest.mark.parametrize("cap", (0, -1000)) +def test_an_output_cap_of_zero_or_below_holds_the_default_output(cap: int) -> None: + uncapped_gemini_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) + uncapped_chat_hold: Final = _chat_hold(GEMINI_CHAT_BODY) + assert uncapped_gemini_hold is not None and uncapped_chat_hold is not None + assert ( + _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": cap}}) + == uncapped_gemini_hold + ) + assert _chat_hold({**GEMINI_CHAT_BODY, "max_tokens": cap}) == uncapped_chat_hold + + def _tiered_deployment(input_cost_per_token: float) -> Deployment: return Deployment( model_name="tiered-group", diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index c7f13651471..72c9ff4ae55 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -110,3 +110,11 @@ def test_a_compaction_item_counts_only_its_summary(model: str) -> None: assert count_input_tokens_for_model(request_body=padded_body, model=model) == count_input_tokens_for_model( request_body=plain_body, model=model ) + + +@pytest.mark.parametrize("input_field", ("prompt", "input", "query")) +def test_a_stray_contents_field_never_replaces_the_request_input_count(input_field: str) -> None: + body: Final = {"model": CL100K_MODEL, input_field: "count every one of these words " * 200} + count: Final = count_input_tokens_for_model(request_body=body, model=CL100K_MODEL) + assert count is not None and count > 1000 + assert count_input_tokens_for_model(request_body={**body, "contents": []}, model=CL100K_MODEL) == count From 84a4505d4248c0eddab99800bec8823a9e4c4519 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 04:13:55 -0700 Subject: [PATCH 11/15] fix(proxy): take the budget hold's request shape from the route On generateContent routes the hold counts only contents and reads its output cap and multiplier from generationConfig, so stray fields of other APIs can no longer hide the prompt. A max_tokens of 0 holds no output again, while negative values and a Gemini maxOutputTokens of 0 or below hold the default output --- .../spend_tracking/budget_reservation.py | 33 +++++----- litellm/proxy/spend_tracking/input_tokens.py | 31 +++++---- .../spend_tracking/test_budget_reservation.py | 66 +++++++++++++++++-- 3 files changed, 95 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 55e59c0f8b0..566fbc886ad 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -44,7 +44,6 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.spend_tracking.input_tokens import ( count_input_tokens, count_input_tokens_for_model, - is_gemini_contents_body, tokenizer_models, ) from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope @@ -1441,7 +1440,7 @@ def _max_cost_for_cost_info( if estimated_input_tokens is None or output_tokens is None: return None - output_multiplier: Final = _get_output_multiplier(request_body=request_body) + output_multiplier: Final = _get_output_multiplier(request_body=request_body, route=route) tiered_pricing: Final = model_info.get("tiered_pricing") if isinstance(tiered_pricing, list) and tiered_pricing: tier: Final = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=estimated_input_tokens) @@ -1611,6 +1610,7 @@ async def count_request_input_tokens( request_body=request_body, raw_body=raw_body, models=tuple(dict.fromkeys(chain.from_iterable(counted_as.values()))), + gemini_route=RouteChecks.is_generate_content_route(route), ) return MappingProxyType( { @@ -1631,7 +1631,9 @@ def _estimate_input_tokens( counted: Final = ( input_tokens if input_tokens is not None - else count_input_tokens_for_model(request_body=request_body, model=model) + else count_input_tokens_for_model( + request_body=request_body, model=model, gemini_route=RouteChecks.is_generate_content_route(route) + ) ) if counted is not None: return counted @@ -1654,7 +1656,7 @@ def _estimate_output_tokens( if _is_input_only_route(route=route): return 0 - requested: Final = _requested_output_tokens(request_body) + requested: Final = _requested_output_tokens(request_body, route) # Clamp at min(requested-or-default, model_max-or-default). Two purposes: # (1) Without an explicit cap we still need a finite reservation so the @@ -1678,33 +1680,28 @@ def _nested_field(request_body: Mapping[str, object], config_field: str, field: def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object, ...]: - if not is_gemini_contents_body(request_body): - return () return tuple( _nested_field(request_body, config_field, field) for config_field, field in product(_GEMINI_CONFIG_FIELDS, ("candidateCount", "candidate_count")) ) -def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: - gemini_output_cap: Final = ( - _nested_field(request_body, "generationConfig", "maxOutputTokens") - if is_gemini_contents_body(request_body) - else None - ) +def _requested_output_tokens(request_body: Mapping[str, object], route: str) -> int | None: + if RouteChecks.is_generate_content_route(route): + gemini_cap: Final = _to_int(_nested_field(request_body, "generationConfig", "maxOutputTokens")) + return gemini_cap if gemini_cap is not None and gemini_cap > 0 else None candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), _nested_field(request_body, "inferenceConfig", "maxTokens"), - gemini_output_cap, ) - return next((tokens for tokens in map(_to_int, candidates) if tokens is not None and tokens > 0), None) + return next((tokens for tokens in map(_to_int, candidates) if tokens is not None and tokens >= 0), None) -def _get_output_multiplier(request_body: Mapping[str, object]) -> int: +def _get_output_multiplier(request_body: Mapping[str, object], route: str) -> int: candidates: Final = ( - request_body.get("n"), - request_body.get("best_of"), - *_gemini_candidate_counts(request_body), + _gemini_candidate_counts(request_body) + if RouteChecks.is_generate_content_route(route) + else (request_body.get("n"), request_body.get("best_of")) ) return max((1, *(count for count in map(_to_int, candidates) if count is not None))) diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 7f33cda7592..30de3e6f596 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -44,7 +44,6 @@ if TYPE_CHECKING: TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 -_NON_GEMINI_INPUT_FIELDS: Final = ("messages", "prompt", "input", "query", "documents") _GEMINI_SYSTEM_INSTRUCTION_FIELDS: Final = ("systemInstruction", "system_instruction") _GEMINI_CACHED_CONTENT_FIELDS: Final = ("cachedContent", "cached_content") _GEMINI_FUNCTION_RESPONSE_FIELDS: Final = ("functionResponse", "function_response") @@ -103,6 +102,7 @@ async def count_input_tokens( request_body: dict[str, object], raw_body: bytes | None, models: Sequence[str], + gemini_route: bool = False, ) -> Mapping[str, int]: """Input-token count per model, sharing one native count across models that select the same tokenizer.""" @@ -110,13 +110,14 @@ async def count_input_tokens( (model, rust_tokenizer(model)) for model in models ) groups: Final[tuple[RustTokenizer | None, ...]] = tuple(dict.fromkeys(tokenizer for _, tokenizer in tokenizers)) - counted_raw_body: Final = None if holds_compaction_summary(request_body.get("input")) else raw_body + counted_raw_body: Final = None if gemini_route or holds_compaction_summary(request_body.get("input")) else raw_body group_counts: Final = [ await _count_group( request_body=request_body, raw_body=counted_raw_body, tokenizer=tokenizer, models=tuple(model for model, selected in tokenizers if selected == tokenizer), + gemini_route=gemini_route, ) for tokenizer in groups ] @@ -130,14 +131,16 @@ async def _count_group( raw_body: bytes | None, tokenizer: RustTokenizer | None, models: tuple[str, ...], + gemini_route: bool, ) -> Mapping[str, int]: async def python() -> Mapping[str, int]: if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: - return _count_input_tokens_for_models(request_body=request_body, models=models) + return _count_input_tokens_for_models(request_body=request_body, models=models, gemini_route=gemini_route) return await asyncio.to_thread( _count_input_tokens_for_models, request_body=request_body, models=models, + gemini_route=gemini_route, ) if tokenizer is None or raw_body is None: @@ -172,18 +175,30 @@ async def _native_counts( def _count_input_tokens_for_models( request_body: dict[str, object], models: Sequence[str], + gemini_route: bool, ) -> Mapping[str, int]: return MappingProxyType( { model: tokens for model in models - if (tokens := count_input_tokens_for_model(request_body=request_body, model=model)) is not None + if ( + tokens := count_input_tokens_for_model( + request_body=request_body, model=model, gemini_route=gemini_route + ) + ) + is not None } ) -def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: +def count_input_tokens_for_model(request_body: dict, model: str, gemini_route: bool = False) -> int | None: try: + if gemini_route: + return ( + _count_gemini_input_tokens(request_body=request_body, model=model) + if "contents" in request_body + else None + ) if "messages" in request_body: try: return litellm.token_counter( @@ -194,8 +209,6 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: ) except ValueError: return _count_text_tokens(model=model, text=request_body.get("messages")) - if is_gemini_contents_body(request_body): - return _count_gemini_input_tokens(request_body=request_body, model=model) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: @@ -214,10 +227,6 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: return None -def is_gemini_contents_body(request_body: Mapping[str, object]) -> bool: - return "contents" in request_body and not any(field in request_body for field in _NON_GEMINI_INPUT_FIELDS) - - def _count_gemini_input_tokens(request_body: Mapping[str, object], model: str) -> int | None: system_instruction: Final = next( (request_body[field] for field in _GEMINI_SYSTEM_INSTRUCTION_FIELDS if field in request_body), None diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index f5e6fa9af48..12d509d5198 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import math from datetime import datetime, timedelta, timezone from typing import Final @@ -25,6 +26,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.spend_tracking.budget_reservation import ( _get_team_member_budget_counter, + estimate_request_input_cost, estimate_request_max_cost, get_budget_window_start, release_unbound_budget_reservation, @@ -324,15 +326,67 @@ def test_gemini_output_fields_on_a_chat_body_leave_its_hold_alone(gemini_output_ @pytest.mark.parametrize("cap", (0, -1000)) -def test_an_output_cap_of_zero_or_below_holds_the_default_output(cap: int) -> None: - uncapped_gemini_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) - uncapped_chat_hold: Final = _chat_hold(GEMINI_CHAT_BODY) - assert uncapped_gemini_hold is not None and uncapped_chat_hold is not None +def test_a_gemini_output_cap_of_zero_or_below_holds_the_default_output(cap: int) -> None: + uncapped_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) + assert uncapped_hold is not None assert ( _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": cap}}) - == uncapped_gemini_hold + == uncapped_hold ) - assert _chat_hold({**GEMINI_CHAT_BODY, "max_tokens": cap}) == uncapped_chat_hold + + +def test_a_negative_max_tokens_holds_the_default_output() -> None: + uncapped_hold: Final = _chat_hold(GEMINI_CHAT_BODY) + assert uncapped_hold is not None + assert _chat_hold({**GEMINI_CHAT_BODY, "max_tokens": -1000}) == uncapped_hold + + +def test_a_text_completion_with_max_tokens_zero_holds_only_its_prompt() -> None: + body: Final = {"model": "gemini-3.8-flash", "prompt": GEMINI_PROMPT * 50, "max_tokens": 0} + prompt_cost: Final = estimate_request_input_cost(request_body=body, route="/v1/completions", llm_router=None) + assert prompt_cost is not None and prompt_cost > 0 + assert estimate_request_max_cost(request_body=body, route="/v1/completions", llm_router=None) == prompt_cost + + +GEMINI_LONG_CONTENTS: Final = {"contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT * 500}]}]} + + +@pytest.mark.parametrize( + "stray_fields", + ( + {"prompt": ""}, + {"messages": [], "n": 5}, + {"input": "", "max_tokens": 1}, + {"query": "", "documents": [], "max_completion_tokens": 1}, + {"inferenceConfig": {"maxTokens": 1}}, + ), +) +def test_fields_of_other_apis_on_a_gemini_route_leave_its_hold_alone(stray_fields: dict[str, object]) -> None: + body: Final = {**GEMINI_LONG_CONTENTS, "generationConfig": {"maxOutputTokens": 4000, "candidateCount": 2}} + gemini_hold: Final = _gemini_hold(body) + assert gemini_hold is not None + assert _gemini_hold({**body, **stray_fields}) == gemini_hold + + +@pytest.mark.parametrize("stray_input", ({"prompt": ""}, {"messages": []}, {"input": ""})) +@pytest.mark.parametrize("model", ("gpt-4", "claude-sonnet-5-5")) +@pytest.mark.asyncio +async def test_a_stray_input_field_on_a_gemini_route_never_hides_the_counted_contents( + model: str, stray_input: dict[str, object] +) -> None: + async def counted(body: dict[str, object]) -> dict[str, int]: + return dict( + await count_request_input_tokens( + request_body=body, + route=f"/v1beta/models/{model}:generateContent", + llm_router=None, + raw_body=json.dumps(body).encode(), + ) + ) + + contents_counts: Final = await counted(GEMINI_LONG_CONTENTS) + assert contents_counts[model] > 1000 + assert await counted({**GEMINI_LONG_CONTENTS, **stray_input}) == contents_counts def _tiered_deployment(input_cost_per_token: float) -> Deployment: From a443e4ab9f81f47c46390e179c74e5b42d13b778 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 09:32:58 -0700 Subject: [PATCH 12/15] fix(proxy): hold one candidate for a count too large to price and read system parts like content parts --- .../google_genai/adapters/transformation.py | 11 +------- .../spend_tracking/budget_reservation.py | 15 ++++++++--- litellm/proxy/spend_tracking/input_tokens.py | 4 ++- .../google_genai/test_google_genai_adapter.py | 25 ++++++++++++++++--- .../spend_tracking/test_budget_reservation.py | 22 ++++++++++++++++ .../proxy/spend_tracking/test_input_tokens.py | 21 ++++++++++++++++ 6 files changed, 79 insertions(+), 19 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index d46dae47f06..40ab377dd5e 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -9,7 +9,6 @@ from typing_extensions import ReadOnly, Required, TypedDict from litellm import verbose_logger from litellm.exceptions import BadRequestError -from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types, normalize_tool_schema @@ -250,16 +249,8 @@ def _genai_tool_call(part: object) -> ChatCompletionAssistantToolCall | None: ) -def _system_part_text(part: object) -> str | None: - fields: Final = as_str_mapping(part) - text: Final = fields.get("text") if fields is not None else None - return text if isinstance(text, str) else None - - def _system_texts(parts: object) -> tuple[str, ...]: - return tuple( - text for text in map(_system_part_text, parts if isinstance(parts, Sequence) else ()) if text is not None - ) + return tuple(part["text"] for part in _genai_message_parts(_genai_content_parts(parts)) if part["type"] == "text") def _genai_content_messages(content: object) -> tuple[AllMessageValues, ...]: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 566fbc886ad..352a166fae0 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1449,8 +1449,8 @@ def _max_cost_for_cost_info( tier_rate(tier, "output_cost_per_token"), tier_rate(tier, "output_cost_per_reasoning_token"), ) - return (estimated_input_tokens * tier_rate(tier, "input_cost_per_token")) + ( - output_tokens * output_multiplier * output_rate + return (estimated_input_tokens * tier_rate(tier, "input_cost_per_token")) + _output_cost( + output_tokens, output_multiplier, output_rate ) input_cost_per_token: Final = _to_float(model_info.get("input_cost_per_token")) @@ -1467,7 +1467,7 @@ def _max_cost_for_cost_info( # under-reserving reasoning-heavy requests. output_rate = max(output_cost_per_token or 0.0, output_cost_per_reasoning_token or 0.0) if output_cost_per_token is not None or output_cost_per_reasoning_token is not None: - cost += output_tokens * output_multiplier * output_rate + cost += _output_cost(output_tokens, output_multiplier, output_rate) elif output_tokens > 0: return None @@ -1706,6 +1706,13 @@ def _get_output_multiplier(request_body: Mapping[str, object], route: str) -> in return max((1, *(count for count in map(_to_int, candidates) if count is not None))) +def _output_cost(output_tokens: int, output_multiplier: int, output_rate: float) -> float: + try: + return output_tokens * output_multiplier * output_rate + except OverflowError: + return output_tokens * output_rate + + def _is_input_only_route(route: str) -> bool: return any( route_part in route @@ -1731,7 +1738,7 @@ def _to_int(value: object) -> int | None: return None try: return int(value) - except (TypeError, ValueError): + except (TypeError, ValueError, OverflowError): return None diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 30de3e6f596..dbf90aadfa0 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -266,7 +266,9 @@ def _gemini_content_is_countable(content: Mapping[str, object]) -> bool: def _gemini_parts_are_countable(parts: object, fields: frozenset[str]) -> bool: - return all(_gemini_part_is_countable(part, fields) for part in iter_items(parts)) + single_part: Final = as_str_mapping(parts) + items: Final = (single_part,) if single_part is not None else iter_items(parts) + return all(_gemini_part_is_countable(part, fields) for part in items) def _gemini_part_is_countable(part: object, fields: frozenset[str]) -> bool: diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 1f581a05031..d622a1b85b9 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -215,14 +215,31 @@ def test_every_system_instruction_text_part_reaches_the_system_message( @pytest.mark.parametrize( ("system_parts", "expected_messages"), [ - ( - [{"text": "Answer tersely"}, {"text": None}, {"text": 3}, "text", {"text": "Use British spelling"}], + pytest.param( + [ + {"text": "Answer tersely"}, + {"text": None}, + {"text": 3}, + {"text": ["x"]}, + {"text": "Use British spelling"}, + ], [{"role": "system", "content": "Answer tersely\nUse British spelling"}, {"role": "user", "content": "hi"}], + id="malformed_parts_skipped_like_contents", ), - (None, [{"role": "user", "content": "hi"}]), + pytest.param( + {"text": "Answer tersely"}, + [{"role": "system", "content": "Answer tersely"}, {"role": "user", "content": "hi"}], + id="single_object_parts", + ), + pytest.param( + ["Answer tersely", {"text": "Use British spelling"}], + [{"role": "system", "content": "Answer tersely\nUse British spelling"}, {"role": "user", "content": "hi"}], + id="bare_string_part_read_like_contents", + ), + pytest.param(None, [{"role": "user", "content": "hi"}], id="no_parts"), ], ) -def test_a_system_instruction_part_without_string_text_is_skipped( +def test_system_instruction_parts_follow_the_contents_part_rules( system_parts: object, expected_messages: list[dict[str, str]] ) -> None: request: Final = GoogleGenAIAdapter().translate_generate_content_to_completion( diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 12d509d5198..a719dbca091 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -254,6 +254,7 @@ def _gemini_hold(body: dict[str, object], input_tokens: int | None = None) -> fl ( {"contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT}, GEMINI_AUDIO_PART]}]}, {"contents": [{"role": "user", "parts": [{"fileData": {"mimeType": "video/mp4", "fileUri": "gs://b/v.mp4"}}]}]}, + {"contents": [{"role": "user", "parts": {"fileData": {"mimeType": "video/mp4", "fileUri": "gs://b/v.mp4"}}}]}, {"contents": [{"role": "user", "parts": [{"inline_data": {"mime_type": "image/png", "data": "iVBORw0K"}}]}]}, {"contents": [{"role": "model", "parts": [{"executableCode": {"language": "PYTHON", "code": "print(1)"}}]}]}, {"contents": [GEMINI_TEXT_CONTENT, {"role": "function", "parts": [{"text": GEMINI_PROMPT * 500}]}]}, @@ -266,6 +267,7 @@ def _gemini_hold(body: dict[str, object], input_tokens: int | None = None) -> fl ] }, {"systemInstruction": {"parts": [GEMINI_AUDIO_PART]}, "contents": [GEMINI_TEXT_CONTENT]}, + {"systemInstruction": {"parts": GEMINI_AUDIO_PART}, "contents": [GEMINI_TEXT_CONTENT]}, {"tools": [{"function_declarations": [{"name": "lookup"}]}], "contents": [GEMINI_TEXT_CONTENT]}, {"tools": [{"googleSearch": {}}], "contents": [GEMINI_TEXT_CONTENT]}, {"cachedContent": "cachedContents/abc", "contents": [GEMINI_TEXT_CONTENT]}, @@ -335,6 +337,26 @@ def test_a_gemini_output_cap_of_zero_or_below_holds_the_default_output(cap: int) ) +UNPRICEABLE_COUNTS: Final = (10**400, "9" * 400, 1e308, math.inf) + + +@pytest.mark.parametrize("count", UNPRICEABLE_COUNTS) +def test_a_candidate_count_too_large_to_price_holds_one_candidate(count: object) -> None: + one_candidate_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) + assert one_candidate_hold is not None + assert ( + _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": count}}) + == one_candidate_hold + ) + + +@pytest.mark.parametrize("count", UNPRICEABLE_COUNTS) +def test_an_n_too_large_to_price_holds_one_choice(count: object) -> None: + one_choice_hold: Final = _chat_hold(GEMINI_CHAT_BODY) + assert one_choice_hold is not None + assert _chat_hold({**GEMINI_CHAT_BODY, "n": count}) == one_choice_hold + + def test_a_negative_max_tokens_holds_the_default_output() -> None: uncapped_hold: Final = _chat_hold(GEMINI_CHAT_BODY) assert uncapped_hold is not None diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index 72c9ff4ae55..6800dd8f91b 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -118,3 +118,24 @@ def test_a_stray_contents_field_never_replaces_the_request_input_count(input_fie count: Final = count_input_tokens_for_model(request_body=body, model=CL100K_MODEL) assert count is not None and count > 1000 assert count_input_tokens_for_model(request_body={**body, "contents": []}, model=CL100K_MODEL) == count + + +GEMINI_TEXT: Final = "count every one of these words " * 200 +GEMINI_GREETING: Final = {"role": "user", "parts": [{"text": "hi"}]} + + +def _gemini_count(body: dict[str, object], model: str) -> int | None: + return count_input_tokens_for_model(request_body=body, model=model, gemini_route=True) + + +@pytest.mark.parametrize("parts", ([GEMINI_TEXT], {"text": GEMINI_TEXT}), ids=("bare_string", "single_object")) +@pytest.mark.parametrize("model", (CL100K_MODEL, "claude-sonnet-5-5")) +def test_gemini_parts_count_the_same_in_every_shape_the_adapter_reads(model: str, parts: object) -> None: + listed_contents: Final = _gemini_count({"contents": [{"role": "user", "parts": [{"text": GEMINI_TEXT}]}]}, model) + listed_system: Final = _gemini_count( + {"systemInstruction": {"parts": [{"text": GEMINI_TEXT}]}, "contents": [GEMINI_GREETING]}, model + ) + assert listed_contents is not None and listed_contents > 1000 + assert listed_system is not None and listed_system > 1000 + assert _gemini_count({"contents": [{"role": "user", "parts": parts}]}, model) == listed_contents + assert _gemini_count({"systemInstruction": {"parts": parts}, "contents": [GEMINI_GREETING]}, model) == listed_system From 0c66c95f4629937dda3f4cc9e31f8caec1e9110c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 14:00:25 -0700 Subject: [PATCH 13/15] fix(proxy): keep the spend already counted when a budget hold is far over the budget --- .../spend_tracking/budget_reservation.py | 25 ++++++++++--------- .../spend_tracking/test_budget_reservation.py | 25 +++++++++++++++++++ tests/unit/proxy/test_budget_reservation.py | 2 +- 3 files changed, 39 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 352a166fae0..c711f8b5a62 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -180,7 +180,7 @@ async def _apply_over_budget_reservation_policy( entry: dict[str, float | str], applied_entries: list[dict[str, float | str]], reservation_cost: float, - current_spend: float, + spend_before_reservation: float, fail_closed_budget_enforcement: bool = False, ) -> float: """ @@ -196,13 +196,13 @@ async def _apply_over_budget_reservation_policy( applied_entries.remove(entry) return reservation_cost - remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost) + remaining_before_reservation: Final = counter.max_budget - spend_before_reservation if remaining_before_reservation <= 1e-12: - _raise_counter_budget_exceeded(counter=counter, current_cost=current_spend) - if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12: + _raise_counter_budget_exceeded(counter=counter, current_cost=spend_before_reservation + reservation_cost) + if fail_closed_budget_enforcement and reservation_cost - remaining_before_reservation > 1e-12: _raise_counter_budget_exceeded( counter=counter, - current_cost=current_spend - reservation_cost, + current_cost=spend_before_reservation, estimated_cost=reservation_cost, ) await _resize_applied_reservation( @@ -992,7 +992,7 @@ async def _reserve_reservable_counters( current_spend_by_counter_key=current_spend_by_counter_key, reservation_cost=reservation_cost, ): - charged_cost = reservation_cost + charged_cost = min(reservation_cost, max(counter.max_budget for counter in group)) entries = tuple(_counter_to_reservation_entry(counter=counter, reserved_cost=charged_cost) for counter in group) applied_entries.extend(entries) reserved_values = await _reserve_counters(counters=group, entries=entries, reservation_cost=charged_cost) @@ -1005,18 +1005,19 @@ async def _reserve_reservable_counters( for counter, entry, reserved_value in zip(group, entries, reserved_values): if entry not in applied_entries: continue - if reserved_value is not None: - current_spend = reserved_value - (charged_cost - reservation_cost) - else: - current_spend = current_spend_by_counter_key[counter.counter_key] + reservation_cost - if current_spend > counter.max_budget: + spend_before_reservation = ( + current_spend_by_counter_key[counter.counter_key] + if reserved_value is None + else reserved_value - charged_cost + ) + if spend_before_reservation + reservation_cost > counter.max_budget: reservation_cost = await _apply_over_budget_reservation_policy( counter=counter, valid_token=valid_token, entry=entry, applied_entries=applied_entries, reservation_cost=reservation_cost, - current_spend=current_spend, + spend_before_reservation=spend_before_reservation, fail_closed_budget_enforcement=fail_closed_budget_enforcement, ) return reservation_cost diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index a719dbca091..744d20ab1c7 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -357,6 +357,31 @@ def test_an_n_too_large_to_price_holds_one_choice(count: object) -> None: assert _chat_hold({**GEMINI_CHAT_BODY, "n": count}) == one_choice_hold +HOLDS_FAR_OVER_A_TINY_BUDGET: Final = ( + (GEMINI_ROUTE, {"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": 1e300}}), + ("/v1/chat/completions", {**GEMINI_CHAT_BODY, "n": 1e300}), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("route", "request_body"), HOLDS_FAR_OVER_A_TINY_BUDGET) +async def test_a_hold_far_over_the_budget_keeps_the_spend_already_counted( + spend_counter_cache: DualCache, route: str, request_body: dict[str, object] +) -> None: + counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}" + counted_spend: Final = 0.004 + hold_estimate: Final = estimate_request_max_cost(request_body=request_body, route=route, llm_router=None) + assert hold_estimate is not None and 1e20 < hold_estimate < math.inf + await spend_counter_cache.async_set_cache(key=counter_key, value=counted_spend) + + reservation: Final = await _reserve_for_tiny_budget_key(route, request_body) + + assert reservation is not None + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.01) + await release_unbound_budget_reservation(reservation) + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(counted_spend) + + def test_a_negative_max_tokens_holds_the_default_output() -> None: uncapped_hold: Final = _chat_hold(GEMINI_CHAT_BODY) assert uncapped_hold is not None diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 535c3da8c3a..25e0ee6a0ea 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -207,7 +207,7 @@ async def test_over_budget_window_counter_tags_clean_entity_id(): entry={"counter_key": counter.counter_key}, applied_entries=[], reservation_cost=0.5, - current_spend=2.0, + spend_before_reservation=1.5, ) assert exc_info.value.entity_type == "key" assert exc_info.value.entity_id == "test-token" From e312dabac21624d2bbe0986d99d1b41b4de54393 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 19:25:58 -0700 Subject: [PATCH 14/15] fix(proxy): count candidateCount in the budget hold only for models that generate the candidates A non-Gemini model reached through /v1beta generateContent never receives candidateCount, so a huge count held the whole budget for the length of the call while the model produced one answer --- .../spend_tracking/budget_reservation.py | 56 ++++++++++++++++--- .../spend_tracking/test_budget_reservation.py | 37 ++++++++++++ 2 files changed, 86 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index c711f8b5a62..f6074762d5f 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -51,6 +51,7 @@ from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.types.router import DeploymentTypedDict +from litellm.utils import ProviderConfigManager @dataclass @@ -1398,12 +1399,16 @@ def _estimate_request_max_cost_for_model( llm_router: Router | None, input_tokens: int | None = None, ) -> float | None: + output_multiplier: Final = _get_output_multiplier( + _requested_output_counts(request_body=request_body, route=route, model=model, llm_router=llm_router) + ) estimates: Final = [ _max_cost_for_cost_info( request_body=request_body, route=route, model=model, model_info=model_info, + output_multiplier=output_multiplier, input_tokens=input_tokens, ) for model_info in _get_model_cost_infos(model=model, llm_router=llm_router) @@ -1417,6 +1422,7 @@ def _max_cost_for_cost_info( route: str, model: str, model_info: Mapping[str, object], + output_multiplier: int, input_tokens: int | None = None, ) -> float | None: image_cost: Final = _estimate_image_generation_cost( @@ -1441,7 +1447,6 @@ def _max_cost_for_cost_info( if estimated_input_tokens is None or output_tokens is None: return None - output_multiplier: Final = _get_output_multiplier(request_body=request_body, route=route) tiered_pricing: Final = model_info.get("tiered_pricing") if isinstance(tiered_pricing, list) and tiered_pricing: tier: Final = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=estimated_input_tokens) @@ -1698,13 +1703,50 @@ def _requested_output_tokens(request_body: Mapping[str, object], route: str) -> return next((tokens for tokens in map(_to_int, candidates) if tokens is not None and tokens >= 0), None) -def _get_output_multiplier(request_body: Mapping[str, object], route: str) -> int: - candidates: Final = ( - _gemini_candidate_counts(request_body) - if RouteChecks.is_generate_content_route(route) - else (request_body.get("n"), request_body.get("best_of")) +def _requested_output_counts( + request_body: Mapping[str, object], + route: str, + model: str, + llm_router: Router | None, +) -> tuple[object, ...]: + if not RouteChecks.is_generate_content_route(route): + return (request_body.get("n"), request_body.get("best_of")) + candidate_counts: Final = _gemini_candidate_counts(request_body) + if all(count is None for count in candidate_counts): + return () + return candidate_counts if _takes_candidate_count(model=model, llm_router=llm_router) else () + + +def _takes_candidate_count(model: str, llm_router: Router | None) -> bool: + deployments: Final = (llm_router.get_model_list(model_name=model) if llm_router is not None else None) or () + deployment_params: Final = tuple( + params for deployment in deployments if (params := as_str_mapping(_get_value(deployment, "litellm_params"))) ) - return max((1, *(count for count in map(_to_int, candidates) if count is not None))) + return any(_serves_generate_content_natively(params) for params in deployment_params or ({"model": model},)) + + +def _serves_generate_content_natively(litellm_params: Mapping[str, object]) -> bool: + model: Final = litellm_params.get("model") + custom_llm_provider: Final = litellm_params.get("custom_llm_provider") + if not isinstance(model, str): + return True + try: + backend_model, provider, _, _ = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + ) + native_config: Final = ProviderConfigManager.get_provider_google_genai_generate_content_config( + model=backend_model, + provider=litellm.LlmProviders(provider), + litellm_params={key: litellm_params[key] for key in ("base_model", "api_base") if key in litellm_params}, + ) + except (litellm.BadRequestError, ValueError): + return True + return native_config is not None + + +def _get_output_multiplier(requested_counts: tuple[object, ...]) -> int: + return max((1, *(count for count in map(_to_int, requested_counts) if count is not None))) def _output_cost(output_tokens: int, output_multiplier: int, output_rate: float) -> float: diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 744d20ab1c7..f5f7b1c6c70 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -357,6 +357,43 @@ def test_an_n_too_large_to_price_holds_one_choice(count: object) -> None: assert _chat_hold({**GEMINI_CHAT_BODY, "n": count}) == one_choice_hold +def _group_generate_content_hold(router: Router, generation_config: dict[str, object]) -> float | None: + return estimate_request_max_cost( + request_body={"contents": [GEMINI_TEXT_CONTENT], "generationConfig": generation_config}, + route="/v1beta/models/assistant:generateContent", + llm_router=router, + input_token_counts={"assistant": 10}, + ) + + +@pytest.mark.parametrize( + ("backend_models", "candidates_are_generated"), + ( + (("vertex_ai/claude-haiku-5-5",), False), + (("anthropic/claude-haiku-5-5",), False), + (("gemini/gemini-3.8-flash",), True), + (("vertex_ai/gemini-3.8-flash",), True), + (("vertex_ai/claude-haiku-5-5", "gemini/gemini-3.8-flash"), True), + ), +) +def test_a_candidate_count_is_held_only_when_a_deployment_generates_the_candidates( + backend_models: tuple[str, ...], candidates_are_generated: bool +) -> None: + router: Final = Router( + model_list=[ + {"model_name": "assistant", "litellm_params": {"model": model, "api_key": "sk-test"}} + for model in backend_models + ] + ) + one_candidate_hold: Final = _group_generate_content_hold(router, {"maxOutputTokens": 100}) + three_candidate_hold: Final = _group_generate_content_hold(router, {"maxOutputTokens": 100, "candidateCount": 3}) + + assert one_candidate_hold is not None + assert three_candidate_hold is not None + assert (three_candidate_hold > one_candidate_hold) is candidates_are_generated + assert three_candidate_hold >= one_candidate_hold + + HOLDS_FAR_OVER_A_TINY_BUDGET: Final = ( (GEMINI_ROUTE, {"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": 1e300}}), ("/v1/chat/completions", {**GEMINI_CHAT_BODY, "n": 1e300}), From ae4f9e72ac55d37bdcd5abdf804a3853039941c0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:32:44 -0700 Subject: [PATCH 15/15] fix(proxy): hold a candidate count only when every deployment generates it, and drop n when every deployment drops it --- .../spend_tracking/budget_reservation.py | 52 +++++++++++--- .../spend_tracking/test_budget_reservation.py | 69 ++++++++++++++++--- 2 files changed, 102 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 962b819b68e..76f0d810e6c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -4,7 +4,7 @@ import asyncio import json import math import time -from collections.abc import AsyncIterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone from itertools import chain, product @@ -1730,41 +1730,71 @@ def _requested_output_counts( llm_router: Router | None, ) -> tuple[object, ...]: if not RouteChecks.is_generate_content_route(route): - return (request_body.get("n"), request_body.get("best_of")) + n: Final = request_body.get("n") + dropped: Final = n is not None and _every_deployment(_drops_n, _deployment_params(model, llm_router)) + return (None if dropped else n, request_body.get("best_of")) candidate_counts: Final = _gemini_candidate_counts(request_body) if all(count is None for count in candidate_counts): return () - return candidate_counts if _takes_candidate_count(model=model, llm_router=llm_router) else () + generated: Final = _every_deployment(_serves_generate_content_natively, _deployment_params(model, llm_router)) + return candidate_counts if generated else () -def _takes_candidate_count(model: str, llm_router: Router | None) -> bool: - deployments: Final = (llm_router.registry().deployments(model) if llm_router is not None else None) or () - deployment_params: Final = tuple( +def _deployment_params(model: str, llm_router: Router | None) -> tuple[Mapping[str, object], ...]: + if llm_router is None: + return ({"model": model},) + deployments: Final = llm_router.registry().deployments(model) or () + return tuple( params for deployment in deployments if (params := as_str_mapping(_get_value(deployment, "litellm_params"))) ) - return any(_serves_generate_content_natively(params) for params in deployment_params or ({"model": model},)) -def _serves_generate_content_natively(litellm_params: Mapping[str, object]) -> bool: +def _every_deployment( + check: Callable[[Mapping[str, object]], bool], deployment_params: tuple[Mapping[str, object], ...] +) -> bool: + return bool(deployment_params) and all(map(check, deployment_params)) + + +def _resolved_provider(litellm_params: Mapping[str, object]) -> tuple[str, str] | None: model: Final = litellm_params.get("model") custom_llm_provider: Final = litellm_params.get("custom_llm_provider") if not isinstance(model, str): - return True + return None try: backend_model, provider, _, _ = litellm.get_llm_provider( model=model, custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, ) + except litellm.BadRequestError: + return None + return backend_model, provider + + +def _serves_generate_content_natively(litellm_params: Mapping[str, object]) -> bool: + resolved: Final = _resolved_provider(litellm_params) + if resolved is None: + return False + backend_model, provider = resolved + try: native_config: Final = ProviderConfigManager.get_provider_google_genai_generate_content_config( model=backend_model, provider=litellm.LlmProviders(provider), litellm_params={key: litellm_params[key] for key in ("base_model", "api_base") if key in litellm_params}, ) - except (litellm.BadRequestError, ValueError): - return True + except ValueError: + return False return native_config is not None +def _drops_n(litellm_params: Mapping[str, object]) -> bool: + resolved: Final = _resolved_provider(litellm_params) + if resolved is None: + return False + backend_model, provider = resolved + supported_params: Final = litellm.get_supported_openai_params(model=backend_model, custom_llm_provider=provider) + return supported_params is not None and "n" not in supported_params + + def _get_output_multiplier(requested_counts: tuple[object, ...]) -> int: return max((1, *(count for count in map(_to_int, requested_counts) if count is not None))) diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 10358d7e6a7..c2627f9faf5 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -366,6 +366,15 @@ def _group_generate_content_hold(router: Router, generation_config: dict[str, ob ) +def _assistant_router(*backend_models: str) -> Router: + return Router( + model_list=[ + {"model_name": "assistant", "litellm_params": {"model": model, "api_key": "sk-test"}} + for model in backend_models + ] + ) + + @pytest.mark.parametrize( ("backend_models", "candidates_are_generated"), ( @@ -373,18 +382,15 @@ def _group_generate_content_hold(router: Router, generation_config: dict[str, ob (("anthropic/claude-haiku-5-5",), False), (("gemini/gemini-3.8-flash",), True), (("vertex_ai/gemini-3.8-flash",), True), - (("vertex_ai/claude-haiku-5-5", "gemini/gemini-3.8-flash"), True), + (("gemini/gemini-3.8-flash", "vertex_ai/gemini-3.8-flash"), True), + (("vertex_ai/claude-haiku-5-5", "gemini/gemini-3.8-flash"), False), + (("gemini/gemini-3.8-flash", "anthropic/claude-haiku-5-5"), False), ), ) -def test_a_candidate_count_is_held_only_when_a_deployment_generates_the_candidates( +def test_a_candidate_count_is_held_only_when_every_deployment_generates_the_candidates( backend_models: tuple[str, ...], candidates_are_generated: bool ) -> None: - router: Final = Router( - model_list=[ - {"model_name": "assistant", "litellm_params": {"model": model, "api_key": "sk-test"}} - for model in backend_models - ] - ) + router: Final = _assistant_router(*backend_models) one_candidate_hold: Final = _group_generate_content_hold(router, {"maxOutputTokens": 100}) three_candidate_hold: Final = _group_generate_content_hold(router, {"maxOutputTokens": 100, "candidateCount": 3}) @@ -394,6 +400,53 @@ def test_a_candidate_count_is_held_only_when_a_deployment_generates_the_candidat assert three_candidate_hold >= one_candidate_hold +def test_a_candidate_count_for_a_gemini_name_no_deployment_serves_holds_one_candidate() -> None: + router: Final = _assistant_router("anthropic/claude-haiku-5-5") + + def hold(generation_config: dict[str, object]) -> float | None: + return estimate_request_max_cost( + request_body={"contents": [GEMINI_TEXT_CONTENT], "generationConfig": generation_config}, + route=GEMINI_ROUTE, + llm_router=router, + input_token_counts={"gemini-3.8-flash": 10}, + ) + + one_candidate_hold: Final = hold({"maxOutputTokens": 100}) + assert one_candidate_hold is not None + assert hold({"maxOutputTokens": 100, "candidateCount": 3}) == one_candidate_hold + + +@pytest.mark.parametrize( + ("backend_models", "choices_are_generated"), + ( + (("anthropic/claude-haiku-5-5",), False), + (("vertex_ai/claude-haiku-5-5", "anthropic/claude-haiku-5-5"), False), + (("gemini/gemini-3.8-flash",), True), + (("anthropic/claude-haiku-5-5", "gemini/gemini-3.8-flash"), True), + ), +) +def test_an_n_is_held_unless_every_deployment_drops_it( + backend_models: tuple[str, ...], choices_are_generated: bool +) -> None: + router: Final = _assistant_router(*backend_models) + + def hold(output_fields: dict[str, object]) -> float | None: + return estimate_request_max_cost( + request_body={"model": "assistant", "messages": GEMINI_CHAT_BODY["messages"], **output_fields}, + route="/v1/chat/completions", + llm_router=router, + input_token_counts={"assistant": 10}, + ) + + one_choice_hold: Final = hold({"max_tokens": 100}) + three_choice_hold: Final = hold({"max_tokens": 100, "n": 3}) + + assert one_choice_hold is not None + assert three_choice_hold is not None + assert (three_choice_hold > one_choice_hold) is choices_are_generated + assert three_choice_hold >= one_choice_hold + + HOLDS_FAR_OVER_A_TINY_BUDGET: Final = ( (GEMINI_ROUTE, {"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": 1e300}}), ("/v1/chat/completions", {**GEMINI_CHAT_BODY, "n": 1e300}),