diff --git a/reme/config/default.yaml b/reme/config/default.yaml index c57c0b3c..7159c0b6 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -72,6 +72,7 @@ jobs: file_catalog: dream scan_days: 2 max_units: 5 + max_chars_per_batch: 40000 - backend: dream_integrate_step - backend: dream_finish_step file_catalog: dream @@ -157,11 +158,16 @@ jobs: type: integer description: "maximum number of extracted memory units" default: 5 + max_chars_per_batch: + type: integer + description: "maximum characters of changed material sent per extraction call" + default: 40000 steps: - backend: dream_extract_step file_catalog: dream scan_days: 2 max_units: 5 + max_chars_per_batch: 40000 - backend: dream_integrate_step - backend: dream_finish_step file_catalog: dream diff --git a/reme/steps/evolve/dream/extract.py b/reme/steps/evolve/dream/extract.py index d3d4c3b7..be27c61c 100644 --- a/reme/steps/evolve/dream/extract.py +++ b/reme/steps/evolve/dream/extract.py @@ -1,6 +1,7 @@ """Global dream extract step.""" import json +from pathlib import Path from ...base_step import BaseStep from ...file_io import refresh_day_index @@ -9,6 +10,7 @@ from ....components import R from ....enumeration import DreamBucketEnum from ....schema import DreamExtractOutput, DreamState from .utils import ( + batch_paths, clean_paths, daily_dir, llm_available, @@ -24,14 +26,55 @@ from .utils import ( _TOOLS = ("read",) +class _UnusableReceipt(RuntimeError): + """The model returned an unusable receipt even after one bounded retry. + + Bisecting the batch cannot help here: the input was delivered fine, the + model just failed to answer in the expected shape. Callers mark the batch + as failed instead of splitting it. + """ + + +def merge_units(units: list[dict]) -> list[dict]: + """Merge same-name units produced by different batches. + + Batched extraction can surface the same abstraction twice; the integration + step writes one digest node per unit, so duplicates would double-write. + Paths are unioned and the richer summary wins. + """ + merged: dict[str, dict] = {} + order: list[str] = [] + for unit in units: + name = str(unit.get("name") or "").strip() + key = name.lower() + if not key: + continue + if key not in merged: + merged[key] = dict(unit) + order.append(key) + continue + current = merged[key] + current["paths"] = sorted({*current.get("paths", []), *unit.get("paths", [])}) + if len(str(unit.get("summary") or "")) > len(str(current.get("summary") or "")): + current["summary"] = unit.get("summary") + return [merged[key] for key in order] + + @R.register("dream_extract_step") class DreamExtractStep(BaseStep): """Scan changed daily files and globally extract merged memory units.""" - def __init__(self, scan_days: int = 2, max_units: int = 5, **kwargs): + def __init__( + self, + scan_days: int = 2, + max_units: int = 5, + max_chars_per_batch: int = 40000, + **kwargs, + ): super().__init__(**kwargs) self.scan_days = scan_days self.max_units = max_units + self.max_chars_per_batch = max_chars_per_batch async def execute(self): assert self.context is not None @@ -113,57 +156,163 @@ class DreamExtractStep(BaseStep): return self._finish(state, False, state.errors[-1]) self.logger.info(f"[{self.name}] agent start changed={len(changed)} dates={len(dates)}") + output_schema = DreamExtractOutput.model_json_schema() + output_schema["required"] = ["units"] + system_prompt = self.prompt_format( + "extract_system_prompt", + workspace_dir=str(workspace), + buckets=", ".join(bucket.value for bucket in DreamBucketEnum), + max_units=max_units, + output_schema_json=json.dumps(output_schema, ensure_ascii=False, indent=2), + ) + raw_max_chars = self.context.get("max_chars_per_batch", self.max_chars_per_batch) + max_chars = max(int(raw_max_chars or self.max_chars_per_batch), 0) + batches = batch_paths(workspace, changed, max_chars=max_chars) + self.logger.info( + f"[{self.name}] batching changed={len(changed)} batches={len(batches)} " f"max_chars_per_batch={max_chars}", + ) + for batch in batches: + remaining = max_units - len(state.units) + if remaining <= 0: + self.logger.info( + f"[{self.name}] stop batching max_units reached units={len(state.units)}", + ) + break + await self._extract_batch( + state, + batch, + workspace=workspace, + day=day, + dates=dates, + hint=hint, + max_units=remaining, + system_prompt=system_prompt, + ) + state.units = merge_units(state.units) + + self.logger.info( + f"[{self.name}] parse done units={len(state.units)} batches={len(batches)} " + f"failed_paths={len(state.failed_paths)} errors={len(state.errors)}", + ) + answer = f"Extracted {len(state.units)} unit(s) from {len(changed)} changed file(s) across {len(dates)} day(s)" + return self._finish(state, not state.errors, answer) + + async def _extract_batch( + self, + state: DreamState, + batch: list[str], + *, + workspace: Path, + day: str, + dates: list[str], + hint: str, + max_units: int, + system_prompt: str, + ) -> None: + """Extract one size-capped batch, bisecting a batch that cannot be handled. + + A transport failure (oversized prompt, backend crash) is never retried + with a *larger* message: the batch is split in half instead, so the + offending input is isolated. A single file that still fails is recorded + in ``failed_paths`` and the remaining batches keep going. + """ + try: + await self._extract_once( + state, + batch, + workspace=workspace, + day=day, + dates=dates, + hint=hint, + max_units=max_units, + system_prompt=system_prompt, + ) + except _UnusableReceipt as e: + error = f"dream extract agent failed after retry: {e}" + state.errors.append(error) + state.failed_paths.extend(batch) + self.logger.error(f"[{self.name}] {error}") + except Exception as e: # noqa: BLE001 + if len(batch) > 1: + middle = len(batch) // 2 + self.logger.warning( + f"[{self.name}] batch failed files={len(batch)}; splitting to isolate the " + f"offending input: {type(e).__name__}: {e}", + ) + for half in (batch[:middle], batch[middle:]): + await self._extract_batch( + state, + half, + workspace=workspace, + day=day, + dates=dates, + hint=hint, + max_units=max(max_units - len(state.units), 0), + system_prompt=system_prompt, + ) + return + target = batch[0] if batch else "(empty batch)" + error = f"dream extract agent failed for {target}: {type(e).__name__}: {e}" + state.errors.append(error) + state.failed_paths.extend(batch) + self.logger.error(f"[{self.name}] {error}") + + async def _extract_once( + self, + state: DreamState, + batch: list[str], + *, + workspace: Path, + day: str, + dates: list[str], + hint: str, + max_units: int, + system_prompt: str, + ) -> None: + """Run one LLM extraction for a single batch and merge its units.""" user_message = self.prompt_format( "extract_user_message", date=day, dates_json=json.dumps(dates, ensure_ascii=False, indent=2), hint=hint or "(none)", max_units=max_units, - changed_paths_json=json.dumps(changed, ensure_ascii=False, indent=2), - material_blob=pack_paths(workspace, changed), + changed_paths_json=json.dumps(batch, ensure_ascii=False, indent=2), + material_blob=pack_paths(workspace, batch), ) - retry_feedback = "" - output_schema = DreamExtractOutput.model_json_schema() - output_schema["required"] = ["units"] - for attempt in range(2): - try: - result = await self.agent_wrapper.reply( - user_message + retry_feedback, - system_prompt=self.prompt_format( - "extract_system_prompt", - workspace_dir=str(workspace), - buckets=", ".join(bucket.value for bucket in DreamBucketEnum), - max_units=max_units, - output_schema_json=json.dumps(output_schema, ensure_ascii=False, indent=2), - ), - job_tools=list(_TOOLS), - ) - self.logger.info(f"[{self.name}] agent done has_result={bool(result.get('result'))}") - raw_result = agent_reply_result_text(result) - state.extract_summary = raw_result - meta = parse_structured_reply(raw_result, strict=True) - self.clean_output(state, meta, max_units=max_units) - break - except Exception as e: # noqa: BLE001 - if attempt == 0: - retry_feedback = "\n\n" + self.prompt_format( - "extract_retry_message", - error=f"{type(e).__name__}: {e}", - ) - self.logger.warning( - f"[{self.name}] extract attempt 1 returned no usable receipt; retrying once: " - f"{type(e).__name__}: {e}", - ) - continue - error = f"dream extract agent failed after retry: {type(e).__name__}: {e}" - state.errors.append(error) - state.failed_paths = list(changed) - self.logger.error(f"[{self.name}] {error}") - return self._finish(state, False, error) - - self.logger.info(f"[{self.name}] parse done keys={','.join(sorted(meta.keys())) if meta else '(none)'}") - answer = f"Extracted {len(state.units)} unit(s) from {len(changed)} changed file(s) across {len(dates)} day(s)" - return self._finish(state, True, answer) + allowed = set(batch) + result = await self.agent_wrapper.reply( + user_message, + system_prompt=system_prompt, + job_tools=list(_TOOLS), + ) + self.logger.info(f"[{self.name}] agent done files={len(batch)} has_result={bool(result.get('result'))}") + raw_result = agent_reply_result_text(result) + state.extract_summary = "\n".join(filter(None, (state.extract_summary, raw_result))) + try: + meta = parse_structured_reply(raw_result, strict=True) + self.clean_output(state, meta, max_units=max_units, allowed=allowed) + return + except Exception as e: # noqa: BLE001 + retry_feedback = "\n\n" + self.prompt_format( + "extract_retry_message", + error=f"{type(e).__name__}: {e}", + ) + self.logger.warning( + f"[{self.name}] extract attempt 1 returned no usable receipt; retrying once: " + f"{type(e).__name__}: {e}", + ) + result = await self.agent_wrapper.reply( + user_message + retry_feedback, + system_prompt=system_prompt, + job_tools=list(_TOOLS), + ) + raw_result = agent_reply_result_text(result) + state.extract_summary = "\n".join(filter(None, (state.extract_summary, raw_result))) + try: + meta = parse_structured_reply(raw_result, strict=True) + self.clean_output(state, meta, max_units=max_units, allowed=allowed) + except Exception as e: # noqa: BLE001 + raise _UnusableReceipt(f"{type(e).__name__}: {e}") from e def _existing(self, workspace, files: list[str]) -> dict[str, float]: out: dict[str, float] = {} @@ -174,8 +323,19 @@ class DreamExtractStep(BaseStep): self.logger.error(f"[{self.name}] stat failed on {rel}: {e}") return out - def clean_output(self, state: DreamState, meta: dict, max_units: int | None = None) -> None: - """Validate the whole receipt before accepting any units or applying the cap.""" + def clean_output( + self, + state: DreamState, + meta: dict, + max_units: int | None = None, + allowed: set[str] | None = None, + ) -> None: + """Validate the whole receipt before accepting any units or applying the cap. + + ``allowed`` scopes the accepted paths to the current batch; by default + every changed path stays valid, which keeps single-batch behaviour + identical to the unbatched implementation. + """ units = meta.get("units") if "units" in meta else meta.get("memory_units") if not isinstance(units, list): raise ValueError("expected a units list (use units: [] only when nothing is worth remembering)") @@ -194,11 +354,11 @@ class DreamExtractStep(BaseStep): raw = {**raw, "bucket": DreamBucketEnum.WIKI.value} normalized.append(raw) output = DreamExtractOutput.model_validate({"units": normalized}) - allowed = set(state.changed_paths) + allowed_paths = set(state.changed_paths) if allowed is None else set(allowed) cleaned = [] for index, unit in enumerate(output.units): name, summary = unit.name.strip(), unit.summary.strip() - paths = clean_paths(unit.paths, allowed) + paths = clean_paths(unit.paths, allowed_paths) if not name or not summary or not paths: raise ValueError(f"units[{index}] requires non-empty name, summary, and paths from changed_paths") cleaned.append({"name": name, "bucket": unit.bucket.value, "summary": summary, "paths": paths}) diff --git a/reme/steps/evolve/dream/utils.py b/reme/steps/evolve/dream/utils.py index bce5c135..a465a4eb 100644 --- a/reme/steps/evolve/dream/utils.py +++ b/reme/steps/evolve/dream/utils.py @@ -113,6 +113,72 @@ def pack_paths( return "\n".join(blocks) +def _file_size(workspace: Path, rel: str) -> int: + """Best-effort byte size of one workspace-relative path.""" + try: + return (workspace / rel).stat().st_size + except OSError: + return 0 + + +def _day_key(rel: str) -> str: + """Group key that keeps one scanned day together (``daily/``).""" + parts = rel.split("/") + return "/".join(parts[:2]) if len(parts) >= 2 else rel + + +def batch_paths( + workspace: Path, + paths: list[str], + *, + max_chars: int, +) -> list[list[str]]: + """Split paths into size-capped batches without dropping any path. + + Files belonging to the same scanned day are kept in one batch whenever the + day fits the budget, so cross-file extraction still sees a whole day. A day + larger than the budget is split by file size, and a single oversized file + keeps a batch of its own (``pack_paths`` still applies its per-file cap). + + Every path lands in exactly one batch; callers must never mark an unseen + path as extracted, which is why this returns a partition and not a slice. + """ + if not paths: + return [] + if max_chars <= 0: + return [list(paths)] + + groups: list[list[str]] = [] + for rel in paths: + if groups and _day_key(groups[-1][0]) == _day_key(rel): + groups[-1].append(rel) + else: + groups.append([rel]) + + batches: list[list[str]] = [] + current: list[str] = [] + total = 0 + for group in groups: + group_size = sum(_file_size(workspace, rel) for rel in group) + if current and total + group_size > max_chars: + batches.append(current) + current, total = [], 0 + if group_size <= max_chars: + current.extend(group) + total += group_size + continue + for rel in group: + size = _file_size(workspace, rel) + if current and total + size > max_chars: + batches.append(current) + current, total = [], 0 + current.append(rel) + total += size + if current: + batches.append(current) + return batches + + def clean_paths(raw_paths, allowed: set[str]) -> list[str]: """Clean paths.""" if not isinstance(raw_paths, list): diff --git a/tests/unit/test_auto_dream.py b/tests/unit/test_auto_dream.py index 5a6a0be8..8475487a 100644 --- a/tests/unit/test_auto_dream.py +++ b/tests/unit/test_auto_dream.py @@ -27,7 +27,12 @@ from reme.steps.evolve.auto_tag import AutoTagStep from reme.steps.evolve.dream.extract import DreamExtractStep from reme.steps.evolve.dream.finish import DreamFinishStep from reme.steps.evolve.dream.integrate import DreamIntegrateStep, _snapshot_digest -from reme.steps.evolve.dream.utils import parse_structured_reply, recent_dates, scan_day_files +from reme.steps.evolve.dream.utils import ( # noqa: E501 + batch_paths, + parse_structured_reply, + recent_dates, + scan_day_files, +) from reme.steps.file_io.frontmatter_update import FrontmatterUpdateStep _MALFORMED_EXTRACT_REPLY = ( @@ -1032,3 +1037,216 @@ async def test_dream_initializes_started_provider_before_availability_check(tmp_ finally: await agent.close() await llm.close() + + +# --------------------------------------------------------------------------- +# Dream extract batching (#659): oversized material must never be sent in one +# call, and nothing may be silently dropped. +# --------------------------------------------------------------------------- + + +def test_batch_paths_partitions_without_dropping(tmp_path): + """Every path lands in exactly one batch and same-day files stay together.""" + day1 = [f"daily/2026-05-28/f{i}.md" for i in range(3)] + day2 = [f"daily/2026-05-27/g{i}.md" for i in range(2)] + for rel in [*day1, *day2]: + _touch(tmp_path / rel, "x" * 1000) + + batches = batch_paths(tmp_path, [*day1, *day2], max_chars=4000) + + assert [sorted(b) for b in batches] == [sorted(day1), sorted(day2)] + flat = [rel for batch in batches for rel in batch] + assert sorted(flat) == sorted([*day1, *day2]) + assert len(flat) == len(set(flat)) + + +def test_batch_paths_keeps_oversize_file_alone(tmp_path): + """A single file larger than the budget gets its own batch, never dropped.""" + _touch(tmp_path / "daily/2026-05-28/huge.md", "x" * 9000) + _touch(tmp_path / "daily/2026-05-28/small.md", "x" * 100) + + batches = batch_paths( + tmp_path, + ["daily/2026-05-28/huge.md", "daily/2026-05-28/small.md"], + max_chars=1000, + ) + + assert ["daily/2026-05-28/huge.md"] in [list(b) for b in batches] + flat = [rel for batch in batches for rel in batch] + assert sorted(flat) == ["daily/2026-05-28/huge.md", "daily/2026-05-28/small.md"] + + +async def _run_dream_extract(tmp_path, agent, **kwargs): + step = DreamExtractStep( + app_context=ApplicationContext(workspace_dir=str(tmp_path)), + agent_wrapper=agent, + **kwargs, + ) + with ( + patch("reme.steps.evolve.dream.extract.refresh_day_index", return_value={}), + patch("reme.steps.evolve.dream.extract.llm_available", return_value=True), + ): + return await step( + RuntimeContext( + date="2026-05-28", + file_catalog=_Catalog(), + file_store=_FileStore(tmp_path), + ), + ) + + +@pytest.mark.asyncio +async def test_extract_splits_oversized_material_across_calls(tmp_path): + """8 changed files must not be packed into a single extraction request.""" + day1 = [f"daily/2026-05-28/f{i}.md" for i in range(4)] + day2 = [f"daily/2026-05-27/g{i}.md" for i in range(4)] + for rel in [*day1, *day2]: + _touch(tmp_path / rel, "x" * 1000) + + agent = _SequenceAgent({"result": '{"units": []}'}, {"result": '{"units": []}'}) + response = await _run_dream_extract( + tmp_path, + agent, + scan_days=2, + max_units=5, + max_chars_per_batch=4000, + ) + + assert response.success is True + assert agent.calls == 2 + seen = [rel for message, _kwargs in agent.requests for rel in [*day1, *day2] if rel in message] + assert sorted(seen) == sorted([*day1, *day2]) + assert not response.metadata["dream"]["failed_paths"] + + +@pytest.mark.asyncio +async def test_extract_bisects_a_failing_batch(tmp_path): + """A transport failure splits the batch instead of retrying a larger prompt.""" + day1 = [f"daily/2026-05-28/f{i}.md" for i in range(2)] + for rel in day1: + _touch(tmp_path / rel, "x" * 1000) + + agent = _SequenceAgent( + RuntimeError("prompt too large"), + {"result": '{"units": []}'}, + {"result": '{"units": []}'}, + ) + response = await _run_dream_extract( + tmp_path, + agent, + scan_days=2, + max_units=5, + max_chars_per_batch=9000, + ) + + assert agent.calls == 3 + assert response.success is True + assert not response.metadata["dream"]["failed_paths"] + + +@pytest.mark.asyncio +async def test_extract_isolates_one_offending_file(tmp_path): + """A file that always fails only blocks itself, not the whole run.""" + day1 = [f"daily/2026-05-28/f{i}.md" for i in range(2)] + for rel in day1: + _touch(tmp_path / rel, "x" * 1000) + + good = { + "result": json.dumps( + { + "units": [ + { + "name": "kept", + "bucket": "wiki", + "summary": "kept unit", + "paths": ["daily/2026-05-28/f0.md"], + }, + ], + }, + ), + } + agent = _SequenceAgent( + RuntimeError("prompt too large"), + good, + RuntimeError("engine crashed"), + ) + response = await _run_dream_extract( + tmp_path, + agent, + scan_days=2, + max_units=5, + max_chars_per_batch=9000, + ) + + dream = response.metadata["dream"] + assert [u["name"] for u in dream["units"]] == ["kept"] + assert dream["failed_paths"] == ["daily/2026-05-28/f1.md"] + assert response.success is False + + +@pytest.mark.asyncio +async def test_extract_merges_same_name_units_across_batches(tmp_path): + """The same abstraction surfaced by two batches becomes one unit.""" + _touch(tmp_path / "daily/2026-05-28/a.md", "x" * 1000) + _touch(tmp_path / "daily/2026-05-27/b.md", "x" * 1000) + + first = { + "result": json.dumps( + { + "units": [ + { + "name": "ci-submodules", + "bucket": "procedure", + "summary": "short", + "paths": ["daily/2026-05-27/b.md"], + }, + ], + }, + ), + } + second = { + "result": json.dumps( + { + "units": [ + { + "name": "CI-Submodules", + "bucket": "procedure", + "summary": "a much longer summary", + "paths": ["daily/2026-05-28/a.md"], + }, + ], + }, + ), + } + agent = _SequenceAgent(first, second) + response = await _run_dream_extract( + tmp_path, + agent, + scan_days=2, + max_units=5, + max_chars_per_batch=1200, + ) + + units = response.metadata["dream"]["units"] + assert agent.calls == 2 + assert len(units) == 1 + assert units[0]["paths"] == ["daily/2026-05-27/b.md", "daily/2026-05-28/a.md"] + assert units[0]["summary"] == "a much longer summary" + + +@pytest.mark.asyncio +async def test_extract_single_small_input_stays_one_call(tmp_path): + """Batching must not change behaviour for inputs that already fit.""" + _touch(tmp_path / "daily/2026-05-28/only.md", "x" * 100) + agent = _SequenceAgent({"result": '{"units": []}'}) + + response = await _run_dream_extract( + tmp_path, + agent, + scan_days=1, + max_units=5, + max_chars_per_batch=40000, + ) + + assert response.success is True + assert agent.calls == 1