Update app.py

app.py

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