Update app.py

app.py

@@ -11,16 +11,12 @@ from qdrant_client.http.models import Distance, VectorParams
 from operator import itemgetter
 import chainlit as cl
 
-# # Load the documents
-# pdf_loader_NIST = PyMuPDFLoader("data/NIST.AI.600-1.pdf").load()
-# pdf_loader_Blueprint = PyMuPDFLoader("data/Blueprint-for-an-AI-Bill-of-Rights.pdf").load()
-# documents = pdf_loader_NIST + pdf_loader_Blueprint
 
-# List to store all the documents
+
+# Load all the documents in the directory
 documents = []
 directory = "data/"
 
-# Iterate through all the files in the directory
 for filename in os.listdir(directory):
     if filename.endswith(".pdf"):  # Check if the file is a PDF
         file_path = os.path.join(directory, filename)
@@ -37,36 +33,24 @@ text_splitter = RecursiveCharacterTextSplitter(
 )
 rag_documents = text_splitter.split_documents(documents)
 
-# 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
-
+embedding = OpenAIEmbeddings(model="text-embedding-3-small")
 
-    embeddings = OpenAIEmbeddings()
-    qdrant_client = QdrantClient(location=LOCATION)
+# Create the vector store
+vectorstore = Qdrant.from_documents(
+    rag_documents,
+    embedding,
+    location=":memory:",
+    collection_name="Implications of AI",
+)
 
-    # Create the collection
-    qdrant_client.create_collection(
-        collection_name=COLLECTION_NAME,
-        vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
-    )
+retriever = vectorstore.as_retriever()
+llm = ChatOpenAI(model="gpt-4")
 
-    # 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.cache_resource
+@cl.on_chat_start
+async def start_chat():
 
-    
     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.
@@ -79,25 +63,23 @@ async def start_chat():
     """
 
     prompt = ChatPromptTemplate.from_template(template)
-    base_llm = ChatOpenAI(model_name="gpt-4", temperature=0)
 
-    retrieval_augmented_qa_chain = (
+    base_chain = (
         {"context": itemgetter("question") | retriever, "question": itemgetter("question")}
-        | RunnablePassthrough.assign(context=itemgetter("context"))
-        | {"response": prompt | base_llm, "context": itemgetter("context")}
+        | prompt | llm | StrOutputParser()
     )
 
-    cl.user_session.set("chain", retrieval_augmented_qa_chain)
+    cl.user_session.set("chain", base_chain)
 
 
 @cl.on_message
 async def main(message):
     chain = cl.user_session.get("chain")
+    result = await chain.invoke({"question":message.content})
 
-    msg = cl.Message(content="")
-    result = await chain.invoke(message.content)
+    msg = cl.Message(content=result)
 
-    async for stream_resp in result["response"]:
-        await msg.stream_token(stream_resp)
+    # async for stream_resp in result["response"]:
+    #     await msg.stream_token(stream_resp)
 
     await msg.send()