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:
1 parent
3d792417cf
commit
674c88e84e
29 files changed
+1097
-356
No files matched your search
@@ -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)]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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"], {}),
|
||||
)
|
||||
@@ -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],
|
||||
)
|
||||
@@ -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
|
||||
@@ -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"], {}),
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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('"', '""') + '"'
|
||||
Reference in new issue
Block a user