diff --git a/core/wren/src/wren/memory/cli.py b/core/wren/src/wren/memory/cli.py index 2a76c45ed3..962ba8b9ed 100644 --- a/core/wren/src/wren/memory/cli.py +++ b/core/wren/src/wren/memory/cli.py @@ -241,12 +241,16 @@ def index( project_path = discover_project_path(explicit=None) md_pairs = load_query_pairs(project_path) - if md_pairs: - # upsert → re-running index converges on the markdown content. - res = mem_store.load_queries(md_pairs, upsert=True) + res = mem_store.sync_markdown_queries(md_pairs) + if res["loaded"] or res["updated"] or res["forgotten"]: typer.echo( f"Indexed {res['loaded'] + res['updated']} pair(s) from " - f"knowledge/sql/.", + f"knowledge/sql/." + + ( + f" Forgot {res['forgotten']} stale pair(s)." + if res["forgotten"] + else "" + ), err=True, ) @@ -255,7 +259,13 @@ def index( raw = queries_file.read_text(encoding="utf-8") doc = yaml.safe_load(raw) if doc and isinstance(doc, dict) and doc.get("pairs"): - load_result = mem_store.load_queries(doc["pairs"], upsert=False) + # Tag legacy=queries.yml explicitly so these rows are + # distinguishable from markdown-backed ones in + # `memory list/dump/forget --source`. They are already + # safe from sync_markdown_queries, whose forget pass is + # scoped to rows carrying its own provenance tag. + legacy_pairs = [{**p, "source": "legacy"} for p in doc["pairs"]] + load_result = mem_store.load_queries(legacy_pairs, upsert=False) loaded = load_result["loaded"] skipped = load_result["skipped"] if loaded: @@ -532,21 +542,28 @@ def check( ) return + from wren.memory.store import _MARKDOWN_SYNC_TAG, _has_tag # noqa: PLC0415 + md_nls = {p["nl"] for p in load_query_pairs(project)} mem_store = idx.store indexed, _ = mem_store.list_queries(limit=1_000_000) indexed_nls = {r.get("nl_query") for r in indexed} - # Only user-sourced pairs come from markdown; seeds/views are derived from - # the manifest and are not expected to have a knowledge/sql/ file. - indexed_user = { + # Markdown-backed means "written by sync_markdown_queries", which is what + # the provenance tag records: the same predicate that method's own forget + # pass uses. Deriving it from `source` instead would put this report out of + # step with the write path in both directions: a `source:legacy`/`view` + # pair exported into knowledge/sql/ would read as permanently "not + # indexed", and a `wren memory load` pair would read as permanently + # "stale" even though no `index` run is able to clear it. + indexed_md = { r.get("nl_query") for r in indexed - if _parse_source(r.get("tags")) not in ("seed", "view") + if _has_tag(r.get("tags"), _MARKDOWN_SYNC_TAG) } - # Compare user pairs only — seed/view rows aren't markdown-backed. - missing = md_nls - indexed_user # in markdown but not indexed as a user pair - stale = indexed_user - md_nls # user-indexed but no longer in markdown + # markdown vs. the rows the sync derived from it + missing = md_nls - indexed_md # in markdown but not indexed as a synced pair + stale = indexed_md - md_nls # synced but no longer in markdown typer.echo( f"knowledge/sql: {len(md_nls)} pair(s); index: {len(indexed_nls)} pair(s)" @@ -558,7 +575,7 @@ def check( typer.echo(f" {len(missing)} not indexed — run `wren memory index`.") if stale: typer.echo( - f" {len(stale)} user pair(s) indexed without markdown — " + f" {len(stale)} indexed pair(s) without markdown, " "stale index, run `wren memory index`." ) @@ -647,13 +664,12 @@ def _reindex() -> None: from wren.memory.markdown import load_query_pairs # noqa: PLC0415 md_pairs = load_query_pairs(project_path) - loaded = 0 - if md_pairs: - res = mem_store.load_queries(md_pairs, upsert=True) - loaded = res["loaded"] + res["updated"] + res = mem_store.sync_markdown_queries(md_pairs) + loaded = res["loaded"] + res["updated"] typer.echo( f"Reindexed {result['schema_items']} schema item(s)" + (f", {loaded} pair(s)" if loaded else "") + + (f", forgot {res['forgotten']} stale pair(s)" if res["forgotten"] else "") + "." ) @@ -881,11 +897,14 @@ def forget( def _parse_source(tags: str | None) -> str: - """Extract source value from a (possibly null/empty) tags string.""" - for part in (tags or "").split(): - if part.startswith("source:"): - return part[len("source:") :] - return "user" + """Source value for display/export, defaulting an untagged row to "user". + + Distinct from ``store._tag_source``, which reports ``None`` for a row with + no ``source:`` token so that ``--source`` filters do not match it. + """ + from wren.memory.store import _tag_source # noqa: PLC0415 + + return _tag_source(tags) or "user" def _pairs_to_yaml(rows: list[dict]) -> str: diff --git a/core/wren/src/wren/memory/index_backend.py b/core/wren/src/wren/memory/index_backend.py index 3dd094ec13..c95d0a87f6 100644 --- a/core/wren/src/wren/memory/index_backend.py +++ b/core/wren/src/wren/memory/index_backend.py @@ -123,9 +123,7 @@ def store(self): def rebuild(self) -> dict: pairs = load_query_pairs(self._project) - if not pairs: - return {"backend": self.name, "loaded": 0, "updated": 0} - res = self._store.load_queries(pairs, upsert=True) + res = self._store.sync_markdown_queries(pairs) return {"backend": self.name, **res} def reset(self) -> None: diff --git a/core/wren/src/wren/memory/store.py b/core/wren/src/wren/memory/store.py index ba38523413..04c14b7fee 100644 --- a/core/wren/src/wren/memory/store.py +++ b/core/wren/src/wren/memory/store.py @@ -26,12 +26,43 @@ _SCHEMA_TABLE = "schema_items" _QUERY_TABLE = "query_history" +# Extra token appended (space-separated, after the `source:` token) to a +# query_history row's tags when it is written by sync_markdown_queries's own +# upsert. `source:user` is not proof a row is markdown-backed: `wren memory +# load` defaults every pair to it, and so does a pre-`source:legacy` import +# of an already-consumed queries.yml, both with no knowledge/sql/*.md file +# behind them (see sync_markdown_queries). Scoping the forget step to rows +# carrying this marker, instead of inferring provenance from `source`, means +# the sync only ever deletes rows it wrote itself. +_MARKDOWN_SYNC_TAG = "origin:markdown-sync" + def _esc(value: str) -> str: """Escape single quotes for LanceDB where-clause literals.""" return value.replace("'", "''") +def _tag_source(tags: str | None) -> str | None: + """Extract the explicit ``source:`` value from a tags string, else ``None``. + + ``None`` means the row carries no ``source:`` token at all, which is what + ``store_query`` writes for a `wren memory store` / MCP ``store_query`` pair, + whose ``tags`` hold the caller's own free-form labels (``"revenue,finance"``) + or nothing. Those rows must not be swept up by a ``--source user`` filter, + least of all by ``forget``, so this stays distinct from the ``"user"`` + *display* default applied in ``cli._parse_source``. + """ + for part in (tags or "").split(): + if part.startswith("source:"): + return part[len("source:") :] + return None + + +def _has_tag(tags: str | None, tag: str) -> bool: + """Whether *tag* is one of the whitespace-separated tokens in *tags*.""" + return tag in (tags or "").split() + + def _schema_items_arrow_schema(dim: int = _DEFAULT_DIM) -> pa.Schema: return pa.schema( [ @@ -427,7 +458,7 @@ def list_queries( # Ensure a clean 0-based index matching the unfiltered table. df = df.reset_index(drop=True) if source: - df = df[df["tags"] == f"source:{source}"] + df = df[df["tags"].map(_tag_source) == source] total = len(df) df = df.sort_values("created_at", ascending=False) rows = df.iloc[offset : offset + limit] @@ -444,7 +475,7 @@ def count_queries_by_source(self, source: str) -> int: return 0 table = self._db.open_table(_QUERY_TABLE) df = table.to_pandas() - return int((df["tags"] == f"source:{source}").sum()) + return int((df["tags"].map(_tag_source) == source).sum()) def forget_queries_by_ids(self, row_ids: list[int]) -> int: """Delete rows at the given positional indices. Returns deleted count.""" @@ -471,14 +502,20 @@ def forget_queries_by_ids(self, row_ids: list[int]) -> int: return len(to_delete) def forget_queries_by_source(self, source: str) -> int: - """Delete all query_history rows matching *source* tag. Returns deleted count.""" + """Delete all query_history rows matching *source* tag. Returns deleted count. + + Goes through :meth:`forget_queries_by_ids` (parsing each row's + ``source:`` token via :func:`_tag_source`) rather than an equality + match on the raw ``tags`` string, so a row carrying an extra token + after its ``source:`` tag (e.g. ``_MARKDOWN_SYNC_TAG``) still matches + on its actual source, exactly as an untagged row would. + """ if _QUERY_TABLE not in _table_names(self._db): return 0 table = self._db.open_table(_QUERY_TABLE) - where = f"tags = 'source:{_esc(source)}'" - before = table.count_rows() - table.delete(where) - return before - table.count_rows() + df = table.to_pandas().reset_index(drop=True) + ids = [i for i, t in enumerate(df["tags"]) if _tag_source(t) == source] + return self.forget_queries_by_ids(ids) if ids else 0 # ── Dump / Load ────────────────────────────────────────────────────── @@ -493,7 +530,7 @@ def dump_queries( table = self._db.open_table(_QUERY_TABLE) df = table.to_pandas() if source: - df = df[df["tags"] == f"source:{source}"] + df = df[df["tags"].map(_tag_source) == source] df = df.sort_values("created_at", ascending=True) return df.drop(columns=["vector"], errors="ignore").to_dict("records") @@ -520,9 +557,19 @@ def _existing_pairs_index( return exact_set, nl_to_rowids def _prepare_query_records( - self, pairs: list[dict], *, tags: str | None = None + self, + pairs: list[dict], + *, + tags: str | None = None, + extra_tag: str | None = None, ) -> list[dict]: - """Prepare a complete query batch without changing its table.""" + """Prepare a complete query batch without changing its table. + + *extra_tag*, when given, is appended as an additional space-separated + token after the row's usual ``source:`` tag (computed or explicit) so + a later scan can recognize the row without touching what ``source:`` + it carries. See ``_MARKDOWN_SYNC_TAG``. + """ if not pairs: return [] texts = [p["nl"] for p in pairs] @@ -534,6 +581,8 @@ def _prepare_query_records( record_tags = ( tags if tags is not None else f"source:{p.get('source', 'user')}" ) + if extra_tag: + record_tags = f"{record_tags} {extra_tag}" records.append( { "text": p["nl"], @@ -568,9 +617,17 @@ def load_queries( *, overwrite: bool = False, upsert: bool = False, + mark_markdown_synced: bool = False, ) -> dict[str, int]: """Batch-import parsed YAML pairs into query_history. + ``mark_markdown_synced`` tags every written row with + ``_MARKDOWN_SYNC_TAG`` in addition to its usual ``source:`` tag, so a + later ``sync_markdown_queries`` forget pass can recognize rows it + wrote itself without inferring that from ``source`` (see + ``sync_markdown_queries``). Only meaningful together with + ``upsert=True``, the only mode ``sync_markdown_queries`` uses. + Returns ``{"loaded": N, "skipped": M, "updated": U}``. """ if overwrite: @@ -590,7 +647,8 @@ def load_queries( seen_nl[p["nl"]] = p deduped = list(seen_nl.values()) - records = self._prepare_query_records(deduped) + extra_tag = _MARKDOWN_SYNC_TAG if mark_markdown_synced else None + records = self._prepare_query_records(deduped, extra_tag=extra_tag) # Batch: collect IDs to delete, then delete once, then insert all. ids_to_delete = [] @@ -624,6 +682,41 @@ def load_queries( return {"loaded": loaded, "skipped": skipped, "updated": 0} + def sync_markdown_queries(self, pairs: list[dict]) -> dict[str, int]: + """Upsert *pairs* and forget any indexed pair absent from them. + + *pairs* must be the complete current ``knowledge/sql/*.md`` set (the + source of truth). The forget step is scoped to rows this method + itself previously wrote (tagged ``_MARKDOWN_SYNC_TAG`` by the upsert + below), never to rows merely matching some inferred "non-markdown" + exclusion list: ``source:user`` is also what ``wren memory load`` + gives a pair with no ``source`` of its own, and what a pre-upgrade + ``queries.yml`` import already carried before ``source:legacy`` + existed, so neither is safe to treat as markdown-backed. A row from + either of those paths is never forgotten here, only a row this sync + wrote on a prior run and no longer sees in *pairs*. + + One consequence of tracking provenance instead of inferring it: a + markdown row that already existed in the index under the old, + untagged scheme is not eligible to be forgotten until the upsert + above has re-written it at least once (which happens on this very + call, for every pair still present). The transition fails toward + keeping a row an extra run rather than losing one it shouldn't. + + Returns ``{"loaded": N, "skipped": M, "updated": U, "forgotten": F}``. + """ + result = self.load_queries(pairs, upsert=True, mark_markdown_synced=True) + current_nls = {p["nl"] for p in pairs} + rows, _ = self.list_queries(limit=1_000_000) + stale_ids = [ + row["_row_id"] + for row in rows + if _has_tag(row.get("tags"), _MARKDOWN_SYNC_TAG) + and row["nl_query"] not in current_nls + ] + forgotten = self.forget_queries_by_ids(stale_ids) if stale_ids else 0 + return {**result, "forgotten": forgotten} + # ── Housekeeping ────────────────────────────────────────────────────── def status(self) -> dict: diff --git a/core/wren/tests/unit/test_memory.py b/core/wren/tests/unit/test_memory.py index 43157bafdf..45eb8b92af 100644 --- a/core/wren/tests/unit/test_memory.py +++ b/core/wren/tests/unit/test_memory.py @@ -1477,6 +1477,484 @@ def test_load_query_pairs_index_is_convergent(self, memory_store, tmp_path): second, _ = memory_store.list_queries(limit=100) assert len(first) == len(second) == 2 + def test_sync_forgets_pair_deleted_from_markdown(self, memory_store, tmp_path): + from wren.memory.markdown import ( # noqa: PLC0415 + load_query_pairs, + write_query_markdown, + ) + + write_query_markdown(tmp_path, "Total revenue", "SELECT SUM(amount) FROM o") + write_query_markdown(tmp_path, "Count orders", "SELECT COUNT(*) FROM o") + memory_store.sync_markdown_queries(load_query_pairs(tmp_path)) + + (tmp_path / "knowledge" / "sql" / "count-orders.md").unlink() + result = memory_store.sync_markdown_queries(load_query_pairs(tmp_path)) + + assert result["forgotten"] == 1 + remaining, total = memory_store.list_queries(limit=100) + assert total == 1 + assert remaining[0]["nl_query"] == "Total revenue" + # the deleted example must no longer surface in recall either + hits = memory_store.recall_queries("count orders", limit=5) + assert all(h["nl_query"] != "Count orders" for h in hits) + + def test_sync_preserves_seed_queries_absent_from_markdown( + self, memory_store, tmp_path + ): + from wren.memory.markdown import ( # noqa: PLC0415 + load_query_pairs, + write_query_markdown, + ) + + memory_store.index_schema(_MANIFEST) # generates seed queries + before, _ = memory_store.list_queries(limit=100, source="seed") + assert before # sanity: seeds exist + + write_query_markdown(tmp_path, "Total revenue", "SELECT SUM(amount) FROM o") + memory_store.sync_markdown_queries(load_query_pairs(tmp_path)) + + after, _ = memory_store.list_queries(limit=100, source="seed") + assert len(after) == len(before) # seeds untouched by the markdown sync + + def test_sync_preserves_legacy_queries_yml_pairs_absent_from_markdown( + self, memory_store + ): + """A pair loaded from a project's legacy queries.yml is not markdown- + backed, so a naive sync_markdown_queries() forgets it as stale on + every run and immediately reloads it, reporting a bogus "forgot 1 + stale pair" each time. Tagging it source:legacy and excluding that + tag from the forget scan (same as seed/view) makes it stable across + repeated syncs with an empty markdown set. + """ + memory_store.load_queries( + [ + { + "nl": "Total revenue", + "sql": "SELECT SUM(amount) FROM o", + "source": "legacy", + } + ], + upsert=False, + ) + before, _ = memory_store.list_queries(limit=100, source="legacy") + assert len(before) == 1 + + result = memory_store.sync_markdown_queries([]) # no markdown pairs + + assert result["forgotten"] == 0 + after, _ = memory_store.list_queries(limit=100, source="legacy") + assert len(after) == 1 + assert after[0]["sql_query"] == "SELECT SUM(amount) FROM o" + + def test_sync_preserves_view_queries_absent_from_markdown(self, memory_store): + """A source:view row is not markdown-backed either (there is no + knowledge/sql/*.md file for a manifest view), so a markdown sync + with an empty pair set must never forget it.""" + memory_store.load_queries( + [ + { + "nl": "Rows in the active orders view", + "sql": "SELECT * FROM active_orders_view", + "source": "view", + } + ], + upsert=False, + ) + before, _ = memory_store.list_queries(limit=100, source="view") + assert len(before) == 1 + + result = memory_store.sync_markdown_queries([]) # no markdown pairs + + assert result["forgotten"] == 0 + after, _ = memory_store.list_queries(limit=100, source="view") + assert len(after) == 1 + assert after[0]["sql_query"] == "SELECT * FROM active_orders_view" + + def test_sync_preserves_user_pair_loaded_from_yaml_with_no_markdown_file( + self, memory_store + ): + """`wren memory load pairs.yml` (no --source given) writes a plain + source:user row backed by no knowledge/sql/*.md file at all: the + YAML file may live anywhere, or be deleted right after the import. + A markdown sync must never treat that row as stale just because its + source happens to be "user" too, which is also what every + markdown-backed pair defaults to. + """ + memory_store.load_queries( + [ + { + "nl": "Revenue by region", + "sql": "SELECT region, SUM(amount) FROM orders GROUP BY region", + "source": "user", + } + ], + upsert=False, + ) + before, _ = memory_store.list_queries(limit=100, source="user") + assert len(before) == 1 + + result = memory_store.sync_markdown_queries([]) # no markdown pairs + + assert result["forgotten"] == 0 + after, _ = memory_store.list_queries(limit=100, source="user") + assert len(after) == 1 + assert after[0]["nl_query"] == "Revenue by region" + + def test_sync_preserves_pre_upgrade_queries_yml_import_once_file_is_gone( + self, memory_store + ): + """Before source:legacy tagging existed, a project's queries.yml was + imported as a plain source:user row (see the CLI's old, untagged + `load_queries(legacy_pairs, upsert=False)` call). A user who + consumed queries.yml on an older version, then deleted the file and + upgraded, must not lose that row on their first post-upgrade + `wren memory index`: there is no markdown file to judge it against, + and the row predates this sync entirely, so it never carries the + provenance tag a markdown sync would need to treat it as its own. + """ + memory_store.load_queries( + [ + { + "nl": "Total revenue", + "sql": "SELECT SUM(amount) FROM o", + "source": "user", + } + ], + upsert=False, + ) + + result = memory_store.sync_markdown_queries([]) # queries.yml is gone + + assert result["forgotten"] == 0 + after, _ = memory_store.list_queries(limit=100, source="user") + assert len(after) == 1 + assert after[0]["nl_query"] == "Total revenue" + + def test_source_filters_still_match_markdown_synced_rows( + self, memory_store, tmp_path + ): + """sync_markdown_queries tags the rows it writes with an extra + provenance token after the source tag (_MARKDOWN_SYNC_TAG), so the + stored tags string is no longer just "source:user". list/count/dump/ + forget filtering by --source must keep matching on the source value + alone, not on an exact match against the whole tags string. + """ + from wren.memory.markdown import ( # noqa: PLC0415 + load_query_pairs, + write_query_markdown, + ) + + write_query_markdown(tmp_path, "Total revenue", "SELECT SUM(amount) FROM o") + memory_store.sync_markdown_queries(load_query_pairs(tmp_path)) + + rows, total = memory_store.list_queries(source="user") + assert total == 1 + assert rows[0]["nl_query"] == "Total revenue" + assert memory_store.count_queries_by_source("user") == 1 + dumped = memory_store.dump_queries(source="user") + assert len(dumped) == 1 + + deleted = memory_store.forget_queries_by_source("user") + assert deleted == 1 + _, total_after = memory_store.list_queries() + assert total_after == 0 + + def test_source_filters_ignore_rows_carrying_no_source_tag(self, memory_store): + """`wren memory store` and the MCP `store_query` tool pass the caller's + own free-form labels through to `tags` (`"revenue,finance"`, or nothing + at all), so a real index holds rows with no `source:` token. A + `--source user` filter must not claim them, least of all `forget`. + """ + memory_store.store_query( + nl_query="A user-tagged pair", sql_query="SELECT 1", tags="revenue,finance" + ) + memory_store.store_query( + nl_query="An untagged pair", sql_query="SELECT 2", tags=None + ) + memory_store.store_query( + nl_query="A source-tagged pair", sql_query="SELECT 3", tags="source:user" + ) + + rows, total = memory_store.list_queries(limit=100, source="user") + assert total == 1 + assert rows[0]["nl_query"] == "A source-tagged pair" + assert memory_store.count_queries_by_source("user") == 1 + assert [r["nl_query"] for r in memory_store.dump_queries(source="user")] == [ + "A source-tagged pair" + ] + + assert memory_store.forget_queries_by_source("user") == 1 + survivors, _ = memory_store.list_queries(limit=100) + assert sorted(r["nl_query"] for r in survivors) == [ + "A user-tagged pair", + "An untagged pair", + ] + + def test_sync_lets_a_markdown_pair_win_over_a_colliding_seed_row( + self, memory_store, tmp_path + ): + """User-authored content wins over an auto-generated seed with the same nl. + + A seed sharing a markdown pair's nl_query is not protected: the sync + upserts the markdown pair over it. This is deliberate, not a gap. The + seed is regenerated by index_schema() on every reindex (self-healing), + while a dropped markdown pair is permanent (its file is still on disk, + so it can never be re-synced, and `wren memory check` would report it + as unfixably out of sync forever). An earlier version of this method + pre-filtered markdown pairs against existing seed/view nl_query values + to "protect" the seed row; that filter made check/index loop on the + markdown pair forever and is why this test asserts the opposite of + what its name once claimed. + """ + from wren.memory.markdown import ( # noqa: PLC0415 + load_query_pairs, + write_query_markdown, + ) + + memory_store.store_query( + nl_query="Total revenue", + sql_query="SELECT SUM(o_totalprice) FROM orders", + tags="source:seed", + ) + before, _ = memory_store.list_queries(limit=100, source="seed") + assert len(before) == 1 + + write_query_markdown(tmp_path, "Total revenue", "SELECT SUM(amount) FROM o") + result = memory_store.sync_markdown_queries(load_query_pairs(tmp_path)) + + # the markdown pair replaced the seed row (an upsert, not a fresh + # load) rather than being skipped + _, total = memory_store.list_queries(limit=100) + assert total == 1 + assert result["updated"] == 1 + remaining, _ = memory_store.list_queries(limit=100) + assert remaining[0]["sql_query"] == "SELECT SUM(amount) FROM o" + + # and the markdown pair now stays in sync on a second run: nothing to + # load, nothing to forget, no unfixable "not indexed" loop + result2 = memory_store.sync_markdown_queries(load_query_pairs(tmp_path)) + assert result2["forgotten"] == 0 + after, _ = memory_store.list_queries(limit=100) + assert len(after) == 1 + assert after[0]["sql_query"] == "SELECT SUM(amount) FROM o" + + def test_sync_with_no_markdown_pairs_forgets_all_user_pairs( + self, memory_store, tmp_path + ): + from wren.memory.markdown import ( # noqa: PLC0415 + load_query_pairs, + write_query_markdown, + ) + + write_query_markdown(tmp_path, "Total revenue", "SELECT SUM(amount) FROM o") + memory_store.sync_markdown_queries(load_query_pairs(tmp_path)) + (tmp_path / "knowledge" / "sql" / "total-revenue.md").unlink() + + result = memory_store.sync_markdown_queries(load_query_pairs(tmp_path)) + assert result["forgotten"] == 1 + _, total = memory_store.list_queries(limit=100) + assert total == 0 + + def test_cli_index_lancedb_forgets_deleted_pair(self, tmp_path, monkeypatch): + """`wren memory index` on the lancedb backend must actually forget a + pair deleted from knowledge/sql/*.md, not just report it as stale.""" + pytest.importorskip("lancedb", reason="wren[memory] extras not installed") + pytest.importorskip( + "sentence_transformers", reason="wren[memory] extras not installed" + ) + from typer.testing import CliRunner # noqa: PLC0415 + + from wren.cli import app # noqa: PLC0415 + from wren.memory.markdown import write_query_markdown # noqa: PLC0415 + from wren.memory.store import MemoryStore # noqa: PLC0415 + + monkeypatch.setenv("WREN_PROJECT_HOME", str(tmp_path)) + monkeypatch.setenv("WREN_MEMORY_BACKEND", "lancedb") + (tmp_path / "target").mkdir() + (tmp_path / "target" / "mdl.json").write_text("{}", encoding="utf-8") + write_query_markdown(tmp_path, "Total revenue", "SELECT SUM(amount) FROM o") + write_query_markdown(tmp_path, "Count orders", "SELECT COUNT(*) FROM o") + + cli = CliRunner() + first = cli.invoke(app, ["memory", "index"]) + assert first.exit_code == 0, first.output + + (tmp_path / "knowledge" / "sql" / "count-orders.md").unlink() + second = cli.invoke(app, ["memory", "index"]) + assert second.exit_code == 0, second.output + assert "Forgot 1 stale pair" in second.output + + store = MemoryStore(path=str(tmp_path / ".wren" / "memory")) + rows, total = store.list_queries(limit=100) + assert total == 1 + assert rows[0]["nl_query"] == "Total revenue" + + def test_cli_load_then_index_does_not_forget_the_loaded_pair( + self, tmp_path, monkeypatch + ): + """`wren memory load .yml` (see PR #2703 review) writes a + source:user row with no knowledge/sql/*.md file behind it. A + `wren memory index` right after must not report or execute a forget + for that row: it was never markdown-backed to begin with.""" + pytest.importorskip("lancedb", reason="wren[memory] extras not installed") + pytest.importorskip( + "sentence_transformers", reason="wren[memory] extras not installed" + ) + from typer.testing import CliRunner # noqa: PLC0415 + + from wren.cli import app # noqa: PLC0415 + from wren.memory.store import MemoryStore # noqa: PLC0415 + + monkeypatch.setenv("WREN_PROJECT_HOME", str(tmp_path)) + monkeypatch.setenv("WREN_MEMORY_BACKEND", "lancedb") + (tmp_path / "target").mkdir() + (tmp_path / "target" / "mdl.json").write_text("{}", encoding="utf-8") + + pairs_file = tmp_path / "pairs.yml" + pairs_file.write_text( + "version: 1\n" + "pairs:\n" + " - nl: Revenue by region\n" + " sql: SELECT region, SUM(amount) FROM orders GROUP BY region\n", + encoding="utf-8", + ) + + cli = CliRunner() + load_result = cli.invoke(app, ["memory", "load", str(pairs_file)]) + assert load_result.exit_code == 0, load_result.output + assert "1 new" in load_result.output + + index_result = cli.invoke(app, ["memory", "index"]) + assert index_result.exit_code == 0, index_result.output + assert "Forgot" not in index_result.output + + # ...and `check` must agree, rather than reporting drift that `index` + # has just proven it cannot clear. + check_result = cli.invoke(app, ["memory", "check"]) + assert check_result.exit_code == 0, check_result.output + assert "In sync." in check_result.output + + store = MemoryStore(path=str(tmp_path / ".wren" / "memory")) + rows, total = store.list_queries(limit=100) + assert total == 1 + assert rows[0]["nl_query"] == "Revenue by region" + + def test_cli_watch_reindex_on_start_lancedb_forgets_deleted_pair( + self, tmp_path, monkeypatch + ): + """`wren memory watch` calls the same reindex path as `index`, so a + deletion must clear on its next forced reindex too.""" + pytest.importorskip("lancedb", reason="wren[memory] extras not installed") + pytest.importorskip( + "sentence_transformers", reason="wren[memory] extras not installed" + ) + from typer.testing import CliRunner # noqa: PLC0415 + + from wren.cli import app # noqa: PLC0415 + from wren.memory.markdown import write_query_markdown # noqa: PLC0415 + from wren.memory.store import MemoryStore # noqa: PLC0415 + + monkeypatch.setenv("WREN_PROJECT_HOME", str(tmp_path)) + monkeypatch.setenv("WREN_MEMORY_BACKEND", "lancedb") + (tmp_path / "wren_project.yml").write_text("name: t\n", encoding="utf-8") + (tmp_path / "target").mkdir() + (tmp_path / "target" / "mdl.json").write_text("{}", encoding="utf-8") + write_query_markdown(tmp_path, "Total revenue", "SELECT SUM(amount) FROM o") + write_query_markdown(tmp_path, "Count orders", "SELECT COUNT(*) FROM o") + + cli = CliRunner() + first = cli.invoke( + app, ["memory", "watch", "--reindex-on-start", "--max-polls", "1"] + ) + assert first.exit_code == 0, first.output + + (tmp_path / "knowledge" / "sql" / "count-orders.md").unlink() + second = cli.invoke( + app, ["memory", "watch", "--reindex-on-start", "--max-polls", "1"] + ) + assert second.exit_code == 0, second.output + + store = MemoryStore(path=str(tmp_path / ".wren" / "memory")) + rows, total = store.list_queries(limit=100) + assert total == 1 + assert rows[0]["nl_query"] == "Total revenue" + + def test_cli_check_does_not_flag_legacy_queries_yml_pairs_as_stale( + self, tmp_path, monkeypatch + ): + """A project with only a legacy queries.yml (no knowledge/sql/*.md for + it) must reach `check`'s "In sync." across repeated index runs. Before + source:legacy was excluded from check()'s own stale filter (which used + to hardcode ("seed", "view") independently of sync_markdown_queries's + set), a legacy pair read as a "user" pair not present in markdown and + was reported as permanent drift no `index` run could clear. + """ + pytest.importorskip("lancedb", reason="wren[memory] extras not installed") + pytest.importorskip( + "sentence_transformers", reason="wren[memory] extras not installed" + ) + from typer.testing import CliRunner # noqa: PLC0415 + + from wren.cli import app # noqa: PLC0415 + + monkeypatch.setenv("WREN_PROJECT_HOME", str(tmp_path)) + monkeypatch.setenv("WREN_MEMORY_BACKEND", "lancedb") + (tmp_path / "target").mkdir() + (tmp_path / "target" / "mdl.json").write_text("{}", encoding="utf-8") + (tmp_path / "queries.yml").write_text( + "pairs:\n - nl: Total revenue\n sql: SELECT SUM(amount) FROM o\n", + encoding="utf-8", + ) + + cli = CliRunner() + for _ in range(2): + index_result = cli.invoke(app, ["memory", "index"]) + assert index_result.exit_code == 0, index_result.output + assert "Forgot" not in index_result.output + + check_result = cli.invoke(app, ["memory", "check"]) + assert check_result.exit_code == 0, check_result.output + assert "In sync." in check_result.output + + def test_cli_export_then_index_reports_in_sync(self, tmp_path, monkeypatch): + """`wren memory export` is the documented one-time migration, and it + preserves each row's source into the markdown frontmatter, so a + legacy queries.yml pair lands in knowledge/sql/ as `source: legacy`. + That file is markdown-backed like any other, so `check` must not read + it as "not indexed". + """ + pytest.importorskip("lancedb", reason="wren[memory] extras not installed") + pytest.importorskip( + "sentence_transformers", reason="wren[memory] extras not installed" + ) + from typer.testing import CliRunner # noqa: PLC0415 + + from wren.cli import app # noqa: PLC0415 + + monkeypatch.setenv("WREN_PROJECT_HOME", str(tmp_path)) + monkeypatch.setenv("WREN_MEMORY_BACKEND", "lancedb") + (tmp_path / "target").mkdir() + (tmp_path / "target" / "mdl.json").write_text("{}", encoding="utf-8") + (tmp_path / "queries.yml").write_text( + "pairs:\n - nl: Total revenue\n sql: SELECT SUM(amount) FROM o\n", + encoding="utf-8", + ) + + cli = CliRunner() + assert cli.invoke(app, ["memory", "index"]).exit_code == 0 + export = cli.invoke(app, ["memory", "export"]) + assert export.exit_code == 0, export.output + assert "source: legacy" in next( + (tmp_path / "knowledge" / "sql").glob("*.md") + ).read_text(encoding="utf-8") + + for _ in range(2): + index_result = cli.invoke(app, ["memory", "index"]) + assert index_result.exit_code == 0, index_result.output + check_result = cli.invoke(app, ["memory", "check"]) + assert check_result.exit_code == 0, check_result.output + assert "In sync." in check_result.output + def test_reset_then_reindex_restores_from_markdown(self, memory_store, tmp_path): from wren.memory.markdown import ( # noqa: PLC0415 load_query_pairs,