feat: add input_type for text_embedding

langgenius · Sep 25, 2024 · bc6ad35 · bc6ad35
1 parent 9fc1362
commit bc6ad35
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/api/core/model_runtime/model_providers/voyage/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/voyage/text_embedding/text_embedding.py
@@ -34,7 +34,7 @@ def _invoke(
         credentials: dict,
         texts: list[str],
         user: Optional[str] = None,
-        input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
+        input_type: EmbeddingInputType = None,
     ) -> TextEmbeddingResult:
         """
         Invoke text embedding model
@@ -43,6 +43,7 @@ def _invoke(
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         api_key = credentials["api_key"]
@@ -54,8 +55,10 @@ def _invoke(
 
         url = base_url + "/embeddings"
         headers = {"Authorization": "Bearer " + api_key, "Content-Type": "application/json"}
-
-        data = {"model": model, "input": texts}
+        voyage_input_type = "null"
+        if input_type is not None:
+            voyage_input_type = input_type.value
+        data = {"model": model, "input": texts, "input_type": voyage_input_type}
 
         try:
             response = requests.post(url, headers=headers, data=dumps(data))