feat: implement Jina embedding through Jina api (#995)

Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.changeset/lovely-flies-press.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Implement Jina embedding through Jina api

examples/multimodal/jina.ts

@@ -0,0 +1,83 @@
+import {
+  ImageDocument,
+  JinaAIEmbedding,
+  similarity,
+  SimilarityType,
+  SimpleDirectoryReader,
+} from "llamaindex";
+import path from "path";
+
+async function main() {
+  const jina = new JinaAIEmbedding({
+    model: "jina-clip-v1",
+  });
+
+  // Get text embeddings
+  const text1 = "a car";
+  const textEmbedding1 = await jina.getTextEmbedding(text1);
+  const text2 = "a football match";
+  const textEmbedding2 = await jina.getTextEmbedding(text2);
+
+  // Get image embedding
+  const image =
+    "https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/football-match.jpg";
+  const imageEmbedding = await jina.getImageEmbedding(image);
+
+  // Calc similarity between text and image
+  const sim1 = similarity(
+    textEmbedding1,
+    imageEmbedding,
+    SimilarityType.DEFAULT,
+  );
+  const sim2 = similarity(
+    textEmbedding2,
+    imageEmbedding,
+    SimilarityType.DEFAULT,
+  );
+
+  console.log(`Similarity between "${text1}" and the image is ${sim1}`);
+  console.log(`Similarity between "${text2}" and the image is ${sim2}`);
+
+  // Get multiple text embeddings
+  const textEmbeddings = await jina.getTextEmbeddings([text1, text2]);
+  const sim3 = similarity(
+    textEmbeddings[0],
+    textEmbeddings[1],
+    SimilarityType.DEFAULT,
+  );
+  console.log(
+    `Similarity between the two texts "${text1}" and "${text2}" is ${sim3}`,
+  );
+
+  // Get multiple image embeddings
+  const catImg1 =
+    "https://i.pinimg.com/600x315/21/48/7e/21487e8e0970dd366dafaed6ab25d8d8.jpg";
+  const catImg2 =
+    "https://i.pinimg.com/736x/c9/f2/3e/c9f23e212529f13f19bad5602d84b78b.jpg";
+  const imageEmbeddings = await jina.getImageEmbeddings([catImg1, catImg2]);
+  const sim4 = similarity(
+    imageEmbeddings[0],
+    imageEmbeddings[1],
+    SimilarityType.DEFAULT,
+  );
+  console.log(`Similarity between the two online cat images is ${sim4}`);
+
+  // Get image embeddings from multiple local files
+  const documents = await new SimpleDirectoryReader().loadData({
+    directoryPath: path.join("multimodal", "data"),
+  });
+  const localImages = documents
+    .filter((doc) => doc instanceof ImageDocument)
+    .slice(0, 2); // Get only the first two images
+  const localImageEmbeddings = await jina.getImageEmbeddings(
+    localImages.map((doc) => (doc as ImageDocument).image),
+  );
+  const sim5 = similarity(
+    localImageEmbeddings[0],
+    localImageEmbeddings[1],
+    SimilarityType.DEFAULT,
+  );
+  console.log(`Similarity between the two local images is ${sim5}`);
+}
+
+void main();

packages/llamaindex/src/embeddings/JinaAIEmbedding.ts

@@ -1,29 +1,127 @@
 import { getEnv } from "@llamaindex/env";
-import { OpenAIEmbedding } from "./OpenAIEmbedding.js";
+import { imageToDataUrl } from "../internal/utils.js";
+import type { ImageType } from "../Node.js";
+import { MultiModalEmbedding } from "./MultiModalEmbedding.js";
 
-export class JinaAIEmbedding extends OpenAIEmbedding {
-  constructor(init?: Partial<OpenAIEmbedding>) {
-    const {
-      apiKey = getEnv("JINAAI_API_KEY"),
-      additionalSessionOptions = {},
-      model = "jina-embeddings-v2-base-en",
-      ...rest
-    } = init ?? {};
+function isLocal(url: ImageType): boolean {
+  if (url instanceof Blob) return true;
+  return new URL(url).protocol === "file:";
+}
 
+export type JinaEmbeddingRequest = {
+  input: Array<{ text: string } | { url: string } | { bytes: string }>;
+  model?: string;
+  encoding_type?: "float" | "binary" | "ubinary";
+};
+
+export type JinaEmbeddingResponse = {
+  model: string;
+  object: string;
+  usage: {
+    total_tokens: number;
+    prompt_tokens: number;
+  };
+  data: Array<{
+    object: string;
+    index: number;
+    embedding: number[];
+  }>;
+};
+
+const JINA_MULTIMODAL_MODELS = ["jina-clip-v1"];
+
+export class JinaAIEmbedding extends MultiModalEmbedding {
+  apiKey: string;
+  model: string;
+  baseURL: string;
+
+  async getTextEmbedding(text: string): Promise<number[]> {
+    const result = await this.getJinaEmbedding({ input: [{ text }] });
+    return result.data[0].embedding;
+  }
+
+  async getImageEmbedding(image: ImageType): Promise<number[]> {
+    const img = await this.getImageInput(image);
+    const result = await this.getJinaEmbedding({ input: [img] });
+    return result.data[0].embedding;
+  }
+
+  // Retrieve multiple text embeddings in a single request
+  getTextEmbeddings = async (texts: string[]): Promise<Array<number[]>> => {
+    const input = texts.map((text) => ({ text }));
+    const result = await this.getJinaEmbedding({ input });
+    return result.data.map((d) => d.embedding);
+  };
+
+  // Retrieve multiple image embeddings in a single request
+  async getImageEmbeddings(images: ImageType[]): Promise<number[][]> {
+    const input = await Promise.all(
+      images.map((img) => this.getImageInput(img)),
+    );
+    const result = await this.getJinaEmbedding({ input });
+    return result.data.map((d) => d.embedding);
+  }
+
+  constructor(init?: Partial<JinaAIEmbedding>) {
+    super();
+    const apiKey = init?.apiKey ?? getEnv("JINAAI_API_KEY");
     if (!apiKey) {
       throw new Error(
         "Set Jina AI API Key in JINAAI_API_KEY env variable. Get one for free or top up your key at https://jina.ai/embeddings",
       );
     }
+    this.apiKey = apiKey;
+    this.model = init?.model ?? "jina-embeddings-v2-base-en";
+    this.baseURL = init?.baseURL ?? "https://api.jina.ai/v1/embeddings";
+    init?.embedBatchSize && (this.embedBatchSize = init?.embedBatchSize);
+  }
 
-    additionalSessionOptions.baseURL =
-      additionalSessionOptions.baseURL ?? "https://api.jina.ai/v1";
+  private async getImageInput(
+    image: ImageType,
+  ): Promise<{ bytes: string } | { url: string }> {
+    if (isLocal(image)) {
+      const base64 = await imageToDataUrl(image);
+      const bytes = base64.split(",")[1];
+      return { bytes };
+    } else {
+      return { url: image.toString() };
+    }
+  }
 
-    super({
-      apiKey,
-      additionalSessionOptions,
-      model,
-      ...rest,
+  private async getJinaEmbedding(
+    input: JinaEmbeddingRequest,
+  ): Promise<JinaEmbeddingResponse> {
+    // if input includes image, check if model supports multimodal embeddings
+    if (
+      input.input.some((i) => "url" in i || "bytes" in i) &&
+      !JINA_MULTIMODAL_MODELS.includes(this.model)
+    ) {
+      throw new Error(
+        `Model ${this.model} does not support image embeddings. Use ${JINA_MULTIMODAL_MODELS.join(", ")}`,
+      );
+    }
+
+    const response = await fetch(this.baseURL, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: `Bearer ${this.apiKey}`,
+      },
+      body: JSON.stringify({
+        model: this.model,
+        encoding_type: "float",
+        ...input,
+      }),
     });
+    if (!response.ok) {
+      throw new Error(
+        `Request ${this.baseURL} failed with status ${response.status}`,
+      );
+    }
+    const result: JinaEmbeddingResponse = await response.json();
+    return {
+      ...result,
+      data: result.data.sort((a, b) => a.index - b.index), // Sort resulting embeddings by index
+    };
   }
 }

packages/llamaindex/src/index.ts

@@ -15,4 +15,5 @@ export { type VertexGeminiSessionOptions } from "./llm/gemini/types.js";
 export { GeminiVertexSession } from "./llm/gemini/vertex.js";
 
 // Expose AzureDynamicSessionTool for node.js runtime only
+export { JinaAIEmbedding } from "./embeddings/JinaAIEmbedding.js";
 export { AzureDynamicSessionTool } from "./tools/AzureDynamicSessionTool.node.js";