| 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}") | |