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 0905f22362e3..c136de9a1e6b 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 @@ -2362,7 +2362,8 @@ def validate_environment( 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 @@ async def async_streaming( 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, diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 8beb19bf1ac0..cb6990b1be78 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -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