feat(ts-sdk): add Vertex AI embedding provider (#5882)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>

mem0-ts/src/oss/tests/config-manager.test.ts

@@ -422,6 +422,57 @@ describe("ConfigManager", () => {
       expect(cfg.vectorStore.config.port).toBe(6333);
     });
   });
+
+  describe("mergeConfig - provider-specific embedder fields", () => {
+    // The embedder config used to be rebuilt from a fixed key list, which
+    // dropped every provider-specific field before the embedder was
+    // constructed. Vertex AI then authenticated against whatever ambient
+    // project ADC resolved to and ignored the configured task types.
+    it("preserves Vertex AI fields through the merge", () => {
+      const cfg = ConfigManager.mergeConfig({
+        embedder: {
+          provider: "vertexai",
+          config: {
+            model: "gemini-embedding-001",
+            googleProjectId: "my-proj",
+            location: "europe-west4",
+            vertexCredentialsJson: "/creds.json",
+            memoryAddEmbeddingType: "SEMANTIC_SIMILARITY",
+          },
+        },
+        vectorStore: { provider: "memory", config: { collectionName: "test" } },
+        llm: { provider: "openai", config: { apiKey: "test-key" } },
+      });
+
+      expect(cfg.embedder.config).toMatchObject({
+        model: "gemini-embedding-001",
+        googleProjectId: "my-proj",
+        location: "europe-west4",
+        vertexCredentialsJson: "/creds.json",
+        memoryAddEmbeddingType: "SEMANTIC_SIMILARITY",
+      });
+    });
+
+    it("still lets normalized values win over the raw user config", () => {
+      const cfg = ConfigManager.mergeConfig({
+        embedder: {
+          provider: "lmstudio",
+          config: {
+            lmstudio_base_url: "http://localhost:1234/v1",
+            embedding_dims: 768,
+          },
+        } as never,
+        vectorStore: { provider: "memory", config: { collectionName: "test" } },
+        llm: { provider: "openai", config: { apiKey: "test-key" } },
+      });
+
+      expect(cfg.embedder.config.baseURL).toBe("http://localhost:1234/v1");
+      expect(cfg.embedder.config.embeddingDims).toBe(768);
+      // Snake_case aliases are normalized, not passed through to the provider.
+      expect(cfg.embedder.config).not.toHaveProperty("lmstudio_base_url");
+      expect(cfg.embedder.config).not.toHaveProperty("embedding_dims");
+    });
+  });
 });
 
 // ─────────────────────────────────────────────────────────────────────────
@@ -629,7 +680,10 @@ describe("Memory – LM Studio end-to-end flow", () => {
       filters: { user_id: "u1" },
     });
 
-    expect(mockEmbedder.embed).toHaveBeenCalledWith("What does the user like?");
+    expect(mockEmbedder.embed).toHaveBeenCalledWith(
+      "What does the user like?",
+      "search",
+    );
     expect(mockVStore.search).toHaveBeenCalled();
     expect(result.results).toHaveLength(1);
     expect(result.results[0].memory).toBe("User likes hiking");

mem0-ts/src/oss/tests/factory.unit.test.ts

@@ -40,6 +40,11 @@ jest.mock("../src/embeddings/lmstudio", () => ({
     .fn()
     .mockImplementation((config) => ({ type: "lmstudio-embedder", config })),
 }));
+jest.mock("../src/embeddings/vertexai", () => ({
+  VertexAIEmbedder: jest
+    .fn()
+    .mockImplementation((config) => ({ type: "vertexai-embedder", config })),
+}));
 jest.mock("../src/embeddings/together", () => ({
   TogetherEmbedder: jest
     .fn()
@@ -241,6 +246,7 @@ describe("EmbedderFactory", () => {
     ["fastembed"],
     ["langchain"],
     ["lmstudio"],
+    ["vertexai"],
     ["together"],
   ])("creates embedder for provider '%s'", (provider) => {
     expect(() =>

mem0-ts/src/oss/tests/memory.embedding-action.test.ts

@@ -0,0 +1,115 @@
+/// <reference types="jest" />
+/**
+ * Verifies the memory pipeline threads the correct memory action
+ * ("add" | "update" | "search") into the embedder. Task-type-aware providers
+ * (e.g. Vertex AI) embed queries and documents differently based on this, and
+ * the argument is silently ignored by every other embedder, so only a
+ * pipeline-level test catches a dropped action.
+ */
+import { Memory } from "../src/memory";
+
+const mockEmbedding = new Array(1536).fill(0.1);
+// Prefixed `mock*` so jest's hoisted module factory may reference them.
+const mockEmbed = jest.fn().mockResolvedValue(mockEmbedding);
+const mockEmbedBatch = jest
+  .fn()
+  .mockImplementation((texts: string[]) =>
+    Promise.resolve(texts.map(() => mockEmbedding)),
+  );
+
+const mockGenerateResponse = jest
+  .fn()
+  .mockResolvedValue(JSON.stringify({ memory: [] }));
+
+jest.mock("../src/embeddings/google", () => ({ GoogleEmbedder: jest.fn() }));
+jest.mock("../src/llms/google", () => ({ GoogleLLM: jest.fn() }));
+jest.mock("../src/llms/openai", () => ({
+  OpenAILLM: jest.fn().mockImplementation(() => ({
+    generateResponse: mockGenerateResponse,
+  })),
+}));
+jest.mock("../src/embeddings/openai", () => ({
+  OpenAIEmbedder: jest.fn().mockImplementation(() => ({
+    embed: mockEmbed,
+    embedBatch: mockEmbedBatch,
+    embeddingDims: 1536,
+  })),
+}));
+
+function createMemory(): Memory {
+  return new Memory({
+    version: "v1.1",
+    embedder: {
+      provider: "openai",
+      config: { apiKey: "test-key", model: "text-embedding-3-small" },
+    },
+    vectorStore: {
+      provider: "memory",
+      config: {
+        collectionName: `test-action-${Date.now()}-${Math.random()}`,
+        dimension: 1536,
+        dbPath: ":memory:",
+      },
+    },
+    llm: {
+      provider: "openai",
+      config: { apiKey: "test-key", model: "gpt-5-mini" },
+    },
+    historyDbPath: ":memory:",
+  });
+}
+
+describe("embedder memory-action threading", () => {
+  let memory: Memory;
+
+  beforeEach(() => {
+    memory = createMemory();
+    mockEmbed.mockClear();
+    mockEmbedBatch.mockClear();
+    mockGenerateResponse.mockResolvedValue(JSON.stringify({ memory: [] }));
+  });
+
+  afterEach(async () => {
+    await memory.reset();
+  });
+
+  test("search() embeds the query with the 'search' action", async () => {
+    await memory.search("what do I like", { filters: { user_id: "u1" } });
+    expect(mockEmbed).toHaveBeenCalledWith("what do I like", "search");
+  });
+
+  test("update() embeds the new value with the 'update' action", async () => {
+    // Missing id: update embeds the value before it throws on the absent row.
+    await memory.update("missing-id", "new value").catch(() => {});
+    expect(mockEmbed).toHaveBeenCalledWith("new value", "update");
+  });
+
+  test("add() batch-embeds extracted memories and entities with the 'add' action", async () => {
+    mockGenerateResponse.mockResolvedValue(
+      JSON.stringify({
+        memory: [
+          { id: "1", text: "John loves sci-fi movies", attributed_to: "user" },
+        ],
+      }),
+    );
+
+    await memory.add("I love sci-fi movies", { userId: "u1" });
+
+    // Phase 1 retrieval embeds the incoming turn as a query.
+    expect(mockEmbed).toHaveBeenCalledWith(
+      expect.stringContaining("I love sci-fi movies"),
+      "search",
+    );
+    // Phase 3 (extracted memories) and phase 7 (linked entities) both batch
+    // embed as documents. Without an explicit action, a task-type-aware
+    // embedder falls back to its own default and silently mis-embeds.
+    expect(mockEmbedBatch).toHaveBeenCalledWith(
+      ["John loves sci-fi movies"],
+      "add",
+    );
+    expect(mockEmbedBatch.mock.calls.length).toBeGreaterThan(0);
+    for (const call of mockEmbedBatch.mock.calls) {
+      expect(call[1]).toBe("add");
+    }
+  });
+});

mem0-ts/src/oss/tests/tsup-externals.test.ts

@@ -47,4 +47,11 @@ describe("tsup.config.ts externals", () => {
   it("should have peerDependencies defined in package.json", () => {
     expect(peerDeps.length).toBeGreaterThan(0);
   });
+
+  it("should not list any dependency twice", () => {
+    const duplicates = externalDeps.filter(
+      (dep, i) => externalDeps.indexOf(dep) !== i,
+    );
+    expect(duplicates).toEqual([]);
+  });
 });

mem0-ts/src/oss/tests/vertexai-embedder.test.ts

@@ -0,0 +1,238 @@
+/// <reference types="jest" />
+
+const mockPredict = jest.fn();
+const mockGetProjectId = jest.fn();
+const mockClientConstructor = jest.fn();
+
+jest.mock("@google-cloud/aiplatform", () => {
+  return {
+    __esModule: true,
+    PredictionServiceClient: jest.fn().mockImplementation((...args) => {
+      mockClientConstructor(...args);
+      return {
+        predict: mockPredict,
+        getProjectId: mockGetProjectId,
+      };
+    }),
+    helpers: {
+      toValue: jest.fn().mockImplementation((val) => val),
+      fromValue: jest.fn().mockImplementation((val) => val),
+    },
+  };
+});
+
+import { VertexAIEmbedder } from "../src/embeddings/vertexai";
+
+const mockEmbedding = [0.1, 0.2, 0.3, 0.4];
+
+/** Echoes one embedding back per instance in the request. */
+function predictEchoingInstances() {
+  return (req: { instances: unknown[] }) =>
+    Promise.resolve([
+      {
+        predictions: req.instances.map(() => ({
+          embeddings: { values: mockEmbedding },
+        })),
+      },
+    ]);
+}
+
+describe("VertexAIEmbedder", () => {
+  beforeEach(() => {
+    mockPredict.mockReset();
+    mockClientConstructor.mockReset();
+    mockGetProjectId.mockReset();
+    mockGetProjectId.mockResolvedValue("adc-project");
+    mockPredict.mockResolvedValue([
+      {
+        predictions: [
+          {
+            embeddings: {
+              values: mockEmbedding,
+            },
+          },
+        ],
+      },
+    ]);
+  });
+
+  describe("basic functionality", () => {
+    it("embed() returns the embedding vector", async () => {
+      const embedder = new VertexAIEmbedder({
+        googleProjectId: "test-project",
+      });
+
+      const result = await embedder.embed("hello");
+      expect(result).toEqual(mockEmbedding);
+      expect(mockPredict).toHaveBeenCalledTimes(1);
+
+      const callArgs = mockPredict.mock.calls[0][0];
+      expect(callArgs.endpoint).toBe(
+        "projects/test-project/locations/us-central1/publishers/google/models/gemini-embedding-001",
+      );
+      // task_type belongs on the instance (snake_case), parameters carries
+      // only outputDimensionality.
+      expect(callArgs.instances).toEqual([
+        { content: "hello", task_type: "SEMANTIC_SIMILARITY" },
+      ]);
+      expect(callArgs.parameters).toEqual({ outputDimensionality: 256 });
+    });
+
+    it("embed() with memory action search uses RETRIEVAL_QUERY", async () => {
+      const embedder = new VertexAIEmbedder({
+        googleProjectId: "test-project",
+      });
+
+      await embedder.embed("hello", "search");
+      const callArgs = mockPredict.mock.calls[0][0];
+      expect(callArgs.instances[0].task_type).toBe("RETRIEVAL_QUERY");
+    });
+
+    it("embed() with memory action add uses RETRIEVAL_DOCUMENT", async () => {
+      const embedder = new VertexAIEmbedder({
+        googleProjectId: "test-project",
+      });
+
+      await embedder.embed("hello", "add");
+      const callArgs = mockPredict.mock.calls[0][0];
+      expect(callArgs.instances[0].task_type).toBe("RETRIEVAL_DOCUMENT");
+    });
+
+    it("throws error when predictions are empty", async () => {
+      mockPredict.mockResolvedValue([{}]);
+      const embedder = new VertexAIEmbedder({
+        googleProjectId: "test-project",
+      });
+
+      await expect(embedder.embed("hello")).rejects.toThrow(
+        "No predictions returned from Vertex AI",
+      );
+    });
+  });
+
+  describe("embedBatch() request sizing", () => {
+    // gemini-embedding-001 (the default model) rejects any predict() call
+    // carrying more than one input text, so the batch loop must degrade to one
+    // request per text rather than the 250-instance chunk the older models take.
+    it("sends one instance per request for gemini-embedding models", async () => {
+      mockPredict.mockImplementation(predictEchoingInstances());
+      const embedder = new VertexAIEmbedder({
+        googleProjectId: "test-project",
+      });
+
+      const texts = ["a", "b", "c"];
+      const result = await embedder.embedBatch(texts);
+
+      expect(result).toEqual(texts.map(() => mockEmbedding));
+      expect(mockPredict).toHaveBeenCalledTimes(3);
+      for (const call of mockPredict.mock.calls) {
+        expect(call[0].instances.length).toBe(1);
+        // batch defaults to the "add" action -> RETRIEVAL_DOCUMENT
+        expect(call[0].instances[0].task_type).toBe("RETRIEVAL_DOCUMENT");
+        expect(call[0].parameters).toEqual({ outputDimensionality: 256 });
+      }
+      expect(
+        mockPredict.mock.calls.map((c) => c[0].instances[0].content),
+      ).toEqual(texts);
+    });
+
+    it("chunks at 250 instances for text-embedding models", async () => {
+      mockPredict.mockImplementation(predictEchoingInstances());
+      const embedder = new VertexAIEmbedder({
+        googleProjectId: "test-project",
+        model: "text-embedding-005",
+      });
+
+      const texts = Array.from({ length: 255 }, (_, i) => `text-${i}`);
+      const result = await embedder.embedBatch(texts, "search");
+
+      expect(result.length).toBe(255);
+      expect(mockPredict).toHaveBeenCalledTimes(2);
+      expect(mockPredict.mock.calls[0][0].instances.length).toBe(250);
+      expect(mockPredict.mock.calls[0][0].instances[0].task_type).toBe(
+        "RETRIEVAL_QUERY",
+      );
+      expect(mockPredict.mock.calls[1][0].instances.length).toBe(5);
+    });
+
+    it("rejects an unknown memory action", async () => {
+      const embedder = new VertexAIEmbedder({
+        googleProjectId: "test-project",
+      });
+
+      await expect(
+        embedder.embedBatch(["a"], "delete" as unknown as "add"),
+      ).rejects.toThrow("Invalid memory action: delete");
+    });
+  });
+
+  describe("client initialization", () => {
+    const PROJECT_ENV_VARS = [
+      "GCP_PROJECT_ID",
+      "GOOGLE_CLOUD_PROJECT",
+      "GCLOUD_PROJECT",
+    ];
+    let savedEnv: Record<string, string | undefined>;
+
+    beforeEach(() => {
+      savedEnv = {};
+      for (const key of PROJECT_ENV_VARS) {
+        savedEnv[key] = process.env[key];
+        delete process.env[key];
+      }
+    });
+
+    afterEach(() => {
+      for (const key of PROJECT_ENV_VARS) {
+        if (savedEnv[key] === undefined) delete process.env[key];
+        else process.env[key] = savedEnv[key];
+      }
+    });
+
+    it("resolves the project ID from credentials when none is configured", async () => {
+      const embedder = new VertexAIEmbedder({});
+
+      await embedder.embed("hello");
+
+      expect(mockGetProjectId).toHaveBeenCalledTimes(1);
+      expect(mockPredict.mock.calls[0][0].endpoint).toBe(
+        "projects/adc-project/locations/us-central1/publishers/google/models/gemini-embedding-001",
+      );
+    });
+
+    it("surfaces a helpful error when no project ID can be resolved", async () => {
+      mockGetProjectId.mockRejectedValue(
+        new Error("Unable to detect a Project Id"),
+      );
+      const embedder = new VertexAIEmbedder({});
+
+      await expect(embedder.embed("hello")).rejects.toThrow(
+        "Vertex AI could not determine a Google Cloud project ID",
+      );
+    });
+
+    it("builds one client for concurrent calls", async () => {
+      const embedder = new VertexAIEmbedder({
+        googleProjectId: "test-project",
+      });
+
+      await Promise.all([
+        embedder.embed("a"),
+        embedder.embed("b"),
+        embedder.embed("c"),
+      ]);
+
+      expect(mockClientConstructor).toHaveBeenCalledTimes(1);
+    });
+
+    it("retries initialization after a failure", async () => {
+      mockGetProjectId.mockRejectedValueOnce(new Error("transient"));
+      const embedder = new VertexAIEmbedder({});
+
+      await expect(embedder.embed("hello")).rejects.toThrow(
+        "Vertex AI could not determine a Google Cloud project ID",
+      );
+      await expect(embedder.embed("hello")).resolves.toEqual(mockEmbedding);
+    });
+  });
+});

mem0-ts/src/oss/tests/vertexai-protobuf.test.ts

@@ -0,0 +1,139 @@
+/// <reference types="jest" />
+
+/**
+ * The sibling `vertexai-embedder.test.ts` stubs `helpers.toValue`/`fromValue`
+ * as identity functions, so it never exercises the real protobuf `Value`
+ * encode/decode path. Here we mock only `PredictionServiceClient` (no
+ * network, no credentials) and keep the REAL `helpers`, so a malformed
+ * instance shape or a broken decode actually fails.
+ */
+
+const mockPredict = jest.fn();
+const mockGetProjectId = jest.fn();
+
+jest.mock("@google-cloud/aiplatform", () => {
+  const actual = jest.requireActual("@google-cloud/aiplatform");
+  return {
+    ...actual,
+    __esModule: true,
+    PredictionServiceClient: jest.fn().mockImplementation(() => ({
+      predict: mockPredict,
+      getProjectId: mockGetProjectId,
+    })),
+  };
+});
+
+import { VertexAIEmbedder } from "../src/embeddings/vertexai";
+import { helpers } from "@google-cloud/aiplatform";
+
+/** A `google.protobuf.Value` always carries `kind`+`structValue`/etc; a plain
+ *  JS object never does. This is what "real encoding happened" looks like. */
+function expectEncodedValue(value: unknown) {
+  expect(value).toEqual(expect.objectContaining({ kind: expect.any(String) }));
+}
+
+describe("VertexAIEmbedder protobuf boundary (real helpers.toValue/fromValue)", () => {
+  beforeEach(() => {
+    mockPredict.mockReset();
+    mockGetProjectId.mockReset();
+    mockGetProjectId.mockResolvedValue("adc-project");
+  });
+
+  it("encodes the instance as a genuine protobuf Value carrying {content, task_type}", async () => {
+    mockPredict.mockResolvedValue([
+      {
+        predictions: [helpers.toValue({ embeddings: { values: [0.1, 0.2] } })],
+      },
+    ]);
+    const embedder = new VertexAIEmbedder({ googleProjectId: "test-project" });
+
+    await embedder.embed("hello world", "search");
+
+    const { instances } = mockPredict.mock.calls[0][0];
+    expect(instances).toHaveLength(1);
+    expectEncodedValue(instances[0]);
+
+    // Decode with the REAL fromValue -- proves the encoded instance is
+    // readable and matches exactly what Vertex expects on the wire.
+    expect(helpers.fromValue(instances[0])).toEqual({
+      content: "hello world",
+      task_type: "RETRIEVAL_QUERY",
+    });
+  });
+
+  it("encodes parameters as a genuine protobuf Value carrying {outputDimensionality}", async () => {
+    mockPredict.mockResolvedValue([
+      {
+        predictions: [helpers.toValue({ embeddings: { values: [0.1] } })],
+      },
+    ]);
+    const embedder = new VertexAIEmbedder({
+      googleProjectId: "test-project",
+      embeddingDims: 768,
+    });
+
+    await embedder.embed("hello");
+
+    const { parameters } = mockPredict.mock.calls[0][0];
+    expectEncodedValue(parameters);
+    expect(helpers.fromValue(parameters)).toEqual({
+      outputDimensionality: 768,
+    });
+  });
+
+  it("decodes a real toValue()-encoded prediction back into the embedding array", async () => {
+    const values = [0.11, -0.22, 0.33, 0.0];
+    mockPredict.mockResolvedValue([
+      {
+        predictions: [helpers.toValue({ embeddings: { values } })],
+      },
+    ]);
+    const embedder = new VertexAIEmbedder({ googleProjectId: "test-project" });
+
+    const result = await embedder.embed("hello");
+
+    expect(result).toEqual(values);
+  });
+
+  it("rejects a prediction that isn't a real encoded protobuf Value", async () => {
+    // A raw JS object (what the identity-stubbed sibling test effectively
+    // assumed `predict()` returns) is not a valid protobuf Value -- the real
+    // fromValue() throws on it instead of silently passing it through.
+    mockPredict.mockResolvedValue([
+      { predictions: [{ embeddings: { values: [0.1, 0.2] } }] },
+    ]);
+    const embedder = new VertexAIEmbedder({ googleProjectId: "test-project" });
+
+    await expect(embedder.embed("hello")).rejects.toThrow();
+  });
+
+  it("embedBatch() round-trips multiple real encoded predictions", async () => {
+    const vectors = [
+      [0.1, 0.2],
+      [0.3, 0.4],
+    ];
+    mockPredict.mockImplementation(
+      (req: { instances: unknown[] }) =>
+        Promise.resolve([
+          {
+            predictions: req.instances.map((_, i) =>
+              helpers.toValue({ embeddings: { values: vectors[i] } }),
+            ),
+          },
+        ]) as any,
+    );
+    const embedder = new VertexAIEmbedder({
+      googleProjectId: "test-project",
+      model: "text-embedding-005",
+    });
+
+    const result = await embedder.embedBatch(["a", "b"]);
+
+    expect(result).toEqual(vectors);
+    const { instances } = mockPredict.mock.calls[0][0];
+    expect(instances.map((i: unknown) => helpers.fromValue(i as any))).toEqual([
+      { content: "a", task_type: "RETRIEVAL_DOCUMENT" },
+      { content: "b", task_type: "RETRIEVAL_DOCUMENT" },
+    ]);
+  });
+});