diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 941ec4ad419..9199687668b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -2569,6 +2569,7 @@ async def make_call( model: str, messages: list, logging_obj, + timeout: float | httpx.Timeout | None = None, ): if gemini_client is not None: client = gemini_client @@ -2578,7 +2579,14 @@ async def make_call( ) try: - response: Final = await client.post(api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj) + response: Final = await client.post( + api_base, + headers=headers, + data=data, + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) response.raise_for_status() except httpx.HTTPStatusError as e: exception_string: Final = str(await e.response.aread()) @@ -2621,13 +2629,21 @@ def make_sync_call( model: str, messages: list, logging_obj, + timeout: float | httpx.Timeout | None = None, ): if gemini_client is not None: client = gemini_client if client is None: client = HTTPHandler() # Create a new client if none provided - response: Final = client.post(api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj) + response: Final = client.post( + api_base, + headers=headers, + data=data, + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) if response.status_code != 200 and response.status_code != 201: raise VertexAIError( @@ -2745,6 +2761,7 @@ class VertexLLM(VertexBase): model=model, messages=messages, logging_obj=logging_obj, + timeout=timeout, ), model=model, custom_llm_provider="vertex_ai_beta", @@ -3021,6 +3038,7 @@ class VertexLLM(VertexBase): messages=messages, logging_obj=logging_obj, headers=headers, + timeout=timeout, ), model=model, custom_llm_provider="vertex_ai_beta", diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index fd735afb16e..82ec2f1a8bc 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -12,7 +12,7 @@ from pydantic import BaseModel import litellm from litellm import ModelResponse, completion from litellm.llms.anthropic.pass_through.messages import handler as anthropic_messages_handler -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -3441,6 +3441,48 @@ async def test_google_ai_studio_forwards_seed_to_generation_config(drop_params: assert response.choices[0].message.content == "seed=42" +def test_google_ai_studio_sync_stream_forwards_request_timeout() -> None: + expected_timeout: Final = 2.5 + + def respond(request: httpx.Request) -> httpx.Response: + assert request.extensions["timeout"]["read"] == expected_timeout + return httpx.Response(200, content=b"", request=request) + + http_client: Final = httpx.Client(transport=httpx.MockTransport(respond), timeout=91.0) + stream: Final = litellm.completion( + model="gemini/gemini-3.8-flash", + messages=[{"role": "user", "content": "hi"}], + stream=True, + timeout=expected_timeout, + api_key="fake-gemini-key", + client=HTTPHandler(client=http_client), + ) + + assert isinstance(stream, CustomStreamWrapper) + stream.fetch_sync_stream() + + +@pytest.mark.asyncio +async def test_google_ai_studio_async_stream_forwards_request_timeout() -> None: + expected_timeout: Final = 3.5 + + def respond(request: httpx.Request) -> httpx.Response: + assert request.extensions["timeout"]["read"] == expected_timeout + return httpx.Response(200, content=b"", request=request) + + stream: Final = await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[{"role": "user", "content": "hi"}], + stream=True, + timeout=httpx.Timeout(expected_timeout), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(timeout=91.0, transport=httpx.MockTransport(respond)), + ) + + assert isinstance(stream, CustomStreamWrapper) + await stream.fetch_stream() + + # ==================== Tool Type Separation Tests ==================== # These tests verify that each Tool object contains exactly one type per Vertex AI API spec # Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1beta1/Tool