less langchain more chroma

src/components/ChatWindow.jsx

@@ -3,9 +3,9 @@ import Image from "next/image";
 import { useCallback, useEffect, useState } from "react";
 import MessageList from './MessageList';
 import {FileLoader} from './FileLoader';
-import { db } from '@/utils/db-client';
 import Loader from "./Loader";
-import { Chroma } from "langchain/vectorstores/chroma";
+const {ChromaClient} = require('chromadb');
+const client = new ChromaClient();
 
 function ChatWindow({
   stopStrings,
@@ -27,10 +27,14 @@ function ChatWindow({
     }
     
     if (fileId) {
-      const similarityMatches = 2;
-      const fileContents = await db.docs.get(fileId);
-      const vectorStore = await Chroma.fromDocuments(fileContents, new XenovaTransformersEmbeddings());
-      const result = await vectorStore.similaritySearch(userInput, similarityMatches);
+      // const fileContents = await db.docs.get(fileId);
+      const collection = await client.getCollection("docs")
+      const result = await collection.query({
+        nResults: 2, 
+        queryTexts: [userPrompt]
+      });
+      console.log(result);
+
       const qaPrompt = 
       `You are an AI assistant providing helpful advice. You are given the following extracted parts of a long document and a question. Provide a conversational answer based on the context provided.
         You should only provide hyperlinks that reference the context below. Do NOT make up hyperlinks.

src/embed/hf.ts

@@ -1,41 +1,62 @@
-import { pipeline } from "@xenova/transformers";
- import { Embeddings, EmbeddingsParams } from "langchain/embeddings/base";
-
- export interface XenovaTransformersEmbeddingsParams extends EmbeddingsParams {
-   model?: string;
- }
-
- export default class XenovaTransformersEmbeddings
-   extends Embeddings
-   implements XenovaTransformersEmbeddingsParams
- {
-   model: string;
-
-   client: any;
-
-   constructor(fields?: XenovaTransformersEmbeddingsParams) {
-     super(fields ?? {});
-     this.model = fields?.model ?? "Xenova/all-MiniLM-L6-v2";
-   }
-
-   async _embed(texts: string[]): Promise<number[][]> {
-     if (!this.client) {
-       this.client = await pipeline("embeddings", this.model);
-     }
-     console.log(this.client, texts);
-
-     return this.caller.call(async () => {
-       return await Promise.all(
-         texts.map(async (t) => (await this.client(t, { pooling: 'mean', normalize: true })).data)
-       );
-     });
-   }
-
-   embedQuery(document: string): Promise<number[]> {
-     return this._embed([document]).then((embeddings) => embeddings[0]);
-   }
-
-   embedDocuments(documents: string[]): Promise<number[][]> {
-     return this._embed(documents);
-   }
- }
+import { IEmbeddingFunction } from "chromadb/src/embeddings/IEmbeddingFunction";
+
+// Dynamically import module
+let TransformersApi: Promise<any>;
+
+export class TransformersEmbeddingFunction implements IEmbeddingFunction {
+  private pipelinePromise: Promise<any> | null;
+
+  /**
+   * TransformersEmbeddingFunction constructor.
+   * @param options The configuration options.
+   * @param options.model The model to use to calculate embeddings. Defaults to 'Xenova/all-MiniLM-L6-v2', which is an ONNX port of `sentence-transformers/all-MiniLM-L6-v2`.
+   * @param options.revision The specific model version to use (can be a branch, tag name, or commit id). Defaults to 'main'.
+   * @param options.quantized Whether to load the 8-bit quantized version of the model. Defaults to `false`.
+   * @param options.progress_callback If specified, this function will be called during model construction, to provide the user with progress updates.
+   */
+  constructor({
+    model = "Xenova/all-MiniLM-L6-v2",
+    revision = "main",
+    quantized = false,
+    progress_callback = null,
+  }: {
+    model?: string;
+    revision?: string;
+    quantized?: boolean;
+    progress_callback?: Function | null;
+  } = {}) {
+    try {
+      // Since Transformers.js is an ESM package, we use the dynamic `import` syntax instead of `require`.
+      // Also, since we use `"module": "commonjs"` in tsconfig.json, we use the following workaround to ensure
+      // the dynamic import is not transpiled to a `require` statement.
+      // For more information, see https://github.com/microsoft/TypeScript/issues/43329#issuecomment-1008361973
+      TransformersApi = Function('return import("@xenova/transformers")')();
+    } catch (e) {
+      throw new Error(
+        "Please install the @xenova/transformers package to use the TransformersEmbeddingFunction, `npm install -S @xenova/transformers`."
+      );
+    }
+
+    // Store a promise that resolves to the pipeline
+    this.pipelinePromise = new Promise(async (resolve, reject) => {
+      try {
+        const { pipeline } = await TransformersApi;
+        resolve(
+          await pipeline("feature-extraction", model, {
+            quantized,
+            revision,
+            progress_callback,
+          })
+        );
+      } catch (e) {
+        reject(e);
+      }
+    });
+  }
+
+  public async generate(texts: string[]): Promise<number[][]> {
+    let pipe = await this.pipelinePromise;
+    let output = await pipe(texts, { pooling: "mean", normalize: true });
+    return output.tolist();
+  }
+}

src/pages/api/docHandle.ts

@@ -1,15 +1,20 @@
 import type { NextApiRequest, NextApiResponse } from 'next';
 import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
-import { Chroma } from "langchain/vectorstores/chroma";
-import XenovaTransformersEmbeddings from '../../embed/hf'
+import { TransformersEmbeddingFunction } from '../../embed/hf';
+const {ChromaClient} = require('chromadb');
+const client = new ChromaClient();
 
 async function handleDocs(text: string) {
   const textSplitter = new RecursiveCharacterTextSplitter({ chunkSize: 1000 });
   const docs = await textSplitter.createDocuments([text]);
-  const vectorStore = await Chroma.fromDocuments(docs, new XenovaTransformersEmbeddings(), {
-    collectionName: 'docs'
+  const collection = await client.createCollection({name: "docs", embeddingFunction: TransformersEmbeddingFunction})
+  await collection.add({
+      ids: [...docs.map((v, k) => k)],
+      metadatas: [...docs.map(doc => doc.metadata)],
+      documents: [...docs.map(doc => doc.pageContent)],
   });
-  return vectorStore;
+
+  return collection;
 }
 
 export default async function handler(