paperless-ngx/src/other/bench_vector_store.py

#!/usr/bin/env python3
"""Head-to-head benchmark: PaperlessLanceVectorStore vs PaperlessSqliteVecVectorStore.

Run from src/ with:
    uv run python bench_vector_store.py [OPTIONS]

Phase 1 (skipped if bench_data.pkl already exists): generate fake documents with
Faker and embed chunks via Ollama; save to disk for reuse.
Phase 2: benchmark both stores against identical data and print a comparison table.

Requires both classes to coexist in paperless_ai.vector_store (Task 3 Phase A).
After Phase B replaces the file, the Lance import fails gracefully and only the
sqlite-vec half runs.
"""

from __future__ import annotations

import argparse
import pickle
import statistics
import tempfile
import time
import uuid
from pathlib import Path

import httpx
from faker import Faker
from llama_index.core.schema import TextNode
from llama_index.core.vector_stores.types import FilterOperator
from llama_index.core.vector_stores.types import MetadataFilter
from llama_index.core.vector_stores.types import MetadataFilters
from llama_index.core.vector_stores.types import VectorStoreQuery

try:
    from paperless_ai.vector_store import PaperlessLanceVectorStore

    _LANCE_OK = True
except ImportError:
    _LANCE_OK = False

from paperless_ai.vector_store import PaperlessSqliteVecVectorStore

DEFAULT_OLLAMA_URL = "http://192.168.1.87:11434"
DEFAULT_EMBED_MODEL = "qwen3-embedding:4b"
DEFAULT_DATA_FILE = "bench_data.pkl"
DEFAULT_N_DOCS = 2000
DEFAULT_CHUNKS_PER_DOC = 3
DEFAULT_QUERY_ITERS = 50
_BATCH = 32


def _embed(texts: list[str], url: str, model: str) -> list[list[float]]:
    r = httpx.post(
        f"{url}/api/embed",
        json={"model": model, "input": texts},
        timeout=120.0,
    )
    r.raise_for_status()
    return r.json()["embeddings"]


def warm_up(url: str, model: str) -> int:
    """Fire one embed call to load the model into GPU; return embedding dim."""
    print(f"Warming up {model}...", end=" ", flush=True)
    dim = len(_embed(["warm"], url, model)[0])
    print(f"dim={dim}")
    return dim


def generate_and_save(
    n_docs: int,
    chunks_per_doc: int,
    url: str,
    model: str,
    out: str,
) -> list[dict]:
    fake = Faker()
    Faker.seed(42)
    print(f"Generating {n_docs} docs ({chunks_per_doc} chunks each)...")
    docs = []
    for i in range(n_docs):
        body = "\n\n".join(fake.paragraph(nb_sentences=8) for _ in range(3))
        clen = max(1, len(body) // chunks_per_doc)
        chunks = []
        for j in range(chunks_per_doc):
            s = j * clen
            e = s + clen if j < chunks_per_doc - 1 else len(body)
            chunks.append(
                {"node_id": str(uuid.uuid4()), "text": body[s:e], "embedding": None},
            )
        docs.append(
            {
                "doc_id": str(i + 1),
                "title": fake.catch_phrase(),
                "modified": fake.date_time_this_decade().isoformat(),
                "chunks": chunks,
            },
        )

    all_texts = [c["text"] for d in docs for c in d["chunks"]]
    print(f"Embedding {len(all_texts)} chunks in batches of {_BATCH}...")
    embeddings: list[list[float]] = []
    for i in range(0, len(all_texts), _BATCH):
        embeddings.extend(_embed(all_texts[i : i + _BATCH], url, model))
        print(
            f"  {min(i + _BATCH, len(all_texts))}/{len(all_texts)}",
            end="\r",
            flush=True,
        )
    print()

    idx = 0
    for d in docs:
        for c in d["chunks"]:
            c["embedding"] = embeddings[idx]
            idx += 1

    with open(out, "wb") as f:
        pickle.dump(docs, f)
    print(f"Saved to {out}")
    return docs


def _build_nodes(docs: list[dict]) -> list[TextNode]:
    nodes = []
    for d in docs:
        for c in d["chunks"]:
            n = TextNode(
                id_=c["node_id"],
                text=c["text"],
                metadata={"document_id": d["doc_id"], "modified": d["modified"]},
            )
            n.relationships = {}
            n.embedding = c["embedding"]
            nodes.append(n)
    return nodes


def _in_filter(ids: list[str]) -> MetadataFilters:
    return MetadataFilters(
        filters=[
            MetadataFilter(key="document_id", operator=FilterOperator.IN, value=ids),
        ],
    )


def _dir_bytes(path: str) -> int:
    return sum(f.stat().st_size for f in Path(path).rglob("*") if f.is_file())


def _sqlite_bytes(uri: str) -> int:
    p = Path(uri) / "llmindex.db"
    return p.stat().st_size if p.exists() else 0


def run_bench(
    store,
    nodes: list[TextNode],
    docs: list[dict],
    q_iters: int,
    is_lance: bool,
) -> dict:
    doc_ids = [d["doc_id"] for d in docs]
    filter_ids = doc_ids[: max(1, len(doc_ids) // 5)]
    q_vecs = [nodes[i * 10 % len(nodes)].embedding for i in range(q_iters)]
    by_doc: dict[str, list[TextNode]] = {}
    for n in nodes:
        by_doc.setdefault(n.metadata["document_id"], []).append(n)
    uri = store._uri

    # insert
    t0 = time.perf_counter()
    store.add(list(nodes))
    r: dict = {"insert": time.perf_counter() - t0}

    # query plain
    times = []
    for emb in q_vecs:
        t0 = time.perf_counter()
        store.query(VectorStoreQuery(query_embedding=emb, similarity_top_k=10))
        times.append(time.perf_counter() - t0)
    r["qp50"] = statistics.median(times)
    r["qp95"] = sorted(times)[int(len(times) * 0.95)]

    # query filtered
    times = []
    flt = _in_filter(filter_ids)
    for emb in q_vecs:
        t0 = time.perf_counter()
        store.query(
            VectorStoreQuery(query_embedding=emb, similarity_top_k=10, filters=flt),
        )
        times.append(time.perf_counter() - t0)
    r["qfp50"] = statistics.median(times)
    r["qfp95"] = sorted(times)[int(len(times) * 0.95)]

    # get_modified_times
    times = []
    for _ in range(20):
        t0 = time.perf_counter()
        store.get_modified_times()
        times.append(time.perf_counter() - t0)
    r["gmt_p50"] = statistics.median(times)

    # upsert (fresh node IDs, same embeddings)
    times = []
    for doc in docs[:q_iters]:
        orig = by_doc.get(doc["doc_id"], [])
        if not orig:
            continue
        fresh = []
        for o in orig:
            fn = TextNode(
                id_=str(uuid.uuid4()),
                text=o.text,
                metadata=o.metadata.copy(),
            )
            fn.relationships = {}
            fn.embedding = o.embedding
            fresh.append(fn)
        t0 = time.perf_counter()
        store.upsert_document(doc["doc_id"], fresh)
        times.append(time.perf_counter() - t0)
    r["up50"] = statistics.median(times) if times else 0.0
    r["up95"] = sorted(times)[int(len(times) * 0.95)] if times else 0.0

    r["size_pre"] = _dir_bytes(uri) if is_lance else _sqlite_bytes(uri)

    # compact
    t0 = time.perf_counter()
    if is_lance:
        store.compact(retention_seconds=0)
    else:
        store.compact(force=True)
    r["compact"] = time.perf_counter() - t0

    r["size_post"] = _dir_bytes(uri) if is_lance else _sqlite_bytes(uri)
    return r


def _pct(lv: float | None, sv: float) -> str:
    if lv is None or lv == 0:
        return "N/A"
    p = (sv - lv) / lv * 100
    return f"{'+' if p > 0 else ''}{p:.0f}%"


def print_results(
    nodes: list[TextNode],
    q_iters: int,
    lance: dict | None,
    sq: dict,
) -> None:
    W = 30
    n, dim = len(nodes), len(nodes[0].embedding)
    print("\n=== Vector Store Benchmark ===")
    print(f"Nodes: {n} | Dim: {dim} | Query iters: {q_iters}\n")
    lh = "LanceDB" if lance else "LanceDB (N/A)"
    print(f"{'Operation':<{W}} {lh:<22} {'sqlite-vec':<22} {'Delta'}")
    print("-" * (W + 66))

    def _s(v: float) -> str:
        return f"{v:.3f}s"

    def _ms(v: float) -> str:
        return f"{v * 1000:.1f}ms"

    def _mb(v: float) -> str:
        return f"{v / 1e6:.1f} MB"

    def row(label: str, lv: float | None, sv: float, fmt) -> None:
        ls = fmt(lv) if lv is not None else "N/A"
        print(f"{label:<{W}} {ls:<22} {fmt(sv):<22} {_pct(lv, sv)}")

    def row2(
        label: str,
        lv1: float | None,
        lv2: float | None,
        sv1: float,
        sv2: float,
    ) -> None:
        def ms_pair(a: float, b: float) -> str:
            return f"{_ms(a)} / {_ms(b)}"

        ls = ms_pair(lv1, lv2) if lv1 is not None else "N/A"
        print(f"{label:<{W}} {ls:<22} {ms_pair(sv1, sv2):<22} {_pct(lv1, sv1)}")

    L = lance
    row(f"insert ({n} nodes)", L["insert"] if L else None, sq["insert"], _s)
    row2(
        "query plain p50/p95",
        L["qp50"] if L else None,
        L["qp95"] if L else None,
        sq["qp50"],
        sq["qp95"],
    )
    row2(
        "query filtered p50/p95",
        L["qfp50"] if L else None,
        L["qfp95"] if L else None,
        sq["qfp50"],
        sq["qfp95"],
    )
    row("get_modified_times p50", L["gmt_p50"] if L else None, sq["gmt_p50"], _ms)
    row2(
        "upsert p50/p95",
        L["up50"] if L else None,
        L["up95"] if L else None,
        sq["up50"],
        sq["up95"],
    )
    row("compact", L["compact"] if L else None, sq["compact"], _s)
    row("file size pre-compact", L["size_pre"] if L else None, sq["size_pre"], _mb)
    row("file size post-compact", L["size_post"] if L else None, sq["size_post"], _mb)


def main() -> None:
    ap = argparse.ArgumentParser(description="Vector store head-to-head benchmark")
    ap.add_argument("--n-docs", type=int, default=DEFAULT_N_DOCS)
    ap.add_argument("--chunks-per-doc", type=int, default=DEFAULT_CHUNKS_PER_DOC)
    ap.add_argument("--data-file", default=DEFAULT_DATA_FILE)
    ap.add_argument("--regenerate", action="store_true")
    ap.add_argument("--ollama-url", default=DEFAULT_OLLAMA_URL)
    ap.add_argument("--embed-model", default=DEFAULT_EMBED_MODEL)
    ap.add_argument("--query-iters", type=int, default=DEFAULT_QUERY_ITERS)
    args = ap.parse_args()

    warm_up(args.ollama_url, args.embed_model)

    data_path = Path(args.data_file)
    if args.regenerate or not data_path.exists():
        docs = generate_and_save(
            args.n_docs,
            args.chunks_per_doc,
            args.ollama_url,
            args.embed_model,
            args.data_file,
        )
    else:
        print(f"Loading {args.data_file}...")
        with open(data_path, "rb") as f:
            docs = pickle.load(f)
        print(f"Loaded {len(docs)} docs ({sum(len(d['chunks']) for d in docs)} nodes)")

    all_nodes = _build_nodes(docs)

    lance_r = None
    if _LANCE_OK:
        print("\nBenchmarking LanceDB...")
        with tempfile.TemporaryDirectory() as d:
            store = PaperlessLanceVectorStore(uri=d)
            lance_r = run_bench(store, all_nodes, docs, args.query_iters, is_lance=True)
    else:
        print("Skipping LanceDB (PaperlessLanceVectorStore not importable).")

    print("\nBenchmarking sqlite-vec...")
    with tempfile.TemporaryDirectory() as d:
        store = PaperlessSqliteVecVectorStore(uri=d)
        sqlite_r = run_bench(store, all_nodes, docs, args.query_iters, is_lance=False)

    print_results(all_nodes, args.query_iters, lance_r, sqlite_r)


if __name__ == "__main__":
    main()