store.py

"""store.py - one VectorStore seam, two backends.

  LocalStore  : numpy matrix on disk. Zero deps beyond numpy. Cosine top-k. Runs
                today; fine for the 541 to ~20k chunk corpus.
  MilvusStore : Milvus-lite (file-local), the NVIDIA RAG-Blueprint store. Same
                interface; flip STORE_MODE=milvus on the Spark (cuVS GPU index is
                the further upgrade). Lazy-imports pymilvus so today's run needs
                nothing installed.

Interface: insert(ids, vecs, contents, metas) ; search(qvec, k) -> [{id, content,
metadata, score}] ; save(path) / load(path) ; len().
"""
from __future__ import annotations

import json
import os

import numpy as np

import config


class LocalStore:
    def __init__(self, dim: int):
        self.dim = dim
        self.ids: list[str] = []
        self.contents: list[str] = []
        self.metas: list[dict] = []
        self.vecs: np.ndarray | None = None

    def insert(self, ids, vecs, contents, metas) -> None:
        vecs = np.asarray(vecs, dtype=np.float32)
        self.vecs = vecs if self.vecs is None else np.vstack([self.vecs, vecs])
        self.ids += list(ids)
        self.contents += list(contents)
        self.metas += list(metas)

    def search(self, qvec, k: int = 8) -> list[dict]:
        if self.vecs is None or not self.ids:
            return []
        q = np.asarray(qvec, dtype=np.float32)
        if q.shape[0] != self.vecs.shape[1]:
            raise ValueError(
                f"query embedding is {q.shape[0]}D but the index is {self.vecs.shape[1]}D - "
                f"re-run `ingest.py` in the same EMBED_MODE you are querying with")
        qn = float(np.linalg.norm(q)) or 1.0
        vn = np.linalg.norm(self.vecs, axis=1)
        vn[vn == 0] = 1.0
        sims = (self.vecs @ q) / (vn * qn)
        idx = np.argsort(-sims)[:k]
        return [{"id": self.ids[i], "content": self.contents[i],
                 "metadata": self.metas[i], "score": float(sims[i])} for i in idx]

    def save(self, path: str) -> None:
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        vecs = self.vecs if self.vecs is not None else np.zeros((0, self.dim), np.float32)
        np.savez(path + ".npz", vecs=vecs)
        json.dump({"dim": self.dim, "ids": self.ids, "contents": self.contents,
                   "metas": self.metas}, open(path + ".json", "w"))

    @classmethod
    def load(cls, path: str) -> "LocalStore":
        meta = json.load(open(path + ".json"))
        s = cls(meta["dim"])
        s.ids, s.contents, s.metas = meta["ids"], meta["contents"], meta["metas"]
        s.vecs = np.load(path + ".npz")["vecs"]
        return s

    def __len__(self) -> int:
        return len(self.ids)


class MilvusStore:
    """Milvus-lite backend (tomorrow / Spark). Same interface as LocalStore."""
    def __init__(self, dim: int, uri: str = "index/milvus.db", collection: str = "london"):
        from pymilvus import MilvusClient   # lazy
        self.dim, self.collection = dim, collection
        os.makedirs(os.path.dirname(uri) or ".", exist_ok=True)
        self.client = MilvusClient(uri)
        if self.client.has_collection(collection):
            self.client.drop_collection(collection)
        self.client.create_collection(collection_name=collection, dimension=dim,
                                      auto_id=False, primary_field_name="pk")
        self._n = 0

    def insert(self, ids, vecs, contents, metas) -> None:
        rows = [{"pk": i, "vector": list(map(float, v)), "content": c, "meta": m}
                for i, v, c, m in zip(range(self._n, self._n + len(ids)), vecs, contents, metas)]
        # keep our string ids in meta so search can return them
        for r, sid in zip(rows, ids):
            r["meta"] = {**r["meta"], "_id": sid}
        self.client.insert(collection_name=self.collection, data=rows)
        self._n += len(rows)

    def search(self, qvec, k: int = 8) -> list[dict]:
        res = self.client.search(collection_name=self.collection,
                                 data=[list(map(float, qvec))], limit=k,
                                 output_fields=["content", "meta"])[0]
        return [{"id": h["entity"]["meta"].get("_id", h["id"]),
                 "content": h["entity"]["content"], "metadata": h["entity"]["meta"],
                 "score": float(h["distance"])} for h in res]

    def save(self, path: str) -> None:
        pass   # Milvus-lite persists to its own file

    def __len__(self) -> int:
        return self._n


def get_store(dim: int):
    return MilvusStore(dim) if config.STORE_MODE == "milvus" else LocalStore(dim)