vector_store.py

import json
from pathlib import Path
from typing import Dict, List, Optional, Union



import chromadb
from sentence_transformers import SentenceTransformer

from models import BoundingBox, Chunk, FigureRef, QueryResult

_TAG_PREFIX = "tag_"


class VectorStore:
    """ChromaDB-backed vector store for :class:`~models.Chunk` objects.

    Chunks are embedded with a sentence-transformers model and stored in a
    persistent local ChromaDB collection.  The ``file_hash`` field on every
    chunk is stored as metadata so that results can later be filtered to a
    single document, and the corresponding DoclingDocument can be retrieved
    from a :class:`~document_store.DocumentStore` for visual grounding.

    Tags are stored as ChromaDB metadata with a ``tag_`` prefix so they can be
    used as filter predicates in :meth:`query`.
    """

    def __init__(
        self,
        collection_name: str = "vector",
        persist_directory: str = "./vector_data/vector/embeddings",
        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
    ):
        persist_path = Path(persist_directory)
        persist_path.mkdir(parents=True, exist_ok=True)

        self.client = chromadb.PersistentClient(path=str(persist_path))
        self.embeddings = SentenceTransformer(embedding_model)
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )
        self.collection_name = collection_name

    # ------------------------------------------------------------------
    # Serialization helpers
    # ------------------------------------------------------------------

    def _to_meta(self, chunk: Chunk, tags: Optional[Dict[str, str]] = None) -> dict:
        """Flatten a Chunk to a ChromaDB-compatible metadata dict (str/int/float values).

        Tags are stored with a ``tag_`` prefix so they can be filtered via
        ChromaDB ``where`` clauses without colliding with built-in fields.
        """
        meta = {
            "index": chunk.index,
            # Use -1 as a sentinel for None; ChromaDB does not support null values.
            "page_number": chunk.page_number if chunk.page_number is not None else -1,
            "page_range": json.dumps(chunk.page_range),
            "headings": json.dumps(chunk.headings),
            "doc_items": json.dumps(chunk.doc_items),
            "bboxes": json.dumps([b.model_dump() for b in chunk.bboxes]),
            "figures": json.dumps([f.model_dump() for f in chunk.figures]),
            "document_name": chunk.document_name,
            "file_hash": chunk.file_hash,
            "file_extension": chunk.file_extension,
        }
        if tags:
            for k, v in tags.items():
                meta[f"{_TAG_PREFIX}{k}"] = v
        return meta

    def _from_meta(self, id: str, text: str, meta: dict) -> Chunk:
        """Reconstruct a Chunk from a ChromaDB result row."""
        page_no = meta["page_number"]
        return Chunk(
            id=id,
            index=meta["index"],
            text=text,
            page_number=page_no if page_no != -1 else None,
            page_range=json.loads(meta.get("page_range", "[]")),
            headings=json.loads(meta["headings"]),
            doc_items=json.loads(meta["doc_items"]),
            bboxes=[BoundingBox(**b) for b in json.loads(meta["bboxes"])],
            figures=[FigureRef(**f) for f in json.loads(meta.get("figures", "[]"))],
            document_name=meta["document_name"],
            file_hash=meta["file_hash"],
            file_extension=meta["file_extension"],
        )

    # ------------------------------------------------------------------
    # CRUD
    # ------------------------------------------------------------------

    def create(self, chunks: List[Chunk], tags: Optional[Dict[str, str]] = None) -> None:
        """Embed and store a list of chunks in the collection.

        Args:
            chunks: Chunks to add.  Their ``id`` values must be unique within
                    the collection (the Chunker generates ``"{file_hash}_{index}"``
                    ids which satisfy this across documents).
            tags: Optional dict of ``key -> value`` tags to attach to every
                  chunk's metadata (stored with a ``tag_`` prefix).
        """
        if not chunks:
            return

        texts = [c.text for c in chunks]
        embeddings = self.embeddings.encode(texts).tolist()
        self.collection.upsert(
            ids=[c.id for c in chunks],
            embeddings=embeddings,
            documents=texts,
            metadatas=[self._to_meta(c, tags=tags) for c in chunks],
        )

    def _fetch_window(self, hit: Chunk, window: int) -> List[Chunk]:
        """Fetch up to *window* chunks before and after *hit* from the same document."""
        if window == 0:
            return [hit]

        start = max(0, hit.index - window)
        end = hit.index + window

        raw = self.collection.get(
            where={
                "$and": [
                    {"file_hash": {"$eq": hit.file_hash}},
                    {"index": {"$gte": start}},
                    {"index": {"$lte": end}},
                ]
            },
            include=["documents", "metadatas"],
        )

        chunks = [
            self._from_meta(
                id=raw["ids"][i],
                text=raw["documents"][i],
                meta=raw["metadatas"][i],
            )
            for i in range(len(raw["ids"]))
        ]
        chunks.sort(key=lambda c: c.index)
        return chunks

    def query(
        self,
        query_text: str,
        top_k: int = 5,
        file_hash: Optional[Union[str, List[str]]] = None,
        window: int = 0,
        tags: Optional[Dict[str, Union[str, List[str]]]] = None,
    ) -> List[QueryResult]:
        """Semantic search over stored chunks.

        Args:
            query_text: Natural-language query.
            top_k: Maximum number of results to return.
            file_hash: Restrict results to one or more documents.
                       Pass a single hash string, a list of hash strings,
                       or ``None`` to search across all documents.
            window: Number of adjacent chunks to include before and after each
                    hit.  ``0`` returns only the matching chunk itself.
            tags: Optional dict of tag filters. Values may be a single string
                  (exact match, ``$eq``) or a list of strings (OR match, ``$in``).
                  All supplied tag keys must match (AND semantics across keys).
                  Tags are stored with a ``tag_`` prefix in ChromaDB metadata.

        Returns:
            List of :class:`~models.QueryResult` objects ranked by cosine
            similarity.  Each result exposes the matching ``chunk`` and a
            ``context`` list of surrounding chunks sorted by index.
        """
        total = self.collection.count()
        if total == 0:
            return []

        query_embedding = self.embeddings.encode(query_text).tolist()

        # Build where clause from file_hash + tags filters
        conditions = []

        if file_hash is not None:
            if isinstance(file_hash, list):
                conditions.append({"file_hash": {"$in": file_hash}})
            else:
                conditions.append({"file_hash": {"$eq": file_hash}})

        if tags:
            for k, v in tags.items():
                field = f"{_TAG_PREFIX}{k}"
                if isinstance(v, list):
                    # OR logic: any of the listed values matches
                    conditions.append({field: {"$in": v}})
                else:
                    conditions.append({field: {"$eq": v}})

        if len(conditions) == 0:
            where = None
        elif len(conditions) == 1:
            where = conditions[0]
        else:
            where = {"$and": conditions}

        # When filtering, the filtered subset may be smaller than top_k.
        # ChromaDB raises if n_results exceeds the number of matching chunks.
        # Fetch at most top_k IDs to clamp n_results without scanning the full collection.
        if where is not None:
            sample = self.collection.get(where=where, limit=top_k)
            filtered_count = len(sample["ids"])
            if filtered_count == 0:
                return []
            n_results = filtered_count
        else:
            n_results = min(top_k, total)

        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=n_results,
            where=where,
        )

        query_results = []
        seen_chunk_ids: set[str] = set()
        for i in range(len(results["documents"][0])):
            hit = self._from_meta(
                id=results["ids"][0][i],
                text=results["documents"][0][i],
                meta=results["metadatas"][0][i],
            )
            # Skip hits whose chunk was already covered by a prior result's
            # context window — avoids returning the same text multiple times.
            if hit.id in seen_chunk_ids:
                continue
            context = self._fetch_window(hit, window)
            for c in context:
                seen_chunk_ids.add(c.id)
            query_results.append(QueryResult(chunk=hit, context=context))

        return query_results

    def list_documents(self) -> List[dict]:
        """Return one summary entry per unique document stored in the collection.

        Returns:
            List of dicts with ``file_hash``, ``document_name``, ``file_extension``,
            and ``chunk_count`` keys, sorted by document_name.
        """
        raw = self.collection.get(include=["metadatas"])
        seen: dict[str, dict] = {}
        for meta in raw["metadatas"]:
            fh = meta["file_hash"]
            if fh not in seen:
                seen[fh] = {
                    "file_hash": fh,
                    "document_name": meta["document_name"],
                    "file_extension": meta["file_extension"],
                    "chunk_count": 0,
                }
            seen[fh]["chunk_count"] += 1
        return sorted(seen.values(), key=lambda d: d["document_name"])

    def delete(self, file_hash: str) -> None:
        """Remove all chunks belonging to a document.

        Args:
            file_hash: The ``binary_hash`` of the document whose chunks should
                       be removed.
        """
        results = self.collection.get(where={"file_hash": file_hash})
        if results["ids"]:
            self.collection.delete(ids=results["ids"])

    def update(self, chunks: List[Chunk], tags: Optional[Dict[str, str]] = None) -> None:
        """Replace all chunks for a document (delete existing, then re-add).

        Args:
            chunks: New chunks for the document.  All chunks must share the
                    same ``file_hash``.
            tags: Optional tags to store on the replacement chunks.
        """
        if not chunks:
            return
        self.delete(chunks[0].file_hash)
        self.create(chunks, tags=tags)

    def remove_tag_metadata(self, file_hash: str, tag_keys: List[str]) -> None:
        """Remove ``tag_<key>`` fields from ChromaDB metadata for all chunks of a document.

        Args:
            file_hash: The document whose chunks should be updated.
            tag_keys: Tag keys to remove (without the ``tag_`` prefix).
        """
        results = self.collection.get(
            where={"file_hash": {"$eq": file_hash}},
            include=["metadatas"],
        )
        if not results["ids"]:
            return
        updated_metas = []
        for m in results["metadatas"]:
            updated = dict(m)
            for k in tag_keys:
                updated.pop(f"{_TAG_PREFIX}{k}", None)
            updated_metas.append(updated)
        self.collection.update(
            ids=results["ids"],
            metadatas=updated_metas,
        )

    def set_tag_metadata(self, file_hash: str, tags: Dict[str, str]) -> None:
        """Merge tag fields into ChromaDB metadata for all chunks of a document.

        Args:
            file_hash: The document whose chunks should be updated.
            tags: Dict of tag key -> value to merge. Keys are stored with the ``tag_`` prefix.
        """
        results = self.collection.get(
            where={"file_hash": {"$eq": file_hash}},
            include=["metadatas"],
        )
        if not results["ids"]:
            return
        updated_metas = []
        for m in results["metadatas"]:
            updated = dict(m)
            for k, v in tags.items():
                updated[f"{_TAG_PREFIX}{k}"] = v
            updated_metas.append(updated)
        self.collection.update(
            ids=results["ids"],
            metadatas=updated_metas,
        )

    def has_chunks(self, file_hash: str) -> bool:
        """Return True if any chunks exist for a document in the collection.

        Uses a limit=1 fetch to avoid scanning the entire collection.
        """
        result = self.collection.get(
            where={"file_hash": {"$eq": file_hash}},
            limit=1,
        )
        return bool(result["ids"])

    def rename_document(self, file_hash: str, new_name: str) -> None:
        """Update the document_name metadata on all chunks for a document.

        No re-embedding required — only the stored metadata field changes.

        Args:
            file_hash: The document to rename.
            new_name: The new display name.
        """
        results = self.collection.get(
            where={"file_hash": {"$eq": file_hash}},
            include=["metadatas"],
        )
        if not results["ids"]:
            return
        updated_metas = [
            {**m, "document_name": new_name} for m in results["metadatas"]
        ]
        self.collection.update(
            ids=results["ids"],
            metadatas=updated_metas,
        )