Spaces:
Running
Running
| import chromadb | |
| from chromadb.config import Settings as ChromaSettings | |
| from app.config import settings | |
| from app.services.embedding_service import get_embedding_service | |
| from typing import List, Dict, Any | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class VectorStore: | |
| def __init__(self): | |
| self._client = chromadb.PersistentClient( | |
| path=settings.chroma_persist_directory, | |
| settings=ChromaSettings(anonymized_telemetry=False), | |
| ) | |
| self._collection = self._client.get_or_create_collection( | |
| name=settings.chroma_collection_name, | |
| metadata={"hnsw:space": "cosine"}, | |
| ) | |
| logger.info(f"ChromaDB collection '{settings.chroma_collection_name}' ready") | |
| def upsert( | |
| self, | |
| ids: List[str], | |
| embeddings: List[List[float]], | |
| documents: List[str], | |
| metadatas: List[Dict[str, Any]], | |
| ) -> None: | |
| self._collection.upsert( | |
| ids=ids, | |
| embeddings=embeddings, | |
| documents=documents, | |
| metadatas=metadatas, | |
| ) | |
| def delete_by_sources(self, sources: List[str]) -> None: | |
| """Remove all chunks belonging to the given source files. Called before | |
| re-ingesting a file so edited documents don't leave stale vectors | |
| behind (chunk IDs are content-hashed, so edits change the IDs).""" | |
| if not sources: | |
| return | |
| self._collection.delete(where={"source": {"$in": sources}}) | |
| def query(self, embedding: List[float], top_k: int) -> Dict[str, Any]: | |
| return self._collection.query( | |
| query_embeddings=[embedding], | |
| n_results=top_k, | |
| include=["documents", "metadatas", "distances"], | |
| ) | |
| def count(self) -> int: | |
| return self._collection.count() | |
| def reset(self) -> None: | |
| self._client.delete_collection(settings.chroma_collection_name) | |
| self._collection = self._client.get_or_create_collection( | |
| name=settings.chroma_collection_name, | |
| metadata={"hnsw:space": "cosine"}, | |
| ) | |
| logger.info("Collection reset") | |
| def list_sources(self) -> List[Dict[str, Any]]: | |
| """Return per-source stats aggregated from all chunk metadata. | |
| Category is read from stored metadata when available, otherwise | |
| derived from the top-level folder in the source path so that | |
| documents ingested before the category field was added still | |
| display correctly.""" | |
| from pathlib import Path as _Path | |
| result = self._collection.get(include=["metadatas"]) | |
| metadatas = result.get("metadatas") or [] | |
| sources: Dict[str, Dict[str, Any]] = {} | |
| for meta in metadatas: | |
| source = meta.get("source", "") | |
| stored_cat = meta.get("category", "") | |
| if not stored_cat: | |
| parts = _Path(source).parts | |
| stored_cat = parts[0] if len(parts) > 1 else "" | |
| if source not in sources: | |
| sources[source] = { | |
| "source": source, | |
| "category": stored_cat, | |
| "title": meta.get("title", ""), | |
| "chunk_count": 0, | |
| } | |
| sources[source]["chunk_count"] += 1 | |
| return sorted(sources.values(), key=lambda x: x["source"]) | |
| def is_healthy(self) -> bool: | |
| try: | |
| self._collection.count() | |
| return True | |
| except Exception: | |
| return False | |
| _instance: VectorStore | None = None | |
| def get_vector_store() -> VectorStore: | |
| global _instance | |
| if _instance is None: | |
| _instance = VectorStore() | |
| return _instance | |