লোকাল ইনডেক্স ও রিট্রিভালের মান

# file: evidencedesk/index.py
"""A local vector index that remembers its embedding model, skips unusable text, filters and persists."""
from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Callable

import numpy as np

from evidencedesk.embeddings import Hit, Passage, embed_all, top_k, toy_embed, words

BatchEmbedder = Callable[[list[str]], np.ndarray]
MIN_CHARS = 3


class IndexMismatch(ValueError):
    """The index and the model that wants to use it are not compatible."""


@dataclass(frozen=True)
class Model:
    """An embedding function plus the identity that must travel with its vectors."""

    name: str
    dim: int
    embed_batch: BatchEmbedder
    normalised: bool = True

    @classmethod
    def toy(cls, dim: int = 256) -> Model:
        return cls(f"toy-trigram-{dim}", dim, lambda texts: embed_all(texts, lambda t: toy_embed(t, dim)))


def counting(model: Model) -> tuple[Model, list[int]]:
    """Wrap a model; the list collects the size of every embed_batch call."""
    sizes: list[int] = []

    def embed_batch(texts: list[str]) -> np.ndarray:
        sizes.append(len(texts))
        return model.embed_batch(texts)

    return Model(model.name, model.dim, embed_batch, model.normalised), sizes


@dataclass(frozen=True)
class IndexMeta:
    model: str
    dim: int
    normalised: bool
    created_at: str  # from an injected clock, so tests do not depend on the date

    def check(self, model: Model) -> None:
        if (self.model, self.dim) != (model.name, model.dim):
            raise IndexMismatch(f"index built by {self.model} ({self.dim} dims), query uses {model.name} "
                                f"({model.dim} dims): rebuild the index")


@dataclass(frozen=True)
class Skipped:
    id: str
    reason: str  # empty | too_short | no_words


def unusable_reason(text: str) -> str | None:
    """Why a text must not be embedded, or None. Never embed whitespace."""
    stripped = text.strip()
    if not stripped:
        return "empty"
    if len(stripped) < MIN_CHARS:
        return "too_short"
    return None if words(stripped) else "no_words"  # only stop words: the toy vector would be all zeros


def embed_batches(texts: list[str], model: Model, batch_size: int = 16) -> np.ndarray:
    """Call the model once per batch of batch_size texts, not once per text."""
    if batch_size < 1:
        raise ValueError("batch_size must be at least 1")
    parts = [model.embed_batch(texts[i : i + batch_size]) for i in range(0, len(texts), batch_size)]
    return np.vstack(parts) if parts else np.zeros((0, model.dim))


def payload_of(p: Passage) -> dict:
    return {"doc_id": p.doc_id, "title": p.title, "version": p.version, "status": p.status, "lang": p.lang}


class VectorIndex:
    def __init__(self, meta: IndexMeta, ids: list[str], matrix: np.ndarray, payloads: list[dict]):
        if matrix.shape != (len(ids), meta.dim) or len(payloads) != len(ids):
            raise IndexMismatch("ids, payloads and vectors do not fit together")
        self.meta, self.ids, self.matrix, self.payloads = meta, ids, matrix, payloads

    @classmethod
    def build(cls, passages: list[Passage], model: Model, clock: Callable[[], str],
              batch_size: int = 16) -> tuple[VectorIndex, list[Skipped]]:
        keep, skipped = [], []
        for p in passages:
            reason = unusable_reason(p.text)
            if reason:
                skipped.append(Skipped(p.id, reason))
            else:
                keep.append(p)
        matrix = embed_batches([p.embed_text for p in keep], model, batch_size)
        meta = IndexMeta(model.name, model.dim, model.normalised, clock())
        return cls(meta, [p.id for p in keep], matrix, [payload_of(p) for p in keep]), skipped

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

    def search(self, query: str, model: Model, k: int = 3, where: dict | None = None) -> list[Hit]:
        """Refuse a foreign model, filter first, then rank only what is allowed. A blank query finds nothing."""
        self.meta.check(model)
        if unusable_reason(query):
            return []
        rows = [i for i, pl in enumerate(self.payloads) if all(pl.get(key) == v for key, v in (where or {}).items())]
        if not rows:
            return []
        return top_k(model.embed_batch([query])[0], self.matrix[rows], [self.ids[i] for i in rows], k)

    def save(self, path: Path) -> None:
        header = json.dumps({"meta": asdict(self.meta), "ids": self.ids, "payloads": self.payloads})
        with open(path, "wb") as f:
            np.savez(f, matrix=self.matrix, header=np.array(header))

    @classmethod
    def load(cls, path: Path) -> VectorIndex:
        with np.load(path, allow_pickle=False) as data:
            header, matrix = json.loads(str(data["header"])), data["matrix"]
        return cls(IndexMeta(**header["meta"]), header["ids"], matrix, header["payloads"])


def ensure_index(index: VectorIndex | None, passages: list[Passage], model: Model, clock: Callable[[], str],
                 batch_size: int = 16) -> tuple[VectorIndex, bool]:
    """Keep a compatible index; otherwise rebuild from the source passages. Returns (index, rebuilt)."""
    if index is not None:
        try:
            index.meta.check(model)
            return index, False
        except IndexMismatch:
            pass
    return VectorIndex.build(passages, model, clock, batch_size)[0], True


def doc_ids(index: VectorIndex, hits: list[Hit]) -> list[str]:
    by_id = dict(zip(index.ids, index.payloads))
    return [by_id[h.passage_id]["doc_id"] for h in hits]


def recall_at_k(index: VectorIndex, model: Model, labels: list[tuple[str, str]], k: int = 3) -> tuple[float, list[str]]:
    """recall@k = questions whose expected document is among the documents of the top k passages / all questions."""
    misses = [q for q, expected in labels if expected not in doc_ids(index, index.search(q, model, k))]
    return 1 - len(misses) / len(labels), misses

from pathlib import Path

from evidencedesk.embeddings import Passage, load_corpus
from evidencedesk.index import IndexMismatch, Model, VectorIndex, counting, ensure_index, recall_at_k

CLOCK = lambda: "2026-10-12T09:00:00Z"  # injected, so the output is the same on every run
show = lambda hits: [(h.passage_id, round(h.score, 3)) for h in hits]
passages = load_corpus(Path("corpus"))
model = Model.toy()
index, skipped = VectorIndex.build(passages, model, CLOCK)
print(len(index), "indexed,", len(skipped), "skipped")
print(index.meta)

36 indexed, 0 skipped
IndexMeta(model='toy-trigram-256', dim=256, normalised=True, created_at='2026-10-12T09:00:00Z')

other = Model.toy(dim=64)  # pretend we switched to a different embedder: new name, new size
try:
    index.search("battery", other)
except IndexMismatch as error:
    print("refused:", error)
rebuilt, did_rebuild = ensure_index(index, passages, other, CLOCK)
print(did_rebuild, rebuilt.meta.model, rebuilt.matrix.shape)
same, did_rebuild = ensure_index(index, passages, model, CLOCK)
print(did_rebuild, same is index)

refused: index built by toy-trigram-256 (256 dims), query uses toy-trigram-64 (64 dims): rebuild the index
True toy-trigram-64 (36, 64)
False True

counted, sizes = counting(model)
VectorIndex.build(passages, counted, CLOCK, batch_size=16)
print("batch sizes:", sizes)
one_by_one, sizes = counting(model)
VectorIndex.build(passages, one_by_one, CLOCK, batch_size=1)
print("calls with batch_size=1:", len(sizes))

batch sizes: [16, 16, 4]
calls with batch_size=1: 36

texts = ["", "  \n ", "ok", "the of a", "Charge it fully before storing."]
odd = [Passage(f"t#s{n}", "t", "Title", "Head", text, 1, "current", "en") for n, text in enumerate(texts)]
counted, sizes = counting(model)
small, skipped = VectorIndex.build(odd, counted, CLOCK)
print("indexed:", small.ids)
print("skipped:", [(s.id, s.reason) for s in skipped])
print("texts sent to the model:", sum(sizes))
print("empty input:", len(VectorIndex.build([], model, CLOCK)[0]), "| blank query:", index.search("   ", model))

indexed: ['t#s4']
skipped: [('t#s0', 'empty'), ('t#s1', 'empty'), ('t#s2', 'too_short'), ('t#s3', 'no_words')]
texts sent to the model: 1
empty input: 0 | blank query: []

LABELS = [  # question -> the document that answers it (a smoke test: 8 questions, not a benchmark)
    ("How long do I have to return a bike?", "kc-returns-v3"),
    ("How long does the battery charge take?", "kc-battery-v1"),
    ("What is the RD2-24 recall?", "kc-recalls-v1"),
    ("How tall must I be for size M?", "kc-sizing-v1"),
    ("When do gift cards expire?", "kc-giftcards-v1"),
    ("How much does delivery cost?", "kc-shipping-v1"),
    ("How do I get my money back?", "kc-returns-v3"),  # paraphrase: shares few letters with the page
    ("Where can I try a bike before buying?", "kc-sizing-v1"),  # the page says "test ride"
]

for k in (1, 3, 5):
    recall, misses = recall_at_k(index, model, LABELS, k)
    print(f"recall@{k} = {recall:.3f}  ({len(misses)} of {len(LABELS)} questions missed)")
print("still missed at k=5:", misses)

recall@1 = 0.750  (2 of 8 questions missed)
recall@3 = 0.750  (2 of 8 questions missed)
recall@5 = 0.875  (1 of 8 questions missed)
still missed at k=5: ['Where can I try a bike before buying?']

import numpy as np

rng = np.random.default_rng(7)  # fixed seed: same numbers on every run
centres = rng.normal(size=(30, 32))  # 30 topics, 100 noisy vectors each: 3000 vectors of 32 numbers
unit = lambda m: m / np.linalg.norm(m, axis=-1, keepdims=True)
data = unit(np.repeat(centres, 100, axis=0) + 0.5 * rng.normal(size=(3000, 32)))
queries = unit(centres[rng.integers(0, 30, 50)] + 0.5 * rng.normal(size=(50, 32)))


def exact(q, k=10):  # compare with every vector
    return set(np.argsort(-(data @ q))[:k])


class Buckets:
    """Classroom model, not a real ANN library: random hyperplanes cut the space into 2**bits buckets."""

    def __init__(self, bits, seed=1):
        self.planes = np.random.default_rng(seed).normal(size=(bits, 32))
        self.weights = 1 << np.arange(bits)
        keys = ((data @ self.planes.T) > 0) @ self.weights
        self.members = {int(key): np.flatnonzero(keys == key) for key in set(keys)}

    def search(self, q, k=10, neighbours=False):
        key = int(((self.planes @ q) > 0) @ self.weights)
        keys = [key] + ([key ^ w for w in self.weights] if neighbours else [])  # own bucket (+ buckets one bit away)
        found = np.concatenate([self.members.get(x, np.array([], dtype=int)) for x in keys])
        best = found[np.argsort(-(data[found] @ q))[:k]]  # exact scoring, but only of the candidates
        return set(best), len(found)


print(f"{'method':<22}{'examined':>9}{'recall@10':>11}")
print(f"{'exact (all vectors)':<22}{'100.0%':>9}{1.0:>11.3f}")
for bits, neighbours in [(4, False), (6, False), (6, True), (8, True)]:
    index_ = Buckets(bits)
    runs = [index_.search(q, neighbours=neighbours) for q in queries]
    recall = np.mean([len(found & exact(q)) / 10 for (found, _), q in zip(runs, queries)])
    share = np.mean([n for _, n in runs]) / len(data)
    name = f"{bits} bits" + (" + neighbours" if neighbours else "")
    print(f"{name:<22}{share:>9.1%}{recall:>11.3f}")

method                 examined  recall@10
exact (all vectors)      100.0%      1.000
4 bits                     7.4%      0.482
6 bits                     2.4%      0.338
6 bits + neighbours       13.2%      0.744
8 bits + neighbours        5.2%      0.566

Q = "How long do I have to return a bike?"
for where in (None, {"status": "current"}, {"lang": "bn"}, {"doc_id": "kc-returns-v2"}):
    print(where, "->", show(index.search(Q, model, 3, where)))
after = [h for h in index.search(Q, model, 2) if index.payloads[index.ids.index(h.passage_id)]["status"] == "current"]
print("filter after top-2 keeps", len(after), "| filter before keeps", len(index.search(Q, model, 2, {"status": "current"})))

None -> [('kc-returns-v3#s1', 0.555), ('kc-returns-v2#s1', 0.463), ('kc-warranty-v2#s1', 0.403)]
{'status': 'current'} -> [('kc-returns-v3#s1', 0.555), ('kc-warranty-v2#s1', 0.403), ('kc-returns-v3#s2', 0.383)]
{'lang': 'bn'} -> [('kc-returns-v3-bn#s2', 0.197), ('kc-returns-v3-bn#s1', 0.137), ('kc-shipping-v1-bn#s2', 0.113)]
{'doc_id': 'kc-returns-v2'} -> [('kc-returns-v2#s1', 0.463), ('kc-returns-v2#s2', 0.234)]
filter after top-2 keeps 1 | filter before keeps 2

import tempfile

import numpy as np

with tempfile.TemporaryDirectory() as folder:
    path = Path(folder) / "kestrel.npz"
    index.save(path)
    reopened = VectorIndex.load(path)  # a new object, read from disk
print("same meta:", reopened.meta == index.meta, "| same vectors:", np.array_equal(reopened.matrix, index.matrix))
print("same ranking:", reopened.search(Q, model, 3) == index.search(Q, model, 3))

same meta: True | same vectors: True
same ranking: True

import tempfile

from qdrant_client import QdrantClient, models

points = [models.PointStruct(id=n, vector=index.matrix[n].tolist(), payload={"pid": pid, **index.payloads[n]})
          for n, pid in enumerate(index.ids)]


def fill(client):
    client.create_collection("kestrel", vectors_config=models.VectorParams(size=model.dim, distance=models.Distance.COSINE))
    client.upsert("kestrel", points=points)


def ask(client, question, status):
    only = models.Filter(must=[models.FieldCondition(key="status", match=models.MatchValue(value=status))])
    res = client.query_points("kestrel", query=model.embed_batch([question])[0].tolist(), limit=3, query_filter=only)
    return [(p.payload["pid"], round(p.score, 3)) for p in res.points]


Q = "How long do I have to return a bike?"
memory = QdrantClient(":memory:")
fill(memory)
print("memory  :", ask(memory, Q, "current"))
print("ours    :", show(index.search(Q, model, 3, {"status": "current"})))
memory.close()
with tempfile.TemporaryDirectory() as folder:
    disk = QdrantClient(path=folder)
    fill(disk)
    disk.close()  # release the folder lock before reopening (and before Windows deletes it)
    again = QdrantClient(path=folder)
    print("reopened:", ask(again, Q, "superseded"))
    again.close()

memory  : [('kc-returns-v3#s1', 0.555), ('kc-warranty-v2#s1', 0.403), ('kc-returns-v3#s2', 0.383)]
ours    : [('kc-returns-v3#s1', 0.555), ('kc-warranty-v2#s1', 0.403), ('kc-returns-v3#s2', 0.383)]
reopened: [('kc-returns-v2#s1', 0.463), ('kc-returns-v2#s2', 0.234)]

python -m pytest -q tests/test_index.py | sed 's/ in [0-9.]*s//'

............                                                             [100%]
12 passed

def confident(index, model, question, min_score=0.3, k=3):
    return [h for h in index.search(question, model, k, {"status": "current"}) if h.score >= min_score]


for question in ("Can I pay in instalments?", "How much does delivery cost?", "How do I get my money back?"):
    print(question, "->", show(confident(index, model, question)))

Can I pay in instalments? -> []
How much does delivery cost? -> [('kc-shipping-v1#s3', 0.539), ('kc-shipping-v1#s1', 0.532), ('kc-shipping-v1#s2', 0.465)]
How do I get my money back? -> []

v2 = Model("toy-trigram-256-v2", 256, model.embed_batch)  # same size, same vectors, but declared a new version
with tempfile.TemporaryDirectory() as folder:
    path = Path(folder) / "kestrel.npz"
    index.save(path)
    old = VectorIndex.load(path)
    try:
        old.search("battery", v2)
    except IndexMismatch as error:
        print("refused:", error)
    fresh, did_rebuild = ensure_index(old, passages, v2, CLOCK)
    fresh.save(path)
    print(did_rebuild, VectorIndex.load(path).meta.model)

refused: index built by toy-trigram-256 (256 dims), query uses toy-trigram-256-v2 (256 dims): rebuild the index
True toy-trigram-256-v2