Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,31 +6,72 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 
 st.set_page_config(page_title="RAG Search", page_icon="🔍")
 
-# --- 1️⃣ Unzip chroma_db.zip if not already extracted ---
-ZIP_PATH = os.path.join(os.path.dirname(__file__), "..", "chroma_db.zip")
-DB_PATH = os.path.join(os.path.dirname(__file__), "..", "chroma_db")
-
-if not os.path.exists(DB_PATH):
-    st.info("📦 Extracting Chroma DB for first-time setup...")
-    with zipfile.ZipFile(ZIP_PATH, "r") as zip_ref:
-        zip_ref.extractall(DB_PATH)
-    st.success("✅ Database extracted successfully!")
-
-# --- 2️⃣ Initialize embedding model ---
-embeddings = HuggingFaceEmbeddings(
-    model_name="mixedbread-ai/mxbai-embed-large-v1",
-    model_kwargs={"device": "cpu"}  # Force CPU for Hugging Face Spaces
-)
-
-# --- 3️⃣ Load Chroma database ---
-vectordb = Chroma(persist_directory=DB_PATH, embedding_function=embeddings)
-
-# --- 4️⃣ Query input & results ---
+# ---------------------------------------------
+# 1️⃣ Locate your Chroma DB zip
+# ---------------------------------------------
+possible_paths = [
+    "chroma_db.zip",
+    os.path.join("src", "chroma_db.zip"),
+    os.path.join(os.path.dirname(__file__), "chroma_db.zip"),
+    os.path.join(os.path.dirname(__file__), "..", "chroma_db.zip"),
+    "/app/chroma_db.zip",
+]
+
+ZIP_PATH = None
+for p in possible_paths:
+    if os.path.exists(p):
+        ZIP_PATH = p
+        break
+
+if ZIP_PATH is None:
+    st.error("❌ Could not find 'chroma_db.zip'. Please ensure it's in your repo root.")
+    st.stop()
+
+DB_PATH = "chroma_db"
+
+# ---------------------------------------------
+# 2️⃣ Extract only once per app session
+# ---------------------------------------------
+if "db_ready" not in st.session_state:
+    if not os.path.exists(DB_PATH):
+        st.info("📦 Extracting Chroma DB for the first time...")
+        with zipfile.ZipFile(ZIP_PATH, "r") as zip_ref:
+            zip_ref.extractall(DB_PATH)
+        st.success("✅ Database extracted successfully!")
+    else:
+        st.info("✅ Chroma DB folder already exists.")
+    st.session_state.db_ready = True  # Mark as ready
+
+# ---------------------------------------------
+# 3️⃣ Load embeddings (CPU safe)
+# ---------------------------------------------
+@st.cache_resource(show_spinner=False)
+def load_embeddings():
+    return HuggingFaceEmbeddings(
+        model_name="mixedbread-ai/mxbai-embed-large-v1",
+        model_kwargs={"device": "cpu"}
+    )
+
+embeddings = load_embeddings()
+
+# ---------------------------------------------
+# 4️⃣ Load Chroma DB (cached)
+# ---------------------------------------------
+@st.cache_resource(show_spinner=False)
+def load_vectordb():
+    return Chroma(persist_directory=DB_PATH, embedding_function=embeddings)
+
+vectordb = load_vectordb()
+
+# ---------------------------------------------
+# 5️⃣ Query + Display Results
+# ---------------------------------------------
 query = st.text_input("Enter your query:", "What is SystemVerilog interface?")
 
 if st.button("Search"):
     st.write("🔎 Searching your local vector database...")
     results = vectordb.similarity_search(query, k=3)
+
     if results:
         for i, doc in enumerate(results):
             st.subheader(f"Result {i+1}")
@@ -38,4 +79,5 @@ if st.button("Search"):
             st.caption(doc.metadata)
             st.markdown("---")
     else:
-        st.warning("⚠️ No results found.")
+        st.warning("⚠️ No matching results found.")
+