From e31592627c1b8e28c9a4d03946d0ef8cd405fc6d Mon Sep 17 00:00:00 2001
From: Anton Dubovik <anton_dubovik@epam.com>
Date: Wed, 17 Jul 2024 18:04:06 +0100
Subject: [PATCH 1/2] feat: added base64 str as embedding representation in the
 response (#138)

---
 aidial_sdk/embeddings/response.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/aidial_sdk/embeddings/response.py b/aidial_sdk/embeddings/response.py
index bec3c69..3f8d559 100644
--- 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

From 608bd4e5de45fae8ea28a294e2d159c2f06c8715 Mon Sep 17 00:00:00 2001
From: Anton Dubovik <anton_dubovik@epam.com>
Date: Thu, 18 Jul 2024 12:04:54 +0100
Subject: [PATCH 2/2] fix: fixed type of tokens in embeddings request from
 float to int (#139)

---
 aidial_sdk/embeddings/request.py |  4 ++--
 tests/test_embeddings.py         | 14 +++++++++++++-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/aidial_sdk/embeddings/request.py b/aidial_sdk/embeddings/request.py
index 4cb8bd6..5fb9d7f 100644
--- 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
index e024821..622fa92 100644
--- 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(