feat: use milvus

app/api/chat/engine/generate.ts

@@ -1,8 +1,12 @@
 import * as dotenv from "dotenv";
 import { getDocuments } from "./loader";
 import { initSettings } from "./settings";
-import { storageContextFromDefaults, VectorStoreIndex } from "llamaindex";
-import { STORAGE_CACHE_DIR } from "@/cl/app/api/chat/engine/shared";
+import {
+  storageContextFromDefaults,
+  VectorStoreIndex,
+  MilvusVectorStore,
+} from "llamaindex";
+import { getMilvusClient } from "@/cl/app/api/chat/engine/shared";
 
 // Load environment variables from local .env.development.local file
 dotenv.config({ path: ".env.development.local" });
@@ -24,14 +28,17 @@ async function generateDatasource() {
   console.log(`Generating storage context for datasource '${datasource}'...`);
   // Split documents, create embeddings and store them in the storage context
   const ms = await getRuntime(async () => {
-    const storageContext = await storageContextFromDefaults({
-      persistDir: `${STORAGE_CACHE_DIR}/${datasource}`,
-    });
     const documents = await getDocuments(datasource);
     //  Set private=false to mark the document as public (required for filtering)
     documents.forEach((doc) => {
       doc.metadata["private"] = "false";
     });
+    const milvusClient = getMilvusClient();
+    const vectorStore = new MilvusVectorStore({
+      milvusClient,
+      collection: datasource,
+    });
+    const storageContext = await storageContextFromDefaults({ vectorStore });
     await VectorStoreIndex.fromDocuments(documents, {
       storageContext,
     });

app/api/chat/engine/index.ts

@@ -1,20 +1,29 @@
-import { SimpleDocumentStore, VectorStoreIndex } from "llamaindex";
-import { storageContextFromDefaults } from "llamaindex/storage/StorageContext";
-import { STORAGE_CACHE_DIR } from "@/cl/app/api/chat/engine/shared";
+import { VectorStoreIndex } from "llamaindex";
+import { MilvusVectorStore } from "llamaindex/storage/vectorStore/MilvusVectorStore";
+import {
+  checkRequiredEnvVars,
+  getMilvusClient,
+} from "@/cl/app/api/chat/engine/shared";
+
+const checkColllectionExist = async (collection: string) => {
+  const milvusClient = getMilvusClient();
+  const isCollectionExist = await milvusClient.hasCollection({
+    collection_name: collection,
+  });
+  return isCollectionExist.value;
+};
 
 export async function getDataSource(datasource: string) {
   console.log(`Using datasource: ${datasource}`);
-  const storageContext = await storageContextFromDefaults({
-    persistDir: `${STORAGE_CACHE_DIR}/${datasource}`,
-  });
+  checkRequiredEnvVars({ checkCollectionEnv: false }); // Do not check for collection env var
+  const milvusClient = getMilvusClient();
 
-  const numberOfDocs = Object.keys(
-    (storageContext.docStore as SimpleDocumentStore).toDict(),
-  ).length;
-  if (numberOfDocs === 0) {
-    return null;
+  // remove this code if you don't want to check collection existence before creating the index
+  // Milvus can automatically create the collection if it does not exist
+  if (!(await checkColllectionExist(datasource))) {
+    throw new Error(`Collection "${datasource}" does not exist`);
   }
-  return await VectorStoreIndex.init({
-    storageContext,
-  });
+
+  const store = new MilvusVectorStore({ milvusClient, collection: datasource });
+  return await VectorStoreIndex.fromVectorStore(store);
 }

create-llama.sh

@@ -20,7 +20,7 @@ npx -y create-llama@0.1.27 \
     --post-install-action none \
     --no-llama-parse \
     --example-file \
-    --vector-db none \
+    --vector-db milvus \
     --use-pnpm \
     -- cl >/dev/null
 

next.config.mjs

@@ -2,7 +2,7 @@ import { withSentryConfig } from "@sentry/nextjs";
 /** @type {import('next').NextConfig} */
 const nextConfig = {
   experimental: {
-    serverComponentsExternalPackages: ["pdf-parse"],
+    serverComponentsExternalPackages: ["pdf-parse", "@zilliz/milvus2-sdk-node"],
     outputFileTracingIncludes: {
       "/*": ["./cache/**/*"],
       "/api/**/*": ["node_modules/tiktoken/tiktoken_bg.wasm"]

package.json

@@ -80,7 +80,8 @@
     "@apidevtools/swagger-parser": "^10.1.0",
     "got": "10.7.0",
     "ajv": "^8.12.0",
-    "tiktoken": "^1.0.15"
+    "tiktoken": "^1.0.15",
+    "@zilliz/milvus2-sdk-node": "^2.4.4"
   },
   "devDependencies": {
     "@types/node": "^20.12.7",