src/retrieval/vectorstore.py

"""ChromaDB vector store with local sentence-transformer embeddings.

Embeddings run locally so retrieval works without any API key and keeps
document content off third-party embedding endpoints.
"""

from __future__ import annotations

from functools import lru_cache

import chromadb

from config import CHROMA_DIR, COLLECTION_NAME, EMBEDDING_MODEL, TOP_K
from src.ingestion.chunker import Chunk


@lru_cache(maxsize=1)
def _embedder():
    from sentence_transformers import SentenceTransformer
    return SentenceTransformer(EMBEDDING_MODEL)


class VectorStore:
    def __init__(self, persist_dir: str = CHROMA_DIR,
                 collection: str = COLLECTION_NAME):
        self._client = chromadb.PersistentClient(path=persist_dir)
        self._name = collection
        self._collection = self._client.get_or_create_collection(
            self._name, metadata={"hnsw:space": "cosine"}
        )
        # in-memory chunk registry for BM25 + citation lookups
        self._chunks: dict[str, Chunk] = {}

    # -- indexing -----------------------------------------------------------

    def add(self, chunks: list[Chunk]) -> int:
        if not chunks:
            return 0
        embeddings = _embedder().encode([c.text for c in chunks], show_progress_bar=False)
        self._collection.upsert(
            ids=[c.chunk_id for c in chunks],
            documents=[c.text for c in chunks],
            embeddings=embeddings.tolist(),
            metadatas=[c.metadata for c in chunks],
        )
        for c in chunks:
            self._chunks[c.chunk_id] = c
        return len(chunks)

    def clear(self):
        self._client.delete_collection(self._name)
        self._collection = self._client.get_or_create_collection(
            self._name, metadata={"hnsw:space": "cosine"}
        )
        self._chunks.clear()

    def drop(self):
        """Delete the collection permanently (session-workspace eviction)."""
        try:
            self._client.delete_collection(self._name)
        except Exception:
            pass
        self._chunks.clear()

    # -- access -------------------------------------------------------------

    @property
    def chunks(self) -> list[Chunk]:
        return list(self._chunks.values())

    def get(self, chunk_id: str) -> Chunk | None:
        return self._chunks.get(chunk_id)

    def doc_ids(self) -> list[str]:
        return sorted({c.doc_id for c in self._chunks.values()})

    # -- search -------------------------------------------------------------

    def semantic_search(self, query: str, k: int = TOP_K,
                        doc_ids: list[str] | None = None) -> list[tuple[str, float]]:
        """Return [(chunk_id, similarity)] — cosine similarity in [0, 1]."""
        if not self._chunks:
            return []
        where = {"doc_id": {"$in": doc_ids}} if doc_ids else None
        embedding = _embedder().encode([query], show_progress_bar=False)
        res = self._collection.query(
            query_embeddings=embedding.tolist(),
            n_results=min(k, len(self._chunks)),
            where=where,
        )
        ids = res["ids"][0]
        distances = res["distances"][0]
        return [(cid, 1.0 - d) for cid, d in zip(ids, distances)]