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