Update app.py

app.py

@@ -1,7 +1,4 @@
 ### Import Section ###
-"""
-IMPORTS HERE
-"""
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import PyMuPDFLoader
 from qdrant_client import QdrantClient
@@ -23,9 +20,6 @@ import chainlit as cl
 
 
 ### Global Section ###
-"""
-GLOBAL CODE HERE
-"""
 text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
 Loader = PyMuPDFLoader
 
@@ -56,10 +50,10 @@ chat_model = ChatOpenAI(model="gpt-4o-mini")
 set_llm_cache(InMemoryCache())
 chat_openai = ChatOpenAI()
 
+
 ### On Chat Start (Session Start) Section ###
 @cl.on_chat_start
 async def on_chat_start():
-    """ SESSION SPECIFIC CODE HERE """
     files = None
 
     # Wait for the user to upload a file
@@ -111,21 +105,21 @@ async def on_chat_start():
     retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 3})
     
     # Create a chain
-    retrieval_chain = RetrievalAugmentedQAPipeline(
-        vector_db_retriever=vector_db,
-        llm=chat_openai
+    retrieval_augmented_qa_chain = (
+        {"context": itemgetter("question") | retriever, "question": itemgetter("question")}
+        | RunnablePassthrough.assign(context=itemgetter("context"))
+        | chat_prompt | chat_model
     )
     
     # Let the user know that the system is ready
     msg.content = f"Processing `{file.name}` done. You can now ask questions!"
     await msg.update()
-    cl.user_session.set("chain", retrieval_chain)
+    cl.user_session.set("chain", retrieval_augmented_qa_chain)
 
 
 ### Rename Chains ###
 @cl.author_rename
 def rename(orig_author: str):
-    """ RENAME CODE HERE """
     rename_dict = {"LLMMathChain": "Albert Einstein", "Chatbot": "Assistant"}
     return rename_dict.get(orig_author, orig_author)
 
@@ -133,9 +127,6 @@ def rename(orig_author: str):
 ### On Message Section ###
 @cl.on_message
 async def main(message: cl.Message):
-    """
-    MESSAGE CODE HERE
-    """
     chain = cl.user_session.get("chain")
 
     msg = cl.Message(content="")