Update app.py

app.py

@@ -3,20 +3,21 @@ 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
 
-# ========== Config ==========
+# Load environment variables
 load_dotenv()
 api_key = os.getenv("GOOGLE_API_KEY")
 genai.configure(api_key=api_key)
 
-# ========== PDF to text (fa+en) ==========
+
+# ✅ Function to read all PDF files (Farsi + English)
 def get_pdf_text(pdf_docs):
     text = ""
     for pdf in pdf_docs:
@@ -27,12 +28,15 @@ def get_pdf_text(pdf_docs):
                     text += page_text + "\n"
     return text
 
-# ========== Chunking ==========
+
+# ✅ Function to split text into chunks
 def get_text_chunks(text):
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    return splitter.split_text(text)
+    splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=300)
+    chunks = splitter.split_text(text)
+    return chunks
+
 
-# ========== Build & save vector store ==========
+# ✅ Function to get embeddings for each chunk and save to vector store
 def get_vector_store(chunks):
     try:
         embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
@@ -41,124 +45,65 @@ def get_vector_store(chunks):
     except Exception as e:
         raise RuntimeError(f"Error creating vector store: {e}")
 
-# ========== Map-Reduce QA chain (بدون load_qa_chain) ==========
-from langchain.prompts import PromptTemplate
-from langchain.chains import LLMChain
-from langchain.chains import MapReduceDocumentsChain, ReduceDocumentsChain, StuffDocumentsChain
-from langchain_google_genai import ChatGoogleGenerativeAI
 
+# ✅ Function to get conversational chain
 def get_conversational_chain():
-    # مرحله 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"],
-    )
-
-    # مرحله Combine: پاسخ نهایی فقط بر اساس خلاصه‌ها
-    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"],
-    )
-
-    # مرحله Collapse: اگر خلاصه‌ها زیاد شد، کوتاه‌شان کن (برای رد شدن از سقف توکن)
-    collapse_prompt_tmpl = PromptTemplate(
-        template=(
-            "Condense the following summaries into a shorter consolidated summary, keeping only "
-            "information relevant to the question.\n\n"
-            "Question:\n{question}\n\n"
-            "Summaries:\n{summaries}\n\n"
-            "Shorter summaries:"
-        ),
-        input_variables=["summaries", "question"],
-    )
-
-    model = ChatGoogleGenerativeAI(model="gemini-1.5-flash", client=genai, temperature=0.3)
-
-    # LLMChain ها
-    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
-    collapse_llm_chain = LLMChain(llm=model, prompt=collapse_prompt_tmpl) # expects: summaries, question
+    prompt_template = """
+    You are a helpful assistant. 
+    Answer the question as detailed as possible using the provided context. 
+    If the answer is unclear, summarize the most relevant part of the context. 
+    Do not make up answers.
 
-    # زنجیره‌های ترکیب و فشرده‌سازی
-    combine_documents_chain = StuffDocumentsChain(
-        llm_chain=combine_llm_chain,
-        document_variable_name="summaries"
-    )
-    collapse_documents_chain = StuffDocumentsChain(
-        llm_chain=collapse_llm_chain,
-        document_variable_name="summaries"
-    )
+    Context:
+    {context}
 
-    # Reduce با عدد صحیح برای token_max (مثلاً 8000)
-    reduce_chain = ReduceDocumentsChain(
-        combine_documents_chain=combine_documents_chain,
-        collapse_documents_chain=collapse_documents_chain,
-        token_max=8000   # ✅ عدد صحیح؛ می‌تونی بر اساس نیازت کم/زیادش کنی
-    )
+    Question:
+    {question}
 
-    # زنجیره‌ی Map → Reduce
-    chain = MapReduceDocumentsChain(
-        llm_chain=map_llm_chain,
-        reduce_documents_chain=reduce_chain,
-        document_variable_name="context",   # باید با ورودی map_prompt یکی باشد
-        return_intermediate_steps=False,
-    )
-    return chain
+    Answer:
+    """
+    try:
+        model = ChatGoogleGenerativeAI(model="gemini-1.5-flash", client=genai, temperature=0.3)
+        prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
+        chain = load_qa_chain(llm=model, chain_type="stuff", prompt=prompt)
+        return chain
+    except Exception as e:
+        raise RuntimeError(f"Error creating conversational chain: {e}")
 
 
-# ========== Chat history utils ==========
+# ✅ Function to clear chat history
 def clear_chat_history():
     st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
 
-# ========== Ask ==========
+
+# ✅ Function to handle user input
 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)
-
-        # برای متن‌های بلند recall را افزایش بده
-        docs = new_db.similarity_search(user_question, k=12)
-
-        # اجرای زنجیره
+        docs = new_db.similarity_search(user_question, k=15)
         chain = get_conversational_chain()
-        response = chain.invoke({"input_documents": docs, "question": user_question})
-        # MapReduceDocumentsChain خروجی را در output_text می‌دهد
+        response = chain({"input_documents": docs, "question": user_question}, return_only_outputs=True)
         return response
     except Exception as e:
         raise RuntimeError(f"Error while answering: {e}")
 
-# ========== UI ==========
+
+# ✅ Main function to run the Streamlit app
 def main():
-    st.set_page_config(page_title="Chatbot", layout="wide", initial_sidebar_state="expanded")
-
-    st.markdown(
-        """
-        <style>
-        body { background-color: #000000; color: #ffffff; }
-        .main { background-color: #333333; padding: 20px; border-radius: 10px; }
-        </style>
-        """,
-        unsafe_allow_html=True
+    st.set_page_config(
+        page_title="Chatbot",
+        layout="wide",
+        initial_sidebar_state="expanded"
     )
 
     if "uploaded" not in st.session_state:
         st.session_state.uploaded = False
 
     if not st.session_state.uploaded:
+        # Upload Page
         st.title("Your personal assistant ...")
-        pdf_docs = st.file_uploader("فایل پی‌دی‌اف(ها) را آپلود کنید", type="pdf", accept_multiple_files=True)
+        pdf_docs = st.file_uploader("فایل پی دی اف مورد نظر را آپلود کنید ", accept_multiple_files=True)
         if st.button("تایید"):
             if pdf_docs:
                 try:
@@ -171,8 +116,9 @@ def main():
                 except RuntimeError as e:
                     st.error(str(e))
             else:
-                st.error("لطفاً حداقل یک فایل را آپلود کنید")
+                st.error("لطفا حداقل یک فایل را آپلود کنید")
     else:
+        # Chat Page
         st.title("Assistant ready ...")
         st.write("میتونین سوالتونو بپرسین 👇")
 
@@ -181,41 +127,53 @@ 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"):
-                        result = user_input(prompt)
-                        if result:
-                            full_response = result.get("output_text", "")
+                        response = user_input(prompt)
+                        if response:
+                            full_response = response['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()