diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 627164abb93..e2aefcd4bb7 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -427,9 +427,6 @@ async def gemini_proxy_route( raise Exception( "Required 'GEMINI_API_KEY'/'GOOGLE_API_KEY' in environment to make pass-through calls to Google AI Studio." ) - # Merge query parameters, giving precedence to those in updated_url - merged_params: Final = dict(request.query_params) - merged_params.update({"key": gemini_api_key}) ## check for streaming is_streaming_request = False @@ -440,9 +437,12 @@ async def gemini_proxy_route( endpoint_func: Final = create_pass_through_route( endpoint=endpoint, target=str(updated_url), + custom_headers={"x-goog-api-key": gemini_api_key}, custom_llm_provider="gemini", is_streaming_request=is_streaming_request, - query_params=merged_params, + query_params=dict(request.query_params), + # `key` is the LiteLLM virtual key; the provider key goes in the x-goog-api-key header. + excluded_query_params=("key",), ) # dynamically construct pass-through endpoint based on incoming path received_value: Final = await endpoint_func( request, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 8111e46c0b7..40a657d0cb8 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -6,7 +6,16 @@ import posixpath import re import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Callable, + Collection, + Iterable, + Mapping, + MutableMapping, + Sequence, +) from dataclasses import dataclass from datetime import datetime from functools import reduce @@ -1206,6 +1215,12 @@ def _unfiltered_headers(headers: Mapping[str, str]) -> Mapping[str, str]: return headers +def _forwarded_query_params( + params: Mapping[str, Any], excluded: Collection[str] +) -> dict[str, Any] | None: # mutable-ok: rewrite_query_ids and the httpx request helpers take a plain dict + return {k: v for k, v in params.items() if k not in excluded} or None + + def _json_object(value: object) -> Mapping[str, object]: return cast("Mapping[str, object]", value) if isinstance(value, Mapping) else MappingProxyType({}) # cast-ok: json @@ -1256,6 +1271,7 @@ async def pass_through_request( merge_query_params: bool | None = False, query_params: dict | None = None, default_query_params: dict | None = None, + excluded_query_params: Collection[str] = (), stream: bool | None = None, cost_per_request: float | None = None, custom_llm_provider: str | None = None, @@ -1277,6 +1293,7 @@ async def pass_through_request( merge_query_params: Whether to merge query params query_params: The query params default_query_params: The default query params to be applied if not overridden by client + excluded_query_params: Query param names dropped before the request is sent upstream stream: Whether to stream the response cost_per_request: Optional field - cost per request to the target endpoint custom_llm_provider: Optional field - custom LLM provider for the endpoint @@ -1323,7 +1340,9 @@ async def pass_through_request( parent_span=user_api_key_dict.parent_otel_span, ) - requested_query_params: dict | None = query_params or dict(request.query_params) or None + requested_query_params: dict[str, Any] | None = _forwarded_query_params( + query_params or request.query_params, excluded_query_params + ) endpoint_type: Final[EndpointType] = HttpPassThroughEndpointHelpers.get_endpoint_type( str(url), custom_llm_provider @@ -2325,6 +2344,7 @@ def create_pass_through_route( is_streaming_request: bool | None = False, query_params: dict | None = None, default_query_params: dict | None = None, + excluded_query_params: Collection[str] = (), guardrails: dict[str, object] | None = None, config_file_path: str | None = None, timeout: float | None = None, @@ -2471,6 +2491,7 @@ def create_pass_through_route( merge_query_params=cast(bool | None, param_merge_query_params), query_params=final_query_params, default_query_params=cast(dict | None, param_default_query_params), + excluded_query_params=excluded_query_params, stream=is_stream, custom_body=final_custom_body, cost_per_request=cast(float | None, param_cost_per_request), diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py index 60b263948b6..088b09eb735 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py @@ -212,9 +212,7 @@ async def test_gemini_pass_through_endpoint_forwards_key_and_body( monkeypatch.setenv("GEMINI_API_KEY", "gemini-test-key") endpoint: Final = "v1beta/models/gemini-2.5-flash:countTokens" - upstream: Final = respx_mock.post( - url__startswith=f"https://generativelanguage.googleapis.com/{endpoint}" - ).mock( + upstream: Final = respx_mock.post(url__startswith=f"https://generativelanguage.googleapis.com/{endpoint}").mock( return_value=httpx.Response( 200, json={"totalTokens": 9}, @@ -252,7 +250,9 @@ async def test_gemini_pass_through_endpoint_forwards_key_and_body( assert upstream.call_count == 1 outbound: Final = upstream.calls[0].request assert outbound.url.path == f"/{endpoint}" - assert outbound.url.params["key"] == "gemini-test-key" + assert outbound.headers["x-goog-api-key"] == "gemini-test-key" + assert "key" not in outbound.url.params + assert "sk-9876" not in str(outbound.url) assert json.loads(outbound.content) == json.loads(body) @pytest.mark.asyncio @@ -546,3 +546,60 @@ def test_vertex_pass_through_forwards_a_request_without_betas_despite_a_malforme assert (response.status_code, response.json()) == (200, _VERTEX_CLAUDE_REPLY) assert "anthropic-beta" not in upstream.calls.last.request.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_key_in_header", [False, True]) +async def test_gemini_pass_through_endpoint_keeps_provider_key_out_of_url( + respx_mock: "respx.MockRouter", + httpx_transport: None, + monkeypatch: pytest.MonkeyPatch, + client_key_in_header: bool, +) -> None: + """The provider key is sent in x-goog-api-key, the LiteLLM key is never forwarded, other query params pass through.""" + import httpx + from fastapi import Request, Response + + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + gemini_proxy_route, + ) + + monkeypatch.setenv("GEMINI_API_KEY", "gemini-test-key") + endpoint: Final = "v1beta/models/gemini-2.5-flash:streamGenerateContent" + upstream: Final = respx_mock.get(url__startswith=f"https://generativelanguage.googleapis.com/{endpoint}").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + # No content-type: the request's own query string is the only source of query params. + headers: Final = [(b"x-goog-api-key", b"sk-9876")] if client_key_in_header else [] + query_string: Final = b"alt=sse" if client_key_in_header else b"alt=sse&key=sk-9876" + scope: Final = { + "type": "http", + "http_version": "1.1", + "method": "GET", + "scheme": "http", + "path": f"/gemini/{endpoint}", + "raw_path": f"/gemini/{endpoint}".encode(), + "root_path": "", + "query_string": query_string, + "headers": headers, + "client": ("127.0.0.1", 51234), + "server": ("proxy.local", 4000), + "state": {}, + } + + async def receive() -> dict: + return {"type": "http.request", "body": b"", "more_body": False} + + response: Final = await gemini_proxy_route( + endpoint=endpoint, + request=Request(scope, receive), + fastapi_response=Response(), + ) + + assert response.status_code == 200 + assert upstream.call_count == 1 + outbound: Final = upstream.calls[0].request + assert outbound.headers["x-goog-api-key"] == "gemini-test-key" + assert dict(outbound.url.params) == {"alt": "sse"} + assert "sk-9876" not in str(outbound.url) + assert "sk-9876" not in str(outbound.headers)