Rank memory search in SQL instead of rebuilding BM25 in Python

The product path pulled every atom into Python and rebuilt a BM25Okapi
index per query: p50 20.16s over 165,973 atoms. It now runs one SQLite
statement with LIMIT/OFFSET.

- text_index.py: one normalizer for both sides of the index. CJK runs are
  split per character (unicode61 cannot segment them), camelCase and
  separators are split so "search atoms" finds SearchAtoms, latin tokens
  get a prefix match. Rows written before this existed hold raw text in
  atom_fts, so schema v3 rebuilds the index on upgrade.
- retrieval.py: filters (holder / scenario / source / modality / since /
  until / min_trust), a base64 cursor, and reciprocal rank fusion in place
  of the 0.55/0.35/0.10 blend whose graph term was always zero.
- An empty result is {status: "no_match"}, not a bare array.
- Go: search_atoms and compile_context rank through POST /api/search. The
  second Go tokenizer is deleted; two implementations could drift, which
  is what the MCP/HTTP parity requirement forbids.
- db.py passed 500 lines, so the row codec moved to row_codec.py.
- Tooling: installed the repo's own pyright. pyrightconfig.json had no
  exclude and swept the downloaded corpus (1171 errors); with excludes the
  count is 62, all pre-existing (evals 48, profiles 4, training 4,
  ingestion 3, runtime 2, cartridge_assets 1). backend/pyproject.toml
  carried a second [tool.pyright] that shadowed the root config; deleted.
- bench/scale_check.py measures the product search at corpus size. The
  100k-atom p50 number is not written until it is measured on an idle
  machine.

Verification: 112 pytest passed; go build/vet/test clean, gofmt clean;
ruff clean on every touched file; mypy clean on the 14 touched Python
files; repo pyright 62 errors, unchanged per-file ledger.
This commit is contained in:
ZacharyZhang-NY committed 2026-10-07 16:02:40 -04:00
1 parent 3d792417cf
commit 674c88e84e
29 files changed
+1097 -356

No files matched your search

+12
View File
@@ -38,6 +38,7 @@ from .ingestion import IngestionService
from .memory import MemoryService
from .profiles import runtime_profile_report
from .research_modules import ResearchModuleService
from .retrieval import RetrievalService
from .runtime import QwenRuntime, RuntimeUnavailableError
from .schemas import (
AdapterManifest,
@@ -53,6 +54,8 @@ from .schemas import (
RememberRequest,
RuntimeProfile,
RuntimeProfileReport,
SearchRequest,
SearchResult,
SupersedeRequest,
SystemStatus,
UploadIngestRequest,
@@ -190,6 +193,15 @@ def create_app() -> FastAPI:
runtime_backend=generation.backend,
)
@app.post("/api/search", response_model=SearchResult)
def search(request: SearchRequest, settings: Annotated[Settings, Depends(settings_dep)]) -> SearchResult:
try:
return RetrievalService(settings).search(
request.query, request.top_k, filters=request.filters, cursor=request.cursor
)
except ValueError as exc:
raise HTTPException(status_code=422, detail=str(exc)) from exc
@app.post("/api/memory", response_model=EvidenceAtom)
def remember(
request: RememberRequest, settings: Annotated[Settings, Depends(settings_dep)]
+33 -2
View File
@@ -24,8 +24,9 @@ from .ingestion import IngestionService
from .memory import MemoryService
from .profiles import runtime_profile_report
from .research_modules import ResearchModuleService
from .retrieval import RetrievalService
from .runtime import QwenRuntime
from .schemas import MemoryOrigin, TimePrecision, TrainingAdapterKind
from .schemas import MemoryOrigin, SearchFilters, TimePrecision, TrainingAdapterKind
from .training import TrainingService
app = typer.Typer(no_args_is_help=True)
@@ -79,6 +80,36 @@ def _moment(value: str, option: str) -> datetime | None:
raise ValueError(f"{option} {value!r} is not an ISO timestamp: {error}") from error
@app.command()
def search(
query: str,
top_k: int = typer.Option(0, "--top-k", help="hits to return; 0 uses the configured default"),
holder: str = typer.Option("", "--holder", help="only this holder's memories may answer"),
scenario: str = typer.Option("", "--scenario"),
source_id: str = typer.Option("", "--source-id"),
since: str = typer.Option("", "--since", help="event time at or after, ISO with offset"),
until: str = typer.Option("", "--until"),
min_trust: float = typer.Option(-1.0, "--min-trust", help="drop memories whose trust falls below this"),
cursor: str = typer.Option("", "--cursor", help="next_cursor from the previous page"),
) -> None:
"""Rank the memories that may answer, and say so when none may."""
settings = get_settings()
result = RetrievalService(settings).search(
query,
top_k or None,
filters=SearchFilters(
holder=holder or None,
scenario=scenario or None,
source_id=source_id or None,
since=_moment(since, "--since"),
until=_moment(until, "--until"),
min_trust=None if min_trust < 0 else min_trust,
),
cursor=cursor or None,
)
typer.echo(result.model_dump_json(indent=2))
@app.command()
def remember(
text: str,
@@ -289,7 +320,7 @@ def build_visual_training_data(
prd_path: Path = Path("PRD.md"),
output_path: Path | None = None,
cartridge_id: str | None = None,
image_path: list[Path] | None = typer.Option(None, "--image-path"),
image_path: list[Path] | None = typer.Option(None, "--image-path"), # noqa: B008 # typer reads option defaults itself
) -> None:
path = TrainingService(get_settings()).build_visual_dataset(
prd_path=prd_path,
+3 -4
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import hashlib
from typing import Iterable
from collections.abc import Iterable
from .cache_fabric import CacheFabric
from .config import Settings, get_settings
@@ -9,7 +9,6 @@ from .db import Database
from .retrieval import RetrievalService
from .schemas import CompiledContext, ContextPlan, Modality, RetrievalHit, RuntimeProfile
PROFILE_TOKEN_BUDGETS = {
RuntimeProfile.live_demo: 128_000,
RuntimeProfile.extreme_context: 1_010_000,
@@ -34,9 +33,9 @@ class ContextCompiler:
return compiled
plan = self.plan(question, profile)
hits = self.retrieval.search(plan.retrieval_query, plan.evidence_budget_atoms)
hits = self.retrieval.search(plan.retrieval_query, plan.evidence_budget_atoms).hits
packed_text, token_total = self._pack_hits(hits, plan.token_budget, plan.retrieval_query)
trace_id = hashlib.sha256(f"{cache_key}:{token_total}".encode("utf-8")).hexdigest()[:16]
trace_id = hashlib.sha256(f"{cache_key}:{token_total}".encode()).hexdigest()[:16]
compiled = CompiledContext(
plan=plan,
hits=hits,
+90 -113
View File
@@ -1,35 +1,27 @@
from __future__ import annotations
import json
import sqlite3
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any
from .config import Settings, get_settings
from .migrations import apply as apply_migrations
from .schemas import EvidenceAtom, MemoryOrigin, SourceRecord, TimePrecision, utc_now
from .row_codec import atom_params, from_json, row_to_atom, row_to_source, to_json
from .schemas import EvidenceAtom, SourceRecord, utc_now
from .text_index import fts_query, index_text
def _json(value: Any) -> str:
return json.dumps(value, ensure_ascii=False, separators=(",", ":"), default=str)
@dataclass(frozen=True)
class FtsPage:
"""One page of BM25 hits plus whether the corpus holds more: the cursor is an offset, and a
ranking this small re-reads rather than keeps a server-side position."""
def _loads(value: str | None, default: Any) -> Any:
if value is None or value == "":
return default
try:
return json.loads(value)
except json.JSONDecodeError as error:
raise ValueError(f"stored JSON is unreadable: {error}") from error
def _datetime(value: str | None) -> datetime | None:
if value is None or value == "":
return None
return datetime.fromisoformat(value)
rows: list[tuple[str, float]]
has_more: bool
class Database:
@@ -80,7 +72,7 @@ class Database:
source.sha256,
source.size_bytes,
source.created_at.isoformat(),
_json(source.metadata),
to_json(source.metadata),
)
for source in sources
],
@@ -123,43 +115,14 @@ class Database:
visibility_json=excluded.visibility_json,
retention_until=excluded.retention_until
""",
[self._atom_params(atom) for atom in atoms],
[atom_params(atom) for atom in atoms],
)
connection.executemany("DELETE FROM atom_fts WHERE atom_id = ?", [(atom.atom_id,) for atom in atoms])
connection.executemany(
"INSERT INTO atom_fts(atom_id, source, text) VALUES (?, ?, ?)",
[(atom.atom_id, atom.source, atom.text) for atom in atoms],
[(atom.atom_id, index_text(atom.source), index_text(atom.text)) for atom in atoms],
)
def _atom_params(self, atom: EvidenceAtom) -> tuple[Any, ...]:
return (
atom.atom_id,
atom.source_id,
atom.modality.value,
atom.source,
atom.ingested_at.isoformat(),
atom.text,
atom.image_ref,
_json(atom.embedding_refs),
_json(atom.layout.model_dump() if atom.layout else None),
_json(atom.relations),
_json(atom.trust.model_dump()),
atom.token_equivalent,
_json(atom.metadata),
atom.holder_id,
atom.origin.value,
atom.scenario,
atom.event_time.isoformat() if atom.event_time else None,
atom.event_time_end.isoformat() if atom.event_time_end else None,
atom.event_time_precision.value,
atom.event_time_source,
atom.valid_from.isoformat() if atom.valid_from else None,
atom.valid_to.isoformat() if atom.valid_to else None,
_json(atom.supersedes),
_json(atom.visibility.model_dump()),
atom.retention_until.isoformat() if atom.retention_until else None,
)
def list_atoms(self, limit: int | None = None) -> list[EvidenceAtom]:
with self.connect() as connection:
if limit is None:
@@ -168,7 +131,7 @@ class Database:
rows = connection.execute(
"SELECT * FROM atoms ORDER BY ingested_at DESC LIMIT ?", (limit,)
).fetchall()
return [self._row_to_atom(row) for row in rows]
return [row_to_atom(row) for row in rows]
def atoms_as_of(self, moment: datetime, holder: str | None = None) -> list[EvidenceAtom]:
"""Memories that were open at `moment`, the question a parameterless view cannot ask.
@@ -195,31 +158,89 @@ class Database:
"ORDER BY ingested_at DESC",
(holder, stamp, stamp),
).fetchall()
return [self._row_to_atom(row) for row in rows]
return [row_to_atom(row) for row in rows]
def get_atoms(self, atom_ids: list[str]) -> list[EvidenceAtom]:
if not atom_ids:
return []
with self.connect() as connection:
rows = connection.execute(
"SELECT * FROM atoms WHERE atom_id IN (SELECT value FROM json_each(?))", (_json(atom_ids),)
"SELECT * FROM atoms WHERE atom_id IN (SELECT value FROM json_each(?))", (to_json(atom_ids),)
).fetchall()
by_id = {row["atom_id"]: self._row_to_atom(row) for row in rows}
by_id = {row["atom_id"]: row_to_atom(row) for row in rows}
return [by_id[atom_id] for atom_id in atom_ids if atom_id in by_id]
def search_fts(self, query: str, limit: int) -> list[tuple[str, float]]:
"""Best BM25 matches over the whole corpus, no filters: the ranking the product needs most days."""
return self.search_atoms(query, limit=limit).rows
def search_atoms(
self,
query: str,
*,
limit: int,
offset: int = 0,
holder: str | None = None,
scenario: str | None = None,
source_id: str | None = None,
modality: str | None = None,
since: datetime | None = None,
until: datetime | None = None,
min_trust: float | None = None,
) -> FtsPage:
"""BM25 over the FTS index joined to memories that are open right now.
Every filter is a `? IS NULL OR` predicate in one literal statement: the ranking never
depends on string-built SQL, and the planner sees one shape. A memory with no event time
cannot answer a time filter, so it drops out when one is asked for.
"""
with self.connect() as connection:
rows = connection.execute(
"""
SELECT atom_id, bm25(atom_fts) AS score
FROM atom_fts
SELECT f.atom_id, bm25(atom_fts) AS score
FROM atom_fts AS f
JOIN v_atoms_current AS a ON a.atom_id = f.atom_id
WHERE atom_fts MATCH ?
AND (? IS NULL OR a.holder_id = ?)
AND (? IS NULL OR a.scenario = ?)
AND (? IS NULL OR a.source_id = ?)
AND (? IS NULL OR a.modality = ?)
AND (? IS NULL OR (a.event_time IS NOT NULL AND a.event_time >= ?))
AND (? IS NULL OR (a.event_time IS NOT NULL AND a.event_time <= ?))
AND (? IS NULL OR (
a.trust_json IS NOT NULL
AND min(
json_extract(a.trust_json, '$.parser'),
json_extract(a.trust_json, '$.ocr'),
json_extract(a.trust_json, '$.model')
) >= ?
))
ORDER BY score
LIMIT ?
LIMIT ? OFFSET ?
""",
(query, limit),
(
fts_query(query),
holder,
holder,
scenario,
scenario,
source_id,
source_id,
modality,
modality,
None if since is None else since.isoformat(),
None if since is None else since.isoformat(),
None if until is None else until.isoformat(),
None if until is None else until.isoformat(),
min_trust,
min_trust,
limit + 1,
offset,
),
).fetchall()
return [(row["atom_id"], float(-row["score"])) for row in rows]
matched = [(str(row["atom_id"]), -float(row["score"])) for row in rows[:limit]]
return FtsPage(rows=matched, has_more=len(rows) > limit)
def source_id_of_atom(self, atom_id: str) -> str | None:
with self.connect() as connection:
@@ -231,7 +252,7 @@ class Database:
rows = connection.execute(
"SELECT * FROM atoms WHERE source_id = ? ORDER BY atom_id", (source_id,)
).fetchall()
return [self._row_to_atom(row) for row in rows]
return [row_to_atom(row) for row in rows]
def close_atoms(self, atom_ids: list[str], closed_at: datetime) -> int:
"""Stop treating these memories as current; the rows stay for time travel."""
@@ -241,7 +262,7 @@ class Database:
cursor = connection.execute(
"UPDATE atoms SET valid_to = ? "
"WHERE valid_to IS NULL AND atom_id IN (SELECT value FROM json_each(?))",
(closed_at.isoformat(), _json(atom_ids)),
(closed_at.isoformat(), to_json(atom_ids)),
)
return int(cursor.rowcount)
@@ -250,10 +271,10 @@ class Database:
return 0
with self.connect() as connection:
connection.execute(
"DELETE FROM atom_fts WHERE atom_id IN (SELECT value FROM json_each(?))", (_json(atom_ids),)
"DELETE FROM atom_fts WHERE atom_id IN (SELECT value FROM json_each(?))", (to_json(atom_ids),)
)
cursor = connection.execute(
"DELETE FROM atoms WHERE atom_id IN (SELECT value FROM json_each(?))", (_json(atom_ids),)
"DELETE FROM atoms WHERE atom_id IN (SELECT value FROM json_each(?))", (to_json(atom_ids),)
)
return int(cursor.rowcount)
@@ -265,7 +286,7 @@ class Database:
"DELETE FROM evidence_edges "
"WHERE from_atom_id IN (SELECT value FROM json_each(?)) "
"OR to_atom_id IN (SELECT value FROM json_each(?))",
(_json(atom_ids), _json(atom_ids)),
(to_json(atom_ids), to_json(atom_ids)),
)
return int(cursor.rowcount)
@@ -277,7 +298,7 @@ class Database:
connection.execute(
"INSERT INTO cache_events(ts, layer, cache_key, event, bytes_estimate, metadata_json) "
"VALUES (?, 'memory', ?, 'invalidated', 0, ?)",
(utc_now().isoformat(), reason, _json({"entries_cleared": cleared})),
(utc_now().isoformat(), reason, to_json({"entries_cleared": cleared})),
)
return cleared
@@ -307,7 +328,7 @@ class Database:
"SELECT source_id, path, kind, mime, sha256, size_bytes, created_at, metadata_json "
"FROM sources ORDER BY created_at DESC"
).fetchall()
return [self._row_to_source(row) for row in rows]
return [row_to_source(row) for row in rows]
def cartridge_count(self) -> int:
with self.connect() as connection:
@@ -376,7 +397,7 @@ class Database:
cache_key,
event,
bytes_estimate,
_json(metadata or {}),
to_json(metadata or {}),
),
)
@@ -395,7 +416,7 @@ class Database:
str(edge["to_atom_id"]),
str(edge["relation"]),
float(edge["confidence"]),
_json(edge.get("metadata") or {}),
to_json(edge.get("metadata") or {}),
)
for edge in edges
],
@@ -420,7 +441,7 @@ class Database:
"to_atom_id": row["to_atom_id"],
"relation": row["relation"],
"confidence": float(row["confidence"]),
"metadata": _loads(row["metadata_json"], {}),
"metadata": from_json(row["metadata_json"], {}),
}
for row in rows
]
@@ -440,47 +461,3 @@ class Database:
GROUP BY layer
"""
).fetchall()
def _row_to_atom(self, row: sqlite3.Row) -> EvidenceAtom:
from .schemas import LayoutBox, Modality, TrustScores, Visibility
layout_data = _loads(row["layout_json"], None)
return EvidenceAtom(
atom_id=row["atom_id"],
modality=Modality(row["modality"]),
source_id=row["source_id"],
source=row["source"],
ingested_at=datetime.fromisoformat(row["ingested_at"]),
text=row["text"],
image_ref=row["image_ref"],
embedding_refs=_loads(row["embedding_refs_json"], []),
layout=LayoutBox(**layout_data) if layout_data else None,
relations=_loads(row["relations_json"], []),
trust=TrustScores(**_loads(row["trust_json"], {"parser": 0.0})),
token_equivalent=int(row["token_equivalent"]),
metadata=_loads(row["metadata_json"], {}),
holder_id=row["holder_id"],
origin=MemoryOrigin(row["origin"]),
scenario=row["scenario"],
event_time=_datetime(row["event_time"]),
event_time_end=_datetime(row["event_time_end"]),
event_time_precision=TimePrecision(row["event_time_precision"]),
event_time_source=row["event_time_source"],
valid_from=_datetime(row["valid_from"]),
valid_to=_datetime(row["valid_to"]),
supersedes=_loads(row["supersedes_json"], []),
visibility=Visibility(**_loads(row["visibility_json"], {})),
retention_until=_datetime(row["retention_until"]),
)
def _row_to_source(self, row: sqlite3.Row) -> SourceRecord:
return SourceRecord(
source_id=row["source_id"],
path=row["path"],
kind=row["kind"],
mime=row["mime"],
sha256=row["sha256"],
size_bytes=int(row["size_bytes"]),
created_at=datetime.fromisoformat(row["created_at"]),
metadata=_loads(row["metadata_json"], {}),
)
+25 -7
View File
@@ -12,7 +12,9 @@ from __future__ import annotations
import sqlite3
SCHEMA_VERSION = 2
from .text_index import index_text
SCHEMA_VERSION = 3
ATOM_COLUMNS_V2 = """
atom_id TEXT PRIMARY KEY,
@@ -85,12 +87,14 @@ _BASE_TABLES = """
metadata_json TEXT NOT NULL
);
CREATE VIRTUAL TABLE IF NOT EXISTS atom_fts USING fts5(
atom_id UNINDEXED,
source,
text,
tokenize = 'unicode61'
);
CREATE VIRTUAL TABLE IF NOT EXISTS atom_fts USING fts5(
atom_id UNINDEXED,
source,
text,
tokenize = 'unicode61'
);
-- source and text hold text_index output, not the atom's own text: unicode61 sees one token
-- per CJK character and per identifier fragment, which is what makes those queries match.
"""
# A memory is usable now when nothing has closed it, it had begun by now, and its
@@ -178,6 +182,20 @@ def apply(connection: sqlite3.Connection) -> None:
_upgrade_v1_to_v2(connection)
else:
_create_v2(connection)
if version < 3:
_rebuild_fts(connection)
connection.executescript(VIEWS)
connection.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
def _rebuild_fts(connection: sqlite3.Connection) -> None:
"""Re-index every atom through text_index: rows indexed before it existed hold raw text, and
FTS5 cannot match 钥匙 inside a token that is a whole Chinese sentence."""
connection.execute("DROP TABLE IF EXISTS atom_fts")
connection.executescript(_BASE_TABLES)
rows = connection.execute("SELECT atom_id, source, text FROM atoms ORDER BY atom_id").fetchall()
connection.executemany(
"INSERT INTO atom_fts(atom_id, source, text) VALUES (?, ?, ?)",
[(row[0], index_text(row[1]), index_text(row[2])) for row in rows],
)
+98 -67
View File
@@ -1,14 +1,23 @@
"""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 re
import base64
from collections import defaultdict
from rank_bm25 import BM25Okapi
from .config import Settings, get_settings
from .db import Database
from .embeddings import EmbeddingService
from .schemas import EvidenceAtom, RetrievalHit
from .schemas import EvidenceAtom, RetrievalHit, SearchFilters, SearchResult
RRF_K = 60
CANDIDATE_WINDOW = 10
class RetrievalService:
@@ -17,52 +26,47 @@ class RetrievalService:
self.db = database or Database(self.settings)
self.embeddings = EmbeddingService(self.settings, self.db)
def search(self, query: str, top_k: int | None = None) -> list[RetrievalHit]:
def search(
self,
query: str,
top_k: int | None = None,
*,
filters: SearchFilters | None = None,
cursor: str | 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
atoms = self.db.list_atoms()
if not atoms:
return []
sparse_scores = self._sparse_search(query, atoms, limit * 3)
dense_scores: dict[str, float] = {}
if limit <= 0:
raise ValueError(f"search needs a positive limit, got {limit}")
offset = decode_cursor(cursor)
page = self.db.search_atoms(
query,
limit=limit * CANDIDATE_WINDOW,
offset=offset,
**(filters or SearchFilters()).model_dump(exclude_none=True),
)
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_scores = self._dense_search(query, limit * 3)
all_ids = set(sparse_scores) | set(dense_scores)
atom_map = {atom.atom_id: atom for atom in atoms}
hits: list[RetrievalHit] = []
for atom_id in all_ids:
atom = atom_map.get(atom_id)
if atom is None:
continue
sparse_score = sparse_scores.get(atom_id, 0.0)
dense_score = dense_scores.get(atom_id, 0.0)
graph_score = self._graph_score(atom, query)
final_score = sparse_score * 0.55 + dense_score * 0.35 + graph_score * 0.10
hits.append(
RetrievalHit(
atom=atom,
sparse_score=sparse_score,
dense_score=dense_score,
graph_score=graph_score,
final_score=final_score,
)
)
hits.sort(key=lambda hit: hit.final_score, reverse=True)
return hits[:limit]
def _sparse_search(self, query: str, atoms: list[EvidenceAtom], limit: int) -> dict[str, float]:
tokenized = [tokenize(atom.text + " " + atom.source) for atom in atoms]
bm25 = BM25Okapi(tokenized)
query_tokens = tokenize(query)
scores = bm25.get_scores(query_tokens)
lexical_scores = [lexical_score(query_tokens, tokens, atoms[index]) for index, tokens in enumerate(tokenized)]
combined = [max(float(scores[index]), lexical_scores[index]) for index in range(len(atoms))]
order = sorted(range(len(combined)), key=lambda index: combined[index], reverse=True)[:limit]
if not order:
return {}
max_score = max(float(combined[index]) for index in order) or 1.0
return {atoms[index].atom_id: float(combined[index]) / max_score for index in order if combined[index] > 0}
dense = self._dense_search(query, limit * CANDIDATE_WINDOW)
rankings.append(list(dense))
for atom_id, cosine in dense.items():
scores[atom_id]["dense_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) -> dict[str, float]:
index = self.embeddings.load_index()
@@ -71,27 +75,38 @@ class RetrievalService:
query_vector = self.embeddings.embed_query(query)
return dict(index.search(query_vector, limit))
def _graph_score(self, atom: EvidenceAtom, query: str) -> float:
query_tokens = set(tokenize(query))
relation_tokens = set()
for relation in atom.relations:
relation_tokens.update(tokenize(relation))
if not query_tokens or not relation_tokens:
return 0.0
return len(query_tokens & relation_tokens) / len(query_tokens)
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),
final_score=fused[atom_id],
)
)
return hits
def tokenize(text: str) -> list[str]:
return [token.lower() for token in re.findall(r"[\w\u4e00-\u9fff]+", text)]
def lexical_score(query_tokens: list[str], document_tokens: list[str], atom: EvidenceAtom) -> float:
if not query_tokens or not document_tokens:
return 0.0
document_set = set(document_tokens)
overlap = sum(1 for token in query_tokens if token in document_set)
phrase_hits = sum(1 for token in query_tokens if token and token in atom.text.lower())
return (overlap + 0.25 * phrase_hits) / len(query_tokens)
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 group_hits_by_source(hits: list[RetrievalHit]) -> dict[str, list[RetrievalHit]]:
@@ -99,3 +114,19 @@ def group_hits_by_source(hits: list[RetrievalHit]) -> dict[str, list[RetrievalHi
for hit in hits:
grouped[hit.atom.source].append(hit)
return dict(grouped)
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
+116
View File
@@ -0,0 +1,116 @@
"""Row ↔ model conversion for the SQLite store, kept out of the query layer.
Every column of an atom is either a scalar, an ISO timestamp, or a JSON blob; one file holds
that mapping so the store can stay about queries and the codec can stay about shapes.
"""
from __future__ import annotations
import json
import sqlite3
from datetime import datetime
from typing import Any
from .schemas import (
EvidenceAtom,
LayoutBox,
MemoryOrigin,
Modality,
SourceRecord,
TimePrecision,
TrustScores,
Visibility,
)
def to_json(value: Any) -> str:
return json.dumps(value, ensure_ascii=False, separators=(",", ":"), default=str)
def from_json(value: str | None, default: Any) -> Any:
if value is None or value == "":
return default
try:
return json.loads(value)
except json.JSONDecodeError as error:
raise ValueError(f"stored JSON is unreadable: {error}") from error
def from_iso(value: str | None) -> datetime | None:
if value is None or value == "":
return None
return datetime.fromisoformat(value)
def atom_params(atom: EvidenceAtom) -> tuple[Any, ...]:
return (
atom.atom_id,
atom.source_id,
atom.modality.value,
atom.source,
atom.ingested_at.isoformat(),
atom.text,
atom.image_ref,
to_json(atom.embedding_refs),
to_json(atom.layout.model_dump() if atom.layout else None),
to_json(atom.relations),
to_json(atom.trust.model_dump()),
atom.token_equivalent,
to_json(atom.metadata),
atom.holder_id,
atom.origin.value,
atom.scenario,
atom.event_time.isoformat() if atom.event_time else None,
atom.event_time_end.isoformat() if atom.event_time_end else None,
atom.event_time_precision.value,
atom.event_time_source,
atom.valid_from.isoformat() if atom.valid_from else None,
atom.valid_to.isoformat() if atom.valid_to else None,
to_json(atom.supersedes),
to_json(atom.visibility.model_dump()),
atom.retention_until.isoformat() if atom.retention_until else None,
)
def row_to_atom(row: sqlite3.Row) -> EvidenceAtom:
layout_data = from_json(row["layout_json"], None)
return EvidenceAtom(
atom_id=row["atom_id"],
modality=Modality(row["modality"]),
source_id=row["source_id"],
source=row["source"],
ingested_at=datetime.fromisoformat(row["ingested_at"]),
text=row["text"],
image_ref=row["image_ref"],
embedding_refs=from_json(row["embedding_refs_json"], []),
layout=LayoutBox(**layout_data) if layout_data else None,
relations=from_json(row["relations_json"], []),
trust=TrustScores(**from_json(row["trust_json"], {"parser": 0.0})),
token_equivalent=int(row["token_equivalent"]),
metadata=from_json(row["metadata_json"], {}),
holder_id=row["holder_id"],
origin=MemoryOrigin(row["origin"]),
scenario=row["scenario"],
event_time=from_iso(row["event_time"]),
event_time_end=from_iso(row["event_time_end"]),
event_time_precision=TimePrecision(row["event_time_precision"]),
event_time_source=row["event_time_source"],
valid_from=from_iso(row["valid_from"]),
valid_to=from_iso(row["valid_to"]),
supersedes=from_json(row["supersedes_json"], []),
visibility=Visibility(**from_json(row["visibility_json"], {})),
retention_until=from_iso(row["retention_until"]),
)
def row_to_source(row: sqlite3.Row) -> SourceRecord:
return SourceRecord(
source_id=row["source_id"],
path=row["path"],
kind=row["kind"],
mime=row["mime"],
sha256=row["sha256"],
size_bytes=int(row["size_bytes"]),
created_at=datetime.fromisoformat(row["created_at"]),
metadata=from_json(row["metadata_json"], {}),
)
+38
View File
@@ -150,6 +150,44 @@ class RetrievalHit(BaseModel):
graph_score: float = 0.0
final_score: float
class SearchFilters(BaseModel):
"""Who may answer, in what setting, and how fresh the memory has to be.
A memory with no event time cannot satisfy since/until, so it drops out when a time bound is
asked for rather than passing as unknown.
"""
holder: str | None = None
scenario: str | None = None
source_id: str | None = None
modality: Modality | None = None
since: datetime | None = None
until: datetime | None = None
min_trust: float | None = Field(default=None, ge=0.0, le=1.0)
@model_validator(mode="after")
def _times_are_ordered(self) -> SearchFilters:
if self.since and self.until and self.since > self.until:
raise ValueError("search since must not be after until")
return self
class SearchRequest(BaseModel):
query: str
top_k: int | None = Field(default=None, gt=0)
filters: SearchFilters | None = None
cursor: str | None = None
class SearchResult(BaseModel):
"""An empty answer says why: no_match is a claim about the corpus, [] is not."""
query: str
status: Literal["hits", "no_match"]
hits: list[RetrievalHit]
next_cursor: str | None = None
class ContextPlan(BaseModel):
question: str
profile: RuntimeProfile
+63
View File
@@ -0,0 +1,63 @@
"""One normalizer for both sides of the FTS index.
FTS5's unicode61 tokenizer splits on whitespace and punctuation only. It reads
「我们的钥匙放在玄关抽屉里」 as one token and `SearchAtoms` as one token, so a search for
钥匙 or atoms finds nothing and reports it as an empty result. Index and query must run the
same transform; keeping them in one file is what stops them from drifting apart.
"""
from __future__ import annotations
import re
from unicodedata import normalize
# Han, its extension A block, the compatibility ideographs, kana, and Hangul syllables:
# scripts written without spaces. One source of truth for both the tokenizer and the query.
_BLOCKS = ((0x3400, 0x4DBF), (0x4E00, 0x9FFF), (0xF900, 0xFAFF), (0x3040, 0x30FF), (0xAC00, 0xD7AF))
_CJK_CLASS = "".join(f"{chr(start)}-{chr(end)}" for start, end in _BLOCKS)
_CJK_CHARS = frozenset("".join("".join(chr(code) for code in range(start, end + 1)) for start, end in _BLOCKS))
_CAMEL_BOUNDARY = re.compile(r"(?<=[a-z0-9])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|(?<=[A-Za-z])(?=\d)")
# Everything a reader would not type as a word: symbols, punctuation, and the underscore
# that joins identifiers. `[^[^\W\s]` looks equivalent but deletes CJK characters on 3.14.
_SEPARATORS = re.compile(r"[^\w\s]|_")
_CJK_RUN = re.compile(f"[{_CJK_CLASS}]+")
_SPACES = re.compile(r"\s+")
def index_text(text: str) -> str:
"""Rewrite text so every searchable unit is a separate FTS5 token."""
if not text:
return ""
return _SPACES.sub(" ", _CJK_RUN.sub(lambda match: " ".join(match.group()) + " ", _bare(text))).strip()
def fts_query(query: str) -> str:
"""Build the MATCH expression for a user query, using the same split as the index.
A CJK run becomes a quoted phrase so its characters stay adjacent: unphrased, a query for
钥匙 would match text that mentions 钥 and 匙 worlds apart. Latin tokens get a prefix match,
which is what lets `atom` reach `atoms`.
"""
terms = _bare(query).split()
if not terms:
raise ValueError(f"search needs at least one searchable word, got {query!r}")
return " ".join(_phrase(term) if _is_cjk(term) else f"{_quote(term)}*" for term in terms)
def _bare(text: str) -> str:
"""Fold width and case, split identifiers, and drop punctuation, leaving CJK runs whole."""
spaced = _CAMEL_BOUNDARY.sub(" ", normalize("NFKC", text))
return _SPACES.sub(" ", _SEPARATORS.sub(" ", spaced)).strip().lower()
def _is_cjk(term: str) -> bool:
return all(char in _CJK_CHARS for char in term)
def _phrase(run: str) -> str:
return '"' + " ".join(run) + '"'
def _quote(term: str) -> str:
return '"' + term.replace('"', '""') + '"'