fix: close remaining consolidation review findings

engraphis/core/consolidate.py

@@ -60,6 +60,10 @@
 # Bound the population that reaches the quadratic fallback clustering pass while
 # allowing the storage scan to page in smaller batches and skip pending rows.
 DISTILL_CLUSTER_LIMIT = 2000
+# Carry a small, bounded set of incomplete clusters into the next sweep.  This keeps
+# interleaved subjects from being permanently split across rotating windows without
+# turning a maintenance pass into an unbounded full-table scan.
+DISTILL_PENDING_LIMIT = 512
 # Cursor name for the bounded episodic sweep; the value is scoped by workspace/repo.
 DISTILL_CURSOR_NAME = "episodic-consolidation"
 
@@ -203,29 +207,54 @@ def _scan_memory_window(store, flt: SearchFilter, *, mtypes: list[MemoryType],
                         batch_size: int, prompt_only: bool = False,
                         max_records: Optional[int] = None,
                         exclude_relation: Optional[str] = None,
-                        start_after_id: str = "") -> tuple[list[MemoryRecord], str]:
+                        start_after_id: str = "", overlap: int = 0,
+                        advance_records: Optional[int] = None,
+                        ) -> tuple[list[MemoryRecord], str]:
     """Read one bounded keyset window and return its next persistent cursor.
 
     ``Store.list_memories_page`` orders by id.  When a bounded window reaches the
     end, the empty cursor deliberately makes the *next* sweep wrap to the start;
     this rotates maintenance over all eligible rows without materializing or
-    clustering the full population on every run.
+    clustering the full population on every run.  ``overlap`` retains a bounded
+    suffix of the raw keyset window for the next sweep, which keeps clusters that
+    straddle a maintenance boundary intact.  ``advance_records`` is a raw-row
+    progress floor used when filtering leaves fewer than ``max_records`` eligible
+    rows; it prevents a bounded sweep from pinning its cursor on an excluded page.
     """
     size = max(1, int(batch_size))
     cap = None if max_records is None else max(0, int(max_records))
-    if cap == 0:
+    advance_cap = (
+        None if advance_records is None else max(0, int(advance_records))
+    )
+    if cap == 0 or advance_cap == 0:
         return [], str(start_after_id or "")
     after_id = str(start_after_id or "")
     records: list[MemoryRecord] = []
+    window_ids: list[str] = []
     scoped = _replace(flt, mtypes=mtypes)
     next_cursor = ""
+
+    def cursor_for_window() -> str:
+        retained = min(max(0, int(overlap)), max(0, len(window_ids) - 1))
+        return window_ids[len(window_ids) - retained - 1]
+
     while True:
-        page = store.list_memories_page(scoped, after_id=after_id, limit=size)
+        if advance_cap is not None and len(window_ids) >= advance_cap:
+            next_cursor = cursor_for_window()
+            break
+        remaining = None if cap is None else max(1, cap - len(records))
+        page_limit = size if remaining is None else min(size, remaining)
+        if advance_cap is not None:
+            page_limit = min(page_limit, advance_cap - len(window_ids))
+        page = store.list_memories_page(
+            scoped, after_id=after_id, limit=page_limit,
+        )
         if not page:
             # The persisted cursor was at the end of the keyspace. Start the next
             # sweep from the beginning instead of retrying an empty tail forever.
             break
         next_after = page[-1].id
+        window_ids.extend(memory.id for memory in page)
         page_size = len(page)
         if exclude_relation:
             excluded = _linked_memory_ids(
@@ -240,12 +269,22 @@ def _scan_memory_window(store, flt: SearchFilter, *, mtypes: list[MemoryType],
         else:
             records.extend(page)
         if cap is not None and len(records) >= cap:
-            next_cursor = next_after
+            next_cursor = cursor_for_window()
+            break
+        if advance_cap is not None and len(window_ids) >= advance_cap:
+            next_cursor = cursor_for_window()
             break
-        if next_after == after_id or page_size < size:
+        if next_after == after_id or page_size < page_limit:
             # End-of-keyspace: clear the cursor for the next invocation.
             break
         after_id = next_after
+    if (
+        not next_cursor
+        and advance_cap is not None
+        and window_ids
+        and len(window_ids) >= max(1, advance_cap - max(0, int(overlap)))
+    ):
+        next_cursor = cursor_for_window()
     records.sort(
         key=lambda memory: (
             memory.ingested_at if memory.ingested_at is not None else float("-inf"),
@@ -256,6 +295,90 @@ def _scan_memory_window(store, flt: SearchFilter, *, mtypes: list[MemoryType],
     return records[:cap] if cap is not None else records, next_cursor
 
 
+def _decode_distill_cursor(value: str) -> tuple[str, list[str]]:
+    """Read a legacy keyset cursor or the current cursor-plus-candidates state."""
+    raw = str(value or "")
+    if not raw.startswith("{"):
+        return raw, []
+    try:
+        state = json.loads(raw)
+    except (TypeError, ValueError):
+        return raw, []
+    if not isinstance(state, dict):
+        return raw, []
+    cursor = state.get("cursor")
+    pending = state.get("pending")
+    if not isinstance(cursor, str) or not isinstance(pending, list):
+        return raw, []
+    ids = list(dict.fromkeys(
+        str(memory_id) for memory_id in pending if str(memory_id or "")
+    ))
+    return cursor, ids[:DISTILL_PENDING_LIMIT]
+
+
+def _encode_distill_cursor(cursor: str, pending_ids: list[str]) -> str:
+    """Persist bounded partial-cluster candidates alongside the scan cursor."""
+    normalized = list(dict.fromkeys(
+        str(memory_id) for memory_id in pending_ids if str(memory_id or "")
+    ))[:DISTILL_PENDING_LIMIT]
+    if not normalized:
+        return str(cursor or "")
+    return json.dumps(
+        {"cursor": str(cursor or ""), "pending": normalized},
+        ensure_ascii=False,
+        separators=(",", ":"),
+    )
+
+
+def _load_distill_candidates(
+    store, flt: SearchFilter, pending_ids: list[str], *, now: float,
+) -> list[MemoryRecord]:
+    """Reload prior partial-cluster sources, dropping deleted or ineligible rows."""
+    from engraphis.core.store import memory_matches_filter
+
+    if not pending_ids:
+        return []
+    records: list[MemoryRecord] = []
+    for memory_id in dict.fromkeys(pending_ids):
+        memory = store.get_memory(memory_id)
+        if (
+            memory is not None
+            and memory.mtype == MemoryType.EPISODIC
+            and memory_matches_filter(memory, flt, at=now)
+            and prompt_eligible(memory.provenance, memory.metadata)
+        ):
+            records.append(memory)
+    if not records:
+        return []
+    linked = _linked_memory_ids(
+        store, [memory.id for memory in records], relation="consolidates",
+    )
+    return [memory for memory in records if memory.id not in linked]
+
+
+def _pending_distill_ids(clusters: list[list[MemoryRecord]], *, min_cluster: int) -> list[str]:
+    """Select bounded candidates whose cluster needs a later sweep to complete.
+
+    Multi-record partial clusters are preferred because they carry positive evidence
+    that a subject is recurring.  Singleton records with an explicit subject key are
+    retained too; unrelated singleton noise is intentionally not allowed to consume
+    the whole carry-over budget.
+    """
+    incomplete = [cluster for cluster in clusters if 0 < len(cluster) < min_cluster]
+    prioritized = [
+        cluster for cluster in incomplete
+        if len(cluster) > 1 or any(memory.subject_key for memory in cluster)
+    ]
+    selected: list[str] = []
+    for cluster in prioritized:
+        for memory in cluster[-max(1, min_cluster - 1):]:
+            if memory.id not in selected:
+                selected.append(memory.id)
+            if len(selected) >= DISTILL_PENDING_LIMIT:
+                return selected
+    return selected
+
+
 def _scan_memories(store, flt: SearchFilter, *, mtypes: list[MemoryType],
                    batch_size: int, prompt_only: bool = False,
                    max_records: Optional[int] = None,
@@ -687,19 +810,30 @@ def consolidate(engine, *, workspace_id: str, repo_id: Optional[str] = None,
                 except Exception as exc:
                     report["errors"].append(_error_entry(retry_cluster, exc))
 
-    distill_cursor = store.get_maintenance_cursor(
+    distill_state = store.get_maintenance_cursor(
         workspace_id, repo_id, DISTILL_CURSOR_NAME,
     )
+    distill_cursor, pending_ids = _decode_distill_cursor(distill_state)
+    distill_overlap = max(0, int(min_cluster) - 1)
     episodic, next_distill_cursor = _scan_memory_window(
         store, flt, mtypes=[MemoryType.EPISODIC],
         batch_size=DISTILL_SCAN_LIMIT, prompt_only=True,
-        max_records=DISTILL_CLUSTER_LIMIT,
+        max_records=DISTILL_CLUSTER_LIMIT + distill_overlap,
         exclude_relation="consolidates",
         start_after_id=distill_cursor,
+        overlap=distill_overlap,
+        advance_records=DISTILL_CLUSTER_LIMIT + distill_overlap,
     )
-    if not dry_run:
-        store.set_maintenance_cursor(
-            workspace_id, repo_id, DISTILL_CURSOR_NAME, next_distill_cursor,
+    prior_candidates = _load_distill_candidates(store, flt, pending_ids, now=now)
+    if prior_candidates:
+        by_id = {memory.id: memory for memory in [*prior_candidates, *episodic]}
+        episodic = sorted(
+            by_id.values(),
+            key=lambda memory: (
+                memory.ingested_at if memory.ingested_at is not None else float("-inf"),
+                memory.id,
+            ),
+            reverse=True,
         )
     # A digest inherits its owner from its first source.  Cluster only records that have
     # the exact same owner, otherwise a workspace sweep could write one repo's digest with
@@ -711,6 +845,14 @@ def consolidate(engine, *, workspace_id: str, repo_id: Optional[str] = None,
             owner_memories, threshold=subject_jaccard, store=store, flt=flt,
         )
     ]
+    if not dry_run:
+        store.set_maintenance_cursor(
+            workspace_id, repo_id, DISTILL_CURSOR_NAME,
+            _encode_distill_cursor(
+                next_distill_cursor,
+                _pending_distill_ids(clusters, min_cluster=min_cluster),
+            ),
+        )
 
     if structured:
         report["structured"] = {"enabled": True, "attempted": 0, "succeeded": 0,
@@ -1424,11 +1566,15 @@ def consolidate_profiles(engine, *, workspace_id: str, repo_id: Optional[str] =
     profile_cursor = store.get_maintenance_cursor(
         workspace_id, repo_id, PROFILE_CURSOR_NAME,
     )
+    profile_overlap = max(0, int(min_mentions) - 1)
     profile_memories, next_profile_cursor = _scan_memory_window(
         store, flt, mtypes=DURABLE_TYPES,
         batch_size=PROFILE_SCAN_LIMIT, prompt_only=True,
-        max_records=PROFILE_MEMORY_LIMIT, exclude_relation=PROFILE_RELATION,
+        max_records=PROFILE_MEMORY_LIMIT + profile_overlap,
+        exclude_relation=PROFILE_RELATION,
         start_after_id=profile_cursor,
+        overlap=profile_overlap,
+        advance_records=PROFILE_MEMORY_LIMIT + profile_overlap,
     )
     if not dry_run:
         store.set_maintenance_cursor(

eval/consolidation_ranking.py

@@ -11,12 +11,14 @@
 from __future__ import annotations
 
 import json
+import math
 from pathlib import Path
 from typing import Any
 
 from engraphis.core.consolidate import consolidate
 from engraphis.core.engine import MemoryEngine
 from engraphis.core.interfaces import MemoryType, SearchFilter
+from engraphis.core.recall import CONSOLIDATION_BONUS
 
 
 DATASET = Path(__file__).with_name("datasets") / "consolidation_ranking.jsonl"
@@ -40,6 +42,24 @@ def load_cases(path: Path = DATASET) -> list[dict[str, Any]]:
         context = case.get("context", [])
         if not isinstance(context, list):
             raise ValueError(f"case {case['id']} context must be a list")
+        probes = case.get("bonus_probes", [])
+        if not isinstance(probes, list):
+            raise ValueError(f"case {case['id']} bonus_probes must be a list")
+        for probe in probes:
+            if (
+                not isinstance(probe, dict)
+                or not isinstance(probe.get("id"), str)
+                or probe.get("expected") not in {"source", "digest"}
+            ):
+                raise ValueError(f"case {case['id']} has an invalid bonus probe")
+            for field in ("source_score", "digest_score"):
+                value = probe.get(field)
+                if (
+                    not isinstance(value, (int, float))
+                    or isinstance(value, bool)
+                    or not math.isfinite(float(value))
+                ):
+                    raise ValueError(f"case {case['id']} has an invalid probe {field}")
         expected = case.get("expected")
         if not isinstance(expected, (str, dict)):
             raise ValueError(f"case {case['id']} needs an expected ranking target")
@@ -100,6 +120,7 @@ def evaluate_case(case: dict[str, Any]) -> dict[str, Any]:
         context_ids[tag] = engine.remember(
             str(item["text"]), workspace_id=workspace_id, repo_id=repo_id,
             mtype=_memory_type(item.get("mtype"), MemoryType.SEMANTIC),
+            importance=float(item.get("importance", 0.0)),
             resolve_conflicts=False,
         )
     report = consolidate(engine, workspace_id=workspace_id, repo_id=repo_id)
@@ -163,11 +184,42 @@ def rank(ids: list[str], target: str) -> int | None:
     }
 
 
+def _evaluate_bonus_probe(probe: dict[str, Any]) -> dict[str, Any]:
+    """Exercise the exact post-normalization bonus at a controlled boundary."""
+    baseline_scores = {
+        "source": float(probe["source_score"]),
+        "digest": float(probe["digest_score"]),
+    }
+    policy_scores = dict(baseline_scores)
+    policy_scores["digest"] += CONSOLIDATION_BONUS
+    baseline = sorted(baseline_scores, key=lambda kind: (-baseline_scores[kind], kind))
+    policy = sorted(policy_scores, key=lambda kind: (-policy_scores[kind], kind))
+    expected = str(probe["expected"])
+    return {
+        "id": probe["id"],
+        "expected": expected,
+        "baseline_top": baseline[0],
+        "policy_top": policy[0],
+        "baseline_hit_at_1": baseline[0] == expected,
+        "policy_hit_at_1": policy[0] == expected,
+        "ranking_changed": baseline != policy,
+    }
+
+
 def evaluate(path: Path = DATASET) -> dict[str, Any]:
     """Return ranking preference and raw-evidence retention metrics."""
-    results = [evaluate_case(case) for case in load_cases(path)]
+    cases = load_cases(path)
+    results = [evaluate_case(case) for case in cases]
+    bonus_probes = [
+        _evaluate_bonus_probe(probe)
+        for case in cases for probe in case.get("bonus_probes", [])
+    ]
     summary = [item for item in results if item["expected_role"] == "digest"]
     details = [item for item in results if item["expected_role"] == "raw"]
+    source_regressions = [
+        item["id"] for item in bonus_probes
+        if item["expected"] == "source" and not item["policy_hit_at_1"]
+    ]
     return {
         "cases": len(results),
         "summary_digest_top1_rate": (
@@ -178,7 +230,31 @@ def evaluate(path: Path = DATASET) -> dict[str, Any]:
             sum(item["baseline_top"] == item["expected_id"] for item in summary)
             / len(summary)
         ),
-        "ranking_changed_rate": sum(item["ranking_changed"] for item in results) / len(results),
+        "production_trace_ranking_changed_rate": (
+            sum(item["ranking_changed"] for item in results) / len(results)
+        ),
+        "bonus_probe_count": len(bonus_probes),
+        "ranking_changed_rate": (
+            sum(item["ranking_changed"] for item in bonus_probes) / len(bonus_probes)
+            if bonus_probes else 0.0
+        ),
+        "bonus_probe_digest_top1_rate": (
+            sum(item["policy_hit_at_1"] for item in bonus_probes
+                if item["expected"] == "digest")
+            / max(1, sum(item["expected"] == "digest" for item in bonus_probes))
+        ),
+        "baseline_bonus_probe_digest_top1_rate": (
+            sum(item["baseline_hit_at_1"] for item in bonus_probes
+                if item["expected"] == "digest")
+            / max(1, sum(item["expected"] == "digest" for item in bonus_probes))
+        ),
+        "bonus_probe_source_top1_rate": (
+            sum(item["policy_hit_at_1"] for item in bonus_probes
+                if item["expected"] == "source")
+            / max(1, sum(item["expected"] == "source" for item in bonus_probes))
+        ),
+        "bonus_probe_source_regressions": source_regressions,
+        "bonus_probes": bonus_probes,
         "expected_hit_at_k": sum(item["expected_hit_at_k"] for item in results) / len(results),
         "raw_detail_hit_at_k": (
             sum(item["expected_hit_at_k"] for item in details) / len(details)

eval/datasets/consolidation_ranking.jsonl

@@ -1,3 +1,3 @@
 # Consolidated digest ranking: summary preference and raw-detail retention.
-{"id":"summary_digest","query":"flaky network integration build failure","expected":"digest","k":5,"cluster":[{"text":"Build failed on the flaky network integration test in CI run 101."},{"text":"Build failed on the flaky network integration test in CI run 202."},{"text":"Build failed on the flaky network integration test in CI run 303."}],"context":[{"tag":"unrelated","text":"The office kitchen orders sourdough every Friday.","mtype":"semantic"}]}
+{"id":"summary_digest","query":"flaky network integration build failure","expected":"digest","k":5,"cluster":[{"text":"Build failed on the flaky network integration test in CI run 101."},{"text":"Build failed on the flaky network integration test in CI run 202."},{"text":"Build failed on the flaky network integration test in CI run 303."}],"context":[{"tag":"unrelated","text":"build failure","mtype":"working","importance":0.5}],"bonus_probes":[{"id":"near-boundary-digest","source_score":0.72,"digest_score":0.70,"expected":"digest"},{"id":"source-margin","source_score":0.80,"digest_score":0.70,"expected":"source"}]}
 {"id":"specific_raw_evidence","query":"What rollback procedure follows a failed API canary?","expected":{"tag":"rollback","role":"raw"},"k":5,"cluster":[{"text":"The API deployment failed a canary health check during the morning release."},{"text":"The API deployment failed a canary health check during the afternoon release."},{"text":"The API deployment failed a canary health check during the evening release."}],"context":[{"tag":"rollback","text":"The exact rollback command is kubectl rollout undo deployment/api after a canary failure.","mtype":"procedural"}]}