Refactor vector database integration to use Qdrant and update related configurations

- Replaced VectorDatabase with QdrantVectorDatabase across the application for improved vector similarity search.
- Updated docker-compose.yml to include Qdrant service with necessary configurations.
- Enhanced README.md to document Qdrant integration and its configuration options.
- Added new method in EmbeddingModel to retrieve embedding dimensions.
- Adjusted API and service files to accommodate Qdrant settings and ensure compatibility.

README.md

@@ -441,4 +441,27 @@ OPENAI_API_KEY=your_openai_api_key
 - Document processing with text chunking
 - Semantic search using embeddings
 - Question answering with LLM (OpenAI models)
-- Real-time chat interface
+- Real-time chat interface
+
+## Vector Database
+
+This application now uses [Qdrant](https://github.com/qdrant/qdrant-client) as its vector database. 
+Qdrant is a high-performance vector similarity search engine that stores both vectors and their metadata.
+
+### Features:
+- Fast vector search with HNSW index
+- Filtering support during search
+- Persisted storage of vectors and metadata
+- Both in-memory and disk-based options
+
+### Configuration:
+The following environment variables can be used to configure Qdrant:
+- `QDRANT_HOST`: Host of the Qdrant server (default: localhost)
+- `QDRANT_PORT`: HTTP port of the Qdrant server (default: 6333)
+- `QDRANT_GRPC_PORT`: gRPC port of the Qdrant server (default: 6334)
+- `QDRANT_PREFER_GRPC`: Whether to prefer gRPC over HTTP (default: true)
+- `QDRANT_COLLECTION`: Base name for collections (default: documents)
+- `QDRANT_IN_MEMORY`: Whether to use in-memory storage (default: true)
+
+When running with Docker, the application automatically connects to the Qdrant service
+defined in the docker-compose.yml file.

aimakerspace/openai_utils/embedding.py

@@ -48,6 +48,21 @@ class EmbeddingModel:
 
         return embedding.data[0].embedding
 
+    def get_embedding_dimension(self) -> int:
+        """Get the dimension of the embedding model
+        
+        Returns:
+            int: Dimension of the embedding model
+        """
+        # Dimensions for OpenAI models
+        dimensions = {
+            "text-embedding-3-small": 1536,
+            "text-embedding-3-large": 3072,
+            "text-embedding-ada-002": 1536,
+        }
+        
+        return dimensions.get(self.embeddings_model_name, 1536)  # Default to 1536
+
 
 if __name__ == "__main__":
     embedding_model = EmbeddingModel()

aimakerspace/qdrant_vectordb.py

@@ -0,0 +1,201 @@
+import numpy as np
+from typing import List, Tuple, Callable, Dict, Any, Optional
+import asyncio
+import uuid
+
+from qdrant_client import QdrantClient, AsyncQdrantClient
+from qdrant_client.http import models
+from qdrant_client.http.models import Distance, VectorParams, PointStruct
+
+from aimakerspace.openai_utils.embedding import EmbeddingModel
+
+
+class QdrantVectorDatabase:
+    """
+    Qdrant vector database implementation that follows the same interface
+    as the in-memory VectorDatabase class.
+    """
+    def __init__(self, 
+                 collection_name: str = "documents", 
+                 embedding_model: EmbeddingModel = None,
+                 host: str = "localhost",
+                 port: int = 6333,
+                 grpc_port: int = 6334,
+                 prefer_grpc: bool = True,
+                 in_memory: bool = True):
+        """
+        Initialize QdrantVectorDatabase
+        
+        Args:
+            collection_name: Name of the collection to use
+            embedding_model: Embedding model to use
+            host: Qdrant server host
+            port: Qdrant server port
+            grpc_port: Qdrant server gRPC port
+            prefer_grpc: Whether to prefer gRPC over HTTP
+            in_memory: Whether to use in-memory storage
+        """
+        self.collection_name = collection_name
+        self.embedding_model = embedding_model or EmbeddingModel()
+        self.in_memory = in_memory
+        
+        if in_memory:
+            self.client = QdrantClient(":memory:")
+            self.async_client = AsyncQdrantClient(":memory:")
+        else:
+            self.client = QdrantClient(
+                host=host, 
+                port=port,
+                grpc_port=grpc_port,
+                prefer_grpc=prefer_grpc
+            )
+            self.async_client = AsyncQdrantClient(
+                host=host, 
+                port=port,
+                grpc_port=grpc_port,
+                prefer_grpc=prefer_grpc
+            )
+            
+        # Store mapping from keys to ids
+        self.key_to_id: Dict[str, str] = {}
+        self.id_to_key: Dict[str, str] = {}
+        
+        # Create collection if it doesn't exist
+        vector_size = self.embedding_model.get_embedding_dimension()
+        self._ensure_collection(vector_size)
+        
+    def _ensure_collection(self, vector_size: int):
+        """Ensure collection exists"""
+        collections = self.client.get_collections().collections
+        collection_names = [c.name for c in collections]
+        
+        if self.collection_name not in collection_names:
+            self.client.create_collection(
+                collection_name=self.collection_name,
+                vectors_config=VectorParams(
+                    size=vector_size,
+                    distance=Distance.COSINE
+                )
+            )
+    
+    def insert(self, key: str, vector: np.array) -> None:
+        """Insert a vector into the database"""
+        # Generate a unique ID for this key
+        point_id = str(uuid.uuid4())
+        
+        # Store the mapping
+        self.key_to_id[key] = point_id
+        self.id_to_key[point_id] = key
+        
+        # Insert the point
+        self.client.upsert(
+            collection_name=self.collection_name,
+            points=[
+                PointStruct(
+                    id=point_id,
+                    vector=vector.tolist(),
+                    payload={"text": key}
+                )
+            ]
+        )
+    
+    def search(
+        self,
+        query_vector: np.array,
+        k: int,
+        distance_measure: Callable = None,  # Ignored, Qdrant uses its own distance measure
+    ) -> List[Tuple[str, float]]:
+        """Search for similar vectors"""
+        # Convert query_vector to list if it's a numpy array
+        if hasattr(query_vector, 'tolist'):
+            query_vector_list = query_vector.tolist()
+        else:
+            # If it's already a list or another iterable, convert to list to be safe
+            query_vector_list = list(query_vector)
+            
+        search_result = self.client.search(
+            collection_name=self.collection_name,
+            query_vector=query_vector_list,
+            limit=k
+        )
+        
+        results = []
+        for scored_point in search_result:
+            point_id = scored_point.id
+            score = scored_point.score
+            # Get the key from the id
+            if point_id in self.id_to_key:
+                key = self.id_to_key[point_id]
+                results.append((key, score))
+        
+        return results
+    
+    def search_by_text(
+        self,
+        query_text: str,
+        k: int,
+        distance_measure: Callable = None,  # Ignored, Qdrant uses its own distance measure
+        return_as_text: bool = False,
+    ) -> List[Tuple[str, float]]:
+        """Search by text query"""
+        query_vector = self.embedding_model.get_embedding(query_text)
+        results = self.search(query_vector, k, distance_measure)
+        return [result[0] for result in results] if return_as_text else results
+    
+    def retrieve_from_key(self, key: str) -> Optional[np.array]:
+        """Retrieve a vector by key"""
+        if key not in self.key_to_id:
+            return None
+        
+        point_id = self.key_to_id[key]
+        points = self.client.retrieve(
+            collection_name=self.collection_name,
+            ids=[point_id]
+        )
+        
+        if not points:
+            return None
+            
+        return np.array(points[0].vector)
+    
+    async def abuild_from_list(self, list_of_text: List[str]) -> "QdrantVectorDatabase":
+        """Build database from a list of texts"""
+        embeddings = await self.embedding_model.async_get_embeddings(list_of_text)
+        
+        # Generate unique IDs for each text
+        point_ids = [str(uuid.uuid4()) for _ in range(len(list_of_text))]
+        
+        # Store mappings
+        for text, point_id in zip(list_of_text, point_ids):
+            self.key_to_id[text] = point_id
+            self.id_to_key[point_id] = text
+        
+        # Prepare points for batch insertion
+        points = [
+            PointStruct(
+                id=point_id,
+                vector=embedding,
+                payload={"text": text}
+            )
+            for point_id, text, embedding in zip(point_ids, list_of_text, embeddings)
+        ]
+        
+        # Use batched upsert for efficiency
+        batch_size = 100
+        for i in range(0, len(points), batch_size):
+            batch = points[i:i+batch_size]
+            await self.async_client.upsert(
+                collection_name=self.collection_name,
+                points=batch
+            )
+        
+        return self
+    
+    def get_all_texts(self) -> List[str]:
+        """
+        Returns all the text documents stored in the vector database.
+        
+        Returns:
+            List[str]: A list of all text documents
+        """
+        return list(self.key_to_id.keys()) 

api/config.py

@@ -10,6 +10,14 @@ STATIC_DIR = "static"
 # Settings
 DEFAULT_NUM_SEARCH_RESULTS = 4  # Number of search results to return by default
 
+# Qdrant settings
+QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
+QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333))
+QDRANT_GRPC_PORT = int(os.getenv("QDRANT_GRPC_PORT", 6334))
+QDRANT_PREFER_GRPC = os.getenv("QDRANT_PREFER_GRPC", "True").lower() == "true"
+QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "documents")
+QDRANT_IN_MEMORY = os.getenv("QDRANT_IN_MEMORY", "True").lower() == "true"
+
 # Get env vars with defaults
 PORT = int(os.getenv("PORT", 8000))
 HOST = os.getenv("HOST", "0.0.0.0") 

api/main.py

@@ -16,13 +16,22 @@ from aimakerspace.openai_utils.prompts import (
     UserRolePrompt,
     SystemRolePrompt
 )
-from aimakerspace.vectordatabase import VectorDatabase
+from aimakerspace.qdrant_vectordb import QdrantVectorDatabase
 from aimakerspace.openai_utils.chatmodel import ChatOpenAI
 
 # API Version information
 API_VERSION = "0.2.0"
 BUILD_DATE = "2024-06-14"  # Update this when making significant changes
 
+# Qdrant settings from environment variables
+import os
+QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
+QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333))
+QDRANT_GRPC_PORT = int(os.getenv("QDRANT_GRPC_PORT", 6334))
+QDRANT_PREFER_GRPC = os.getenv("QDRANT_PREFER_GRPC", "True").lower() == "true"
+QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "documents")
+QDRANT_IN_MEMORY = os.getenv("QDRANT_IN_MEMORY", "True").lower() == "true"
+
 app = FastAPI(
     title="Quick Understand API",
     description="RAG-based question answering API for document understanding",
@@ -199,7 +208,14 @@ async def upload_file(
             texts = text_splitter.split_texts(documents)
             
             # Create vector database
-            vector_db = VectorDatabase()
+            vector_db = QdrantVectorDatabase(
+                collection_name=f"{QDRANT_COLLECTION}_{session_id}",
+                host=QDRANT_HOST,
+                port=QDRANT_PORT,
+                grpc_port=QDRANT_GRPC_PORT,
+                prefer_grpc=QDRANT_PREFER_GRPC,
+                in_memory=QDRANT_IN_MEMORY
+            )
             vector_db = await vector_db.abuild_from_list(texts)
             
             # Create chat model
@@ -766,7 +782,7 @@ async def catch_all(path: str):
     return FileResponse("static/index.html")
 
 class RetrievalAugmentedQAPipeline:
-    def __init__(self, llm: ChatOpenAI, vector_db_retriever: VectorDatabase, 
+    def __init__(self, llm: ChatOpenAI, vector_db_retriever: QdrantVectorDatabase, 
                 system_template: str = DEFAULT_SYSTEM_TEMPLATE, 
                 user_template: str = DEFAULT_USER_TEMPLATE) -> None:
         self.llm = llm

api/requirements.txt

@@ -6,4 +6,5 @@ openai==1.59.9
 pydantic==2.10.1
 pypdf2==3.0.1
 python-dotenv==1.0.1
-websockets==14.2 
+websockets==14.2
+qdrant-client==1.13.0 

api/routers/document.py

@@ -5,8 +5,10 @@ from fastapi import APIRouter, UploadFile, File, Form, HTTPException, Request, R
 from typing import Dict, List
 
 from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader, PDFLoader
+from aimakerspace.openai_utils.embedding import EmbeddingModel
 from aimakerspace.openai_utils.chatmodel import ChatOpenAI
-from aimakerspace.vectordatabase import VectorDatabase
+from aimakerspace.qdrant_vectordb import QdrantVectorDatabase
+from api.config import QDRANT_HOST, QDRANT_PORT, QDRANT_GRPC_PORT, QDRANT_PREFER_GRPC, QDRANT_COLLECTION, QDRANT_IN_MEMORY
 
 from api.models.pydantic_models import DocumentSummaryRequest, DocumentSummaryResponse
 from api.services.pipeline import RetrievalAugmentedQAPipeline
@@ -66,7 +68,14 @@ async def upload_file(
             texts = text_splitter.split_texts(documents)
             
             # Create vector database
-            vector_db = VectorDatabase()
+            vector_db = QdrantVectorDatabase(
+                collection_name=f"{QDRANT_COLLECTION}_{session_id}",
+                host=QDRANT_HOST,
+                port=QDRANT_PORT,
+                grpc_port=QDRANT_GRPC_PORT,
+                prefer_grpc=QDRANT_PREFER_GRPC,
+                in_memory=QDRANT_IN_MEMORY
+            )
             vector_db = await vector_db.abuild_from_list(texts)
             
             # Create chat model

api/services/pipeline.py

@@ -1,10 +1,10 @@
 from typing import List, Dict, Any
 from aimakerspace.openai_utils.prompts import SystemRolePrompt, UserRolePrompt
 from aimakerspace.openai_utils.chatmodel import ChatOpenAI
-from aimakerspace.vectordatabase import VectorDatabase
+from aimakerspace.qdrant_vectordb import QdrantVectorDatabase
 
 class RetrievalAugmentedQAPipeline:
-    def __init__(self, llm: ChatOpenAI, vector_db_retriever: VectorDatabase, 
+    def __init__(self, llm: ChatOpenAI, vector_db_retriever: QdrantVectorDatabase, 
                 system_template: str, 
                 user_template: str) -> None:
         self.llm = llm

app.py

@@ -8,7 +8,7 @@ from aimakerspace.openai_utils.prompts import (
     AssistantRolePrompt,
 )
 from aimakerspace.openai_utils.embedding import EmbeddingModel
-from aimakerspace.vectordatabase import VectorDatabase
+from aimakerspace.qdrant_vectordb import QdrantVectorDatabase
 from aimakerspace.openai_utils.chatmodel import ChatOpenAI
 import chainlit as cl
 
@@ -26,7 +26,7 @@ Question:
 user_role_prompt = UserRolePrompt(user_prompt_template)
 
 class RetrievalAugmentedQAPipeline:
-    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None:
+    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: QdrantVectorDatabase) -> None:
         self.llm = llm
         self.vector_db_retriever = vector_db_retriever
 
@@ -108,7 +108,7 @@ async def on_chat_start():
     print(f"Processing {len(texts)} text chunks")
 
     # Create a dict vector store
-    vector_db = VectorDatabase()
+    vector_db = QdrantVectorDatabase(collection_name="chainlit_documents")
     vector_db = await vector_db.abuild_from_list(texts)
     
     chat_openai = ChatOpenAI()

docker-compose.yml

@@ -1,6 +1,17 @@
 version: '3.8'
 
 services:
+  qdrant:
+    image: qdrant/qdrant:latest
+    ports:
+      - "6333:6333"
+      - "6334:6334"
+    volumes:
+      - qdrant_data:/qdrant/storage
+    environment:
+      - QDRANT_ALLOW_ORIGIN=*
+    restart: unless-stopped
+
   rag-app:
     build: 
       context: .
@@ -12,6 +23,17 @@ services:
       - OPENAI_API_KEY=${OPENAI_API_KEY}
       - PORT=7860
       - HOST=0.0.0.0
+      - QDRANT_HOST=qdrant
+      - QDRANT_PORT=6333
+      - QDRANT_GRPC_PORT=6334
+      - QDRANT_PREFER_GRPC=true
+      - QDRANT_COLLECTION=documents
+      - QDRANT_IN_MEMORY=false
     env_file:
       - .env
-    restart: unless-stopped 
+    depends_on:
+      - qdrant
+    restart: unless-stopped
+
+volumes:
+  qdrant_data: 

scripts/run_qdrant.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+
+echo "Starting Qdrant vector database..."
+
+# Check if container is already running
+if docker ps | grep -q "qdrant-server"; then
+    echo "Qdrant is already running."
+else
+    # Create a Docker volume for persistence
+    docker volume create qdrant_data
+
+    # Run Qdrant
+    docker run -d --name qdrant-server \
+        -p 6333:6333 \
+        -p 6334:6334 \
+        -v qdrant_data:/qdrant/storage \
+        -e QDRANT_ALLOW_ORIGIN="*" \
+        qdrant/qdrant:latest
+
+    echo "Qdrant started on ports 6333 (HTTP) and 6334 (gRPC)"
+fi
+
+echo "Qdrant is now available at http://localhost:6333"
+echo "Use Ctrl+C to exit this script (Qdrant will continue running in the background)"
+echo "To stop Qdrant later, run: docker stop qdrant-server"
+echo "To remove the container, run: docker rm qdrant-server"