Update app.py

app.py

@@ -1,13 +1,14 @@
+import streamlit as st
+from PyPDF2 import PdfReader
+import os
 import asyncio
+
+# ------------ Ensure an asyncio loop (required by Gemini libs) ------------
 try:
     asyncio.get_running_loop()
 except RuntimeError:
     asyncio.set_event_loop(asyncio.new_event_loop())
 
-import streamlit as st
-from PyPDF2 import PdfReader
-import os
-
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_google_genai import GoogleGenerativeAIEmbeddings
 from langchain_community.vectorstores import FAISS
@@ -15,79 +16,100 @@ from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain.chains.question_answering import load_qa_chain
 from langchain.prompts import PromptTemplate
 
+# --------------------------- CONFIG ---------------------------------------
 GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY", "")
-PDF_PATH = "librarianship.pdf"
-INDEX_PATH = "faiss_index"   # CHANGED from /tmp/faiss_index to /data/faiss_index
+PDF_PATH  = "librarianship.pdf"
+INDEX_DIR = "/tmp/faiss_index"            # FAISS saves as INDEX_DIR.index
 
-def get_pdf_text(pdf_path):
-    text = ""
+# --------------------- Helper functions -----------------------------------
+def get_pdf_text(pdf_path: str) -> str:
+    """Extract full text from a single PDF."""
     reader = PdfReader(pdf_path)
-    for page in reader.pages:
-        page_text = page.extract_text()
-        if page_text:
-            text += page_text
-    return text
+    return "".join(page.extract_text() or "" for page in reader.pages)
 
-def get_text_chunks(text):
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
-    return text_splitter.split_text(text)
+def get_text_chunks(text: str):
+    splitter = RecursiveCharacterTextSplitter(chunk_size=10_000, chunk_overlap=1_000)
+    return splitter.split_text(text)
 
 def build_and_save_vector_store(text_chunks, api_key):
-    embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=api_key)
+    embeddings   = GoogleGenerativeAIEmbeddings(
+        model="models/embedding-001", google_api_key=api_key
+    )
     vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
-    vector_store.save_local(INDEX_PATH)
+    vector_store.save_local(INDEX_DIR)
 
-@st.cache_resource(show_spinner=False)
 def load_vector_store(api_key):
-    embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=api_key)
-    return FAISS.load_local(INDEX_PATH, embeddings, allow_dangerous_deserialization=True)
+    """Load FAISS index from disk once and keep in session memory."""
+    if "vector_store" not in st.session_state:
+        embeddings = GoogleGenerativeAIEmbeddings(
+            model="models/embedding-001", google_api_key=api_key
+        )
+        st.session_state.vector_store = FAISS.load_local(
+            INDEX_DIR, embeddings, allow_dangerous_deserialization=True
+        )
+    return st.session_state.vector_store
 
-@st.cache_resource(show_spinner=False)
 def get_conversational_chain(api_key):
-    prompt_template = """
-    You are a helpful assistant that only answers based on the context provided from the PDF document.
-    Do not use any external knowledge or assumptions. If the answer is not found in the context below, reply with "I don't know."
-    Context:
-    {context}
-    Question:
-    {question}
-    Answer:
-    """
-    model = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0, google_api_key=api_key)
-    prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
-    chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)
-    return chain
-
-def user_input(user_question, api_key):
-    db = load_vector_store(api_key)
-    docs = db.similarity_search(user_question)
-    chain = get_conversational_chain(api_key)
-    response = chain({"input_documents": docs, "question": user_question}, return_only_outputs=True)
-    st.write("**Reply:**", response["output_text"])
+    """Create the QA chain once and reuse it."""
+    if "qa_chain" not in st.session_state:
+        prompt_template = """
+        You are a helpful assistant that only answers based on the context provided
+        from the PDF document below. If the answer is not in the context, reply with
+        "I don't know."
+
+        Context:
+        {context}
+
+        Question:
+        {question}
+
+        Answer:
+        """
+        model  = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash", temperature=0, google_api_key=api_key
+        )
+        prompt = PromptTemplate(
+            template=prompt_template, input_variables=["context", "question"]
+        )
+        st.session_state.qa_chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)
+    return st.session_state.qa_chain
+
+def answer_question(question, api_key):
+    db     = load_vector_store(api_key)
+    docs   = db.similarity_search(question)
+    chain  = get_conversational_chain(api_key)
+    result = chain(
+        {"input_documents": docs, "question": question},
+        return_only_outputs=True,
+    )
+    st.write("**Reply:**", result["output_text"])
 
+# ----------------------------- MAIN APP ------------------------------------
 def main():
-    st.set_page_config(page_title="Chat with librarianship")
+    st.set_page_config(page_title="Chat librarianship.pdf")
     st.header("Librarianship AI Application (Gemini 2.0)")
     st.markdown("---")
 
-    # --- Ensure API key is present ---
+    # --- 0. Check API key ---------------------------------------------------
     if not GOOGLE_API_KEY:
-        st.error("Please set the GOOGLE_API_KEY environment variable in your Hugging Face Space secrets or .env file.")
+        st.error(
+            "Please set the GOOGLE_API_KEY environment variable "
+            "in your Hugging Face Space secrets or .env file."
+        )
         st.stop()
 
-    # --- Build FAISS index if not present ---
-    if not os.path.exists(INDEX_PATH + ".index"):
-        with st.spinner(f"Indexing {PDF_PATH}..."):
-            raw_text = get_pdf_text(PDF_PATH)
-            text_chunks = get_text_chunks(raw_text)
+    # --- 1. Build FAISS index (only if missing) ----------------------------
+    if not os.path.exists(INDEX_DIR + ".index"):
+        with st.spinner(f"Indexing {PDF_PATH} …"):
+            text_chunks = get_text_chunks(get_pdf_text(PDF_PATH))
             build_and_save_vector_store(text_chunks, GOOGLE_API_KEY)
-        st.success(f"Indexed {PDF_PATH}. You can now ask questions.")
+        st.success("Index built! Ask away 👇")
 
-    # --- Ask questions UI ---
+    # --- 2. Chat UI --------------------------------------------------------
     st.subheader("Ask a question about librarianship")
-    user_question = st.text_input("Ask a question")
-    if user_question:
-        user_input(user_question, GOOGLE_API_KEY)
+    q = st.text_input("Type your question here")
+    if q:
+        answer_question(q, GOOGLE_API_KEY)
 
 if __name__ == "__main__":
     main()