feat: add progress callback to embeddings (#2098)

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

.changeset/spotty-files-behave.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+add progress callback to embeddings

packages/core/src/embeddings/base.ts

@@ -17,6 +17,7 @@ export type EmbeddingInfo = {
 
 export type BaseEmbeddingOptions = {
   logProgress?: boolean;
+  progressCallback?: (current: number, total: number) => void;
 };
 
 export abstract class BaseEmbedding extends TransformComponent<
@@ -138,9 +139,11 @@ export async function batchEmbeddings<T>(
       const embeddings = await embedFunc(curBatch);
 
       resultEmbeddings.push(...embeddings);
-
+      if (options?.progressCallback) {
+        options?.progressCallback?.(i + 1, queue.length);
+      }
       if (options?.logProgress) {
-        console.log(`getting embedding progress: ${i} / ${queue.length}`);
+        console.log(`getting embedding progress: ${i + 1} / ${queue.length}`);
       }
 
       curBatch.length = 0;

packages/core/tests/embeddings.test.ts

@@ -1,4 +1,9 @@
-import { truncateMaxTokens } from "@llamaindex/core/embeddings";
+import {
+  BaseEmbedding,
+  batchEmbeddings,
+  truncateMaxTokens,
+  type BaseEmbeddingOptions,
+} from "@llamaindex/core/embeddings";
 import { Tokenizers, tokenizers } from "@llamaindex/env/tokenizers";
 import { describe, expect, test } from "vitest";
 
@@ -27,3 +32,77 @@ describe("truncateMaxTokens", () => {
     expect(t.includes("�")).toBe(false);
   });
 });
+
+describe("BaseEmbedding progressCallback", () => {
+  const mockEmbedFunc = async (text: string): Promise<number[]> => {
+    return Array.from({ length: 10 }, () => Math.random());
+  };
+  const mockBatchEmbedFunc = async (
+    texts: string[],
+  ): Promise<Array<number[]>> => {
+    return await Promise.all(texts.map(mockEmbedFunc));
+  };
+  const mockProgressCallback = (current: number, total: number) => {
+    console.log(`Progress: ${current}/${total}`);
+  };
+  const mockLogProgress = true;
+
+  const mockOptions = {
+    logProgress: mockLogProgress,
+    progressCallback: mockProgressCallback,
+  };
+
+  class MockEmbedding extends BaseEmbedding {
+    constructor(options: BaseEmbeddingOptions) {
+      super();
+      this.options = options;
+    }
+
+    private options: BaseEmbeddingOptions;
+
+    async getTextEmbedding(text: string): Promise<number[]> {
+      return await mockEmbedFunc(text);
+    }
+
+    getTextEmbeddings = async (texts: string[]): Promise<Array<number[]>> => {
+      return await mockBatchEmbedFunc(texts);
+    };
+
+    async getTextEmbeddingsBatch(
+      texts: string[],
+      options?: BaseEmbeddingOptions,
+    ): Promise<Array<number[]>> {
+      const mergedOptions = { ...this.options, ...options };
+
+      expect(mergedOptions.progressCallback).toBeDefined();
+
+      return await batchEmbeddings(
+        texts,
+        this.getTextEmbeddings,
+        this.embedBatchSize,
+        mergedOptions,
+      );
+    }
+  }
+
+  test("should call progressCallback with correct values", async () => {
+    // Import and use a real embedding class instead
+
+    const progressCalls: Array<{ current: number; total: number }> = [];
+    const progressCallback = (current: number, total: number) => {
+      progressCalls.push({ current, total });
+    };
+    const texts = ["text1", "text2", "text3"];
+    const embedding = new MockEmbedding({ progressCallback: progressCallback });
+    embedding.embedBatchSize = 1; // Set batch size to 1 for testing
+    // so that progressCallback is called for each item
+    // (otherwise, we'd only get a callback for 3/3, which is fine but less clear)
+    await embedding.getTextEmbeddingsBatch(texts);
+
+    expect(progressCalls).toEqual([
+      { current: 1, total: 3 },
+      { current: 2, total: 3 },
+      { current: 3, total: 3 },
+    ]);
+  });
+});