src/shared_db.py

import chromadb
from pathlib import Path
import shutil

import threading

# Use absolute path for ChromaDB to avoid issues with working directory
BASE_DIR = Path(__file__).parent.parent
CHROMA_PATH = str(BASE_DIR / "db")

_client = None
_lock = threading.Lock()

def get_chroma_client():
    global _client
    if _client is None:
        with _lock:
            if _client is None:
                try:
                    print(f"Initializing ChromaDB PersistentClient at {CHROMA_PATH}")
                    # Ensure directory exists
                    Path(CHROMA_PATH).mkdir(exist_ok=True, parents=True)
                    _client = chromadb.PersistentClient(path=CHROMA_PATH)
                    print("ChromaDB Client initialized successfully.")
                except Exception as e:
                    print(f"FAILED to initialize ChromaDB Client: {e}")
                    raise e
    return _client

def get_collection(name: str):
    client = get_chroma_client()
    try:
        return client.get_or_create_collection(
            name=name,
            metadata={"hnsw:space": "cosine"}
        )
    except Exception as e:
        print(f"Error getting/creating collection '{name}': {e}")
        raise e


def clear_chroma_storage():
    """Remove all persistent chroma data on disk and reset the client."""
    global _client
    with _lock:
        # Drop the old client reference so SQLite file handles are released
        _client = None

        try:
            shutil.rmtree(CHROMA_PATH)
            Path(CHROMA_PATH).mkdir(exist_ok=True, parents=True)
            print(f"Cleared chroma storage directory: {CHROMA_PATH}")
        except Exception as e:
            print(f"Error clearing chroma storage: {e}")
            raise

        # Eagerly create a fresh client so subsequent calls don't hit a stale db
        try:
            _client = chromadb.PersistentClient(path=CHROMA_PATH)
            print("Fresh ChromaDB Client re-initialized after clear.")
        except Exception as e:
            print(f"Warning: could not re-init ChromaDB Client: {e}")