🚀 feat: Qdrant

.gitignore

@@ -6,3 +6,5 @@ uploads/
 myenv/
 venv/
 *.pyc
+env/
+env2/

.vscode/launch.json

@@ -13,6 +13,19 @@
               "--reload"
           ],
           "jinja": true
-      }
+      },
+      {
+        "name": "Debug FastAPI (debugpy)",
+        "type": "debugpy",
+        "request": "launch",
+        "module": "uvicorn",
+        "args": [
+            "main:app",
+            "--port",
+            "8000",
+            "--reload"
+        ],
+        "jinja": true
+    }
   ]
 }

config.py

@@ -4,6 +4,7 @@
 import logging
 from enum import Enum
 from datetime import datetime
+from threading import TIMEOUT_MAX
 from dotenv import find_dotenv, load_dotenv
 from langchain_community.embeddings import (
     HuggingFaceEmbeddings,
@@ -12,14 +13,15 @@
 )
 from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
 from starlette.middleware.base import BaseHTTPMiddleware
-from store_factory import get_vector_store
+from store_factory import QdrantConfig, get_vector_store
 
 load_dotenv(find_dotenv())
 
 
 class VectorDBType(Enum):
     PGVECTOR = "pgvector"
     ATLAS_MONGO = "atlas-mongo"
+    QDRANT = "qdrant"
 
 
 class EmbeddingsProvider(Enum):
@@ -63,9 +65,22 @@ def get_env_variable(
 MONGO_VECTOR_COLLECTION = get_env_variable(
     "MONGO_VECTOR_COLLECTION", "vector_collection"
 )
+QDRANT_HOST= get_env_variable("QDRANT_HOST", "127.0.0.1:6333")
+COLLECTION_NAME = get_env_variable(
+    "COLLECTION_NAME", "vector_collection"
+)
+EMBEDDINGS_DIMENSION = get_env_variable("EMBEDDINGS_DIMENSION", "768")
+QDRANT_API_KEY = get_env_variable("QDRANT_API_KEY")
 
 CHUNK_SIZE = int(get_env_variable("CHUNK_SIZE", "1500"))
 CHUNK_OVERLAP = int(get_env_variable("CHUNK_OVERLAP", "100"))
+maxChunks = get_env_variable("MAX_CHUNKS")
+MAX_CHUNKS = int(maxChunks) if maxChunks else None 
+embeddingTimeout = get_env_variable("EMBEDDING_TIMEOUT")
+EMBEDDING_TIMEOUT = int(get_env_variable("EMBEDDING_TIMEOUT",100000))#default 100 second timeout
+
+BATCH_SIZE = int(get_env_variable("BATCH_SIZE","75"))
+CONCURRENT_LIMIT = int(get_env_variable("CONCURRENT_LIMIT","10"))
 
 env_value = get_env_variable("PDF_EXTRACT_IMAGES", "False").lower()
 PDF_EXTRACT_IMAGES = True if env_value == "true" else False
@@ -228,12 +243,13 @@ def init_embeddings(provider, model):
 logger.info(f"Initialized embeddings of type: {type(embeddings)}")
 
 # Vector store
+print(VECTOR_DB_TYPE)
 if VECTOR_DB_TYPE == VectorDBType.PGVECTOR:
     vector_store = get_vector_store(
         connection_string=CONNECTION_STRING,
         embeddings=embeddings,
         collection_name=COLLECTION_NAME,
-        mode="async",
+        mode="async-PGVector",
     )
 elif VECTOR_DB_TYPE == VectorDBType.ATLAS_MONGO:
     logger.warning("Using Atlas MongoDB as vector store is not fully supported yet.")
@@ -243,9 +259,21 @@ def init_embeddings(provider, model):
         collection_name=MONGO_VECTOR_COLLECTION,
         mode="atlas-mongo",
     )
+elif VECTOR_DB_TYPE == VectorDBType.QDRANT:
+    vector_store = get_vector_store(
+        connection_string=QDRANT_HOST,
+        embeddings=embeddings,
+        collection_name=COLLECTION_NAME,
+        mode="qdrant",
+        additional_kwargs = QdrantConfig(
+            qdrant_api_key=QDRANT_API_KEY,
+            embeddings_dimension=EMBEDDINGS_DIMENSION
+        )
+)
 else:
     raise ValueError(f"Unsupported vector store type: {VECTOR_DB_TYPE}")
 
+
 retriever = vector_store.as_retriever()
 
 known_source_ext = [

main.py

@@ -1,4 +1,5 @@
 import os
+import time
 import hashlib
 import aiofiles
 import aiofiles.os
@@ -12,6 +13,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from langchain_core.runnables.config import run_in_executor
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from uuid import uuid1
 from fastapi import (
     File,
     Form,
@@ -49,7 +51,8 @@
 from middleware import security_middleware
 from mongo import mongo_health_check
 from constants import ERROR_MESSAGES
-from store import AsyncPgVector
+from store import AsyncPgVector, AsyncQdrant, async_DB 
+from process_docs import store_documents
 
 load_dotenv(find_dotenv())
 
@@ -58,18 +61,21 @@
     debug_mode,
     CHUNK_SIZE,
     CHUNK_OVERLAP,
+    EMBEDDING_TIMEOUT,
+    MAX_CHUNKS,
+    BATCH_SIZE,
+    CONCURRENT_LIMIT,
     vector_store,
     RAG_UPLOAD_DIR,
     known_source_ext,
     PDF_EXTRACT_IMAGES,
     LogMiddleware,
     RAG_HOST,
     RAG_PORT,
-    VectorDBType,
+    VECTOR_DB_TYPE,
     # RAG_EMBEDDING_MODEL,
     # RAG_EMBEDDING_MODEL_DEVICE_TYPE,
     # RAG_TEMPLATE,
-    VECTOR_DB_TYPE,
 )
 
 
@@ -85,6 +91,7 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -105,7 +112,7 @@ async def lifespan(app: FastAPI):
 @app.get("/ids")
 async def get_all_ids():
     try:
-        if isinstance(vector_store, AsyncPgVector):
+        if isinstance(vector_store, async_DB):
             ids = await vector_store.get_all_ids()
         else:
             ids = vector_store.get_all_ids()
@@ -116,12 +123,7 @@ async def get_all_ids():
 
 
 def isHealthOK():
-    if VECTOR_DB_TYPE == VectorDBType.PGVECTOR:
-        return pg_health_check()
-    if VECTOR_DB_TYPE == VectorDBType.ATLAS_MONGO:
-        return mongo_health_check()
-    else:
-        return True
+    return pg_health_check()
 
 
 @app.get("/health")
@@ -135,7 +137,7 @@ async def health_check():
 @app.get("/documents", response_model=list[DocumentResponse])
 async def get_documents_by_ids(ids: list[str] = Query(...)):
     try:
-        if isinstance(vector_store, AsyncPgVector):
+        if isinstance(vector_store, async_DB):
             existing_ids = await vector_store.get_all_ids()
             documents = await vector_store.get_documents_by_ids(ids)
         else:
@@ -145,13 +147,6 @@ async def get_documents_by_ids(ids: list[str] = Query(...)):
         # Ensure all requested ids exist
         if not all(id in existing_ids for id in ids):
             raise HTTPException(status_code=404, detail="One or more IDs not found")
-
-        # Ensure documents list is not empty
-        if not documents:
-            raise HTTPException(
-                status_code=404, detail="No documents found for the given IDs"
-            )
-
         return documents
     except HTTPException as http_exc:
         raise http_exc
@@ -160,19 +155,17 @@ async def get_documents_by_ids(ids: list[str] = Query(...)):
 
 
 @app.delete("/documents")
-async def delete_documents(document_ids: List[str] = Body(...)):
+async def delete_documents(ids: list[str]):
     try:
-        if isinstance(vector_store, AsyncPgVector):
+        if isinstance(vector_store, async_DB):
             existing_ids = await vector_store.get_all_ids()
-            await vector_store.delete(ids=document_ids)
+            await vector_store.delete(ids=ids)
+            if not all(id in existing_ids for id in ids):
+                raise HTTPException(status_code=404, detail="One or more IDs not found")
         else:
             existing_ids = vector_store.get_all_ids()
-            vector_store.delete(ids=document_ids)
-
-        if not all(id in existing_ids for id in document_ids):
-            raise HTTPException(status_code=404, detail="One or more IDs not found")
-
-        file_count = len(document_ids)
+            vector_store.delete(ids=ids)
+        file_count = len(ids)
         return {
             "message": f"Documents for {file_count} file{'s' if file_count > 1 else ''} deleted successfully"
         }
@@ -188,9 +181,8 @@ async def query_embeddings_by_file_id(body: QueryRequestBody, request: Request):
     authorized_documents = []
 
     try:
-        embedding = vector_store.embedding_function.embed_query(body.query)
-
-        if isinstance(vector_store, AsyncPgVector):
+        if isinstance(vector_store, async_DB):
+            embedding = vector_store.embedding_function.embed_query(body.query)
             documents = await run_in_executor(
                 None,
                 vector_store.similarity_search_with_score_by_vector,
@@ -199,6 +191,7 @@ async def query_embeddings_by_file_id(body: QueryRequestBody, request: Request):
                 filter={"file_id": body.file_id},
             )
         else:
+            embedding = vector_store.embedding_function.embed_query(body.query)
             documents = vector_store.similarity_search_with_score_by_vector(
                 embedding, k=body.k, filter={"file_id": body.file_id}
             )
@@ -234,12 +227,18 @@ async def store_data_in_vector_db(
     file_id: str,
     user_id: str = "",
     clean_content: bool = False,
-) -> bool:
+) -> dict:
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=app.state.CHUNK_SIZE, chunk_overlap=app.state.CHUNK_OVERLAP
     )
+    ids = []
     documents = text_splitter.split_documents(data)
 
+    if MAX_CHUNKS and len(documents) > MAX_CHUNKS:
+        raise HTTPException(
+            status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
+            detail=f"Too big file, Attempted to process {len(documents)} chunks, but MAX_CHUNKS is set to {MAX_CHUNKS} chunks",
+        )
     # If `clean_content` is True, clean the page_content of each document (remove null bytes)
     if clean_content:
         for doc in documents:
@@ -258,21 +257,21 @@ async def store_data_in_vector_db(
         )
         for doc in documents
     ]
+
 
     try:
-        if isinstance(vector_store, AsyncPgVector):
-            ids = await vector_store.aadd_documents(
-                docs, ids=[file_id] * len(documents)
-            )
+        uuids = [str(uuid1()) for _ in docs]
+        if isinstance(vector_store, async_DB):
+            ids = await store_documents(docs, ids=uuids)
         else:
-            ids = vector_store.add_documents(docs, ids=[file_id] * len(documents))
+            ids = vector_store.add_documents(docs, ids=uuids)
 
         return {"message": "Documents added successfully", "ids": ids}
 
     except Exception as e:
         logger.error(e)
         return {"message": "An error occurred while adding documents.", "error": str(e)}
-
+    
 
 def get_loader(filename: str, file_content_type: str, filepath: str):
     file_ext = filename.split(".")[-1].lower()
@@ -303,8 +302,6 @@ def get_loader(filename: str, file_content_type: str, filepath: str):
         "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
     ] or file_ext in ["xls", "xlsx"]:
         loader = UnstructuredExcelLoader(filepath)
-    elif file_content_type == "application/json" or file_ext == "json":
-        loader = TextLoader(filepath, autodetect_encoding=True)
     elif file_ext in known_source_ext or (
         file_content_type and file_content_type.find("text/") >= 0
     ):
@@ -375,7 +372,6 @@ async def embed_file(
         user_id = "public"
     else:
         user_id = request.state.user.get("id")
-
     temp_base_path = os.path.join(RAG_UPLOAD_DIR, user_id)
     os.makedirs(temp_base_path, exist_ok=True)
     temp_file_path = os.path.join(RAG_UPLOAD_DIR, user_id, file.filename)
@@ -443,7 +439,7 @@ async def embed_file(
 async def load_document_context(id: str):
     ids = [id]
     try:
-        if isinstance(vector_store, AsyncPgVector):
+        if isinstance(vector_store, async_DB):
             existing_ids = await vector_store.get_all_ids()
             documents = await vector_store.get_documents_by_ids(ids)
         else:
@@ -531,14 +527,20 @@ async def query_embeddings_by_file_ids(body: QueryMultipleBody):
         if isinstance(vector_store, AsyncPgVector):
             documents = await run_in_executor(
                 None,
-                vector_store.similarity_search_with_score_by_vector,
+                vector_store.asimilarity_search_with_score_by_vector,
                 embedding,
                 k=body.k,
-                filter={"file_id": {"$in": body.file_ids}},
+                filter={"custom_id": {"$in": body.file_ids}},
             )
+        elif isinstance(vector_store, AsyncQdrant):
+            documents = await vector_store.similarity_search_many(
+                embedding=embedding,
+                k=body.k,
+                ids=body.file_ids)
+
         else:
             documents = vector_store.similarity_search_with_score_by_vector(
-                embedding, k=body.k, filter={"file_id": {"$in": body.file_ids}}
+                embedding, k=body.k, filter={"custom_id": {"$in": body.file_ids}}
             )
 
         # Ensure documents list is not empty
@@ -556,4 +558,4 @@ async def query_embeddings_by_file_ids(body: QueryMultipleBody):
     app.include_router(router=pgvector_router)
 
 if __name__ == "__main__":
-    uvicorn.run(app, host=RAG_HOST, port=RAG_PORT, log_config=None)
+    uvicorn.run(app, host=RAG_HOST, port=RAG_PORT, log_config=None)

process_docs.py

@@ -0,0 +1,33 @@
+from langchain.schema import Document
+import asyncio
+from config import  (
+    CONCURRENT_LIMIT,
+    BATCH_SIZE,
+    EMBEDDING_TIMEOUT,
+    vector_store
+)
+
+
+#Prepare documents to be async added to vectorstore async in batches
+async def store_documents(
+    docs: list[Document], ids:list[str]
+):
+    semaphore = asyncio.Semaphore(CONCURRENT_LIMIT)
+    tasks = []
+    #logger.info(f"Processing list of documents of length: {len(docs)}")
+    for i in range(0, len(docs), BATCH_SIZE):
+        batch = docs[i : min(i + BATCH_SIZE, len(docs))]
+        #logger.info(f"Sending batch {i} to {i+len(batch)} / {len(docs)}")
+        task = asyncio.create_task(process_batch(batch, ids, semaphore))
+        tasks.append(task)
+    try:
+        idList = await asyncio.wait_for(asyncio.gather(*tasks), timeout=(EMBEDDING_TIMEOUT/1000)) 
+    except asyncio.TimeoutError:
+        raise Exception(f"TIMEOUT: embedding process took over the time limit of {EMBEDDING_TIMEOUT}ms. Partially added to database")
+    return [id for sublist in idList for id in sublist]
+
+#Helper for process_documents
+async def process_batch(batch: list[Document], ids: list[str], semaphore):
+    async with semaphore:
+        return await vector_store.aadd_documents(batch,ids=ids)
+