From d0f3a2dfdb65385e8359db8953d6ccce53e79605 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 14 Oct 2026 10:51:18 -0700 Subject: [PATCH] fix(vertex_ai): pass the per-request timeout to Gemini streaming calls (#46349) * fix(vertex_ai): pass the per-request timeout to Gemini streaming calls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(vertex_ai): pass timeout to make_call and make_sync_call in the http 201 tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex_ai): default timeout to None on make_call and make_sync_call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../vertex_and_google_ai_studio_gemini.py | 12 +- .../test_vertex_gemini_stream_timeout_wire.py | 147 ++++++++++++++++++ 2 files changed, 157 insertions(+), 2 deletions(-) create mode 100644 tests/integration/sdk/test_vertex_gemini_stream_timeout_wire.py 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 517ac9e98b1..109dddf44ed 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 @@ -2748,6 +2748,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 @@ -2757,7 +2758,9 @@ 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()) @@ -2800,13 +2803,16 @@ 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( @@ -2924,6 +2930,7 @@ class VertexLLM(VertexBase): model=model, messages=messages, logging_obj=logging_obj, + timeout=timeout, ), model=model, custom_llm_provider="vertex_ai_beta", @@ -3200,6 +3207,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/integration/sdk/test_vertex_gemini_stream_timeout_wire.py b/tests/integration/sdk/test_vertex_gemini_stream_timeout_wire.py new file mode 100644 index 00000000000..68d98961dce --- /dev/null +++ b/tests/integration/sdk/test_vertex_gemini_stream_timeout_wire.py @@ -0,0 +1,147 @@ +from __future__ import annotations + +import asyncio +import json +import threading +import time +from collections.abc import Callable +from concurrent.futures import Future, ThreadPoolExecutor +from concurrent.futures import TimeoutError as FutureTimeout +from typing import Final + +import pytest +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server + +import litellm +from litellm.exceptions import MidStreamFallbackError + +_MODEL: Final = "gemini-3.8-flash" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_MESSAGES: Final = [{"role": "user", "content": "does the per-request timeout reach a stalled stream?"}] +_REQUEST_TIMEOUT_SECONDS: Final = 1.0 +_HANG_DEADLINE_SECONDS: Final = 15.0 +_TOKEN_REPLY: Final = Reply( + body=json.dumps({"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600}).encode() +) + + +def _sse(text: str, finish_reason: str | None) -> bytes: + candidate: Final = { + "content": {"role": "model", "parts": [{"text": text}]}, + **({"finishReason": finish_reason} if finish_reason else {}), + } + usage: Final = {"promptTokenCount": 9, "candidatesTokenCount": 2, "totalTokenCount": 11} + return b"data: " + json.dumps({"candidates": [candidate], "usageMetadata": usage}).encode() + b"\n\n" + + +def _stream_that_stalls_after_the_first_chunk() -> Reply: + return Reply( + content_type="text/event-stream", + chunks=(_sse("partial", None), _sse("", "STOP")), + gate_after_first=threading.Event(), + ) + + +def _stream_that_completes() -> Reply: + return Reply(content_type="text/event-stream", chunks=(_sse("scripted ", None), _sse("answer", "STOP"))) + + +def _body_that_stalls_half_way() -> Reply: + return Reply(chunks=(b'{"candidates": [', b"]}"), gate_after_first=threading.Event()) + + +def _peer(reply: Reply) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == "/_oauth/token": + return _TOKEN_REPLY + assert request.method == "POST", request.method + return reply + + return respond + + +def _deployment(route: str, wire: Wire) -> dict[str, str]: + if route == "gemini": + return {"model": f"gemini/{_MODEL}", "api_base": wire.url, "api_key": "scripted-key"} + return { + "model": f"vertex_ai/{_MODEL}", + "api_base": wire.url, + "vertex_project": _PROJECT, + "vertex_location": _LOCATION, + "vertex_credentials": service_account_json(_PROJECT, wire.url), + } + + +def _text_of(chunk: object) -> str: + choices: Final = getattr(chunk, "choices", None) or () + return "".join(str(choice.delta.content or "") for choice in choices) + + +def _complete_sync(route: str, wire: Wire, stream: bool) -> str: + response: Final = litellm.completion( + **_deployment(route, wire), messages=_MESSAGES, stream=stream, timeout=_REQUEST_TIMEOUT_SECONDS + ) + if stream: + return "".join(_text_of(chunk) for chunk in response) + return str(response.choices[0].message.content) + + +async def _complete_async(route: str, wire: Wire, stream: bool) -> str: + response: Final = await litellm.acompletion( + **_deployment(route, wire), messages=_MESSAGES, stream=stream, timeout=_REQUEST_TIMEOUT_SECONDS + ) + if stream: + parts: Final = [_text_of(chunk) async for chunk in response] + return "".join(parts) + return str(response.choices[0].message.content) + + +def _complete(client: str, route: str, wire: Wire, stream: bool) -> str: + if client == "async": + return asyncio.run(_complete_async(route, wire, stream)) + return _complete_sync(route, wire, stream) + + +def _before_the_hang_deadline(call: Callable[[], str]) -> str: + pool: Final = ThreadPoolExecutor(max_workers=1) + future: Final[Future[str]] = pool.submit(call) + pool.shutdown(wait=False) + try: + return future.result(timeout=_HANG_DEADLINE_SECONDS) + except FutureTimeout: + pytest.fail( + f"still waiting on the peer {_HANG_DEADLINE_SECONDS}s after a {_REQUEST_TIMEOUT_SECONDS}s request timeout" + ) + + +_CLIENTS: Final = ("sync", "async") +_ROUTES: Final = ("vertex_ai", "gemini") + + +@pytest.mark.parametrize("route", _ROUTES) +@pytest.mark.parametrize("client", _CLIENTS) +def test_stalled_stream_raises_at_the_request_timeout(client: str, route: str) -> None: + with wire_server(_peer(_stream_that_stalls_after_the_first_chunk())) as wire: + started: Final = time.monotonic() + with pytest.raises(MidStreamFallbackError): + _before_the_hang_deadline(lambda: _complete(client, route, wire, stream=True)) + elapsed: Final = time.monotonic() - started + assert elapsed >= _REQUEST_TIMEOUT_SECONDS, elapsed + assert len(wire.drain()) >= 1 + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_stalled_body_raises_at_the_request_timeout(client: str) -> None: + with wire_server(_peer(_body_that_stalls_half_way())) as wire: + started: Final = time.monotonic() + with pytest.raises(litellm.Timeout): + _before_the_hang_deadline(lambda: _complete(client, "vertex_ai", wire, stream=False)) + assert time.monotonic() - started >= _REQUEST_TIMEOUT_SECONDS + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_stream_with_a_request_timeout_still_completes(client: str) -> None: + with wire_server(_peer(_stream_that_completes())) as wire: + assert _complete(client, "vertex_ai", wire, stream=True) == "scripted answer"