Update app.py

app.py

@@ -1,64 +1,85 @@
 import os
 import chainlit as cl
-from langchain_community.vectorstores import FAISS
 from langchain_community.document_loaders import PyPDFLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.prompts import PromptTemplate
+from langchain_community.llms import HuggingFaceEndpoint
 from langchain.chains import RetrievalQA
-from langchain.chat_models import ChatOpenAI
-
-# Initialize embedding model
-embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
 
-# Use this variable to store retriever globally
-retriever = None
+# PDF processing
+def process_file(file_path):
+    loader = PyPDFLoader(file_path)
+    documents = loader.load()
 
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    texts = text_splitter.split_documents(documents)
 
-@cl.on_message
-async def handle_message(message: cl.Message):
-    global retriever
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    vectorstore = FAISS.from_documents(texts, embeddings)
 
-    if not retriever:
-        await cl.Message(content="📂 Please upload a PDF file first.").send()
-        return
+    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 you don't know.
 
-    # Setup QA chain
-    qa_chain = RetrievalQA.from_chain_type(
-        llm=ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.2),
-        chain_type="stuff",
-        retriever=retriever,
-        return_source_documents=True,
+    {context}
+    Question: {question}
+    Helpful Answer:
+    """
+    prompt = PromptTemplate(
+        template=custom_prompt_template,
+        input_variables=["context", "question"]
     )
 
-    # Ask question
-    response = qa_chain(message.content)
-
-    # Respond in UI
-    await cl.Message(content=response["result"]).send()
-
+    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")
+    )
 
-@cl.on_file_upload
-async def on_file_upload(files: list[cl.UploadedFile]):
-    global retriever
+    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
+    )
 
-    # Handle only the first file
-    uploaded_file = files[0]
+    return qa
 
-    # Save file locally
-    file_path = f"/tmp/{uploaded_file.name}"
-    with open(file_path, "wb") as f:
-        f.write(uploaded_file.content)
+# Chat start
+@cl.on_chat_start
+async def start():
+    await cl.Message("👋 Welcome! Please upload a PDF to get started.").send()
 
-    # Load and split PDF
-    loader = PyPDFLoader(file_path)
-    docs = loader.load()
-    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    chunks = splitter.split_documents(docs)
+    # Let user upload file
+    msg = await cl.AskFileMessage(
+        content="📄 Upload your PDF file here:",
+        accept=["application/pdf"],
+        max_size_mb=20,
+        timeout=180
+    ).send()
 
-    # Create FAISS index
-    db = FAISS.from_documents(chunks, embedding_model)
-    retriever = db.as_retriever()
+    if msg and msg.files:
+        file = msg.files[0]
+        qa_chain = process_file(file.path)
 
-    await cl.Message(content="✅ PDF uploaded and processed! You can now ask questions.").send()
+        # Save QA chain to session
+        cl.user_session.set("qa_chain", qa_chain)
+        cl.user_session.set("ready", True)
+        await cl.Message(f"✅ File `{file.name}` processed. You can now ask your questions.").send()
 
+# On user message
+@cl.on_message
+async def handle_message(message: cl.Message):
+    if not cl.user_session.get("ready"):
+        await cl.Message("⚠️ Please upload a PDF first.").send()
+        return
 
+    qa_chain = cl.user_session.get("qa_chain")
+    response = qa_chain(message.content)
+    answer = response["result"]
+    await cl.Message(content=answer).send()