Update app.py

app.py

@@ -1,80 +1,43 @@
 import os
-import chainlit as cl
-from langchain.document_loaders import PyPDFLoader
-from langchain.embeddings import OpenAIEmbeddings
-from langchain.vectorstores import FAISS
-from langchain.chat_models import ChatOpenAI
-from langchain.chains import RetrievalQA
+os.environ["CHAINLIT_STORAGE_DIR"] = "/app/.chainlit"
 
-# Optional: Switch to Groq API
-USE_GROQ = False  # Set to True if using Groq
-
-if USE_GROQ:
-    import httpx
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.chains.question_answering import load_qa_chain
+from langchain.llms import OpenAI  # replace with Groq if needed
 
-def load_pdf(file_path):
-    loader = PyPDFLoader(file_path)
-    pages = loader.load_and_split()
-    return pages
+import chainlit as cl
 
 @cl.on_chat_start
 async def start():
-    files = await cl.AskFileMessage(
-        content="📄 Please upload a PDF file to begin.",
-        accept=["application/pdf"],
-        max_size_mb=10
-    ).send()
+    await cl.Message("👋 Upload a PDF to ask questions!").send()
+
+@cl.on_message
+async def handle_message(message: cl.Message):
+    files = await cl.AskFileMessage("📎 Upload your PDF", accept=["application/pdf"]).send()
+    if not files:
+        await cl.Message("❌ No file received.").send()
+        return
 
     file = files[0]
     file_path = file.path
 
-    pages = load_pdf(file_path)
-
-    # Create vector store
-    embeddings = OpenAIEmbeddings()
-    vectordb = FAISS.from_documents(pages, embeddings)
-
-    retriever = vectordb.as_retriever()
-
-    if USE_GROQ:
-        # Custom RAG using Groq API
-        cl.user_session.set("retriever", retriever)
-    else:
-        # Standard LangChain RetrievalQA with OpenAI
-        qa_chain = RetrievalQA.from_chain_type(
-            llm=ChatOpenAI(temperature=0.2),
-            retriever=retriever
-        )
-        cl.user_session.set("chain", qa_chain)
+    loader = PyPDFLoader(file_path)
+    documents = loader.load()
 
-@cl.on_message
-async def chat(msg: cl.Message):
-    if USE_GROQ:
-        retriever = cl.user_session.get("retriever")
-        docs = retriever.get_relevant_documents(msg.content)
-        context = "\n".join([doc.page_content for doc in docs][:3])
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    chunks = splitter.split_documents(documents)
 
-        headers = {
-            "Authorization": f"Bearer {os.getenv('GROQ_API_KEY')}",
-            "Content-Type": "application/json"
-        }
+    embeddings = HuggingFaceEmbeddings()
+    vectordb = FAISS.from_documents(chunks, embeddings)
 
-        body = {
-            "messages": [
-                {"role": "system", "content": f"You are a helpful assistant answering questions from a PDF. Use this context:\n{context}"},
-                {"role": "user", "content": msg.content}
-            ],
-            "model": "llama3-8b-8192",
-            "temperature": 0.5
-        }
+    qa_chain = load_qa_chain(OpenAI(temperature=0), chain_type="stuff")  # or Groq
 
-        async with httpx.AsyncClient() as client:
-            res = await client.post("https://api.groq.com/openai/v1/chat/completions", headers=headers, json=body)
+    query = message.content
+    docs = vectordb.similarity_search(query)
+    answer = qa_chain.run(input_documents=docs, question=query)
 
-        reply = res.json()["choices"][0]["message"]["content"]
-        await cl.Message(content=reply).send()
+    await cl.Message(answer).send()
 
-    else:
-        chain = cl.user_session.get("chain")
-        response = chain.run(msg.content)
-        await cl.Message(content=response).send()