Update app.py

app.py

@@ -1,44 +1,95 @@
 import os
 import chainlit as cl
-from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_groq import ChatGroq
 from langchain_community.document_loaders import PyPDFLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.chains.question_answering import load_qa_chain
-from langchain_community.llms import HuggingFaceHub
-from dotenv import load_dotenv
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.prompts import PromptTemplate
+from langchain.llms import HuggingFaceEndpoint
+from langchain.chains import RetrievalQA
 
-load_dotenv()
+# Global variable to store the QA chain
+qa_chain = None
 
-@cl.on_message
-async def main(message: cl.Message):
-    # Load PDF
-    loader = PyPDFLoader("document.pdf")
-    pages = loader.load_and_split()
+# Load and process the PDF file
+def process_file(file_path):
+    loader = PyPDFLoader(file_path)
+    documents = loader.load()
 
-    # Split text
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    documents = text_splitter.split_documents(pages)
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    texts = text_splitter.split_documents(documents)
 
-    # Embeddings & Vector DB
-    embeddings = HuggingFaceEmbeddings()
-    vectorstore = FAISS.from_documents(documents, embeddings)
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    vectorstore = FAISS.from_documents(texts, embeddings)
 
-    # Get relevant docs
-    query = message.content
-    matching_docs = vectorstore.similarity_search(query)
+    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.
+
+    {context}
+
+    Question: {question}
+
+    Helpful Answer:
+    """
+    prompt = PromptTemplate(
+        template=custom_prompt_template,
+        input_variables=["context", "question"]
+    )
+
+    llm = HuggingFaceEndpoint(
+        repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
+        temperature=0.5,
+        max_new_tokens=500,
+        top_k=50,
+        huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+    )
 
-    # LLM setup
-    llm = ChatGroq(
-        api_key=os.getenv("GROQ_API_KEY"),
-        model_name="LLaMA3-8b-8192"
+    qa = RetrievalQA.from_chain_type(
+        llm=llm,
+        chain_type="stuff",
+        retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
+        chain_type_kwargs={"prompt": prompt},
+        return_source_documents=True
     )
 
-    # QA chain
-    chain = load_qa_chain(llm, chain_type="stuff")
-    answer = chain.run(input_documents=matching_docs, question=query)
+    return qa
+
+# Chainlit on_chat_start event
+@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()
 
-    await cl.Message(content=answer).send()
+# Chainlit file upload handler
+@cl.on_file_upload
+def on_file_upload(file):
+    global qa_chain
+    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
+@cl.on_message
+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()
+        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()
+        return
+
+    query = message.content
+    response = qa_chain(query)
+    answer = response["result"]
+    cl.Message(content=answer).send()