Move to milvus from chroma

Author: Eroxl

backend/package.json

@@ -36,10 +36,10 @@
     "@types/ioredis": "^4.28.8",
     "@types/jsonwebtoken": "^8.5.8",
     "@types/node": "^17.0.10",
+    "@zilliz/milvus2-sdk-node": "^2.2.7",
     "add": "^2.0.6",
     "bcrypt": "^5.0.1",
     "body-parser": "^1.19.1",
-    "chromadb": "^1.3.1",
     "commonjs": "^0.0.1",
     "cors": "^2.8.5",
     "dotenv": "^14.2.0",

backend/src/helpers/clients/ChromaClient.ts

@@ -1,7 +1,7 @@
 import { Client } from '@elastic/elasticsearch';
 
-if (!process.env.ELASTICSEARCH_URL) {
-  throw new Error('ELASTICSEARCH_URL is not defined');
+if (!process.env.ELASTICSEARCH_URL && !process.env.ELASTICSEARCH_CLOUD_ID) {
+  throw new Error('ELASTICSEARCH_URL is not defined or ELASTICSEARCH_CLOUD_ID is not defined');
 }
 
 if (!process.env.ELASTICSEARCH_PASSWORD) {
@@ -12,8 +12,18 @@ if (!process.env.ELASTICSEARCH_USERNAME) {
   throw new Error('ELASTICSEARCH_USERNAME is not defined');
 }
 
+const URLConnection = process.env.ELASTICSEARCH_URL !== undefined
+  ? {
+    node: process.env.ELASTICSEARCH_URL as string,
+  }
+  : {
+    cloud: {
+      id: process.env.ELASTICSEARCH_CLOUD_ID as string,
+    }
+  }
+
 const client = new Client({
-  node: process.env.ELASTICSEARCH_URL,
+  ...URLConnection,
   auth: {
     password: process.env.ELASTICSEARCH_PASSWORD,
     username: process.env.ELASTICSEARCH_USERNAME,

backend/src/helpers/clients/ElasticSearchClient.ts

@@ -0,0 +1,78 @@
+import { MilvusClient, DataType } from "@zilliz/milvus2-sdk-node";
+
+if (!process.env.MILVUS_URL) {
+  throw new Error('MILVUS_URL is not defined');
+}
+
+const EMBEDDING_DIM = 1536;
+
+const BLOCK_FIELDS = [
+  {
+    name: 'block_id',
+    description: 'The ID of the block (Same as the ID of the block in mongoDB)',
+    data_type: DataType.VarChar,
+    max_length: 24,
+    is_primary_key: true,
+  },
+  {
+    name: 'page_id',
+    description: 'The ID of the page that the block belongs to',
+    max_length: 24,
+    data_type: DataType.VarChar,
+  },
+  {
+    name: 'embedding',
+    description: 'The embedding of the block (Generated by OpenAI)',
+    data_type: DataType.FloatVector,
+    dim: EMBEDDING_DIM,
+  },
+  {
+    name: 'content',
+    description: 'The content of the block',
+    data_type: DataType.VarChar,
+    max_length: 1000,
+  },
+  {
+    name: 'context',
+    description: 'The context of the block',
+    data_type: DataType.VarChar,
+    max_length: (24 + 2) * 11 + 2,
+  }
+]
+
+const milvusClient = new MilvusClient(
+  process.env.MILVUS_URL,
+  true,
+  process.env.MILVUS_USERNAME,
+  process.env.MILVUS_PASSWORD
+);
+
+(async () => {
+  const doesCollectionExist = await milvusClient.hasCollection({
+    collection_name: 'blocks',
+  });
+  
+  if (doesCollectionExist.status.error_code !== 'Success') {
+    throw new Error(doesCollectionExist.status.reason);
+  }
+
+  if (doesCollectionExist.value) return;
+
+  await milvusClient.createCollection({
+    collection_name: 'blocks',
+    description: 'Collection for storing block embeddings',
+    fields: BLOCK_FIELDS,
+  });
+
+  await milvusClient.createIndex({
+    collection_name: 'blocks',
+    field_name: 'embedding',
+    index_type: 'IVF_FLAT',
+    params: {
+      M: 8,
+      efConstruction: 64
+    }
+  })
+})();
+
+export default milvusClient;

backend/src/helpers/clients/MilvusClient.ts

@@ -1,7 +1,9 @@
+import { DataType } from '@zilliz/milvus2-sdk-node/dist/milvus';
+import type { ChatCompletionRequestMessage } from 'openai';
+
 import OpenAIClient from './clients/OpenAIClient';
-import ChromaClient from './clients/ChromaClient';
+import MilvusClient from './clients/MilvusClient';
 
-import type { ChatCompletionRequestMessage } from 'openai';
 
 const RELATIVE_TEXT_COUNT = 3;
 
@@ -13,30 +15,15 @@ const PRE_PROMPT = `You are a helpful AI assistant. Use the following pieces of
 If you don't know the answer, just say you don't know. DO NOT try to make up an answer.
 Try to keep your answers helpful, short and to the point using markdown formatting.`;
 
-interface QueryResponse {
-  ids: string[][];
-  documents: string[][];
-  metadatas: {
-    userID: string;
-    context: string[];
-  }[][];
-}
-
-interface GetQueryResponse {
-  ids: string[];
-  documents: string[];
-  metadatas: {
-    userID: string;
-    context: string[];
-  }[];
-}
-
 const getChatResponse = async (
   messages: ChatCompletionRequestMessage[],
   question: string,
-  user: string
+  page: string
 ): Promise<string> => {
-  const blockCollection = await ChromaClient.getCollection('blocks');
+  // ~ Load the block collection
+  await MilvusClient.loadCollection({
+    collection_name: 'blocks',
+  });
 
   // ~ Create an embedding for the latest message
   const embeddings = await OpenAIClient.createEmbedding({
@@ -45,45 +32,54 @@ const getChatResponse = async (
   });
 
   try {
-    // ~ Query the database for the most similar message
-    const similarMessages = await blockCollection.query(
-      embeddings.data.data[0].embedding,
-      RELATIVE_TEXT_COUNT,
-      {
-        userID: user,
+    const similarMessages = await MilvusClient.search({
+      collection_name: 'blocks',
+      limit: RELATIVE_TEXT_COUNT,
+      vector_type: DataType.FloatVector,
+      params: {
+        anns_field: 'block_id',
+        topk: `${RELATIVE_TEXT_COUNT}`,
+        metric_type: "L2",
+        params: JSON.stringify({ nprobe: 10 }),
       },
-    ) as QueryResponse;
+      vector: embeddings.data.data[0].embedding,
+      expr: `page_id == ${page}`,
+      output_fields: ['content', 'context']
+    })
 
     // ~ If there are no similar messages, return a default message
-    if (!similarMessages.documents) {
+    if (!similarMessages.results) {
       return 'I don\'t know what to say.';
     }
 
+    console.log(JSON.stringify(
+      similarMessages
+    ))
+
     // ~ Get the context messages
     const contextIDs = Array.from(
       new Set(
-        similarMessages.metadatas
+        similarMessages.results
           .flat()
-          .flatMap((metadata) => metadata.context)
+          .flatMap((metadata) => JSON.parse(metadata.context) as string[])
       )
-    )
+    );
 
-    const contextMessages = await blockCollection.get(
-      contextIDs,
-      {
-        userID: user,
-      }
-    ) as GetQueryResponse;
+    const contextMessages = await MilvusClient.query({
+      collection_name: 'blocks',
+      expr: `block_id in [${contextIDs.join(', ')}]`,
+      output_fields: ['content', 'block_id'],
+    })
 
     const contextMessagesMap: Record<string, string> = {};
 
-    contextMessages.documents.forEach((document, index) => {
-      contextMessagesMap[contextMessages.ids[index]] =  document;
+    contextMessages.data.forEach((result) => {
+      contextMessagesMap[result.block_id] =  result.content;
     });
 
-    const context = similarMessages.metadatas
-      .flat(2)
-      .flatMap((metadata) => metadata.context)
+    const context = similarMessages.results
+      .map((result) => result.context)
+      .flat()
       .map((id) => contextMessagesMap[id])
       .filter((message) => message)
       .map((message) => message.trim())

backend/src/helpers/getChatResponse.ts

@@ -1,4 +1,4 @@
-import ChromaClient from './clients/ChromaClient';
+import MilvusClient from './clients/MilvusClient';
 import OpenAIClient from './clients/OpenAIClient';
 
 export interface EmbedOperation {
@@ -14,15 +14,10 @@ export interface EmbedOperation {
  * @param page The page to update the embeds for.
  * @param pageData The page data of the page that is being updated.
  */
-const refreshEmbeds = async (updates: EmbedOperation[], page: string, pageOwner: string) => {
-  const blockCollection = await ChromaClient.getCollection('blocks');
-
-  await blockCollection.delete(
-    updates.map((operation) => `${page}.${operation.id}`),
-    {
-      userID: pageOwner,
-    }
-  );
+const refreshEmbeds = async (updates: EmbedOperation[], page: string) => {
+  await MilvusClient.loadCollection({
+    collection_name: 'blocks'
+  });
 
   const updateOperations = updates.filter((update) => update.type === 'update');
 
@@ -31,19 +26,32 @@ const refreshEmbeds = async (updates: EmbedOperation[], page: string, pageOwner:
     model: 'text-embedding-ada-002',
   });
 
-  const metaData = new Array(updateOperations.length)
+  const fieldsData = new Array(updateOperations.length)
     .fill(undefined)
     .map((_, index) => ({
-      userID: pageOwner,
-      context: updateOperations[index].context.map((context) => `${page}-${context}`),
+      block_id: updateOperations[index].id,
+      page_id: page,
+      embedding: embeddings.data.data[index].embedding,
+      content: updateOperations[index].value,
+      context: JSON.stringify(updateOperations[index].context),
     }));
 
-  await blockCollection.add(
-    updateOperations.map((operation) => operation.id),
-    embeddings.data.data.map((embedding) => embedding.embedding),
-    metaData,
-    updateOperations.map((operation) => operation?.value || ''),
-  );
+
+  await MilvusClient.insert({
+    collection_name: 'blocks',
+    fields_data: fieldsData
+  });
+
+  const deleteOperations = updates.filter((update) => update.type === 'delete');
+
+  await MilvusClient.deleteEntities({
+    collection_name: 'blocks',
+    expr: `block_id in [${deleteOperations.map((operation) => operation.id).join(', ')}]`,
+  });
+
+  await MilvusClient.flush({
+    collection_names: ['blocks'],
+  });
 };
 
 export default refreshEmbeds;

backend/src/helpers/refreshEmbeds.ts

@@ -10,11 +10,15 @@ import setupAuth from './setupAuth';
 
 // -=- Connect to MongoDB with dotenv file -=-
 dotenv.config();
-mongoose.connect(process.env.MONGO_URL ?? '');
+mongoose.connect(
+  process.env.MONGO_URL ?? '',
+).catch((err) => {
+  console.log(err);
+});
 
 // -=- Setup express -=-
 const app = express();
-const port = 8000;
+const port = 8000;  
 
 // -=- Setup Super Tokens -=-
 setupAuth();

backend/src/index.ts

@@ -12,7 +12,7 @@ router.get(
   '/chat',
   verifySession(),
   async (req: SessionRequest, res) => {
-    const { message, previousMessages } = req.query;
+    const { message, previousMessages, pageID } = req.query;
 
     if (typeof message !== 'string') {
       res.statusCode = 401;
@@ -23,6 +23,15 @@ router.get(
       return;
     }
 
+    if (typeof pageID !== 'string') {
+      res.statusCode = 401;
+      res.json({
+        status: 'error',
+        message: 'Please enter a pageID!',
+      });
+      return;
+    }
+
     if (previousMessages && typeof previousMessages !== 'string') {
       res.statusCode = 401;
       res.json({
@@ -55,7 +64,7 @@ router.get(
       messages = messages.slice(-10);
     }
 
-    const response = await getChatResponse(messages, message, req.session!.getUserId());
+    const response = await getChatResponse(messages, message, pageID);
 
     messages.push(
       {