Update app.py

app.py

@@ -1,80 +1,87 @@
 import sys
 import subprocess
 from flask import Flask, render_template, request, flash, redirect, url_for
-from datasets import load_dataset
 import torch
 from transformers import AutoTokenizer, AutoModel
-import numpy as np
 import os
+import chromadb
+from huggingface_hub import snapshot_download
 
 # --- 1. Initialize Flask App ---
 app = Flask(__name__)
-# A secret key is needed for flashing messages to the user's session
 app.secret_key = os.urandom(24)
 
 # --- 2. Configuration & Resource Loading ---
 print("Starting application...")
 
-# Point this to the Hugging Face Dataset repository you want to create/use.
-# This MUST match the DATASET_REPO in build_rag.py
-DATASET_REPO = "broadfield-dev/bible-rag-dataset-gemma"
-MODEL_NAME = "google/embeddinggemma-300m" # Use a consistent model for embedding and searching
+# --- Configuration (Must match build_rag.py) ---
+CHROMA_PATH = "chroma_db"
+COLLECTION_NAME = "bible_verses"
+MODEL_NAME = "google/embeddinggemma-300m"
+DATASET_REPO = "broadfield-dev/bible-chromadb-gemma"
 
-# Global variables for the dataset and models
-rag_dataset = None
+# --- Global variables for resources ---
+chroma_collection = None
 tokenizer = None
 embedding_model = None
 
 def load_resources():
     """
-    Attempts to load the dataset and models from the Hugging Face Hub.
-    Returns True on success, False on failure.
+    Downloads the DB from the Hub if not present, then loads it and the model.
     """
-    global rag_dataset, tokenizer, embedding_model
-    if rag_dataset:
+    global chroma_collection, tokenizer, embedding_model
+    if chroma_collection and embedding_model:
         return True
     
-    print(f"Attempting to load resources: {DATASET_REPO} and {MODEL_NAME}")
+    print("Attempting to load resources...")
     try:
-        # Load the pre-built dataset with the FAISS index
-        rag_dataset = load_dataset(DATASET_REPO)['train']
-        
-        # Load the Gemma model and tokenizer
+        # 1. Download the ChromaDB files from the Hugging Face Hub
+        # This will only download if the folder doesn't already exist.
+        print(f"Ensuring database is available locally from '{DATASET_REPO}'...")
+        snapshot_download(
+            repo_id=DATASET_REPO,
+            repo_type="dataset",
+            local_dir=CHROMA_PATH,
+            local_dir_use_symlinks=False # Recommended for Spaces
+        )
+        print("Database files are present locally.")
+
+        # 2. Initialize ChromaDB client from the downloaded files
+        client = chromadb.PersistentClient(path=CHROMA_PATH)
+        collection = client.get_collection(name=COLLECTION_NAME)
+
+        if collection.count() == 0:
+            print(f"Warning: Database collection is empty.")
+            return False
+            
+        chroma_collection = collection
+        print(f"Successfully connected to DB with {collection.count()} items.")
+
+        # 3. Load the embedding model
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
         embedding_model = AutoModel.from_pretrained(MODEL_NAME)
+        print(f"Embedding model '{MODEL_NAME}' loaded successfully.")
         
-        print("Models and dataset loaded successfully!")
         return True
     except Exception as e:
-        print(f"Could not load RAG dataset from '{DATASET_REPO}'. It may not exist yet.")
+        print(f"Could not load resources. The database may not be built yet.")
         print(f"Error: {e}")
-        # Reset globals to ensure a clean state
-        rag_dataset = None
-        tokenizer = None
-        embedding_model = None
         return False
 
-# Try to load resources on startup. The app can still run if this fails.
+# Try to load resources on startup.
 resources_loaded = load_resources()
 
-# --- 3. Define App Routes ---
-
+# --- 3. Define App Routes (Unchanged from previous ChromaDB version) ---
 @app.route('/')
 def home():
     if not resources_loaded:
-        flash(f"Welcome! The required RAG dataset '{DATASET_REPO}' is not loaded. Please use the 'Build RAG Dataset' button to create and upload it.", "warning")
+        flash(f"Welcome! Database not ready. Use the admin panel to build it.", "warning")
     return render_template('index.html')
 
 @app.route('/build-rag', methods=['POST'])
 def build_rag_route():
-    """
-    Triggers the build_rag.py script as a background process.
-    NOTE: This requires a Hugging Face token with 'write' permissions
-    to be saved as a secret named HF_TOKEN in the Space settings.
-    """
-    print("RAG build process requested.")
+    print("Vector database build process requested.")
     try:
-        # Use Popen to run the script in the background without blocking the app.
         process = subprocess.Popen(
             [sys.executable, "build_rag.py"],
             stdout=subprocess.PIPE,
@@ -82,51 +89,47 @@ def build_rag_route():
             text=True
         )
         print(f"Started build process with PID: {process.pid}")
-        flash("RAG build process initiated! This will run in the background and can take several minutes. Please check the Space logs for progress. Once complete, you can start searching.", "info")
+        flash("Database build & push initiated! This can take several minutes. Check logs for progress. The app will be ready when it completes.", "info")
     except Exception as e:
         print(f"Failed to start build process: {e}")
-        flash(f"An error occurred while trying to start the build process: {e}", "error")
-
+        flash(f"An error occurred: {e}", "error")
     return redirect(url_for('home'))
 
 @app.route('/search', methods=['POST'])
 def search():
     global resources_loaded
-    # If resources weren't loaded, try again in case the build just finished.
     if not resources_loaded:
-        print("Resources not loaded. Attempting to reload for search...")
+        print("Reloading resources for search...")
         resources_loaded = load_resources()
         if not resources_loaded:
-            flash("The RAG dataset is not ready yet. Please wait for the build process to complete or check the logs for errors.", "error")
+            flash("Database not ready. Please wait for the build process to finish.", "error")
             return redirect(url_for('home'))
 
     user_query = request.form['query']
     if not user_query:
         return render_template('index.html', results=[])
 
-    # --- Create embedding for the user's query ---
     inputs = tokenizer(user_query, return_tensors="pt")
     with torch.no_grad():
         outputs = embedding_model(**inputs)
     query_embedding = outputs.last_hidden_state.mean(dim=1).cpu().numpy()
-    
-    query_embedding = np.float32(query_embedding)
 
-    # --- Search the FAISS index ---
-    scores, retrieved_examples = rag_dataset.get_nearest_examples(
-        'embeddings',
-        query_embedding,
-        k=10 # Get top 10 results
+    search_results = chroma_collection.query(
+        query_embeddings=query_embedding.tolist(),
+        n_results=10
     )
 
-    # --- Format results for display ---
     results_list = []
-    for i in range(len(scores)):
+    documents = search_results['documents'][0]
+    metadatas = search_results['metadatas'][0]
+    distances = search_results['distances'][0]
+
+    for i in range(len(documents)):
         results_list.append({
-            'score': scores[i],
-            'text': retrieved_examples['text'][i],
-            'reference': retrieved_examples['reference'][i],
-            'version': retrieved_examples['version'][i]
+            'score': distances[i],
+            'text': documents[i],
+            'reference': metadatas[i].get('reference', 'N/A'),
+            'version': metadatas[i].get('version', 'N/A')
         })
     
     return render_template('index.html', results=results_list, query=user_query)