Merge branch 'development' into feat/supported-extensible-routes

epam · Jul 19, 2024 · 8152f4c · 8152f4c
2 parents 9585418 + 608bd4e
commit 8152f4c
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 7 deletions.
diff --git a/aidial_sdk/embeddings/request.py b/aidial_sdk/embeddings/request.py
@@ -2,14 +2,14 @@
 
 from aidial_sdk.chat_completion.request import Attachment
 from aidial_sdk.deployment.from_request_mixin import FromRequestDeploymentMixin
-from aidial_sdk.pydantic_v1 import StrictFloat, StrictInt, StrictStr
+from aidial_sdk.pydantic_v1 import StrictInt, StrictStr
 from aidial_sdk.utils.pydantic import ExtraForbidModel
 
 
 class AzureEmbeddingsRequest(ExtraForbidModel):
     model: Optional[StrictStr] = None
     input: Union[
-        StrictStr, List[StrictStr], List[StrictFloat], List[List[StrictFloat]]
+        StrictStr, List[StrictStr], List[StrictInt], List[List[StrictInt]]
     ]
     encoding_format: Literal["float", "base64"] = "float"
     dimensions: Optional[StrictInt] = None

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
diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py
@@ -33,7 +33,19 @@
         app,
         deployment,
         "embeddings",
-        {"input": "a", "custom_fields": {"type": "query"}},
+        {
+            "input": "a",
+            "custom_fields": {
+                "type": "query",
+                "instruction": "instruction",
+            },
+        },
+        expected_response_1,
+    ),
+    TestCase(
+        simple,
+        "embeddings",
+        {"input": [15339]},
         expected_response_1,
     ),
     TestCase(