Update app.py

app.py

@@ -2,16 +2,13 @@ import os
 import chainlit as cl
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import FAISS
-from langchain.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.prompts import PromptTemplate
-from langchain.llms import HuggingFaceEndpoint
+from langchain_community.llms import HuggingFaceEndpoint
 from langchain.chains import RetrievalQA
 
-# Global variable to store the QA chain
-qa_chain = None
-
-# Load and process the PDF file
+# Helper function to process PDF
 def process_file(file_path):
     loader = PyPDFLoader(file_path)
     documents = loader.load()
@@ -22,19 +19,13 @@ def process_file(file_path):
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     vectorstore = FAISS.from_documents(texts, embeddings)
 
-    custom_prompt_template = """
-    Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
-
+    prompt_template = """
+    Use the following pieces of context to answer the question at the end. If you don't know the answer, just say you don't know.
     {context}
-
     Question: {question}
-
     Helpful Answer:
     """
-    prompt = PromptTemplate(
-        template=custom_prompt_template,
-        input_variables=["context", "question"]
-    )
+    prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)
 
     llm = HuggingFaceEndpoint(
         repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
@@ -44,7 +35,7 @@ def process_file(file_path):
         huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
     )
 
-    qa = RetrievalQA.from_chain_type(
+    qa_chain = RetrievalQA.from_chain_type(
         llm=llm,
         chain_type="stuff",
         retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
@@ -52,44 +43,46 @@ def process_file(file_path):
         return_source_documents=True
     )
 
-    return qa
+    return qa_chain
 
-# Chainlit on_chat_start event
+# Chat start
 @cl.on_chat_start
-def start():
-    global qa_chain
-    qa_chain = None
-    cl.user_session.set("ready", False)
-    cl.user_session.set("file_name", "")
-    cl.user_session.set("qa_chain", None)
-    cl.Message(content="👋 Welcome! Upload a PDF file to get started.").send()
-
-# Chainlit file upload handler
-@cl.on_file_upload
-def on_file_upload(file):
-    global qa_chain
+async def start():
+    await cl.Message(content="👋 Welcome! Please upload a PDF file to begin.").send()
+    files = await cl.AskFileMessage(
+        content="📄 Upload a PDF to get started.",
+        accept=["application/pdf"],
+        max_size_mb=10,
+        timeout=180
+    ).send()
+
+    if not files:
+        await cl.Message(content="⚠️ No file was uploaded.").send()
+        return
+
+    file = files[0]
     file_path = file.path
+
     qa_chain = process_file(file_path)
     cl.user_session.set("qa_chain", qa_chain)
-    cl.user_session.set("file_name", file.name)
     cl.user_session.set("ready", True)
-    cl.Message(content=f"✅ File `{file.name}` uploaded and processed. You can now ask questions.").send()
 
-# Chainlit message handler
+    await cl.Message(content=f"✅ `{file.name}` processed successfully. You can now ask questions.").send()
+
+# Handle questions
 @cl.on_message
-def handle_message(message: cl.Message):
+async def handle_message(message: cl.Message):
     if not cl.user_session.get("ready"):
-        cl.Message(content="⚠️ Please upload a PDF file before asking questions.").send()
+        await cl.Message(content="⚠️ Please upload a PDF file first.").send()
         return
 
     qa_chain = cl.user_session.get("qa_chain")
-    if qa_chain is None:
-        cl.Message(content="⚠️ Something went wrong. Please upload the PDF again.").send()
+    if not qa_chain:
+        await cl.Message(content="⚠️ Error: QA chain not initialized.").send()
         return
 
     query = message.content
     response = qa_chain(query)
-    answer = response["result"]
-    cl.Message(content=answer).send()
+    await cl.Message(content=response["result"]).send()