diff --git a/src/everos/memory/reflection/orchestrator.py b/src/everos/memory/reflection/orchestrator.py index 58b04d50c..c46653afd 100644 --- a/src/everos/memory/reflection/orchestrator.py +++ b/src/everos/memory/reflection/orchestrator.py @@ -10,7 +10,6 @@ from __future__ import annotations -import asyncio import datetime as _dt import json import uuid @@ -32,13 +31,25 @@ from everos.core.observability.tracing import memory_span from everos.core.persistence import MemoryRoot from everos.infra.ome.context import StrategyContext -from everos.infra.persistence.index import all_of, eq, is_null +from everos.infra.persistence.index import all_of, eq, is_null, one_of from everos.memory._partition_locks import get_partition_lock from everos.memory.events import EpisodeExtracted logger = get_logger(__name__) _MAX_CLUSTERS_PER_RUN = 10 + +# Upper bound on the episodes one merge consumes. A cluster with more members +# is merged incrementally: its existing merged episode plus the oldest sources +# go into this run, the remaining members stay in the cluster and are folded +# in by later runs in update mode. Keeps the merged narrative, and the +# cascade upsert it triggers, inside the LanceDB write-lock budget. +_MAX_SOURCES_PER_MERGE = 50 + +# Entry ids per LanceDB ``update`` when deprecating sources. One update per +# batch instead of one per row keeps the number of write-lock acquisitions +# small for large clusters. +_DEPRECATE_BATCH_SIZE = 100 _WAIT_TIMEOUT_SECONDS = 120.0 @@ -354,7 +365,52 @@ async def _load_cluster_episodes( app_id=app_id, project_id=project_id, ) - return members, episodes + return self._cap_cluster_sources( + cluster_id=cluster_id, members=members, episodes=episodes + ) + + @staticmethod + def _cap_cluster_sources( + *, + cluster_id: str, + members: list[tuple[str, str]], + episodes: list[Any], + ) -> tuple[list[tuple[str, str]], list[Any]]: + """Limit one merge to ``_MAX_SOURCES_PER_MERGE`` episodes. + + The cluster's existing merged episodes (``parent_type == "cluster"``) + are always kept so update mode still sees the previous narrative; + the remaining budget goes to the oldest source episodes. Members + whose episode is deferred are dropped from this run's snapshot so + they are neither merged nor deprecated now. + + Args: + cluster_id: Target cluster identifier (for logging). + members: Cluster members ``(member_id, member_type)``. + episodes: Episode rows sorted by timestamp ascending. + + Returns: + ``(members, episodes)`` restricted to this run. + """ + if len(episodes) <= _MAX_SOURCES_PER_MERGE: + return members, episodes + + merged = [e for e in episodes if e.parent_type == "cluster"] + sources = [e for e in episodes if e.parent_type != "cluster"] + budget = max(_MAX_SOURCES_PER_MERGE - len(merged), 1) + kept = merged + sources[:budget] + kept.sort(key=lambda e: e.timestamp) + deferred = sources[budget:] + deferred_keys = {e.entry_id for e in deferred} | {e.parent_id for e in deferred} + kept_members = [(mid, mt) for mid, mt in members if mid not in deferred_keys] + logger.info( + "reflection_cluster_capped", + cluster_id=cluster_id, + member_count=len(members), + merged_count=len(kept_members), + deferred_count=len(members) - len(kept_members), + ) + return kept_members, kept async def _write_and_reextract( self, @@ -732,7 +788,7 @@ async def _execute_deprecation( Returns: A ReflectionReport on success, ``None`` when no members to deprecate. """ - to_deprecate = await self._resolve_deprecation_targets( + to_deprecate, remaining = await self._resolve_deprecation_targets( cluster_id=cluster_id, original_members=original_members, ) @@ -753,6 +809,7 @@ async def _execute_deprecation( merged_entry_id=merged_entry_id, algo_result=algo_result, episodes=episodes, + member_count=remaining + 1, ) report = await self._create_reflection_report( cluster_id=cluster_id, @@ -783,7 +840,14 @@ async def _apply_deprecation_writes( project_id: str, merged_entry_id: str, ) -> tuple[int, int]: - """Patch md frontmatter and mark episodes/facts deprecated in LanceDB. + """Mark episodes/facts deprecated in LanceDB, then patch md frontmatter. + + The LanceDB writes go first and the markdown record last, so the + durable record only lists deprecations that were actually applied. + If any write fails, the ``deprecated_by`` values this run already + wrote are reverted before the error propagates, so a failed run + leaves no sources half-deprecated toward a merge that was never + committed to the cluster. Args: episodes: Source episode rows (for md patching). @@ -796,31 +860,90 @@ async def _apply_deprecation_writes( Returns: ``(deprecated_episode_count, deprecated_fact_count)``. """ - await self._patch_md_frontmatter( - episodes=episodes, - to_deprecate=to_deprecate, - merged_entry_id=merged_entry_id, - ) - deprecated_ep_count = await self._deprecate_lance_episodes( - entry_ids=to_deprecate, - owner_id=owner_id, - app_id=app_id, - project_id=project_id, - merged_entry_id=merged_entry_id, - ) - deprecated_fact_count = await self._deprecate_lance_facts( - parent_ids=to_deprecate, - owner_id=owner_id, + try: + deprecated_ep_count = await self._deprecate_lance_episodes( + entry_ids=to_deprecate, + owner_id=owner_id, + app_id=app_id, + project_id=project_id, + merged_entry_id=merged_entry_id, + ) + deprecated_fact_count = await self._deprecate_lance_facts( + parent_ids=to_deprecate, + owner_id=owner_id, + merged_entry_id=merged_entry_id, + ) + await self._patch_md_frontmatter( + episodes=episodes, + to_deprecate=to_deprecate, + merged_entry_id=merged_entry_id, + ) + except BaseException: + await self._revert_lance_deprecation( + entry_ids=to_deprecate, + owner_id=owner_id, + app_id=app_id, + project_id=project_id, + merged_entry_id=merged_entry_id, + ) + raise + return deprecated_ep_count, deprecated_fact_count + + async def _revert_lance_deprecation( + self, + *, + entry_ids: set[str], + owner_id: str, + app_id: str, + project_id: str, + merged_entry_id: str, + ) -> None: + """Best-effort compensation: clear ``deprecated_by`` written by this run. + + Only rows that point at ``merged_entry_id`` are touched, so + deprecations from earlier, committed merges are left alone. + Failures are logged, never raised: the original error is the one + the caller needs to see. + """ + reverted = True + for store, field_name, label in ( + (self._episode_store, "entry_id", "episode"), + (self._atomic_fact_store, "parent_id", "atomic_fact"), + ): + for batch in _batched(sorted(entry_ids), _DEPRECATE_BATCH_SIZE): + where = all_of( + one_of(field_name, batch), + eq("owner_id", owner_id), + eq("deprecated_by", merged_entry_id), + ) + if label == "episode": + where = all_of( + where, eq("app_id", app_id), eq("project_id", project_id) + ) + try: + await store.update({"deprecated_by": None}, where=where) + except Exception: + reverted = False + logger.error( + "reflection_deprecate_revert_failed", + table=label, + merged_entry_id=merged_entry_id, + batch_size=len(batch), + exc_info=True, + ) + logger.warning( + "reflection_deprecate_reverted", merged_entry_id=merged_entry_id, + source_count=len(entry_ids), + complete=reverted, ) - return deprecated_ep_count, deprecated_fact_count async def _resolve_deprecation_targets( self, *, cluster_id: str, original_members: list[tuple[str, str]], - ) -> set[str]: + ) -> tuple[set[str], int]: """Re-read cluster members and intersect with the original snapshot. Args: @@ -828,12 +951,16 @@ async def _resolve_deprecation_targets( original_members: Snapshot ``(member_id, member_type)`` from selection. Returns: - Set of member IDs safe to deprecate (present in both snapshots). + ``(to_deprecate, remaining)``: the member IDs safe to deprecate + (present in both snapshots) and the number of current members + that stay in the cluster (deferred by the per-merge cap, or + added after the snapshot). """ current_members = await self._cluster_repo.get_members_with_type(cluster_id) current_ids = {mid for mid, _ in current_members} original_ids = {mid for mid, _ in original_members} - return original_ids & current_ids + to_deprecate = original_ids & current_ids + return to_deprecate, len(current_ids - to_deprecate) async def _deprecate_lance_episodes( self, @@ -846,24 +973,25 @@ async def _deprecate_lance_episodes( ) -> int: """Mark deprecated episodes in LanceDB by entry_id. + One ``update`` per batch of ``_DEPRECATE_BATCH_SIZE`` ids, issued + sequentially: every call takes the table's write lock, so N + concurrent single-row updates queue behind each other and, on a + large cluster, blow the lock deadline. + Returns: - Number of LanceDB update calls issued. + Number of episodes targeted. """ - coros: list[Any] = [ - self._episode_store.update( + for batch in _batched(sorted(entry_ids), _DEPRECATE_BATCH_SIZE): + await self._episode_store.update( {"deprecated_by": merged_entry_id}, where=all_of( - eq("entry_id", eid), + one_of("entry_id", batch), eq("owner_id", owner_id), eq("app_id", app_id), eq("project_id", project_id), ), ) - for eid in entry_ids - ] - if coros: - await asyncio.gather(*coros) - return len(coros) + return len(entry_ids) async def _deprecate_lance_facts( self, @@ -880,24 +1008,21 @@ async def _deprecate_lance_facts( merged_entry_id: Entry ID of the replacement merged episode. Returns: - Total number of LanceDB update calls issued. + Number of parents whose facts were targeted. """ if not parent_ids: return 0 - coros = [ - self._atomic_fact_store.update( + for batch in _batched(sorted(parent_ids), _DEPRECATE_BATCH_SIZE): + await self._atomic_fact_store.update( {"deprecated_by": merged_entry_id}, where=all_of( - eq("parent_id", pid), + one_of("parent_id", batch), eq("owner_id", owner_id), is_null("deprecated_by"), ), ) - for pid in parent_ids - ] - await asyncio.gather(*coros) - return len(coros) + return len(parent_ids) async def _update_cluster_after_merge( self, @@ -907,6 +1032,7 @@ async def _update_cluster_after_merge( merged_entry_id: str, algo_result: AlgoEpisode, episodes: list[Any], + member_count: int = 1, ) -> None: """Remove old members, add merged, and recompute centroid. @@ -916,6 +1042,8 @@ async def _update_cluster_after_merge( merged_entry_id: Entry ID of the newly merged episode. algo_result: Algo reflector output (episode text for centroid). episodes: Source episode rows (for last timestamp). + member_count: Members left in the cluster after the merge + (the merged episode plus any deferred or newly added ones). """ await self._cluster_repo.remove_members(cluster_id, to_deprecate) await self._cluster_repo.add_member(cluster_id, merged_entry_id, "episode") @@ -926,7 +1054,7 @@ async def _update_cluster_after_merge( await self._cluster_repo.update_metadata( cluster_id, centroid_blob=centroid_blob, - count=1, + count=member_count, last_ts_ms=last_ts_ms, preview_json=json.dumps([algo_result.episode[:200]], ensure_ascii=False), ) @@ -1017,6 +1145,11 @@ async def _patch_md_frontmatter( ) +def _batched(items: list[str], size: int) -> list[list[str]]: + """Split ``items`` into consecutive lists of at most ``size`` elements.""" + return [items[i : i + size] for i in range(0, len(items), size)] + + def _to_algo_episodes(episodes: list[Any]) -> list[AlgoEpisode]: """Convert LanceDB episode rows to algo Episode objects. diff --git a/tests/unit/test_memory/test_reflection/test_orchestrator.py b/tests/unit/test_memory/test_reflection/test_orchestrator.py index a413e4a0c..76bdf2620 100644 --- a/tests/unit/test_memory/test_reflection/test_orchestrator.py +++ b/tests/unit/test_memory/test_reflection/test_orchestrator.py @@ -16,8 +16,11 @@ import pytest +from everos.core.errors import VectorStoreBusyError from everos.infra.ome.testing import FakeStrategyContext +from everos.infra.persistence.predicate import All, Comparison, In, Predicate from everos.memory._partition_locks import _reset_for_tests +from everos.memory.reflection import orchestrator as orchestrator_module from everos.memory.reflection.orchestrator import ( _MAX_CLUSTERS_PER_RUN, ReflectionOrchestrator, @@ -483,3 +486,180 @@ async def test_call_reflector_emits_consolidate_generation_span() -> None: spans["everos.reflect.consolidate"].attributes["langfuse.observation.type"] == "generation" ) + + +# ── Batched deprecation, compensation and per-merge cap ─────────────────── + + +def _in_values(predicate: Predicate, field_name: str) -> list[str]: + """Collect the values of every ``IN`` clause on ``field_name``.""" + if isinstance(predicate, In): + return list(predicate.values) if predicate.field == field_name else [] + if isinstance(predicate, All): + return [ + v for child in predicate.children for v in _in_values(child, field_name) + ] + return [] + + +def _eq_value(predicate: Predicate, field_name: str) -> object | None: + if isinstance(predicate, Comparison): + return predicate.value if predicate.field == field_name else None + if isinstance(predicate, All): + for child in predicate.children: + found = _eq_value(child, field_name) + if found is not None: + return found + return None + + +async def test_deprecate_lance_episodes_batches_updates(monkeypatch) -> None: + """One update per batch of ids instead of one update per row.""" + monkeypatch.setattr(orchestrator_module, "_DEPRECATE_BATCH_SIZE", 100) + episode_store = MagicMock() + episode_store.update = AsyncMock() + orch = _build_orchestrator(episode_store=episode_store) + entry_ids = {f"ep_{i:04d}" for i in range(250)} + + count = await orch._deprecate_lance_episodes( + entry_ids=entry_ids, + owner_id="u_alice", + app_id="default", + project_id="default", + merged_entry_id="ep_merged", + ) + + assert count == 250 + assert episode_store.update.await_count == 3 + seen: list[str] = [] + for call in episode_store.update.await_args_list: + updates, where = call.args[0], call.kwargs["where"] + assert updates == {"deprecated_by": "ep_merged"} + batch = _in_values(where, "entry_id") + assert 0 < len(batch) <= 100 + assert _eq_value(where, "owner_id") == "u_alice" + seen.extend(batch) + assert sorted(seen) == sorted(entry_ids) + + +async def test_deprecate_failure_reverts_applied_writes(monkeypatch) -> None: + """A failed LanceDB write reverts what this run already wrote and leaves + the markdown record untouched, then propagates the error.""" + monkeypatch.setattr(orchestrator_module, "_DEPRECATE_BATCH_SIZE", 100) + episode_store = MagicMock() + episode_store.update = AsyncMock( + side_effect=[None, VectorStoreBusyError("write lock deadline"), None, None] + ) + atomic_fact_store = MagicMock() + atomic_fact_store.update = AsyncMock() + episode_writer = MagicMock() + episode_writer.patch_frontmatter = AsyncMock() + orch = _build_orchestrator( + episode_store=episode_store, + atomic_fact_store=atomic_fact_store, + episode_writer=episode_writer, + ) + entry_ids = {f"ep_{i:04d}" for i in range(150)} + episodes = [_make_episode_row(entry_id=eid) for eid in sorted(entry_ids)] + + with pytest.raises(VectorStoreBusyError): + await orch._apply_deprecation_writes( + episodes=episodes, + to_deprecate=entry_ids, + owner_id="u_alice", + app_id="default", + project_id="default", + merged_entry_id="ep_merged", + ) + + # The markdown record was never patched. + episode_writer.patch_frontmatter.assert_not_awaited() + # Two forward batches (second failed) + two revert batches on episodes. + assert episode_store.update.await_count == 4 + reverts = episode_store.update.await_args_list[2:] + reverted: list[str] = [] + for call in reverts: + updates, where = call.args[0], call.kwargs["where"] + assert updates == {"deprecated_by": None} + assert _eq_value(where, "deprecated_by") == "ep_merged" + reverted.extend(_in_values(where, "entry_id")) + assert sorted(reverted) == sorted(entry_ids) + # Facts were never deprecated, only (harmlessly) reverted. + for call in atomic_fact_store.update.await_args_list: + assert call.args[0] == {"deprecated_by": None} + + +async def test_run_caps_sources_per_merge_and_keeps_the_rest(monkeypatch) -> None: + """A cluster above the cap merges only its oldest sources; the deferred + members stay in the cluster and the count reflects them.""" + monkeypatch.setattr(orchestrator_module, "_MAX_SOURCES_PER_MERGE", 2) + cluster_repo = MagicMock() + episode_store = MagicMock() + atomic_fact_store = MagicMock() + episode_writer = MagicMock() + report_repo = MagicMock() + reflector = MagicMock() + embedder = MagicMock() + + report_repo.list_reflected_cluster_ids = AsyncMock(return_value=set()) + cluster_repo.list_ids_and_member_counts = AsyncMock(return_value=[("cl_abc", 3)]) + episode_store.find_where = AsyncMock(return_value=[]) + members = [ + ("ep_20260601_0001", "episode"), + ("ep_20260601_0002", "episode"), + ("ep_20260601_0003", "episode"), + ] + cluster_repo.get_members_with_type = AsyncMock(return_value=members) + rows = [ + _make_episode_row( + entry_id=f"ep_20260601_000{i}", + parent_id=f"mc_00{i}", + timestamp=_dt.datetime(2026, 6, i, tzinfo=_dt.UTC), + ) + for i in (3, 1, 2) + ] + episode_store.find_by_owner_entries = AsyncMock(return_value=rows) + reflector.areflect = AsyncMock( + return_value=_FakeAlgoResult( + owner_id=None, episode="merged", subject="s", timestamp=1717200000000 + ) + ) + episode_writer.append_entries = AsyncMock( + return_value=[_make_entry_id("ep_20260614_0001")] + ) + episode_writer.patch_frontmatter = AsyncMock() + cluster_repo.remove_members = AsyncMock() + cluster_repo.add_member = AsyncMock() + cluster_repo.update_metadata = AsyncMock() + embedder.embed = AsyncMock(return_value=[0.1] * 1024) + episode_store.update = AsyncMock() + atomic_fact_store.update = AsyncMock() + report_repo.create = AsyncMock() + + orch = _build_orchestrator( + cluster_repo=cluster_repo, + episode_store=episode_store, + atomic_fact_store=atomic_fact_store, + episode_writer=episode_writer, + report_repo=report_repo, + reflector=reflector, + embedder=embedder, + ) + reports = await orch.run(ctx=FakeStrategyContext(), owner_id="u_alice") + + # Only the two oldest sources were reflected. + reflected = reflector.areflect.await_args.args[0] + assert [ep.timestamp for ep in reflected] == [ + _ts_to_ms(_dt.datetime(2026, 6, 1, tzinfo=_dt.UTC)), + _ts_to_ms(_dt.datetime(2026, 6, 2, tzinfo=_dt.UTC)), + ] + # ... and only those were deprecated and removed from the cluster. + cluster_repo.remove_members.assert_awaited_once_with( + "cl_abc", {"ep_20260601_0001", "ep_20260601_0002"} + ) + deprecated = _in_values(episode_store.update.await_args.kwargs["where"], "entry_id") + assert sorted(deprecated) == ["ep_20260601_0001", "ep_20260601_0002"] + # The deferred member plus the merged episode remain. + assert cluster_repo.update_metadata.await_args.kwargs["count"] == 2 + assert len(reports) == 1 + assert reports[0].source_count == 2