Update app.py

app.py

@@ -3,22 +3,20 @@ import fitz  # PyMuPDF
 import streamlit as st
 import google.generativeai as genai
 from dotenv import load_dotenv
-from google.api_core.exceptions import GoogleAPIError, InvalidArgument
 
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
 from langchain_community.vectorstores import FAISS
-from langchain.chains.question_answering import load_qa_chain
 from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
+from langchain.chains import MapReduceDocumentsChain, ReduceDocumentsChain, StuffDocumentsChain
 
-
-# Load environment variables
+# ========== Config ==========
 load_dotenv()
 api_key = os.getenv("GOOGLE_API_KEY")
 genai.configure(api_key=api_key)
 
-
-# ✅ Function to read all PDF files (Farsi + English)
+# ========== PDF to text (fa+en) ==========
 def get_pdf_text(pdf_docs):
     text = ""
     for pdf in pdf_docs:
@@ -29,15 +27,12 @@ def get_pdf_text(pdf_docs):
                     text += page_text + "\n"
     return text
 
-
-# ✅ Function to split text into chunks
+# ========== Chunking ==========
 def get_text_chunks(text):
     splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    chunks = splitter.split_text(text)
-    return chunks
+    return splitter.split_text(text)
 
-
-# ✅ Function to get embeddings for each chunk and save to vector store
+# ========== Build & save vector store ==========
 def get_vector_store(chunks):
     try:
         embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
@@ -46,94 +41,89 @@ def get_vector_store(chunks):
     except Exception as e:
         raise RuntimeError(f"Error creating vector store: {e}")
 
-
-# ✅ Function to get conversational chain (map_reduce برای متن‌های بلند)
+# ========== Map-Reduce QA chain (بدون load_qa_chain) ==========
 def get_conversational_chain():
-    map_prompt = """
-    You are a helpful assistant. Summarize the following context to capture the key points 
-    that are relevant for answering the final question. 
-    Do NOT add extra info, just summarize faithfully.
-
-    Context:
-    {context}
-
-    Summary:
-    """
-
-    combine_prompt = """
-    You are a helpful assistant. Answer the question as detailed as possible 
-    using ONLY the provided summaries. 
-    If the answer is not in the summaries, say: "answer is not available in the context".
-    Do not make up answers.
+    # Map: برای هر تکه، خلاصه/نکات مرتبط با سؤال را بیرون بکش
+    map_prompt_tmpl = PromptTemplate(
+        template=(
+            "You are a helpful assistant. Extract a concise summary strictly from the context "
+            "that would help answer the final question. Do NOT add outside info.\n\n"
+            "Context:\n{context}\n\n"
+            "Question:\n{question}\n\n"
+            "Summary:"
+        ),
+        input_variables=["context", "question"],
+    )
+    # Reduce/Combine: پاسخ نهایی با اتکا به خلاصه‌های مرحله Map
+    combine_prompt_tmpl = PromptTemplate(
+        template=(
+            "You are a helpful assistant. Answer the question using ONLY the provided summaries. "
+            "If the answer is not in the summaries, say: \"answer is not available in the context\".\n\n"
+            "Summaries:\n{summaries}\n\n"
+            "Question:\n{question}\n\n"
+            "Answer:"
+        ),
+        input_variables=["summaries", "question"],
+    )
 
-    Summaries:
-    {summaries}
+    model = ChatGoogleGenerativeAI(model="gemini-2.5-pro", client=genai, temperature=0.3)
 
-    Question:
-    {question}
+    # LLM chains برای دو مرحله
+    map_llm_chain = LLMChain(llm=model, prompt=map_prompt_tmpl)                      # expects: context, question
+    combine_llm_chain = LLMChain(llm=model, prompt=combine_prompt_tmpl)              # expects: summaries, question
 
-    Answer:
-    """
+    # StuffDocumentsChain برای ترکیب خلاصه‌ها
+    combine_documents_chain = StuffDocumentsChain(
+        llm_chain=combine_llm_chain,
+        document_variable_name="summaries"  # باید با ورودی پرامپت combine یکی باشد
+    )
 
-    try:
-        model = ChatGoogleGenerativeAI(model="gemini-2.5-pro", client=genai, temperature=0.3)
-
-        map_prompt_template = PromptTemplate(template=map_prompt, input_variables=["context"])
-        combine_prompt_template = PromptTemplate(template=combine_prompt, input_variables=["summaries", "question"])
-
-        chain = load_qa_chain(
-            llm=model,
-            chain_type="map_reduce",
-            map_prompt=map_prompt_template,
-            combine_prompt=combine_prompt_template
-        )
-        return chain
-    except Exception as e:
-        raise RuntimeError(f"Error creating conversational chain: {e}")
+    # Reduce chain
+    reduce_chain = ReduceDocumentsChain(
+        combine_documents_chain=combine_documents_chain,
+        collapse_documents_chain=None,  # می‌تونی در صورت خیلی زیاد بودن خلاصه‌ها یک collapse هم اضافه کنی
+        token_max=None
+    )
 
+    # Map-Reduce chain اصلی
+    chain = MapReduceDocumentsChain(
+        llm_chain=map_llm_chain,                   # expects: context, question
+        reduce_documents_chain=reduce_chain,
+        document_variable_name="context",          # باید با ورودی پرامپت map یکی باشد
+        return_intermediate_steps=False,
+    )
+    return chain
 
-# ✅ Function to clear chat history
+# ========== Chat history utils ==========
 def clear_chat_history():
     st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
 
-
-# ✅ Function to handle user input
+# ========== Ask ==========
 def user_input(user_question):
     try:
         embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
         new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
 
-        # similarity_search with more results (better recall for long docs)
+        # برای متن‌های بلند recall را افزایش بده
         docs = new_db.similarity_search(user_question, k=12)
 
+        # اجرای زنجیره
         chain = get_conversational_chain()
-        response = chain({"input_documents": docs, "question": user_question}, return_only_outputs=True)
+        response = chain.invoke({"input_documents": docs, "question": user_question})
+        # MapReduceDocumentsChain خروجی را در output_text می‌دهد
         return response
     except Exception as e:
         raise RuntimeError(f"Error while answering: {e}")
 
-
-# ✅ Main function to run the Streamlit app
+# ========== UI ==========
 def main():
-    st.set_page_config(
-        page_title="Chatbot",
-        layout="wide",
-        initial_sidebar_state="expanded"
-    )
+    st.set_page_config(page_title="Chatbot", layout="wide", initial_sidebar_state="expanded")
 
-    # Dark theme styling
     st.markdown(
         """
         <style>
-        body {
-            background-color: #000000;
-            color: #ffffff;
-        }
-        .main {
-            background-color: #333333;
-            padding: 20px;
-            border-radius: 10px;
-        }
+        body { background-color: #000000; color: #ffffff; }
+        .main { background-color: #333333; padding: 20px; border-radius: 10px; }
         </style>
         """,
         unsafe_allow_html=True
@@ -143,9 +133,8 @@ def main():
         st.session_state.uploaded = False
 
     if not st.session_state.uploaded:
-        # Upload Page
         st.title("Your personal assistant ...")
-        pdf_docs = st.file_uploader("فایل پی دی اف مورد نظر را آپلود کنید ", accept_multiple_files=True)
+        pdf_docs = st.file_uploader("فایل پی‌دی‌اف(ها) را آپلود کنید", type="pdf", accept_multiple_files=True)
         if st.button("تایید"):
             if pdf_docs:
                 try:
@@ -158,9 +147,8 @@ def main():
                 except RuntimeError as e:
                     st.error(str(e))
             else:
-                st.error("لطفا حداقل یک فایل را آپلود کنید")
+                st.error("لطفاً حداقل یک فایل را آپلود کنید")
     else:
-        # Chat Page
         st.title("Assistant ready ...")
         st.write("میتونین سوالتونو بپرسین 👇")
 
@@ -169,53 +157,41 @@ def main():
             clear_chat_history()
             st.rerun()
 
-        st.button('حذف مکالمه', on_click=clear_chat_history)
+        st.button("حذف مکالمه", on_click=clear_chat_history)
 
         if "messages" not in st.session_state:
             st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
 
-        # Show history
+        # نمایش تاریخچه
         for message in st.session_state.messages:
             with st.chat_message(message["role"]):
                 st.markdown(
-                    f"""
-                    <div style="direction: rtl; text-align: right; font-size: 16px;">
-                        {message["content"]}
-                    </div>
-                    """,
+                    f"<div style='direction: rtl; text-align: right; font-size: 16px;'>{message['content']}</div>",
                     unsafe_allow_html=True
                 )
 
+        # ورودی چت
         if prompt := st.chat_input():
             st.session_state.messages.append({"role": "user", "content": prompt})
             with st.chat_message("user"):
                 st.markdown(
-                    f"""
-                    <div style="direction: rtl; text-align: right; font-size: 16px;">
-                        {prompt}
-                    </div>
-                    """,
+                    f"<div style='direction: rtl; text-align: right; font-size: 16px;'>{prompt}</div>",
                     unsafe_allow_html=True
                 )
 
             if st.session_state.messages[-1]["role"] != "assistant":
                 try:
                     with st.chat_message("assistant"):
-                        response = user_input(prompt)
-                        if response:
-                            full_response = response['output_text']
+                        result = user_input(prompt)
+                        if result:
+                            full_response = result.get("output_text", "")
                             st.markdown(
-                                f"""
-                                <div style="direction: rtl; text-align: right; font-size: 16px;">
-                                    {full_response}
-                                </div>
-                                """,
+                                f"<div style='direction: rtl; text-align: right; font-size: 16px;'>{full_response}</div>",
                                 unsafe_allow_html=True
                             )
                             st.session_state.messages.append({"role": "assistant", "content": full_response})
                 except RuntimeError as e:
                     st.error(str(e))
 
-
 if __name__ == "__main__":
     main()