fix: fixed type of tokens in embeddings request from float to int (#139)

epam · Jul 18, 2024 · 608bd4e · 608bd4e
1 parent e315926
commit 608bd4e
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 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/tests/test_embeddings.py b/tests/test_embeddings.py
@@ -30,7 +30,19 @@
     TestCase(
         simple,
         "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(