phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix get_async_httpx_client vertex

Browse source
This commit is contained in:
Ishaan Jaff 2024-11-21 10:22:30 -08:00
parent f7f9e8c41f
commit 0a10b1ef1c
2 changed files with 18 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

11
litellm/llms/vertex_ai_and_google_ai_studio/image_generation/image_generation_handler.py
View file

@ -5,7 +5,11 @@ import httpx
from openai.types.image import Image
import litellm
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
get_async_httpx_client,
)
from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import (
VertexLLM,
)
@ -156,7 +160,10 @@ class VertexImageGeneration(VertexLLM):
else:
_params["timeout"] = httpx.Timeout(timeout=600.0, connect=5.0)
self.async_handler = AsyncHTTPHandler(**_params) # type: ignore
self.async_handler = get_async_httpx_client(
llm_provider=litellm.LlmProviders.VERTEX_AI,
params={"timeout": timeout},
)
else:
self.async_handler = client # type: ignore

11
litellm/llms/vertex_ai_and_google_ai_studio/multimodal_embeddings/embedding_handler.py
View file

@ -5,7 +5,11 @@ import httpx
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
get_async_httpx_client,
)
from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import (
VertexAIError,
VertexLLM,
@ -172,7 +176,10 @@ class VertexMultimodalEmbedding(VertexLLM):
if isinstance(timeout, float) or isinstance(timeout, int):
timeout = httpx.Timeout(timeout)
_params["timeout"] = timeout
client = AsyncHTTPHandler(**_params) # type: ignore
client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.VERTEX_AI,
params={"timeout": timeout},
)
else:
client = client # type: ignore