pull chroma from s3

Dockerfile

@@ -6,8 +6,15 @@ RUN apt-get update && apt-get install -y \
     build-essential \
     curl \
     git \
+    unzip \
     && rm -rf /var/lib/apt/lists/*
 
+# Install AWS CLI
+RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" \
+    && unzip awscliv2.zip \
+    && ./aws/install \
+    && rm -rf aws awscliv2.zip
+
 # Install uv
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
 

pyproject.toml

@@ -5,6 +5,7 @@ description = "Congressional Legislation Query and Analysis Tool"
 readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
+    "boto3>=1.35.0",
     "chromadb>=1.1.0",
     "datasets>=3.0.0",
     "langchain>=0.3.27",

src/legisqa_local/app.py

@@ -1,7 +1,7 @@
 """Main Streamlit application for LegisQA"""
 
 import streamlit as st
-from legisqa_local.config.settings import STREAMLIT_CONFIG, setup_environment
+from legisqa_local.config.settings import STREAMLIT_CONFIG, setup_environment, setup_chromadb
 from legisqa_local.components.sidebar import render_sidebar
 from legisqa_local.tabs.rag_tab import RAGTab
 from legisqa_local.tabs.rag_sbs_tab import RAGSideBySideTab
@@ -16,6 +16,9 @@ def main():
     # Setup environment
     setup_environment()
     
+    # Setup ChromaDB (download if needed)
+    setup_chromadb()
+    
     # Main content
     st.title(":classical_building: LegisQA :classical_building:")
     st.header("Query Congressional Bills")

src/legisqa_local/config/settings.py

@@ -34,6 +34,127 @@ def get_chroma_config():
         "collection_name": os.getenv("CHROMA_COLLECTION_NAME", "usc")
     }
 
+def setup_chromadb():
+    """Setup ChromaDB - use persistent storage (/data) or download from S3 if needed"""
+    chroma_config = get_chroma_config()
+    chroma_path = chroma_config["persist_directory"]
+    
+    # For HF Spaces with persistent storage, prefer /data directory
+    persistent_chroma_path = "/data/chromadb"
+    if os.path.exists("/data"):
+        print("HF Spaces persistent storage detected at /data")
+        
+        # Check if ChromaDB exists in persistent storage
+        if os.path.exists(persistent_chroma_path) and os.listdir(persistent_chroma_path):
+            print(f"✅ ChromaDB found in persistent storage: {persistent_chroma_path}")
+            # Update environment variable to point to persistent storage
+            os.environ["CHROMA_PERSIST_DIRECTORY"] = persistent_chroma_path
+            return persistent_chroma_path
+        
+        # Download from S3 to persistent storage
+        s3_bucket = os.getenv("CHROMA_S3_BUCKET", "")
+        s3_prefix = os.getenv("CHROMA_S3_PREFIX", "")
+        
+        if s3_bucket and s3_prefix:
+            print(f"📥 Downloading ChromaDB from S3 to persistent storage...")
+            print(f"   Source: s3://{s3_bucket}/{s3_prefix}")
+            print(f"   Target: {persistent_chroma_path}")
+            
+            success = download_chromadb_from_s3(s3_bucket, s3_prefix, persistent_chroma_path)
+            if success:
+                # Update environment variable to point to persistent storage
+                os.environ["CHROMA_PERSIST_DIRECTORY"] = persistent_chroma_path
+                return persistent_chroma_path
+        else:
+            print("❌ No S3 configuration found (CHROMA_S3_BUCKET, CHROMA_S3_PREFIX)")
+    
+    # Fallback: check if ChromaDB exists at configured path (local development)
+    if os.path.exists(chroma_path) and os.listdir(chroma_path):
+        print(f"✅ ChromaDB found at {chroma_path}")
+        return chroma_path
+    
+    print(f"⚠️  Using default ChromaDB path: {chroma_path}")
+    return chroma_path
+
+def download_chromadb_from_s3(bucket: str, prefix: str, local_path: str) -> bool:
+    """Download ChromaDB from S3"""
+    try:
+        import subprocess
+        import os
+        
+        # Ensure target directory exists
+        os.makedirs(local_path, exist_ok=True)
+        
+        # Use AWS CLI to sync from S3 (no credentials needed for public buckets)
+        s3_url = f"s3://{bucket}/{prefix}"
+        cmd = ["aws", "s3", "sync", s3_url, local_path, "--no-sign-request"]
+        
+        print(f"Running: {' '.join(cmd)}")
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        
+        if result.returncode == 0:
+            print("✅ ChromaDB download from S3 complete!")
+            return True
+        else:
+            print(f"❌ S3 download failed: {result.stderr}")
+            return False
+            
+    except FileNotFoundError:
+        print("❌ AWS CLI not found. Trying with boto3...")
+        return download_chromadb_from_s3_boto3(bucket, prefix, local_path)
+    except Exception as e:
+        print(f"❌ Error downloading from S3: {e}")
+        return False
+
+def download_chromadb_from_s3_boto3(bucket: str, prefix: str, local_path: str) -> bool:
+    """Download ChromaDB from S3 using boto3 (fallback)"""
+    try:
+        import boto3
+        from botocore import UNSIGNED
+        from botocore.config import Config
+        import os
+        
+        print("📦 Using boto3 for S3 download...")
+        
+        # Create S3 client with no credentials (for public buckets)
+        s3 = boto3.client('s3', config=Config(signature_version=UNSIGNED))
+        
+        # List objects in the S3 prefix
+        paginator = s3.get_paginator('list_objects_v2')
+        pages = paginator.paginate(Bucket=bucket, Prefix=prefix)
+        
+        os.makedirs(local_path, exist_ok=True)
+        
+        file_count = 0
+        for page in pages:
+            if 'Contents' in page:
+                for obj in page['Contents']:
+                    key = obj['Key']
+                    # Get relative path by removing prefix
+                    relative_path = key[len(prefix):].lstrip('/')
+                    if relative_path:  # Skip empty paths
+                        local_file_path = os.path.join(local_path, relative_path)
+                        
+                        # Create directory if needed
+                        os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+                        
+                        # Download file
+                        file_count += 1
+                        if file_count % 10 == 0:
+                            print(f"Downloaded {file_count} files...")
+                        
+                        s3.download_file(bucket, key, local_file_path)
+        
+        print(f"✅ ChromaDB download from S3 (boto3) complete! Downloaded {file_count} files.")
+        return True
+        
+    except ImportError:
+        print("❌ boto3 not available. Please install: pip install boto3")
+        return False
+    except Exception as e:
+        print(f"❌ Error downloading from S3 with boto3: {e}")
+        return False
+
 # Embedding model configuration
 EMBEDDING_MODEL = "sentence-transformers/static-retrieval-mrl-en-v1"
 EMBEDDING_DEVICE = "cpu"

uv.lock

@@ -191,6 +191,34 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/10/cb/f2ad4230dc2eb1a74edf38f1a38b9b52277f75bef262d8908e60d957e13c/blinker-1.9.0-py3-none-any.whl", hash = "sha256:ba0efaa9080b619ff2f3459d1d500c57bddea4a6b424b60a91141db6fd2f08bc", size = 8458, upload-time = "2024-11-08T17:25:46.184Z" },
 ]
 
+[[package]]
+name = "boto3"
+version = "1.40.40"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "botocore" },
+    { name = "jmespath" },
+    { name = "s3transfer" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/3c/12/1a31b36802d0f33bc6982ab8b7e6437d75ef3c179abe6c53d4d8f7b4248f/boto3-1.40.40.tar.gz", hash = "sha256:f384d3a0410d0f1a4d4ae7aa69c41d0549c6ca5a76667dc25fc97d50ad6db740", size = 111606, upload-time = "2025-09-26T19:23:46.923Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/90/69/c65566dbdaaea3af0c23f7731ab0f185a38b593fd449d2423374150dbfe0/boto3-1.40.40-py3-none-any.whl", hash = "sha256:385904de68623e1c341bdc095d94a30006843032c912adeb1e0752a343632ec6", size = 139340, upload-time = "2025-09-26T19:23:45.557Z" },
+]
+
+[[package]]
+name = "botocore"
+version = "1.40.40"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "jmespath" },
+    { name = "python-dateutil" },
+    { name = "urllib3" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/83/5a/43a7fea503ad14fa79819f2b3103a38977fb587a3663d1ac6e958fccf592/botocore-1.40.40.tar.gz", hash = "sha256:78eb121a16a6481ed0f6e1aebe53a4f23aa121f34466846c13a5ca48fa980e31", size = 14363370, upload-time = "2025-09-26T19:23:37.853Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ed/5e/3bbf6d34cbf307c1b9e58e0204ceba2d35bbc0c93b4e3b3cc895aae0a5fd/botocore-1.40.40-py3-none-any.whl", hash = "sha256:68506142b3cde93145ef3ee0268f2444f2b68ada225a151f714092bbd3d6516a", size = 14031738, upload-time = "2025-09-26T19:23:35.475Z" },
+]
+
 [[package]]
 name = "build"
 version = "1.3.0"
@@ -817,6 +845,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/af/22/7ab7b4ec3a1c1f03aef376af11d23b05abcca3fb31fbca1e7557053b1ba2/jiter-0.11.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6e2bbf24f16ba5ad4441a9845e40e4ea0cb9eed00e76ba94050664ef53ef4406", size = 347102, upload-time = "2025-09-15T09:20:20.16Z" },
 ]
 
+[[package]]
+name = "jmespath"
+version = "1.0.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/00/2a/e867e8531cf3e36b41201936b7fa7ba7b5702dbef42922193f05c8976cd6/jmespath-1.0.1.tar.gz", hash = "sha256:90261b206d6defd58fdd5e85f478bf633a2901798906be2ad389150c5c60edbe", size = 25843, upload-time = "2022-06-17T18:00:12.224Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/31/b4/b9b800c45527aadd64d5b442f9b932b00648617eb5d63d2c7a6587b7cafc/jmespath-1.0.1-py3-none-any.whl", hash = "sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980", size = 20256, upload-time = "2022-06-17T18:00:10.251Z" },
+]
+
 [[package]]
 name = "joblib"
 version = "1.5.2"
@@ -1076,6 +1113,7 @@ name = "legisqa-local"
 version = "0.1.0"
 source = { editable = "." }
 dependencies = [
+    { name = "boto3" },
     { name = "chromadb" },
     { name = "datasets" },
     { name = "langchain" },
@@ -1094,6 +1132,7 @@ dependencies = [
 
 [package.metadata]
 requires-dist = [
+    { name = "boto3", specifier = ">=1.35.0" },
     { name = "chromadb", specifier = ">=1.1.0" },
     { name = "datasets", specifier = ">=3.0.0" },
     { name = "langchain", specifier = ">=0.3.27" },
@@ -2351,6 +2390,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/64/8d/0133e4eb4beed9e425d9a98ed6e081a55d195481b7632472be1af08d2f6b/rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762", size = 34696, upload-time = "2025-04-16T09:51:17.142Z" },
 ]
 
+[[package]]
+name = "s3transfer"
+version = "0.14.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "botocore" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/62/74/8d69dcb7a9efe8baa2046891735e5dfe433ad558ae23d9e3c14c633d1d58/s3transfer-0.14.0.tar.gz", hash = "sha256:eff12264e7c8b4985074ccce27a3b38a485bb7f7422cc8046fee9be4983e4125", size = 151547, upload-time = "2025-09-09T19:23:31.089Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/48/f0/ae7ca09223a81a1d890b2557186ea015f6e0502e9b8cb8e1813f1d8cfa4e/s3transfer-0.14.0-py3-none-any.whl", hash = "sha256:ea3b790c7077558ed1f02a3072fb3cb992bbbd253392f4b6e9e8976941c7d456", size = 85712, upload-time = "2025-09-09T19:23:30.041Z" },
+]
+
 [[package]]
 name = "safetensors"
 version = "0.6.2"