Update app.py

app.py

@@ -1,177 +1,131 @@
-import os
-import fitz  # PyMuPDF
 import streamlit as st
+from PyPDF2 import PdfReader
+from langchain.prompts import PromptTemplate
+from langchain.chains import MapReduceDocumentsChain, ReduceDocumentsChain, StuffDocumentsChain
+from langchain.chains.question_answering import load_qa_chain
+from langchain_google_genai import ChatGoogleGenerativeAI
 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.embeddings import HuggingFaceEmbeddings
+import os
+from dotenv import load_dotenv
 
-# Load environment variables
+# --------------------
+# تنظیمات اولیه
+# --------------------
 load_dotenv()
-api_key = os.getenv("GOOGLE_API_KEY")
-genai.configure(api_key=api_key)
+genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
 
+st.set_page_config(page_title="چت‌بات اسناد PDF", layout="centered")
 
-# ✅ Read all PDFs (supports Farsi + English)
-def get_pdf_text(pdf_docs):
+# --------------------
+# بارگذاری PDF ها
+# --------------------
+def load_pdfs(pdf_files):
     text = ""
-    for pdf in pdf_docs:
-        with fitz.open(stream=pdf.read(), filetype="pdf") as doc:
-            for page in doc:
-                page_text = page.get_text("text")
-                if page_text:
-                    text += page_text + "\n"
+    for pdf_file in pdf_files:
+        pdf_reader = PdfReader(pdf_file)
+        for page in pdf_reader.pages:
+            text += page.extract_text() or ""
     return text
 
-
-# ✅ Split text into chunks
-def get_text_chunks(text):
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    return splitter.split_text(text)
-
-
-# ✅ Get embeddings + save vector store
-def get_vector_store(chunks):
-    embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
-    vector_store = FAISS.from_texts(chunks, embedding=embeddings)
-    vector_store.save_local("faiss_index")
-
-
-# ✅ Conversational chain
+# --------------------
+# ساخت زنجیره پرسش‌و‌پاسخ
+# --------------------
 def get_conversational_chain():
-    map_prompt_template = """
-    لطفاً بر اساس متن زیر فقط به سؤال پاسخ دهید.
-    اگر جواب دقیق نبود، بگویید: "اطلاعات کافی در متن موجود نیست".
+    # پرامپت برای هر تیکه (map)
+    map_prompt = PromptTemplate(
+        input_variables=["context", "question"],
+        template="""
+        لطفاً بر اساس متن زیر فقط به سؤال پاسخ دهید.
+        اگر جواب دقیق نبود، بگویید: "اطلاعات کافی در متن موجود نیست".
 
-    --- متن:
-    {context}
+        --- متن:
+        {context}
 
-    --- سوال:
-    {question}
+        --- سوال:
+        {question}
 
-    --- پاسخ:
-    """
-    combine_prompt_template = """
-    شما یک دستیار هوشمند هستید. پاسخ نهایی را بر اساس نتایج تیکه‌های مختلف بنویس.
-    اگر پاسخی در متن نبود، بگویید: "اطلاعات کافی در متن موجود نیست".
+        --- پاسخ:
+        """
+    )
 
-    --- نتایج جزئی:
-    {summaries}
+    # پرامپت برای جمع‌بندی (combine)
+    combine_prompt = PromptTemplate(
+        input_variables=["summaries", "question"],
+        template="""
+        شما یک دستیار هوشمند هستید. پاسخ نهایی را بر اساس نتایج تیکه‌های مختلف بنویس.
+        اگر پاسخی در متن نبود، بگویید: "اطلاعات کافی در متن موجود نیست".
 
-    --- سوال:
-    {question}
+        --- نتایج جزئی:
+        {summaries}
 
-    --- پاسخ نهایی:
-    """
+        --- سوال:
+        {question}
 
-    model = ChatGoogleGenerativeAI(model="gemini-2.5-pro", client=genai, temperature=0.3)
+        --- پاسخ نهایی:
+        """
+    )
 
-    map_prompt = PromptTemplate(template=map_prompt_template, input_variables=["context", "question"])
-    combine_prompt = PromptTemplate(template=combine_prompt_template, input_variables=["summaries", "question"])
+    # مدل Gemini
+    model = ChatGoogleGenerativeAI(model="gemini-2.0-pro", temperature=0.3)
 
-    chain = load_qa_chain(
-        llm=model,
-        chain_type="map_reduce",
-        map_prompt=map_prompt,
-        combine_prompt=combine_prompt
+    # chain برای combine (مرحله آخر)
+    combine_documents_chain = StuffDocumentsChain(
+        llm_chain=load_qa_chain(model, chain_type="stuff", prompt=combine_prompt),
+        document_variable_name="summaries"
     )
-    return chain
 
+    # chain اصلی (map → reduce)
+    chain = MapReduceDocumentsChain(
+        llm_chain=load_qa_chain(model, chain_type="stuff", prompt=map_prompt),
+        reduce_documents_chain=ReduceDocumentsChain(
+            combine_documents_chain=combine_documents_chain
+        ),
+        document_variable_name="context",
+        return_intermediate_steps=False,
+    )
 
+    return chain
 
-# ✅ Clear history
-def clear_chat_history():
-    st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
+# --------------------
+# اینترفیس استریم‌لیت
+# --------------------
+def main():
+    st.title("🤖 چت‌بات PDF با Gemini")
 
+    # بارگذاری فایل‌ها
+    pdf_files = st.file_uploader("📂 فایل‌های PDF خود را بارگذاری کنید", type="pdf", accept_multiple_files=True)
 
-# ✅ Handle user input
-def user_input(user_question):
-    embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
-    new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
-    docs = new_db.similarity_search(user_question, k=8)  # ⬅️ بیشتر شده
-    chain = get_conversational_chain()
-    response = chain({"input_documents": docs, "question": user_question}, return_only_outputs=True)
-    return response
+    if pdf_files:
+        text = load_pdfs(pdf_files)
 
+        if text.strip() == "":
+            st.warning("❌ هیچ متنی از PDF استخراج نشد.")
+            return
 
-# ✅ Main Streamlit app
-def main():
-    st.set_page_config(page_title="Chatbot", layout="wide", initial_sidebar_state="expanded")
+        # دریافت سؤال کاربر
+        question = st.text_input("❓ پرسش خود را وارد کنید:")
 
-    st.markdown(
-        """
-        <style>
-        body { background-color: #000000; color: #ffffff; }
-        .main { background-color: #333333; padding: 20px; border-radius: 10px; }
-        .rtl { direction: rtl; text-align: right; }
-        </style>
-        """,
-        unsafe_allow_html=True
-    )
+        if question:
+            with st.spinner("در حال پردازش..."):
+                try:
+                    chain = get_conversational_chain()
+                    docs = [{"page_content": text}]  # کل متن به عنوان یک داکیومنت
+
+                    response = chain.invoke({
+                        "input_documents": docs,
+                        "question": question
+                    })
 
-    if "uploaded" not in st.session_state:
-        st.session_state.uploaded = False
-
-    if not st.session_state.uploaded:
-        # Upload Page
-        st.title("دستیار شخصی شما ...")
-        pdf_docs = st.file_uploader("فایل(های) PDF را آپلود کنید", accept_multiple_files=True)
-        if st.button("تایید"):
-            if pdf_docs:
-                st.info("در حال پردازش ...")
-                raw_text = get_pdf_text(pdf_docs)
-                text_chunks = get_text_chunks(raw_text)
-                get_vector_store(text_chunks)
-                st.session_state.uploaded = True
-                st.success("پردازش موفق شد ✅")
-            else:
-                st.error("لطفاً حداقل یک فایل انتخاب کنید")
-    else:
-        # Chat Page
-        st.title("دستیار آماده است ...")
-        st.write("سؤالتان را بپرسید 👇")
-
-        if st.button("بازگشت به صفحه آپلود"):
-            st.session_state.uploaded = False
-            clear_chat_history()
-            st.rerun()
-
-        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"]):
-                # تشخیص ساده فارسی برای نمایش راست به چپ
-                if any("\u0600" <= ch <= "\u06FF" for ch in message["content"]):
-                    st.markdown(f"<div class='rtl'>{message['content']}</div>", unsafe_allow_html=True)
-                else:
-                    st.write(message["content"])
-
-        if prompt := st.chat_input():
-            st.session_state.messages.append({"role": "user", "content": prompt})
-            with st.chat_message("user"):
-                st.markdown(f"<div class='rtl'>{prompt}</div>", unsafe_allow_html=True) if any(
-                    "\u0600" <= ch <= "\u06FF" for ch in prompt) else st.write(prompt)
-
-            if st.session_state.messages[-1]["role"] != "assistant":
-                with st.chat_message("assistant"):
-                    response = user_input(prompt)
-                    if response:
-                        full_response = response['output_text']
-                        if any("\u0600" <= ch <= "\u06FF" for ch in full_response):
-                            st.markdown(f"<div class='rtl'>{full_response}</div>", unsafe_allow_html=True)
-                        else:
-                            st.write(full_response)
-                        st.session_state.messages.append({"role": "assistant", "content": full_response})
+                    st.markdown(
+                        f"<div style='direction: rtl; text-align: right; font-size: 16px;'>"
+                        f"{response['output_text']}</div>",
+                        unsafe_allow_html=True
+                    )
 
+                except Exception as e:
+                    st.error(f"⚠️ خطا: {e}")
 
 if __name__ == "__main__":
     main()