File size: 2,087 Bytes
317eaf3 ce5b66d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | 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}")
|