Update app.py

app.py

@@ -1,131 +1,176 @@
+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 langchain_community.vectorstores import FAISS
-from langchain.embeddings import HuggingFaceEmbeddings
-import os
 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
+
+# Load environment variables
 load_dotenv()
-genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
+api_key = os.getenv("GOOGLE_API_KEY")
+genai.configure(api_key=api_key)
 
-st.set_page_config(page_title="چت‌بات اسناد PDF", layout="centered")
 
-# --------------------
-# بارگذاری PDF ها
-# --------------------
-def load_pdfs(pdf_files):
+# ✅ Function to read all PDF files (Farsi + English)
+def get_pdf_text(pdf_docs):
     text = ""
-    for pdf_file in pdf_files:
-        pdf_reader = PdfReader(pdf_file)
-        for page in pdf_reader.pages:
-            text += page.extract_text() or ""
+    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"
     return text
 
-# --------------------
-# ساخت زنجیره پرسش‌و‌پاسخ
-# --------------------
-def get_conversational_chain():
-    # پرامپت برای هر تیکه (map)
-    map_prompt = PromptTemplate(
-        input_variables=["context", "question"],
-        template="""
-        لطفاً بر اساس متن زیر فقط به سؤال پاسخ دهید.
-        اگر جواب دقیق نبود، بگویید: "اطلاعات کافی در متن موجود نیست".
 
-        --- متن:
-        {context}
+# ✅ Function to split text into chunks
+def get_text_chunks(text):
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    chunks = splitter.split_text(text)
+    return chunks
 
-        --- سوال:
-        {question}
 
-        --- پاسخ:
-        """
-    )
+# ✅ 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")
+        vector_store = FAISS.from_texts(chunks, embedding=embeddings)
+        vector_store.save_local("faiss_index")
+    except Exception as e:
+        raise RuntimeError(f"Error creating vector store: {e}")
 
-    # پرامپت برای جمع‌بندی (combine)
-    combine_prompt = PromptTemplate(
-        input_variables=["summaries", "question"],
-        template="""
-        شما یک دستیار هوشمند هستید. پاسخ نهایی را بر اساس نتایج تیکه‌های مختلف بنویس.
-        اگر پاسخی در متن نبود، بگویید: "اطلاعات کافی در متن موجود نیست".
 
-        --- نتایج جزئی:
-        {summaries}
-
-        --- سوال:
-        {question}
-
-        --- پاسخ نهایی:
-        """
-    )
-
-    # مدل Gemini
-    model = ChatGoogleGenerativeAI(model="gemini-2.0-pro", temperature=0.3)
-
-    # chain برای combine (مرحله آخر)
-    combine_documents_chain = StuffDocumentsChain(
-        llm_chain=load_qa_chain(model, chain_type="stuff", prompt=combine_prompt),
-        document_variable_name="summaries"
+# ✅ Function to get conversational chain
+def get_conversational_chain():
+    prompt_template = """
+    You are a helpful assistant. Do NOT reveal your identity (Gemini) or the company (Google).
+    Answer the question as detailed as possible using ONLY the provided context. 
+    If the answer is not in the context, say: "answer is not available in the context".
+    Do not make up answers.
+
+    Context:
+    {context}
+
+    Question:
+    {question}
+
+    Answer:
+    """
+    try:
+        model = ChatGoogleGenerativeAI(model="gemini-2.5-pro", 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}")
+
+
+# ✅ Function to clear chat history
+def clear_chat_history():
+    st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
+
+
+# ✅ 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)
+        docs = new_db.similarity_search(user_question, k=4)
+        chain = get_conversational_chain()
+        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}")
+
+
+# ✅ Main function to run the Streamlit app
+def main():
+    st.set_page_config(
+        page_title="Chatbot",
+        layout="wide",
+        initial_sidebar_state="expanded"
     )
 
-    # 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,
+    # Dark theme styling
+    st.markdown(
+        """
+        <style>
+        body {
+            background-color: #000000;
+            color: #ffffff;
+        }
+        .main {
+            background-color: #333333;
+            padding: 20px;
+            border-radius: 10px;
+        }
+        </style>
+        """,
+        unsafe_allow_html=True
     )
 
-    return chain
-
-# --------------------
-# اینترفیس استریم‌لیت
-# --------------------
-def main():
-    st.title("🤖 چت‌بات PDF با Gemini")
-
-    # بارگذاری فایل‌ها
-    pdf_files = st.file_uploader("📂 فایل‌های PDF خود را بارگذاری کنید", type="pdf", accept_multiple_files=True)
-
-    if pdf_files:
-        text = load_pdfs(pdf_files)
-
-        if text.strip() == "":
-            st.warning("❌ هیچ متنی از PDF استخراج نشد.")
-            return
+    if "uploaded" not in st.session_state:
+        st.session_state.uploaded = False
 
-        # دریافت سؤال کاربر
-        question = st.text_input("❓ پرسش خود را وارد کنید:")
-
-        if question:
-            with st.spinner("در حال پردازش..."):
+    if not st.session_state.uploaded:
+        # Upload Page
+        st.title("Your personal assistant ...")
+        pdf_docs = st.file_uploader("فایل پی دی اف مورد نظر را آپلود کنید ", accept_multiple_files=True)
+        if st.button("تایید"):
+            if pdf_docs:
+                try:
+                    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("پردازش موفق شد ✅")
+                except RuntimeError as e:
+                    st.error(str(e))
+            else:
+                st.error("لطفا حداقل یک فایل را آپلود کنید")
+    else:
+        # Chat Page
+        st.title("Assistant ready ...")
+        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"]):
+                st.write(message["content"])
+
+        if prompt := st.chat_input():
+            st.session_state.messages.append({"role": "user", "content": prompt})
+            with st.chat_message("user"):
+                st.write(prompt)
+
+            if st.session_state.messages[-1]["role"] != "assistant":
                 try:
-                    chain = get_conversational_chain()
-                    docs = [{"page_content": text}]  # کل متن به عنوان یک داکیومنت
-
-                    response = chain.invoke({
-                        "input_documents": docs,
-                        "question": question
-                    })
-
-                    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}")
+                    with st.chat_message("assistant"):
+                        response = user_input(prompt)
+                        if response:
+                            full_response = response['output_text']
+                            st.write(full_response)
+                            st.session_state.messages.append({"role": "assistant", "content": full_response})
+                except RuntimeError as e:
+                    st.error(str(e))
+
 
 if __name__ == "__main__":
     main()