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

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -2362,14 +2362,17 @@ 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,
     model: str,
     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,

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