updated

RagPipeline.py

@@ -0,0 +1,41 @@
+from aimakerspace.openai_utils.prompts import (
+    UserRolePrompt,
+    SystemRolePrompt,
+    AssistantRolePrompt,
+)
+from aimakerspace.vectordatabase import VectorDatabase
+from aimakerspace.openai_utils.chatmodel import ChatOpenAI
+
+class RetrievalAugmentedQAPipeline:
+    def __init__(
+        self,
+        system_role_prompt: SystemRolePrompt,
+        user_role_prompt: UserRolePrompt,
+        llm: ChatOpenAI(),
+        vector_db_retriever: VectorDatabase,
+    ) -> None:
+        self.system_role_prompt = system_role_prompt
+        self.user_role_prompt = user_role_prompt
+        self.llm = llm
+        self.vector_db_retriever = vector_db_retriever
+
+    async def arun_pipeline(self, user_query: str):
+        context_list = self.vector_db_retriever.search_by_text(user_query, k=4)
+
+        context_prompt = ""
+        for context in context_list[0]:
+            context_prompt += context[0] + "\n"
+
+        formatted_system_prompt = self.system_role_prompt.create_message()
+
+        formatted_user_prompt = self.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}

aimakerspace/openai_utils/embedding.py

@@ -28,13 +28,6 @@ class EmbeddingModel:
 
         return [embeddings.embedding for embeddings in embedding_response.data]
 
-    async def async_get_embeddings_openai(self, list_of_text: List[str]) :
-        embedding_response = await self.async_client.embeddings.create(
-            input=list_of_text, model=self.embeddings_model_name, dimensions=self.dimensions
-        )
-
-        return embedding_response
-
     async def async_get_embedding(self, text: str) -> List[float]:
         embedding = await self.async_client.embeddings.create(
             input=text, model=self.embeddings_model_name, dimensions=self.dimensions

aimakerspace/text_utils.py

@@ -45,7 +45,7 @@ class PdfFileLoader:
     def load(self):
         if os.path.isdir(self.path):
             self.load_directory()
-        elif os.path.isfile(self.path) and self.path.endswith(".pdf"):
+        elif self.path.endswith(".pdf"):
             self.load_file()
         else:
             raise ValueError(

aimakerspace/vectordatabase.py

@@ -1,5 +1,6 @@
 from enum import Enum
 import numpy as np
+import uuid
 from collections import defaultdict
 from typing import List, Tuple, Callable
 from aimakerspace.openai_utils.embedding import EmbeddingModel
@@ -73,9 +74,31 @@ class VectorDatabase:
             self.vectors = defaultdict(np.array)
         if vector_db_options == VectorDatabaseOptions.QDRANT:
             self.qdrant_client = QdrantClient(":memory:")
+            vector_params = VectorParams(
+                size=embedding_model.dimensions,  # vector size 
+                distance="Cosine"  # distance metric
+            )
+            self.qdrant_client.recreate_collection(
+                collection_name=collection_name,
+                vectors_config={"default": vector_params},
+            )
 
     def insert(self, key: str, vector: np.array) -> None:
-        self.vectors[key] = vector
+        idx = str(uuid.uuid4())
+        payload = {"text": key}  
+
+        point = PointStruct(
+            id=idx,
+            vector={"default": vector.tolist()},
+            payload=payload
+        )
+        # Insert the vector into Qdrant with the associated document
+        self.qdrant_client.upsert(
+            collection_name=collection_name,
+            points=[point]
+        )
+        print(f"Inserted vector with ID {idx}: {vector}")
+        
 
     def search(
         self,
@@ -83,11 +106,18 @@ class VectorDatabase:
         k: int,
         distance_measure: Callable = cosine_similarity,
     ) -> List[Tuple[str, float]]:
-        scores = [
-            (key, distance_measure(query_vector, vector))
-            for key, vector in self.vectors.items()
-        ]
-        return sorted(scores, key=lambda x: x[1], reverse=True)[:k]
+        # if isinstance(query_vector, list):
+        #     query_vector = np.array(query_vector)
+        print(f"Searching in collection: {collection_name} with vector: {query_vector}")
+        collection_info = self.qdrant_client.get_collection(collection_name)
+        print(f"Collection info: {collection_info}")
+
+        search_results = self.qdrant_client.search(
+            collection_name=collection_name,
+            query_vector=query_vector,
+            limit=k
+        )
+        return [(result.payload['text'], result.score) for result in search_results]
 
     def search_by_text(
         self,
@@ -97,39 +127,17 @@ class VectorDatabase:
         return_as_text: bool = False,
     ) -> List[Tuple[str, float]]:
         query_vector = self.embedding_model.get_embedding(query_text)
-        if self.vector_db_options == VectorDatabaseOptions.DICTIONARY:
-            results = self.search(query_vector, k, distance_measure)
-            return [result[0] for result in results] if return_as_text else results
-        if self.vector_db_options == VectorDatabaseOptions.QDRANT:
-            search_result = self.qdrant_client.search(collection_name,query_vector=query_vector)
-            return [(point.payload["text"],point.score) for point in search_result]
+        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) -> np.array:
         return self.vectors.get(key, None)
 
     async def abuild_from_list(self, list_of_text: List[str]) -> "VectorDatabase":
-        if self.vector_db_options == VectorDatabaseOptions.DICTIONARY:
-            embeddings = await self.embedding_model.async_get_embeddings(list_of_text)
-            for text, embedding in zip(list_of_text, embeddings):
-                self.insert(text, np.array(embedding))
-        if self.vector_db_options == VectorDatabaseOptions.QDRANT:
-            embeddings_response = await self.embedding_model.async_get_embeddings_openai(list_of_text)
-            points = [
-                PointStruct(
-                    id=idx,
-                    vector=data.embedding,
-                    payload={"text": text},
-                )
-                for idx, (data, text) in enumerate(zip(embeddings_response.data, list_of_text))
-            ]
-            self.qdrant_client.create_collection(
-                collection_name,
-                vectors_config=VectorParams(
-                    size=self.embedding_model.dimensions,
-                    distance=Distance.COSINE,
-                ),
-            )
-            self.qdrant_client.upsert(collection_name, points)
+        embeddings = await self.embedding_model.async_get_embeddings(list_of_text)
+        for text, embedding in zip(list_of_text, embeddings):
+            self.insert(text, np.array(embedding))
         return self
 
 

app.py

@@ -1,5 +1,6 @@
 import os
 from openai import AsyncOpenAI
+from RagPipeline import RetrievalAugmentedQAPipeline
 from typing import List
 from chainlit.types import AskFileResponse
 from chainlit.cli import run_chainlit
@@ -57,25 +58,28 @@ def process_text_file(file: AskFileResponse) -> List[str]:
     import tempfile
 
     with tempfile.NamedTemporaryFile(
-        mode="w", delete=False, suffix=".txt"
+        mode="wb", delete=False, suffix=".txt"
     ) as temp_file:
         temp_file_path = temp_file.name
-
-    with open(file.path, "r", encoding="utf-8") as f:
-        text = f.read()
-
-    with open(temp_file_path, "w") as f:
-        f.write(text)
+        temp_file.write(file.content)
 
     text_loader = TextFileLoader(temp_file_path)
     documents = text_loader.load_documents()
     texts = []
     for doc in documents:
-        texts.append(text_splitter.split_text(doc))
+        texts += text_splitter.split_text(doc)
     return texts
 
 def process_pdf_file(file: AskFileResponse) -> List[str]:
-    pdf_loader = PdfFileLoader(file.path)
+    import tempfile
+    with tempfile.NamedTemporaryFile(
+        mode="wb", delete=False, suffix=".pdf"
+    ) as temp_file:
+        temp_file_path = temp_file.name
+        temp_file.write(file.content)
+
+
+    pdf_loader = PdfFileLoader(temp_file_path)
     texts = pdf_loader.load_documents() # Also handles splitting the text in this case pages
     return texts
 
@@ -108,9 +112,9 @@ async def on_chat_start():
     texts : List[str] = []
     for file in files:
         if file.type == "application/pdf":
-            texts.extend(process_pdf_file(file))
+            texts += process_pdf_file(file)
         if file.type == "text/plain":
-            texts.extend(process_text_file(file))
+            texts += process_text_file(file)
 
     # await send_new_message(content=f"Processing `{file.name}`...")
     msg = cl.Message(content=f"Processing `{file.name}`...")

requirements copy.txt

@@ -0,0 +1,7 @@
+numpy
+chainlit==0.7.700
+openai
+langchain-text-splitters
+pypdf
+langchain-community
+qdrant-client

requirements.txt

@@ -1,6 +1,6 @@
 numpy==1.26.4
 chainlit==0.7.700   # 1.1.402 
-openai==1.3.5
+openai
 qdrant-client==1.11.0
 langchain-text-splitters
 langchain-community