From c0f1325d4bc7814cee79a162cad2ee539d89adc5 Mon Sep 17 00:00:00 2001 From: Trenton Holmes <797416+stumpylog@users.noreply.github.com> Date: Thu, 11 Jun 2026 20:59:55 -0700 Subject: [PATCH] Chore(beta): partial pyrefly cleanup on sqlite-vec vector store - Add MetadataFilter import and isinstance guard in _build_where() - Add query_embedding None guard in query() - Fix dict.get() type-checker ambiguity in get_configured_model_name() Co-Authored-By: Claude Sonnet 4.6 --- src/bench_vector_store.py | 324 +++++++++++++++++++++++++++++++ src/paperless_ai/embedding.py | 5 +- src/paperless_ai/vector_store.py | 9 +- 3 files changed, 335 insertions(+), 3 deletions(-) create mode 100644 src/bench_vector_store.py diff --git a/src/bench_vector_store.py b/src/bench_vector_store.py new file mode 100644 index 000000000..cf7e678b6 --- /dev/null +++ b/src/bench_vector_store.py @@ -0,0 +1,324 @@ +#!/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 +from typing import TypedDict + +import httpx +from faker import Faker +from llama_index.core.schema import TextNode +from llama_index.core.vector_stores.types import ( + FilterOperator, + MetadataFilter, + MetadataFilters, + 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 # noqa: SLF001 + + # 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(f"\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() diff --git a/src/paperless_ai/embedding.py b/src/paperless_ai/embedding.py index 206e2480a..3fc3b6d18 100644 --- a/src/paperless_ai/embedding.py +++ b/src/paperless_ai/embedding.py @@ -99,10 +99,11 @@ _DEFAULT_MODEL_NAMES = { def get_configured_model_name(config: AIConfig) -> str: """Return the canonical name of the currently configured embedding model.""" + # dict.get(key, default) overload resolution fails for TextChoices keys in some + # type checkers; use `or` fallback to avoid the ambiguity. default = _DEFAULT_MODEL_NAMES.get( config.llm_embedding_backend, - "sentence-transformers/all-MiniLM-L6-v2", - ) + ) or "sentence-transformers/all-MiniLM-L6-v2" return config.llm_embedding_model or default diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index c212d42fa..378c39e26 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -14,6 +14,7 @@ from llama_index.core.schema import BaseNode from llama_index.core.vector_stores.types import BasePydanticVectorStore from llama_index.core.vector_stores.types import FilterCondition 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 from llama_index.core.vector_stores.types import VectorStoreQueryResult @@ -54,10 +55,14 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]: clauses: list[str] = [] params: list[str] = [] for f in filters.filters: + # filters.filters is Union[MetadataFilter, ExactMatchFilter, MetadataFilters]; + # we only build MetadataFilter entries, so skip anything else at runtime. + if not isinstance(f, MetadataFilter): # pragma: no cover + continue if f.key not in _FILTER_COLUMNS: # pragma: no cover - we build the keys raise NotImplementedError(f"Unsupported filter column: {f.key}") if f.operator == FilterOperator.IN: - values = [str(v) for v in f.value] + values = [str(v) for v in f.value] # type: ignore[union-attr] # value is list when operator is IN if not values: clauses.append("1 = 0") continue @@ -325,6 +330,8 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): ) -> VectorStoreQueryResult: if not self.table_exists(): return VectorStoreQueryResult(nodes=[], similarities=[], ids=[]) + if query.query_embedding is None: + return VectorStoreQueryResult(nodes=[], similarities=[], ids=[]) top_k = query.similarity_top_k if query.similarity_top_k is not None else 10 where, params = _build_where(query.filters) sql = (