ZacharyZhang-NY
cd9851a12c
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.
206 lines
8.7 KiB
Python
206 lines
8.7 KiB
Python
"""Retrieval: SQLite FTS5 does the matching, RRF merges the signals that exist.
|
|
|
|
The first version pulled every atom into Python and rebuilt a BM25 index per query, so search
|
|
cost grew with the corpus instead of with the answer, and a Chinese query matched nothing
|
|
because the tokenizer had no CJK rule. Both are fixed below the score line: the index and the
|
|
query share `text_index`, and the ranking is a SQL query.
|
|
"""
|
|
|
|
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
|
|
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
|
|
CANDIDATE_WINDOW = 10
|
|
# Words decide most searches; a picture vector is only worth its encoding cost when words fell short.
|
|
THIN_SPARSE_PAGE = 2
|
|
|
|
|
|
class RetrievalService:
|
|
def __init__(self, settings: Settings | None = None, database: Database | None = None) -> None:
|
|
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(
|
|
self,
|
|
query: str,
|
|
top_k: int | None = None,
|
|
*,
|
|
filters: SearchFilters | None = None,
|
|
cursor: str | None = None,
|
|
image: Path | None = None,
|
|
) -> SearchResult:
|
|
"""Answer one query against the memories that are open right now.
|
|
|
|
The candidate window is wider than the page so fusion has something to fuse; the cursor
|
|
advances one page at a time and re-reads the tail of the window it came from, which is
|
|
why no result is ever skipped between pages.
|
|
"""
|
|
limit = top_k or self.settings.retrieval_top_k
|
|
if limit <= 0:
|
|
raise ValueError(f"search needs a positive limit, got {limit}")
|
|
if not query.strip() and image is None:
|
|
raise ValueError("search needs words or a picture to ask with, and got neither")
|
|
offset = decode_cursor(cursor)
|
|
criteria = (filters or SearchFilters()).model_dump(exclude_none=True)
|
|
scope = criteria.get("cartridge_ids")
|
|
if scope is None:
|
|
scope = cartridge_store.mounted_ids(self.db) or None
|
|
criteria["cartridge_ids"] = scope
|
|
# A picture-only question has no words for BM25, so the word ranking is skipped rather than matching an empty expression.
|
|
page = (
|
|
search_atoms(self.db, query, limit=limit * CANDIDATE_WINDOW, offset=offset, **criteria)
|
|
if query.strip()
|
|
else FtsPage(rows=[], has_more=False)
|
|
)
|
|
rankings: list[list[str]] = [[atom_id for atom_id, _ in page.rows]]
|
|
scores: dict[str, dict[str, float]] = defaultdict(dict)
|
|
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, criteria)
|
|
rankings.append(list(dense))
|
|
for atom_id, cosine in dense.items():
|
|
scores[atom_id]["dense_score"] = cosine
|
|
visual = self._visual_search(
|
|
query,
|
|
image,
|
|
limit * CANDIDATE_WINDOW,
|
|
scope,
|
|
thin=len(rankings[0]) < limit * THIN_SPARSE_PAGE,
|
|
)
|
|
if visual:
|
|
rankings.append(list(visual))
|
|
for atom_id, cosine in visual.items():
|
|
scores[atom_id]["visual_score"] = cosine
|
|
hits = self._rank(rankings, scores, limit)
|
|
more = page.has_more or len(rankings[0]) > limit
|
|
return SearchResult(
|
|
query=query,
|
|
status="hits" if hits else "no_match",
|
|
hits=hits,
|
|
next_cursor=encode_cursor(offset + limit) if hits and more else None,
|
|
)
|
|
|
|
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 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,
|
|
query: str,
|
|
image: Path | None,
|
|
limit: int,
|
|
scope: list[str] | None,
|
|
*,
|
|
thin: bool,
|
|
) -> dict[str, float]:
|
|
"""Ask the picture index when words are not enough.
|
|
|
|
A screenshot with no text on it has nothing for FTS to match, so the sparse page comes
|
|
back short; that is the signal to spend the encoding. An image query always asks for it.
|
|
"""
|
|
if image is None and not thin:
|
|
return {}
|
|
if not self.visual.available() or self.visual.load() is None:
|
|
return {}
|
|
if image is not None:
|
|
try:
|
|
ranked = self.visual.search_image(image, limit)
|
|
except (OSError, ValueError) as error:
|
|
raise ValueError(f"the query image {image} cannot be read: {error}") from error
|
|
else:
|
|
ranked = self.visual.search_text(query, limit)
|
|
if scope:
|
|
allowed = cartridge_store.atoms_of_cartridges(self.db, scope)
|
|
ranked = [(atom_id, score) for atom_id, score in ranked if atom_id in allowed]
|
|
return dict(ranked)
|
|
|
|
def _rank(self, rankings: list[list[str]], scores: dict[str, dict[str, float]], limit: int) -> list[RetrievalHit]:
|
|
"""Reciprocal rank fusion: rank position is comparable across signals, raw scores are not.
|
|
|
|
The old blend multiplied BM25 by 0.55 and cosine by 0.35, which compared a term-frequency
|
|
count to a bounded cosine and let the larger number decide. graph_score stays 0.0 until
|
|
evidence edges are retrieved as their own ranking.
|
|
"""
|
|
fused: dict[str, float] = {}
|
|
for ranking in rankings:
|
|
for position, atom_id in enumerate(ranking):
|
|
fused[atom_id] = fused.get(atom_id, 0.0) + 1.0 / (RRF_K + position + 1)
|
|
order = sorted(fused, key=lambda atom_id: (-fused[atom_id], atom_id))[:limit]
|
|
atoms = {atom.atom_id: atom for atom in self.db.get_atoms(order)}
|
|
hits: list[RetrievalHit] = []
|
|
for atom_id in order:
|
|
atom = atoms.get(atom_id)
|
|
if atom is None:
|
|
raise ValueError(f"search returned {atom_id} but the atom is not in the store")
|
|
signal = scores.get(atom_id, {})
|
|
hits.append(
|
|
RetrievalHit(
|
|
atom=atom,
|
|
sparse_score=signal.get("sparse_score", 0.0),
|
|
dense_score=signal.get("dense_score", 0.0),
|
|
visual_score=signal.get("visual_score", 0.0),
|
|
final_score=fused[atom_id],
|
|
)
|
|
)
|
|
return hits
|
|
|
|
def recall(self, query: str, top_k: int | None = None, *, filters: SearchFilters | None = None) -> list[EvidenceAtom]:
|
|
"""The atoms a search would put in front of a model, no scoring attached."""
|
|
return [hit.atom for hit in self.search(query, top_k, filters=filters).hits]
|
|
|
|
|
|
def encode_cursor(offset: int) -> str:
|
|
return base64.urlsafe_b64encode(str(offset).encode()).decode()
|
|
|
|
|
|
def decode_cursor(cursor: str | None) -> int:
|
|
if cursor is None or cursor == "":
|
|
return 0
|
|
try:
|
|
offset = int(base64.urlsafe_b64decode(cursor.encode()).decode())
|
|
except (ValueError, UnicodeDecodeError) as error:
|
|
raise ValueError(f"search cursor is unreadable: {cursor!r}") from error
|
|
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()
|