mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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:
parent
8665e92aa8
commit
872554df42
2 changed files with 76 additions and 2 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue