From 60564d5d543d5f6748c27ea39a4fcdcabc8aa140 Mon Sep 17 00:00:00 2001 From: sa-buc Date: Tue, 4 Aug 2026 16:22:11 +0800 Subject: [PATCH] refactor(benchmark): make session compression config-driven via compress_session flag Move session-transcript compression from LME hard-coded injection to a runtime context flag set by evaluation.compress_session in each benchmark config. Compression is off by default for both BEAM and LME, and BaseAgenticAnswerStep now conditionally injects the _search compress payload only when the flag is truthy. --- benchmark/beam/config.yaml | 1 + benchmark/beam/run.py | 5 ++++- benchmark/longmemeval/config.yaml | 1 + benchmark/longmemeval/run.py | 8 +++++++- reme/steps/benchmark/base/agentic_answer.py | 17 +++++++++++++++-- reme/steps/benchmark/lme/agentic_answer.py | 12 +++--------- 6 files changed, 31 insertions(+), 13 deletions(-) diff --git a/benchmark/beam/config.yaml b/benchmark/beam/config.yaml index 66e63d04..cad4b6fd 100644 --- a/benchmark/beam/config.yaml +++ b/benchmark/beam/config.yaml @@ -11,6 +11,7 @@ dataset: evaluation: num_workers: 18 # 0 = auto; 1 = sequential; >1 = parallel (per-case) + compress_session: false # true = compress session chunks in search_v2 (query-aware); false = no compression reme: config: "beam.yaml" # reme config (in reme/config/) diff --git a/benchmark/beam/run.py b/benchmark/beam/run.py index 7227f74e..a75f2ddf 100644 --- a/benchmark/beam/run.py +++ b/benchmark/beam/run.py @@ -247,7 +247,7 @@ def get_available_cases(beam_root: Path, chat_size: str) -> list[str]: # --------------------------------------------------------------------------- # Answer generation # --------------------------------------------------------------------------- -async def answer_question_agentic(app, question: str) -> tuple[str, dict]: +async def answer_question_agentic(app, question: str, compress_session: bool = False) -> tuple[str, dict]: """Answer a probing question using ReMe's agentic_answer job. Returns (answer, metadata) @@ -264,6 +264,7 @@ async def answer_question_agentic(app, question: str) -> tuple[str, dict]: query_resp = await app.run_job( "agentic_answer", query=question, + compress_session=compress_session, ) answer = (query_resp.answer or "").strip() @@ -323,6 +324,7 @@ async def evaluate_case(eval_config: dict, case_id: str, eval_only: bool = False dataset_cfg = eval_config["dataset"] chat_size = dataset_cfg["chat_size"] + compress_session = bool(eval_config["evaluation"].get("compress_session", False)) beam_root = _PROJECT_ROOT / dataset_cfg.get("beam_root", "benchmark/beam/dataset/BEAM") chat_path = beam_root / "chats" / chat_size / case_id / "chat.json" probing_questions_path = beam_root / "chats" / chat_size / case_id / "probing_questions" / "probing_questions.json" @@ -452,6 +454,7 @@ async def evaluate_case(eval_config: dict, case_id: str, eval_only: bool = False agentic_answer, agentic_meta = await answer_question_agentic( app, question, + compress_session=compress_session, ) except Exception as e: logger.error(f"[Case {case_id}] Agentic answer failed: {e}") diff --git a/benchmark/longmemeval/config.yaml b/benchmark/longmemeval/config.yaml index d79b41ea..72ee7b75 100644 --- a/benchmark/longmemeval/config.yaml +++ b/benchmark/longmemeval/config.yaml @@ -15,6 +15,7 @@ evaluation: # Judgment is always binary (yes/no) — defined in lme/llm_judge.yaml num_workers: 32 # 0 = auto (cpu_count - 2, min 1); 1 = sequential; >1 = parallel filter_future_sessions: true # true = only ingest sessions with timestamp <= question_date + compress_session: false # true = compress session chunks in search_v2 (query-aware); false = no compression reme: config: "lme.yaml" # reme config to use (in reme/config/) diff --git a/benchmark/longmemeval/run.py b/benchmark/longmemeval/run.py index c3da2d63..9e48af43 100644 --- a/benchmark/longmemeval/run.py +++ b/benchmark/longmemeval/run.py @@ -426,6 +426,7 @@ async def evaluate_item(item: dict, eval_config: dict, item_index: int, eval_onl # ── Phase 4: Ask question via agentic_answer job (ReAct agent) ── question = item["question"] + compress_session = bool(eval_config["evaluation"].get("compress_session", False)) question_date_raw = item.get("question_date", "") question_dt = parse_haystack_date(question_date_raw) if question_date_raw else None query_time = to_iso(question_dt) if question_dt else "" @@ -440,7 +441,12 @@ async def evaluate_item(item: dict, eval_config: dict, item_index: int, eval_onl app.context, ) as token_usages, ): - query_resp = await app.run_job("agentic_answer", query=question, query_time=query_time) + query_resp = await app.run_job( + "agentic_answer", + query=question, + query_time=query_time, + compress_session=compress_session, + ) agentic_tool_counts = tool_counts agentic_token_usage = token_usages["bench"] agentic_response = (query_resp.answer or "").strip() diff --git a/reme/steps/benchmark/base/agentic_answer.py b/reme/steps/benchmark/base/agentic_answer.py index 9e47e117..eb06a55c 100644 --- a/reme/steps/benchmark/base/agentic_answer.py +++ b/reme/steps/benchmark/base/agentic_answer.py @@ -41,8 +41,21 @@ class BaseAgenticAnswerStep(BaseStep): Overridable hook: subclasses can extend the static ``INJECTED_JOB_KWARGS`` with per-request values derived from ``query``. + + When the runtime context carries a truthy ``compress_session`` flag, + session-transcript compression is enabled in ``search_v2_step`` by + injecting a ``_search._compress.session`` marker plus the current + ``query`` as the query-aware relevance filter. Default (falsy) leaves + session chunks uncompressed. """ - return dict(self.INJECTED_JOB_KWARGS) + injected = dict(self.INJECTED_JOB_KWARGS) + if self.context is not None and self.context.get("compress_session"): + injected["_search"] = { + "_compress": {"session": "true"}, + "queries": [query], + "type": "query-aware", + } + return injected async def execute(self): assert self.context is not None @@ -72,7 +85,7 @@ class BaseAgenticAnswerStep(BaseStep): "react_config": {"max_iters": self.MAX_ITERATION}, "tool_context_id": tool_context_id, } - if injected_job_kwargs := self._injected_job_kwargs(f"Query: {query}, query time: {query_time}"): + if injected_job_kwargs := self._injected_job_kwargs(f"{query}(query time: {query_time})"): wrapper_kwargs["injected_job_kwargs"] = injected_job_kwargs if self.context.stream: diff --git a/reme/steps/benchmark/lme/agentic_answer.py b/reme/steps/benchmark/lme/agentic_answer.py index bba031c8..70f8926f 100644 --- a/reme/steps/benchmark/lme/agentic_answer.py +++ b/reme/steps/benchmark/lme/agentic_answer.py @@ -12,15 +12,9 @@ class LmeAgenticAnswerStep(BaseAgenticAnswerStep): ``search`` job tool to retrieve relevant memory chunks before generating a final answer. - Every job tool call carries an injected ``_search`` payload that enables - session-transcript compression in ``search_v2_step`` and forwards the - original benchmark query as the compression relevance filter - (query-aware compression). + Session-transcript compression in ``search_v2_step`` is controlled by the + ``compress_session`` flag in the runtime context (set by the benchmark + runner from ``evaluation.compress_session``); it is off by default. """ TOOL_CONTEXT_PREFIX = "lme_agentic_answer" - - def _injected_job_kwargs(self, query: str) -> dict: - injected = super()._injected_job_kwargs(query) - injected["_search"] = {"_compress": {"session": "true"}, "queries": [query], "type": "query-aware"} - return injected