diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 35cd54d65f6..b0c7c4ead04 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -158,6 +158,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): logging_obj=logging_obj, optional_params=optional_params, litellm_params=litellm_params, + client= client, timeout=timeout, encoding=encoding, ) @@ -172,6 +173,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): logging_obj=logging_obj, optional_params=optional_params, litellm_params=litellm_params, + client= client, timeout=timeout, encoding=encoding, ) @@ -187,11 +189,12 @@ class VertexGemmaConfig(OpenAIGPTConfig): logging_obj: Any, optional_params: dict, litellm_params: dict, + client= Optional[Any], timeout: Optional[Union[float, httpx.Timeout]], encoding: Any, ): """Synchronous completion request""" - from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import convert_to_model_response_object # Check if streaming is requested (will be faked) @@ -223,7 +226,12 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) # Make the HTTP request - http_handler = HTTPHandler(concurrent_limit=1) + if client is None or isinstance(client, (httpx.AsyncClient, AsyncHTTPHandler)): + http_handler = HTTPHandler(concurrent_limit=1) + elif isinstance(client, httpx.Client): + http_handler = HTTPHandler(client=client) + else: + http_handler = client response = http_handler.post( url=api_base, headers=headers, @@ -279,11 +287,16 @@ class VertexGemmaConfig(OpenAIGPTConfig): logging_obj: Any, optional_params: dict, litellm_params: dict, + client: Optional[Any], timeout: Optional[Union[float, httpx.Timeout]], encoding: Any, ): """Asynchronous completion request""" - from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + get_async_httpx_client, + ) from litellm.types.utils import LlmProviders from litellm.utils import convert_to_model_response_object @@ -316,9 +329,13 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) # Make the HTTP request - http_handler = get_async_httpx_client( - llm_provider=LlmProviders.VERTEX_AI, - ) + if client is None or isinstance(client, (httpx.Client, HTTPHandler)): + http_handler = get_async_httpx_client( + llm_provider=LlmProviders.VERTEX_AI, + ) + else: + http_handler = client + response = await http_handler.post( url=api_base, headers=headers,