add Qdrant

.gitignore

@@ -1 +1,2 @@
 __pycache__/
+.env

app.py

@@ -1,4 +1,5 @@
 import os
+import numpy as np
 from typing import List
 from chainlit.types import AskFileResponse
 from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader
@@ -11,12 +12,19 @@ from aimakerspace.openai_utils.embedding import EmbeddingModel
 from aimakerspace.vectordatabase import VectorDatabase
 from aimakerspace.openai_utils.chatmodel import ChatOpenAI
 import chainlit as cl
-import fitz
-
+from qdrant_client import QdrantClient
+from qdrant_client.models import VectorParams, Distance
+from chainlit.input_widget import Select
+from qdrant_client.models import PointStruct
+#Qdrant client
+client = None
+
+#System Chat Prompt
 system_template = """\
 Use the following context to answer a users question. If you cannot find the answer in the context, say you don't know the answer."""
 system_role_prompt = SystemRolePrompt(system_template)
 
+#User Prompt for chat
 user_prompt_template = """\
 Context:
 {context}
@@ -26,6 +34,16 @@ Question:
 """
 user_role_prompt = UserRolePrompt(user_prompt_template)
 
+#Categorization of VectorDatabase
+system_template_db = """\
+You are an expert in categorization. Given the last user response determine if he or she wants to use the Qdrant database. If yes return the output single word 'QDrant' without any other phrases. If no return the only the word 'AI Makerspace'.
+
+"""
+system_role_prompt_db = SystemRolePrompt(system_template_db)
+
+user_prompt_template_db = "User Input:\n{user_input}"
+user_role_prompt_db = UserRolePrompt(user_prompt_template)
+
 class RetrievalAugmentedQAPipeline:
     def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None:
         self.llm = llm
@@ -48,6 +66,34 @@ class RetrievalAugmentedQAPipeline:
 
         return {"response": generate_response(), "context": context_list}
 
+class RetrievalAugmentedQAPipelineQdrant:
+    def __init__(self, llm: ChatOpenAI(), vector_db_retriever) -> None:
+        self.llm = llm
+        self.vector_db_retriever = vector_db_retriever
+        self.embedding_model = EmbeddingModel()
+
+    async def arun_pipeline(self, user_query: str):
+        query_vector = self.embedding_model.get_embedding(user_query)
+        context_list = self.vector_db_retriever.search(
+            collection_name="my_collection",
+            query_vector=query_vector, 
+            limit=4
+        )
+
+        context_prompt = ""
+        for context in context_list:
+            context_prompt += context.payload['text'] + "\n"
+
+        formatted_system_prompt = system_role_prompt.create_message()
+
+        formatted_user_prompt = user_role_prompt.create_message(question=user_query, context=context_prompt)
+
+        async def generate_response():
+            async for chunk in self.llm.astream([formatted_system_prompt, formatted_user_prompt]):
+                yield chunk
+
+        return {"response": generate_response(), "context": context_list}
+
 text_splitter = CharacterTextSplitter()
 
 
@@ -79,10 +125,42 @@ def process_pdf_file(file: AskFileResponse):
     texts = text_splitter.split_texts(documents)
     return texts
 
-
+async def initialize_qdrant(text):
+    client = QdrantClient(":memory:")
+    if not client.collection_exists("my_collection"):
+        client.create_collection(
+            collection_name="my_collection",
+            vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
+        )
+    embedding_model = EmbeddingModel()
+    embeddings = await embedding_model.async_get_embeddings(text)
+    i = 0
+    for text, embedding in zip(text, embeddings):
+        insert(text, np.array(embedding), i, client)
+        i+=1
+   
+    return client
+
+def insert(text, vector, idx, client):
+    point= PointStruct(
+                id=idx,
+                vector=vector.tolist(),
+                payload={"text": text}
+    )
+    client.upsert(
+        collection_name="my_collection",
+        points=[point]
+    )
+    
+    
+def choose_db(llm, user_input):
+    formatted_system_prompt_db = system_role_prompt_db.create_message()
+    formatted_user_prompt_db = user_role_prompt_db.create_message(question=user_input)
+    return llm.run([formatted_system_prompt_db, formatted_user_prompt_db])
 
 @cl.on_chat_start
 async def on_chat_start():
+    global client
     files = None
 
     # Wait for the user to upload a file
@@ -110,18 +188,32 @@ async def on_chat_start():
         texts = process_text_file(file)
 
     print(f"Processing {len(texts)} text chunks")
-
-    # Create a dict vector store
-    vector_db = VectorDatabase()
-    vector_db = await vector_db.abuild_from_list(texts)
     
     chat_openai = ChatOpenAI()
 
+    res = await cl.AskUserMessage(content="Do you want to use the QDrant vector database or AI Makerspace's?").send()
+    
+    if res:
+        chosen_db = choose_db(chat_openai, res['content'])
+        await cl.Message(
+            content=f"You have chosen {chosen_db}. Please start asking questions!",
+        ).send()
+
     # Create a chain
-    retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
-        vector_db_retriever=vector_db,
-        llm=chat_openai
-    )
+    retrieval_augmented_qa_pipeline = None
+    if chosen_db.lower() == 'qdrant': 
+       client = await initialize_qdrant(texts)
+       retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipelineQdrant(
+            vector_db_retriever=client,
+            llm=chat_openai
+        )
+    else:  
+        vector_db = VectorDatabase()
+        vector_db = await vector_db.abuild_from_list(texts)
+        retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
+            vector_db_retriever=vector_db,
+            llm=chat_openai
+        )
     
     # Let the user know that the system is ready
     msg.content = f"Processing `{file.name}` done. You can now ask questions!"

requirements.txt

@@ -1,4 +1,5 @@
 numpy
 chainlit==0.7.700
 openai
-pymupdf==1.24.9
+pymupdf==1.24.9
+qdrant-client==1.11.0