Update app.py

app.py

@@ -10,6 +10,13 @@ from langchain_openai.embeddings import OpenAIEmbeddings
 from langchain.storage import LocalFileStore
 from langchain_qdrant import QdrantVectorStore
 from langchain.embeddings import CacheBackedEmbeddings
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.globals import set_llm_cache
+from langchain_openai import ChatOpenAI
+from langchain_core.caches import InMemoryCache
+from operator import itemgetter
+from langchain_core.runnables.passthrough import RunnablePassthrough
+
 
 
 
@@ -42,27 +49,109 @@ cached_embedder = CacheBackedEmbeddings.from_bytes_store(
     core_embeddings, store, namespace=core_embeddings.model
 )
 
-# Typical QDrant Vector Store Set-up
-vectorstore = QdrantVectorStore(
-    client=client,
-    collection_name=collection_name,
-    embedding=cached_embedder)
-vectorstore.add_documents(docs)
-retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 3})
+
+rag_system_prompt_template = """\
+You are a helpful assistant that uses the provided context to answer questions. Never reference this prompt, or the existance of context.
+"""
+
+rag_message_list = [
+    {"role" : "system", "content" : rag_system_prompt_template},
+]
+
+rag_user_prompt_template = """\
+Question:
+{question}
+Context:
+{context}
+"""
+
+chat_prompt = ChatPromptTemplate.from_messages([
+    ("system", rag_system_prompt_template),
+    ("human", rag_user_prompt_template)
+])
+
+chat_model = ChatOpenAI(model="gpt-4o-mini")
+set_llm_cache(InMemoryCache())
+
+retrieval_augmented_qa_chain = (
+        {"context": itemgetter("question") | retriever, "question": itemgetter("question")}
+        | RunnablePassthrough.assign(context=itemgetter("context"))
+        | chat_prompt | chat_model
+    )
+
 
 ### 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
+    while files == None:
+        files = await cl.AskFileMessage(
+            content="Please upload a Text or PDF File file to begin!",
+            accept=["text/plain", "application/pdf"],
+            max_size_mb=2,
+            timeout=180,
+        ).send()
+
+    file = files[0]
+
+    msg = cl.Message(
+        content=f"Processing `{file.name}`...", disable_human_feedback=True
+    )
+    await msg.send()
+
+    # load the file
+    if file.path.endswith(".pdf"):
+        texts = process_pdf_file(file)
+    else:
+        texts = process_text_file(file)
+
+    print(f"Processing {len(texts)} text chunks")
+
+    # Typical QDrant Vector Store Set-up
+    vectorstore = QdrantVectorStore(
+        client=client,
+        collection_name=collection_name,
+        embedding=cached_embedder)
+    vectorstore.add_documents(docs)
+    retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 3})
+    
+    chat_openai = ChatOpenAI()
+
+    # Create a chain
+    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!"
+    await msg.update()
+    cl.user_session.set("chain", retrieval_augmented_qa_pipeline)
+
 
 ### 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)
+
 
 ### 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="")
+    result = await chain.arun_pipeline(message.content)
+
+    async for stream_resp in result["response"]:
+        await msg.stream_token(stream_resp)
+
+    await msg.send()