update

src/legisqa_local/config/settings.py

@@ -3,6 +3,8 @@
 import os
 import logging
 import streamlit as st
+import chromadb
+from chromadb.config import Settings, APIVersion
 
 logger = logging.getLogger(__name__)
 
@@ -12,6 +14,10 @@ STREAMLIT_CONFIG = {
     "page_title": "LegisQA"
 }
 
+# Embedding model configuration
+EMBEDDING_MODEL = "sentence-transformers/static-retrieval-mrl-en-v1"
+EMBEDDING_DEVICE = "cpu"
+
 def get_secret(key: str, default=None):
     """Get secret from Streamlit secrets or environment variables"""
     try:
@@ -33,253 +39,26 @@ def setup_environment():
 def get_chroma_config():
     """Get ChromaDB configuration from environment variables"""
     return {
-        "persist_directory": os.getenv("CHROMA_PERSIST_DIRECTORY", "./chromadb"),
         "collection_name": os.getenv("CHROMA_COLLECTION_NAME", "usc")
     }
 
-def setup_chromadb():
-    """Setup ChromaDB - use persistent storage (/data) or download from HF Dataset if needed"""
-    # Use a file-based lock to prevent duplicate runs (more reliable than session state in HF Spaces)
-    lock_file = "/tmp/chromadb_setup.lock"
-    
-    # Check if setup is already in progress or complete
-    if os.path.exists(lock_file):
-        logger.info("ChromaDB setup already in progress or complete (lock file exists)")
-        
-        # Even if setup is complete, ensure environment variable is correctly set
-        # Check if ChromaDB exists at /data/chromadb (HF Spaces persistent storage)
-        if os.path.exists("/data/chromadb") and os.listdir("/data/chromadb"):
-            logger.info("🔧 Ensuring CHROMA_PERSIST_DIRECTORY points to /data/chromadb")
-            os.environ["CHROMA_PERSIST_DIRECTORY"] = "/data/chromadb"
-        
-        chroma_config = get_chroma_config()
-        return chroma_config["persist_directory"]
-    
-    # Create lock file
-    try:
-        with open(lock_file, 'w') as f:
-            f.write("ChromaDB setup in progress")
-        logger.info("Created ChromaDB setup lock file")
-    except Exception as e:
-        logger.warning(f"Could not create lock file: {e}")
-        # Continue anyway
-    
-    logger.info("=== ChromaDB Setup Starting ===")
-    
-    chroma_config = get_chroma_config()
-    chroma_path = chroma_config["persist_directory"]
-    logger.info(f"Initial ChromaDB path: {chroma_path}")
-    
-    # Check for force refresh flag
-    force_refresh = os.getenv("CHROMA_FORCE_REFRESH", "").lower() in ("true", "1", "yes")
-    if force_refresh:
-        logger.info("🔄 CHROMA_FORCE_REFRESH enabled - will clear and re-download ChromaDB")
-    
-    # For HF Spaces with persistent storage, prefer /data directory
-    persistent_chroma_path = "/data/chromadb"
-    
-    # Check if we're in HF Spaces with persistent storage
-    if os.path.exists("/data"):
-        logger.info("🚀 HF Spaces persistent storage detected at /data")
-        
-        # Always clear /data completely to free up maximum space
-        logger.info("🧹 Clearing entire /data directory to free up space...")
-        import shutil
-        try:
-            # Check initial disk space
-            statvfs = os.statvfs("/data")
-            free_space_gb = (statvfs.f_bavail * statvfs.f_frsize) / (1024 * 1024 * 1024)
-            logger.info(f"💾 Initial free space: {free_space_gb:.2f} GB")
-            
-            for item in os.listdir("/data"):
-                item_path = os.path.join("/data", item)
-                if os.path.isdir(item_path):
-                    shutil.rmtree(item_path)
-                    logger.info(f"   Removed directory: {item}")
-                else:
-                    os.remove(item_path)
-                    logger.info(f"   Removed file: {item}")
-            
-            # Check free space after cleanup
-            statvfs = os.statvfs("/data")
-            free_space_gb = (statvfs.f_bavail * statvfs.f_frsize) / (1024 * 1024 * 1024)
-            logger.info(f"✅ /data directory cleared successfully")
-            logger.info(f"💾 Free space after cleanup: {free_space_gb:.2f} GB")
-        except Exception as e:
-            logger.error(f"❌ Error clearing /data directory: {e}")
-            logger.info("Continuing with download anyway...")
-        
-        # ChromaDB will always need to be downloaded fresh now
-        
-        # ChromaDB not found in persistent storage, try to download from HF Dataset
-        logger.info("ChromaDB not found in persistent storage, checking HF Dataset configuration...")
-        
-        dataset_repo = os.getenv("CHROMA_DATASET_REPO", "hyperdemocracy/usc-chroma-vecs-v1-chunks-v1-s8192-o512-sentence-transformers-static-retrieval-mrl-en-v1")
-        logger.info(f"HF Dataset repo: {dataset_repo}")
-        
-        if dataset_repo:
-            logger.info(f"📥 Downloading ChromaDB from HF Dataset to persistent storage...")
-            logger.info(f"   Source: {dataset_repo}")
-            logger.info(f"   Target: {persistent_chroma_path}")
-            
-            success = download_chromadb_from_hf_dataset(dataset_repo, persistent_chroma_path)
-            if success:
-                # Update environment variable to point to persistent storage
-                os.environ["CHROMA_PERSIST_DIRECTORY"] = persistent_chroma_path
-                logger.info(f"✅ ChromaDB download successful! Updated path to: {persistent_chroma_path}")
-                # Inspect the downloaded ChromaDB
-                inspect_chromadb(persistent_chroma_path)
-                return persistent_chroma_path
-            else:
-                logger.error("❌ ChromaDB download from HF Dataset failed!")
-        else:
-            logger.error("❌ No HF Dataset configuration found (CHROMA_DATASET_REPO)")
-            logger.info("Available environment variables:")
-            for key, value in os.environ.items():
-                if "CHROMA" in key:
-                    logger.info(f"  {key}={value}")
-    else:
-        logger.info("No /data directory found (not in HF Spaces with persistent storage)")
-    
-    # Fallback: check if ChromaDB exists at configured path (local development)
-    logger.info(f"Checking fallback path: {chroma_path}")
-    if os.path.exists(chroma_path):
-        if os.listdir(chroma_path):
-            logger.info(f"✅ ChromaDB found at {chroma_path}")
-            # Inspect the fallback ChromaDB
-            inspect_chromadb(chroma_path)
-            return chroma_path
-        else:
-            logger.warning(f"ChromaDB directory exists but is empty: {chroma_path}")
-    else:
-        logger.warning(f"ChromaDB directory does not exist: {chroma_path}")
-    
-    logger.warning(f"⚠️  Using default ChromaDB path: {chroma_path}")
-    logger.info("=== ChromaDB Setup Complete ===")
-    return chroma_path
+def create_chroma_client():
+    """Create and return a ChromaDB client"""
+    host = get_secret("CHROMA_PROXY_BASE")
+    token = get_secret("CHROMA_AUTH_TOKEN")
+    if not host or not token:
+        raise ValueError("Set CHROMA_PROXY_BASE and CHROMA_AUTH_TOKEN")
 
-def download_chromadb_from_hf_dataset(dataset_repo: str, local_path: str) -> bool:
-    """Download ChromaDB from HuggingFace Dataset"""
-    logger.info(f"Starting HF Dataset download: {dataset_repo} -> {local_path}")
-    
-    try:
-        from huggingface_hub import snapshot_download
-        import os
-        
-        # Ensure target directory exists
-        logger.info(f"Creating target directory: {local_path}")
-        os.makedirs(local_path, exist_ok=True)
-        
-        # Download the dataset using snapshot_download
-        logger.info(f"Downloading dataset: {dataset_repo}")
-        logger.info("This may take several minutes for large datasets...")
-        
-        # Download to a unique temporary location first to avoid conflicts
-        import uuid
-        temp_download_path = f"{local_path}_temp_{uuid.uuid4().hex[:8]}"
-        
-        # Use /tmp for cache (we cleared /data completely)
-        cache_dir = "/tmp/hf_chromadb_cache"
-        
-        downloaded_path = snapshot_download(
-            repo_id=dataset_repo,
-            repo_type="dataset",
-            local_dir=temp_download_path,
-            cache_dir=cache_dir
-            # Note: resume_download and local_dir_use_symlinks are now handled automatically
-        )
-        
-        logger.info(f"✅ ChromaDB download from HF Dataset complete!")
-        logger.info(f"Downloaded to: {downloaded_path}")
-        
-        # The HF dataset contains a 'chromadb' subdirectory with the actual ChromaDB files
-        chromadb_subdir = os.path.join(temp_download_path, "chromadb")
-        
-        if os.path.exists(chromadb_subdir):
-            logger.info(f"📁 Found ChromaDB subdirectory: {chromadb_subdir}")
-            
-            try:
-                # Move the ChromaDB files from the subdirectory to the target location
-                import shutil
-                if os.path.exists(local_path):
-                    logger.info(f"Removing existing target directory: {local_path}")
-                    shutil.rmtree(local_path)
-                
-                logger.info(f"Moving ChromaDB from {chromadb_subdir} to {local_path}")
-                shutil.move(chromadb_subdir, local_path)
-                
-                # Clean up the temporary download directory
-                logger.info(f"Cleaning up temporary directory: {temp_download_path}")
-                shutil.rmtree(temp_download_path)
-                
-                # Clean up HF cache to save disk space
-                if os.path.exists(cache_dir):
-                    logger.info(f"Cleaning up HF cache directory: {cache_dir}")
-                    try:
-                        shutil.rmtree(cache_dir)
-                        logger.info("✅ HF cache cleaned up successfully")
-                    except Exception as e:
-                        logger.warning(f"Could not clean up HF cache: {e}")
-                
-                logger.info(f"✅ ChromaDB files moved to: {local_path}")
-            except Exception as e:
-                logger.error(f"❌ Error moving ChromaDB files: {e}")
-                # Clean up temporary directory on error
-                try:
-                    shutil.rmtree(temp_download_path)
-                except:
-                    pass
-                return False
-        else:
-            logger.error(f"❌ ChromaDB subdirectory not found in downloaded data: {chromadb_subdir}")
-            # List what we actually downloaded for debugging
-            if os.path.exists(temp_download_path):
-                logger.info(f"Contents of {temp_download_path}:")
-                for item in os.listdir(temp_download_path):
-                    logger.info(f"  {item}")
-            # Clean up temporary directory
-            try:
-                import shutil
-                shutil.rmtree(temp_download_path)
-            except:
-                pass
-            return False
-        
-        # Verify the final ChromaDB structure
-        if os.path.exists(local_path) and os.listdir(local_path):
-            file_count = sum(len(files) for _, _, files in os.walk(local_path))
-            total_size = sum(
-                os.path.getsize(os.path.join(dirpath, filename))
-                for dirpath, _, filenames in os.walk(local_path)
-                for filename in filenames
-            ) / (1024 * 1024 * 1024)  # Convert to GB
-            
-            logger.info(f"📊 ChromaDB verification:")
-            logger.info(f"   Files: {file_count}")
-            logger.info(f"   Total size: {total_size:.2f} GB")
-            
-            # Check for key ChromaDB files
-            sqlite_file = os.path.join(local_path, "chroma.sqlite3")
-            if os.path.exists(sqlite_file):
-                sqlite_size = os.path.getsize(sqlite_file) / (1024 * 1024 * 1024)
-                logger.info(f"   SQLite database: {sqlite_size:.2f} GB")
-                return True
-            else:
-                logger.error("❌ chroma.sqlite3 not found in ChromaDB directory")
-                return False
-        else:
-            logger.error("❌ ChromaDB directory is empty after processing")
-            return False
-            
-    except ImportError:
-        logger.error("❌ huggingface_hub not available. Please install: pip install huggingface_hub")
-        return False
-    except Exception as e:
-        logger.error(f"❌ Error downloading from HF Dataset: {e}")
-        logger.error(f"Exception type: {type(e).__name__}")
-        return False
+    headers = {"Authorization": f"Bearer {token}"}
+    settings = Settings(
+        chroma_server_api_default_path=APIVersion.V2, anonymized_telemetry=False
+    )
+    client = chromadb.HttpClient(host=host, headers=headers, settings=settings)
+
+    return client
 
-def inspect_chromadb(chroma_path: str):
+
+def inspect_chromadb(client: chromadb.HttpClient):
     """Inspect ChromaDB collection to verify it's working correctly"""
     logger.info("🔍 === ChromaDB Collection Inspection ===")
     
@@ -291,11 +70,6 @@ def inspect_chromadb(chroma_path: str):
         collection_name = chroma_config["collection_name"]
         
         logger.info(f"📋 Collection name: {collection_name}")
-        logger.info(f"📁 ChromaDB path: {chroma_path}")
-        
-        # Create ChromaDB client
-        client = PersistentClient(path=chroma_path)
-        logger.info("✅ ChromaDB client created successfully")
         
         # List all collections
         collections = client.list_collections()
@@ -354,6 +128,4 @@ def inspect_chromadb(chroma_path: str):
     
     logger.info("🔍 === ChromaDB Inspection Complete ===")
 
-# Embedding model configuration
-EMBEDDING_MODEL = "sentence-transformers/static-retrieval-mrl-en-v1"
-EMBEDDING_DEVICE = "cpu"
+

src/legisqa_local/core/embeddings.py

@@ -1,14 +1,12 @@
 """Embedding functionality for LegisQA"""
 
-from langchain_huggingface import HuggingFaceEmbeddings
+from chromadb.utils import embedding_functions
 from legisqa_local.config.settings import EMBEDDING_MODEL, EMBEDDING_DEVICE
 
 
 def load_embeddings():
     """Load and return the embedding function"""
-    model_kwargs = {"device": EMBEDDING_DEVICE}
-    emb_fn = HuggingFaceEmbeddings(
-        model_name=EMBEDDING_MODEL,
-        model_kwargs=model_kwargs,
+    emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
+        model_name=EMBEDDING_MODEL
     )
     return emb_fn

src/legisqa_local/core/vectorstore.py

@@ -5,7 +5,7 @@ import os
 import streamlit as st
 from langchain_chroma import Chroma
 from legisqa_local.core.embeddings import load_embeddings
-from legisqa_local.config.settings import get_chroma_config
+from legisqa_local.config.settings import get_chroma_config, create_chroma_client
 
 logger = logging.getLogger(__name__)
 
@@ -16,10 +16,12 @@ def load_vectorstore():
     config = get_chroma_config()
     emb_fn = load_embeddings()
     
+    client = create_chroma_client()
     vectorstore = Chroma(
-        persist_directory=config["persist_directory"],
+        client=client,
+#        persist_directory=config["persist_directory"],
         collection_name=config["collection_name"],
-        embedding_function=emb_fn,
+#        embedding_function=emb_fn,
     )
     
     logger.info("✅ Vectorstore loaded successfully")