Update app.py

app.py

@@ -4,14 +4,14 @@ from langchain_community.document_loaders import PyPDFLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain.chains import RetrievalQA
+
 from langchain_openai import ChatOpenAI
+from langchain_core.prompts import ChatPromptTemplate
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.chains import create_retrieval_chain
 
 
-def run_qa(pdf_path, question):
-    if pdf_path is None or question.strip() == "":
-        return "Please upload a PDF and enter a question."
-
+def build_chain(pdf_path: str):
     loader = PyPDFLoader(pdf_path)
     docs = loader.load()
 
@@ -20,20 +20,38 @@ def run_qa(pdf_path, question):
 
     embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
     vectordb = FAISS.from_documents(chunks, embeddings)
+    retriever = vectordb.as_retriever(search_kwargs={"k": 4})
 
     llm = ChatOpenAI(temperature=0)
 
-    qa = RetrievalQA.from_chain_type(
-        llm=llm,
-        retriever=vectordb.as_retriever(),
-        return_source_documents=True
+    prompt = ChatPromptTemplate.from_template(
+        """You are a helpful assistant. Answer the question using ONLY the provided context.
+If the answer is not in the context, say you don't know.
+
+<context>
+{context}
+</context>
+
+Question: {input}
+"""
     )
 
-    result = qa.invoke({"query": question})
+    doc_chain = create_stuff_documents_chain(llm, prompt)
+    retrieval_chain = create_retrieval_chain(retriever, doc_chain)
+    return retrieval_chain
+
+
+def run_qa(pdf_path, question):
+    if pdf_path is None or question.strip() == "":
+        return "Please upload a PDF and enter a question."
+
+    chain = build_chain(pdf_path)
+    result = chain.invoke({"input": question})
+
+    answer_text = result.get("answer", "")
+    ctx_docs = result.get("context", []) or []
 
-    answer_text = result.get("result", "")
-    source_docs = result.get("source_documents", [])
-    sources = "\n\n".join([d.page_content[:500] for d in source_docs[:2]])
+    sources = "\n\n".join([d.page_content[:500] for d in ctx_docs[:2]])
 
     return f"### Answer\n{answer_text}\n\n---\n### Sources\n{sources}"