use chromadb to run in browser

package-lock.json

@@ -15,6 +15,7 @@
                 "@types/react": "18.2.6",
                 "@types/react-dom": "18.2.4",
                 "@xenova/transformers": "^2.1.1",
+                "chromadb": "^1.5.2",
                 "dexie": "^3.2.4",
                 "eslint": "8.40.0",
                 "eslint-config-next": "13.4.2",
@@ -2212,9 +2213,7 @@
         "node_modules/chromadb": {
             "version": "1.5.2",
             "resolved": "https://registry.npmjs.org/chromadb/-/chromadb-1.5.2.tgz",
-            "integrity": "sha512-x/rOD7Oo1RiYA+vPK+Ma7CliCHlx26OjUt5J7Z9HZ5Ud1qDrPlvctBycK9Il3zqza96yeUoPQ7gCXHVKNoyvRQ==",
-            "optional": true,
-            "peer": true
+            "integrity": "sha512-x/rOD7Oo1RiYA+vPK+Ma7CliCHlx26OjUt5J7Z9HZ5Ud1qDrPlvctBycK9Il3zqza96yeUoPQ7gCXHVKNoyvRQ=="
         },
         "node_modules/client-only": {
             "version": "0.0.1",

package.json

@@ -15,11 +15,11 @@
         "@types/react": "18.2.6",
         "@types/react-dom": "18.2.4",
         "@xenova/transformers": "^2.1.1",
+        "chromadb": "^1.5.2",
         "dexie": "^3.2.4",
         "eslint": "8.40.0",
         "eslint-config-next": "13.4.2",
         "fs-extra": "^11.1.1",
-        "hnswlib-node": "^1.4.2",
         "langchain": "^0.0.90",
         "next": "13.4.2",
         "pdfjs-dist": "^3.7.107",

src/pages/api/docHandle.ts

@@ -1,15 +1,14 @@
 import type { NextApiRequest, NextApiResponse } from 'next';
-import {
-  vectorStoreToHNSWLibModel,
-} from '@/utils/file-handler';
 import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
-import { HNSWLib } from 'langchain/vectorstores/hnswlib';
+import { Chroma } from "langchain/vectorstores/chroma";
 import XenovaTransformersEmbeddings from '../../embed/hf'
 
 async function handleDocs(text: string) {
   const textSplitter = new RecursiveCharacterTextSplitter({ chunkSize: 1000 });
   const docs = await textSplitter.createDocuments([text]);
-  const vectorStore = await HNSWLib.fromDocuments(docs, new XenovaTransformersEmbeddings());
+  const vectorStore = await Chroma.fromDocuments(docs, new XenovaTransformersEmbeddings(), {
+    collectionName: 'docs'
+  });
   return vectorStore;
 }
 
@@ -25,9 +24,8 @@ export default async function handler(
   }
 
   const vectorStore = await handleDocs(text);
-  const model = await vectorStoreToHNSWLibModel(vectorStore);
   res.status(200).send({
-    ...model,
+    model: vectorStore,
   });
 }
 

src/utils/file-handler.ts

@@ -1,76 +0,0 @@
-import type XenovaTransformersEmbeddings from '@/embed/hf';
-// import { HuggingFaceInferenceEmbeddings } from 'langchain/embeddings/hf';
-import fs from 'fs-extra';
-import {
-  HNSWLib,
-  type HNSWLib as StoreTypeHNSWLib,
-} from 'langchain/vectorstores/hnswlib';
-import path from 'path';
-
-const ifDev = process.env.NODE_ENV === 'development';
-// in prod mode, only allowed to write to /tmp/
-// https://vercel.com/guides/how-can-i-use-files-in-serverless-functions
-export const storesDir = ifDev ? 'tmp/hnswlib-stores' : '/tmp/hnswlib-stores';
-
-type HNSWLibModel = {
-  args: string;
-  docstore: string;
-  hnswlibIndex: string;
-};
-
-const HNSWLibModelFilesName = {
-  args: 'args.json',
-  docstore: 'docstore.json',
-  hnswlibIndex: 'hnswlib.index',
-};
-
-// looking forward to a better way to transfrom hnswlibStore <=> indexes
-export async function HNSWLibModelToVectorStore(
-  model: HNSWLibModel,
-  embeddings: XenovaTransformersEmbeddings,
-) {
-  await saveHNSWLibModelToLocal(model);
-  // load from dir
-  const vectorStore = await HNSWLib.load(storesDir, embeddings);
-  return vectorStore;
-}
-
-export async function saveHNSWLibModelToLocal(model: HNSWLibModel) {
-  // save model to /tmp/
-  await Promise.all(
-    Object.keys(HNSWLibModelFilesName).map((key) => {
-      const fullPath = path.join(
-        storesDir,
-        (HNSWLibModelFilesName as Record<string, string>)[key],
-      );
-      console.log(fullPath);
-      const data = (model as Record<string, string>)[key];
-      console.log(data);
-
-      return fs.writeFile(fullPath, data);
-    }),
-  );
-}
-
-export async function vectorStoreToHNSWLibModel(
-  store: StoreTypeHNSWLib,
-): Promise<HNSWLibModel> {
-  await store.save(storesDir);
-  return await readHNSWLibModelFromLocal();
-}
-
-export async function readHNSWLibModelFromLocal(): Promise<HNSWLibModel> {
-  const [args, docstore, hnswlibIndex] = await Promise.all([
-    fs.readFile(path.join(storesDir, HNSWLibModelFilesName.args), 'utf-8'),
-    fs.readFile(path.join(storesDir, HNSWLibModelFilesName.docstore), 'utf-8'),
-    fs.readFile(
-      path.join(storesDir, HNSWLibModelFilesName.hnswlibIndex),
-      'hex',
-    ),
-  ]);
-  return {
-    args,
-    docstore,
-    hnswlibIndex,
-  };
-}