Fix: User specified async client ignored with Gemini streaming+async

The user-specified async client was being overwritten by
`litellm.module_level_aclient` in `streaming_handler.py` when using
async+streaming with Gemini.

This fix adds a `gemini_client` parameter to `make_call()` (matching
the existing pattern in `make_sync_call()`) so the user's custom client
is preserved and not overwritten.

Fixes #17148
This commit is contained in:
Chesars 2025-12-05 16:50:46 -03:00
parent 8665e92aa8
commit 872554df42
2 changed files with 76 additions and 2 deletions

View file

@ -2362,7 +2362,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
async def make_call(
client: Optional[AsyncHTTPHandler],
client: Optional[AsyncHTTPHandler], # module-level client
gemini_client: Optional[AsyncHTTPHandler], # if passed by user
api_base: str,
headers: dict,
data: str,
@ -2370,6 +2371,8 @@ async def make_call(
messages: list,
logging_obj,
):
if gemini_client is not None:
client = gemini_client
if client is None:
client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.VERTEX_AI,
@ -2541,7 +2544,11 @@ class VertexLLM(VertexBase):
completion_stream=None,
make_call=partial(
make_call,
client=client,
gemini_client=(
client
if client is not None and isinstance(client, AsyncHTTPHandler)
else None
),
api_base=api_base,
headers=headers,
data=request_body_str,

View file

@ -3724,3 +3724,70 @@ def test_vertex_ai_usage_metadata_video_tokens_with_caching():
assert result.prompt_tokens_details.text_tokens == 9
assert result.prompt_tokens_details.audio_tokens == 200
def test_async_streaming_uses_custom_client():
"""
Test that user-specified async client is correctly passed to make_call
for async streaming calls.
Fixes: https://github.com/BerriAI/litellm/issues/17148
"""
from functools import partial
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
make_call,
)
# Create a mock async client
mock_client = MagicMock(spec=AsyncHTTPHandler)
# Create a partial function like the code does in async_streaming
partial_make_call = partial(
make_call,
gemini_client=mock_client,
api_base="https://example.com",
headers={},
data="{}",
model="gemini-pro",
messages=[],
logging_obj=MagicMock(),
)
# Verify that gemini_client is in the partial's keywords
assert "gemini_client" in partial_make_call.keywords
assert partial_make_call.keywords["gemini_client"] is mock_client
def test_sync_streaming_uses_custom_client():
"""
Test that user-specified sync client is correctly passed to make_sync_call
for sync streaming calls.
This verifies the existing behavior that we want to match for async.
"""
from functools import partial
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
make_sync_call,
)
# Create a mock sync client
mock_client = MagicMock(spec=HTTPHandler)
# Create a partial function like the code does in sync streaming
partial_make_sync_call = partial(
make_sync_call,
gemini_client=mock_client,
api_base="https://example.com",
headers={},
data="{}",
model="gemini-pro",
messages=[],
logging_obj=MagicMock(),
)
# Verify that gemini_client is in the partial's keywords
assert "gemini_client" in partial_make_sync_call.keywords
assert partial_make_sync_call.keywords["gemini_client"] is mock_client