fix(utils.py): fix linting errors

This commit is contained in:
Krrish Dholakia 2023-11-25 14:19:14 -08:00
parent ec81b393e2
commit e43e1e9ab1

View file

@ -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