From 17e07423346a244c545bccfe65a8d9ae3ab1fb54 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 18 Mar 2025 12:31:44 -0700 Subject: [PATCH] fix typing errors --- litellm/main.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 64049c31d1..85aa0e96a9 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -25,6 +25,7 @@ from functools import partial from typing import ( Any, Callable, + Coroutine, Dict, List, Literal, @@ -3288,7 +3289,7 @@ def embedding( # noqa: PLR0915 litellm_call_id=None, logger_fn=None, **kwargs, -) -> EmbeddingResponse: +) -> Union[EmbeddingResponse, Coroutine[Any, Any, EmbeddingResponse]]: """ Embedding function that calls an API to generate embeddings for the given input. @@ -3409,7 +3410,9 @@ def embedding( # noqa: PLR0915 if mock_response is not None: return mock_embedding(model=model, mock_response=mock_response) try: - response: Optional[EmbeddingResponse] = None + response: Optional[ + Union[EmbeddingResponse, Coroutine[Any, Any, EmbeddingResponse]] + ] = None if azure is True or custom_llm_provider == "azure": # azure configs @@ -3901,7 +3904,11 @@ def embedding( # noqa: PLR0915 raise LiteLLMUnknownProvider( model=model, custom_llm_provider=custom_llm_provider ) - if response is not None and hasattr(response, "_hidden_params"): + if ( + response is not None + and hasattr(response, "_hidden_params") + and isinstance(response, EmbeddingResponse) + ): response._hidden_params["custom_llm_provider"] = custom_llm_provider if response is None: