Update app.py

app.py

@@ -42,8 +42,13 @@ def get_vector_store(chunks):
         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
+
 def get_conversational_chain():
-    # Map: برای هر تکه، خلاصه/نکات مرتبط با سؤال را بیرون بکش
+    # مرحله Map: خلاصه یا نِکات مرتبط از هر تکه
     map_prompt_tmpl = PromptTemplate(
         template=(
             "You are a helpful assistant. Extract a concise summary strictly from the context "
@@ -54,7 +59,8 @@ def get_conversational_chain():
         ),
         input_variables=["context", "question"],
     )
-    # Reduce/Combine: پاسخ نهایی با اتکا به خلاصه‌های مرحله Map
+
+    # مرحله Combine: پاسخ نهایی فقط بر اساس خلاصه‌ها
     combine_prompt_tmpl = PromptTemplate(
         template=(
             "You are a helpful assistant. Answer the question using ONLY the provided summaries. "
@@ -66,34 +72,52 @@ def get_conversational_chain():
         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-2.5-pro", client=genai, temperature=0.3)
 
-    # 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
+    # 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
 
-    # StuffDocumentsChain برای ترکیب خلاصه‌ها
+    # زنجیره‌های ترکیب و فشرده‌سازی
     combine_documents_chain = StuffDocumentsChain(
         llm_chain=combine_llm_chain,
-        document_variable_name="summaries"  # باید با ورودی پرامپت combine یکی باشد
+        document_variable_name="summaries"
+    )
+    collapse_documents_chain = StuffDocumentsChain(
+        llm_chain=collapse_llm_chain,
+        document_variable_name="summaries"
     )
 
-    # Reduce chain
+    # Reduce با عدد صحیح برای token_max (مثلاً 8000)
     reduce_chain = ReduceDocumentsChain(
         combine_documents_chain=combine_documents_chain,
-        collapse_documents_chain=None,  # می‌تونی در صورت خیلی زیاد بودن خلاصه‌ها یک collapse هم اضافه کنی
-        token_max=None
+        collapse_documents_chain=collapse_documents_chain,
+        token_max=8000   # ✅ عدد صحیح؛ می‌تونی بر اساس نیازت کم/زیادش کنی
     )
 
-    # Map-Reduce chain اصلی
+    # زنجیره‌ی Map → Reduce
     chain = MapReduceDocumentsChain(
-        llm_chain=map_llm_chain,                   # expects: context, question
+        llm_chain=map_llm_chain,
         reduce_documents_chain=reduce_chain,
-        document_variable_name="context",          # باید با ورودی پرامپت map یکی باشد
+        document_variable_name="context",   # باید با ورودی map_prompt یکی باشد
         return_intermediate_steps=False,
     )
     return chain
 
+
 # ========== Chat history utils ==========
 def clear_chat_history():
     st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]