[ENH]: Switching Ollama JS EF to ollama js client

clients/js/src/embeddings/OllamaEmbeddingFunction.ts

@@ -1,35 +1,77 @@
 import { IEmbeddingFunction } from "./IEmbeddingFunction";
-
+const DEFAULT_MODEL = "chroma/all-minilm-l6-v2-f32";
+const DEFAULT_LOCAL_URL = "http://localhost:11434";
 export class OllamaEmbeddingFunction implements IEmbeddingFunction {
-  private readonly url: string;
+  private readonly url?: string | undefined;
   private readonly model: string;
+  private ollamaClient: any;
 
-  constructor({ url, model }: { url: string; model: string }) {
+  constructor(
+    {
+      url = DEFAULT_LOCAL_URL,
+      model = DEFAULT_MODEL,
+    }: { url?: string; model?: string } = {
+      url: DEFAULT_LOCAL_URL,
+      model: DEFAULT_MODEL,
+    },
+  ) {
     // we used to construct the client here, but we need to async import the types
     // for the openai npm package, and the constructor can not be async
-    this.url = url;
-    this.model = model;
+    if (url && url.endsWith("/api/embeddings")) {
+      this.url = url.slice(0, -"/api/embeddings".length);
+    } else {
+      this.url = url || DEFAULT_LOCAL_URL;
+    }
+    this.model = model || DEFAULT_MODEL;
   }
 
-  public async generate(texts: string[]) {
-    let embeddings: number[][] = [];
-    for (let text of texts) {
-      const response = await fetch(this.url, {
-        method: "POST",
-        headers: {
-          "Content-Type": "application/json",
-        },
-        body: JSON.stringify({ model: this.model, prompt: text }),
+  private async initClient() {
+    if (this.ollamaClient) return;
+    try {
+      // @ts-ignore
+      this.ollamaClient = await import("ollama/browser").then((ollama) => {
+        // @ts-ignore
+        return new ollama.Ollama({ host: this.url });
       });
-
-      if (!response.ok) {
+    } catch (e) {
+      // @ts-ignore
+      if (e.code === "MODULE_NOT_FOUND") {
         throw new Error(
-          `Failed to generate embeddings: ${response.status} (${response.statusText})`,
+          "Please install the ollama package to use the CohereEmbeddingFunction, `npm install -S ollama`",
         );
       }
-      let finalResponse = await response.json();
-      embeddings.push(finalResponse["embedding"]);
+      throw e;
     }
-    return embeddings;
+  }
+
+  public async generate(texts: string[]) {
+    await this.initClient();
+    return await this.ollamaClient
+      .embed({
+        model: this.model,
+        input: texts,
+      })
+      .then((response: any) => {
+        return response.embeddings;
+      });
+    // let embeddings: number[][] = [];
+    // for (let text of texts) {
+    //   const response = await fetch(this.url, {
+    //     method: "POST",
+    //     headers: {
+    //       "Content-Type": "application/json",
+    //     },
+    //     body: JSON.stringify({ model: this.model, prompt: text }),
+    //   });
+    //
+    //   if (!response.ok) {
+    //     throw new Error(
+    //       `Failed to generate embeddings: ${response.status} (${response.statusText})`,
+    //     );
+    //   }
+    //   let finalResponse = await response.json();
+    //   embeddings.push(finalResponse["embedding"]);
+    // }
+    // return embeddings;
   }
 }

clients/js/test/add.collections.test.ts

@@ -214,30 +214,4 @@ describe("add collections", () => {
       expect(e.message).toMatch("got empty embedding at pos");
     }
   });
-
-  if (!process.env.OLLAMA_SERVER_URL) {
-    test.skip("it should use ollama EF, OLLAMA_SERVER_URL not defined", async () => {});
-  } else {
-    test("it should use ollama EF", async () => {
-      const embedder = new OllamaEmbeddingFunction({
-        url:
-          process.env.OLLAMA_SERVER_URL ||
-          "http://127.0.0.1:11434/api/embeddings",
-        model: "nomic-embed-text",
-      });
-      const collection = await client.createCollection({
-        name: "test",
-        embeddingFunction: embedder,
-      });
-      const embeddings = await embedder.generate(DOCUMENTS);
-      await collection.add({ ids: IDS, embeddings: embeddings });
-      const count = await collection.count();
-      expect(count).toBe(3);
-      var res = await collection.get({
-        ids: IDS,
-        include: [IncludeEnum.Embeddings],
-      });
-      expect(res.embeddings).toEqual(embeddings); // reverse because of the order of the ids
-    });
-  }
 });

clients/js/test/embeddings/ollama.test.ts

@@ -0,0 +1,55 @@
+import { OllamaEmbeddingFunction } from "../../src/index";
+import { describe, expect, test } from "@jest/globals";
+import { DOCUMENTS } from "../data";
+
+describe("ollama embedding function", () => {
+  if ((globalThis as any).ollamaAvailable) {
+    test("it should embed with defaults", async () => {
+      const embedder = new OllamaEmbeddingFunction({
+        url: process.env.OLLAMA_URL,
+      });
+      const embeddings = await embedder.generate(DOCUMENTS);
+      expect(embeddings).toBeDefined();
+      expect(embeddings.length).toBe(DOCUMENTS.length);
+      expect(embeddings[0]).toBeDefined();
+      expect(embeddings[0].length).toBe(384);
+    });
+    test("it should embed with model", async () => {
+      const embedder = new OllamaEmbeddingFunction({
+        url: process.env.OLLAMA_URL,
+        model: "nomic-embed-text",
+      });
+      const embeddings = await embedder.generate(DOCUMENTS);
+      expect(embeddings).toBeDefined();
+      expect(embeddings.length).toBe(DOCUMENTS.length);
+      expect(embeddings[0]).toBeDefined();
+      expect(embeddings[0].length).toBe(768);
+    });
+
+    test("it should fail with unknown model", async () => {
+      const model_name = "not-a-real-model" + Math.floor(Math.random() * 1000);
+      const embedder = new OllamaEmbeddingFunction({
+        url: process.env.OLLAMA_URL,
+        model: model_name,
+      });
+      try {
+        await embedder.generate(DOCUMENTS);
+      } catch (e: any) {
+        expect(e.message).toContain(`model \"${model_name}\" not found`);
+      }
+    });
+
+    test("it should fail wrong host", async () => {
+      const embedder = new OllamaEmbeddingFunction({
+        url: "https://example.com:1234",
+      });
+      try {
+        await embedder.generate(DOCUMENTS);
+      } catch (e: any) {
+        expect(e.message).toContain(`fetch failed`);
+      }
+    });
+  } else {
+    test.skip("ollama not installed", async () => {});
+  }
+});

clients/js/test/embeddings/startOllamaContainer.ts

@@ -0,0 +1,62 @@
+import { GenericContainer, Wait } from "testcontainers";
+const OLLAMA_PORT = 11434;
+
+// checking model name to prevent command injection
+function isValidModel(model: string): boolean {
+  const regex =
+    /^[a-zA-Z0-9\-]+((:[a-zA-Z0-9\-]+)?|(\/[a-zA-Z0-9\-]+)?(:[a-zA-Z0-9\-]+)?)$/; // matches: <model> or <model>:<version> or <user>/<model> or <user>/<model>:<version>
+  return regex.test(model);
+}
+
+export async function startOllamaContainer(
+  {
+    model = "chroma/all-minilm-l6-v2-f32",
+  }: {
+    model?: string;
+  } = { model: "chroma/all-minilm-l6-v2-f32" },
+) {
+  let container: GenericContainer;
+  if (process.env.PREBUILT_CHROMADB_IMAGE) {
+    container = new GenericContainer(process.env.PREBUILT_CHROMADB_IMAGE);
+  } else {
+    container = new GenericContainer("ollama/ollama:latest");
+  }
+
+  const env: Record<string, string> = {};
+
+  const startedContainer = await container
+    // uncomment to see container logs
+    // .withLogConsumer((stream) => {
+    //   stream.on("data", (line) => console.log(line));
+    //   stream.on("err", (line) => console.error(line));
+    //   stream.on("end", () => console.log("Stream closed"));
+    // })
+    .withExposedPorts(OLLAMA_PORT)
+    .withWaitStrategy(Wait.forListeningPorts())
+    .withStartupTimeout(120_000)
+    .withEnvironment(env)
+    .start();
+
+  const ollamaUrl = `http://${startedContainer.getHost()}:${startedContainer.getMappedPort(
+    OLLAMA_PORT,
+  )}`;
+  if (!model) {
+    throw new Error("Model name is required");
+  }
+  if (model && !isValidModel(model)) {
+    throw new Error("Invalid model name");
+  }
+  await startedContainer.exec(["ollama", "pull", model]);
+  await startedContainer.exec(["ollama", "pull", "nomic-embed-text"]);
+  await startedContainer.exec([
+    "ollama",
+    "pull",
+    "chroma/all-minilm-l6-v2-f32",
+  ]);
+  return {
+    ollamaUrl: ollamaUrl,
+    host: startedContainer.getHost(),
+    port: startedContainer.getMappedPort(OLLAMA_PORT),
+    ollamaContainer: startedContainer,
+  };
+}

clients/js/test/testEnvSetup.ts

@@ -1,7 +1,24 @@
 import { startChromaContainer } from "./startChromaContainer";
-
+import { execSync } from "child_process";
+import { startOllamaContainer } from "./embeddings/startOllamaContainer";
 export default async function testSetup() {
   const { container, url } = await startChromaContainer();
   process.env.DEFAULT_CHROMA_INSTANCE_URL = url;
   (globalThis as any).chromaContainer = container;
+  (globalThis as any).ollamaAvailable = false;
+  try {
+    execSync("npm ls ollama", { stdio: "ignore" });
+    const { ollamaUrl, ollamaContainer } = await startOllamaContainer();
+    process.env.OLLAMA_URL = ollamaUrl;
+    (globalThis as any).chromaContainer = ollamaContainer;
+    (globalThis as any).ollamaAvailable = true;
+    console.log(
+      "ollama is installed and Ollama container is running. Running tests...",
+    );
+  } catch (error) {
+    console.log(
+      "Ollama package not installed or failed to start ollama. Skipping tests: " +
+        error,
+    );
+  }
 }

docs/docs.trychroma.com/markdoc/content/integrations/embedding-models/ollama.md

@@ -35,7 +35,7 @@ embeddings = ollama_ef(["This is my first text to embed",
 ```typescript
 import { OllamaEmbeddingFunction } from "chromadb";
 const embedder = new OllamaEmbeddingFunction({
-    url: "http://127.0.0.1:11434/api/embeddings",
+    url: "http://127.0.0.1:11434/",
     model: "llama2"
 })