scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-04 07:30:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 10

fix(dream): batch oversized extraction input instead of one giant prompt (#659)

Browse source 
The dream extract step packed every changed file into a single material blob and sent it in one LLM call, so a large scan window could exceed the backend context and crash it (issue #659: 19.4K-token prompt -> vkAllocateMemory failure).

Split changed paths into size-capped batches (max_chars_per_batch, default 40000), scoped clean_paths to the current batch, merged same-name units across batches, and bisected a batch that fails in transport so the offending input is isolated instead of retried with an even larger prompt.
This commit is contained in:
BeiMu-new 2026-10-23 19:56:49 +08:00
parent c773fa7b77
commit 0ccb364bd6
4 changed files with 500 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
reme/config/default.yaml
View file

@ -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

258
reme/steps/evolve/dream/extract.py
View file

@ -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})

66
reme/steps/evolve/dream/utils.py
View file

@ -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/<date>``)."""
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):

220
tests/unit/test_auto_dream.py
View file

@ -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