Update app.py

app.py

@@ -1,68 +1,65 @@
 import os
 import fitz  # PyMuPDF
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_google_genai import GoogleGenerativeAIEmbeddings
 import streamlit as st
 import google.generativeai as genai
-from langchain.vectorstores import FAISS
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain.chains.question_answering import load_qa_chain
-from langchain.prompts import PromptTemplate
 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()
 api_key = os.getenv("GOOGLE_API_KEY")
 genai.configure(api_key=api_key)
 
 
-# ---------------- PDF Extraction ----------------
+# ✅ Function to read all PDF files (Farsi + English)
 def get_pdf_text(pdf_docs):
-    """Extract text from uploaded PDFs (supports Farsi + English)."""
     text = ""
     for pdf in pdf_docs:
-        doc = fitz.open(stream=pdf.read(), filetype="pdf")
-        for page in doc:
-            page_text = page.get_text("text")
-            if page_text:
-                text += page_text + "\n"
+        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
 
 
-# ---------------- Text Chunking ----------------
+# ✅ Function to split text into chunks
 def get_text_chunks(text):
-    """Split text into smaller chunks for embedding."""
-    splitter = RecursiveCharacterTextSplitter(
-        chunk_size=800,
-        chunk_overlap=150,
-        separators=["\n\n", "\n", " ", ""]
-    )
-    return splitter.split_text(text)
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    chunks = splitter.split_text(text)
+    return chunks
 
 
-# ---------------- FAISS Vector Store ----------------
+# ✅ Function to get embeddings for each chunk and save to vector store
 def get_vector_store(chunks):
-    """Create and store FAISS index in session state."""
     try:
         embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
         vector_store = FAISS.from_texts(chunks, embedding=embeddings)
-        st.session_state.vector_store = vector_store
-    except (GoogleAPIError, InvalidArgument):
-        raise RuntimeError("Error processing embeddings. Please try again in a minute.")
+        vector_store.save_local("faiss_index")
+    except Exception as e:
+        raise RuntimeError(f"Error creating vector store: {e}")
 
 
-# ---------------- Conversational Chain ----------------
+# ✅ Function to get conversational chain
 def get_conversational_chain():
-    """Create QA chain with custom prompt."""
     prompt_template = """
-    You are a helpful assistant. 
-    Always answer in the same language as the question (Farsi or English).  
-    If the answer is not in the provided context, reply with: 
-    "پاسخ در متن موجود نیست" for Farsi OR "answer is not available in the context" for English.  
+    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}
 
-    Context:\n {context}\n
-    Question:\n {question}\n
     Answer:
     """
     try:
@@ -70,28 +67,29 @@ def get_conversational_chain():
         prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
         chain = load_qa_chain(llm=model, chain_type="stuff", prompt=prompt)
         return chain
-    except (GoogleAPIError, InvalidArgument):
-        raise RuntimeError("Error creating conversational chain. Please try again in a minute.")
+    except Exception as e:
+        raise RuntimeError(f"Error creating conversational chain: {e}")
 
 
-# ---------------- Chat Functions ----------------
+# ✅ Function to clear chat history
 def clear_chat_history():
-    st.session_state.messages = [
-        {"role": "assistant", "content": "در خدمتیم"}]
+    st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
 
 
+# ✅ Function to handle user input
 def user_input(user_question):
-    """Handle user query and return response."""
     try:
-        docs = st.session_state.vector_store.similarity_search(user_question, k=4)
+        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['output_text']
-    except (GoogleAPIError, InvalidArgument):
-        raise RuntimeError("لطفا پس از چند لحظه دوباره امتحان کنید ")
+        return response
+    except Exception as e:
+        raise RuntimeError(f"Error while answering: {e}")
 
 
-# ---------------- Main App ----------------
+# ✅ Main function to run the Streamlit app
 def main():
     st.set_page_config(
         page_title="Chatbot",
@@ -99,7 +97,7 @@ def main():
         initial_sidebar_state="expanded"
     )
 
-    # Custom CSS for better design
+    # Dark theme styling
     st.markdown(
         """
         <style>
@@ -117,7 +115,6 @@ def main():
         unsafe_allow_html=True
     )
 
-    # Initialize session state variables
     if "uploaded" not in st.session_state:
         st.session_state.uploaded = False
 
@@ -128,12 +125,12 @@ def main():
         if st.button("تایید"):
             if pdf_docs:
                 try:
-                    st.info("در حال پردازش")
+                    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("پردازش موفق")
+                    st.success("پردازش موفق شد ✅")
                 except RuntimeError as e:
                     st.error(str(e))
             else:
@@ -141,40 +138,36 @@ def main():
     else:
         # Chat Page
         st.title("Assistant ready ...")
-        st.write("میتونین سوالتونو بپرسین")
+        st.write("میتونین سوالتونو بپرسین 👇")
 
-        # Add a "Return" button to go back to the upload page
         if st.button("بازگشت به صفحه آپلود"):
-            st.session_state.uploaded = False  
-            st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}] 
+            st.session_state.uploaded = False
+            clear_chat_history()
             st.rerun()
 
         st.button('حذف مکالمه', on_click=clear_chat_history)
 
-        # Initialize chat history
         if "messages" not in st.session_state:
-            st.session_state.messages = [
-                {"role": "assistant", "content": "در خدمتیم"}]
+            st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
 
-        # Display chat messages
+        # Show history
         for message in st.session_state.messages:
             with st.chat_message(message["role"]):
                 st.write(message["content"])
 
-        # Chat input
         if prompt := st.chat_input():
             st.session_state.messages.append({"role": "user", "content": prompt})
             with st.chat_message("user"):
                 st.write(prompt)
 
-            # Generate bot response
             if st.session_state.messages[-1]["role"] != "assistant":
                 try:
                     with st.chat_message("assistant"):
                         response = user_input(prompt)
-                        st.write(response)
-                        message = {"role": "assistant", "content": response}
-                        st.session_state.messages.append(message)
+                        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))