Spaces:
Sleeping
Sleeping
File size: 3,635 Bytes
e7a3876 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """
ChromaDB vector store wrapper.
ChromaDB stores: the raw text, the embedding vector, and metadata (source, page).
On disk so the vector store survives between runs — you don't re-embed on every restart.
Interview concept to own:
Chroma uses HNSW (Hierarchical Navigable Small World) indexing under the hood.
HNSW is an approximate nearest neighbour algorithm — it trades a tiny bit of
recall for massive speed gains at scale. At 10M vectors, brute-force cosine
search is unusable; HNSW is O(log n).
"""
import chromadb
from chromadb.config import Settings
from src.utils.config import config
from src.utils.logger import logger
class VectorStore:
def __init__(self, collection_name: str = None, persist_dir: str = None):
collection_name = collection_name or config.collection_name
persist_dir = persist_dir or config.chroma_persist_dir
self.client = chromadb.PersistentClient(
path=persist_dir,
settings=Settings(anonymized_telemetry=False),
)
self.collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"hnsw:space": "cosine"}, # Use cosine distance, not L2
)
logger.info(f"Vector store ready: '{collection_name}' at {persist_dir}")
logger.info(f" Current document count: {self.collection.count()}")
def add_chunks(self, chunks: list[dict], embeddings: list[list[float]]) -> None:
"""
Store chunks + their embeddings in Chroma.
Args:
chunks: Output from chunker.chunk_pages()
embeddings: Parallel list of embedding vectors from Embedder.embed_texts()
"""
if not chunks:
logger.warning("No chunks to add.")
return
self.collection.add(
ids=[c["chunk_id"] for c in chunks],
embeddings=embeddings,
documents=[c["text"] for c in chunks],
metadatas=[{"source": c["source"], "page": c["page"]} for c in chunks],
)
logger.info(f"Added {len(chunks)} chunks. Total in store: {self.collection.count()}")
def query(self, query_embedding: list[float], top_k: int = None) -> list[dict]:
"""
Find the top-k most similar chunks to a query embedding.
Returns:
List of dicts: [{"text": str, "source": str, "page": int, "score": float}]
Ordered by similarity (most similar first).
"""
top_k = top_k or config.top_k
results = self.collection.query(
query_embeddings=[query_embedding],
n_results=top_k,
include=["documents", "metadatas", "distances"],
)
chunks = []
for doc, meta, dist in zip(
results["documents"][0],
results["metadatas"][0],
results["distances"][0],
):
chunks.append({
"text": doc,
"source": meta["source"],
"page": meta["page"],
"score": round(1 - dist, 4), # Convert cosine distance → similarity
})
return chunks
def reset(self) -> None:
"""Delete and recreate the collection. Use during development to re-ingest."""
name = self.collection.name
self.client.delete_collection(name)
self.collection = self.client.get_or_create_collection(
name=name,
metadata={"hnsw:space": "cosine"},
)
logger.warning(f"Collection '{name}' has been reset.")
|