From e43e1e9ab142726bf2f0b4fafb28c09219e4623f Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 25 Nov 2023 14:19:14 -0800 Subject: [PATCH] fix(utils.py): fix linting errors --- litellm/utils.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 0da420239a..6e64a3b8a0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -357,7 +357,7 @@ class EmbeddingResponse(OpenAIObject): model: Optional[str] = None """The model used for embedding.""" - data: Optional[Embedding] = None + data: Optional[List] = None """The actual embedding value""" object: str @@ -3323,15 +3323,17 @@ def convert_to_model_response_object(response_object: Optional[dict]=None, model if "object" in response_object: model_response_object.object = response_object["object"] - data = [] - for idx, embedding in enumerate(response_object["data"]): # type: ignore - embedding_obj = Embedding( # type: ignore + embedding_data = [] + for idx, embedding in enumerate(response_object["data"]): + embedding_obj = Embedding( embedding=embedding.get("embedding", None), index = embedding.get("index", idx), object=embedding.get("object", "embedding") ) - data.append(embedding_obj) # type: ignore - model_response_object.data = data + if isinstance(embedding_obj, EmbeddingResponse): + embedding_data.append(embedding_obj) + if len(embedding_data) > 0: + model_response_object.data = embedding_data if "usage" in response_object and response_object["usage"] is not None: model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) # type: ignore