mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(vertex_ai): forward Gemini streaming timeouts
This commit is contained in:
parent
22b36cbcf6
commit
4a62921ed5
2 changed files with 63 additions and 3 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue