From cd9851a12c000199ce5c211f937ef6c1e0138efa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Thu, 8 Oct 2026 14:05:27 -0400 Subject: [PATCH] Keep the vectors beside the atoms instead of rebuilding them per query The dense leg used to be a pair of files next to the database: a full .npy re-encoded from every atom, re-normalized on every query, and stale the moment anything was ingested. It is now a vec0 table inside the same SQLite file, written in the same transaction as the atom it describes, so a memory is queryable the instant it is remembered. - vector_index.py maintains the table; the writer inserts, the reader never rebuilds. A repair scan (ely-eye build-dense-index) catches what a crash hid, and it is not on the search path. - Ranking stays in the engine; every boundary (closed, retired, not yet true, another holder's private memory, the mounted cartridge) is checked once against the atoms table afterwards. A visibility subquery inside the KNN cost 48.9s on 153,857 vectors; the same check over 200 candidates costs 0.2ms. When the boundaries reject enough candidates the ask widens once, never twice. - The Go mirror and the web panel carry the new status fields. - bench/dense_scale.py and bench/quant_probe.py measure the leg at corpus size; bench/scale_check.py now works on a copy so the frozen baseline stays frozen. Measured on the frozen 165,973-atom corpus. bench/reports/DENSE.md, QUANT.md and SCALE.md carry the numbers and the machine they were taken on. Gates: 285 pytest, ruff clean, pyright 0 on the touched files, go build/vet/test and gofmt clean, frontend build and eslint clean. mypy is 25 findings against the parent's 24: the new one is "Skipping analyzing sqlite_vec: missing library stubs", the same class as the two existing qwen_vl_utils lines, and the repo keeps no type: ignore anywhere. --- backend/pyproject.toml | 1 + backend/src/ely_eye/cli_senses.py | 8 +- backend/src/ely_eye/config.py | 1 + backend/src/ely_eye/db.py | 3 + backend/src/ely_eye/embeddings.py | 98 +++--- backend/src/ely_eye/ingestion.py | 4 + backend/src/ely_eye/maintenance.py | 2 +- backend/src/ely_eye/memory.py | 33 +- backend/src/ely_eye/migrations.py | 10 +- backend/src/ely_eye/retrieval.py | 38 ++- backend/src/ely_eye/schemas.py | 3 +- backend/src/ely_eye/senses.py | 19 +- backend/src/ely_eye/status_schemas.py | 2 +- backend/src/ely_eye/vector_index.py | 424 ++++++++++++++++++++++++ backend/tests/conftest.py | 41 +++ backend/tests/test_maintenance.py | 26 ++ backend/tests/test_memory_write.py | 15 +- backend/tests/test_vector_index.py | 291 ++++++++++++++++ backend/tests/test_visual_index.py | 2 +- bench/dense_scale.py | 318 ++++++++++++++++++ bench/quant_probe.py | 332 +++++++++++++++++++ bench/scale_check.py | 36 +- ely-eye-mcp/internal/domain/types.go | 3 +- frontend/src/components/MemoryPanel.tsx | 3 +- frontend/src/components/SensesPanel.tsx | 2 +- frontend/src/types.ts | 5 +- 26 files changed, 1602 insertions(+), 118 deletions(-) create mode 100644 backend/src/ely_eye/vector_index.py create mode 100644 backend/tests/test_vector_index.py create mode 100644 bench/dense_scale.py create mode 100644 bench/quant_probe.py diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 1dded57..ab9cc7a 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -34,6 +34,7 @@ dependencies = [ "rapidocr-onnxruntime>=1.2.3,<2", "safetensors>=0.7.0", "sentence-transformers>=5.2.0", + "sqlite-vec>=0.1.9", "torch>=2.12.0", "torchvision>=0.27.0", "transformers>=4.57.0", diff --git a/backend/src/ely_eye/cli_senses.py b/backend/src/ely_eye/cli_senses.py index 44ad1d5..5ed249d 100644 --- a/backend/src/ely_eye/cli_senses.py +++ b/backend/src/ely_eye/cli_senses.py @@ -18,7 +18,7 @@ from .cli_common import json_dump from .config import Settings from .db import Database from .describe import DEFERRED, VisionDescriber -from .embeddings import EmbeddingService +from .vector_index import VectorIndex from .visual_index import VisualIndex @@ -73,7 +73,5 @@ def register(app: typer.Typer, settings: Callable[[], Settings]) -> None: @app.command("build-dense-index") def build_dense_index() -> None: - """Encode every atom into the dense index the text retriever falls back to.""" - index = EmbeddingService(settings()).build_index() - dimensions = int(index.vectors.shape[1]) if index.vectors.size else 0 - typer.echo(json_dump({"atoms": len(index.atom_ids), "dimensions": dimensions})) + """Encode the memories the dense index is missing; its vectors live beside the atoms.""" + typer.echo(json_dump(VectorIndex(settings()).sync())) diff --git a/backend/src/ely_eye/config.py b/backend/src/ely_eye/config.py index b556c88..af0b2ec 100644 --- a/backend/src/ely_eye/config.py +++ b/backend/src/ely_eye/config.py @@ -18,6 +18,7 @@ class Settings(BaseSettings): kv_cache_dtype: str = "auto" planner_model_id: str = "Qwen/Qwen3.5-0.8B" embedding_model_id: str = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2" + embedding_model_dir: Path | None = None clip_model_dir: Path | None = None vlm_model_dir: Path | None = None vlm_max_new_tokens: int = 96 diff --git a/backend/src/ely_eye/db.py b/backend/src/ely_eye/db.py index 3282ad1..cfc9c81 100644 --- a/backend/src/ely_eye/db.py +++ b/backend/src/ely_eye/db.py @@ -23,8 +23,11 @@ class Database: @contextmanager def connect(self) -> Iterator[sqlite3.Connection]: + from .vector_index import load_extension + connection = sqlite3.connect(self.path) connection.row_factory = sqlite3.Row + load_extension(connection) connection.execute("PRAGMA foreign_keys = ON") connection.execute("PRAGMA journal_mode = WAL") try: diff --git a/backend/src/ely_eye/embeddings.py b/backend/src/ely_eye/embeddings.py index fc9af9a..acf35dd 100644 --- a/backend/src/ely_eye/embeddings.py +++ b/backend/src/ely_eye/embeddings.py @@ -1,7 +1,8 @@ +"""The text side of the senses: one model, one unit vector per memory that has words.""" + from __future__ import annotations -import json -from dataclasses import dataclass +from pathlib import Path from typing import Any import numpy as np @@ -11,80 +12,63 @@ from .db import Database from .schemas import EvidenceAtom -@dataclass -class DenseIndex: - atom_ids: list[str] - vectors: np.ndarray - - def search(self, query_vector: np.ndarray, top_k: int) -> list[tuple[str, float]]: - if self.vectors.size == 0: - return [] - query = normalize(query_vector.reshape(1, -1))[0] - vectors = normalize(self.vectors) - scores = vectors @ query - order = np.argsort(scores)[::-1][:top_k] - return [(self.atom_ids[index], float(scores[index])) for index in order] - - class EmbeddingService: - def __init__(self, settings: Settings | None = None, database: Database | None = None) -> None: + """`model` is injectable so a test can rank without loading a 500 MB snapshot.""" + + def __init__( + self, settings: Settings | None = None, database: Database | None = None, model: Any | None = None + ) -> None: self.settings = settings or get_settings() self.db = database or Database(self.settings) - self.index_path = self.settings.cache_dir / "dense_index.npy" - self.map_path = self.settings.cache_dir / "dense_index_atoms.json" - self._model: Any | None = None + self._model = model + self._dim: int | None = None def enabled(self) -> bool: return self.settings.dense_embeddings_enabled - def build_index(self, atoms: list[EvidenceAtom] | None = None) -> DenseIndex: - atoms = atoms or self.db.list_atoms() - if not atoms: - index = DenseIndex(atom_ids=[], vectors=np.zeros((0, 0), dtype=np.float32)) - self._write_index(index) - return index + def encode_texts(self, texts: list[str]) -> np.ndarray: + """Unit-normalized float32 rows, in the order given.""" + if not texts: + raise ValueError("encode_texts was asked to encode nothing") + vectors = self._load_model().encode( + [{"text": text} for text in texts], normalize_embeddings=True, show_progress_bar=False + ) + matrix = np.asarray(vectors, dtype=np.float32) + if matrix.ndim != 2: + raise ValueError(f"the embedding model returned a {matrix.shape} array, expected rows of vectors") + return normalize(matrix) - model = self._load_model() - payloads = [self._atom_payload(atom) for atom in atoms] - vectors = model.encode(payloads, normalize_embeddings=True, show_progress_bar=True) - index = DenseIndex(atom_ids=[atom.atom_id for atom in atoms], vectors=np.asarray(vectors, dtype=np.float32)) - self._write_index(index) - return index - - def load_index(self) -> DenseIndex | None: - if not self.index_path.exists() or not self.map_path.exists(): - return None - atom_ids = json.loads(self.map_path.read_text(encoding="utf-8")) - vectors = np.load(self.index_path) - return DenseIndex(atom_ids=atom_ids, vectors=vectors) + def encode_atoms(self, atoms: list[EvidenceAtom]) -> np.ndarray: + return self.encode_texts([_text_of(atom) for atom in atoms]) def embed_query(self, query: str) -> np.ndarray: - model = self._load_model() - vector = model.encode([{"text": query}], normalize_embeddings=True) - return np.asarray(vector[0], dtype=np.float32) + return self.encode_texts([query])[0] + + def probe_dim(self) -> int: + """The width this model produces, asked before the vector table is created.""" + if self._dim is None: + self._dim = int(self.encode_texts(["dimension probe"]).shape[1]) + return self._dim def _load_model(self) -> Any: if self._model is None: from sentence_transformers import SentenceTransformer - self._model = SentenceTransformer(self.settings.embedding_model_id, trust_remote_code=True) + self._model = SentenceTransformer(str(self._model_ref()), trust_remote_code=True) return self._model - def _write_index(self, index: DenseIndex) -> None: - self.index_path.parent.mkdir(parents=True, exist_ok=True) - np.save(self.index_path, index.vectors) - self.map_path.write_text(json.dumps(index.atom_ids, ensure_ascii=False, indent=2), encoding="utf-8") + def _model_ref(self) -> str | Path: + """A local snapshot when one is set; otherwise the hub id, which SentenceTransformer caches.""" + directory = self.settings.embedding_model_dir + if directory is None: + return self.settings.embedding_model_id + if not directory.exists(): + raise SystemExit(f"{directory}: the embedding model has not been downloaded") + return directory - def _atom_payload(self, atom: EvidenceAtom) -> dict[str, Any]: - if atom.image_ref: - from .storage import ObjectStore - image_path = ObjectStore(self.settings).resolve(atom.image_ref) - payload: dict[str, Any] = {"image": str(image_path)} - if atom.text: - payload["text"] = atom.text - return payload - return {"text": atom.text or atom.source} +def _text_of(atom: EvidenceAtom) -> str: + return atom.text or "" def normalize(vectors: np.ndarray) -> np.ndarray: diff --git a/backend/src/ely_eye/ingestion.py b/backend/src/ely_eye/ingestion.py index 05e42a3..8957dce 100644 --- a/backend/src/ely_eye/ingestion.py +++ b/backend/src/ely_eye/ingestion.py @@ -37,6 +37,7 @@ from .storage import ( sha256_file, token_equivalent, ) +from .vector_index import VectorIndex class IngestionService(MediaIngest, RepositoryIngest): @@ -45,12 +46,14 @@ class IngestionService(MediaIngest, RepositoryIngest): settings: Settings | None = None, database: Database | None = None, object_store: ObjectStore | None = None, + vectors: VectorIndex | None = None, ) -> None: self.settings = settings or get_settings() self.db = database or Database(self.settings) self.objects = object_store or ObjectStore(self.settings) self.describer = VisionDescriber(self.settings) self.transcriber = SpeechTranscriber(self.settings) + self.vectors = vectors or VectorIndex(self.settings, self.db) def ingest_path(self, path: Path, cartridge_name: str | None = None) -> IngestResponse: path = path.expanduser().resolve() @@ -77,6 +80,7 @@ class IngestionService(MediaIngest, RepositoryIngest): self.db.upsert_sources(batch.sources) self.db.upsert_atoms(batch.atoms) + self.vectors.sync_atoms([atom.atom_id for atom in batch.atoms]) cartridge_id = None if cartridge_name: diff --git a/backend/src/ely_eye/maintenance.py b/backend/src/ely_eye/maintenance.py index 102e6b6..15ddeaa 100644 --- a/backend/src/ely_eye/maintenance.py +++ b/backend/src/ely_eye/maintenance.py @@ -201,7 +201,7 @@ def _check_indexes(facts: SenseStatus, warnings: list[str]) -> None: warnings.append("dense embeddings are on but no index exists; run `ely-eye build-dense-index`") elif facts.dense == "behind": warnings.append( - f"the dense index covers {facts.dense_covered} atoms while the store holds {facts.store_atoms}; " + f"the dense index holds {facts.dense_covered} of the {facts.dense_expected} memories with words; " "run `ely-eye build-dense-index`" ) diff --git a/backend/src/ely_eye/memory.py b/backend/src/ely_eye/memory.py index 06e696d..f2e9d0d 100644 --- a/backend/src/ely_eye/memory.py +++ b/backend/src/ely_eye/memory.py @@ -27,6 +27,7 @@ from .schemas import ( utc_now, ) from .storage import token_equivalent +from .vector_index import VectorIndex from .visual_index import VisualIndex # Trust scores describe extraction fidelity, not truthfulness: a remembered sentence is stored verbatim, so origin carries the epistemic weight. @@ -34,9 +35,16 @@ MEMORY_TRUST = TrustScores(parser=1.0, ocr=0.0, model=0.0) class MemoryService: - def __init__(self, settings: Settings | None = None, database: Database | None = None) -> None: + def __init__( + self, + settings: Settings | None = None, + database: Database | None = None, + vectors: VectorIndex | None = None, + ) -> None: self.settings = settings or get_settings() self.db = database or Database(self.settings) + self.vectors = vectors or VectorIndex(self.settings, self.db) + self.pictures = VisualIndex(self.settings, self.db) def remember( self, @@ -107,6 +115,7 @@ class MemoryService: retention_until=retention_until, ) self.db.upsert_atoms([atom]) + self.vectors.sync_atoms([atom.atom_id]) return atom def supersede( @@ -148,6 +157,7 @@ class MemoryService: valid_from=moment, ) self.db.close_atoms([atom_id], moment) + self.vectors.sync_atoms([atom_id]) return replacement def forget( @@ -175,26 +185,27 @@ class MemoryService: edges_deleted = self.db.delete_edges_touching(doomed_ids) atoms_deleted = self.db.delete_atoms(doomed_ids) cache_cleared = self.db.clear_cache_events(f"forget:{target_source}") - vectors_dropped = self._drop_vector_index() + vectors_dropped, visual_dropped = self._drop_derived_vectors(doomed_ids) return ForgetReport( source_id=target_source, atom_ids=doomed_ids, atoms_deleted=atoms_deleted, edges_deleted=edges_deleted, cache_entries_cleared=cache_cleared, - vector_index_dropped=vectors_dropped, + vectors_dropped=vectors_dropped, + visual_index_dropped=visual_dropped, affected_cartridges=self._cartridges_holding(target_source), reason=reason, ) - def _drop_vector_index(self) -> bool: - """The text and picture indexes are derived data; a stale one keeps answering with deleted memories.""" - dropped = False - for path in (self.settings.cache_dir / "dense_index.npy", self.settings.cache_dir / "dense_index_atoms.json"): - if path.exists(): - path.unlink() - dropped = True - return VisualIndex(self.settings, self.db).drop() or dropped + def _drop_derived_vectors(self, atom_ids: list[str]) -> tuple[int, bool]: + """A forgotten memory must stop answering from the derived indexes too. + + The dense vectors live in the database beside the atoms, so they go one memory at a + time. The picture index is one file for the whole store, so it is dropped and rebuilt. + """ + dropped = self.vectors.drop_atoms(atom_ids) + return dropped, self.pictures.drop() def _cartridges_holding(self, source_id: str) -> list[dict[str, Any]]: """Name the cartridges that exported this source, so they can be rebuilt. diff --git a/backend/src/ely_eye/migrations.py b/backend/src/ely_eye/migrations.py index f0057fd..4d046e6 100644 --- a/backend/src/ely_eye/migrations.py +++ b/backend/src/ely_eye/migrations.py @@ -16,7 +16,7 @@ import sqlite3 from .storage import token_equivalent, visual_token_equivalent from .text_index import index_text -SCHEMA_VERSION = 7 +SCHEMA_VERSION = 8 ATOM_COLUMNS_V2 = """ atom_id TEXT PRIMARY KEY, @@ -125,6 +125,14 @@ _BASE_TABLES = """ documents INTEGER NOT NULL ); + -- What the dense vector table was built from. vec0 cannot be altered in place, so when the + -- model or the width changes the table is dropped and re-encoded; these three values are + -- how the next start notices it has to. + CREATE TABLE IF NOT EXISTS vector_index_meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + ); + CREATE VIRTUAL TABLE IF NOT EXISTS atom_fts USING fts5( atom_id UNINDEXED, source, diff --git a/backend/src/ely_eye/retrieval.py b/backend/src/ely_eye/retrieval.py index 02f2b87..c8a5c37 100644 --- a/backend/src/ely_eye/retrieval.py +++ b/backend/src/ely_eye/retrieval.py @@ -10,7 +10,9 @@ from __future__ import annotations import base64 from collections import defaultdict +from datetime import datetime from pathlib import Path +from typing import Any from . import cartridge_store from .config import Settings, get_settings @@ -18,6 +20,7 @@ from .db import Database from .embeddings import EmbeddingService from .schemas import EvidenceAtom, RetrievalHit, SearchFilters, SearchResult from .store_search import FtsPage, search_atoms +from .vector_index import VectorIndex from .visual_index import VisualIndex RRF_K = 60 @@ -31,6 +34,7 @@ class RetrievalService: self.settings = settings or get_settings() self.db = database or Database(self.settings) self.embeddings = EmbeddingService(self.settings, self.db) + self.vectors = VectorIndex(self.settings, self.db, self.embeddings) self.visual = VisualIndex(self.settings, self.db) def search( @@ -70,10 +74,7 @@ class RetrievalService: for atom_id, bm25 in page.rows: scores[atom_id]["sparse_score"] = bm25 if self.embeddings.enabled(): - dense = self._dense_search(query, limit * CANDIDATE_WINDOW) - if scope: - allowed = cartridge_store.atoms_of_cartridges(self.db, scope) - dense = {atom_id: score for atom_id, score in dense.items() if atom_id in allowed} + dense = self._dense_search(query, limit * CANDIDATE_WINDOW, criteria) rankings.append(list(dense)) for atom_id, cosine in dense.items(): scores[atom_id]["dense_score"] = cosine @@ -97,12 +98,26 @@ class RetrievalService: next_cursor=encode_cursor(offset + limit) if hits and more else None, ) - def _dense_search(self, query: str, limit: int) -> dict[str, float]: - index = self.embeddings.load_index() - if index is None: - index = self.embeddings.build_index() + def _dense_search(self, query: str, limit: int, criteria: dict[str, Any]) -> dict[str, float]: + """The vectors sit beside the atoms, so a memory written a moment ago is findable. + + Every filter is applied inside the engine, in the same statement as the nearest-neighbour + walk, so a scoped search cannot return a memory the caller was never allowed to see. The + index catches up inside `search`, so nothing here has to rebuild it first. + """ query_vector = self.embeddings.embed_query(query) - return dict(index.search(query_vector, limit)) + return self.vectors.search( + query_vector, + limit, + holder=criteria.get("holder"), + scenario=criteria.get("scenario"), + source_id=criteria.get("source_id"), + modality=criteria.get("modality"), + since=_text(criteria.get("since")), + until=_text(criteria.get("until")), + cartridge_ids=criteria.get("cartridge_ids"), + as_holder=criteria.get("as_holder"), + ) def _visual_search( self, @@ -183,3 +198,8 @@ def decode_cursor(cursor: str | None) -> int: if offset < 0: raise ValueError(f"search cursor is unreadable: {cursor!r}") return offset + + +def _text(moment: datetime | None) -> str | None: + """A time bound reaches SQLite as the same ISO text the atoms table stores.""" + return None if moment is None else moment.isoformat() diff --git a/backend/src/ely_eye/schemas.py b/backend/src/ely_eye/schemas.py index 3582146..b823bea 100644 --- a/backend/src/ely_eye/schemas.py +++ b/backend/src/ely_eye/schemas.py @@ -431,7 +431,8 @@ class ForgetReport(BaseModel): atoms_deleted: int edges_deleted: int cache_entries_cleared: int - vector_index_dropped: bool + vectors_dropped: int + visual_index_dropped: bool affected_cartridges: list[dict[str, Any]] reason: str diff --git a/backend/src/ely_eye/senses.py b/backend/src/ely_eye/senses.py index 98e86ff..69480b3 100644 --- a/backend/src/ely_eye/senses.py +++ b/backend/src/ely_eye/senses.py @@ -6,8 +6,8 @@ from pathlib import Path from .config import Settings from .db import Database -from .embeddings import EmbeddingService from .status_schemas import DenseState, SenseState, SenseStatus, VisualState +from .vector_index import VectorIndex, expected_vectors from .visual_index import VisualIndex @@ -20,7 +20,7 @@ def sense_facts(settings: Settings, database: Database, visual: VisualIndex | No return SenseStatus( dense=dense, dense_covered=covered, - store_atoms=total, + dense_expected=total, picture_atoms=claimed, picture_bytes_missing=max(0, claimed - len(on_disk)), visual=_visual_state(index, on_disk), @@ -34,14 +34,15 @@ def sense_facts(settings: Settings, database: Database, visual: VisualIndex | No def _dense_facts(settings: Settings, database: Database) -> tuple[DenseState, int, int]: - total = database.atom_count() + expected = expected_vectors(database) if not settings.dense_embeddings_enabled: - return "off", 0, total - loaded = EmbeddingService(settings, database).load_index() - if loaded is None: - return "missing", 0, total - covered = len(loaded.atom_ids) - return ("current" if covered == total else "behind"), covered, total + return "off", 0, expected + facts = VectorIndex(settings, database).stats() + if not facts["built"]: + return "missing", 0, expected + # A memory with no words belongs to the picture index, so the honest denominator is the + # memories this index is meant to hold, not every row in the store. + return ("current" if facts["pending"] == 0 else "behind"), int(facts["vectors"]), expected def _visual_state(visual: VisualIndex, pictures: set[str]) -> VisualState: diff --git a/backend/src/ely_eye/status_schemas.py b/backend/src/ely_eye/status_schemas.py index 2d3a3dd..b5b3d4e 100644 --- a/backend/src/ely_eye/status_schemas.py +++ b/backend/src/ely_eye/status_schemas.py @@ -52,7 +52,7 @@ class SenseStatus(BaseModel): dense: DenseState dense_covered: int - store_atoms: int + dense_expected: int picture_atoms: int picture_bytes_missing: int visual: VisualState diff --git a/backend/src/ely_eye/vector_index.py b/backend/src/ely_eye/vector_index.py new file mode 100644 index 0000000..126677a --- /dev/null +++ b/backend/src/ely_eye/vector_index.py @@ -0,0 +1,424 @@ +"""The dense leg lives in the database, so a memory is findable the moment it is written. + +The first version kept the vectors in `cache/dense_index.npy` plus a JSON list of atom ids. +That pair had no owner: nothing noticed when an atom was added, changed or forgotten, so a +search answered from a snapshot of the past and `forget` had to delete the whole file. Here the +vectors sit in the same SQLite file as the atoms, in a `vec0` table the `sqlite-vec` extension +answers, and the only rule the index has is "one row per current atom, re-encode what changed". +""" + +from __future__ import annotations + +import hashlib +import sqlite3 +import time +from typing import Any + +import numpy as np + +from .cartridge_store import atoms_of_cartridges +from .config import Settings, get_settings +from .db import Database +from .embeddings import EmbeddingService +from .row_codec import to_json + +# The vec0 schema cannot be altered in place, so the format tag is load-bearing: changing the +# vector width or the model means dropping the table and re-encoding, not migrating it. +FORMAT = "ely-eye-dense-vectors-v1" +TABLE = "vec_atoms" +META_TABLE = "vector_index_meta" +BATCH = 64 +# A vec0 table has no index on atom_id, so a delete by id scans the table: ids go in as sets. +DELETE_CHUNK = 20_000 +# The engine is asked wider than the answer: the boundary check afterwards drops candidates. +CANDIDATE_FLOOR = 100 +# A selective filter thins the wide ask, so the search widens once — never twice. +CANDIDATE_CEILING = 1_024 +# A scope is narrowed inside the engine by id list only while it is small enough to name. +SCOPE_NARROW_MAX = 2_000 + + +def load_extension(connection: sqlite3.Connection) -> None: + """Register sqlite-vec and the stamp function on one connection; every connection that + touches vec0 needs both, and the stamp is compared in SQL so a sync never reads the + whole corpus of text into Python.""" + if not hasattr(connection, "enable_load_extension"): + raise SystemExit( + "this Python's sqlite3 cannot load extensions, and the dense index needs sqlite-vec; " + "use a sqlite3 built with SQLITE_ENABLE_LOAD_EXTENSION" + ) + connection.enable_load_extension(True) + try: + import sqlite_vec + except ImportError as error: + raise SystemExit( + "the dense index needs the sqlite-vec extension; install it with `pip install sqlite-vec`" + ) from error + sqlite_vec.load(connection) + connection.create_function("atom_stamp", 2, atom_stamp, deterministic=True) + + +class VectorIndex: + """One float32 vector per current atom: the writer keeps it, the reader asks it questions.""" + + def __init__( + self, settings: Settings | None = None, database: Database | None = None, embeddings: EmbeddingService | None = None + ) -> None: + self.settings = settings or get_settings() + self.db = database or Database(self.settings) + self.embeddings = embeddings or EmbeddingService(self.settings, self.db) + + def ensure(self, connection: sqlite3.Connection, dim: int) -> None: + """Create the table, or rebuild it when the model behind the vectors changed.""" + meta = self._meta(connection) + wanted = {"format": FORMAT, "dim": str(dim), "model": self.settings.embedding_model_id} + if meta and all(meta.get(key) == value for key, value in wanted.items()): + return + if meta: + self._rebuild(connection) + connection.executescript( + """ + CREATE VIRTUAL TABLE IF NOT EXISTS vec_atoms USING vec0( + embedding float[{dim}] distance_metric=cosine, + atom_id TEXT, + stamp TEXT + ); + """.replace("{dim}", str(dim)) + ) + for key, value in wanted.items(): + connection.execute( + "INSERT INTO vector_index_meta (key, value) VALUES (?, ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value", + (key, value), + ) + + def sync(self) -> dict[str, Any]: + """Encode what is missing or stale; the repair path, not the per-query path.""" + started = time.monotonic() + with self.db.connect() as connection: + dim = self.embeddings.probe_dim() + self.ensure(connection, dim) + stale = self._stale(connection) + if not stale: + return {"encoded": 0, "dim": dim, "model": self.settings.embedding_model_id, "seconds": 0.0} + vectors = self.embeddings.encode_atoms([atom for _, atom in stale]) + self._write(connection, stale, vectors) + return { + "encoded": len(stale), + "dim": int(dim), + "model": self.settings.embedding_model_id, + "seconds": round(time.monotonic() - started, 2), + } + + def sync_atoms(self, atom_ids: list[str]) -> dict[str, Any]: + """Bring the index in step with memories as they are written. + + Nothing on the search path rebuilds, so this is the half that keeps "written now, + findable now" true; a memory that has closed or been forgotten loses its vector here too. + """ + wanted = [atom_id for atom_id in atom_ids if atom_id] + if not wanted or not self.settings.dense_embeddings_enabled: + return {"encoded": 0, "dropped": 0} + started = time.monotonic() + with self.db.connect() as connection: + dim = self.embeddings.probe_dim() + self.ensure(connection, dim) + current = self._current(connection, wanted) + kept = {atom_id for atom_id, _ in current} + dropped = self._delete(connection, [atom_id for atom_id in wanted if atom_id not in kept]) + if current: + self._write(connection, current, self.embeddings.encode_atoms([atom for _, atom in current])) + return {"encoded": len(current), "dropped": dropped, "seconds": round(time.monotonic() - started, 3)} + + def search( + self, + query_vector: np.ndarray, + limit: int, + *, + holder: str | None = None, + scenario: str | None = None, + source_id: str | None = None, + modality: str | None = None, + since: str | None = None, + until: str | None = None, + cartridge_ids: list[str] | None = None, + as_holder: str | None = None, + ) -> dict[str, float]: + """Closest memories a search is allowed to show, cosine distance turned into a score. + + The engine answers nearest-neighbour and nothing else; every boundary — closed, retired, + not yet true, another holder's private memory, the cartridge scope — is checked once + against the atoms table afterwards. Measured on 153,857 stored vectors: folding those + boundaries into the KNN as a subquery over `v_atoms_current` costs 48.9 seconds a search, + and checking them afterwards over the candidates costs 0.2 milliseconds. + + The writer keeps the index current (`sync_atoms`), so a search never waits for a rebuild; + `stats` and `doctor` say when a store has memories the index has never seen. + """ + if limit <= 0: + raise ValueError(f"dense search needs a positive limit, got {limit}") + scope = None if cartridge_ids is None else sorted(atoms_of_cartridges(self.db, cartridge_ids)) + if scope == []: + return {} + narrow = to_json(scope) if scope is not None and len(scope) <= SCOPE_NARROW_MAX else None + asked = max(limit, CANDIDATE_FLOOR) + vector = np.asarray(query_vector, dtype=np.float32).tobytes() + with self.db.connect() as connection: + if self._dim(connection) is None: + return {} + while True: + ranked = self._ranked(connection, vector, asked, narrow) + if not ranked: + return {} + allowed = self._allowed( + connection, + [atom_id for atom_id, _ in ranked], + holder=holder, + scenario=scenario, + source_id=source_id, + modality=modality, + since=since, + until=until, + cartridge_ids=cartridge_ids, + as_holder=as_holder, + ) + hits = {atom_id: 1.0 - distance for atom_id, distance in ranked if atom_id in allowed} + # A filter that keeps little of the wide ask means there are more answers further + # out; one wider round buys them, and the ceiling keeps the round bounded. + if len(hits) >= limit or len(ranked) < asked or asked >= CANDIDATE_CEILING: + break + asked = min(asked * 4, CANDIDATE_CEILING) + return dict(sorted(hits.items(), key=lambda item: -item[1])[:limit]) + + def _ranked( + self, connection: sqlite3.Connection, vector: bytes, asked: int, narrow: str | None + ) -> list[tuple[str, float]]: + """The engine's own answer: nearest vectors, narrowed to a scope only when the scope is small.""" + return [ + (str(row["atom_id"]), float(row["distance"])) + for row in connection.execute( + """ + SELECT atom_id, distance FROM vec_atoms + WHERE embedding MATCH ? AND k = ? + AND (? IS NULL OR atom_id IN (SELECT value FROM json_each(?))) + ORDER BY distance + """, + (vector, asked, narrow, narrow), + ).fetchall() + ] + + def _allowed( + self, + connection: sqlite3.Connection, + atom_ids: list[str], + *, + holder: str | None, + scenario: str | None, + source_id: str | None, + modality: str | None, + since: str | None, + until: str | None, + cartridge_ids: list[str] | None, + as_holder: str | None, + ) -> set[str]: + """Which candidates this search may show, asked of the atoms table in one statement. + + The currency rules repeat `v_atoms_current` here on purpose: the view answers this + question with a full scan (77ms over 153,857 atoms) while the atoms table answers it + through the atom_id index in 0.2ms, and a search cannot afford the scan. The scope is + checked here rather than in the engine so the rule lives in one place; the engine's id + list is only a narrowing. + """ + rows = connection.execute( + """ + SELECT atom_id FROM atoms + WHERE atom_id IN (SELECT value FROM json_each(?)) + AND (valid_to IS NULL OR valid_to > strftime('%Y-%m-%dT%H:%M:%f+00:00', 'now')) + AND (valid_from IS NULL OR valid_from <= strftime('%Y-%m-%dT%H:%M:%f+00:00', 'now')) + AND (retention_until IS NULL OR retention_until >= strftime('%Y-%m-%dT%H:%M:%f+00:00', 'now')) + AND (? IS NULL OR holder_id = ?) + AND (? IS NULL OR scenario = ?) + AND (? IS NULL OR source_id = ?) + AND (? IS NULL OR modality = ?) + AND (? IS NULL OR (event_time IS NOT NULL AND event_time >= ?)) + AND (? IS NULL OR (event_time IS NOT NULL AND event_time <= ?)) + AND (? IS NULL OR EXISTS ( + SELECT 1 FROM cartridge_atoms AS ca + WHERE ca.atom_id = atoms.atom_id + AND ca.cartridge_id IN (SELECT value FROM json_each(?)) + )) + AND (? IS NULL OR holder_id = '' + OR json_extract(visibility_json, '$.scope') = 'everyone' + OR holder_id = ? + OR EXISTS ( + SELECT 1 FROM json_each(json_extract(visibility_json, '$.holders')) + WHERE value = ? + )) + """, + ( + to_json(atom_ids), + holder, + holder, + scenario, + scenario, + source_id, + source_id, + modality, + modality, + since, + since, + until, + until, + None if cartridge_ids is None else to_json(cartridge_ids), + None if cartridge_ids is None else to_json(cartridge_ids), + as_holder, + as_holder, + as_holder, + ), + ).fetchall() + return {str(row["atom_id"]) for row in rows} + + def drop_atoms(self, atom_ids: list[str]) -> int: + """A forgotten memory must stop answering from the vector side too.""" + with self.db.connect() as connection: + return self._delete(connection, atom_ids) + + def stats(self) -> dict[str, Any]: + with self.db.connect() as connection: + meta = self._meta(connection) + if not meta: + return {"built": False, "format": None, "dim": 0, "vectors": 0, "expected": 0, "pending": 0} + rows = connection.execute("SELECT count(*) AS n FROM vec_atoms").fetchone()["n"] + expected = _expected(connection) + pending = len(self._stale(connection)) + return { + "built": True, + "format": meta.get("format"), + "dim": int(meta.get("dim", "0")), + "model": meta.get("model"), + "vectors": int(rows), + "expected": int(expected), + "pending": pending, + } + + def _meta(self, connection: sqlite3.Connection) -> dict[str, str]: + if not _table_exists(connection, META_TABLE): + return {} + return { + str(row["key"]): str(row["value"]) + for row in connection.execute("SELECT key, value FROM vector_index_meta") + } + + def _dim(self, connection: sqlite3.Connection) -> int | None: + dim = self._meta(connection).get("dim") + return None if dim is None else int(dim) + + def _rebuild(self, connection: sqlite3.Connection) -> None: + """The vectors no longer describe what is stored: start the table over.""" + connection.execute("DROP TABLE IF EXISTS vec_atoms") + + def _delete(self, connection: sqlite3.Connection, atom_ids: list[str]) -> int: + if not atom_ids or self._dim(connection) is None: + return 0 + removed = 0 + for start in range(0, len(atom_ids), DELETE_CHUNK): + stop = min(start + DELETE_CHUNK, len(atom_ids)) + removed += connection.execute( + "DELETE FROM vec_atoms WHERE atom_id IN (SELECT value FROM json_each(?))", + (to_json(atom_ids[start:stop]),), + ).rowcount + return removed + + def _current(self, connection: sqlite3.Connection, atom_ids: list[str]) -> list[tuple[str, Any]]: + """The written memories that still deserve a vector: open, and with words.""" + from .row_codec import row_to_atom + + rows = connection.execute( + """ + SELECT * FROM v_atoms_current + WHERE atom_id IN (SELECT value FROM json_each(?)) AND text IS NOT NULL AND text != '' + """, + (to_json(atom_ids),), + ).fetchall() + return [(str(row["atom_id"]), row_to_atom(row)) for row in rows] + + def _write(self, connection: sqlite3.Connection, items: list[tuple[str, Any]], vectors: np.ndarray) -> None: + """Replace these atoms' rows: the ids are deleted as sets, then the vectors go in. + + Deleting one memory at a time was quadratic on a real corpus — a vec0 table has no index + on atom_id, so each delete re-scans every vector already stored. + """ + if len(vectors) != len(items): + raise ValueError(f"sync encoded {len(vectors)} vectors for {len(items)} atoms") + matrix = np.asarray(vectors, dtype=np.float32) + if matrix.ndim != 2 or matrix.shape[1] != self._dim(connection): + raise ValueError(f"sync produced a {matrix.shape} matrix, expected one row per atom at the stored width") + ids = [atom_id for atom_id, _ in items] + for start in range(0, len(ids), DELETE_CHUNK): + stop = min(start + DELETE_CHUNK, len(ids)) + connection.execute( + "DELETE FROM vec_atoms WHERE atom_id IN (SELECT value FROM json_each(?))", (to_json(ids[start:stop]),) + ) + for start in range(0, len(items), BATCH): + stop = min(start + BATCH, len(items)) + connection.executemany( + "INSERT INTO vec_atoms (embedding, atom_id, stamp) VALUES (?, ?, ?)", + [ + (matrix[index].tobytes(), atom_id, _stamp(atom)) + for index, (atom_id, atom) in enumerate(items[start:stop], start=start) + ], + ) + + def _stale(self, connection: sqlite3.Connection) -> list[tuple[str, Any]]: + """Current memories the index does not hold, or holds from an older version of the memory. + + A memory with no words is the visual index's job, so it never reaches this table. The + stored stamps are copied into a temp table first: a vec0 table has no index on atom_id, + so joining it straight against the atoms scans every vector once per atom — hours on a + real corpus where the copy makes it one scan each. + """ + from .row_codec import row_to_atom + + connection.executescript( + "DROP TABLE IF EXISTS temp.dense_seen;\n" + "CREATE TEMP TABLE dense_seen (atom_id TEXT PRIMARY KEY, stamp TEXT);\n" + "INSERT INTO dense_seen SELECT atom_id, stamp FROM vec_atoms;" + ) + rows = connection.execute( + """ + SELECT a.* FROM v_atoms_current AS a + LEFT JOIN dense_seen AS d ON d.atom_id = a.atom_id + WHERE a.text IS NOT NULL AND a.text != '' + AND (d.atom_id IS NULL OR d.stamp IS NULL OR d.stamp != atom_stamp(a.ingested_at, a.text)) + ORDER BY a.atom_id + """ + ).fetchall() + return [(str(row["atom_id"]), row_to_atom(row)) for row in rows] + + +def atom_stamp(ingested_at: str, text: str | None) -> str: + digest = hashlib.sha256((text or "").encode("utf-8", "replace")).hexdigest()[:16] + return f"{ingested_at}:{digest}" + + +def expected_vectors(database: Database) -> int: + """Memories this index is meant to hold: the open ones that have words.""" + with database.connect() as connection: + return _expected(connection) + + +def _expected(connection: sqlite3.Connection) -> int: + row = connection.execute( + "SELECT count(*) AS n FROM v_atoms_current WHERE text IS NOT NULL AND text != ''" + ).fetchone() + return int(row["n"]) + + +def _stamp(atom: Any) -> str: + return atom_stamp(atom.ingested_at.isoformat(), atom.text) + + +def _table_exists(connection: sqlite3.Connection, name: str) -> bool: + return ( + connection.execute("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", (name,)).fetchone() is not None + ) diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py index af3fffa..6ce2ee5 100644 --- a/backend/tests/conftest.py +++ b/backend/tests/conftest.py @@ -2,11 +2,13 @@ from __future__ import annotations +import hashlib import json from collections.abc import Callable from pathlib import Path from typing import Any +import numpy as np import pytest ROPE: dict[str, Any] = { @@ -51,3 +53,42 @@ def write_model_fixture() -> Callable[..., Path]: @pytest.fixture def model_dir(tmp_path: Path) -> Path: return write_model(tmp_path / "model") + + +class HashEmbeddings: + """A stand-in for the embedding model: the same words give the same vector, no download. + + Different words land in different directions, which is all a retrieval test needs; the + numbers mean nothing to a human, so no test should assert on a score's value. + """ + + def __init__(self, dim: int = 16) -> None: + self.dim = dim + + def enabled(self) -> bool: + return True + + def encode_texts(self, texts: list[str]) -> np.ndarray: + if not texts: + raise ValueError("encode_texts was asked to encode nothing") + rows = [ + np.frombuffer(hashlib.sha256(text.encode("utf-8")).digest()[: self.dim], dtype=np.int8).astype(np.float32) + for text in texts + ] + matrix = np.stack(rows) + norms = np.linalg.norm(matrix, axis=1, keepdims=True) + return matrix / np.where(norms == 0, 1.0, norms) + + def encode_atoms(self, atoms: list[Any]) -> np.ndarray: + return self.encode_texts([atom.text or "" for atom in atoms]) + + def embed_query(self, query: str) -> np.ndarray: + return self.encode_texts([query])[0] + + def probe_dim(self) -> int: + return self.dim + + +@pytest.fixture +def fake_embeddings() -> Callable[[], HashEmbeddings]: + return HashEmbeddings diff --git a/backend/tests/test_maintenance.py b/backend/tests/test_maintenance.py index 7d4750f..b820ccf 100644 --- a/backend/tests/test_maintenance.py +++ b/backend/tests/test_maintenance.py @@ -5,7 +5,9 @@ from __future__ import annotations import hashlib import json import sqlite3 +from collections.abc import Callable from pathlib import Path +from typing import Any import pytest from typer.testing import CliRunner @@ -20,6 +22,7 @@ from ely_eye.memory import MemoryService from ely_eye.migrations import SCHEMA_VERSION from ely_eye.schemas import EvidenceAtom, MemoryOrigin, Modality, SourceRecord, TrustScores, utc_now from ely_eye.storage import ObjectStore +from ely_eye.vector_index import VectorIndex KEYS = "the keys are on the shelf by the door" PASSPORT = "the passport is in the blue box" @@ -85,6 +88,29 @@ def test_a_restore_brings_the_memories_back(tmp_path: Path) -> None: assert {atom.text for atom in database.list_atoms()} == {KEYS, PASSPORT} +def test_a_backup_carries_the_vector_index_and_a_restore_answers_from_it( + tmp_path: Path, fake_embeddings: Callable[[], Any] +) -> None: + source = Settings(home=tmp_path / "home", default_holder="wife", dense_embeddings_enabled=True) + database = Database(source) + database.init() + index = VectorIndex(source, database, fake_embeddings()) + memory = MemoryService(source, database, vectors=index) + for statement in (KEYS, PASSPORT): + memory.remember(statement, holder_id="wife", origin=MemoryOrigin.self_) + assert index.stats()["vectors"] == 2 + + maintenance.backup(source, tmp_path / "backup") + destination = Settings(home=tmp_path / "elsewhere", default_holder="wife", dense_embeddings_enabled=True) + maintenance.restore(destination, tmp_path / "backup") + + restored = VectorIndex(destination, Database(destination), fake_embeddings()) + assert restored.stats()["vectors"] == 2, "the vectors are data, not a cache to redo" + hits = restored.search(restored.embeddings.embed_query(KEYS), 10) + best = next(iter(hits)) + assert Database(destination).get_atoms([best])[0].text == KEYS + + def test_a_restore_keeps_the_store_it_replaces(tmp_path: Path) -> None: settings = _settings(tmp_path) _store(settings) diff --git a/backend/tests/test_memory_write.py b/backend/tests/test_memory_write.py index 62f78e4..b5c7f55 100644 --- a/backend/tests/test_memory_write.py +++ b/backend/tests/test_memory_write.py @@ -2,8 +2,10 @@ from __future__ import annotations +from collections.abc import Callable from datetime import UTC, datetime, timedelta from pathlib import Path +from typing import Any import pytest @@ -18,6 +20,7 @@ from ely_eye.schemas import ( TimePrecision, TrustScores, ) +from ely_eye.vector_index import VectorIndex KEYS = "我们的钥匙放在玄关抽屉里" NOW = datetime(2026, 5, 1, 12, 0, tzinfo=UTC) @@ -124,8 +127,8 @@ def test_time_travel_reads_the_memory_that_was_true_then(service: MemoryService) } -def test_forget_reaches_the_row_the_index_the_edges_the_cache_and_the_vectors( - service: MemoryService, +def test_forget_reaches_every_copy_of_a_memory( + service: MemoryService, fake_embeddings: Callable[[], Any] ) -> None: atom = _seed_atom("seed-1:text", "seed-1") _seed(service.db, atom) @@ -133,23 +136,21 @@ def test_forget_reaches_the_row_the_index_the_edges_the_cache_and_the_vectors( [{"from_atom_id": "seed-1:text", "to_atom_id": "other", "relation": "cites", "confidence": 1.0}] ) service.db.log_cache_event(layer="compiler", cache_key="k", event="hit", bytes_estimate=10) - vectors = service.settings.cache_dir / "dense_index.npy" - vectors.write_bytes(b"stale") - (service.settings.cache_dir / "dense_index_atoms.json").write_text("[]", encoding="utf-8") + VectorIndex(service.settings, service.db, fake_embeddings()).sync() report = service.forget(atom_id="seed-1:text", reason="the drawer is empty") assert report.atoms_deleted == 1 assert report.edges_deleted == 1 assert report.cache_entries_cleared == 1 - assert report.vector_index_dropped is True + assert report.vectors_dropped == 1 + assert VectorIndex(service.settings, service.db, fake_embeddings()).stats()["vectors"] == 0 assert service.db.atom_count() == 0 with service.db.connect() as connection: assert connection.execute("SELECT COUNT(*) FROM atom_fts").fetchone()[0] == 0 assert connection.execute("SELECT COUNT(*) FROM evidence_edges").fetchone()[0] == 0 trace = connection.execute("SELECT layer, event FROM cache_events").fetchall() assert [(row["layer"], row["event"]) for row in trace] == [("memory", "invalidated")] - assert not vectors.exists() def test_forget_names_the_cartridge_that_exported_the_source(service: MemoryService) -> None: diff --git a/backend/tests/test_vector_index.py b/backend/tests/test_vector_index.py new file mode 100644 index 0000000..6c669d1 --- /dev/null +++ b/backend/tests/test_vector_index.py @@ -0,0 +1,291 @@ +"""The dense index answers from the database, so a memory is findable the moment it is written. + +Every test here runs with a stand-in model: what is being proved is the index's freshness and +its boundaries, not the quality of a 500 MB snapshot. +""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any + +import pytest + +from ely_eye.config import Settings +from ely_eye.db import Database +from ely_eye.memory import MemoryService +from ely_eye.retrieval import RetrievalService +from ely_eye.schemas import MemoryOrigin, SearchFilters +from ely_eye.vector_index import VectorIndex, expected_vectors + +KEYS = "我们的钥匙放在玄关抽屉里" +PASSPORT = "护照在书房第二个抽屉" + + +def _settings(tmp_path: Path, model_id: str = "hash-mini") -> Settings: + return Settings(home=tmp_path, dense_embeddings_enabled=True, embedding_model_id=model_id) + + +def _index(settings: Settings, database: Database, fake_embeddings: Callable[[], Any]) -> VectorIndex: + return VectorIndex(settings, database, fake_embeddings()) + + +def _memory(settings: Settings, database: Database, index: VectorIndex) -> MemoryService: + return MemoryService(settings, database, vectors=index) + + +def test_a_memory_written_now_is_findable_now(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + _memory(settings, database, index).remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + + hits = index.search(index.embeddings.embed_query(KEYS), 10) + + assert list(hits) == ["mem-" + _digest(KEYS, "wife") + ":text"], "no rebuild step stands between writing and finding" + + +def _digest(statement: str, holder: str) -> str: + import hashlib + + return hashlib.sha256(f"{holder}\n{statement}".encode()).hexdigest()[:16] + + +def test_a_closed_memory_stops_answering(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + memory = _memory(settings, database, index) + old = memory.remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + + assert index.search(index.embeddings.embed_query(KEYS), 10) == {old.atom_id: pytest.approx(1.0)} + replacement = memory.supersede(old.atom_id, text=PASSPORT, holder_id="wife") + + after = index.search(index.embeddings.embed_query(KEYS), 10) + assert old.atom_id not in after, "a closed memory stops answering" + assert set(after) == {replacement.atom_id} + + +def test_a_rewritten_memory_is_encoded_again(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + atom = _memory(settings, database, index).remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + index.sync() + + rewritten = atom.model_copy(update={"text": PASSPORT}) + database.upsert_atoms([rewritten]) + index.sync() + + stale = index.search(index.embeddings.embed_query(KEYS), 10) + fresh = index.search(index.embeddings.embed_query(PASSPORT), 10) + assert stale[rewritten.atom_id] < 0.9, "the old wording still holds the old vector" + assert fresh[rewritten.atom_id] > 0.99 + + +def test_a_forgetting_leaves_no_vector_behind(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + memory = _memory(settings, database, index) + memory.remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + index.sync() + + report = memory.forget(atom_id="mem-" + _digest(KEYS, "wife") + ":text", reason="moved house") + + assert report.vectors_dropped == 1 + assert index.stats()["vectors"] == 0 + + +def test_somebody_elses_memory_is_not_mine_to_find(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + memory = _memory(settings, database, index) + hers_atom = memory.remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + his_atom = memory.remember(PASSPORT, holder_id="husband", origin=MemoryOrigin.self_) + index.sync() + + hers = index.search(index.embeddings.embed_query(KEYS), 10, as_holder="wife") + his = index.search(index.embeddings.embed_query(KEYS), 10, as_holder="husband") + + assert set(hers) == {hers_atom.atom_id} + assert set(his) == {his_atom.atom_id}, "he answers from his own memory, not from hers" + + +def test_a_shared_memory_is_findable_by_both(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + from ely_eye.schemas import Visibility + + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + memory = _memory(settings, database, index) + memory.remember( + KEYS, holder_id="wife", origin=MemoryOrigin.self_, visibility=Visibility(scope="visible_to", holders=["husband"]) + ) + index.sync() + + assert index.search(index.embeddings.embed_query(KEYS), 10, as_holder="wife") + assert index.search(index.embeddings.embed_query(KEYS), 10, as_holder="husband") + assert index.search(index.embeddings.embed_query(KEYS), 10, as_holder="daughter") == {} + + +def test_a_cartridge_scope_cannot_be_walked_around(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + from datetime import UTC, datetime + + from ely_eye import cartridge_store + + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + memory = _memory(settings, database, index) + kept = memory.remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + other = memory.remember(PASSPORT, holder_id="wife", origin=MemoryOrigin.self_) + index.sync() + database.write_cartridge( + cartridge_id="cart-1", + name="household", + root_path=tmp_path / "cart-1", + manifest_json="{}", + dna="d" * 64, + token_equivalent=1, + created_at=datetime(2026, 5, 1, tzinfo=UTC), + ) + cartridge_store.add_membership(database, "cart-1", [kept.atom_id]) + + hits = index.search(index.embeddings.embed_query(KEYS + " " + PASSPORT), 10, cartridge_ids=["cart-1"]) + + assert set(hits) == {kept.atom_id}, "the scope is applied inside the engine, not after it" + assert other.atom_id not in hits + + +def test_a_different_model_rebuilds_instead_of_mixing_vectors(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + _memory(settings, database, index).remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + index.sync() + assert index.stats()["vectors"] == 1 + + swapped = VectorIndex(_settings(tmp_path, model_id="other-model"), database, fake_embeddings()) + swapped.sync() + + facts = swapped.stats() + assert facts["model"] == "other-model" + assert facts["vectors"] == 1, "the table was rebuilt, not appended to" + + +def test_a_picture_with_no_words_is_the_other_indexs_job(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + from datetime import UTC, datetime + + from ely_eye.schemas import EvidenceAtom, Modality, SourceRecord, TrustScores + + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + now = datetime(2026, 5, 1, tzinfo=UTC) + database.upsert_sources( + [ + SourceRecord( + source_id="pic-1", + path="/tmp/pic.png", + kind="image", + mime="image/png", + sha256="0" * 64, + size_bytes=4, + created_at=now, + ) + ] + ) + database.upsert_atoms( + [ + EvidenceAtom( + atom_id="pic-1:image", + modality=Modality.image, + source_id="pic-1", + source="pic-1", + ingested_at=now, + text="", + image_ref="pic-1", + trust=TrustScores(parser=1.0), + token_equivalent=1, + holder_id="wife", + ) + ] + ) + + assert index.sync()["encoded"] == 0 + assert expected_vectors(database) == 0 + assert index.stats()["vectors"] == 0, "a wordless picture is the visual index's job" + + +def test_the_index_survives_a_reopen(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + settings = _settings(tmp_path) + first = Database(settings) + index = _index(settings, first, fake_embeddings) + _memory(settings, first, index).remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + index.sync() + + reopened = VectorIndex(settings, Database(settings), fake_embeddings()) + + assert reopened.stats()["vectors"] == 1 + assert reopened.search(reopened.embeddings.embed_query(KEYS), 10) + + +def test_superseding_moves_the_vector_rather_than_keeping_two(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + memory = _memory(settings, database, index) + old = memory.remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + + memory.supersede(old.atom_id, text=PASSPORT, holder_id="wife") + + facts = index.stats() + assert facts["vectors"] == 1, "the closed memory lost its vector and the replacement took its place" + assert facts["pending"] == 0 + + +def test_a_store_that_never_built_its_index_says_so_instead_of_stalling_a_search( + tmp_path: Path, fake_embeddings: Callable[[], Any] +) -> None: + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + atom = _memory(settings, database, index).remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + index.drop_atoms([atom.atom_id]) + + facts = index.stats() + assert (facts["vectors"], facts["expected"], facts["pending"]) == (0, 1, 1) + assert index.search(index.embeddings.embed_query(KEYS), 10) == {}, "a search never rebuilds" + assert index.sync()["encoded"] == 1 + assert index.search(index.embeddings.embed_query(KEYS), 10) + + +def test_the_dense_leg_hides_what_the_view_hides(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + """The search repeats the view's currency rules, so the two must keep answering alike.""" + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + atom = _memory(settings, database, index).remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + lapsed = (datetime.now(UTC) - timedelta(days=1)).isoformat() + with database.connect() as connection: + connection.execute("UPDATE atoms SET retention_until = ? WHERE atom_id = ?", (lapsed, atom.atom_id)) + visible = [row["atom_id"] for row in connection.execute("SELECT atom_id FROM v_atoms_current")] + + assert atom.atom_id not in visible, "the view has retired the memory" + assert index.stats()["vectors"] == 1, "the index still holds its vector" + assert index.search(index.embeddings.embed_query(KEYS), 10) == {}, "a search retires it too" + + +def test_a_query_asks_for_a_real_number_of_answers(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + index = _index(settings, database, fake_embeddings) + with pytest.raises(ValueError, match="positive limit"): + index.search(index.embeddings.embed_query(KEYS), 0) + + +def test_search_answers_from_the_index_without_a_rebuild(tmp_path: Path, fake_embeddings: Callable[[], Any]) -> None: + settings, database = _settings(tmp_path), Database(_settings(tmp_path)) + embeddings = fake_embeddings() + service = RetrievalService(settings, database) + service.embeddings = embeddings # type: ignore[assignment] + index = VectorIndex(settings, database, embeddings) + service.vectors = index + _memory(settings, database, index).remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_) + + result = service.search(KEYS, filters=SearchFilters(as_holder="wife")) + + assert result.hits, "the needle was written and found in the same breath" + assert result.hits[0].atom.atom_id == "mem-" + _digest(KEYS, "wife") + ":text" + assert result.hits[0].dense_score > 0.0 diff --git a/backend/tests/test_visual_index.py b/backend/tests/test_visual_index.py index 507b786..0a5f12e 100644 --- a/backend/tests/test_visual_index.py +++ b/backend/tests/test_visual_index.py @@ -180,7 +180,7 @@ def test_forgetting_a_picture_voids_the_picture_index(tmp_path: Path) -> None: report = MemoryService(settings, database).forget(source_id=atom.source_id) - assert report.vector_index_dropped + assert report.visual_index_dropped assert not index.index_path.exists() and not index.map_path.exists() assert index.load() is None, "a forgotten picture must not keep answering from memory" diff --git a/bench/dense_scale.py b/bench/dense_scale.py new file mode 100644 index 0000000..d43950f --- /dev/null +++ b/bench/dense_scale.py @@ -0,0 +1,318 @@ +"""What the dense leg costs at the size the library actually reaches, and what a million would. + +The baseline store holds 165,973 atoms. This check builds the product's own vector table inside +a copy of that store, times real searches, proves the engine's answers match a brute-force scan, +and then states what one million memories would cost on the same machine — with the assumption +written down instead of implied. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import sqlite3 +import time +from collections.abc import Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np + +from bench.query_set import load_queries, scale_probes +from bench.report import git_commit +from bench.run import latency +from bench.systems import DENSE_MODEL, MODELS, machine_facts + +if TYPE_CHECKING: + from ely_eye.db import Database + +DEFAULT_HOME = Path("bench/runs/baseline") +DEFAULT_SCRATCH = Path("bench/runs/dense_scale") +MIN_ATOMS = 100_000 +TOP_K = 10 +CORRECTNESS_PROBES = 40 +EXTRAPOLATE_TO = 1_000_000 +WRITE_PROBES = 50 + + +def _one_write(index: Any, atom_id: str) -> float: + """Time one memory going into an index that already holds the whole corpus.""" + began = time.perf_counter() + index.sync_atoms([atom_id]) + return time.perf_counter() - began + + +def vector_footprint(database: Any, table: str) -> tuple[int, str]: + """What the vector table costs on disk, and how that was counted. + + The database file itself is a poor measure: free pages absorb a table this size, so the file + can stay exactly the same length while the table grows. dbstat counts the pages the vec0 + shadow tables actually hold, and it works without the vec0 extension loaded. + """ + with database.connect() as connection: + try: + pages = connection.execute( + "SELECT COALESCE(SUM(pgsize), 0) FROM dbstat WHERE name LIKE ?", (f"{table}%",) + ).fetchone()[0] + if int(pages): + return int(pages), "dbstat page sizes over the vec0 shadow tables" + except sqlite3.OperationalError: + pass + payload = connection.execute( + "SELECT COALESCE(SUM(length(embedding)), 0) FROM vec_atoms" + ).fetchone()[0] + return int(payload), "vector bytes only, this sqlite3 has no dbstat" + + +def copy_store(home: Path, scratch: Path) -> tuple[Path, float]: + """Work on a copy: the baseline store is frozen evidence and must not be migrated in place.""" + source = home / "data" / "ely_eye.sqlite" + if not source.is_file(): + raise SystemExit(f"{source}: no product database, build it with python -m bench.ingest_baseline") + data = scratch / "data" + data.mkdir(parents=True, exist_ok=True) + began = time.perf_counter() + for suffix in ("", "-wal", "-shm"): + if (path := Path(str(source) + suffix)).is_file(): + shutil.copyfile(path, data / (source.name + suffix)) + return scratch, time.perf_counter() - began + + +def stored_vectors(database: Database) -> tuple[list[str], np.ndarray]: + """Every vector the engine holds, for a scan it can be checked against.""" + with database.connect() as connection: + rows = connection.execute( + """ + SELECT atom_id, embedding FROM vec_atoms ORDER BY atom_id + """ + ).fetchall() + if not rows: + raise SystemExit("the vector table is empty, so there is nothing to measure") + ids = [str(row["atom_id"]) for row in rows] + matrix = np.stack([np.frombuffer(row["embedding"], dtype=np.float32) for row in rows]) + return ids, matrix + + +def agreement(ids: list[str], matrix: np.ndarray, probes: Sequence[tuple[np.ndarray, list[str]]]) -> dict[str, Any]: + """How often the engine's top ten agrees with a full scan, counting a tie as a tie.""" + if not probes: + raise SystemExit("correctness needs at least one probe") + matched = 0 + compared = 0 + for vector, returned in probes: + cosine = matrix @ vector + order = np.argsort(-cosine, kind="stable") + threshold = cosine[order[TOP_K - 1]] + correct = {ids[index] for index in order if cosine[index] >= threshold} + matched += sum(1 for atom_id in returned if atom_id in correct) + compared += min(len(returned), TOP_K) + return { + "probes": len(probes), + "top_k": TOP_K, + "tie_aware_agreement": matched / compared if compared else 0.0, + "note": "a returned atom counts when its true cosine is at least the k-th best, so ties cannot fail the check", + } + + +def measure(home: Path, scratch: Path, probes: int, model_dir: Path) -> dict[str, Any]: + from ely_eye.config import Settings + from ely_eye.db import Database + from ely_eye.embeddings import EmbeddingService + from ely_eye.vector_index import VectorIndex, expected_vectors + + copy_home, copy_seconds = copy_store(home, scratch) + settings = Settings( + home=copy_home, + dense_embeddings_enabled=True, + embedding_model_dir=model_dir, + embedding_model_id=DENSE_MODEL, + ) + database = Database(settings) + atoms = database.atom_count() + if atoms < MIN_ATOMS: + raise SystemExit(f"{database.path}: {atoms} atoms, this check needs at least {MIN_ATOMS}") + + embeddings = EmbeddingService(settings, database) + index = VectorIndex(settings, database, embeddings) + print(f"{atoms} atoms copied in {copy_seconds:.1f}s, encoding with {DENSE_MODEL}", flush=True) + began = time.perf_counter() + index.sync() + build_seconds = time.perf_counter() - began + vec_bytes, bytes_method = vector_footprint(database, "vec_atoms") + print(f"index built in {build_seconds:.1f}s", flush=True) + + began = time.perf_counter() + index.sync() + repair_seconds = time.perf_counter() - began + print(f"nothing to repair, the scan took {repair_seconds:.1f}s", flush=True) + + queries = scale_probes(load_queries(), probes) + index.search(embeddings.embed_query(queries[0].text), TOP_K) + search_times: list[float] = [] + encode_times: list[float] = [] + returned: list[tuple[str, list[str]]] = [] + for query in queries: + began = time.perf_counter() + vector = embeddings.embed_query(query.text) + encode_times.append(time.perf_counter() - began) + began = time.perf_counter() + hits = index.search(vector, TOP_K) + search_times.append(time.perf_counter() - began) + returned.append((query.text, list(hits))) + if len(search_times) % 100 == 0: + print(f"searched {len(search_times)}/{len(queries)}, p50 {sorted(search_times)[len(search_times) // 2] * 1000:.1f}ms", flush=True) + + ids, matrix = stored_vectors(database) + print(f"{len(ids)} vectors read back, checking the engine against a brute-force scan", flush=True) + checked_probes = [] + for query in queries[:CORRECTNESS_PROBES]: + vector = embeddings.embed_query(query.text) + checked_probes.append((vector, list(index.search(vector, TOP_K)))) + checked = agreement(ids, matrix, checked_probes) + print(f"agreement {checked['tie_aware_agreement']:.3f} over {checked['probes']} probes", flush=True) + + written = [hit[0] for _, hit in returned[:WRITE_PROBES] if hit] + writes = latency([_one_write(index, atom_id) for atom_id in written]) + print(f"{len(written)} writes into a full index, p50 {writes['p50'] * 1000:.0f}ms", flush=True) + + facts = index.stats() + encoded = int(facts["vectors"]) + per_vector = vec_bytes / encoded if encoded else 0.0 + texts_per_second = encoded / build_seconds if build_seconds else 0.0 + timing = latency(search_times) + return { + "commit": git_commit(), + "database": str(database.path), + "atoms": atoms, + "memories_with_words": expected_vectors(database), + "vectors": encoded, + "pending": facts["pending"], + "dim": facts["dim"], + "model": facts["model"], + "copy_seconds": copy_seconds, + "build_seconds": build_seconds, + "texts_per_second": texts_per_second, + "repair_seconds": repair_seconds, + "writes": len(written), + "write_latency": writes, + "probes": len(queries), + "probe_modalities": sorted({query.modality for query in queries}), + "latency": timing, + "encode_latency": latency(encode_times), + "correctness": checked, + "bytes_per_vector": per_vector, + "vector_bytes": vec_bytes, + "bytes_method": bytes_method, + "extrapolation": { + "to": EXTRAPOLATE_TO, + "assumption": ( + "sqlite-vec 0.1.9 scans every vector, so latency and bytes grow linearly with the " + "count; the model, the machine and the thread budget stay the same" + ), + "p50_seconds": timing["p50"] * EXTRAPOLATE_TO / atoms, + "p95_seconds": timing["p95"] * EXTRAPOLATE_TO / atoms, + "write_p50_seconds": writes["p50"] * EXTRAPOLATE_TO / atoms, + "gibibytes": per_vector * EXTRAPOLATE_TO / 1024**3, + "build_seconds": EXTRAPOLATE_TO / texts_per_second if texts_per_second else None, + }, + "machine": machine_facts(), + } + + +def write_markdown(report: dict[str, Any], path: Path) -> None: + timing = report["latency"] + encode = report["encode_latency"] + write = report["write_latency"] + machine = report["machine"] + future = report["extrapolation"] + path.write_text( + "\n".join( + [ + f"# DENSE: the vector leg at corpus size (commit {report['commit'][:9]})", + "", + f"- database: `{report['database']}`", + ( + f"- atoms: {report['atoms']}, memories with words: {report['memories_with_words']}, " + f"vectors: {report['vectors']}, pending: {report['pending']}" + ), + f"- model: `{report['model']}`, width {report['dim']}", + ( + f"- build: {report['build_seconds']:.1f}s for {report['vectors']} memories " + f"({report['texts_per_second']:.0f} per second)" + ), + ( + f"- writing one memory into a full index: p50 {write['p50'] * 1000:.0f}ms · " + f"p95 {write['p95'] * 1000:.0f}ms over {write['queries']} writes" + ), + f"- the repair scan (`ely-eye build-dense-index`) over a full index: {report['repair_seconds']:.1f}s, and it is not on the search path", + ( + f"- search p50 {timing['p50'] * 1000:.1f}ms · p95 {timing['p95'] * 1000:.1f}ms · " + f"max {timing['max'] * 1000:.1f}ms over {timing['queries']} queries" + ), + ( + f"- of which encoding the question: p50 {encode['p50'] * 1000:.1f}ms · " + f"p95 {encode['p95'] * 1000:.1f}ms" + ), + ( + f"- against a brute-force scan: {report['correctness']['tie_aware_agreement']:.3f} " + f"of the top {report['correctness']['top_k']} over {report['correctness']['probes']} probes " + f"({report['correctness']['note']})" + ), + ( + f"- bytes per stored vector: {report['bytes_per_vector']:.0f} " + f"({report['vector_bytes']} bytes for {report['vectors']} vectors, {report['bytes_method']})" + ), + ( + f"- machine: {machine['system']} {machine['release']}, {machine['cpus']} cores, " + f"{machine['threads_per_query']} threads per query, load {machine['load_average_1_5_15']}" + ), + "", + "## One million memories on this machine", + "", + f"- p50 {future['p50_seconds'] * 1000:.0f}ms · p95 {future['p95_seconds'] * 1000:.0f}ms", + f"- vector table: {future['gibibytes']:.1f} GiB", + f"- first build: {future['build_seconds'] / 60:.1f} minutes", + f"- writing one memory: {future['write_p50_seconds'] * 1000:.0f}ms — a vec0 delete by id scans the table, so this grows with the count too", + f"- assumption: {future['assumption']}", + "", + ( + "_Latency is wall time per search on the machine named above with the thread budget the " + "benchmark pins; a run on a loaded machine is not comparable._" + ), + "", + ] + ), + encoding="utf-8", + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--home", type=Path, default=DEFAULT_HOME) + parser.add_argument("--scratch", type=Path, default=DEFAULT_SCRATCH) + parser.add_argument("--probes", type=int, default=300) + parser.add_argument("--model-dir", type=Path, default=MODELS / "paraphrase-multilingual-MiniLM-L12-v2") + parser.add_argument("--out", type=Path, default=Path("bench/reports/DENSE.md")) + parser.add_argument("--json-out", type=Path, default=Path("bench/reports/dense.json")) + args = parser.parse_args() + if args.probes <= 0: + raise SystemExit(f"--probes must be positive, got {args.probes}") + if not args.model_dir.is_dir(): + raise SystemExit(f"{args.model_dir}: no embedding snapshot, run python -m bench.fetch_models") + + report = measure(args.home, args.scratch, args.probes, args.model_dir) + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8") + write_markdown(report, args.out) + timing = report["latency"] + print( + f"{report['vectors']} vectors: build {report['build_seconds']:.1f}s, " + f"search p50 {timing['p50'] * 1000:.1f}ms p95 {timing['p95'] * 1000:.1f}ms, " + f"agreement {report['correctness']['tie_aware_agreement']:.3f}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/bench/quant_probe.py b/bench/quant_probe.py new file mode 100644 index 0000000..eccb5aa --- /dev/null +++ b/bench/quant_probe.py @@ -0,0 +1,332 @@ +"""What each way of storing a vector actually buys: latency, recall past the ties, and bytes. + +The product stores float32 today. This check answers, on the real corpus and with the thread +budget the benchmark pins, whether int8, a sign bit, or a shorter vector would buy anything — +and where the crossover is, so the answer does not have to be guessed at a million memories. +""" + +from __future__ import annotations + +import argparse +import json +import random +import sqlite3 +import time +from pathlib import Path +from typing import Any + +import numpy as np +import sqlite_vec + +from bench.report import git_commit +from bench.run import latency +from bench.systems import machine_facts + +DEFAULT_VECTORS = Path( + "bench/runs/baseline/index/paraphrase-multilingual-MiniLM-L12-v2-7a66920843343254.npy" +) +PROBES = 200 +TOP = 10 +K = 10 +COARSE = 200 +SEED = 7 +TRUNCATIONS = (192, 128, 64) +BYTES_PER_VECTOR = {"float32": 1536, "int8": 384, "bit": 48} + + +def load_matrix(path: Path) -> np.ndarray: + if not path.is_file(): + raise SystemExit(f"{path}: no frozen vector matrix, build it with python -m bench.run") + matrix = np.load(path) + if matrix.ndim != 2: + raise SystemExit(f"{path}: expected a matrix of vectors, got shape {matrix.shape}") + return matrix + + +def tie_aware_hits(cosine: np.ndarray, threshold: np.ndarray, returned: list[int], probe: int) -> int: + """The corpus holds byte-identical atoms, so the top-K id set is ambiguous: a returned atom + counts when its true cosine is at least the k-th best, which is what a caller would accept.""" + return sum(1 for row in returned[:K] if cosine[row][probe] >= threshold[probe] - 1e-6) + + +def run_arm( + name: str, + arm: Any, + count: int, + warm: int, + cosine: np.ndarray, + threshold: np.ndarray, +) -> dict[str, Any]: + for index in range(warm): + arm(index) + times: list[float] = [] + hits = 0 + for index in range(count): + began = time.perf_counter() + returned = arm(index) + times.append(time.perf_counter() - began) + hits += tie_aware_hits(cosine, threshold, returned, index) + timing = latency(times) + return { + "arm": name, + "queries": count, + "p50_ms": round(timing["p50"] * 1000, 1), + "p95_ms": round(timing["p95"] * 1000, 1), + f"recall_at_{K}": round(hits / (count * K), 3), + } + + +def vec_table(kind: str, count: int, dim: int, payload: Any) -> sqlite3.Connection: + connection = sqlite3.connect(":memory:") + connection.enable_load_extension(True) + sqlite_vec.load(connection) + if kind == "float": + connection.executescript( + f""" + CREATE VIRTUAL TABLE vv USING vec0( + embedding float[{dim}] distance_metric=cosine, atom_id text, holder_id text + ); + """ + ) + insert = """ + INSERT INTO vv(rowid, embedding, atom_id, holder_id) VALUES (?, ?, ?, ?) + """ + elif kind == "int8": + connection.executescript( + f""" + CREATE VIRTUAL TABLE vv USING vec0( + embedding int8[{dim}], atom_id text, holder_id text + ); + """ + ) + insert = """ + INSERT INTO vv(rowid, embedding, atom_id, holder_id) VALUES (?, vec_int8(?), ?, ?) + """ + else: + connection.executescript( + f""" + CREATE VIRTUAL TABLE vv USING vec0( + embedding bit[{dim}], atom_id text, holder_id text + ); + """ + ) + insert = """ + INSERT INTO vv(rowid, embedding, atom_id, holder_id) VALUES (?, vec_bit(?), ?, ?) + """ + for start in range(0, count, 20_000): + stop = min(start + 20_000, count) + connection.executemany( + insert, + [(row, payload(row), str(row), f"h{row % 2}") for row in range(start, stop)], + ) + return connection + + +def knn(connection: sqlite3.Connection, kind: str, query: Any, top: int, holder: str | None = None) -> list[int]: + """Nearest rowids. Every statement here is written out in full: the only choice is which one.""" + if holder is None: + if kind == "float": + rows = connection.execute( + """ + SELECT rowid FROM vv WHERE embedding MATCH ? AND k = ? ORDER BY distance + """, + (query, top), + ).fetchall() + elif kind == "int8": + rows = connection.execute( + """ + SELECT rowid FROM vv WHERE embedding MATCH vec_int8(?) AND k = ? ORDER BY distance + """, + (query, top), + ).fetchall() + else: + rows = connection.execute( + """ + SELECT rowid FROM vv WHERE embedding MATCH vec_bit(?) AND k = ? ORDER BY distance + """, + (query, top), + ).fetchall() + elif kind == "float": + rows = connection.execute( + """ + SELECT rowid FROM vv WHERE embedding MATCH ? AND k = ? AND holder_id = ? ORDER BY distance + """, + (query, top, holder), + ).fetchall() + elif kind == "int8": + rows = connection.execute( + """ + SELECT rowid FROM vv WHERE embedding MATCH vec_int8(?) AND k = ? AND holder_id = ? ORDER BY distance + """, + (query, top, holder), + ).fetchall() + else: + rows = connection.execute( + """ + SELECT rowid FROM vv WHERE embedding MATCH vec_bit(?) AND k = ? AND holder_id = ? ORDER BY distance + """, + (query, top, holder), + ).fetchall() + return [int(row[0]) for row in rows] + + +def measure(vectors: np.ndarray, probes: int) -> dict[str, Any]: + count, dim = vectors.shape + if probes > count: + raise SystemExit(f"--probes {probes} exceeds the {count} vectors in the corpus") + chosen = random.Random(SEED).sample(range(count), probes) + queries = vectors[chosen].copy() + cosine = vectors @ queries.T + threshold = -np.sort(-cosine, axis=0)[K - 1] + + int8 = np.ascontiguousarray(np.clip(np.round(vectors * 127), -127, 127).astype(np.int8)) + queries_int8 = np.ascontiguousarray(np.clip(np.round(queries * 127), -127, 127).astype(np.int8)) + packed = np.packbits((vectors >= 0).astype(np.uint8), axis=1) + queries_packed = np.packbits((queries >= 0).astype(np.uint8), axis=1) + unpacked = np.unpackbits(packed, axis=1) + queries_unpacked = np.unpackbits(queries_packed, axis=1) + population = unpacked.sum(axis=1).astype(np.int32) + + rows: list[dict[str, Any]] = [] + tables: dict[str, sqlite3.Connection] = {} + for kind, payload, query_bytes in ( + ("float", lambda row: vectors[row].tobytes(), lambda probe: queries[probe].tobytes()), + ("int8", lambda row: int8[row].tobytes(), lambda probe: queries_int8[probe].tobytes()), + ("bit", lambda row: packed[row].tobytes(), lambda probe: queries_packed[probe].tobytes()), + ): + tables[kind] = vec_table(kind, count, dim, payload) + rows.append( + run_arm( + f"sqlite-vec {kind}", + lambda probe, kind=kind, query_bytes=query_bytes: knn(tables[kind], kind, query_bytes(probe), TOP), + probes, + 20, + cosine, + threshold, + ) + ) + rows.append( + run_arm( + f"sqlite-vec {kind} + holder filter", + lambda probe, kind=kind, query_bytes=query_bytes: knn( + tables[kind], kind, query_bytes(probe), TOP, holder="h0" + ), + probes, + 20, + cosine, + threshold, + ) + ) + + rows.append( + run_arm( + "numpy float32 scan (the truth)", + lambda probe: list(np.argsort(-(vectors @ queries[probe]))[:TOP]), + probes, + 10, + cosine, + threshold, + ) + ) + + def hamming(probe: int) -> list[int]: + shared = unpacked @ queries_unpacked[probe] + return list(np.argsort(-(2 * shared - population))[:TOP]) + + rows.append(run_arm("numpy hamming (sign bit)", hamming, probes, 10, cosine, threshold)) + + for dims in TRUNCATIONS: + rows.append( + run_arm( + f"numpy float32 first {dims} dims", + lambda probe, dims=dims: list(np.argsort(-(vectors[:, :dims] @ queries[probe][:dims]))[:TOP]), + min(probes, 50), + 5, + cosine, + threshold, + ) + ) + + def rerank_outside(probe: int) -> list[int]: + coarse = knn(tables["bit"], "bit", queries_packed[probe].tobytes(), COARSE) + order = np.argsort(-(vectors[coarse] @ queries[probe]))[:TOP] + return [coarse[row] for row in order] + + rows.append(run_arm(f"bit coarse@{COARSE} -> float32 rerank in numpy", rerank_outside, probes, 10, cosine, threshold)) + + return { + "commit": git_commit(), + "vectors": count, + "dimensions": dim, + "probes": probes, + "top_k": TOP, + "recall_note": ( + "probes are drawn from the corpus itself, so these arms are compared with each other, " + "not quoted as retrieval quality; a hit is any returned atom whose true cosine is at " + f"least the {K}-th best, because byte-identical atoms make the id set ambiguous" + ), + "bytes_per_vector": BYTES_PER_VECTOR, + "arms": rows, + "machine": machine_facts(), + } + + +def write_markdown(report: dict[str, Any], path: Path) -> None: + machine = report["machine"] + lines = [ + f"# QUANT: what each vector encoding buys (commit {report['commit'][:9]})", + "", + ( + f"- corpus: {report['vectors']} vectors, {report['dimensions']} dimensions, " + f"{report['probes']} probes, top {report['top_k']}" + ), + ( + f"- machine: {machine['system']} {machine['release']}, {machine['cpus']} cores, " + f"{machine['threads_per_query']} threads per query, load {machine['load_average_1_5_15']}" + ), + f"- {report['recall_note']}", + "", + f"| arm | p50 ms | p95 ms | recall@{K} |", + "| --- | --- | --- | --- |", + ] + for row in report["arms"]: + lines.append( + f"| {row['arm']} | {row['p50_ms']} | {row['p95_ms']} | {row[f'recall_at_{K}']} |" + ) + lines += [ + "", + "| encoding | bytes per vector |", + "| --- | --- |", + ] + lines += [f"| {name} | {size} |" for name, size in report["bytes_per_vector"].items()] + lines += [ + "", + ( + "_Latency is wall time per query on the machine named above with the thread budget the " + "benchmark pins; a run on a loaded machine is not comparable._" + ), + "", + ] + path.write_text("\n".join(lines), encoding="utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--vectors", type=Path, default=DEFAULT_VECTORS) + parser.add_argument("--probes", type=int, default=PROBES) + parser.add_argument("--out", type=Path, default=Path("bench/reports/QUANT.md")) + parser.add_argument("--json-out", type=Path, default=Path("bench/reports/quant.json")) + args = parser.parse_args() + if args.probes <= 0: + raise SystemExit(f"--probes must be positive, got {args.probes}") + + report = measure(load_matrix(args.vectors), args.probes) + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8") + write_markdown(report, args.out) + for row in report["arms"]: + print(f"{row['arm']:42} p50 {row['p50_ms']:7.1f}ms recall@{K} {row[f'recall_at_{K}']}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/bench/scale_check.py b/bench/scale_check.py index d81e971..341f9ba 100644 --- a/bench/scale_check.py +++ b/bench/scale_check.py @@ -14,33 +14,35 @@ import time from pathlib import Path from typing import Any +from bench.dense_scale import copy_store from bench.query_set import load_queries, scale_probes +from bench.report import git_commit from bench.run import latency from bench.systems import machine_facts DEFAULT_HOME = Path("bench/runs/baseline") +DEFAULT_SCRATCH = Path("bench/runs/scale") MIN_ATOMS = 100_000 DEFAULT_PROBES = 300 -def measure(home: Path, probes: int) -> dict[str, Any]: - """Open the product store (migrating it in place if it is old) and time real searches.""" +def measure(home: Path, scratch: Path, probes: int) -> dict[str, Any]: + """Open a copy of the product store (migrating it if it is old) and time real searches.""" from ely_eye.config import Settings from ely_eye.db import Database from ely_eye.retrieval import RetrievalService - if not (home / "data" / "ely_eye.sqlite").is_file(): - raise SystemExit( - f"{home / 'data' / 'ely_eye.sqlite'}: no product database, " - "build it with python -m bench.ingest_baseline" - ) - settings = Settings(home=home) + copy_home, copy_seconds = copy_store(home, scratch) + settings = Settings(home=copy_home) started = time.perf_counter() database = Database(settings) migration_seconds = time.perf_counter() - started atoms = database.atom_count() if atoms < MIN_ATOMS: raise SystemExit(f"{database.path}: {atoms} atoms, this check needs at least {MIN_ATOMS}") + with database.connect() as connection: + schema = int(connection.execute("PRAGMA user_version").fetchone()[0]) + print(f"{atoms} atoms copied in {copy_seconds:.1f}s, migrated in {migration_seconds:.1f}s", flush=True) queries = scale_probes(load_queries(), probes) service = RetrievalService(settings, database) @@ -50,10 +52,15 @@ def measure(home: Path, probes: int) -> dict[str, Any]: began = time.perf_counter() service.search(query.text) times.append(time.perf_counter() - began) + if len(times) % 100 == 0: + print(f"searched {len(times)}/{len(queries)}", flush=True) return { + "commit": git_commit(), "database": str(database.path), "atoms": atoms, + "copy_seconds": copy_seconds, + "schema_version": schema, "migration_seconds": migration_seconds, "dense_embeddings_enabled": settings.dense_embeddings_enabled, "probes": len(queries), @@ -71,13 +78,18 @@ def write_markdown(report: dict[str, Any], path: Path) -> None: [ "# SCALE: the product's own search at corpus size", "", + f"- commit: `{report['commit']}`", f"- database: {report['database']}", + ( + f"- copied from the frozen baseline in {report['copy_seconds']:.1f}s and migrated " + f"in {report['migration_seconds']:.1f}s to schema {report['schema_version']}; " + "the baseline store itself is left untouched" + ), ( f"- atoms: {report['atoms']}, probes: {report['probes']} " f"({', '.join(report['probe_modalities'])})" ), f"- dense embeddings: {report['dense_embeddings_enabled']}", - f"- opening the store (migrations included): {report['migration_seconds']:.3f}s", ( f"- p50 {timing['p50']:.4f}s · p95 {timing['p95']:.4f}s · max {timing['max']:.4f}s " f"over {timing['queries']} queries" @@ -88,6 +100,9 @@ def write_markdown(report: dict[str, Any], path: Path) -> None: f"load {machine['load_average_1_5_15']}" ), "", + "The baseline recorded p50 20.16s on this corpus because the old search pulled every", + "atom into Python and rebuilt a BM25 index per query.", + "", ( "_Latency is wall time per search call on an idle machine with the thread budget " "the benchmark pins; a run on a loaded machine is not comparable._" @@ -102,6 +117,7 @@ def write_markdown(report: dict[str, Any], path: Path) -> None: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--home", type=Path, default=DEFAULT_HOME) + parser.add_argument("--scratch", type=Path, default=DEFAULT_SCRATCH) parser.add_argument("--probes", type=int, default=DEFAULT_PROBES) parser.add_argument("--out", type=Path, default=Path("bench/reports/SCALE.md")) parser.add_argument("--json-out", type=Path, default=Path("bench/reports/scale.json")) @@ -109,7 +125,7 @@ def main() -> None: if args.probes <= 0: raise SystemExit(f"--probes must be positive, got {args.probes}") - report = measure(args.home, args.probes) + report = measure(args.home, args.scratch, args.probes) args.json_out.parent.mkdir(parents=True, exist_ok=True) args.json_out.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8") write_markdown(report, args.out) diff --git a/ely-eye-mcp/internal/domain/types.go b/ely-eye-mcp/internal/domain/types.go index f4bd700..05f4f2c 100644 --- a/ely-eye-mcp/internal/domain/types.go +++ b/ely-eye-mcp/internal/domain/types.go @@ -27,7 +27,8 @@ type ForgetReport struct { AtomsDeleted int `json:"atoms_deleted"` EdgesDeleted int `json:"edges_deleted"` CacheEntriesCleared int `json:"cache_entries_cleared"` - VectorIndexDropped bool `json:"vector_index_dropped"` + VectorsDropped int `json:"vectors_dropped"` + VisualIndexDropped bool `json:"visual_index_dropped"` AffectedCartridges []any `json:"affected_cartridges"` Reason string `json:"reason"` } diff --git a/frontend/src/components/MemoryPanel.tsx b/frontend/src/components/MemoryPanel.tsx index 48b9b00..7a0dbec 100644 --- a/frontend/src/components/MemoryPanel.tsx +++ b/frontend/src/components/MemoryPanel.tsx @@ -103,7 +103,8 @@ export function MemoryPanel({ onChanged }: { onChanged: () => void }) { setNotice( `forgotten: ${report.atoms_deleted} atom${report.atoms_deleted === 1 ? '' : 's'}, ` + `${report.edges_deleted} links, ${report.cache_entries_cleared} cache entries` + - (report.vector_index_dropped ? ', dense index dropped' : ''), + (report.vectors_dropped ? `, ${report.vectors_dropped} dense vectors dropped` : '') + + (report.visual_index_dropped ? ', picture index dropped' : ''), ); setResult(null); await loadBrowse(); diff --git a/frontend/src/components/SensesPanel.tsx b/frontend/src/components/SensesPanel.tsx index c634ced..e466cb6 100644 --- a/frontend/src/components/SensesPanel.tsx +++ b/frontend/src/components/SensesPanel.tsx @@ -13,7 +13,7 @@ export function SensesPanel({ senses }: { senses: SenseStatus | null }) { >; reason: string; }