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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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}), From 0a560278683b81a6a669151fe326341e7f2c3154 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 01:36:19 -0700 Subject: [PATCH 16/23] fix(proxy): count candidateCount when any deployment generates it and always count n The stricter rules from the previous commit let a key spend past max_budget through a Gemini name no deployment serves, an allowed_openai_params n, and a router fallback. Restore the earlier rules and add regression tests for those three shapes. --- .../spend_tracking/budget_reservation.py | 52 ++------- .../spend_tracking/test_budget_reservation.py | 103 ++++++++++-------- 2 files changed, 70 insertions(+), 85 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 76f0d810e6c..962b819b68e 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, Callable, Mapping, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone from itertools import chain, product @@ -1730,71 +1730,41 @@ def _requested_output_counts( llm_router: Router | None, ) -> tuple[object, ...]: if not RouteChecks.is_generate_content_route(route): - 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")) + 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 () - generated: Final = _every_deployment(_serves_generate_content_natively, _deployment_params(model, llm_router)) - return candidate_counts if generated else () + return candidate_counts if _takes_candidate_count(model=model, llm_router=llm_router) else () -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( +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( 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 _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: +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 None + 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, ) - 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 ValueError: - return False + except (litellm.BadRequestError, ValueError): + return True 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 c2627f9faf5..2fe064e21ae 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -366,15 +366,6 @@ 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"), ( @@ -382,15 +373,18 @@ def _assistant_router(*backend_models: str) -> Router: (("anthropic/claude-haiku-5-5",), False), (("gemini/gemini-3.8-flash",), True), (("vertex_ai/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), + (("vertex_ai/claude-haiku-5-5", "gemini/gemini-3.8-flash"), True), ), ) -def test_a_candidate_count_is_held_only_when_every_deployment_generates_the_candidates( +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 = _assistant_router(*backend_models) + 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}) @@ -400,51 +394,72 @@ def test_a_candidate_count_is_held_only_when_every_deployment_generates_the_cand 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 test_a_candidate_count_for_a_gemini_name_no_deployment_serves_is_held() -> None: + router: Final = Router( + model_list=[ + {"model_name": "assistant", "litellm_params": {"model": "anthropic/claude-haiku-5-5", "api_key": "sk-test"}} + ] + ) + one_candidate_hold: Final = estimate_request_max_cost( + request_body={"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": 100}}, + route=GEMINI_ROUTE, + llm_router=router, + input_token_counts={"gemini-3.8-flash": 10}, + ) + three_candidate_hold: Final = estimate_request_max_cost( + request_body={ + "contents": [GEMINI_TEXT_CONTENT], + "generationConfig": {"maxOutputTokens": 100, "candidateCount": 3}, + }, + route=GEMINI_ROUTE, + llm_router=router, + input_token_counts={"gemini-3.8-flash": 10}, + ) - 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 + assert three_candidate_hold is not None + assert three_candidate_hold > 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), +N_SENDING_GROUPS: Final = ( + Router( + model_list=[ + {"model_name": "assistant", "litellm_params": {"model": "anthropic/claude-haiku-5-5", "api_key": "sk"}}, + {"model_name": "backup", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk"}}, + ], + fallbacks=[{"assistant": ["backup"]}], + ), + Router( + model_list=[ + { + "model_name": "assistant", + "litellm_params": { + "model": "mistral/mistral-small-latest", + "api_key": "sk", + "allowed_openai_params": ["n"], + }, + } + ] ), ) -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( + +@pytest.mark.parametrize("router", N_SENDING_GROUPS) +def test_an_n_is_held_when_a_fallback_or_an_allowlist_can_still_send_it(router: Router) -> None: + holds: Final = tuple( + 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}, ) + for output_fields in ({"max_tokens": 100}, {"max_tokens": 100, "n": 3}) + ) - one_choice_hold: Final = hold({"max_tokens": 100}) - three_choice_hold: Final = hold({"max_tokens": 100, "n": 3}) - + one_choice_hold, three_choice_hold = holds 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 + assert three_choice_hold > one_choice_hold HOLDS_FAR_OVER_A_TINY_BUDGET: Final = ( From 819fd1616667c7a0ae3c113b3d1bbd006cd0116e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:03:04 -0700 Subject: [PATCH 17/23] fix(proxy): cap the candidateCount hold at main's default output hold A candidateCount like 1e300 sent to a mixed Gemini and Claude group, or through a key alias from a Gemini name to Claude, held a whole key while Claude served the call and answered a strict key 422. On generateContent routes the counted candidates now raise the output hold at most to the default output hold main used for the model, so no count makes a call hold more than on main. Chat n is unchanged. --- .../spend_tracking/budget_reservation.py | 26 ++++++++++++++---- .../spend_tracking/test_budget_reservation.py | 27 +++++++++++-------- 2 files changed, 37 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 962b819b68e..9f5a48831fc 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1451,6 +1451,9 @@ def _max_cost_for_cost_info( ) if estimated_input_tokens is None or output_tokens is None: return None + held_output_tokens: Final = _held_output_tokens( + route=route, model_info=model_info, output_tokens=output_tokens, output_multiplier=output_multiplier + ) tiered_pricing: Final = model_info.get("tiered_pricing") if isinstance(tiered_pricing, list) and tiered_pricing: @@ -1461,7 +1464,7 @@ def _max_cost_for_cost_info( tier_rate(tier, "output_cost_per_reasoning_token"), ) return (estimated_input_tokens * tier_rate(tier, "input_cost_per_token")) + _output_cost( - output_tokens, output_multiplier, output_rate + held_output_tokens, output_tokens, output_rate ) input_cost_per_token: Final = _to_float(model_info.get("input_cost_per_token")) @@ -1478,7 +1481,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_cost(output_tokens, output_multiplier, output_rate) + cost += _output_cost(held_output_tokens, output_tokens, output_rate) elif output_tokens > 0: return None @@ -1692,7 +1695,7 @@ def _estimate_output_tokens( # the reservation up to remaining team headroom and pin the counter # at the cap — the model can only physically emit max_output_tokens # anyway, so reserving more is both wasteful and a DoS surface. - model_ceiling: Final = _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK + model_ceiling: Final = _model_output_ceiling(model_info) return min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK if requested is None else requested, model_ceiling) @@ -1769,13 +1772,26 @@ 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: +def _held_output_tokens( + route: str, model_info: Mapping[str, object], output_tokens: int, output_multiplier: int +) -> int: + if not RouteChecks.is_generate_content_route(route): + return output_tokens * output_multiplier + default_output_tokens: Final = min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, _model_output_ceiling(model_info)) + return max(output_tokens, min(output_tokens * output_multiplier, default_output_tokens)) + + +def _output_cost(held_output_tokens: int, output_tokens: int, output_rate: float) -> float: try: - return output_tokens * output_multiplier * output_rate + return held_output_tokens * output_rate except OverflowError: return output_tokens * output_rate +def _model_output_ceiling(model_info: Mapping[str, object]) -> int: + return _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK + + def _is_input_only_route(route: str) -> bool: return any( route_part in route diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 2fe064e21ae..2a3258e7555 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -196,8 +196,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}}, {"n": 3}), - ({"generationConfig": {"candidate_count": 3}}, {"n": 3}), + ({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {}), + ({"generationConfig": {"maxOutputTokens": 5, "candidate_count": 3}}, {"max_tokens": 5, "n": 3}), ({}, {}), ), ) @@ -340,14 +340,20 @@ 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( + ("generation_config", "same_hold_config"), + ( + *(({"candidateCount": count}, {}) for count in (*UNPRICEABLE_COUNTS, 8)), + ({"maxOutputTokens": 8000, "candidateCount": 4}, {}), + ({"maxOutputTokens": 20000, "candidateCount": 2}, {"maxOutputTokens": 20000}), + ), +) +def test_candidates_hold_at_most_one_candidate_at_the_default_or_requested_output( + generation_config: dict[str, object], same_hold_config: dict[str, object] +) -> None: + expected_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": same_hold_config}) + assert expected_hold is not None + assert _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": generation_config}) == expected_hold @pytest.mark.parametrize("count", UNPRICEABLE_COUNTS) @@ -463,7 +469,6 @@ def test_an_n_is_held_when_a_fallback_or_an_allowlist_can_still_send_it(router: 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 1fd653ce1fc513701b5f64344d6d42ef6770e3da Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:12:21 -0700 Subject: [PATCH 18/23] fix(proxy): hold every Gemini candidate up to the merge base's whole hold The previous cap stopped the candidateCount output hold at the default 16384 tokens. Because the hold now counts the Gemini prompt instead of the whole input window, a strict key admitted a call like 4 candidates of 8000 tokens and then spent past max_budget, where the merge base refused it. Every candidate is now held at its output cap, like chat n, and the whole hold is clamped at what the merge base held for the same call (the model's input window plus the default output), never below one candidate's hold. A huge count still holds no more than the merge base did. Chat n is unchanged --- .../spend_tracking/budget_reservation.py | 69 +++++++++------ .../spend_tracking/test_budget_reservation.py | 86 ++++++++++++++++--- 2 files changed, 119 insertions(+), 36 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 9f5a48831fc..78968dd897d 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1451,41 +1451,67 @@ def _max_cost_for_cost_info( ) if estimated_input_tokens is None or output_tokens is None: return None - held_output_tokens: Final = _held_output_tokens( - route=route, model_info=model_info, output_tokens=output_tokens, output_multiplier=output_multiplier + candidates_hold: Final = _priced_hold( + model_info=model_info, + input_tokens=estimated_input_tokens, + held_output_tokens=output_tokens * output_multiplier, + output_tokens=output_tokens, + ) + if candidates_hold is None or output_multiplier == 1 or not RouteChecks.is_generate_content_route(route): + return candidates_hold + one_candidate_hold: Final = _priced_hold( + model_info=model_info, + input_tokens=estimated_input_tokens, + held_output_tokens=output_tokens, + output_tokens=output_tokens, + ) + main_hold: Final = _main_generate_content_hold(model_info) + if one_candidate_hold is None or main_hold is None: + return one_candidate_hold + return max(one_candidate_hold, min(candidates_hold, main_hold)) + + +def _main_generate_content_hold(model_info: Mapping[str, object]) -> float | None: + context_window: Final = _to_int(model_info.get("max_input_tokens")) + if context_window is None: + return None + default_output_tokens: Final = min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, _model_output_ceiling(model_info)) + return _priced_hold( + model_info=model_info, + input_tokens=context_window, + held_output_tokens=default_output_tokens, + output_tokens=default_output_tokens, ) + +def _priced_hold( + model_info: Mapping[str, object], input_tokens: int, held_output_tokens: int, output_tokens: int +) -> float | None: 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) + tier: Final = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=input_tokens) if tier is not None: - output_rate = max( + tier_output_rate: Final = max( 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_cost( - held_output_tokens, output_tokens, output_rate + return (input_tokens * tier_rate(tier, "input_cost_per_token")) + _output_cost( + held_output_tokens, output_tokens, tier_output_rate ) input_cost_per_token: Final = _to_float(model_info.get("input_cost_per_token")) output_cost_per_token: Final = _to_float(model_info.get("output_cost_per_token")) output_cost_per_reasoning_token: Final = _to_float(model_info.get("output_cost_per_reasoning_token")) - cost = 0.0 - if input_cost_per_token is not None: - cost += estimated_input_tokens * input_cost_per_token - elif estimated_input_tokens > 0: + if input_cost_per_token is None and input_tokens > 0: + return None + if output_cost_per_token is None and output_cost_per_reasoning_token is None and output_tokens > 0: return None # The reasoning-token share is unknown before the request runs, so reserve every # output token at the higher of the standard and reasoning rates to avoid # 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_cost(held_output_tokens, output_tokens, output_rate) - elif output_tokens > 0: - return None - - return cost + output_rate: Final = max(output_cost_per_token or 0.0, output_cost_per_reasoning_token or 0.0) + return input_tokens * (input_cost_per_token or 0.0) + _output_cost(held_output_tokens, output_tokens, output_rate) def _estimate_image_generation_cost( @@ -1772,15 +1798,6 @@ 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 _held_output_tokens( - route: str, model_info: Mapping[str, object], output_tokens: int, output_multiplier: int -) -> int: - if not RouteChecks.is_generate_content_route(route): - return output_tokens * output_multiplier - default_output_tokens: Final = min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, _model_output_ceiling(model_info)) - return max(output_tokens, min(output_tokens * output_multiplier, default_output_tokens)) - - def _output_cost(held_output_tokens: int, output_tokens: int, output_rate: float) -> float: try: return held_output_tokens * output_rate diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 2a3258e7555..3dc7ef7c226 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -25,6 +25,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_reservation_cache_key, ) from litellm.proxy.spend_tracking.budget_reservation import ( + DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, _get_team_member_budget_counter, estimate_request_input_cost, estimate_request_max_cost, @@ -196,8 +197,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}}, {}), - ({"generationConfig": {"maxOutputTokens": 5, "candidate_count": 3}}, {"max_tokens": 5, "n": 3}), + ({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {"n": 3}), + ({"generationConfig": {"candidate_count": 3}}, {"n": 3}), ({}, {}), ), ) @@ -340,20 +341,85 @@ 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 + ) + + +WIDE_WINDOW: Final = {"max_input_tokens": 1_000_000, "max_output_tokens": 100_000} +LOW_OUTPUT_CEILING: Final = {**WIDE_WINDOW, "max_output_tokens": 1000} + + +def _priced_gemini_hold( + model_info: dict[str, object], generation_config: dict[str, object], input_tokens: int = 10 +) -> float | None: + router: Final = Router( + model_list=[ + { + "model_name": "assistant", + "litellm_params": { + "model": "vertex_ai/gemini-3.8-flash", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 4e-6, + }, + "model_info": model_info, + } + ] + ) + 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": input_tokens}, + ) + + @pytest.mark.parametrize( - ("generation_config", "same_hold_config"), + ("model_info", "generation_config", "same_hold_model_info", "same_hold_config"), ( - *(({"candidateCount": count}, {}) for count in (*UNPRICEABLE_COUNTS, 8)), - ({"maxOutputTokens": 8000, "candidateCount": 4}, {}), - ({"maxOutputTokens": 20000, "candidateCount": 2}, {"maxOutputTokens": 20000}), + (WIDE_WINDOW, {"maxOutputTokens": 8000, "candidateCount": 4}, WIDE_WINDOW, {"maxOutputTokens": 32000}), + (WIDE_WINDOW, {"maxOutputTokens": 20000, "candidate_count": 2}, WIDE_WINDOW, {"maxOutputTokens": 40000}), + (WIDE_WINDOW, {"candidateCount": 2}, WIDE_WINDOW, {"maxOutputTokens": 2 * DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK}), + (LOW_OUTPUT_CEILING, {"candidateCount": 4}, {**WIDE_WINDOW, "max_output_tokens": 4000}, {}), ), ) -def test_candidates_hold_at_most_one_candidate_at_the_default_or_requested_output( - generation_config: dict[str, object], same_hold_config: dict[str, object] +def test_candidates_hold_what_one_candidate_writing_all_their_output_holds( + model_info: dict[str, object], + generation_config: dict[str, object], + same_hold_model_info: dict[str, object], + same_hold_config: dict[str, object], ) -> None: - expected_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": same_hold_config}) + expected_hold: Final = _priced_gemini_hold(same_hold_model_info, same_hold_config) assert expected_hold is not None - assert _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": generation_config}) == expected_hold + assert _priced_gemini_hold(model_info, generation_config) == expected_hold + + +@pytest.mark.parametrize( + ("model_info", "generation_config"), + ( + (WIDE_WINDOW, {"candidateCount": 1e300}), + (WIDE_WINDOW, {"maxOutputTokens": 100_000, "candidateCount": 8}), + (LOW_OUTPUT_CEILING, {"candidateCount": 1e300}), + ), +) +def test_candidates_never_hold_more_than_the_whole_context_window_hold( + model_info: dict[str, object], generation_config: dict[str, object] +) -> None: + context_window_hold: Final = _priced_gemini_hold(model_info, {}, input_tokens=1_000_000) + assert context_window_hold is not None + assert _priced_gemini_hold(model_info, generation_config) == context_window_hold + + +def test_candidates_never_hold_less_than_one_candidate() -> None: + narrow_window: Final = {"max_input_tokens": 10, "max_output_tokens": 100_000} + one_candidate_hold: Final = _priced_gemini_hold(narrow_window, {"maxOutputTokens": 50_000}) + assert one_candidate_hold is not None + assert _priced_gemini_hold(narrow_window, {"maxOutputTokens": 50_000, "candidateCount": 2}) == one_candidate_hold @pytest.mark.parametrize("count", UNPRICEABLE_COUNTS) From 6bb359064325bd94f0d95a57a210961e32ad01e6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:25:15 -0700 Subject: [PATCH 19/23] fix(proxy): let a Gemini output cap that /v1beta drops only raise the hold The Vertex and Gemini pass-throughs forward every spelling of the output cap, but /v1beta passes only generationConfig.maxOutputTokens to other models. The hold read only that spelling, so a snake_case cap above the default output held less than the pass-through could spend. A cap in any other spelling now raises the hold to it and never lowers it --- .../spend_tracking/budget_reservation.py | 27 +++++++++++++++++-- .../spend_tracking/test_budget_reservation.py | 25 +++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 78968dd897d..f45f9028994 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1743,8 +1743,7 @@ def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object 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 + return _gemini_requested_output_tokens(request_body) candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), _nested_field(request_body, "inferenceConfig", "maxTokens"), @@ -1752,6 +1751,30 @@ 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) +_GEMINI_PASS_THROUGH_ONLY_CAPS: Final = ( + ("generationConfig", "max_output_tokens"), + ("generation_config", "maxOutputTokens"), + ("generation_config", "max_output_tokens"), +) + + +def _gemini_requested_output_tokens(request_body: Mapping[str, object]) -> int | None: + gemini_cap: Final = _positive_int(_nested_field(request_body, "generationConfig", "maxOutputTokens")) + pass_through_caps: Final = tuple( + cap + for config_field, field in _GEMINI_PASS_THROUGH_ONLY_CAPS + if (cap := _positive_int(_nested_field(request_body, config_field, field))) is not None + ) + if not pass_through_caps: + return gemini_cap + return max(gemini_cap or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, *pass_through_caps) + + +def _positive_int(value: object) -> int | None: + number: Final = _to_int(value) + return number if number is not None and number > 0 else None + + def _requested_output_counts( request_body: Mapping[str, object], route: str, diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 3dc7ef7c226..f3227130974 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -338,6 +338,31 @@ def test_a_gemini_output_cap_of_zero_or_below_holds_the_default_output(cap: int) ) +@pytest.mark.parametrize( + ("output_fields", "same_hold_fields"), + ( + ({"generationConfig": {"max_output_tokens": 30000}}, {"generationConfig": {"maxOutputTokens": 30000}}), + ({"generation_config": {"maxOutputTokens": 30000}}, {"generationConfig": {"maxOutputTokens": 30000}}), + ({"generation_config": {"max_output_tokens": 30000}}, {"generationConfig": {"maxOutputTokens": 30000}}), + ( + {"generationConfig": {"maxOutputTokens": 8000}, "generation_config": {"max_output_tokens": 30000}}, + {"generationConfig": {"maxOutputTokens": 30000}}, + ), + ({"generation_config": {"max_output_tokens": 5}}, {}), + ( + {"generationConfig": {"maxOutputTokens": 8000, "max_output_tokens": 5}}, + {"generationConfig": {"maxOutputTokens": 8000}}, + ), + ), +) +def test_an_output_cap_spelling_v1beta_drops_can_only_raise_the_hold( + output_fields: dict[str, object], same_hold_fields: dict[str, object] +) -> None: + expected_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], **same_hold_fields}) + assert expected_hold is not None + assert _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], **output_fields}) == expected_hold + + UNPRICEABLE_COUNTS: Final = (10**400, "9" * 400, 1e308, math.inf) From a2f15098f784cd120856d7d1bea2d96a0736655e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:09:18 -0700 Subject: [PATCH 20/23] fix(proxy): hold Gemini candidates of a model with no context window up to the default output A Gemini model the cost map has no context window for, priced in its deployment, held one candidate for any candidateCount, so parallel calls could pass their budget holds. Its candidates are held again up to its counted prompt plus the default output hold --- .../spend_tracking/budget_reservation.py | 12 ++++----- .../spend_tracking/test_budget_reservation.py | 27 +++++++++++++++++-- 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index f45f9028994..a2e0337ce02 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1465,20 +1465,18 @@ def _max_cost_for_cost_info( held_output_tokens=output_tokens, output_tokens=output_tokens, ) - main_hold: Final = _main_generate_content_hold(model_info) - if one_candidate_hold is None or main_hold is None: + hold_ceiling: Final = _candidates_hold_ceiling(model_info=model_info, counted_input_tokens=estimated_input_tokens) + if one_candidate_hold is None or hold_ceiling is None: return one_candidate_hold - return max(one_candidate_hold, min(candidates_hold, main_hold)) + return max(one_candidate_hold, min(candidates_hold, hold_ceiling)) -def _main_generate_content_hold(model_info: Mapping[str, object]) -> float | None: +def _candidates_hold_ceiling(model_info: Mapping[str, object], counted_input_tokens: int) -> float | None: context_window: Final = _to_int(model_info.get("max_input_tokens")) - if context_window is None: - return None default_output_tokens: Final = min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, _model_output_ceiling(model_info)) return _priced_hold( model_info=model_info, - input_tokens=context_window, + input_tokens=counted_input_tokens if context_window is None else context_window, held_output_tokens=default_output_tokens, output_tokens=default_output_tokens, ) diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index f3227130974..5baddde6b39 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -381,14 +381,17 @@ LOW_OUTPUT_CEILING: Final = {**WIDE_WINDOW, "max_output_tokens": 1000} def _priced_gemini_hold( - model_info: dict[str, object], generation_config: dict[str, object], input_tokens: int = 10 + model_info: dict[str, object], + generation_config: dict[str, object], + input_tokens: int = 10, + model: str = "vertex_ai/gemini-3.8-flash", ) -> float | None: router: Final = Router( model_list=[ { "model_name": "assistant", "litellm_params": { - "model": "vertex_ai/gemini-3.8-flash", + "model": model, "input_cost_per_token": 1e-6, "output_cost_per_token": 4e-6, }, @@ -447,6 +450,26 @@ def test_candidates_never_hold_less_than_one_candidate() -> None: assert _priced_gemini_hold(narrow_window, {"maxOutputTokens": 50_000, "candidateCount": 2}) == one_candidate_hold +NO_WINDOW_MODEL: Final = "vertex_ai/gemini-without-cost-map-entry" + + +@pytest.mark.parametrize( + ("generation_config", "same_hold_config"), + ( + ({"maxOutputTokens": 4000, "candidateCount": 2}, {"maxOutputTokens": 8000}), + ({"maxOutputTokens": 8000, "candidateCount": 4}, {}), + ({"candidateCount": 1_000_000}, {}), + ), +) +def test_candidates_of_a_model_with_no_context_window_hold_up_to_one_uncapped_candidate( + generation_config: dict[str, object], same_hold_config: dict[str, object] +) -> None: + no_window: Final = {"max_output_tokens": 100_000} + expected_hold: Final = _priced_gemini_hold(no_window, same_hold_config, model=NO_WINDOW_MODEL) + assert expected_hold is not None + assert _priced_gemini_hold(no_window, generation_config, model=NO_WINDOW_MODEL) == expected_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) From 82ac5a1ef9b4dd52868eacb690dd43c23a8eb8af Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:23:23 -0700 Subject: [PATCH 21/23] fix(proxy): keep main's Gemini hold unless the request caps its output An uncapped generateContent call, or one capped only by a snake_case spelling, now holds what main holds (the context window plus the default output). A call with generationConfig.maxOutputTokens holds its counted prompt plus cap times candidateCount, never clamped to the deployment's max_output_tokens and never more than main's hold. A model with no context window keeps main's no-hold, and candidateCount always counts on generateContent routes. Also keeps get_litellm_virtual_key in pass_through_endpoints/common_utils.py, so an out-of-tree importer does not break --- .../pass_through_endpoints/common_utils.py | 18 +++ .../spend_tracking/budget_reservation.py | 113 +++++++----------- ...test_passthrough_endpoints_common_utils.py | 40 +++++++ .../spend_tracking/test_budget_reservation.py | 108 +++++++++-------- 4 files changed, 156 insertions(+), 123 deletions(-) create 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 new file mode 100644 index 00000000000..3ae83cbe906 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/common_utils.py @@ -0,0 +1,18 @@ +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 a2e0337ce02..173ba3f637c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -55,7 +55,6 @@ from litellm.router import Router from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.types.router import DeploymentTypedDict from litellm.types.utils import CallTypes -from litellm.utils import ProviderConfigManager @dataclass @@ -1404,9 +1403,7 @@ 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) - ) + output_multiplier: Final = _get_output_multiplier(_requested_output_counts(request_body=request_body, route=route)) estimates: Final = [ _max_cost_for_cost_info( request_body=request_body, @@ -1444,39 +1441,54 @@ def _max_cost_for_cost_info( model_info=model_info, input_tokens=input_tokens, ) + if estimated_input_tokens is None: + return None + if RouteChecks.is_generate_content_route(route): + return _generate_content_hold( + request_body=request_body, + model_info=model_info, + input_tokens=estimated_input_tokens, + output_multiplier=output_multiplier, + ) output_tokens: Final = _estimate_output_tokens( request_body=request_body, route=route, model_info=model_info, ) - if estimated_input_tokens is None or output_tokens is None: + if output_tokens is None: return None - candidates_hold: Final = _priced_hold( + return _priced_hold( model_info=model_info, input_tokens=estimated_input_tokens, held_output_tokens=output_tokens * output_multiplier, output_tokens=output_tokens, ) - if candidates_hold is None or output_multiplier == 1 or not RouteChecks.is_generate_content_route(route): - return candidates_hold - one_candidate_hold: Final = _priced_hold( + + +def _generate_content_hold( + request_body: Mapping[str, object], model_info: Mapping[str, object], input_tokens: int, output_multiplier: int +) -> float | None: + main_hold: Final = _main_generate_content_hold(model_info) + requested_output_tokens: Final = _gemini_requested_output_tokens(request_body) + if main_hold is None or requested_output_tokens is None: + return main_hold + capped_hold: Final = _priced_hold( model_info=model_info, - input_tokens=estimated_input_tokens, - held_output_tokens=output_tokens, - output_tokens=output_tokens, + input_tokens=input_tokens, + held_output_tokens=requested_output_tokens * output_multiplier, + output_tokens=requested_output_tokens, ) - hold_ceiling: Final = _candidates_hold_ceiling(model_info=model_info, counted_input_tokens=estimated_input_tokens) - if one_candidate_hold is None or hold_ceiling is None: - return one_candidate_hold - return max(one_candidate_hold, min(candidates_hold, hold_ceiling)) + return main_hold if capped_hold is None else min(capped_hold, main_hold) -def _candidates_hold_ceiling(model_info: Mapping[str, object], counted_input_tokens: int) -> float | None: +def _main_generate_content_hold(model_info: Mapping[str, object]) -> float | None: context_window: Final = _to_int(model_info.get("max_input_tokens")) + if context_window is None: + return None default_output_tokens: Final = min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, _model_output_ceiling(model_info)) return _priced_hold( model_info=model_info, - input_tokens=counted_input_tokens if context_window is None else context_window, + input_tokens=context_window, held_output_tokens=default_output_tokens, output_tokens=default_output_tokens, ) @@ -1709,7 +1721,7 @@ def _estimate_output_tokens( if _is_input_only_route(route=route): return 0 - requested: Final = _requested_output_tokens(request_body, route) + requested: Final = _requested_output_tokens(request_body) # 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 @@ -1739,9 +1751,7 @@ def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object ) -def _requested_output_tokens(request_body: Mapping[str, object], route: str) -> int | None: - if RouteChecks.is_generate_content_route(route): - return _gemini_requested_output_tokens(request_body) +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"), @@ -1749,7 +1759,7 @@ 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) -_GEMINI_PASS_THROUGH_ONLY_CAPS: Final = ( +_OTHER_GEMINI_CAP_SPELLINGS: Final = ( ("generationConfig", "max_output_tokens"), ("generation_config", "maxOutputTokens"), ("generation_config", "max_output_tokens"), @@ -1758,14 +1768,13 @@ _GEMINI_PASS_THROUGH_ONLY_CAPS: Final = ( def _gemini_requested_output_tokens(request_body: Mapping[str, object]) -> int | None: gemini_cap: Final = _positive_int(_nested_field(request_body, "generationConfig", "maxOutputTokens")) - pass_through_caps: Final = tuple( - cap - for config_field, field in _GEMINI_PASS_THROUGH_ONLY_CAPS - if (cap := _positive_int(_nested_field(request_body, config_field, field))) is not None + if gemini_cap is None: + return None + other_caps: Final = ( + _positive_int(_nested_field(request_body, config_field, field)) + for config_field, field in _OTHER_GEMINI_CAP_SPELLINGS ) - if not pass_through_caps: - return gemini_cap - return max(gemini_cap or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, *pass_through_caps) + return max((gemini_cap, *(cap for cap in other_caps if cap is not None))) def _positive_int(value: object) -> int | None: @@ -1773,46 +1782,10 @@ def _positive_int(value: object) -> int | None: return number if number is not None and number > 0 else None -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.registry().deployments(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 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 _requested_output_counts(request_body: Mapping[str, object], route: str) -> tuple[object, ...]: + if RouteChecks.is_generate_content_route(route): + return _gemini_candidate_counts(request_body) + return (request_body.get("n"), request_body.get("best_of")) def _get_output_multiplier(requested_counts: tuple[object, ...]) -> int: 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 54161b4883b..078bd4dd402 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,6 +1,46 @@ +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 5baddde6b39..fd02c3bbc65 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -25,7 +25,6 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_reservation_cache_key, ) from litellm.proxy.spend_tracking.budget_reservation import ( - DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, _get_team_member_budget_counter, estimate_request_input_cost, estimate_request_max_cost, @@ -197,9 +196,11 @@ 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}}, {"n": 3}), - ({"generationConfig": {"candidate_count": 3}}, {"n": 3}), - ({}, {}), + ({"generationConfig": {"maxOutputTokens": 5, "candidate_count": 3}}, {"max_tokens": 5, "n": 3}), + ( + {"generationConfig": {"maxOutputTokens": 5}, "generation_config": {"candidate_count": 3}}, + {"max_tokens": 5, "n": 3}, + ), ), ) @pytest.mark.parametrize("system_texts", (GEMINI_SYSTEM_PARTS[:1], GEMINI_SYSTEM_PARTS)) @@ -239,6 +240,7 @@ def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window 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}} +GEMINI_OUTPUT_CAP: Final = {"generationConfig": {"maxOutputTokens": 5}} def _gemini_hold(body: dict[str, object], input_tokens: int | None = None) -> float | None: @@ -277,7 +279,9 @@ def _gemini_hold(body: dict[str, object], input_tokens: int | None = None) -> fl 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) + assert _gemini_hold({**body, **GEMINI_OUTPUT_CAP}) == _gemini_hold( + {"contents": [GEMINI_TEXT_CONTENT], **GEMINI_OUTPUT_CAP}, input_tokens=context_window + ) def test_gemini_tool_turns_and_thought_signatures_count_as_prompt() -> None: @@ -290,10 +294,12 @@ def test_gemini_tool_turns_and_thought_signatures_count_as_prompt() -> None: {"role": "model", "parts": [{"functionCall": {"name": "lookup", "args": {}}, "thoughtSignature": "c2ln"}]}, {"role": "user", "parts": [{"functionResponse": {"name": "lookup", "response": {"answer": "hi"}}}]}, ], + **GEMINI_OUTPUT_CAP, } - text_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) + text_body: Final = {"contents": [GEMINI_TEXT_CONTENT], **GEMINI_OUTPUT_CAP} + text_hold: Final = _gemini_hold(text_body) tool_turns_hold: Final = _gemini_hold(tool_turns) - context_window_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}, input_tokens=context_window) + context_window_hold: Final = _gemini_hold(text_body, 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 @@ -328,22 +334,34 @@ def test_gemini_output_fields_on_a_chat_body_leave_its_hold_alone(gemini_output_ assert _chat_hold({**GEMINI_CHAT_BODY, **gemini_output_fields}) == plain_hold -@pytest.mark.parametrize("cap", (0, -1000)) -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 +@pytest.mark.parametrize( + "generation_config", + ({}, {"candidateCount": 4}, {"maxOutputTokens": 0}, {"maxOutputTokens": -1000}, {"maxOutputTokens": "many"}), +) +def test_a_gemini_call_without_an_output_cap_holds_the_whole_context_window( + generation_config: dict[str, object], +) -> None: + context_window: Final = litellm.get_model_info("gemini-3.8-flash")["max_input_tokens"] + assert context_window is not None + window_filling_chat_hold: Final = estimate_request_max_cost( + request_body=GEMINI_CHAT_BODY, + route="/chat/completions", + llm_router=None, + input_token_counts={"gemini-3.8-flash": context_window}, + ) + assert window_filling_chat_hold is not None assert ( - _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": cap}}) - == uncapped_hold + _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": generation_config}) + == window_filling_chat_hold ) @pytest.mark.parametrize( ("output_fields", "same_hold_fields"), ( - ({"generationConfig": {"max_output_tokens": 30000}}, {"generationConfig": {"maxOutputTokens": 30000}}), - ({"generation_config": {"maxOutputTokens": 30000}}, {"generationConfig": {"maxOutputTokens": 30000}}), - ({"generation_config": {"max_output_tokens": 30000}}, {"generationConfig": {"maxOutputTokens": 30000}}), + ({"generationConfig": {"max_output_tokens": 30000}}, {}), + ({"generation_config": {"maxOutputTokens": 30000}}, {}), + ({"generation_config": {"max_output_tokens": 30000}}, {}), ( {"generationConfig": {"maxOutputTokens": 8000}, "generation_config": {"max_output_tokens": 30000}}, {"generationConfig": {"maxOutputTokens": 30000}}, @@ -355,7 +373,7 @@ def test_a_gemini_output_cap_of_zero_or_below_holds_the_default_output(cap: int) ), ), ) -def test_an_output_cap_spelling_v1beta_drops_can_only_raise_the_hold( +def test_an_output_cap_spelled_other_than_max_output_tokens_can_only_raise_the_hold( output_fields: dict[str, object], same_hold_fields: dict[str, object] ) -> None: expected_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], **same_hold_fields}) @@ -368,10 +386,12 @@ 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]}) + one_candidate_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], **GEMINI_OUTPUT_CAP}) assert one_candidate_hold is not None assert ( - _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": count}}) + _gemini_hold( + {"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": 5, "candidateCount": count}} + ) == one_candidate_hold ) @@ -412,8 +432,8 @@ def _priced_gemini_hold( ( (WIDE_WINDOW, {"maxOutputTokens": 8000, "candidateCount": 4}, WIDE_WINDOW, {"maxOutputTokens": 32000}), (WIDE_WINDOW, {"maxOutputTokens": 20000, "candidate_count": 2}, WIDE_WINDOW, {"maxOutputTokens": 40000}), - (WIDE_WINDOW, {"candidateCount": 2}, WIDE_WINDOW, {"maxOutputTokens": 2 * DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK}), - (LOW_OUTPUT_CEILING, {"candidateCount": 4}, {**WIDE_WINDOW, "max_output_tokens": 4000}, {}), + (LOW_OUTPUT_CEILING, {"maxOutputTokens": 8000, "candidateCount": 4}, WIDE_WINDOW, {"maxOutputTokens": 32000}), + (LOW_OUTPUT_CEILING, {"maxOutputTokens": 8000}, WIDE_WINDOW, {"maxOutputTokens": 8000}), ), ) def test_candidates_hold_what_one_candidate_writing_all_their_output_holds( @@ -433,9 +453,10 @@ def test_candidates_hold_what_one_candidate_writing_all_their_output_holds( (WIDE_WINDOW, {"candidateCount": 1e300}), (WIDE_WINDOW, {"maxOutputTokens": 100_000, "candidateCount": 8}), (LOW_OUTPUT_CEILING, {"candidateCount": 1e300}), + (LOW_OUTPUT_CEILING, {"maxOutputTokens": 100_000, "candidateCount": 8}), ), ) -def test_candidates_never_hold_more_than_the_whole_context_window_hold( +def test_a_gemini_hold_never_exceeds_the_whole_context_window_hold( model_info: dict[str, object], generation_config: dict[str, object] ) -> None: context_window_hold: Final = _priced_gemini_hold(model_info, {}, input_tokens=1_000_000) @@ -443,31 +464,14 @@ def test_candidates_never_hold_more_than_the_whole_context_window_hold( assert _priced_gemini_hold(model_info, generation_config) == context_window_hold -def test_candidates_never_hold_less_than_one_candidate() -> None: - narrow_window: Final = {"max_input_tokens": 10, "max_output_tokens": 100_000} - one_candidate_hold: Final = _priced_gemini_hold(narrow_window, {"maxOutputTokens": 50_000}) - assert one_candidate_hold is not None - assert _priced_gemini_hold(narrow_window, {"maxOutputTokens": 50_000, "candidateCount": 2}) == one_candidate_hold - - NO_WINDOW_MODEL: Final = "vertex_ai/gemini-without-cost-map-entry" @pytest.mark.parametrize( - ("generation_config", "same_hold_config"), - ( - ({"maxOutputTokens": 4000, "candidateCount": 2}, {"maxOutputTokens": 8000}), - ({"maxOutputTokens": 8000, "candidateCount": 4}, {}), - ({"candidateCount": 1_000_000}, {}), - ), + "generation_config", ({}, {"maxOutputTokens": 4000}, {"maxOutputTokens": 8000, "candidateCount": 4}) ) -def test_candidates_of_a_model_with_no_context_window_hold_up_to_one_uncapped_candidate( - generation_config: dict[str, object], same_hold_config: dict[str, object] -) -> None: - no_window: Final = {"max_output_tokens": 100_000} - expected_hold: Final = _priced_gemini_hold(no_window, same_hold_config, model=NO_WINDOW_MODEL) - assert expected_hold is not None - assert _priced_gemini_hold(no_window, generation_config, model=NO_WINDOW_MODEL) == expected_hold +def test_a_gemini_model_with_no_context_window_is_not_held(generation_config: dict[str, object]) -> None: + assert _priced_gemini_hold({"max_output_tokens": 100_000}, generation_config, model=NO_WINDOW_MODEL) is None @pytest.mark.parametrize("count", UNPRICEABLE_COUNTS) @@ -487,18 +491,17 @@ def _group_generate_content_hold(router: Router, generation_config: dict[str, ob @pytest.mark.parametrize( - ("backend_models", "candidates_are_generated"), + "backend_models", ( - (("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), + ("vertex_ai/claude-haiku-5-5",), + ("anthropic/claude-haiku-5-5",), + ("openai/gpt-5.4-mini",), + ("gemini/gemini-3.8-flash",), + ("vertex_ai/gemini-3.8-flash",), + ("vertex_ai/claude-haiku-5-5", "gemini/gemini-3.8-flash"), ), ) -def test_a_candidate_count_is_held_only_when_a_deployment_generates_the_candidates( - backend_models: tuple[str, ...], candidates_are_generated: bool -) -> None: +def test_a_candidate_count_is_held_whatever_deployment_serves_the_name(backend_models: tuple[str, ...]) -> None: router: Final = Router( model_list=[ {"model_name": "assistant", "litellm_params": {"model": model, "api_key": "sk-test"}} @@ -510,8 +513,7 @@ def test_a_candidate_count_is_held_only_when_a_deployment_generates_the_candidat 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 + assert three_candidate_hold > one_candidate_hold def test_a_candidate_count_for_a_gemini_name_no_deployment_serves_is_held() -> None: From f418a83f4962b811d2378dd89c3718c0c3dce07e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 00:23:47 -0700 Subject: [PATCH 22/23] fix(proxy): keep main's budget hold on generateContent routes The PR now leaves the generateContent budget estimate as main has it. A live risk check showed the counted Gemini hold under-reserves (the prompt counter undercounts Gemini bodies, tiered and image output rates are not read, and request fields the handler merges later are not counted), so that estimate moves to a follow-up. The single auth pass on the Vertex pass-through and the systemInstruction adapter fix stay --- .../spend_tracking/budget_reservation.py | 182 ++----- litellm/proxy/spend_tracking/input_tokens.py | 121 +---- .../spend_tracking/test_budget_reservation.py | 486 ------------------ .../proxy/spend_tracking/test_input_tokens.py | 29 -- tests/unit/proxy/test_budget_reservation.py | 2 +- 5 files changed, 53 insertions(+), 767 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 173ba3f637c..fc6fb266d04 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -7,7 +7,7 @@ import time from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from itertools import chain, product +from itertools import chain from types import MappingProxyType from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast @@ -18,7 +18,6 @@ import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -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 ( @@ -184,7 +183,7 @@ async def _apply_over_budget_reservation_policy( entry: dict[str, float | str], applied_entries: list[dict[str, float | str]], reservation_cost: float, - spend_before_reservation: float, + current_spend: float, fail_closed_budget_enforcement: bool = False, ) -> float: """ @@ -200,13 +199,13 @@ async def _apply_over_budget_reservation_policy( applied_entries.remove(entry) return reservation_cost - remaining_before_reservation: Final = counter.max_budget - spend_before_reservation + remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost) if remaining_before_reservation <= 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) + if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12: _raise_counter_budget_exceeded( counter=counter, - current_cost=spend_before_reservation, + current_cost=current_spend - reservation_cost, estimated_cost=reservation_cost, ) await _resize_applied_reservation( @@ -997,7 +996,7 @@ async def _reserve_reservable_counters( current_spend_by_counter_key=current_spend_by_counter_key, reservation_cost=reservation_cost, ): - charged_cost = min(reservation_cost, max(counter.max_budget for counter in group)) + charged_cost = reservation_cost 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) @@ -1010,19 +1009,18 @@ async def _reserve_reservable_counters( for counter, entry, reserved_value in zip(group, entries, reserved_values): if entry not in applied_entries: continue - 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: + 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: reservation_cost = await _apply_over_budget_reservation_policy( counter=counter, valid_token=valid_token, entry=entry, applied_entries=applied_entries, reservation_cost=reservation_cost, - spend_before_reservation=spend_before_reservation, + current_spend=current_spend, fail_closed_budget_enforcement=fail_closed_budget_enforcement, ) return reservation_cost @@ -1403,14 +1401,12 @@ 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)) 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) @@ -1424,7 +1420,6 @@ 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,87 +1436,46 @@ def _max_cost_for_cost_info( model_info=model_info, input_tokens=input_tokens, ) - if estimated_input_tokens is None: - return None - if RouteChecks.is_generate_content_route(route): - return _generate_content_hold( - request_body=request_body, - model_info=model_info, - input_tokens=estimated_input_tokens, - output_multiplier=output_multiplier, - ) output_tokens: Final = _estimate_output_tokens( request_body=request_body, route=route, model_info=model_info, ) - if output_tokens is None: + if estimated_input_tokens is None or output_tokens is None: return None - return _priced_hold( - model_info=model_info, - input_tokens=estimated_input_tokens, - held_output_tokens=output_tokens * output_multiplier, - output_tokens=output_tokens, - ) - -def _generate_content_hold( - request_body: Mapping[str, object], model_info: Mapping[str, object], input_tokens: int, output_multiplier: int -) -> float | None: - main_hold: Final = _main_generate_content_hold(model_info) - requested_output_tokens: Final = _gemini_requested_output_tokens(request_body) - if main_hold is None or requested_output_tokens is None: - return main_hold - capped_hold: Final = _priced_hold( - model_info=model_info, - input_tokens=input_tokens, - held_output_tokens=requested_output_tokens * output_multiplier, - output_tokens=requested_output_tokens, - ) - return main_hold if capped_hold is None else min(capped_hold, main_hold) - - -def _main_generate_content_hold(model_info: Mapping[str, object]) -> float | None: - context_window: Final = _to_int(model_info.get("max_input_tokens")) - if context_window is None: - return None - default_output_tokens: Final = min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, _model_output_ceiling(model_info)) - return _priced_hold( - model_info=model_info, - input_tokens=context_window, - held_output_tokens=default_output_tokens, - output_tokens=default_output_tokens, - ) - - -def _priced_hold( - model_info: Mapping[str, object], input_tokens: int, held_output_tokens: int, output_tokens: int -) -> float | None: + output_multiplier: Final = _get_output_multiplier(request_body=request_body) 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=input_tokens) + tier: Final = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=estimated_input_tokens) if tier is not None: - tier_output_rate: Final = max( + output_rate = max( tier_rate(tier, "output_cost_per_token"), tier_rate(tier, "output_cost_per_reasoning_token"), ) - return (input_tokens * tier_rate(tier, "input_cost_per_token")) + _output_cost( - held_output_tokens, output_tokens, tier_output_rate + return (estimated_input_tokens * tier_rate(tier, "input_cost_per_token")) + ( + output_tokens * output_multiplier * output_rate ) input_cost_per_token: Final = _to_float(model_info.get("input_cost_per_token")) output_cost_per_token: Final = _to_float(model_info.get("output_cost_per_token")) output_cost_per_reasoning_token: Final = _to_float(model_info.get("output_cost_per_reasoning_token")) - if input_cost_per_token is None and input_tokens > 0: - return None - if output_cost_per_token is None and output_cost_per_reasoning_token is None and output_tokens > 0: + cost = 0.0 + if input_cost_per_token is not None: + cost += estimated_input_tokens * input_cost_per_token + elif estimated_input_tokens > 0: return None # The reasoning-token share is unknown before the request runs, so reserve every # output token at the higher of the standard and reasoning rates to avoid # under-reserving reasoning-heavy requests. - output_rate: Final = max(output_cost_per_token or 0.0, output_cost_per_reasoning_token or 0.0) - return input_tokens * (input_cost_per_token or 0.0) + _output_cost(held_output_tokens, output_tokens, output_rate) + 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 + elif output_tokens > 0: + return None + + return cost def _estimate_image_generation_cost( @@ -1663,7 +1617,6 @@ async def count_request_input_tokens( request_body=request_body, raw_body=raw_body, models=tuple(dict.fromkeys(chain.from_iterable(counted_models.values()))), - gemini_route=RouteChecks.is_generate_content_route(route), ) chat_pdf_files: Final = ( inline_pdf_file_count(request_body.get("messages")) @@ -1696,9 +1649,7 @@ 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, gemini_route=RouteChecks.is_generate_content_route(route) - ) + else count_input_tokens_for_model(request_body=request_body, model=model) ) if counted is not None: return counted @@ -1731,76 +1682,29 @@ def _estimate_output_tokens( # the reservation up to remaining team headroom and pin the counter # at the cap — the model can only physically emit max_output_tokens # anyway, so reserving more is both wasteful and a DoS surface. - model_ceiling: Final = _model_output_ceiling(model_info) + model_ceiling: Final = _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK return min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK if requested is None else requested, model_ceiling) _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: - config: Final = as_str_mapping(request_body.get(config_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: + inference_config: Final = request_body.get("inferenceConfig") candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), - _nested_field(request_body, "inferenceConfig", "maxTokens"), + inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None, ) - 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), None) -_OTHER_GEMINI_CAP_SPELLINGS: Final = ( - ("generationConfig", "max_output_tokens"), - ("generation_config", "maxOutputTokens"), - ("generation_config", "max_output_tokens"), -) - - -def _gemini_requested_output_tokens(request_body: Mapping[str, object]) -> int | None: - gemini_cap: Final = _positive_int(_nested_field(request_body, "generationConfig", "maxOutputTokens")) - if gemini_cap is None: - return None - other_caps: Final = ( - _positive_int(_nested_field(request_body, config_field, field)) - for config_field, field in _OTHER_GEMINI_CAP_SPELLINGS - ) - return max((gemini_cap, *(cap for cap in other_caps if cap is not None))) - - -def _positive_int(value: object) -> int | None: - number: Final = _to_int(value) - return number if number is not None and number > 0 else None - - -def _requested_output_counts(request_body: Mapping[str, object], route: str) -> tuple[object, ...]: - if RouteChecks.is_generate_content_route(route): - return _gemini_candidate_counts(request_body) - return (request_body.get("n"), request_body.get("best_of")) - - -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(held_output_tokens: int, output_tokens: int, output_rate: float) -> float: - try: - return held_output_tokens * output_rate - except OverflowError: - return output_tokens * output_rate - - -def _model_output_ceiling(model_info: Mapping[str, object]) -> int: - return _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK +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 _is_input_only_route(route: str) -> bool: @@ -1828,7 +1732,7 @@ def _to_int(value: object) -> int | None: return None try: return int(value) - except (TypeError, ValueError, OverflowError): + except (TypeError, ValueError): return None diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 3a5ec23bf51..9f37d2965bc 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -16,15 +16,11 @@ import json from collections.abc import Iterator, Mapping, Sequence from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Final, TypeAlias - -from pydantic import TypeAdapter +from typing import TYPE_CHECKING, Final import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import OPENAI_CHAT_PDF_FILE_TOKENS, PDF_DATA_URL_PREFIX -from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter -from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items from litellm.litellm_core_utils.deployment_pricing import deployment_pricing_models, str_keyed from litellm.litellm_core_utils.token_counter import ( deployment_openai_model, @@ -53,30 +49,7 @@ TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 _OPENAI_PREFIX: Final = "openai/" -_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 = ( - "messages", - "prompt", - "input", - "query", - "documents", - "contents", - *_GEMINI_SYSTEM_INSTRUCTION_FIELDS, - "tools", - "tool_choice", -) +_INPUT_SIZE_FIELDS: Final = ("messages", "prompt", "input", "query", "documents", "tools", "tool_choice") def _approximate_input_size(request_body: Mapping[str, object]) -> int: @@ -89,10 +62,9 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int: async def count_input_tokens( - request_body: dict[str, object], + request_body: dict, 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.""" @@ -100,14 +72,13 @@ 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 gemini_route or holds_compaction_summary(request_body.get("input")) else raw_body + counted_raw_body: Final = None if 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 ] @@ -117,20 +88,18 @@ async def count_input_tokens( async def _count_group( - request_body: dict[str, object], + request_body: dict, 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, gemini_route=gemini_route) + return _count_input_tokens_for_models(request_body=request_body, models=models) 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: @@ -163,20 +132,14 @@ async def _native_counts( def _count_input_tokens_for_models( - request_body: dict[str, object], + request_body: dict, 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, gemini_route=gemini_route - ) - ) - is not None + if (tokens := count_input_tokens_for_model(request_body=request_body, model=model)) is not None } ) @@ -258,14 +221,8 @@ def _served_by_openai(deployment: object) -> bool: return served_provider == "openai" -def count_input_tokens_for_model(request_body: dict, model: str, gemini_route: bool = False) -> int | None: +def count_input_tokens_for_model(request_body: dict, model: str) -> 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( @@ -294,66 +251,6 @@ def count_input_tokens_for_model(request_body: dict, model: str, gemini_route: b 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 - ) - 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=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: - 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: - 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/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index fd02c3bbc65..a87ac7fa7c0 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -1,6 +1,5 @@ from __future__ import annotations -import json import math from datetime import datetime, timedelta, timezone from typing import Final @@ -26,7 +25,6 @@ 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, @@ -179,490 +177,6 @@ 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_PARTS: Final = ("Answer tersely", "Use British spelling") - - -@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}), - ({"generationConfig": {"maxOutputTokens": 5, "candidateCount": 3}}, {"max_tokens": 5, "n": 3}), - ({"generationConfig": {"maxOutputTokens": 5, "candidate_count": 3}}, {"max_tokens": 5, "n": 3}), - ( - {"generationConfig": {"maxOutputTokens": 5}, "generation_config": {"candidate_count": 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], - system_texts: tuple[str, ...], -) -> None: - native_cost: Final = estimate_request_max_cost( - request_body={ - "systemInstruction": {"parts": [{"text": text} for text in system_texts]}, - "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": "\n".join(system_texts)}, - {"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 - - -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}} -GEMINI_OUTPUT_CAP: Final = {"generationConfig": {"maxOutputTokens": 5}} - - -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": {"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]}, - {"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_OUTPUT_CAP}) == _gemini_hold( - {"contents": [GEMINI_TEXT_CONTENT], **GEMINI_OUTPUT_CAP}, 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"}}}]}, - ], - **GEMINI_OUTPUT_CAP, - } - text_body: Final = {"contents": [GEMINI_TEXT_CONTENT], **GEMINI_OUTPUT_CAP} - text_hold: Final = _gemini_hold(text_body) - tool_turns_hold: Final = _gemini_hold(tool_turns) - context_window_hold: Final = _gemini_hold(text_body, 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 - - -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( - "generation_config", - ({}, {"candidateCount": 4}, {"maxOutputTokens": 0}, {"maxOutputTokens": -1000}, {"maxOutputTokens": "many"}), -) -def test_a_gemini_call_without_an_output_cap_holds_the_whole_context_window( - generation_config: dict[str, object], -) -> None: - context_window: Final = litellm.get_model_info("gemini-3.8-flash")["max_input_tokens"] - assert context_window is not None - window_filling_chat_hold: Final = estimate_request_max_cost( - request_body=GEMINI_CHAT_BODY, - route="/chat/completions", - llm_router=None, - input_token_counts={"gemini-3.8-flash": context_window}, - ) - assert window_filling_chat_hold is not None - assert ( - _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": generation_config}) - == window_filling_chat_hold - ) - - -@pytest.mark.parametrize( - ("output_fields", "same_hold_fields"), - ( - ({"generationConfig": {"max_output_tokens": 30000}}, {}), - ({"generation_config": {"maxOutputTokens": 30000}}, {}), - ({"generation_config": {"max_output_tokens": 30000}}, {}), - ( - {"generationConfig": {"maxOutputTokens": 8000}, "generation_config": {"max_output_tokens": 30000}}, - {"generationConfig": {"maxOutputTokens": 30000}}, - ), - ({"generation_config": {"max_output_tokens": 5}}, {}), - ( - {"generationConfig": {"maxOutputTokens": 8000, "max_output_tokens": 5}}, - {"generationConfig": {"maxOutputTokens": 8000}}, - ), - ), -) -def test_an_output_cap_spelled_other_than_max_output_tokens_can_only_raise_the_hold( - output_fields: dict[str, object], same_hold_fields: dict[str, object] -) -> None: - expected_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], **same_hold_fields}) - assert expected_hold is not None - assert _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], **output_fields}) == expected_hold - - -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], **GEMINI_OUTPUT_CAP}) - assert one_candidate_hold is not None - assert ( - _gemini_hold( - {"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": 5, "candidateCount": count}} - ) - == one_candidate_hold - ) - - -WIDE_WINDOW: Final = {"max_input_tokens": 1_000_000, "max_output_tokens": 100_000} -LOW_OUTPUT_CEILING: Final = {**WIDE_WINDOW, "max_output_tokens": 1000} - - -def _priced_gemini_hold( - model_info: dict[str, object], - generation_config: dict[str, object], - input_tokens: int = 10, - model: str = "vertex_ai/gemini-3.8-flash", -) -> float | None: - router: Final = Router( - model_list=[ - { - "model_name": "assistant", - "litellm_params": { - "model": model, - "input_cost_per_token": 1e-6, - "output_cost_per_token": 4e-6, - }, - "model_info": model_info, - } - ] - ) - 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": input_tokens}, - ) - - -@pytest.mark.parametrize( - ("model_info", "generation_config", "same_hold_model_info", "same_hold_config"), - ( - (WIDE_WINDOW, {"maxOutputTokens": 8000, "candidateCount": 4}, WIDE_WINDOW, {"maxOutputTokens": 32000}), - (WIDE_WINDOW, {"maxOutputTokens": 20000, "candidate_count": 2}, WIDE_WINDOW, {"maxOutputTokens": 40000}), - (LOW_OUTPUT_CEILING, {"maxOutputTokens": 8000, "candidateCount": 4}, WIDE_WINDOW, {"maxOutputTokens": 32000}), - (LOW_OUTPUT_CEILING, {"maxOutputTokens": 8000}, WIDE_WINDOW, {"maxOutputTokens": 8000}), - ), -) -def test_candidates_hold_what_one_candidate_writing_all_their_output_holds( - model_info: dict[str, object], - generation_config: dict[str, object], - same_hold_model_info: dict[str, object], - same_hold_config: dict[str, object], -) -> None: - expected_hold: Final = _priced_gemini_hold(same_hold_model_info, same_hold_config) - assert expected_hold is not None - assert _priced_gemini_hold(model_info, generation_config) == expected_hold - - -@pytest.mark.parametrize( - ("model_info", "generation_config"), - ( - (WIDE_WINDOW, {"candidateCount": 1e300}), - (WIDE_WINDOW, {"maxOutputTokens": 100_000, "candidateCount": 8}), - (LOW_OUTPUT_CEILING, {"candidateCount": 1e300}), - (LOW_OUTPUT_CEILING, {"maxOutputTokens": 100_000, "candidateCount": 8}), - ), -) -def test_a_gemini_hold_never_exceeds_the_whole_context_window_hold( - model_info: dict[str, object], generation_config: dict[str, object] -) -> None: - context_window_hold: Final = _priced_gemini_hold(model_info, {}, input_tokens=1_000_000) - assert context_window_hold is not None - assert _priced_gemini_hold(model_info, generation_config) == context_window_hold - - -NO_WINDOW_MODEL: Final = "vertex_ai/gemini-without-cost-map-entry" - - -@pytest.mark.parametrize( - "generation_config", ({}, {"maxOutputTokens": 4000}, {"maxOutputTokens": 8000, "candidateCount": 4}) -) -def test_a_gemini_model_with_no_context_window_is_not_held(generation_config: dict[str, object]) -> None: - assert _priced_gemini_hold({"max_output_tokens": 100_000}, generation_config, model=NO_WINDOW_MODEL) is None - - -@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 _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", - ( - ("vertex_ai/claude-haiku-5-5",), - ("anthropic/claude-haiku-5-5",), - ("openai/gpt-5.4-mini",), - ("gemini/gemini-3.8-flash",), - ("vertex_ai/gemini-3.8-flash",), - ("vertex_ai/claude-haiku-5-5", "gemini/gemini-3.8-flash"), - ), -) -def test_a_candidate_count_is_held_whatever_deployment_serves_the_name(backend_models: tuple[str, ...]) -> 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 - - -def test_a_candidate_count_for_a_gemini_name_no_deployment_serves_is_held() -> None: - router: Final = Router( - model_list=[ - {"model_name": "assistant", "litellm_params": {"model": "anthropic/claude-haiku-5-5", "api_key": "sk-test"}} - ] - ) - one_candidate_hold: Final = estimate_request_max_cost( - request_body={"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": 100}}, - route=GEMINI_ROUTE, - llm_router=router, - input_token_counts={"gemini-3.8-flash": 10}, - ) - three_candidate_hold: Final = estimate_request_max_cost( - request_body={ - "contents": [GEMINI_TEXT_CONTENT], - "generationConfig": {"maxOutputTokens": 100, "candidateCount": 3}, - }, - route=GEMINI_ROUTE, - llm_router=router, - input_token_counts={"gemini-3.8-flash": 10}, - ) - - assert one_candidate_hold is not None - assert three_candidate_hold is not None - assert three_candidate_hold > one_candidate_hold - - -N_SENDING_GROUPS: Final = ( - Router( - model_list=[ - {"model_name": "assistant", "litellm_params": {"model": "anthropic/claude-haiku-5-5", "api_key": "sk"}}, - {"model_name": "backup", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk"}}, - ], - fallbacks=[{"assistant": ["backup"]}], - ), - Router( - model_list=[ - { - "model_name": "assistant", - "litellm_params": { - "model": "mistral/mistral-small-latest", - "api_key": "sk", - "allowed_openai_params": ["n"], - }, - } - ] - ), -) - - -@pytest.mark.parametrize("router", N_SENDING_GROUPS) -def test_an_n_is_held_when_a_fallback_or_an_allowlist_can_still_send_it(router: Router) -> None: - holds: Final = tuple( - 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}, - ) - for output_fields in ({"max_tokens": 100}, {"max_tokens": 100, "n": 3}) - ) - - one_choice_hold, three_choice_hold = holds - assert one_choice_hold is not None - assert three_choice_hold is not None - assert three_choice_hold > one_choice_hold - - -HOLDS_FAR_OVER_A_TINY_BUDGET: Final = ( - ("/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 - 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, - team_id=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: 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 87192a47447..40aee7bb00d 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -280,32 +280,3 @@ 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 - - -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 diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 25e0ee6a0ea..535c3da8c3a 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, - spend_before_reservation=1.5, + current_spend=2.0, ) assert exc_info.value.entity_type == "key" assert exc_info.value.entity_id == "test-token" From 02723a28a662851112dcfeaa8a5185bf7e96b1d0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:25:29 -0800 Subject: [PATCH 23/23] fix(vertex_ai): keep main's Gemini adapter and pass a missing Authorization as no key The Gemini adapter change for later systemInstruction parts goes back to main, since it is not the double-auth fix this PR is for. The Vertex auth dependency now hands a missing Authorization header to the auth check as no key, like main's first auth pass, so custom auth and the 401 text see what they see on main --- .../google_genai/adapters/transformation.py | 10 ++-- .../llm_passthrough_endpoints.py | 4 +- .../google_genai/test_google_genai_adapter.py | 56 ------------------- .../test_llm_pass_through_endpoints.py | 19 ++++++- 4 files changed, 24 insertions(+), 65 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 40ab377dd5e..c1895376675 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -249,10 +249,6 @@ def _genai_tool_call(part: object) -> ChatCompletionAssistantToolCall | None: ) -def _system_texts(parts: object) -> tuple[str, ...]: - 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, ...]: fields: Final = _validated(_CONFIG_FIELDS, content) if fields is None: @@ -636,9 +632,11 @@ class GoogleGenAIAdapter: system_instruction: _GenAISystemInstruction | None = None, ) -> list[AllMessageValues]: """Transform Google GenAI contents to OpenAI messages format""" - system_texts: Final = _system_texts(system_instruction.get("parts") if system_instruction else None) + system_parts: Final = system_instruction.get("parts", ()) if system_instruction else () system_messages: Final = ( - (ChatCompletionSystemMessage(role="system", content="\n".join(system_texts)),) if system_texts else () + (ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"]),) + if system_parts and "text" in system_parts[0] + else () ) return [*system_messages, *chain.from_iterable(_genai_content_messages(content) for content in contents)] diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 9ffca3957b2..30b5c9965d8 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3203,7 +3203,9 @@ async def _authenticate_vertex_caller( 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 + request=request, + api_key=authorization, # pyright: ignore[reportArgumentType] # its own Security default is None when unset + custom_litellm_key_header=litellm_api_key, ) diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index d622a1b85b9..4ab642dc99d 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -194,62 +194,6 @@ 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"), "Answer tersely\nUse British spelling"), - ], -) -def test_every_system_instruction_text_part_reaches_the_system_message( - system_texts: tuple[str, ...], expected_content: 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} - - -@pytest.mark.parametrize( - ("system_parts", "expected_messages"), - [ - 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", - ), - 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_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( - 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 2f7a87d54eb..2aa0ae118aa 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 @@ -8490,9 +8490,9 @@ class TestAzureSpeechRawBodyThroughRealAuth: class _CountingCustomAuth: def __init__(self, only_accepts: str | None = None) -> None: self.only_accepts: Final = only_accepts - self.api_keys_seen: tuple[str, ...] = () + self.api_keys_seen: tuple[str | None, ...] = () - async def __call__(self, request: Request, api_key: str) -> UserAPIKeyAuth: + async def __call__(self, request: Request, api_key: str | None) -> UserAPIKeyAuth: 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") @@ -8694,6 +8694,21 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: assert not google.called + @pytest.mark.parametrize("path", [f"/vertex_ai/{GENERATE_CONTENT}", f"/vertex_ai/discovery/{DISCOVERY_SEARCH}"]) + def test_request_without_a_key_hands_custom_auth_none( + self, monkeypatch: pytest.MonkeyPatch, path: 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, path, {}, custom_auth, ("us-central1", "global")) + + assert response.status_code == 401, response.text + assert custom_auth.api_keys_seen == (None,) + assert recorder.served_user_ids == () + assert not google.called + class TestTypeSafePassthroughRoute: @staticmethod def _request(body: object, query_params: Mapping[str, str] | None = None) -> MagicMock: