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.
This commit is contained in:
1 parent
be08e7bf9c
commit
cd9851a12c
26 files changed
+1602
-118
No files matched your search
@@ -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",
|
||||
|
||||
@@ -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()))
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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`"
|
||||
)
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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"
|
||||
|
||||
|
||||
Reference in new issue
Block a user