From bfc0ccdad82a351147dda00de1bc811719be2654 Mon Sep 17 00:00:00 2001 From: sa-buc Date: Tue, 28 Jul 2026 12:06:56 +0800 Subject: [PATCH] refactor(beam): rename max_chunk_words to max_segment_words, drop one-off pilot scripts --- benchmark/beam/_pilot_ingest.py | 65 ------------ benchmark/beam/check_source_markers.py | 128 ----------------------- reme/config/beam.yaml | 9 +- reme/steps/benchmark/beam/auto_memory.py | 77 +++++++------- tests/unit/test_beam_auto_memory.py | 68 ++++++------ 5 files changed, 78 insertions(+), 269 deletions(-) delete mode 100644 benchmark/beam/_pilot_ingest.py delete mode 100644 benchmark/beam/check_source_markers.py diff --git a/benchmark/beam/_pilot_ingest.py b/benchmark/beam/_pilot_ingest.py deleted file mode 100644 index 6387654d..00000000 --- a/benchmark/beam/_pilot_ingest.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Pilot: ingest a few BEAM sessions and inspect source-marker quality. - -Temporary helper for prompt tuning — ingests the first N sessions of one -case into an isolated pilot workspace (never the real benchmark one). - -Usage: - python benchmark/beam/_pilot_ingest.py [case_id] [n_sessions] -""" - -import asyncio -import shutil -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).parent)) - -from run import _PROJECT_ROOT, load_beam_chat, load_eval_config # noqa: E402 - -PILOT_ROOT = _PROJECT_ROOT / "benchmark/memory_workspaces/beam_pilot" - - -async def main(case_id: str, n_sessions: int): - from reme import Application - from reme.config import resolve_app_config - - eval_config = load_eval_config() - chat_size = eval_config["dataset"]["chat_size"] - beam_root = _PROJECT_ROOT / eval_config["dataset"].get("beam_root", "benchmark/datasets/BEAM") - chat_path = beam_root / "chats" / chat_size / case_id / "chat.json" - - case_dir = PILOT_ROOT / f"{chat_size}_{case_id}" - if case_dir.exists(): - shutil.rmtree(case_dir) - case_dir.mkdir(parents=True, exist_ok=True) - - cfg = resolve_app_config( - config=eval_config["reme"]["config"], - workspace_dir=str(case_dir / ".reme"), - log_to_console=False, - log_to_file=False, - enable_logo=False, - ) - app = Application(**cfg) - await app.start() - try: - sessions = load_beam_chat(chat_path, chat_size, case_id)[:n_sessions] - for i, session in enumerate(sessions): - print(f"[pilot] ingesting session {i + 1}/{len(sessions)}: " - f"{session['session_id']} msgs={len(session['messages'])}", flush=True) - resp = await app.run_job( - "auto_memory", - messages=session["messages"], - session_id=session["session_id"], - date=session["date"], - ) - print(f"[pilot] success={resp.success} chunks={resp.metadata.get('chunks')} " - f"answer={str(resp.answer)[:300]}", flush=True) - await app.run_job("index_update") - finally: - await app.close() - - -if __name__ == "__main__": - asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "1", - int(sys.argv[2]) if len(sys.argv) > 2 else 2)) diff --git a/benchmark/beam/check_source_markers.py b/benchmark/beam/check_source_markers.py deleted file mode 100644 index c0087f64..00000000 --- a/benchmark/beam/check_source_markers.py +++ /dev/null @@ -1,128 +0,0 @@ -"""Check source markers in BEAM auto-memory notes. - -Scans daily notes under one or more case workspaces and validates every -source marker against the required format:: - - [[#L-L(,L-L)*]] - -Reported problems: - -- ``legacy_format``: old single-bracket ``[path:1-2]`` markers. -- ``bad_anchor``: ``[[...jsonl#...]]`` whose anchor is not L-prefixed - hyphen ranges (colon separators, missing ``L``, spaces, bare line...). -- ``labelled``: marker prefixed with ``Source``/``来源`` labels. -- ``missing_file``: linked session file does not exist in the workspace. -- ``out_of_range``: cited line numbers exceed the session file length, - or start > end. - -Usage: - python benchmark/beam/check_source_markers.py [workspace_root ...] - (default: benchmark/memory_workspaces/beam) -""" - -import re -import sys -from pathlib import Path - -_PROJECT_ROOT = Path(__file__).parent.parent.parent - -WIKILINK_RE = re.compile(r"\[\[([^\[\]\|\n]+?)\]\]") -VALID_ANCHOR_RE = re.compile(r"^L(\d+)-L(\d+)(,L(\d+)-L(\d+))*$") -RANGE_RE = re.compile(r"L(\d+)-L(\d+)") -# Old format: [session/dialog/x.jsonl:1-2,5-6] (single brackets, colon) -LEGACY_RE = re.compile(r"(? str: - """Return the note body, skipping the leading YAML frontmatter block.""" - if text.startswith("---\n"): - end = text.find("\n---", 4) - if end != -1: - return text[end + 4 :] - return text - - -def _session_line_count(workspace: Path, rel_path: str) -> int | None: - f = workspace / rel_path - if not f.is_file(): - return None - with open(f, encoding="utf-8") as fh: - return sum(1 for _ in fh) - - -def check_note(note: Path, workspace: Path) -> tuple[list[dict], int]: - """Return (problems, valid_marker_count) for one note file.""" - text = _split_frontmatter(note.read_text(encoding="utf-8")) - rel_note = str(note.relative_to(workspace)) - problems: list[dict] = [] - valid = 0 - - def _add(kind: str, marker: str): - problems.append({"note": rel_note, "kind": kind, "marker": marker.strip()[:160]}) - - for m in LEGACY_RE.finditer(text): - _add("legacy_format", m.group(0)) - - for m in LABEL_RE.finditer(text): - _add("labelled", m.group(0)) - - line_counts: dict[str, int | None] = {} - for m in WIKILINK_RE.finditer(text): - inner = m.group(1) - if "#" not in inner: - continue # plain file link (e.g. day-index), not a line marker - target, anchor = inner.split("#", 1) - if not target.endswith(".jsonl"): - continue - if not VALID_ANCHOR_RE.match(anchor): - _add("bad_anchor", m.group(0)) - continue - if target not in line_counts: - line_counts[target] = _session_line_count(workspace, target) - n_lines = line_counts[target] - if n_lines is None: - _add("missing_file", m.group(0)) - continue - ok = True - for r in RANGE_RE.finditer(anchor): - start, end = int(r.group(1)), int(r.group(2)) - if start < 1 or end > n_lines or start > end: - _add("out_of_range", f"{m.group(0)} (file has {n_lines} lines)") - ok = False - break - if ok: - valid += 1 - return problems, valid - - -def main(roots: list[str]) -> int: - """Scan every case workspace under *roots* and print a problem report.""" - total_valid = 0 - all_problems: list[dict] = [] - n_notes = 0 - for root in roots: - root_path = Path(root) - if not root_path.is_absolute(): - root_path = _PROJECT_ROOT / root - for workspace in sorted(root_path.glob("*/.reme")): - for note in sorted(workspace.rglob("daily/*/*.md")): - n_notes += 1 - problems, valid = check_note(note, workspace) - total_valid += valid - all_problems.extend({**p, "workspace": workspace.parent.name} for p in problems) - - print(f"Scanned {n_notes} notes: {total_valid} valid markers, {len(all_problems)} problems") - by_kind: dict[str, int] = {} - for p in all_problems: - by_kind[p["kind"]] = by_kind.get(p["kind"], 0) + 1 - for kind, count in sorted(by_kind.items(), key=lambda kv: -kv[1]): - print(f" {kind}: {count}") - print() - for p in all_problems: - print(f"[{p['kind']}] {p['workspace']}/{p['note']}\n {p['marker']}") - return 1 if all_problems else 0 - - -if __name__ == "__main__": - sys.exit(main(sys.argv[1:] or ["benchmark/memory_workspaces/beam"])) diff --git a/reme/config/beam.yaml b/reme/config/beam.yaml index 1aae28ed..9769ad1e 100644 --- a/reme/config/beam.yaml +++ b/reme/config/beam.yaml @@ -164,10 +164,11 @@ jobs: - messages steps: - backend: beam_auto_memory_step - # Long sessions are split into turn-aligned chunks and fed to the - # agent incrementally; each chunk holds at most this many words. - # <= 0 disables splitting. - max_chunk_words: 10000 + # Long sessions are split into turn-aligned segments and fed to the + # agent incrementally; each segment holds at most this many words. + # ("segment" here is a slice of the message list, unrelated to file + # chunking in the index.) <= 0 disables splitting. + max_segment_words: 10000 # ── Reindex (full rebuild) ── reindex: diff --git a/reme/steps/benchmark/beam/auto_memory.py b/reme/steps/benchmark/beam/auto_memory.py index dd57ab12..cccb0c28 100644 --- a/reme/steps/benchmark/beam/auto_memory.py +++ b/reme/steps/benchmark/beam/auto_memory.py @@ -8,8 +8,8 @@ from ...evolve.auto_memory import AutoMemoryStep, _normalize_msg_timestamp from ...file_io import validate_session_id from ....components import R -# Runtime-context key carrying the 0-based line offset of the current chunk -# inside the full session file (chunked ingestion of long sessions). +# Runtime-context key carrying the 0-based line offset of the current segment +# inside the full session file (segmented ingestion of long sessions). _LINE_OFFSET_KEY = "beam_line_offset" @@ -18,16 +18,16 @@ def _msg_word_count(msg: Msg) -> int: return len((msg.get_text_content() or "").split()) -def split_turn_chunks(messages: list[Msg], max_words: int) -> list[tuple[int, list[Msg]]]: - """Split *messages* into chunks of complete turns bounded by *max_words*. +def split_turn_segments(messages: list[Msg], max_words: int) -> list[tuple[int, list[Msg]]]: + """Split *messages* into segments of complete turns bounded by *max_words*. A turn starts at every ``user`` message and spans all following non-user - messages (assistant replies). Chunks only break at turn boundaries, so a - ``user + assistant`` exchange is never split across two chunks. A single - turn larger than *max_words* still becomes its own (oversized) chunk. + messages (assistant replies). Segments only break at turn boundaries, so a + ``user + assistant`` exchange is never split across two segments. A single + turn larger than *max_words* still becomes its own (oversized) segment. Returns ``[(offset, msgs), ...]`` where ``offset`` is the 0-based index of - the chunk's first message in *messages* — i.e. its line number in the + the segment's first message in *messages* — i.e. its line number in the session file minus 1. ``max_words <= 0`` disables splitting. """ if not messages: @@ -43,22 +43,22 @@ def split_turn_chunks(messages: list[Msg], max_words: int) -> list[tuple[int, li else: turns[-1].append(msg) - chunks: list[tuple[int, list[Msg]]] = [] + segments: list[tuple[int, list[Msg]]] = [] current: list[Msg] = [] current_words = 0 offset = 0 for turn in turns: turn_words = sum(_msg_word_count(m) for m in turn) if current and current_words + turn_words > max_words: - chunks.append((offset, current)) + segments.append((offset, current)) offset += len(current) current = [] current_words = 0 current.extend(turn) current_words += turn_words if current: - chunks.append((offset, current)) - return chunks + segments.append((offset, current)) + return segments def _parse_iso_seconds(value: str) -> datetime | None: @@ -176,50 +176,51 @@ class BeamAutoMemoryStep(AutoMemoryStep): agent wrapper's server-owned ``injected_job_kwargs``. Long sessions are ingested incrementally: the message list is split into - chunks of complete turns (``max_chunk_words`` step config, default 10000 - words per chunk) and each chunk runs one agent pass — the first pass - creates the daily note, later passes merge into it. Line numbers shown to - the agent always refer to the ORIGINAL session file, not the chunk. + segments of complete turns (``max_segment_words`` step config, default + 10000 words per segment) and each segment runs one agent pass — the first + pass creates the daily note, later passes merge into it. Line numbers + shown to the agent always refer to the ORIGINAL session file, not the + segment. """ async def execute(self): assert self.context is not None raw_messages = self.context.get("messages") or [] session_id: str = self.context.get("session_id", "") - max_words = int(self.kwargs.get("max_chunk_words", 10000) or 0) + max_words = int(self.kwargs.get("max_segment_words", 10000) or 0) messages = self._build_messages(raw_messages) - chunks = split_turn_chunks(messages, max_words) + segments = split_turn_segments(messages, max_words) - if len(chunks) <= 1: + if len(segments) <= 1: self.context[_LINE_OFFSET_KEY] = 0 await super().execute() return - # Persist the FULL session file up-front so every chunk's [Ln] labels + # Persist the FULL session file up-front so every segment's [Ln] labels # (and the note's source markers) match the final on-disk layout. The - # per-chunk saves inside the base execute become no-op appends. + # per-segment saves inside the base execute become no-op appends. if session_id and validate_session_id(session_id) is None: await self._save_session_messages(session_id, messages) base_hint = self.context.get("memory_hint", "") or "" answers: list[str] = [] - failed_chunks = 0 + failed_segments = 0 self.logger.info( - f"[{self.name}] chunked ingestion session_id={session_id!r} " - f"messages={len(messages)} chunks={len(chunks)} max_chunk_words={max_words}", + f"[{self.name}] segmented ingestion session_id={session_id!r} " + f"messages={len(messages)} segments={len(segments)} max_segment_words={max_words}", ) - for part, (offset, chunk) in enumerate(chunks, start=1): + for part, (offset, segment) in enumerate(segments, start=1): part_hint = ( - f"This is part {part}/{len(chunks)} of one long session " - f"(lines {offset + 1}-{offset + len(chunk)} of the session file). " + f"This is part {part}/{len(segments)} of one long session " + f"(lines {offset + 1}-{offset + len(segment)} of the session file). " "Earlier parts were already recorded; extract ONLY from the turns shown below." ) - self.context["messages"] = chunk + self.context["messages"] = segment self.context["memory_hint"] = f"{base_hint}\n{part_hint}".strip() self.context[_LINE_OFFSET_KEY] = offset - # One bad chunk (e.g. transiently corrupted note frontmatter) must - # not lose the remaining chunks of the session. + # One bad segment (e.g. transiently corrupted note frontmatter) must + # not lose the remaining segments of the session. try: await super().execute() ok = bool(self.context.response.success) @@ -228,19 +229,19 @@ class BeamAutoMemoryStep(AutoMemoryStep): ok = False summary = f"Error: {exc}" self.logger.warning( - f"[{self.name}] chunk {part}/{len(chunks)} failed session_id={session_id!r}: {exc}", + f"[{self.name}] segment {part}/{len(segments)} failed session_id={session_id!r}: {exc}", ) if not ok: - failed_chunks += 1 - answers.append(f"[part {part}/{len(chunks)}] {summary}") + failed_segments += 1 + answers.append(f"[part {part}/{len(segments)}] {summary}") # Restore caller-visible context and aggregate the response. self.context["messages"] = raw_messages self.context["memory_hint"] = base_hint - self.context.response.success = failed_chunks < len(chunks) + self.context.response.success = failed_segments < len(segments) self.context.response.answer = "\n".join(answers) self.context.response.metadata.update( - {"n_messages": len(messages), "chunks": len(chunks), "failed_chunks": failed_chunks}, + {"n_messages": len(messages), "segments": len(segments), "failed_segments": failed_segments}, ) def _build_messages(self, raw_messages: list) -> list[Msg]: @@ -258,9 +259,9 @@ class BeamAutoMemoryStep(AutoMemoryStep): # [[#L-L]]. A BEAM batch is already # in chronological order, matching the one-message-per-line layout # written by AutoMemoryStep._save_session_messages, so number each turn - # by its position plus the chunk's offset into the full session file - # (no re-sorting). When the session is ingested in chunks, the offset - # keeps [Ln] labels aligned with the ORIGINAL file, not the chunk. + # by its position plus the segment's offset into the full session file + # (no re-sorting). When the session is ingested in segments, the offset + # keeps [Ln] labels aligned with the ORIGINAL file, not the segment. session_id = self.context.get("session_id", "") if self.context is not None else "" line_offset = int(self.context.get(_LINE_OFFSET_KEY, 0) or 0) if self.context is not None else 0 source_path = self._session_source_path(session_id) diff --git a/tests/unit/test_beam_auto_memory.py b/tests/unit/test_beam_auto_memory.py index 193a04b5..bddb3e79 100644 --- a/tests/unit/test_beam_auto_memory.py +++ b/tests/unit/test_beam_auto_memory.py @@ -1,4 +1,4 @@ -"""Unit tests for BEAM auto-memory chunked ingestion and line numbering.""" +"""Unit tests for BEAM auto-memory segmented ingestion and line numbering.""" # pylint: disable=missing-class-docstring,missing-function-docstring,protected-access @@ -7,7 +7,7 @@ from agentscope.message import Msg from reme.components.runtime_context import RuntimeContext from reme.steps.benchmark.beam.auto_memory import ( BeamAutoMemoryStep, - split_turn_chunks, + split_turn_segments, ) @@ -28,40 +28,40 @@ def _dialog(n_turns: int, words_per_msg: int) -> list[Msg]: return messages -class TestSplitTurnChunks: +class TestSplitTurnSegments: def test_empty(self): - assert not split_turn_chunks([], 100) + assert not split_turn_segments([], 100) - def test_disabled_returns_single_chunk(self): + def test_disabled_returns_single_segment(self): messages = _dialog(3, 10) - chunks = split_turn_chunks(messages, 0) - assert len(chunks) == 1 - assert chunks[0] == (0, messages) + segments = split_turn_segments(messages, 0) + assert len(segments) == 1 + assert segments[0] == (0, messages) - def test_under_limit_single_chunk(self): + def test_under_limit_single_segment(self): messages = _dialog(3, 10) # 60 words total - chunks = split_turn_chunks(messages, 100) - assert len(chunks) == 1 - assert chunks[0][0] == 0 - assert chunks[0][1] == messages + segments = split_turn_segments(messages, 100) + assert len(segments) == 1 + assert segments[0][0] == 0 + assert segments[0][1] == messages def test_splits_at_turn_boundaries(self): - # 4 turns x 20 words each; limit 40 -> 2 turns per chunk + # 4 turns x 20 words each; limit 40 -> 2 turns per segment messages = _dialog(4, 10) - chunks = split_turn_chunks(messages, 40) - assert len(chunks) == 2 - offsets = [offset for offset, _ in chunks] + segments = split_turn_segments(messages, 40) + assert len(segments) == 2 + offsets = [offset for offset, _ in segments] assert offsets == [0, 4] - # Every chunk starts with a user message and ends with an assistant - for _, chunk in chunks: - assert chunk[0].role == "user" - assert chunk[-1].role == "assistant" + # Every segment starts with a user message and ends with an assistant + for _, segment in segments: + assert segment[0].role == "user" + assert segment[-1].role == "assistant" # No message lost or duplicated, order preserved - flattened = [m for _, chunk in chunks for m in chunk] + flattened = [m for _, segment in segments for m in segment] assert flattened == messages def test_never_splits_inside_a_turn(self): - # One turn alone exceeds the limit -> becomes its own oversized chunk + # One turn alone exceeds the limit -> becomes its own oversized segment messages = [ _msg("user", 5), _msg("assistant", 5), @@ -70,17 +70,17 @@ class TestSplitTurnChunks: _msg("user", 5), _msg("assistant", 5), ] - chunks = split_turn_chunks(messages, 60) - assert [offset for offset, _ in chunks] == [0, 2, 4] - assert [len(chunk) for _, chunk in chunks] == [2, 2, 2] + segments = split_turn_segments(messages, 60) + assert [offset for offset, _ in segments] == [0, 2, 4] + assert [len(segment) for _, segment in segments] == [2, 2, 2] def test_offsets_are_original_indices(self): messages = _dialog(5, 30) # 60 words per turn - chunks = split_turn_chunks(messages, 120) - # 2 turns per chunk -> offsets 0, 4, 8 - assert [offset for offset, _ in chunks] == [0, 4, 8] - for offset, chunk in chunks: - for i, msg in enumerate(chunk): + segments = split_turn_segments(messages, 120) + # 2 turns per segment -> offsets 0, 4, 8 + assert [offset for offset, _ in segments] == [0, 4, 8] + for offset, segment in segments: + for i, msg in enumerate(segment): assert msg is messages[offset + i] def test_multi_assistant_turn_stays_together(self): @@ -91,9 +91,9 @@ class TestSplitTurnChunks: _msg("user", 10), _msg("assistant", 10), ] - chunks = split_turn_chunks(messages, 30) - assert [offset for offset, _ in chunks] == [0, 3] - assert len(chunks[0][1]) == 3 + segments = split_turn_segments(messages, 30) + assert [offset for offset, _ in segments] == [0, 3] + assert len(segments[0][1]) == 3 class TestFormatHistoryLineNumbers: