From c19565b6d184fe36d39bba3902bdbb72e34e340a Mon Sep 17 00:00:00 2001
From: Mohammad Kalim Akram <kalimakram@gmail.com>
Date: Mon, 16 Sep 2024 13:52:52 +0200
Subject: [PATCH] refactor: support jina embeddings v3 (#39)

This PR should allow pymilvus to support `jina-embeddings-v3:

with 2 additional fields:

task: to properly load adapters
dimensions: allow the user to truncate dimension with MRL.

This PR is the continuation of
https://github.com/milvus-io/milvus-model/pull/38 and we can close that.

Signed-off-by: Mohammad Kalim Akram <kalim.akram@jina.ai>
---
 milvus_model/dense/jinaai.py | 26 ++++++++++++++++++--------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/milvus_model/dense/jinaai.py b/milvus_model/dense/jinaai.py
index c458cee..3ca1d65 100644
--- a/milvus_model/dense/jinaai.py
+++ b/milvus_model/dense/jinaai.py
@@ -12,8 +12,10 @@
 class JinaEmbeddingFunction(BaseEmbeddingFunction):
     def __init__(
         self,
-        model_name: str = "jina-embeddings-v2-base-en",
+        model_name: str = "jina-embeddings-v3",
         api_key: Optional[str] = None,
+        task: str = 'retrieval.passage',
+        dimensions: Optional[int] = None,
         **kwargs,
     ):
         if api_key is None:
@@ -33,8 +35,8 @@ def __init__(
         self._session.headers.update(
             {"Authorization": f"Bearer {self.api_key}", "Accept-Encoding": "identity"}
         )
-        self.model_name = model_name
-        self._dim = None
+        self.task = task
+        self._dim = dimensions
 
     @property
     def dim(self):
@@ -43,17 +45,25 @@ def dim(self):
         return self._dim
 
     def encode_queries(self, queries: List[str]) -> List[np.array]:
-        return self._call_jina_api(queries)
+        return self._call_jina_api(queries, task='retrieval.query')
 
     def encode_documents(self, documents: List[str]) -> List[np.array]:
-        return self._call_jina_api(documents)
+        return self._call_jina_api(documents, task='retrieval.passage')
 
     def __call__(self, texts: List[str]) -> List[np.array]:
-        return self._call_jina_api(texts)
+        return self._call_jina_api(texts, task=self.task)
 
-    def _call_jina_api(self, texts: List[str]):
+    def _call_jina_api(self, texts: List[str], task: Optional[str] = None):
+        data = {
+            "input": texts,
+            "model": self.model_name,
+            "task": task,
+        }
+        if self._dim is not None:
+            data["dimensions"] = self._dim
         resp = self._session.post(  # type: ignore[assignment]
-            API_URL, json={"input": texts, "model": self.model_name},
+            API_URL,
+            json=data,
         ).json()
         if "data" not in resp:
             raise RuntimeError(resp["detail"])