From 3618a97c20e386aa49a39c1f6af55ce7414ffc5c Mon Sep 17 00:00:00 2001
From: Aaron Ji <127167174+DresAaron@users.noreply.github.com>
Date: Mon, 23 Sep 2024 13:45:09 +0800
Subject: [PATCH] feat: extend api params for Jina Embeddings V3 (#8657)

---
 .../model_providers/jina/jina.yaml            | 43 +++++++++++++++++++
 .../jina/text_embedding/text_embedding.py     | 42 ++++++++++++++----
 2 files changed, 76 insertions(+), 9 deletions(-)

diff --git a/api/core/model_runtime/model_providers/jina/jina.yaml b/api/core/model_runtime/model_providers/jina/jina.yaml
index 23e18ad75f..9c70d6ff33 100644
--- a/api/core/model_runtime/model_providers/jina/jina.yaml
+++ b/api/core/model_runtime/model_providers/jina/jina.yaml
@@ -67,3 +67,46 @@ model_credential_schema:
       required: false
       type: text-input
       default: '8192'
+    - variable: task
+      label:
+        zh_Hans: 下游任务
+        en_US: Downstream task
+      placeholder:
+        zh_Hans: 选择将使用向量模型的下游任务。模型将返回针对该任务优化的向量。
+        en_US: Select the downstream task for which the embeddings will be used. The model will return the optimized embeddings for that task.
+      required: false
+      type: select
+      options:
+        - value: retrieval.query
+          label:
+            en_US: retrieval.query
+        - value: retrieval.passage
+          label:
+            en_US: retrieval.passage
+        - value: separation
+          label:
+            en_US: separation
+        - value: classification
+          label:
+            en_US: classification
+        - value: text-matching
+          label:
+            en_US: text-matching
+    - variable: dimensions
+      label:
+        zh_Hans: 输出维度
+        en_US: Output dimensions
+      placeholder:
+        zh_Hans: 输入您的输出维度
+        en_US: Enter output dimensions
+      required: false
+      type: text-input
+    - variable: late_chunking
+      label:
+        zh_Hans: 后期分块
+        en_US: Late chunking
+      placeholder:
+        zh_Hans: 应用后期分块技术来利用模型的长上下文功能来生成上下文块向量化。
+        en_US: Apply the late chunking technique to leverage the model's long-context capabilities for generating contextual chunk embeddings.
+      required: false
+      type: switch
diff --git a/api/core/model_runtime/model_providers/jina/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/jina/text_embedding/text_embedding.py
index ceb79567d5..6c96699ea2 100644
--- a/api/core/model_runtime/model_providers/jina/text_embedding/text_embedding.py
+++ b/api/core/model_runtime/model_providers/jina/text_embedding/text_embedding.py
@@ -27,6 +27,38 @@ class JinaTextEmbeddingModel(TextEmbeddingModel):
 
     api_base: str = "https://api.jina.ai/v1"
 
+    def _to_payload(self, model: str, texts: list[str], credentials: dict) -> dict:
+        """
+        Parse model credentials
+
+        :param model: model name
+        :param credentials: model credentials
+        :param texts: texts to embed
+        :return: parsed credentials
+        """
+
+        def transform_jina_input_text(model, text):
+            if model == "jina-clip-v1":
+                return {"text": text}
+            return text
+
+        data = {"model": model, "input": [transform_jina_input_text(model, text) for text in texts]}
+
+        task = credentials.get("task")
+        dimensions = credentials.get("dimensions")
+        late_chunking = credentials.get("late_chunking")
+
+        if task is not None:
+            data["task"] = task
+
+        if dimensions is not None:
+            data["dimensions"] = int(dimensions)
+
+        if late_chunking is not None:
+            data["late_chunking"] = late_chunking
+
+        return data
+
     def _invoke(
         self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None
     ) -> TextEmbeddingResult:
@@ -49,15 +81,7 @@ class JinaTextEmbeddingModel(TextEmbeddingModel):
         url = base_url + "/embeddings"
         headers = {"Authorization": "Bearer " + api_key, "Content-Type": "application/json"}
 
-        def transform_jina_input_text(model, text):
-            if model == "jina-clip-v1":
-                return {"text": text}
-            return text
-
-        data = {"model": model, "input": [transform_jina_input_text(model, text) for text in texts]}
-
-        if model == "jina-embeddings-v3":
-            data["task"] = "text-matching"
+        data = self._to_payload(model=model, texts=texts, credentials=credentials)
 
         try:
             response = post(url, headers=headers, data=dumps(data))