Update app.py

app.py

@@ -1,100 +1,80 @@
-from typing import List
-from aimakerspace.text_utils import CharacterTextSplitter, PDFFileLoader
-from aimakerspace.openai_utils.prompts import (
-    UserRolePrompt,
-    SystemRolePrompt
-)
-from aimakerspace.vectordatabase import VectorDatabase
-from aimakerspace.openai_utils.chatmodel import ChatOpenAI
-from langchain_community.embeddings import OpenAIEmbeddings
-from langchain_community.vectorstores import Chroma
+from langchain_community.document_loaders import PyMuPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from langchain_qdrant import QdrantVectorStore
+from langchain.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import Distance, VectorParams
+from operator import itemgetter
 import chainlit as cl
-import nest_asyncio
-nest_asyncio.apply()
-
-# # pdf_loader_NIST = PDFFileLoader("data/NIST.AI.600-1.pdf")
-# # pdf_loader_Blueprint = PDFFileLoader("data/Blueprint-for-an-AI-Bill-of-Rights.pdf")
-# # documents_NIST = pdf_loader_NIST.load_documents()
-# # documents_Blueprint = pdf_loader_Blueprint.load_documents()
-
-# text_splitter = CharacterTextSplitter()
-# # text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=40)
-# split_documents_NIST = text_splitter.split_texts(documents_NIST)
-# split_documents_Blueprint = text_splitter.split_texts(documents_Blueprint)
-loader = PDFFileLoader("data/")
-loader.load()
-splitter = CharacterTextSplitter()
-chunks = splitter.split_texts(loader.documents)
-
-# rag_documents = split_documents_NIST + split_documents_Blueprint
-
-RAG_PROMPT_TEMPLATE = """ \
-Use the provided context to answer the user's query.
-You may not answer the user's query unless there is specific context in the following text.
-If you do not know the answer, or cannot answer, please respond with "I don't know".
-"""
-
-rag_prompt = SystemRolePrompt(RAG_PROMPT_TEMPLATE)
-
-USER_PROMPT_TEMPLATE = """ \
-Context:
-{context}
-Question:
-{question}
-"""
-
-user_prompt = UserRolePrompt(USER_PROMPT_TEMPLATE)
-
-class RetrievalAugmentedQAPipeline:
-    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: Chroma) -> None:
-        self.llm = llm
-        self.vector_db_retriever = vector_db_retriever
-
-    async def arun_pipeline(self, question: str):
-        context_list = self.vector_db_retriever.search_by_text(question, k=4)
-
-        context_prompt = ""
-        for context in context_list:
-            context_prompt += context[0] + "\n"
 
-        formatted_system_prompt = rag_prompt.create_message()
+# Load the documents
+pdf_loader_NIST = PyMuPDFLoader(file_path="data/NIST.AI.600-1.pdf").load()
+pdf_loader_Blueprint = PyMuPDFLoader(file_path="data/Blueprint-for-an-AI-Bill-of-Rights.pdf").load()
+documents = pdf_loader_NIST + pdf_loader_Blueprint
+
+# Split the documents
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=500,
+    chunk_overlap=40,
+    length_function=len,
+    is_separator_regex=False
+)
+rag_documents = text_splitter.split_documents(documents)
 
-        formatted_user_prompt = user_prompt.create_message(user_query=user_query, context=context_prompt)
+# Create the vector store
+# @cl.cache_resource
+@cl.on_chat_start
+async def start_chat():
+    LOCATION = ":memory:"
+    COLLECTION_NAME = "Implications of AI"
+    VECTOR_SIZE = 1536
 
-        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}
+    embeddings = OpenAIEmbeddings()
+    qdrant_client = QdrantClient(location=LOCATION)
 
+    # Create the collection
+    qdrant_client.create_collection(
+        collection_name=COLLECTION_NAME,
+        vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
+    )
 
-# ------------------------------------------------------------
+    # Create the vector store
+    vectorstore = QdrantVectorStore(
+        client=qdrant_client,
+        collection_name=COLLECTION_NAME,
+        embedding=embeddings
+    )
 
+    # Load and add documents
+    vectorstore.add_documents(rag_documents)
+    retriever = vectorstore.as_retriever()
 
-@cl.on_chat_start
-async def start_chat():
-    settings = {
-        "model": "gpt-4o-mini"
-    }
-    cl.user_session.set("settings", settings)
-
-    # Create a vector store
-    # vector_db = VectorDatabase()
-    # vector_db = await vector_db.abuild_from_list(split_documents_NIST)
-    # vector_db = await vector_db.abuild_from_list(split_documents_Blueprint)
     
-    embeddings = OpenAIEmbeddings()
-    vector_db = Chroma.from_texts(chunks, embeddings)
-
-    # Create a chain
-    retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
-        vector_db_retriever=vector_db,
-        llm=chat_openai
+    template = """
+    Use the provided context to answer the user's query.
+    You may not answer the user's query unless there is specific context in the following text.
+    If you do not know the answer, or cannot answer, please respond with "I don't know".
+    Question:
+    {question}
+    Context:
+    {context}
+    Answer:
+    """
+
+    prompt = ChatPromptTemplate.from_template(template)
+    base_llm = ChatOpenAI(model_name="gpt-4", temperature=0)
+
+    retrieval_augmented_qa_chain = (
+        {"context": itemgetter("question") | retriever, "question": itemgetter("question")}
+        | RunnablePassthrough.assign(context=itemgetter("context"))
+        | {"response": prompt | base_llm, "context": itemgetter("context")}
     )
 
-    cl.user_session.set("chain", retrieval_augmented_qa_pipeline)
+    cl.user_session.set("chain", retrieval_augmented_qa_chain)
 
 
 @cl.on_message
@@ -102,7 +82,7 @@ async def main(message):
     chain = cl.user_session.get("chain")
 
     msg = cl.Message(content="")
-    result = await chain.arun_pipeline(message.content)
+    result = await chain.invoke(message.content)
 
     async for stream_resp in result["response"]:
         await msg.stream_token(stream_resp)