Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-from PyPDF2 import PdfReader
+import fitz  # PyMuPDF
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_google_genai import GoogleGenerativeAIEmbeddings
 import streamlit as st
@@ -10,43 +10,57 @@ 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 google.cloud import aiplatform
 
 # Load environment variables
 load_dotenv()
 api_key = os.getenv("GOOGLE_API_KEY")
 genai.configure(api_key=api_key)
 
-# Function to read all PDF files and return text
+
+# ---------------- PDF Extraction ----------------
 def get_pdf_text(pdf_docs):
+    """Extract text from uploaded PDFs (supports Farsi + English)."""
     text = ""
     for pdf in pdf_docs:
-        pdf_reader = PdfReader(pdf)
-        for page in pdf_reader.pages:
-            text += page.extract_text()
+        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"
     return text
 
-# Function to split text into chunks
+
+# ---------------- Text Chunking ----------------
 def get_text_chunks(text):
-    splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
-    chunks = splitter.split_text(text)
-    return chunks  
+    """Split text into smaller chunks for embedding."""
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=800,
+        chunk_overlap=150,
+        separators=["\n\n", "\n", " ", ""]
+    )
+    return splitter.split_text(text)
 
-# Function to get embeddings for each chunk and save to vector store
+
+# ---------------- FAISS 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)
-        vector_store.save_local("faiss_index")
+        st.session_state.vector_store = vector_store
     except (GoogleAPIError, InvalidArgument):
         raise RuntimeError("Error processing embeddings. Please try again in a minute.")
 
-# Function to get conversational chain
+
+# ---------------- Conversational Chain ----------------
 def get_conversational_chain():
+    """Create QA chain with custom prompt."""
     prompt_template = """
-    You are a personal assistant.Do NOT reveal your identity like your name (Gemini) or the company (Google). Answer the question as detailed as possible from the provided context. 
-    If the answer is not in the provided context, just say, "answer is not available in the context".
-    Don't provide the wrong answer.\n\n
+    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.  
+
     Context:\n {context}\n
     Question:\n {question}\n
     Answer:
@@ -59,24 +73,25 @@ def get_conversational_chain():
     except (GoogleAPIError, InvalidArgument):
         raise RuntimeError("Error creating conversational chain. Please try again in a minute.")
 
-# Function to clear chat history
+
+# ---------------- Chat Functions ----------------
 def clear_chat_history():
     st.session_state.messages = [
         {"role": "assistant", "content": "در خدمتیم"}]
 
-# Function to handle user input
+
 def user_input(user_question):
+    """Handle user query and return response."""
     try:
-        embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")  
-        new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
-        docs = new_db.similarity_search(user_question, k=4)
+        docs = st.session_state.vector_store.similarity_search(user_question, k=4)
         chain = get_conversational_chain()
         response = chain({"input_documents": docs, "question": user_question}, return_only_outputs=True)
-        return response
+        return response['output_text']
     except (GoogleAPIError, InvalidArgument):
         raise RuntimeError("لطفا پس از چند لحظه دوباره امتحان کنید ")
 
-# Main function to run the Streamlit app
+
+# ---------------- Main App ----------------
 def main():
     st.set_page_config(
         page_title="Chatbot",
@@ -157,13 +172,12 @@ def main():
                 try:
                     with st.chat_message("assistant"):
                         response = user_input(prompt)
-                        if response:
-                            full_response = ''.join(response['output_text'])
-                            st.write(full_response)
-                            message = {"role": "assistant", "content": full_response}
-                            st.session_state.messages.append(message)
+                        st.write(response)
+                        message = {"role": "assistant", "content": response}
+                        st.session_state.messages.append(message)
                 except RuntimeError as e:
                     st.error(str(e))
 
+
 if __name__ == "__main__":
-    main()
+    main()