Merge 4a62921ed5 into f553c80cd3

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -2568,6 +2568,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
@@ -2577,7 +2578,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())
@@ -2620,13 +2628,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(
@@ -2744,6 +2760,7 @@ class VertexLLM(VertexBase):
                 model=model,
                 messages=messages,
                 logging_obj=logging_obj,
+                timeout=timeout,
             ),
             model=model,
             custom_llm_provider="vertex_ai_beta",
@@ -3020,6 +3037,7 @@ class VertexLLM(VertexBase):
                     messages=messages,
                     logging_obj=logging_obj,
                     headers=headers,
+                    timeout=timeout,
                 ),
                 model=model,
                 custom_llm_provider="vertex_ai_beta",

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