refactor: improve Ollama embedder, normalize model names, add error handling, update tests (#4403)

mem0-ts/src/oss/src/embeddings/ollama.ts

@@ -27,14 +27,18 @@ export class OllamaEmbedder implements Embedder {
     } catch (err) {
       logger.error(`Error ensuring model exists: ${err}`);
     }
-    // Ollama's Go server requires prompt to be a string. Coerce defensively
-    // since callers may pass values parsed from untrusted LLM JSON output.
-    const prompt = typeof text === "string" ? text : JSON.stringify(text);
-    const response = await this.ollama.embeddings({
+    // Coerce defensively since callers may pass values parsed from untrusted LLM JSON output.
+    const input = typeof text === "string" ? text : JSON.stringify(text);
+    const response = await this.ollama.embed({
       model: this.model,
-      prompt,
+      input,
     });
-    return response.embedding;
+    if (!response.embeddings || response.embeddings.length === 0) {
+      throw new Error(
+        `Ollama embed() returned no embeddings for model '${this.model}'`,
+      );
+    }
+    return response.embeddings[0];
   }
 
   async embedBatch(texts: string[]): Promise<number[][]> {
@@ -42,12 +46,21 @@ export class OllamaEmbedder implements Embedder {
     return response;
   }
 
+  private static normalizeModelName(name: string): string {
+    return name.includes(":") ? name : `${name}:latest`;
+  }
+
   private async ensureModelExists(): Promise<boolean> {
     if (this.initialized) {
       return true;
     }
     const local_models = await this.ollama.list();
-    if (!local_models.models.find((m: any) => m.name === this.model)) {
+    const target = OllamaEmbedder.normalizeModelName(this.model);
+    if (
+      !local_models.models.find(
+        (m: any) => OllamaEmbedder.normalizeModelName(m.name) === target,
+      )
+    ) {
       logger.info(`Pulling model ${this.model}...`);
       await this.ollama.pull({ model: this.model });
     }

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

@@ -0,0 +1,121 @@
+/// <reference types="jest" />
+/**
+ * Ollama Embedder — unit tests (mocked Ollama client).
+ */
+
+import { OllamaEmbedder } from "../src/embeddings/ollama";
+
+const mockEmbedding = [0.1, 0.2, 0.3, 0.4, 0.5];
+const mockEmbed = jest.fn().mockResolvedValue({
+  model: "nomic-embed-text:latest",
+  embeddings: [mockEmbedding],
+});
+const mockList = jest.fn().mockResolvedValue({
+  models: [{ name: "nomic-embed-text:latest" }],
+});
+const mockPull = jest.fn().mockResolvedValue({});
+
+jest.mock("ollama", () => ({
+  Ollama: jest.fn().mockImplementation(() => ({
+    embed: mockEmbed,
+    list: mockList,
+    pull: mockPull,
+  })),
+}));
+
+describe("OllamaEmbedder (unit)", () => {
+  beforeEach(() => {
+    mockEmbed.mockClear();
+    mockList.mockClear();
+    mockPull.mockClear();
+  });
+
+  it("embed() calls ollama.embed with model and input, returns first embedding", async () => {
+    const embedder = new OllamaEmbedder({
+      model: "nomic-embed-text:latest",
+    });
+
+    const result = await embedder.embed("Sample text to embed.");
+
+    expect(mockEmbed).toHaveBeenCalledTimes(1);
+    expect(mockEmbed.mock.calls[0][0]).toEqual({
+      model: "nomic-embed-text:latest",
+      input: "Sample text to embed.",
+    });
+    expect(result).toEqual(mockEmbedding);
+  });
+
+  it("embed() coerces non-string input to JSON string", async () => {
+    const embedder = new OllamaEmbedder({
+      model: "nomic-embed-text:latest",
+    });
+
+    // Force a non-string through the type boundary
+    await embedder.embed(42 as any);
+
+    expect(mockEmbed.mock.calls[0][0].input).toBe("42");
+  });
+
+  it("embedBatch() returns vectors for multiple inputs", async () => {
+    const embedder = new OllamaEmbedder({
+      model: "nomic-embed-text:latest",
+    });
+
+    const result = await embedder.embedBatch(["text1", "text2"]);
+
+    expect(mockEmbed).toHaveBeenCalledTimes(2);
+    expect(result).toEqual([mockEmbedding, mockEmbedding]);
+  });
+
+  it("ensureModelExists() does not pull when model is already present", async () => {
+    const embedder = new OllamaEmbedder({
+      model: "nomic-embed-text:latest",
+    });
+
+    await embedder.embed("trigger ensureModelExists");
+
+    expect(mockList).toHaveBeenCalled();
+    expect(mockPull).not.toHaveBeenCalled();
+  });
+
+  it("ensureModelExists() pulls model when not found locally", async () => {
+    mockList.mockResolvedValueOnce({ models: [] });
+
+    const embedder = new OllamaEmbedder({
+      model: "nomic-embed-text:latest",
+    });
+
+    await embedder.embed("trigger ensureModelExists");
+
+    expect(mockPull).toHaveBeenCalledWith({ model: "nomic-embed-text:latest" });
+  });
+
+  it("ensureModelExists() normalizes model name with :latest tag", async () => {
+    mockList.mockResolvedValue({
+      models: [{ name: "nomic-embed-text:latest" }],
+    });
+
+    const embedder = new OllamaEmbedder({
+      model: "nomic-embed-text",
+    });
+
+    await embedder.embed("trigger ensureModelExists");
+
+    expect(mockPull).not.toHaveBeenCalled();
+  });
+
+  it("embed() throws when embeddings array is empty", async () => {
+    mockEmbed.mockResolvedValueOnce({
+      model: "nomic-embed-text:latest",
+      embeddings: [],
+    });
+
+    const embedder = new OllamaEmbedder({
+      model: "nomic-embed-text:latest",
+    });
+
+    await expect(embedder.embed("text")).rejects.toThrow(
+      "Ollama embed() returned no embeddings",
+    );
+  });
+});