feat: added base64 str as embedding representation in the response (#138

)
epam · Jul 17, 2024 · e315926 · e315926
1 parent 8fbc61a
commit e315926
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/aidial_sdk/embeddings/response.py b/aidial_sdk/embeddings/response.py
@@ -1,10 +1,10 @@
-from typing import List, Literal
+from typing import List, Literal, Union
 
 from aidial_sdk.utils.pydantic import ExtraForbidModel
 
 
 class Embedding(ExtraForbidModel):
-    embedding: List[float]
+    embedding: Union[str, List[float]]
     index: int
     object: Literal["embedding"] = "embedding"
 
@@ -14,11 +14,11 @@ class Usage(ExtraForbidModel):
     total_tokens: int
 
 
-class CreateEmbeddingResponse(ExtraForbidModel):
+class EmbeddingResponse(ExtraForbidModel):
     data: List[Embedding]
     model: str
     object: Literal["list"] = "list"
     usage: Usage
 
 
-Response = CreateEmbeddingResponse
+Response = EmbeddingResponse