diff --git a/src/grimoire/mcp/tools/retrieval_review.py b/src/grimoire/mcp/tools/retrieval_review.py index 71c0f4ac..ec717f52 100644 --- a/src/grimoire/mcp/tools/retrieval_review.py +++ b/src/grimoire/mcp/tools/retrieval_review.py @@ -42,7 +42,11 @@ async def menus_to_review( Args: source: the arm to review. Only arms whose search can be re-run - exactly are offered; today that is `auto_inject`. + exactly are offered: `auto_inject`, and the rule arms + `prompt_rule`, `pre_tool_rule` and `write_path_rule`. A rule + line's passage is the trigger the arm matched, and for the + command and write arms `within_budget` also needs the line to + sit inside the band of the top hit, as it did when shown. n: how many calls to sample, 1–20. Random among the calls you have not judged, so a review does not describe one afternoon. days: how far back to sample from. diff --git a/src/grimoire/services/retrieval_review.py b/src/grimoire/services/retrieval_review.py index a2c0afcd..0b96f447 100644 --- a/src/grimoire/services/retrieval_review.py +++ b/src/grimoire/services/retrieval_review.py @@ -46,13 +46,18 @@ from grimoire.models.base import iso from grimoire.models.note_usage import PULLED, NoteUsageEvent from grimoire.models.retrieval_judgment import VERDICTS, RetrievalJudgment from grimoire.models.retrieval_log import RetrievalLog +from grimoire.models.rule_usage import PULLED as RULE_PULLED +from grimoire.models.rule_usage import RuleUsageEvent logger = logging.getLogger(__name__) # The arms a review can re-run FAITHFULLY. A re-run with different parameters # from the arm's would judge a menu nobody was shown, so a surface joins this -# only with its own search spelled out in `_rerun`. -REVIEWABLE = ("auto_inject",) +# only with its own search spelled out in `_rerun`. The rule arms named here +# are the ones that annotate: each logs its query, floor, budget and project, +# which is everything its search takes. +RULE_REVIEWABLE = ("prompt_rule", "pre_tool_rule", "write_path_rule") +REVIEWABLE = ("auto_inject", *RULE_REVIEWABLE) DEPTH_DEFAULT = 6 DEPTH_MAX = 10 @@ -94,7 +99,58 @@ def _after(stamp, call) -> bool: async def _rerun( user_id: int, row: RetrievalLog, depth: int, report: dict | None = None, ) -> list[dict]: - """The logged call's query, searched again the way its arm searches. + """The logged call's query, searched again the way its arm searches.""" + if row.source in RULE_REVIEWABLE: + return await _rerun_rules(user_id, row, depth, report) + return await _rerun_notes(user_id, row, depth, report) + + +async def _rerun_rules( + user_id: int, row: RetrievalLog, depth: int, report: dict | None = None, +) -> list[dict]: + """A rule arm's call, searched again: its query against rule triggers, + with the floor and project it logged. + + `within_budget` is the arm's own cut: a line inside the budget, and, for + an arm that bands, inside the band of the top hit. The preference slot is + left out — it logs its own source. Rules created after the call are + dropped and named in `report["postdated"]`, as for notes. + """ + from grimoire.services import retrieval_pipeline as rp + from grimoire.services.embeddings import semantic_search_rules + + arm = next(a for a in rp.RULE_ARMS if a.source == row.source) + hits = await semantic_search_rules( + user_id, row.query or "", + limit=depth + POSTDATED_SLACK, + threshold=row.threshold if row.threshold is not None else 0.0, + project_id=row.project_id or None, + ) + postdated = [int(r.id) for _s, r in hits if _after(r.created_at, row.created_at)] + if report is not None: + report["postdated"] = postdated + hits = [h for h in hits if int(h[1].id) not in postdated][:depth] + budget = int(row.limit_n or 0) + cut = hits[:budget] + kept = {int(r.id) for _s, r in (rp._rule_band(cut) if arm.band else cut)} + return [{ + "rank": rank, + "record_id": int(rule.id), + "score": round(float(score), 4), + "within_budget": int(rule.id) in kept, + "changed_since_call": _after(rule.updated_at, row.created_at), + "kind": rule.kind or "rule", + "name": rule.title, + # The trigger is what the arm matched against and what its line + # quotes, so it is the passage the reader saw. + "passage": rule.when_to_apply or rule.statement or None, + } for rank, (score, rule) in enumerate(hits, start=1)] + + +async def _rerun_notes( + user_id: int, row: RetrievalLog, depth: int, report: dict | None = None, +) -> list[dict]: + """A notes arm's call, searched again. `auto_inject`'s parameters, read from its spec (`retrieval_pipeline. AUTO_INJECT`) rather than restated: which kinds, whose records, and the @@ -151,19 +207,25 @@ async def _opened_after(user_id: int, row: RetrievalLog, ids: list[int]) -> set[ Agent pulls only — the `mcp_` prefix, per NoteUsageEvent's convention — because the question is whether the line moved the session, not whether - the operator later browsed to the record. + the operator later browsed to the record. A rule arm's lines are read + from the rule table: a note and a rule can share an id. """ if not ids or row.created_at is None: return set() + table, id_col, pulled = ( + (RuleUsageEvent, RuleUsageEvent.rule_id, RULE_PULLED) + if row.source in RULE_REVIEWABLE + else (NoteUsageEvent, NoteUsageEvent.note_id, PULLED) + ) async with async_session() as session: found = await session.execute( - select(NoteUsageEvent.note_id).distinct().where( - NoteUsageEvent.user_id == user_id, - NoteUsageEvent.event == PULLED, - NoteUsageEvent.source.startswith("mcp_", autoescape=True), - NoteUsageEvent.note_id.in_(ids), - NoteUsageEvent.created_at >= row.created_at, - NoteUsageEvent.created_at <= row.created_at + OPENED_WINDOW, + select(id_col).distinct().where( + table.user_id == user_id, + table.event == pulled, + table.source.startswith("mcp_", autoescape=True), + id_col.in_(ids), + table.created_at >= row.created_at, + table.created_at <= row.created_at + OPENED_WINDOW, ) ) return {int(i) for (i,) in found.all()} @@ -175,7 +237,7 @@ async def _logged_after(row: RetrievalLog, ids: list[int]) -> set[int]: (a closed task takes no claim stamp), so it is read here.""" from grimoire.models.task_log import TaskLog - if not ids or row.created_at is None: + if not ids or row.created_at is None or row.source in RULE_REVIEWABLE: return set() async with async_session() as session: found = await session.execute( diff --git a/tests/test_retrieval_review.py b/tests/test_retrieval_review.py index 5b53e9ba..d2fb46e2 100644 --- a/tests/test_retrieval_review.py +++ b/tests/test_retrieval_review.py @@ -81,7 +81,7 @@ async def test_the_re_run_searches_the_way_the_arm_does(): rerun_kw.update(kw) return [] - row = SimpleNamespace(query="q", threshold=0.5, limit_n=3, project_id=None, + row = SimpleNamespace(source="auto_inject", query="q", threshold=0.5, limit_n=3, project_id=None, created_at=CALL) with patch("grimoire.services.embeddings.semantic_search_notes", side_effect=rerun_search): await review._rerun(1, row, 3) @@ -99,7 +99,7 @@ BEFORE = CALL - timedelta(days=1) @pytest.mark.asyncio async def test_the_re_run_ranks_from_one_and_marks_the_budget_cut(): - row = SimpleNamespace(query="q", threshold=0.55, limit_n=2, project_id=None, + row = SimpleNamespace(source="auto_inject", query="q", threshold=0.55, limit_n=2, project_id=None, created_at=CALL) hits = [ (0.81, fake_note(id=11, title="first", created_at=BEFORE, updated_at=BEFORE)), @@ -127,7 +127,7 @@ async def test_a_record_written_after_the_call_is_dropped_before_ranking(): """The session that made a call goes on to write about it, and that record outranks everything in a re-run. The call could never have been offered it, so it must not take a rank — or a verdict — from what was.""" - row = SimpleNamespace(query="q", threshold=0.55, limit_n=1, project_id=None, + row = SimpleNamespace(source="auto_inject", query="q", threshold=0.55, limit_n=1, project_id=None, created_at=CALL) after = CALL + timedelta(hours=1) hits = [ @@ -147,6 +147,103 @@ async def test_a_record_written_after_the_call_is_dropped_before_ranking(): for ln in lines] == [(1, 22, True, False), (2, 23, False, True)] +# ── the rule arms ──────────────────────────────────────────────────────── + + +def _rule(rid, title, created=BEFORE, updated=BEFORE): + return SimpleNamespace( + id=rid, title=title, kind="rule", when_to_apply=f"when {title}", + statement="s", created_at=created, updated_at=updated, + ) + + +def test_the_reviewable_rule_arms_are_the_ones_that_annotate(): + """Every rule arm that shows lines can be reviewed; the reply backstop + shows none, so there is no menu to judge.""" + from grimoire.services import retrieval_pipeline as rp + + assert set(review.RULE_REVIEWABLE) == { + a.source for a in rp.RULE_ARMS if not a.stop_only + } + + +@pytest.mark.asyncio +async def test_the_rule_re_run_searches_the_way_the_arm_does(): + """The floor and the project the arm searched with, run and compared.""" + from grimoire.services import retrieval_pipeline as rp + + arm_kw: dict = {} + + async def arm_search(_uid, _q, **kw): + if "kind" not in kw: + arm_kw.update(kw) + return [] + + await rp.run_rule_arm( + rp.PROMPT_RULE, rp.RuleMoment(user_id=1, query="q", project_id=7), + floor=0.7, budget=3, + io=rp.RuleIO(search=arm_search, record_retrieval=MagicMock(), + record_rule_surfaced=MagicMock()), + ) + rerun_kw: dict = {} + + async def rerun_search(_uid, _q, **kw): + rerun_kw.update(kw) + return [] + + row = SimpleNamespace(source="prompt_rule", query="q", threshold=0.7, limit_n=3, + project_id=7, created_at=CALL) + with patch("grimoire.services.embeddings.semantic_search_rules", side_effect=rerun_search): + await review._rerun(1, row, 3) + + for key in ("threshold", "project_id"): + assert rerun_kw[key] == arm_kw[key], key + assert "kind" not in rerun_kw + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("source", "within"), [ + # No band: the budget alone is the cut. + ("prompt_rule", [True, True, True, False]), + # Banded: the third sits outside the band of the top hit, so it was + # not shown though the budget had room for it. + ("pre_tool_rule", [True, True, False, False]), +]) +async def test_a_rule_line_is_inside_the_cut_only_where_the_arm_showed_it(source, within): + row = SimpleNamespace(source=source, query="git push", threshold=0.7, limit_n=3, + project_id=None, created_at=CALL) + hits = [(0.86, _rule(1, "one")), (0.84, _rule(2, "two")), + (0.76, _rule(3, "three")), (0.74, _rule(4, "four"))] + + async def search(_uid, _q, **kw): + assert kw["threshold"] == 0.7 and kw["project_id"] is None + return hits + + with patch("grimoire.services.embeddings.semantic_search_rules", side_effect=search): + lines = await review._rerun(1, row, 4) + assert [ln["within_budget"] for ln in lines] == within + assert lines[0]["passage"] == "when one" + assert lines[0]["kind"] == "rule" + + +@pytest.mark.asyncio +async def test_a_rule_written_after_the_call_is_dropped_before_ranking(): + row = SimpleNamespace(source="write_path_rule", query="q", threshold=0.7, limit_n=2, + project_id=None, created_at=CALL) + after = CALL + timedelta(hours=1) + hits = [(0.9, _rule(31, "written after", created=after, updated=after)), + (0.8, _rule(32, "held then"))] + + async def search(_uid, _q, **kw): + return hits + + rep: dict = {} + with patch("grimoire.services.embeddings.semantic_search_rules", side_effect=search): + lines = await review._rerun(1, row, 2, report=rep) + assert rep["postdated"] == [31] + assert [(ln["rank"], ln["record_id"]) for ln in lines] == [(1, 32)] + + # ── against Postgres ─────────────────────────────────────────────────────