From 1be69a5aa4a979fa137917edda7fc91f2347b1a6 Mon Sep 17 00:00:00 2001
From: Anush <anushshetty90@gmail.com>
Date: Wed, 10 Apr 2024 08:07:33 +0530
Subject: [PATCH] feat: Qdrant support (#42)

---
 .changeset/lovely-olives-carry.md             |  5 +++
 helpers/env-variables.ts                      | 17 +++++++++
 helpers/types.ts                              |  3 +-
 questions.ts                                  |  1 +
 .../vectordbs/python/qdrant/__init__.py       |  0
 .../vectordbs/python/qdrant/generate.py       | 37 +++++++++++++++++++
 .../vectordbs/python/qdrant/index.py          | 20 ++++++++++
 .../vectordbs/typescript/qdrant/generate.ts   | 37 +++++++++++++++++++
 .../vectordbs/typescript/qdrant/index.ts      | 13 +++++++
 .../vectordbs/typescript/qdrant/shared.ts     | 32 ++++++++++++++++
 10 files changed, 164 insertions(+), 1 deletion(-)
 create mode 100644 .changeset/lovely-olives-carry.md
 create mode 100644 templates/components/vectordbs/python/qdrant/__init__.py
 create mode 100644 templates/components/vectordbs/python/qdrant/generate.py
 create mode 100644 templates/components/vectordbs/python/qdrant/index.py
 create mode 100644 templates/components/vectordbs/typescript/qdrant/generate.ts
 create mode 100644 templates/components/vectordbs/typescript/qdrant/index.ts
 create mode 100644 templates/components/vectordbs/typescript/qdrant/shared.ts

diff --git a/.changeset/lovely-olives-carry.md b/.changeset/lovely-olives-carry.md
new file mode 100644
index 00000000..99741125
--- /dev/null
+++ b/.changeset/lovely-olives-carry.md
@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+Add Qdrant support
diff --git a/helpers/env-variables.ts b/helpers/env-variables.ts
index 65782179..ecaed32b 100644
--- a/helpers/env-variables.ts
+++ b/helpers/env-variables.ts
@@ -108,6 +108,23 @@ const getVectorDBEnvs = (vectorDb: TemplateVectorDB) => {
           description: "The name of the collection in your Astra database",
         },
       ];
+    case "qdrant":
+      return [
+        {
+          name: "QDRANT_URL",
+          description:
+            "The qualified REST URL of the Qdrant server. Eg: http://localhost:6333",
+        },
+        {
+          name: "QDRANT_COLLECTION",
+          description: "The name of Qdrant collection to use.",
+        },
+        {
+          name: "QDRANT_API_KEY",
+          description:
+            "Optional API key for authenticating requests to Qdrant.",
+        },
+      ];
     default:
       return [];
   }
diff --git a/helpers/types.ts b/helpers/types.ts
index 52a300a7..0db557ed 100644
--- a/helpers/types.ts
+++ b/helpers/types.ts
@@ -10,7 +10,8 @@ export type TemplateVectorDB =
   | "pg"
   | "pinecone"
   | "milvus"
-  | "astra";
+  | "astra"
+  | "qdrant";
 export type TemplatePostInstallAction =
   | "none"
   | "VSCode"
diff --git a/questions.ts b/questions.ts
index 4afd0085..6efbc45d 100644
--- a/questions.ts
+++ b/questions.ts
@@ -102,6 +102,7 @@ const getVectorDbChoices = (framework: TemplateFramework) => {
     { title: "Pinecone", value: "pinecone" },
     { title: "Milvus", value: "milvus" },
     { title: "Astra", value: "astra" },
+    { title: "Qdrant", value: "qdrant" },
   ];
 
   const vectordbLang = framework === "fastapi" ? "python" : "typescript";
diff --git a/templates/components/vectordbs/python/qdrant/__init__.py b/templates/components/vectordbs/python/qdrant/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/templates/components/vectordbs/python/qdrant/generate.py b/templates/components/vectordbs/python/qdrant/generate.py
new file mode 100644
index 00000000..db7c055e
--- /dev/null
+++ b/templates/components/vectordbs/python/qdrant/generate.py
@@ -0,0 +1,37 @@
+import logging
+import os
+from app.engine.loaders import get_documents
+from app.settings import init_settings
+from dotenv import load_dotenv
+from llama_index.core.indices import VectorStoreIndex
+from llama_index.core.storage import StorageContext
+from llama_index.vector_stores.qdrant import QdrantVectorStore
+load_dotenv()
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger()
+
+
+def generate_datasource():
+    init_settings()
+    logger.info("Creating new index with Qdrant")
+    # load the documents and create the index
+    documents = get_documents()
+    store = QdrantVectorStore(
+        collection_name=os.getenv("QDRANT_COLLECTION"),
+        url=os.getenv("QDRANT_URL"),
+        api_key=os.getenv("QDRANT_API_KEY"),
+    )
+    storage_context = StorageContext.from_defaults(vector_store=store)
+    VectorStoreIndex.from_documents(
+        documents,
+        storage_context=storage_context,
+        show_progress=True,  # this will show you a progress bar as the embeddings are created
+    )
+    logger.info(
+        f"Successfully uploaded documents to the {os.getenv('QDRANT_COLLECTION')} collection."
+    )
+
+
+if __name__ == "__main__":
+    generate_datasource()
diff --git a/templates/components/vectordbs/python/qdrant/index.py b/templates/components/vectordbs/python/qdrant/index.py
new file mode 100644
index 00000000..0a388d8a
--- /dev/null
+++ b/templates/components/vectordbs/python/qdrant/index.py
@@ -0,0 +1,20 @@
+import logging
+import os
+
+from llama_index.core.indices import VectorStoreIndex
+from llama_index.vector_stores.qdrant import QdrantVectorStore
+
+
+logger = logging.getLogger("uvicorn")
+
+
+def get_index():
+    logger.info("Connecting to Qdrant collection..")
+    store = QdrantVectorStore(
+        collection_name=os.getenv("QDRANT_COLLECTION"),
+        url=os.getenv("QDRANT_URL"),
+        api_key=os.getenv("QDRANT_API_KEY"),
+    )
+    index = VectorStoreIndex.from_vector_store(store)
+    logger.info("Finished connecting to Qdrant collection.")
+    return index
diff --git a/templates/components/vectordbs/typescript/qdrant/generate.ts b/templates/components/vectordbs/typescript/qdrant/generate.ts
new file mode 100644
index 00000000..f0e7fa40
--- /dev/null
+++ b/templates/components/vectordbs/typescript/qdrant/generate.ts
@@ -0,0 +1,37 @@
+/* eslint-disable turbo/no-undeclared-env-vars */
+import * as dotenv from "dotenv";
+import {
+  QdrantVectorStore,
+  VectorStoreIndex,
+  storageContextFromDefaults,
+} from "llamaindex";
+import { getDocuments } from "./loader";
+import { initSettings } from "./settings";
+import { checkRequiredEnvVars, getQdrantClient } from "./shared";
+
+dotenv.config();
+
+const collectionName = process.env.QDRANT_COLLECTION;
+
+async function loadAndIndex() {
+  // load objects from storage and convert them into LlamaIndex Document objects
+  const documents = await getDocuments();
+
+  // Connect to Qdrant
+  const vectorStore = new QdrantVectorStore(collectionName, getQdrantClient());
+
+  const storageContext = await storageContextFromDefaults({ vectorStore });
+  await VectorStoreIndex.fromDocuments(documents, {
+    storageContext: storageContext,
+  });
+  console.log(
+    `Successfully upload embeddings to Qdrant collection ${collectionName}.`,
+  );
+}
+
+(async () => {
+  checkRequiredEnvVars();
+  initSettings();
+  await loadAndIndex();
+  console.log("Finished generating storage.");
+})();
diff --git a/templates/components/vectordbs/typescript/qdrant/index.ts b/templates/components/vectordbs/typescript/qdrant/index.ts
new file mode 100644
index 00000000..229d7a39
--- /dev/null
+++ b/templates/components/vectordbs/typescript/qdrant/index.ts
@@ -0,0 +1,13 @@
+import * as dotenv from "dotenv";
+import { QdrantVectorStore, VectorStoreIndex } from "llamaindex";
+import { checkRequiredEnvVars, getQdrantClient } from "./shared";
+
+dotenv.config();
+
+export async function getDataSource() {
+  checkRequiredEnvVars();
+  const collectionName = process.env.QDRANT_COLLECTION;
+  const store = new QdrantVectorStore(collectionName, getQdrantClient());
+
+  return await VectorStoreIndex.fromVectorStore(store);
+}
diff --git a/templates/components/vectordbs/typescript/qdrant/shared.ts b/templates/components/vectordbs/typescript/qdrant/shared.ts
new file mode 100644
index 00000000..e1475eee
--- /dev/null
+++ b/templates/components/vectordbs/typescript/qdrant/shared.ts
@@ -0,0 +1,32 @@
+import { QdrantClient } from "@qdrant/js-client-rest";
+
+const REQUIRED_ENV_VARS = ["QDRANT_URL", "QDRANT_COLLECTION"]; // QDRANT_API_KEY is optional
+
+export function getQdrantClient() {
+  const url = process.env.QDRANT_URL;
+  if (!url) {
+    throw new Error("QDRANT_URL environment variable is required");
+  }
+  const apiKey = process.env?.QDRANT_API_KEY;
+  return new QdrantClient({
+    url,
+    apiKey,
+  });
+}
+
+export function checkRequiredEnvVars() {
+  const missingEnvVars = REQUIRED_ENV_VARS.filter((envVar) => {
+    return !process.env[envVar];
+  });
+
+  if (missingEnvVars.length > 0) {
+    console.log(
+      `The following environment variables are required but missing: ${missingEnvVars.join(
+        ", ",
+      )}`,
+    );
+    throw new Error(
+      `Missing environment variables: ${missingEnvVars.join(", ")}`,
+    );
+  }
+}
-- 
GitLab