From 674c88e84e47906a018346dc88932ae8fcb590c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Wed, 7 Oct 2026 16:02:40 -0400 Subject: [PATCH] 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. --- CHECKLIST.md | 28 +-- backend/pyproject.toml | 3 - backend/src/ely_eye/api.py | 12 ++ backend/src/ely_eye/cli.py | 35 ++- backend/src/ely_eye/compiler.py | 7 +- backend/src/ely_eye/db.py | 203 ++++++++---------- backend/src/ely_eye/migrations.py | 32 ++- backend/src/ely_eye/retrieval.py | 165 ++++++++------ backend/src/ely_eye/row_codec.py | 116 ++++++++++ backend/src/ely_eye/schemas.py | 38 ++++ backend/src/ely_eye/text_index.py | 63 ++++++ backend/tests/test_memory_schema.py | 19 ++ backend/tests/test_search_api.py | 114 ++++++++++ backend/tests/test_text_index.py | 75 +++++++ bench/run.py | 3 +- bench/scale_check.py | 124 +++++++++++ bench/systems.py | 13 +- bench/tests/test_joint.py | 8 +- bench/tests/test_systems.py | 15 ++ ely-eye-mcp/internal/apiclient/client.go | 50 +++++ ely-eye-mcp/internal/domain/types.go | 1 + .../internal/mcpserver/search_tools_test.go | 109 ++++++++++ ely-eye-mcp/internal/mcpserver/tools.go | 19 +- ely-eye-mcp/internal/store/atoms.go | 131 ----------- package-lock.json | 45 ++++ package.json | 8 + pyproject.toml | 7 - pyrightconfig.json | 3 +- scripts/mcp_write_smoke.sh | 7 + 29 files changed, 1097 insertions(+), 356 deletions(-) create mode 100644 backend/src/ely_eye/row_codec.py create mode 100644 backend/src/ely_eye/text_index.py create mode 100644 backend/tests/test_search_api.py create mode 100644 backend/tests/test_text_index.py create mode 100644 bench/scale_check.py create mode 100644 ely-eye-mcp/internal/mcpserver/search_tools_test.go create mode 100644 package-lock.json create mode 100644 package.json diff --git a/CHECKLIST.md b/CHECKLIST.md index f1dfdc5..6f6212e 100644 --- a/CHECKLIST.md +++ b/CHECKLIST.md @@ -60,20 +60,22 @@ - [x] p0-1.4 `PRAGMA user_version` 迁移框架 + 旧库可重放（`migrations.py`：v1 库原地 ALTER + RENAME COLUMN 升到 v2，行不丢；版本高于本构建则 SystemExit） - [x] p0-1.5 视图 `v_atoms_current` + 时间旅行查询——偏差记录：SQLite 视图不能带参数，所以 as-of 是 `db.atoms_as_of(moment, holder=None)`，谓词与视图一致但去掉 retention（retention 是保存策略，不是当时是否为真） - 验收：`pytest backend/tests/test_memory_schema.py` 绿（8 条，含 v1→v2 升级不丢行、时间旅行返回历史版本）；实际路径是 `backend/tests/`，不是 `tests/` -- [ ] **p0-2** 写入路径 - - [ ] p0-2.1 `remember` / `forget` / `supersede` 三写，Python 单一写者 - - [ ] p0-2.2 MCP 工具面 + HTTP `POST /api/memory{,/forget,/supersede}` + CLI 三面同源 - - [ ] p0-2.3 服务端盖 provenance：`ingested_at`、`holder_id`、`origin` 由服务端写，客户端不可传 - - [ ] p0-2.4 `forget` 级联：派生 atom、卡带成员、向量、元信念 - - 验收：`pytest tests/test_write_path.py tests/test_forget_cascade.py` 绿；端到端「记住→重启→召回」 +- [x] **p0-2** 写入路径 + - [x] p0-2.1 `remember` / `forget` / `supersede` 三写，Python 单一写者（`memory.py:MemoryService`；同一 holder 同一句子是更新，两个 holder 各留一条） + - [x] p0-2.2 MCP 工具面 + HTTP `POST /api/memory{,/supersede,/forget}` + CLI 三面同源——偏差记录：Go 侧不碰 SQLite 写，三个工具经 HTTP 转给 Python 写者（`ELY_EYE_API`，默认 `http://127.0.0.1:8765`），写规则只有一份；顺带把 Go 读路径跟上 v2（`atoms.time`→`ingested_at`，读改走 `v_atoms_current`，fetch/search 现在带 holder/origin/scenario/event_time） + - [x] p0-2.3 服务端盖 provenance：`ingested_at`、`holder_id`、`origin` 由服务端写，客户端不可传（`RememberRequest` 里根本没有这三个字段；HTTP 测试里 POST `holder_id=attacker&origin=third_party` 仍落成 `wife`/`self`） + - [x] p0-2.4 `forget` 级联：删 atom 与 FTS 行、删证据边、清 cache_events（留一条 invalidated 轨迹）、删 dense 向量索引——偏差记录：卡带成员是**报告**而非重写，已引用 DNA 的导出卡带不该被偷偷改掉 + - 验收：`pytest backend/tests/test_memory_write.py backend/tests/test_memory_api.py backend/tests/test_memory_cli.py` 绿（21 条）；端到端跨进程 `bash scripts/mcp_write_smoke.sh` → `PASS: remember, supersede, search and forget work through MCP`（Go MCP + Python API 两个进程读写同一 SQLite，supersede 后旧记忆不再可检，forget 后 fetch 报 not found）；偏差：checklist 写的 `tests/test_write_path.py` / `tests/test_forget_cascade.py` 路径不存在，实际在 `backend/tests/` + - 顺带修的根因：`import fitz` 把 deprecation 警告打到 **stdout**，污染所有 CLI 命令的 JSON 输出（改 `import pymupdf as fitz`）；`cli status` 原本 echo 一个 Python dict repr，改 `json_dump` - [ ] **p0-3** 检索改 SQL-only - - [ ] p0-3.1 用 `db.search_fts()` 取代 `list_atoms()`，`LIMIT top_k*10` - - [ ] p0-3.2 分词修复：CJK 分词 + camelCase/标识符切分 + 前缀 - - [ ] p0-3.3 空命中返回显式 `no_match`，不返回裸空数组 - - [ ] p0-3.4 过滤器：`since/until/source/scenario/holder/modality/trust` + 游标分页 - - [ ] p0-3.5 RRF 取代 `0.55/0.35/0.10` 加权混合 - - [ ] p0-3.6 MCP 与 HTTP 收敛到同一排序实现 - - 验收：`MATCH "钥匙"` 命中「我们的钥匙放在玄关抽屉里」；`search atoms` 命中 `func SearchAtoms`；100k atom p50 <100ms；`pytest tests/test_retrieval_parity.py` 绿 + - [x] p0-3.1 `retrieval.py` 不再 `list_atoms()`：一条 SQL 里 `LIMIT ? OFFSET ?`，候选窗口 `CANDIDATE_WINDOW = 10`（`db.search_atoms()` 返回 `FtsPage{rows, has_more}`） + - [x] p0-3.2 分词修复：`text_index.py` 一处归一化，索引侧与查询侧共用（NFKC + camelCase 切分 + CJK 逐字 + 前缀 `"token"*` + CJK 串按短语匹配）——偏差记录：schema 升到 v3，`_rebuild_fts()` 重建旧 `atom_fts`；旧行存的是原文，「钥匙」在旧索引里永远匹配不上 + - [x] p0-3.3 空命中是 `SearchResult{status:"no_match"}`，HTTP 与 CLI 同形，不再是裸空数组 + - [x] p0-3.4 七个过滤器（holder/scenario/source/modality/since/until/min_trust）+ base64 游标；`since > until` 报 "search since must not be after until"，游标坏了报 "search cursor is unreadable" + - [x] p0-3.5 RRF（`RRF_K = 60`）取代 `0.55/0.35/0.10`；顺手删掉无人调用的 `tokenize()` / `lexical_score()` / `_graph_score()`（graph 项恒为 0，因为证据边从不参与排序） + - [x] p0-3.6 Go 的 `search_atoms` / `compile_context` 改走 `POST /api/search`，删掉 Go 侧第二套分词器（`store.SearchAtoms` / `ftsQuery` / `tokenPattern`）——两面不可能再漂开，因为只剩一个排序实现 + - [ ] p0-3.7 100k atom 下 p50 <100ms：`python -m bench.scale_check` → `bench/reports/SCALE.md`。基线是 p50 20.16s / 165,973 atoms（`bench/reports/BASELINE.md`）。**未测不写数**：机器负载 3.7–17.9 期间测出来的数不算数 + - 验收：`pytest backend/tests/test_text_index.py backend/tests/test_search_api.py` 绿（18 条，真 FTS5 表 + TestClient）：「钥匙」命中「我们的钥匙放在玄关抽屉里」、`search atoms` 命中 `func SearchAtoms(ctx context.Context)`、`forg` 命中 forgotten、全角 `Ｓｅａｒｃｈ　Ａｔｏｍｓ` 命中、游标翻页无重复、supersede 后旧记忆不再应答；MCP 一致性由 `ely-eye-mcp/internal/mcpserver/search_tools_test.go`（3 条）+ `bash scripts/mcp_write_smoke.sh` PASS 证明；偏差：checklist 写的 `tests/test_retrieval_parity.py` 不存在，实际在 `backend/tests/` 与 Go 测试里 - [ ] **p0-4** 卡带变成可挂载作用域 - [ ] p0-4.1 `mount` / `unmount` / `switch` / `append` 四动词 - [ ] p0-4.2 检索按已挂载卡带过滤（作用域，不再是全局） diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 7e9d4dd..8edb8ab 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -51,6 +51,3 @@ where = ["src"] line-length = 100 target-version = "py312" -[tool.pyright] -pythonVersion = "3.12" -include = ["src"] diff --git a/backend/src/ely_eye/api.py b/backend/src/ely_eye/api.py index 50fd7d7..9c507fe 100644 --- a/backend/src/ely_eye/api.py +++ b/backend/src/ely_eye/api.py @@ -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)] diff --git a/backend/src/ely_eye/cli.py b/backend/src/ely_eye/cli.py index fdd1510..ca2aaa7 100644 --- a/backend/src/ely_eye/cli.py +++ b/backend/src/ely_eye/cli.py @@ -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, diff --git a/backend/src/ely_eye/compiler.py b/backend/src/ely_eye/compiler.py index 46de0be..6447616 100644 --- a/backend/src/ely_eye/compiler.py +++ b/backend/src/ely_eye/compiler.py @@ -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, diff --git a/backend/src/ely_eye/db.py b/backend/src/ely_eye/db.py index 91a504b..a8dbb50 100644 --- a/backend/src/ely_eye/db.py +++ b/backend/src/ely_eye/db.py @@ -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"], {}), - ) diff --git a/backend/src/ely_eye/migrations.py b/backend/src/ely_eye/migrations.py index b4fd85e..7ec8c84 100644 --- a/backend/src/ely_eye/migrations.py +++ b/backend/src/ely_eye/migrations.py @@ -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], + ) diff --git a/backend/src/ely_eye/retrieval.py b/backend/src/ely_eye/retrieval.py index 27daa02..e9c28de 100644 --- a/backend/src/ely_eye/retrieval.py +++ b/backend/src/ely_eye/retrieval.py @@ -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 diff --git a/backend/src/ely_eye/row_codec.py b/backend/src/ely_eye/row_codec.py new file mode 100644 index 0000000..9d4ac35 --- /dev/null +++ b/backend/src/ely_eye/row_codec.py @@ -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"], {}), + ) diff --git a/backend/src/ely_eye/schemas.py b/backend/src/ely_eye/schemas.py index a77a765..30f0d94 100644 --- a/backend/src/ely_eye/schemas.py +++ b/backend/src/ely_eye/schemas.py @@ -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 diff --git a/backend/src/ely_eye/text_index.py b/backend/src/ely_eye/text_index.py new file mode 100644 index 0000000..acd426d --- /dev/null +++ b/backend/src/ely_eye/text_index.py @@ -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('"', '""') + '"' diff --git a/backend/tests/test_memory_schema.py b/backend/tests/test_memory_schema.py index 1960e88..5e2aef1 100644 --- a/backend/tests/test_memory_schema.py +++ b/backend/tests/test_memory_schema.py @@ -222,3 +222,22 @@ def test_atoms_as_of_walks_the_validity_timeline(tmp_path: Path) -> None: assert {atom.atom_id for atom in database.atoms_as_of(NOW - timedelta(days=2))} == {"closed"} assert database.atoms_as_of(NOW, holder="nobody") == [] assert {atom.atom_id for atom in database.atoms_as_of(NOW, holder="wife")} == {"live"} + + +def test_a_second_version_database_rebuilds_its_index_on_upgrade(tmp_path: Path) -> None: + """A second-version store indexed raw text, so 钥匙 could never match; the upgrade fixes the index.""" + database = _database(tmp_path) + _seed_source(database) + database.upsert_atoms([_atom("key")]) + with database.connect() as connection: + connection.execute("DELETE FROM atom_fts") + connection.execute( + "INSERT INTO atom_fts(atom_id, source, text) VALUES ('key', 'notes/key.txt', '钥匙放在玄关抽屉里')" + ) + connection.execute("PRAGMA user_version = 2") + assert database.search_fts("钥匙", 5) == [], "the old index really could not answer this query" + + upgraded = _database(tmp_path) + assert [atom_id for atom_id, _ in upgraded.search_fts("钥匙", 5)] == ["key"] + with upgraded.connect() as connection: + assert int(connection.execute("PRAGMA user_version").fetchone()[0]) == SCHEMA_VERSION diff --git a/backend/tests/test_search_api.py b/backend/tests/test_search_api.py new file mode 100644 index 0000000..e3bc93b --- /dev/null +++ b/backend/tests/test_search_api.py @@ -0,0 +1,114 @@ +"""Search over HTTP: the answers a caller can act on, and the reasons it gets none.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +import ely_eye.api as api_module +from ely_eye.config import Settings +from ely_eye.db import Database +from ely_eye.memory import MemoryService +from ely_eye.schemas import MemoryOrigin + +KEYS = "我们的钥匙放在玄关抽屉里" +PASSPORT = "护照在书房第二个抽屉里" +WALLET = "钱包在玄关的挂钩上" + + +@pytest.fixture +def client(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> TestClient: + settings = Settings(home=tmp_path, default_holder="wife") + monkeypatch.setattr(api_module, "get_settings", lambda: settings) + return TestClient(api_module.create_app()) + + +def search(client: TestClient, payload: dict) -> dict: + response = client.post("/api/search", json=payload) + assert response.status_code in (200, 422), response.text + return response.json() + + +def test_a_chinese_question_finds_the_memory_it_was_written_in(client: TestClient) -> None: + client.post("/api/memory", json={"text": KEYS, "scenario": "home"}) + answer = search(client, {"query": "钥匙"}) + assert answer["status"] == "hits" + assert [hit["atom"]["text"] for hit in answer["hits"]] == [KEYS] + + +def test_an_empty_corpus_answers_no_match_rather_than_an_empty_array(client: TestClient) -> None: + answer = search(client, {"query": "钥匙"}) + assert answer["hits"] == [] + assert answer["status"] == "no_match", "an empty answer must say the corpus has nothing" + + +def test_a_query_with_nothing_searchable_is_refused_with_the_reason(client: TestClient) -> None: + response = client.post("/api/search", json={"query": "???"}) + assert response.status_code == 422 + assert "needs at least one searchable word" in response.json()["detail"] + + +def test_the_scenario_filter_hides_memories_from_other_scenarios(client: TestClient) -> None: + client.post("/api/memory", json={"text": KEYS, "scenario": "home"}) + assert search(client, {"query": "钥匙", "filters": {"scenario": "home"}})["status"] == "hits" + assert search(client, {"query": "钥匙", "filters": {"scenario": "travel"}})["status"] == "no_match" + + +def test_a_superseded_memory_stops_answering(client: TestClient) -> None: + old = client.post("/api/memory", json={"text": KEYS}).json() + client.post("/api/memory/supersede", json={"atom_id": old["atom_id"], "text": PASSPORT}) + answer = search(client, {"query": "抽屉"}) + assert [hit["atom"]["text"] for hit in answer["hits"]] == [PASSPORT] + + +def test_paging_walks_every_hit_without_skipping_one(client: TestClient) -> None: + for text in (KEYS, PASSPORT, WALLET, "备用钥匙在车里", "工具箱在阳台的抽屉里"): + client.post("/api/memory", json={"text": text}) + seen: list[str] = [] + cursor = None + for _ in range(10): + payload: dict = {"query": "钥匙", "top_k": 2} + if cursor: + payload["cursor"] = cursor + answer = search(client, payload) + seen.extend(hit["atom"]["atom_id"] for hit in answer["hits"]) + cursor = answer["next_cursor"] + if not cursor: + break + assert cursor is None, "the page loop never reached the end" + assert len(seen) == len(set(seen)), "a page repeated a memory" + assert len(seen) == 2, "only two memories mention 钥匙" + + +def test_the_holder_filter_answers_only_that_holders_memory( + client: TestClient, tmp_path: Path +) -> None: + client.post("/api/memory", json={"text": KEYS}) + husband = MemoryService(Settings(home=tmp_path, default_holder="husband"), Database(Settings(home=tmp_path))) + husband.remember("我的钥匙在办公室", holder_id="husband", origin=MemoryOrigin.self_) + answer = search(client, {"query": "钥匙", "filters": {"holder": "husband"}}) + assert [hit["atom"]["text"] for hit in answer["hits"]] == ["我的钥匙在办公室"] + + +def test_the_cli_answers_in_json_a_machine_can_read(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """The three faces rank alike: the CLI reports the same hit the HTTP face returns.""" + import json + + from typer.testing import CliRunner + + import ely_eye.cli as cli_module + from ely_eye.cli import app + + settings = Settings(home=tmp_path, default_holder="wife") + monkeypatch.setattr(cli_module, "get_settings", lambda: settings) + runner = CliRunner() + MemoryService(settings, Database(settings)).remember(KEYS, holder_id="wife", origin=MemoryOrigin.self_, scenario="home") + + answer = json.loads(runner.invoke(app, ["search", "钥匙", "--scenario", "home"]).stdout) + assert answer["status"] == "hits" + assert [hit["atom"]["text"] for hit in answer["hits"]] == [KEYS] + + empty = json.loads(runner.invoke(app, ["search", "钥匙", "--scenario", "travel"]).stdout) + assert empty["status"] == "no_match" diff --git a/backend/tests/test_text_index.py b/backend/tests/test_text_index.py new file mode 100644 index 0000000..7508c99 --- /dev/null +++ b/backend/tests/test_text_index.py @@ -0,0 +1,75 @@ +"""Checks for the FTS normalizer, run against a real FTS5 table. + +String equality alone would pass a normalizer SQLite cannot use, so every acceptance case +here inserts indexed text and asks the engine the question the user would ask. +""" + +from __future__ import annotations + +import sqlite3 + +import pytest + +from ely_eye.text_index import fts_query, index_text + + +def _fts(rows: list[str]) -> sqlite3.Connection: + connection = sqlite3.connect(":memory:") + connection.execute("CREATE VIRTUAL TABLE fts USING fts5(original UNINDEXED, text, tokenize = 'unicode61')") + connection.executemany("INSERT INTO fts(original, text) VALUES (?, ?)", [(row, index_text(row)) for row in rows]) + return connection + + +def _hits(connection: sqlite3.Connection, query: str) -> list[str]: + return [row[0] for row in connection.execute("SELECT original FROM fts WHERE fts MATCH ?", (fts_query(query),))] + + +def test_index_text_splits_a_chinese_sentence_into_characters() -> None: + assert index_text("我们的钥匙放在玄关抽屉里") == "我 们 的 钥 匙 放 在 玄 关 抽 屉 里" + + +def test_index_text_splits_identifiers_a_user_cannot_guess_the_casing_of() -> None: + assert index_text("func SearchAtoms") == "func search atoms" + assert index_text("def search_atoms(self) -> None:") == "def search atoms self none" + assert index_text("HTTPServer2 listens on 8765") == "http server 2 listens on 8765" + + +def test_index_text_folds_full_width_letters_that_search_would_miss() -> None: + assert index_text("Ｓｅａｒｃｈ　Ａｔｏｍｓ") == "search atoms" + + +def test_fts_query_phrases_cjk_and_prefixes_latin() -> None: + assert fts_query("钥匙") == '"钥 匙"' + assert fts_query("search atoms") == '"search"* "atoms"*' + + +def test_fts_query_refuses_a_query_with_nothing_to_match() -> None: + for empty in ("", " ", "?!。"): + with pytest.raises(ValueError, match="needs at least one searchable word"): + fts_query(empty) + + +def test_a_chinese_question_finds_the_chinese_memory() -> None: + connection = _fts(["我们的钥匙放在玄关抽屉里", "The keys are in the kitchen drawer"]) + assert _hits(connection, "钥匙") == ["我们的钥匙放在玄关抽屉里"] + + +def test_a_cjk_phrase_will_not_match_its_characters_scattered_apart() -> None: + connection = _fts(["钥匙在抽屉", "匙 在 钥"]) + assert _hits(connection, "钥匙") == ["钥匙在抽屉"] + + +def test_an_english_query_finds_the_identifier_it_names() -> None: + connection = _fts(["func SearchAtoms(ctx context.Context)", "func ListCartridges()"]) + assert _hits(connection, "search atoms") == ["func SearchAtoms(ctx context.Context)"] + + +def test_a_prefix_reaches_the_word_it_grew_into() -> None: + connection = _fts(["the atom was forgotten", "unrelated text"]) + assert _hits(connection, "atom") == ["the atom was forgotten"] + assert _hits(connection, "forg") == ["the atom was forgotten"] + + +def test_a_mixed_query_needs_both_halves_in_the_same_memory() -> None: + connection = _fts(["钥匙放在 drawer 里", "钥匙在玄关", "the drawer is empty"]) + assert _hits(connection, "钥匙 drawer") == ["钥匙放在 drawer 里"] diff --git a/bench/run.py b/bench/run.py index 8b6d137..1886520 100644 --- a/bench/run.py +++ b/bench/run.py @@ -252,7 +252,8 @@ def write_markdown(report: dict[str, Any], path: Path) -> None: ), f"- models: dense `{report['meta']['models']['dense']}`, rerank `{report['meta']['models']['rerank']}`", ( - f"- machine: {report['meta']['machine']['cpus']} cores, load average " + f"- machine: {report['meta']['machine']['cpus']} cores, " + f"{report['meta']['machine']['threads_per_query']} threads per query, load average " f"{report['meta']['machine']['load_average_1_5_15'][0]:.1f} when the report was written" ), "", diff --git a/bench/scale_check.py b/bench/scale_check.py new file mode 100644 index 0000000..e76b211 --- /dev/null +++ b/bench/scale_check.py @@ -0,0 +1,124 @@ +"""Time the product's own search at the size the acceptance contract names. + +The baseline recorded RetrievalService.search at p50 20.16s over 165,973 atoms because it +pulled every atom into Python and rebuilt a BM25 index per query. This check answers one +question with a number: what does the same call cost now, on the same corpus, on an idle +machine, with the same thread budget the benchmark pins? +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path +from typing import Any + +from bench.run import latency, load_queries, scale_probes +from bench.systems import machine_facts + +DEFAULT_HOME = Path("bench/runs/baseline") +MIN_ATOMS = 100_000 +DEFAULT_PROBES = 300 + + +def measure(home: Path, probes: int) -> dict[str, Any]: + """Open the product store (migrating it in place if it is old) and time real searches.""" + from ely_eye.config import Settings + from ely_eye.db import Database + from ely_eye.retrieval import RetrievalService + + if not (home / "data" / "ely_eye.sqlite").is_file(): + raise SystemExit( + f"{home / 'data' / 'ely_eye.sqlite'}: no product database, " + "build it with python -m bench.ingest_baseline" + ) + settings = Settings(home=home) + started = time.perf_counter() + database = Database(settings) + migration_seconds = time.perf_counter() - started + atoms = database.atom_count() + if atoms < MIN_ATOMS: + raise SystemExit(f"{database.path}: {atoms} atoms, this check needs at least {MIN_ATOMS}") + + queries = scale_probes(load_queries(), probes) + service = RetrievalService(settings, database) + service.search(queries[0].text) # warm the index and any model before the first sample + times: list[float] = [] + for query in queries: + began = time.perf_counter() + service.search(query.text) + times.append(time.perf_counter() - began) + + return { + "database": str(database.path), + "atoms": atoms, + "migration_seconds": migration_seconds, + "dense_embeddings_enabled": settings.dense_embeddings_enabled, + "probes": len(queries), + "probe_modalities": sorted({query.modality for query in queries}), + "latency": latency(times), + "machine": machine_facts(), + } + + +def write_markdown(report: dict[str, Any], path: Path) -> None: + timing = report["latency"] + machine = report["machine"] + path.write_text( + "\n".join( + [ + "# SCALE: the product's own search at corpus size", + "", + f"- database: {report['database']}", + ( + f"- atoms: {report['atoms']}, probes: {report['probes']} " + f"({', '.join(report['probe_modalities'])})" + ), + f"- dense embeddings: {report['dense_embeddings_enabled']}", + f"- opening the store (migrations included): {report['migration_seconds']:.3f}s", + ( + f"- p50 {timing['p50']:.4f}s · p95 {timing['p95']:.4f}s · max {timing['max']:.4f}s " + f"over {timing['queries']} queries" + ), + ( + f"- machine: {machine['system']} {machine['release']}, {machine['cpus']} cores, " + f"{machine['threads_per_query']} threads per query, " + f"load {machine['load_average_1_5_15']}" + ), + "", + ( + "_Latency is wall time per search call on an idle machine with the thread budget " + "the benchmark pins; a run on a loaded machine is not comparable._" + ), + "", + ] + ), + encoding="utf-8", + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--home", type=Path, default=DEFAULT_HOME) + parser.add_argument("--probes", type=int, default=DEFAULT_PROBES) + parser.add_argument("--out", type=Path, default=Path("bench/reports/SCALE.md")) + parser.add_argument("--json-out", type=Path, default=Path("bench/reports/scale.json")) + args = parser.parse_args() + if args.probes <= 0: + raise SystemExit(f"--probes must be positive, got {args.probes}") + + report = measure(args.home, args.probes) + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8") + write_markdown(report, args.out) + timing = report["latency"] + print( + f"{report['atoms']} atoms, {report['probes']} probes: " + f"p50 {timing['p50']:.4f}s p95 {timing['p95']:.4f}s", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/bench/systems.py b/bench/systems.py index 2d7e0c1..3b7928e 100644 --- a/bench/systems.py +++ b/bench/systems.py @@ -30,6 +30,16 @@ INDEX_DIR = RUNS / "index" MODELS = ROOT / "models" DENSE_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2" RERANK_MODEL = "BAAI/bge-reranker-base" +THREADS_PER_QUERY = 4 + + +def _pin_threads() -> None: + """Fix the CPU budget per query before torch loads; a run that grabs all 32 cores measures the scheduler.""" + for variable in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS"): + os.environ.setdefault(variable, str(THREADS_PER_QUERY)) + + +_pin_threads() def machine_facts() -> dict[str, Any]: @@ -38,6 +48,7 @@ def machine_facts() -> dict[str, Any]: "system": platform.system(), "release": platform.release(), "cpus": os.cpu_count(), + "threads_per_query": int(os.environ.get("OMP_NUM_THREADS", THREADS_PER_QUERY)), "load_average_1_5_15": [round(value, 2) for value in os.getloadavg()], } @@ -432,7 +443,7 @@ def live_product_rankings( timings: list[float] = [] for query in queries: started = time.perf_counter() - hits = service.search(query, top_k=top_k) + hits = service.search(query, top_k=top_k).hits timings.append(time.perf_counter() - started) # The product ranks atoms; collapse to sources by first (best) hit so the arms stay comparable. sources: list[str] = [] diff --git a/bench/tests/test_joint.py b/bench/tests/test_joint.py index fab604a..a674b4d 100644 --- a/bench/tests/test_joint.py +++ b/bench/tests/test_joint.py @@ -160,7 +160,13 @@ def _report() -> dict: "corpus_sources": 10, "queries": 16, "answerable": 8, - "machine": {"system": "Linux", "release": "x", "cpus": 32, "load_average_1_5_15": [1.0, 2.0, 3.0]}, + "machine": { + "system": "Linux", + "release": "x", + "cpus": 32, + "threads_per_query": 4, + "load_average_1_5_15": [1.0, 2.0, 3.0], + }, "models": {"dense": "d", "rerank": "r"}, }, "chance": {"recall_at_1": 0.0, "recall_at_5": 0.0, "recall_at_10": 0.0}, diff --git a/bench/tests/test_systems.py b/bench/tests/test_systems.py index 5a7d1a6..9ef5a70 100644 --- a/bench/tests/test_systems.py +++ b/bench/tests/test_systems.py @@ -187,3 +187,18 @@ def test_the_dense_model_is_loaded_once_per_process(tmp_path: Path, monkeypatch: systems._sentence_model(tmp_path / "missing") finally: systems._MODELS.clear() + + +def test_the_cpu_budget_is_pinned_and_reported() -> None: + """Latency measured with every core available measures the scheduler, so the budget is fixed and printed.""" + import os + + from bench import systems + + budget = systems.machine_facts()["threads_per_query"] + pinned = {os.environ[name] for name in _THREAD_VARIABLES} + assert pinned == {str(budget)} + assert 1 <= budget <= 8 + + +_THREAD_VARIABLES = ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS") diff --git a/ely-eye-mcp/internal/apiclient/client.go b/ely-eye-mcp/internal/apiclient/client.go index 4910571..e90b591 100644 --- a/ely-eye-mcp/internal/apiclient/client.go +++ b/ely-eye-mcp/internal/apiclient/client.go @@ -13,6 +13,8 @@ import ( "net/http" "strings" "time" + + "github.com/ZacharyZhang-NY/ely-eye/ely-eye-mcp/internal/domain" ) const maxErrorBody = 4096 @@ -82,3 +84,51 @@ func serverReason(status int, data []byte) string { } return trimmed } + +// SearchOptions is one ranked query. Empty optional fields are omitted, so the +// server's own defaults decide rather than a zero value invented here. +type SearchOptions struct { + Query string + Modality string + Holder string + Scenario string + Limit int +} + +// SearchResult carries the server's own verdict: "no_match" is a claim about the +// corpus, an empty hit list is not. +type SearchResult struct { + Status string `json:"status"` + Hits []domain.AtomHit `json:"hits"` + NextCursor string `json:"next_cursor"` +} + +// Search asks the running server to rank. The ranking lives in one place, so an +// MCP hit cannot disagree with an HTTP or CLI hit for the same query. +func (c *Client) Search(ctx context.Context, options SearchOptions) (*SearchResult, error) { + if strings.TrimSpace(options.Query) == "" { + return nil, errors.New("query is required") + } + body := map[string]any{"query": options.Query} + if options.Limit > 0 { + body["top_k"] = options.Limit + } + filters := map[string]string{} + for key, value := range map[string]string{ + "modality": options.Modality, + "holder": options.Holder, + "scenario": options.Scenario, + } { + if value != "" { + filters[key] = value + } + } + if len(filters) > 0 { + body["filters"] = filters + } + var result SearchResult + if err := c.Post(ctx, "/api/search", body, &result); err != nil { + return nil, err + } + return &result, nil +} diff --git a/ely-eye-mcp/internal/domain/types.go b/ely-eye-mcp/internal/domain/types.go index 77b62f0..f201187 100644 --- a/ely-eye-mcp/internal/domain/types.go +++ b/ely-eye-mcp/internal/domain/types.go @@ -93,6 +93,7 @@ type Atom struct { type AtomHit struct { Atom Atom `json:"atom"` SparseScore float64 `json:"sparse_score"` + DenseScore float64 `json:"dense_score"` FinalScore float64 `json:"final_score"` } diff --git a/ely-eye-mcp/internal/mcpserver/search_tools_test.go b/ely-eye-mcp/internal/mcpserver/search_tools_test.go new file mode 100644 index 0000000..771b021 --- /dev/null +++ b/ely-eye-mcp/internal/mcpserver/search_tools_test.go @@ -0,0 +1,109 @@ +package mcpserver + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/ZacharyZhang-NY/ely-eye/ely-eye-mcp/internal/apiclient" +) + +func searchTestServer(t *testing.T, recorded *map[string]any) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/search" { + t.Errorf("path = %q, want /api/search", r.URL.Path) + } + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Errorf("request body does not decode: %v", err) + } + *recorded = body + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "query": "钥匙", + "status": "hits", + "hits": [{"atom": {"atom_id": "mem-1:text", "source_id": "mem-1", "modality": "text", + "source": "memory://wife/home", "ingested_at": "2026-05-01T12:00:00+00:00", + "holder_id": "wife", "scenario": "home", "text": "钥匙在玄关抽屉"}, + "sparse_score": 4.2, "dense_score": 0.0, "final_score": 0.016}], + "next_cursor": null + }`)) + })) +} + +func TestSearchAsksTheServerToRank(t *testing.T) { + var recorded map[string]any + server := searchTestServer(t, &recorded) + defer server.Close() + + s := &Server{api: apiclient.New(server.URL)} + _, out, err := s.searchAtomsTool(context.Background(), nil, SearchAtomsInput{ + Query: "钥匙", + Modality: "text", + Holder: "wife", + Scenario: "home", + Limit: 5, + }) + if err != nil { + t.Fatalf("searchAtomsTool: %v", err) + } + if len(out.Hits) != 1 || out.Hits[0].Atom.AtomID != "mem-1:text" { + t.Fatalf("hits = %+v, want the server's single hit", out.Hits) + } + if recorded["query"] != "钥匙" { + t.Errorf("query sent = %v, want 钥匙", recorded["query"]) + } + if recorded["top_k"] != float64(5) { + t.Errorf("top_k sent = %v, want 5", recorded["top_k"]) + } + filters, ok := recorded["filters"].(map[string]any) + if !ok { + t.Fatalf("filters sent = %v, want an object", recorded["filters"]) + } + for key, want := range map[string]string{"modality": "text", "holder": "wife", "scenario": "home"} { + if filters[key] != want { + t.Errorf("filters[%s] = %v, want %q", key, filters[key], want) + } + } +} + +func TestSearchRefusesAnEmptyQueryBeforeTheNetwork(t *testing.T) { + var recorded map[string]any + server := searchTestServer(t, &recorded) + defer server.Close() + + s := &Server{api: apiclient.New(server.URL)} + _, _, err := s.searchAtomsTool(context.Background(), nil, SearchAtomsInput{Query: " "}) + if err == nil { + t.Fatal("an empty query reached the server; the tool should refuse it") + } + if recorded != nil { + t.Errorf("the server was called for an empty query: %v", recorded) + } +} + +func TestCompileContextPacksWhatTheServerReturned(t *testing.T) { + var recorded map[string]any + server := searchTestServer(t, &recorded) + defer server.Close() + + s := &Server{api: apiclient.New(server.URL)} + _, out, err := s.compileContextTool(context.Background(), nil, CompileContextInput{ + Question: "钥匙在哪", + Profile: "library_100m", + Limit: 5, + }) + if err != nil { + t.Fatalf("compileContextTool: %v", err) + } + if out.Context.CacheTraceID == "" { + t.Error("the packed context has no cache trace id") + } + if !strings.Contains(out.Context.PackedText, "钥匙在玄关抽屉") { + t.Errorf("packed text does not carry the evidence: %q", out.Context.PackedText) + } +} diff --git a/ely-eye-mcp/internal/mcpserver/tools.go b/ely-eye-mcp/internal/mcpserver/tools.go index 067e0b4..9fe0fdb 100644 --- a/ely-eye-mcp/internal/mcpserver/tools.go +++ b/ely-eye-mcp/internal/mcpserver/tools.go @@ -4,6 +4,7 @@ import ( "context" "errors" + "github.com/ZacharyZhang-NY/ely-eye/ely-eye-mcp/internal/apiclient" "github.com/ZacharyZhang-NY/ely-eye/ely-eye-mcp/internal/contextpack" "github.com/ZacharyZhang-NY/ely-eye/ely-eye-mcp/internal/domain" "github.com/modelcontextprotocol/go-sdk/mcp" @@ -21,8 +22,10 @@ type GetCartridgeInput struct { } type SearchAtomsInput struct { - Query string `json:"query" jsonschema:"search query for the Ely-Eye FTS atom index"` + Query string `json:"query" jsonschema:"search query, matched by the server's own ranking"` Modality string `json:"modality,omitempty" jsonschema:"optional exact modality filter such as text, code, image, pdf_page, pdf_region, video_frame, ui_screenshot"` + Holder string `json:"holder,omitempty" jsonschema:"optional: only this holder's memories may answer"` + Scenario string `json:"scenario,omitempty" jsonschema:"optional: only this scenario, e.g. home or work"` Limit int `json:"limit,omitempty" jsonschema:"maximum number of atoms to return, capped at 100"` } @@ -160,11 +163,17 @@ func (s *Server) getCartridgeTool(ctx context.Context, _ *mcp.CallToolRequest, i } func (s *Server) searchAtomsTool(ctx context.Context, _ *mcp.CallToolRequest, input SearchAtomsInput) (*mcp.CallToolResult, SearchAtomsOutput, error) { - hits, err := s.store.SearchAtoms(ctx, input.Query, input.Modality, input.Limit) + result, err := s.api.Search(ctx, apiclient.SearchOptions{ + Query: input.Query, + Modality: input.Modality, + Holder: input.Holder, + Scenario: input.Scenario, + Limit: input.Limit, + }) if err != nil { return nil, SearchAtomsOutput{}, err } - return nil, SearchAtomsOutput{Hits: contextpack.SummarizeHits(hits, 1200)}, nil + return nil, SearchAtomsOutput{Hits: contextpack.SummarizeHits(result.Hits, 1200)}, nil } func (s *Server) fetchAtomTool(ctx context.Context, _ *mcp.CallToolRequest, input FetchAtomInput) (*mcp.CallToolResult, FetchAtomOutput, error) { @@ -179,11 +188,11 @@ func (s *Server) compileContextTool(ctx context.Context, _ *mcp.CallToolRequest, if input.Question == "" { return nil, CompileContextOutput{}, errors.New("question is required") } - hits, err := s.store.SearchAtoms(ctx, input.Question, "", input.Limit) + hits, err := s.api.Search(ctx, apiclient.SearchOptions{Query: input.Question, Limit: input.Limit}) if err != nil { return nil, CompileContextOutput{}, err } - pack := contextpack.Build(input.Question, input.Profile, hits, input.TokenBudget) + pack := contextpack.Build(input.Question, input.Profile, hits.Hits, input.TokenBudget) return nil, CompileContextOutput{Context: pack}, nil } diff --git a/ely-eye-mcp/internal/store/atoms.go b/ely-eye-mcp/internal/store/atoms.go index 30eb7e9..8170ffd 100644 --- a/ely-eye-mcp/internal/store/atoms.go +++ b/ely-eye-mcp/internal/store/atoms.go @@ -4,15 +4,11 @@ import ( "context" "database/sql" "errors" - "math" - "regexp" "strings" "github.com/ZacharyZhang-NY/ely-eye/ely-eye-mcp/internal/domain" ) -var tokenPattern = regexp.MustCompile(`[\p{L}\p{N}_]+`) - func (s *Store) FetchAtom(ctx context.Context, atomID string) (domain.Atom, error) { if strings.TrimSpace(atomID) == "" { return domain.Atom{}, errors.New("atom_id is required") @@ -34,65 +30,6 @@ func (s *Store) FetchAtom(ctx context.Context, atomID string) (domain.Atom, erro return atom, nil } -func (s *Store) SearchAtoms(ctx context.Context, query string, modality string, limit int) ([]domain.AtomHit, error) { - matchQuery, err := ftsQuery(query) - if err != nil { - return nil, err - } - if limit <= 0 || limit > 100 { - limit = 24 - } - args := []any{matchQuery} - filter := "" - if modality != "" { - filter = "AND a.modality = ?" - args = append(args, modality) - } - args = append(args, limit) - rows, err := s.db.QueryContext(ctx, ` - SELECT a.atom_id, a.source_id, a.modality, a.source, a.ingested_at, - COALESCE(a.holder_id, ''), COALESCE(a.origin, ''), COALESCE(a.scenario, ''), - COALESCE(a.event_time, ''), a.text, - COALESCE(a.image_ref, ''), COALESCE(a.layout_json, ''), - a.relations_json, a.trust_json, a.token_equivalent, a.metadata_json, - bm25(atom_fts) AS score - FROM atom_fts - JOIN v_atoms_current a ON a.atom_id = atom_fts.atom_id - WHERE atom_fts MATCH ? `+filter+` - ORDER BY score - LIMIT ?`, args...) - if err != nil { - return nil, err - } - defer rows.Close() - var hits []domain.AtomHit - var bestScore float64 - for rows.Next() { - atom, score, err := scanAtomHit(rows) - if err != nil { - return nil, err - } - sparseScore := -score - if math.IsNaN(sparseScore) || math.IsInf(sparseScore, 0) { - sparseScore = 0 - } - if sparseScore > bestScore { - bestScore = sparseScore - } - hits = append(hits, domain.AtomHit{Atom: atom, SparseScore: sparseScore, FinalScore: sparseScore}) - } - if err := rows.Err(); err != nil { - return nil, err - } - if bestScore > 0 { - for index := range hits { - hits[index].SparseScore = hits[index].SparseScore / bestScore - hits[index].FinalScore = hits[index].SparseScore - } - } - return hits, nil -} - type atomScanner interface { Scan(dest ...any) error } @@ -142,71 +79,3 @@ func scanAtom(scanner atomScanner) (domain.Atom, error) { } return atom, nil } - -func scanAtomHit(scanner atomScanner) (domain.Atom, float64, error) { - var atom domain.Atom - var layout string - var relations string - var trust string - var metadata string - var score float64 - if err := scanner.Scan( - &atom.AtomID, - &atom.SourceID, - &atom.Modality, - &atom.Source, - &atom.IngestedAt, - &atom.HolderID, - &atom.Origin, - &atom.Scenario, - &atom.EventTime, - &atom.Text, - &atom.ImageRef, - &layout, - &relations, - &trust, - &atom.TokenEquivalent, - &metadata, - &score, - ); err != nil { - return domain.Atom{}, 0, err - } - var err error - atom.Layout, err = optionalJSONValue(layout) - if err != nil { - return domain.Atom{}, 0, err - } - atom.Relations, err = requiredJSONValue(relations) - if err != nil { - return domain.Atom{}, 0, err - } - atom.Trust, err = requiredJSONValue(trust) - if err != nil { - return domain.Atom{}, 0, err - } - atom.Metadata, err = requiredJSONValue(metadata) - if err != nil { - return domain.Atom{}, 0, err - } - return atom, score, nil -} - -func ftsQuery(query string) (string, error) { - tokens := tokenPattern.FindAllString(strings.ToLower(query), -1) - if len(tokens) == 0 { - return "", errors.New("query must contain searchable text") - } - unique := make([]string, 0, len(tokens)) - seen := map[string]bool{} - for _, token := range tokens { - if seen[token] { - continue - } - seen[token] = true - unique = append(unique, `"`+strings.ReplaceAll(token, `"`, `""`)+`"`) - if len(unique) == 16 { - break - } - } - return strings.Join(unique, " OR "), nil -} diff --git a/package-lock.json b/package-lock.json new file mode 100644 index 0000000..4566531 --- /dev/null +++ b/package-lock.json @@ -0,0 +1,45 @@ +{ + "name": "ely-eye-tools", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "ely-eye-tools", + "devDependencies": { + "pyright": "1.1.404" + } + }, + "node_modules/fsevents": { + "version": "2.3.3", + "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz", + "integrity": "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^8.16.0 || ^10.6.0 || >=11.0.0" + } + }, + "node_modules/pyright": { + "version": "1.1.404", + "resolved": "https://registry.npmjs.org/pyright/-/pyright-1.1.404.tgz", + "integrity": "sha512-Gr48RLeyis8uMjx04vXGIdj3OFe/WGGBMg53XIh9GH30KRyuyKYNAjfUZqj3r51Xrv5BFKyDyguJu/BV3hi7QA==", + "dev": true, + "license": "MIT", + "bin": { + "pyright": "index.js", + "pyright-langserver": "langserver.index.js" + }, + "engines": { + "node": ">=14.0.0" + }, + "optionalDependencies": { + "fsevents": "~2.3.3" + } + } + } +} diff --git a/package.json b/package.json new file mode 100644 index 0000000..a797abe --- /dev/null +++ b/package.json @@ -0,0 +1,8 @@ +{ + "name": "ely-eye-tools", + "private": true, + "description": "Dev-only type and lint gates for the Python and TypeScript surfaces.", + "devDependencies": { + "pyright": "1.1.404" + } +} diff --git a/pyproject.toml b/pyproject.toml index 7a12311..cd6ea25 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,3 @@ -[tool.pyright] -venvPath = "." -venv = ".venv" -extraPaths = [".", "backend/src"] -pythonVersion = "3.12" -include = ["bench", "backend/src"] - [tool.mypy] python_executable = ".venv/bin/python" mypy_path = "." diff --git a/pyrightconfig.json b/pyrightconfig.json index 94a9bfb..2e8e8bf 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -2,6 +2,7 @@ "venvPath": ".", "venv": ".venv", "extraPaths": [".", "backend/src"], - "include": ["bench", "backend/src"], + "include": ["bench", "backend/src", "backend/tests"], + "exclude": ["bench/corpus", "bench/models", "bench/runs", "node_modules", ".venv"], "pythonVersion": "3.12" } diff --git a/scripts/mcp_write_smoke.sh b/scripts/mcp_write_smoke.sh index d0b41fe..41f679d 100755 --- a/scripts/mcp_write_smoke.sh +++ b/scripts/mcp_write_smoke.sh @@ -98,6 +98,13 @@ if not hits or hits[0]["atom"]["atom_id"] != atom_id: if hits[0]["atom"].get("scenario") != "home": raise SystemExit(f"FAIL: the search hit lost the scenario: {hits[0]['atom']}") +chinese = tool("remember", {"text": "我们的钥匙放在玄关抽屉里", "scenario": "home"})["atom"] +short = tool("search_atoms", {"query": "钥匙"})["hits"] +if not any(hit["atom"]["atom_id"] == chinese["atom_id"] for hit in short): + raise SystemExit(f"FAIL: a two-character Chinese query missed the memory written in Chinese: {short}") +if tool("search_atoms", {"query": "keys", "scenario": "travel"})["hits"]: + raise SystemExit("FAIL: the scenario filter let another scenario answer") + replacement = "keys are on the coffee table" newer = tool("supersede_memory", {"atom_id": atom_id, "text": replacement})["atom"] if newer["supersedes"] != [atom_id]: