Update app.py

app.py

@@ -16,18 +16,40 @@ from langchain_huggingface import HuggingFacePipeline, HuggingFaceEmbeddings
 from langchain_chroma import Chroma
 from typing import Dict, Any, List
 
-# --- 2. UNZIP DATABASE (The Fix) ---
+# --- 2. UNZIP & AUTO-DETECT PATH ---
 print("⏳ Checking for Database...")
 
-# If the folder doesn't exist, but the ZIP does, unzip it!
-if not os.path.exists("./chroma_db") and os.path.exists("./chroma_db.zip"):
-    print("📦 Found chroma_db.zip! Unzipping...")
+# Unzip if the zip exists
+if os.path.exists("./chroma_db.zip"):
+    print("📦 Found zip file! Unzipping...")
     with zipfile.ZipFile("./chroma_db.zip", 'r') as zip_ref:
         zip_ref.extractall(".")
     print("✅ Unzip complete.")
 
-if not os.path.exists("./chroma_db"):
-    raise ValueError("❌ Error: Could not find 'chroma_db' folder or 'chroma_db.zip'. Please upload the zip file.")
+# SMART DETECTION: Find where the database went
+db_path = ""
+
+if os.path.exists("./chroma_db/chroma.sqlite3"):
+    # Case A: It's inside the folder (Perfect)
+    db_path = "./chroma_db"
+    print(f"📂 Found database in folder: {db_path}")
+
+elif os.path.exists("./chroma.sqlite3"):
+    # Case B: It spilled into the root directory
+    db_path = "."
+    print(f"📂 Found database in root directory: {db_path}")
+
+elif os.path.exists("./content/chroma_db/chroma.sqlite3"):
+    # Case C: It's inside a 'content' folder (Common Colab issue)
+    db_path = "./content/chroma_db"
+    print(f"📂 Found database in content folder: {db_path}")
+
+else:
+    # Case D: Panic
+    # Let's list the files to debug
+    print("❌ ERROR: Cannot find chroma.sqlite3. Current files in folder:")
+    print(os.listdir("."))
+    raise ValueError("Could not find the database file after unzipping!")
 
 # --- 3. MODEL SETUP ---
 print("⏳ Loading Embeddings...")
@@ -36,9 +58,9 @@ embedding_function = HuggingFaceEmbeddings(
     model_kwargs={"trust_remote_code": True, "device": "cpu"}
 )
 
-print("⏳ Loading Database...")
+print(f"⏳ Loading Database from {db_path}...")
 vector_db = Chroma(
-    persist_directory="./chroma_db", 
+    persist_directory=db_path, 
     embedding_function=embedding_function
 )