[ENH] Update Jina embedding function to support all models and configurations

jairad26 · jairad26 · commit 5bcd97d940bd · 2025-04-09T08:27:36.000-07:00
diff --git a/chromadb/utils/embedding_functions/jina_embedding_function.py b/chromadb/utils/embedding_functions/jina_embedding_function.py
@@ -8,14 +8,20 @@
 class JinaEmbeddingFunction(EmbeddingFunction[Documents]):
     """
     This class is used to get embeddings for a list of texts using the Jina AI API.
-    It requires an API key and a model name. The default model name is "jina-embeddings-v2-base-en".
+    It requires an API key and a model name. The default model name is "jina-embeddings-v3".
     """
 
     def __init__(
         self,
         api_key: Optional[str] = None,
-        model_name: str = "jina-embeddings-v2-base-en",
+        model_name: str = "jina-embeddings-v3",
         api_key_env_var: str = "CHROMA_JINA_API_KEY",
+        task: Optional[str] = None,
+        late_chunking: Optional[bool] = None,
+        truncate: Optional[bool] = None,
+        dimensions: Optional[int] = None,
+        embedding_type: Optional[str] = None,
+        normalized: Optional[bool] = None,
     ):
         """
         Initialize the JinaEmbeddingFunction.
@@ -24,7 +30,7 @@ def __init__(
             api_key_env_var (str, optional): Environment variable name that contains your API key for the Jina AI API.
                 Defaults to "CHROMA_JINA_API_KEY".
             model_name (str, optional): The name of the model to use for text embeddings.
-                Defaults to "jina-embeddings-v2-base-en".
+                Defaults to "jina-embeddings-v3".
         """
         try:
             import httpx
@@ -40,6 +46,14 @@ def __init__(
 
         self.model_name = model_name
 
+        # Initialize optional attributes to None
+        self.task = task
+        self.late_chunking = late_chunking
+        self.truncate = truncate
+        self.dimensions = dimensions
+        self.embedding_type = embedding_type
+        self.normalized = normalized
+
         self._api_url = "https://api.jina.ai/v1/embeddings"
         self._session = httpx.Client()
         self._session.headers.update(
@@ -51,7 +65,7 @@ def __call__(self, input: Documents) -> Embeddings:
         Get the embeddings for a list of texts.
 
         Args:
-            input (Documents): A list of texts or images to get embeddings for.
+            input (Documents): A list of texts to get embeddings for.
 
         Returns:
             Embeddings: The embeddings for the texts.
@@ -64,10 +78,31 @@ def __call__(self, input: Documents) -> Embeddings:
         if not all(isinstance(item, str) for item in input):
             raise ValueError("Jina AI only supports text documents, not images")
 
+        payload: Dict[str, Any] = {
+            "input": input,
+            "model": self.model_name,
+        }
+
+        if self.task is not None:
+            payload["task"] = self.task
+
+        if self.late_chunking is not None:
+            payload["late_chunking"] = self.late_chunking
+
+        if self.truncate is not None:
+            payload["truncate"] = self.truncate
+
+        if self.dimensions is not None:
+            payload["dimensions"] = self.dimensions
+
+        if self.embedding_type is not None:
+            payload["embedding_type"] = self.embedding_type
+
+        if self.normalized is not None:
+            payload["normalized"] = self.normalized
+
         # Call Jina AI Embedding API
-        resp = self._session.post(
-            self._api_url, json={"input": input, "model": self.model_name}
-        ).json()
+        resp = self._session.post(self._api_url, json=payload).json()
 
         if "data" not in resp:
             raise RuntimeError(resp.get("detail", "Unknown error"))
@@ -97,16 +132,38 @@ def supported_spaces(self) -> List[Space]:
     def build_from_config(config: Dict[str, Any]) -> "EmbeddingFunction[Documents]":
         api_key_env_var = config.get("api_key_env_var")
         model_name = config.get("model_name")
+        task = config.get("task")
+        late_chunking = config.get("late_chunking")
+        truncate = config.get("truncate")
+        dimensions = config.get("dimensions")
+        embedding_type = config.get("embedding_type")
+        normalized = config.get("normalized")
 
         if api_key_env_var is None or model_name is None:
-            assert False, "This code should not be reached"
+            assert False, "This code should not be reached"  # this is for type checking
 
         return JinaEmbeddingFunction(
-            api_key_env_var=api_key_env_var, model_name=model_name
+            api_key_env_var=api_key_env_var,
+            model_name=model_name,
+            task=task,
+            late_chunking=late_chunking,
+            truncate=truncate,
+            dimensions=dimensions,
+            embedding_type=embedding_type,
+            normalized=normalized,
         )
 
     def get_config(self) -> Dict[str, Any]:
-        return {"api_key_env_var": self.api_key_env_var, "model_name": self.model_name}
+        return {
+            "api_key_env_var": self.api_key_env_var,
+            "model_name": self.model_name,
+            "task": self.task,
+            "late_chunking": self.late_chunking,
+            "truncate": self.truncate,
+            "dimensions": self.dimensions,
+            "embedding_type": self.embedding_type,
+            "normalized": self.normalized,
+        }
 
     def validate_config_update(
         self, old_config: Dict[str, Any], new_config: Dict[str, Any]
diff --git a/clients/js/packages/chromadb-core/src/embeddings/JinaEmbeddingFunction.ts b/clients/js/packages/chromadb-core/src/embeddings/JinaEmbeddingFunction.ts
@@ -3,24 +3,59 @@ import { validateConfigSchema } from "../schemas/schemaUtils";
 type StoredConfig = {
   api_key_env_var: string;
   model_name: string;
+  task?: string;
+  late_chunking?: boolean;
+  truncate?: boolean;
+  dimensions?: number;
+  embedding_type?: string;
+  normalized?: boolean;
 };
 
+interface JinaRequestBody {
+  input: string[];
+  model: string;
+  task?: string;
+  late_chunking?: boolean;
+  truncate?: boolean;
+  dimensions?: number;
+  embedding_type?: string;
+  normalized?: boolean;
+}
+
 export class JinaEmbeddingFunction implements IEmbeddingFunction {
   name = "jina";
 
   private api_key_env_var: string;
   private model_name: string;
   private api_url: string;
   private headers: { [key: string]: string };
+  private task: string | undefined;
+  private late_chunking: boolean | undefined;
+  private truncate: boolean | undefined;
+  private dimensions: number | undefined;
+  private embedding_type: string | undefined;
+  private normalized: boolean | undefined;
 
   constructor({
     jinaai_api_key,
     model_name = "jina-embeddings-v2-base-en",
     api_key_env_var = "JINAAI_API_KEY",
+    task,
+    late_chunking,
+    truncate,
+    dimensions,
+    embedding_type,
+    normalized,
   }: {
     jinaai_api_key?: string;
     model_name?: string;
     api_key_env_var: string;
+    task?: string;
+    late_chunking?: boolean;
+    truncate?: boolean;
+    dimensions?: number;
+    embedding_type?: string;
+    normalized?: boolean;
   }) {
     const apiKey = jinaai_api_key ?? process.env[api_key_env_var];
     if (!apiKey) {
@@ -31,6 +66,12 @@ export class JinaEmbeddingFunction implements IEmbeddingFunction {
 
     this.model_name = model_name;
     this.api_key_env_var = api_key_env_var;
+    this.task = task;
+    this.late_chunking = late_chunking;
+    this.truncate = truncate;
+    this.dimensions = dimensions;
+    this.embedding_type = embedding_type;
+    this.normalized = normalized;
 
     this.api_url = "https://api.jina.ai/v1/embeddings";
     this.headers = {
@@ -41,14 +82,40 @@ export class JinaEmbeddingFunction implements IEmbeddingFunction {
   }
 
   public async generate(texts: string[]) {
+    let json_body: JinaRequestBody = {
+      input: texts,
+      model: this.model_name,
+    };
+
+    if (this.task) {
+      json_body.task = this.task;
+    }
+
+    if (this.late_chunking) {
+      json_body.late_chunking = this.late_chunking;
+    }
+
+    if (this.truncate) {
+      json_body.truncate = this.truncate;
+    }
+
+    if (this.dimensions) {
+      json_body.dimensions = this.dimensions;
+    }
+
+    if (this.embedding_type) {
+      json_body.embedding_type = this.embedding_type;
+    }
+
+    if (this.normalized) {
+      json_body.normalized = this.normalized;
+    }
+
     try {
       const response = await fetch(this.api_url, {
         method: "POST",
         headers: this.headers,
-        body: JSON.stringify({
-          input: texts,
-          model: this.model_name,
-        }),
+        body: JSON.stringify(json_body),
       });
 
       const data = (await response.json()) as { data: any[]; detail: string };
@@ -73,13 +140,25 @@ export class JinaEmbeddingFunction implements IEmbeddingFunction {
     return new JinaEmbeddingFunction({
       model_name: config.model_name,
       api_key_env_var: config.api_key_env_var,
+      task: config.task,
+      late_chunking: config.late_chunking,
+      truncate: config.truncate,
+      dimensions: config.dimensions,
+      embedding_type: config.embedding_type,
+      normalized: config.normalized,
     });
   }
 
   getConfig(): StoredConfig {
     return {
       api_key_env_var: this.api_key_env_var,
       model_name: this.model_name,
+      task: this.task,
+      late_chunking: this.late_chunking,
+      truncate: this.truncate,
+      dimensions: this.dimensions,
+      embedding_type: this.embedding_type,
+      normalized: this.normalized,
     };
   }
 
diff --git a/schemas/embedding_functions/jina.json b/schemas/embedding_functions/jina.json
@@ -12,6 +12,30 @@
     "api_key_env_var": {
       "type": "string",
       "description": "Parameter api_key_env_var for the jina embedding function"
+    },
+    "task": {
+      "type": "string",
+      "description": "Parameter task for the jina embedding function"
+    },
+    "late_chunking": {
+      "type": "boolean",
+      "description": "Parameter late_chunking for the jina embedding function"
+    },
+    "truncate": {
+      "type": "boolean",
+      "description": "Parameter truncate for the jina embedding function"
+    },
+    "dimensions": {
+      "type": "integer",
+      "description": "Parameter dimensions for the jina embedding function"
+    },
+    "embedding_type": {
+      "type": "string",
+      "description": "Parameter embedding_type for the jina embedding function"
+    },
+    "normalized": {
+      "type": "boolean",
+      "description": "Parameter normalized for the jina embedding function"
     }
   },
   "required": [
