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(