From ed74122cc8cdbb109c358ebd9487dbf67a59e838 Mon Sep 17 00:00:00 2001 From: WQS Date: Mon, 26 Oct 2026 15:07:16 +0800 Subject: [PATCH] feat: add ReMe Image evaluation on MemLens (#640) * feat(benchmark): add MemLens-256K reproduction * fix(benchmark): declare Pillow for MemLens images * fix(benchmark): retain MemLens usage and auditable failures * fix(benchmark): preserve MemLens diagnostic schemas during redaction * refactor(memlens): complete review cleanup and condense regression tests * fix(memlens): preserve nested usage fields during diagnostic redaction * test(memlens): simplify response schema collision coverage * refactor(memlens): simplify diagnostics and summary aggregation --------- Co-authored-by: jinli.yl --- benchmark/memlens/README.md | 107 +++ benchmark/memlens/README_ZH.md | 69 ++ benchmark/memlens/artifacts.py | 163 +++++ benchmark/memlens/config.yaml | 4 + benchmark/memlens/run.py | 621 ++++++++++++++++++ plugins/memlens-judge/LICENSE-CODE | 21 + plugins/memlens-judge/README.md | 30 + plugins/memlens-judge/pyproject.toml | 22 + .../src/judge_memlens/__init__.py | 5 + .../judge_memlens/_memlens_judge_prompts.json | 29 + .../src/judge_memlens/llm_judge.py | 72 ++ .../src/judge_memlens/plugin.yaml | 38 ++ .../src/judge_memlens/protocol.py | 139 ++++ plugins/memlens/README.md | 48 ++ plugins/memlens/pyproject.toml | 25 + plugins/memlens/src/reme_memlens/__init__.py | 5 + .../src/reme_memlens/agentic_answer.py | 136 ++++ .../src/reme_memlens/agentic_answer.yaml | 10 + .../memlens/src/reme_memlens/auto_memory.py | 34 + .../memlens/src/reme_memlens/auto_memory.yaml | 25 + .../memlens/src/reme_memlens/benchmark.yaml | 30 + plugins/memlens/src/reme_memlens/plugin.yaml | 75 +++ tests/unit/test_memlens_benchmark.py | 335 ++++++++++ tests/unit/test_memlens_judge.py | 101 +++ tests/unit/test_memlens_runner_failures.py | 423 ++++++++++++ 25 files changed, 2567 insertions(+) create mode 100644 benchmark/memlens/README.md create mode 100644 benchmark/memlens/README_ZH.md create mode 100644 benchmark/memlens/artifacts.py create mode 100644 benchmark/memlens/config.yaml create mode 100644 benchmark/memlens/run.py create mode 100644 plugins/memlens-judge/LICENSE-CODE create mode 100644 plugins/memlens-judge/README.md create mode 100644 plugins/memlens-judge/pyproject.toml create mode 100644 plugins/memlens-judge/src/judge_memlens/__init__.py create mode 100644 plugins/memlens-judge/src/judge_memlens/_memlens_judge_prompts.json create mode 100644 plugins/memlens-judge/src/judge_memlens/llm_judge.py create mode 100644 plugins/memlens-judge/src/judge_memlens/plugin.yaml create mode 100644 plugins/memlens-judge/src/judge_memlens/protocol.py create mode 100644 plugins/memlens/README.md create mode 100644 plugins/memlens/pyproject.toml create mode 100644 plugins/memlens/src/reme_memlens/__init__.py create mode 100644 plugins/memlens/src/reme_memlens/agentic_answer.py create mode 100644 plugins/memlens/src/reme_memlens/agentic_answer.yaml create mode 100644 plugins/memlens/src/reme_memlens/auto_memory.py create mode 100644 plugins/memlens/src/reme_memlens/auto_memory.yaml create mode 100644 plugins/memlens/src/reme_memlens/benchmark.yaml create mode 100644 plugins/memlens/src/reme_memlens/plugin.yaml create mode 100644 tests/unit/test_memlens_benchmark.py create mode 100644 tests/unit/test_memlens_judge.py create mode 100644 tests/unit/test_memlens_runner_failures.py diff --git a/benchmark/memlens/README.md b/benchmark/memlens/README.md new file mode 100644 index 00000000..73dcebfa --- /dev/null +++ b/benchmark/memlens/README.md @@ -0,0 +1,107 @@ +# ReMe Image on MemLens-128K / 256K + +Evaluate ReMe Image on MemLens's official 195-question memory-agent subset. +Each length uses its own full histories; every question has an isolated workspace. + +## Run + +```bash +pip install -e ".[as]" -e plugins/memlens -e plugins/memlens-judge +export MEMLENS_DATA=/path/to/MEMLENS +export EXPERIMENT_DIR=/path/to/new-experiment +python -m benchmark.memlens.run --model qwen3.8-max --context-length 256k +``` + +Supply `LLM_API_KEY` and `LLM_BASE_URL` through the environment, outside Git. +The dataset directory contains `dataset_128k.json`, `dataset_256k.json`, +`agent_subset_195.json` and `release_images/`. Select 128K with +`--context-length 128k`; `--limit-scenes 1 --output /path/to/separate-experiment` +runs one complete history. Concurrency is configured in `config.yaml`. + +## Evaluation + +Original text, timestamps, Image IDs and interleaved image bytes go to native +Auto Memory. Supplied BLIP captions and answer annotations are not inputs. +BM25 indexes daily notes, digest notes and original session transcripts. +Each question starts a fresh 30-iteration Agent with `search`, `read` and +`read_image`; no embedding or auto-tagging is used. + +| Setting | Memory / compression | Answer | +| --- | --- | --- | +| Model | Selected candidate | Same candidate | +| Thinking / temperature / streaming | On / 0 / off | On / 0 / off | +| Output budget | 65,536 | 16,384 | +| Model context size | 200,000 | 400,000 | + +The [candidate preset](../../plugins/memlens/src/reme_memlens/benchmark.yaml) +retains 30 Agent iterations, context trigger/reserve ratios 0.8/0.1, a +50,000-token tool-result limit and 4,096-image retention. Native requests use +three retries, a 60-second delay and a 600-second HTTP timeout; nested retries +are disabled. The runner does not repair model outputs. + +The independent [MemLens Judge](../../plugins/memlens-judge/README.md) uses the +official category/subtype-specific binary accuracy protocol, Qwen3.8-Max by default, +thinking off, temperature 0 and 256 output tokens. SubEM and F1 follow the +official code. Reference answers, question categories and subtypes enter only scoring. +`JUDGE_MODEL_NAME` overrides the judge model independently of the candidate. +`settings.json` records the effective judge model, generation parameters and +protocol after plugin defaults and environment expansion, without credentials. + +## Results + +`/memlens-//` contains settings and summary JSON. +`reme_image/scenes//` retains ingestion responses, a native workspace, +answers, judgments, metrics and coverage status. The workspace's `mem_session/` +contains native Agent JSONL trajectories. Existing output groups are not overwritten. + +New runs use usage policy `successful-http-requests-v2`: usage includes observed +successful HTTP requests in memory/answer scopes, including compression and requests +completed before a later Job or history failure. Each scope is counted once. +Judge requests and unsuccessful HTTP attempts are excluded. This is measured request +usage, not a full-sample cost estimate or provider billing total. Unknown counters +remain `null`; `records` and per-counter `coverage` show how many requests were measured. +Expected, scored and unscored question counts make coverage gaps explicit. + +Technical failures are coverage gaps; valid empty or iteration-limit answers are scored. +Failure artifacts retain the failed Response before validation and record the stage, +scene/session/question IDs, sanitized error chain and traceback locations without locals. +These artifacts omit credentials, request headers, raw configuration and image data. +Failed responses use separate `diagnostics/*-response.json` files; `answer.json` continues to +indicate a successful answer Job and retains the original prediction. Per-scope +`diagnostics/*-usage.json` receipts survive direct exceptions and cancellation, while +`status.json.memory_complete` records successful ingestion and indexing independently. +Cleanup errors retain scores and usage, increment `cleanup_error_count`, and cause +a nonzero CLI exit even when all questions were scored. +Historical reports using selected successful Job/history usage retain their original +values and scope; they are not rewritten under the new policy. + +## Sources and license + +[Dataset](https://huggingface.co/datasets/xiyuRenBill/MEMLENS); +[official code](https://github.com/xrenaf/MEMLENS/tree/77f3ab9a52fa2d6a17978e2dffe80438a4ecced2). +The dataset retains its own data/image terms. Adapted query and metric code carries +the following upstream notice: + +```text +MIT License + +Copyright (c) 2026 Xiyu Ren, Zhaowei Wang, and the MemLens Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +``` diff --git a/benchmark/memlens/README_ZH.md b/benchmark/memlens/README_ZH.md new file mode 100644 index 00000000..5a384859 --- /dev/null +++ b/benchmark/memlens/README_ZH.md @@ -0,0 +1,69 @@ +# 在 MemLens-128K / 256K 上评测 ReMe Image + +使用官方固定的 195 题 memory-agent 子集。128K、256K 分别读取自己的完整历史， +每题使用独立 workspace。 + +## 运行 + +```bash +pip install -e ".[as]" -e plugins/memlens -e plugins/memlens-judge +export MEMLENS_DATA=/path/to/MEMLENS +export EXPERIMENT_DIR=/path/to/new-experiment +python -m benchmark.memlens.run --model qwen3.8-max --context-length 256k +``` + +通过环境变量提供 `LLM_API_KEY` 和 `LLM_BASE_URL`，不要提交凭据。 +数据目录包含 `dataset_128k.json`、`dataset_256k.json`、 +`agent_subset_195.json` 和 `release_images/`。使用 +`--context-length 128k` 切换长度；用 `--limit-scenes 1` 和独立的 +`--output` 路径小批量验证，不裁剪场景历史。并发配置在 `config.yaml`。 + +## 评测设置 + +原始文本、时间、Image ID 和交错的原图输入原生 Auto Memory； +不使用附带的 BLIP caption 或答案标注。BM25 索引 daily、digest 和原始 +session 正文。每题新建 30 轮回答 Agent，工具为 `search`、`read`、 +`read_image`，不启用 embedding 或 auto-tag。 + +| 设置 | 记忆／压缩 | 回答 | +| --- | --- | --- | +| 模型 | 选定的候选模型 | 同一候选模型 | +| Thinking／temperature／stream | 开启／0／关闭 | 开启／0／关闭 | +| 输出预算 | 65,536 | 16,384 | +| 模型 context size | 200,000 | 400,000 | + +[候选插件配置](../../plugins/memlens/src/reme_memlens/benchmark.yaml) 保留 +30 轮预算、0.8／0.1 上下文触发／预留比例、50,000 token 的工具结果限制和 +4,096 张图片保留上限。原生请求最多重试 3 次，间隔 60 秒，HTTP timeout +为 600 秒；不叠加 SDK／Agent 重试，也不修补模型输出。 + +独立的 [MemLens Judge](../../plugins/memlens-judge/README.md) 使用官方按 +category／subtype 区分的二元评分：默认 Qwen3.8-Max、thinking 关闭、temperature +0、输出预算 256。辅助指标为官方 SubEM、F1。答案与题型字段仅进入评分。 +`JUDGE_MODEL_NAME` 可独立覆盖 Judge 模型；`settings.json` 记录插件默认值和 +环境变量展开后实际生效的 Judge 模型、生成参数与协议，不保存凭据或服务端点。 + +## 结果 + +结果保存到 `/memlens-//`。该目录的 JSON 记录配置与汇总； +`reme_image/scenes//` 保存摄入响应、原生 workspace、回答、judge、 +指标及完成状态。workspace 下的 `mem_session/` 保留原生 Agent JSONL 轨迹。 +已有结果目录不会被覆盖。 + +新运行使用 `successful-http-requests-v2` 口径：统计记忆／回答作用域中观测到的 +成功 HTTP 请求，包括压缩，以及所属 Job 或历史处理随后失败前已成功的请求； +每个作用域只计一次。不计 Judge 请求及未成功的 HTTP 尝试。这是实测请求用量， +并非完整样本成本估计或供应商账单。未知计数保持 `null`，`records` 与各计数的 +`coverage` 标示观测覆盖；预期、已评分、未评分题数分别记录。 + +技术失败计为覆盖缺口，合法空答案与达到轮数上限的回答仍评分。失败诊断记录阶段、 +场景／session／题目 ID、脱敏错误链和不含局部变量的堆栈位置，不保存凭据、 +请求头、原始配置或图片数据。失败 Response 在检查前保存到独立的 +`diagnostics/*-response.json`；`answer.json` 仍只表示成功回答，保留原始预测。 +`diagnostics/*-usage.json` 在直接异常或取消时也保留已观测用量； +`status.json.memory_complete` 独立表示历史摄入和索引完成。 +清理失败保留评分及用量、增加 `cleanup_error_count`，并使 CLI 返回非零状态。 + +历史报告沿用其原有的选定成功 Job／历史统计口径，不按新口径重写数据。 + +数据来源、上游版本和完整授权文本见[英文 README](README.md#sources-and-license)。 diff --git a/benchmark/memlens/artifacts.py b/benchmark/memlens/artifacts.py new file mode 100644 index 00000000..18c48749 --- /dev/null +++ b/benchmark/memlens/artifacts.py @@ -0,0 +1,163 @@ +"""Small, credential-free snapshots and diagnostics for benchmark artifacts.""" + +import os +from pathlib import Path +import re +import traceback +from urllib.parse import urlsplit, urlunsplit + +_PRIVATE_KEY = re.compile( + r"credential|api[_-]?key|private[_-]?key|authorization|password|secret|cookie|access[_-]?token|refresh[_-]?token|" + r"(^|[_-])token$|headers|request[_-]?body|^environment$|^config(?:uration)?$|^messages$", + re.IGNORECASE, +) +_ASSIGNMENT = re.compile( + r"(?i)([\"']?(?:api[_-]?key|private[_-]?key|authorization|password|secret|(?:access|refresh|id)[_-]?token|token)" + r"[\"']?\s*[:=]\s*)(?:\"[^\"]*\"|'[^']*'|[^\s,;}]+)", +) +_MODEL_FIELDS = ( + "backend", + "model", + "stream", + "context_size", + "max_retries", + "retry_delay", +) +_MODEL_OPTIONS = { + "parameters": ( + "temperature", + "top_p", + "max_tokens", + "max_completion_tokens", + "thinking_enable", + ), + "extra_body": ("enable_thinking",), + "client_kwargs": ("timeout", "max_retries"), +} +_USAGE_FIELDS = { + "usage", + "coverage", + "records", + "input_tokens", + "output_tokens", + "total_tokens", + "cached_input_tokens", + "noncached_input_tokens", + "reasoning_tokens", +} + + +def _safe_text(value: str) -> str: + for key, secret in sorted(os.environ.items(), key=lambda item: len(item[1]), reverse=True): + if secret and _PRIVATE_KEY.search(key): + if len(secret) < 12: + value = re.sub(r"(?]+", safe_url, value) + value = re.sub(r"(?s)-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----", "[REDACTED_KEY]", value) + value = re.sub(r"(?im)(\b(?:set[-_])?cookie[\"']?\s*[:=]\s*)[^\r\n]+", r"\1[REDACTED]", value) + value = re.sub(r"(?i)\b(?:Bearer|Basic)\s+[A-Za-z0-9._~+/=-]+", "[REDACTED_AUTH]", value) + value = _ASSIGNMENT.sub(r"\1[REDACTED]", value) + value = re.sub(r"(?i)data:[^;,\s]+;base64,[A-Za-z0-9+/=\s]+", "[REDACTED_IMAGE]", value) + value = re.sub(r"\b[A-Za-z0-9+/]{80,}={0,2}", "[REDACTED_DATA]", value) + # Exceptions sometimes embed complete request/config dictionaries in their text. + return re.sub( + r"(?is)(\b(?:config|configuration|request_headers|headers|request_body|messages)\b[\"']?\s*[:=]).*", + r"\1[REDACTED]", + value, + ) + + +def safe_value(value, *, fixed_keys=()): + """Redact diagnostic data recursively without serializing arbitrary objects.""" + if isinstance(value, dict): + return { + (key if key in fixed_keys else _safe_text(str(key))): ( + "[REDACTED]" + if _PRIVATE_KEY.search(str(key)) or (key == "data" and value.get("type") == "base64") + else safe_value(item, fixed_keys=fixed_keys) + ) + for key, item in value.items() + } + if isinstance(value, (list, tuple)): + return [safe_value(item, fixed_keys=fixed_keys) for item in value] + if isinstance(value, str): + return _safe_text(value) + if value is None or isinstance(value, (bool, int, float)): + return value + return f"<{type(value).__name__}>" + + +def status_record(status: dict) -> dict: + """Keep runner-owned schema/enums/counters intact; errors are already sanitized.""" + return { + key: safe_value(value) if key in {"scene_id", "session_id", "question_id"} else value + for key, value in status.items() + } + + +def response_record(response) -> dict: + """Keep the Response schema while sanitizing its untrusted payload and extras.""" + raw = response.model_dump() + return { + **safe_value({key: value for key, value in raw.items() if key not in {"success", "answer", "metadata"}}), + "success": raw["success"], + "answer": safe_value(raw["answer"]), + "metadata": safe_value(raw["metadata"], fixed_keys=_USAGE_FIELDS), + } + + +def error_details(exc: BaseException, seen: set[int] | None = None) -> dict: + """Keep the error chain and frame locations, never frame locals or source text.""" + seen = set() if seen is None else seen + seen.add(id(exc)) + result = { + "type": type(exc).__name__, + "message": _safe_text(str(exc)), + "traceback": [ + { + "file": _safe_text(Path(frame.filename).name), + "line": frame.lineno, + "function": _safe_text(frame.name), + } + for frame in traceback.extract_tb(exc.__traceback__) + ], + } + cause = exc.__cause__ or (None if exc.__suppress_context__ else exc.__context__) + if cause is not None and id(cause) not in seen: + result["cause"] = error_details(cause, seen) + return result + + +def model_settings(config) -> dict: + """Whitelist reproducibility fields from a validated ApplicationConfig.""" + models = {} + for name, component in config.components.get("as_llm", {}).items(): + spec = component.model_dump() + snapshot = {key: safe_value(spec[key]) for key in _MODEL_FIELDS if key in spec} + for key, allowed in _MODEL_OPTIONS.items(): + options = spec.get(key) or {} + snapshot[key] = {name: safe_value(options[name]) for name in allowed if name in options} + models[_safe_text(name)] = snapshot + return models + + +def judge_settings(app) -> dict: + """Record the actual plugin/env-expanded judge model and official protocol.""" + from judge_memlens.protocol import PROTOCOL_ID, SOURCE_REVISION + + return { + "protocol_id": PROTOCOL_ID, + "source_revision": SOURCE_REVISION, + "models": model_settings(app.config), + } diff --git a/benchmark/memlens/config.yaml b/benchmark/memlens/config.yaml new file mode 100644 index 00000000..77d961c8 --- /dev/null +++ b/benchmark/memlens/config.yaml @@ -0,0 +1,4 @@ +dataset: ${MEMLENS_DATA} +output: ${EXPERIMENT_DIR} +scenario_concurrency: 4 +query_concurrency: 16 diff --git a/benchmark/memlens/run.py b/benchmark/memlens/run.py new file mode 100644 index 00000000..433270fc --- /dev/null +++ b/benchmark/memlens/run.py @@ -0,0 +1,621 @@ +"""Evaluate ReMe Image on the official MemLens memory-agent subset.""" + +import argparse +import base64 +from collections import Counter +from contextlib import contextmanager +from contextvars import ContextVar +from datetime import datetime +import io +import asyncio +import json +import os +from pathlib import Path +import re +import subprocess +import string +import sys + +import yaml + +from benchmark.memlens.artifacts import ( + error_details, + judge_settings, + model_settings, + response_record, + safe_value, + status_record, +) + +# Adapted query/metric code retains its upstream MIT notice in README.md. +SYSTEM_PROMPT = ( + "Provide answers based on the given conversation history. " + "If the question cannot be answered based on the given conversation, " + 'respond with "Insufficient information".' +) +INSTRUCTION = "Directly output the answer with no extra output." + + +def _image_path(root: Path, source: str) -> Path: + """Resolve only original images inside the released image directory.""" + path = (root / source).resolve() + if Path(source).is_absolute() or "\\" in source or not path.is_relative_to(root): + raise ValueError("MemLens image must stay inside release_images") + return path + + +def load_memlens(root: Path, context_length: str = "128k") -> list[dict]: + """Select the official 195 IDs in dataset order, with each question's own history.""" + root = root.expanduser().resolve() + source = root / f"dataset_{context_length}.json" + selected = json.loads((root / "agent_subset_195.json").read_text(encoding="utf-8"))["question_ids"] + if len(selected) != 195 or len(set(selected)) != 195: + raise ValueError("MemLens requires its official 195-question agent subset") + selected = set(selected) + images = (root / "release_images").resolve() + if not images.is_relative_to(root): + raise ValueError("MemLens images must stay inside the dataset directory") + result = [] + for row in json.loads(source.read_text(encoding="utf-8")): + if row["question_id"] not in selected: + continue + sessions = [] + for turns, date, session_id in zip( + row["haystack_sessions"], + row["haystack_dates"], + row["haystack_session_ids"], + strict=True, + ): + for turn in turns: + for image in turn.get("images", []): + _image_path(images, image["file"]) + sessions.append( + { + "id": session_id, + "timestamp": datetime.strptime(date, "%Y/%m/%d (%a) %H:%M").isoformat(timespec="minutes"), + "turns": turns, + "question_id": row["question_id"], + "image_root": images, + }, + ) + question = { + "id": row["question_id"], + "text": row["question"], + "question_date": row["question_date"], + "answer": row["answer"], + "category": row["question_type"], + "subtype": row.get("question_subtype", ""), + "old_answer": row.get("old_answer") or row.get("question_content", {}).get("old_answer", ""), + } + result.append({"id": row["question_id"], "sessions": sessions, "questions": [question]}) + if len(result) != 195 or {scene["id"] for scene in result} != selected: + raise ValueError("MemLens dataset must contain every selected question exactly once") + return result + + +def image_block(path: Path) -> dict: + """Encode untouched image bytes with their actual decoded MIME.""" + from PIL import Image + + if path.stat().st_size > 50 * 1024 * 1024: + raise ValueError(f"Oversized original image: {path.name}") + data = path.read_bytes() + if len(data) > 50 * 1024 * 1024: + raise ValueError(f"Oversized original image: {path.name}") + with Image.open(io.BytesIO(data)) as decoded: + if decoded.width * decoded.height > 50_000_000: + raise ValueError(f"Oversized original image dimensions: {path.name}") + media_type = decoded.get_format_mimetype() + if not media_type or not media_type.startswith("image/"): + raise ValueError(f"Missing image MIME: {path.name}") + decoded.load() + return { + "type": "data", + "source": { + "type": "base64", + "data": base64.b64encode(data).decode("ascii"), + "media_type": media_type, + }, + } + + +def session_messages(session: dict) -> list[dict]: + """Keep source roles, IDs, timestamps and the official placement.""" + messages = [] + for index, turn in enumerate(session["turns"]): + text, images = turn.get("content", ""), turn.get("images", []) + if not images: + stripped = text.replace("", "").strip() + parts = [stripped] if stripped else [] + elif "" not in text: + parts = [*images, *([text.strip()] if text.strip() else [])] + else: + parts, image_index = [], 0 + for part in re.split(r"()", text): + if part == "": + if image_index < len(images): + parts.append(images[image_index]) + image_index += 1 + elif part.strip(): + parts.append(part) + parts.extend(images[image_index:]) + content = [] + for part in parts: + if isinstance(part, str): + content.append({"type": "text", "text": part}) + else: + source = part["file"] + content.extend( + [ + {"type": "text", "text": f"[Image ID: {Path(source).stem}]"}, + image_block(_image_path(session["image_root"], source)), + ], + ) + messages.append( + { + "id": f"{session['question_id']}:{session['id']}:turn:{index}", + "name": turn["role"], + "role": turn["role"], + "content": content, + "created_at": session["timestamp"], + }, + ) + return messages + + +def query_prompt(question: dict, context: str) -> str: + """Keep the official dated query template unchanged.""" + return ( + f"{SYSTEM_PROMPT}\nConversation:\n{context}\n\n" + f"{INSTRUCTION}\nQuestion Date: {question['question_date']}\nQuestion: {question['text']}\n" + ) + + +def _normalize(text: str) -> str: + """Native SubEM/F1 normalization, preserving decimal points and date slashes.""" + text = text.lower() + text = "".join( + ( + char + if char not in string.punctuation + or (char in "./" and 0 < index < len(text) - 1 and text[index - 1].isdigit() and text[index + 1].isdigit()) + else " " + ) + for index, char in enumerate(text) + ) + text = re.sub(r"\ba(?=\s+[a-z])", " ", text) + text = re.sub(r"\b(an|the)\b", " ", text) + return " ".join(text.split()) + + +def score_answer(question: dict, prediction: str) -> dict: + """Official inline auxiliary SubEM/F1; primary accuracy comes from the LLM judge.""" + from judge_memlens.protocol import parse_model_output + + predicted, expected = _normalize(parse_model_output(prediction)), _normalize(question["answer"]) + prediction_tokens, answer_tokens = predicted.split(), expected.split() + overlap = sum((Counter(prediction_tokens) & Counter(answer_tokens)).values()) + return { + "sub_em": float(bool(predicted and expected) and expected in predicted), + "f1": 2 * overlap / (len(prediction_tokens) + len(answer_tokens)) if overlap else 0.0, + } + + +_CURRENT = ContextVar("memlens_usage", default=None) +_KEYS = ( + "input_tokens", + "output_tokens", + "total_tokens", + "cached_input_tokens", + "noncached_input_tokens", + "reasoning_tokens", +) + + +def normalized_usage(raw: dict) -> dict: + """Keep unknown counters unknown; cache/reasoning are subsets, not extra tokens.""" + incoming = raw.get("input_tokens", raw.get("prompt_tokens")) + outgoing = raw.get("output_tokens", raw.get("completion_tokens")) + total = raw.get("total_tokens") + if total is None and incoming is not None and outgoing is not None: + total = incoming + outgoing + cached = raw.get( + "cached_input_tokens", + raw.get("cached_tokens", (raw.get("prompt_tokens_details") or {}).get("cached_tokens")), + ) + reasoning = raw.get("reasoning_tokens", (raw.get("completion_tokens_details") or {}).get("reasoning_tokens")) + return { + "input_tokens": incoming, + "output_tokens": outgoing, + "total_tokens": total, + "cached_input_tokens": cached, + "noncached_input_tokens": incoming - cached if incoming is not None and cached is not None else None, + "reasoning_tokens": reasoning, + } + + +def total_usage(records: list[dict]) -> dict: + """Aggregate measured requests, recording coverage of optional counters.""" + return { + key: ( + sum(row[key] for row in records if row.get(key) is not None) + if any(row.get(key) is not None for row in records) + else None + ) + for key in _KEYS + } | { + "records": len(records), + "coverage": {key: sum(row.get(key) is not None for row in records) for key in _KEYS}, + } + + +@contextmanager +def collect_usage(target: list | None = None, receipt: Path | None = None): + """Keep concurrent calls' usage separate and include their compression calls.""" + records = [] + token = _CURRENT.set(records) + try: + yield records + finally: + _CURRENT.reset(token) + if target is not None: + target.extend(records) + if receipt is not None: + save(receipt, total_usage(records)) + + +def observe_usage(app) -> None: + """Observe owned clients and distinguish native interrupted replies from answers.""" + for wrapper in app.context.components.get("agent_wrapper", {}).values(): + original_reply = wrapper.reply + + async def reply(*args, _reply=original_reply, **kwargs): + result = await _reply(*args, **kwargs) + last = result.get("last_message") or {} + if last.get("finished_reason") == "interrupted" or last.get("error") is not None: + raise RuntimeError("Native Agent was interrupted; see the application log.") + return result + + wrapper.reply = reply + for component in app.context.components.get("as_llm", {}).values(): + client = getattr(getattr(component.model, "client", None), "_client", None) + if client is None or getattr(client, "_memlens_usage_observer", False): + continue + original = client.send + + async def send(request, *args, _send=original, **kwargs): + response = await _send(request, *args, **kwargs) + records = _CURRENT.get() + if records is not None and response.status_code == 200 and response.is_stream_consumed: + try: + payload = response.json() + except (ValueError, UnicodeError): + payload = {} + records.append(normalized_usage(payload.get("usage") or {})) + return response + + client.send = send + client._memlens_usage_observer = True # pylint: disable=protected-access + + +def save(path: Path, value) -> None: + """Save source/results unchanged; callers sanitize diagnostics and config snapshots.""" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, ensure_ascii=False, indent=2, default=str) + "\n", encoding="utf-8") + + +def app_config(workspace: Path, model: str, *, judge: bool = False) -> dict: + """Resolve standard plugin settings through ReMe's configuration parser.""" + from reme.config import resolve_app_config + + values = resolve_app_config( + config="benchmark" if judge else "memlens-benchmark", + plugins=["memlens-judge" if judge else "memlens"], + workspace_dir=str(workspace), + log_config=False, + enable_logo=False, + log_to_console=False, + log_to_file=True, + ) + # BM25-only: do not start unused embedding clients. + values["components"].pop("as_embedding", None) + values["components"].pop("embedding_store", None) + values["components"]["file_store"]["default"]["embedding_store"] = None + if not judge: + for name in ("default", "bench", "compressor"): + values["components"]["as_llm"][name]["model"] = model + else: + values["jobs"] = {"answer_judge": {}} + values["components"] = {} + return values + + +def check_response(response, diagnostic_path: Path) -> None: + """Save failed responses before validation; keep execution failures unscored.""" + if not response.success: + save(diagnostic_path, {"response": response_record(response)}) + if asyncio.current_task().cancelling(): + raise asyncio.CancelledError + if not response.success: + detail = response.answer or response.metadata.get("error") or "No failure detail provided." + raise RuntimeError(f"ReMe Job failed; response detail: {safe_value(detail)}") + + +async def close_application(app) -> None: + """Finish owned cleanup even when cancellation arrives during close.""" + closing = asyncio.create_task(app.close()) + cancellation = None + while True: + try: + await asyncio.shield(closing) + break + except asyncio.CancelledError as exc: + if closing.done(): + raise + cancellation = exc + except Exception as exc: + if cancellation is not None: + raise cancellation from exc + raise + if cancellation is not None: + raise cancellation + + +async def evaluate_scene(scene, number, model, group, judge, slots): + """Accumulate one isolated history, then answer and judge its question.""" + from reme import Application + + label = re.sub(r"[^\w.-]+", "_", scene["id"])[:100] + case = group / "reme_image" / "scenes" / f"{number:04d}-{label}" + app = None + memory, answers, results = [], [], [] + status = {"complete": False, "memory_complete": False, "stage": "initialize", "scene_id": scene["id"]} + try: + save(case / "source.json", {"id": scene["id"], "sessions": len(scene["sessions"])}) + app = Application(**app_config(case / "ingestion/attempt-001/workspace", model)) + status["stage"] = "start" + await app.start() + observe_usage(app) + for ordinal, session in enumerate(scene["sessions"]): + directory = case / "ingestion/attempt-001" + status.update(stage="preprocess", session_id=session["id"]) + async with slots: + with collect_usage(memory, directory / "diagnostics" / f"{ordinal:05d}-usage.json") as requests: + messages = await asyncio.to_thread(session_messages, session) + status["stage"] = "auto_memory" + response = await app.run_job( + "auto_memory", + messages=messages, + session_id=session["id"], + date=session["timestamp"][:10], + include_images=True, + include_source_sessions=True, + ) + check_response(response, directory / "diagnostics" / f"{ordinal:05d}-response.json") + save( + directory / "sessions" / f"{ordinal:05d}.json", + { + "session_id": session["id"], + "response": response.model_dump(), + "usage": total_usage(requests), + }, + ) + status.pop("session_id", None) + status["stage"] = "index_update" + response = await app.run_job("index_update") + check_response(response, case / "ingestion/attempt-001/diagnostics/index-response.json") + status["memory_complete"] = True + for index, question in enumerate(scene["questions"]): + directory = case / "questions" / f"{index:04d}" + status.update(stage="agentic_answer", question_id=question["id"]) + async with slots: + with collect_usage(answers, directory / "diagnostics/answer-usage.json") as requests: + response = await app.run_job( + "agentic_answer", + query=query_prompt(question, "Use the available tools to retrieve relevant evidence."), + query_time=question["question_date"], + query_images=[], + scene_id=scene["id"], + include_images=True, + include_source_sessions=True, + max_iters=30, + ) + check_response(response, directory / "diagnostics/answer-response.json") + save( + directory / "answer.json", + { + "question_id": question["id"], + "prediction": response.answer, + "usage": total_usage(requests), + "agent": response.metadata, + }, + ) + status["stage"] = "judge" + async with slots: + judgment = await judge.run_job( + "answer_judge", + question=question["text"], + reference=question["answer"], + prediction=response.answer, + category=question["category"], + subtype=question.get("subtype", ""), + old_answer=question.get("old_answer", ""), + ) + check_response(judgment, directory / "diagnostics/judge-response.json") + status["stage"] = "score" + metrics = score_answer(question, response.answer) + metrics["llm_judge"] = judgment.metadata["llm_judge_score"] + save(directory / "judge.json", judgment.metadata) + record = { + "question_id": question["id"], + "category": question["category"], + "question": question["text"], + "reference": question["answer"], + "prediction": response.answer, + "metrics": metrics, + } + save(directory / "result.json", record) + results.append(record) + status.update(complete=True, stage="complete") + except asyncio.CancelledError as exc: + status["error"] = error_details(exc) + raise + except Exception as exc: # pylint: disable=broad-exception-caught + status["error"] = error_details(exc) + finally: + try: + if app is not None: + await close_application(app) + except (Exception, asyncio.CancelledError) as exc: # pylint: disable=broad-exception-caught + if "error" in status: + status["cleanup_error"] = error_details(exc) + else: + status.update(complete=False, stage="cleanup", error=error_details(exc)) + if isinstance(exc, asyncio.CancelledError): + raise + finally: + status.update(memory_usage=total_usage(memory), answer_usage=total_usage(answers)) + if "error" in status: + status["error_type"] = status["error"]["type"] + save(case / "status.json", status_record(status)) + return results, memory, answers, status + + +async def run(settings: dict, model: str, length: str, limit: int | None = None) -> dict: + """Run one model in a new output group; existing runs are never overwritten.""" + from reme import Application + from reme.schema import ApplicationConfig + + group = Path(settings["output"]).expanduser().resolve() / f"memlens-{length}" / model + group.mkdir(parents=True, exist_ok=False) + judge = None + status = {"complete": False, "stage": "preprocess"} + try: + scenes = load_memlens(Path(settings["dataset"]), length) + if limit is not None: + scenes = scenes[:limit] + status["stage"] = "configure" + effective = ApplicationConfig(**app_config(group / "unused-workspace", model)) + judge = Application(**app_config(group / "judge-workspace", model, judge=True)) + save( + group / "settings.json", + { + "method": "reme_image", + "model": model, + "context_length": length, + "dataset": str(Path(settings["dataset"]).expanduser().resolve()), + "limit_scenes": limit, + "include_source_sessions": True, + "include_images": True, + "answer_max_iters": 30, + "concurrency": {key: settings[key] for key in ("scenario_concurrency", "query_concurrency")}, + "models": model_settings(effective), + "judge": judge_settings(judge), + "usage_policy_version": "successful-http-requests-v2", + "reme_commit": subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + }, + ) + slots = asyncio.Semaphore(settings["query_concurrency"]) + scene_slots = asyncio.Semaphore(settings["scenario_concurrency"]) + status["stage"] = "judge_start" + await judge.start() + status["stage"] = "evaluate" + + async def evaluate(number, scene): + async with scene_slots: + print(f"Starting {length} {model} {scene['id']}", flush=True) + result = await evaluate_scene(scene, number, model, group, judge, slots) + print(f"Finished {scene['id']}: {len(result[0])}/{len(scene['questions'])} scored", flush=True) + return result + + tasks = [asyncio.create_task(evaluate(index, scene)) for index, scene in enumerate(scenes)] + try: + outputs = await asyncio.gather(*tasks) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + status["stage"] = "summary" + results = [row for output in outputs for row in output[0]] + memory = [row for output in outputs for row in output[1]] + answers = [row for output in outputs for row in output[2]] + expected = sum(len(scene["questions"]) for scene in scenes) + statuses = [output[3] for output in outputs] + cleanup_errors = sum("cleanup_error" in row or row["stage"] == "cleanup" for row in statuses) + summary = { + "method": "reme_image", + "model": model, + "context_length": length, + "expected_scenes": len(scenes), + "completed_scenes": sum(row["complete"] for row in statuses), + "memory_complete_scenes": sum(row["memory_complete"] for row in statuses), + "cleanup_error_count": cleanup_errors, + "expected_questions": expected, + "scored_questions": len(results), + "unscored_questions": expected - len(results), + "metrics": { + key: (sum(row["metrics"][key] for row in results) / len(results) if results else None) + for key in ("llm_judge", "sub_em", "f1") + }, + "memory_usage": total_usage(memory), + "answer_usage": total_usage(answers), + "token_usage": total_usage(memory + answers), + "usage_policy_version": "successful-http-requests-v2", + "usage_policy": ( + "observed successful HTTP requests in memory/answer scopes, including compression and requests " + "before Job/history failure; judge excluded; not full-sample or billed cost" + ), + } + save(group / "summary.json", summary) + status.update(complete=len(results) == expected and not cleanup_errors, stage="complete") + except (Exception, asyncio.CancelledError) as exc: + status["error"] = error_details(exc) + raise + finally: + try: + if judge is not None: + await close_application(judge) + except (Exception, asyncio.CancelledError) as exc: + if "error" in status: + status["cleanup_error"] = error_details(exc) + if isinstance(exc, asyncio.CancelledError): + raise + else: + status.update(complete=False, stage="judge_cleanup", error=error_details(exc)) + raise + finally: + save(group / "status.json", status_record(status)) + return summary + + +def main(): + """Run one model at the requested MemLens history length.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, default=Path(__file__).with_name("config.yaml")) + parser.add_argument("--model", default="qwen3.8-max") + parser.add_argument("--context-length", choices=("256k", "128k"), default="256k") + parser.add_argument("--limit-scenes", type=int) + parser.add_argument("--output", type=Path) + arguments = parser.parse_args() + try: + settings = yaml.safe_load(os.path.expandvars(arguments.config.read_text(encoding="utf-8"))) + if arguments.output: + settings["output"] = str(arguments.output) + if arguments.limit_scenes is not None and arguments.limit_scenes < 1: + parser.error("--limit-scenes must be positive") + for key in ("query_concurrency", "scenario_concurrency"): + if not isinstance(settings[key], int) or settings[key] < 1: + parser.error(f"{key} must be positive") + summary = asyncio.run(run(settings, arguments.model, arguments.context_length, arguments.limit_scenes)) + except (Exception, asyncio.CancelledError, KeyboardInterrupt) as exc: # pylint: disable=broad-exception-caught + print(json.dumps({"complete": False, "error": error_details(exc)}, ensure_ascii=False), file=sys.stderr) + return 130 if isinstance(exc, (asyncio.CancelledError, KeyboardInterrupt)) else 2 + print(json.dumps(summary, ensure_ascii=False, indent=2)) + return ( + 0 if summary["scored_questions"] == summary["expected_questions"] and not summary["cleanup_error_count"] else 2 + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/plugins/memlens-judge/LICENSE-CODE b/plugins/memlens-judge/LICENSE-CODE new file mode 100644 index 00000000..a6bed67f --- /dev/null +++ b/plugins/memlens-judge/LICENSE-CODE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Xiyu Ren, Zhaowei Wang, and the MemLens Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/plugins/memlens-judge/README.md b/plugins/memlens-judge/README.md new file mode 100644 index 00000000..da0b4e0b --- /dev/null +++ b/plugins/memlens-judge/README.md @@ -0,0 +1,30 @@ +# ReMe MemLens Judge + +Independent `answer_judge` Job for saved MemLens predictions. Install with +`pip install -e plugins/memlens-judge` and enable the `memlens-judge` plugin +in a separate evaluation application/workspace, not the candidate memory application. + +Inputs are `question`, `reference`, `prediction`, and `category`, with optional +`subtype` and `old_answer`. The result is available as +`response.metadata.llm_judge_score` and `acc`. Metadata includes the raw response, +task rubric, native answer-preprocessing diagnostics and judge token usage. +Judge usage is evaluation cost, not method cost. + +The official category/subtype rubrics, outdated-answer handling, verdict parser, +long-answer preprocessing and over-500-word zero rule are preserved. A successful +but unparseable judge output receives the official zero score; API errors instead +remain failed jobs for resumable judging of the same saved prediction. + +The default judge is `qwen3.8-max`, with temperature 0 and thinking disabled. +Set `JUDGE_MODEL_NAME` to override the judge model independently of the candidate. +The benchmark's `settings.json` records the effective judge model, generation +parameters and protocol from the validated application configuration, after plugin +defaults and environment expansion. The snapshot excludes credentials and endpoints. +The output +budget is 256 tokens from the released OpenAI-compatible API judge branch +(`max_completion_tokens`), not the separate vLLM branch's 2,048-token setting. +These judge settings are independent of the candidate reader's thinking mode. + +Source: [official MemLens judge](https://github.com/xrenaf/MEMLENS/blob/77f3ab9a52fa2d6a17978e2dffe80438a4ecced2/llm_judge.py). +Verbatim prompt strings and their MIT notice are packaged in `_memlens_judge_prompts.json`; +see also `LICENSE-CODE`. The plugin is independent of candidate code and other benchmarks. diff --git a/plugins/memlens-judge/pyproject.toml b/plugins/memlens-judge/pyproject.toml new file mode 100644 index 00000000..8782eb5b --- /dev/null +++ b/plugins/memlens-judge/pyproject.toml @@ -0,0 +1,22 @@ +[project] +name = "reme-memlens-judge" +version = "0.1.0" +description = "Official memlens judge protocol as a ReMe plugin." +readme = "README.md" +license = "Apache-2.0" +requires-python = ">=3.11" +dependencies = ["reme-ai[as]>=0.4.1.11"] + +[project.entry-points."reme.plugins"] +memlens-judge = "judge_memlens" + +[tool.setuptools.packages.find] +where = ["src"] +include = ["judge_memlens*"] + +[tool.setuptools.package-data] +judge_memlens = ["plugin.yaml", "*.json"] + +[build-system] +requires = ["setuptools>=77", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/plugins/memlens-judge/src/judge_memlens/__init__.py b/plugins/memlens-judge/src/judge_memlens/__init__.py new file mode 100644 index 00000000..39c8fd4a --- /dev/null +++ b/plugins/memlens-judge/src/judge_memlens/__init__.py @@ -0,0 +1,5 @@ +"""Independent memlens evaluation-only plugin.""" + +from .llm_judge import MemLensJudgeStep + +__all__ = ["MemLensJudgeStep"] diff --git a/plugins/memlens-judge/src/judge_memlens/_memlens_judge_prompts.json b/plugins/memlens-judge/src/judge_memlens/_memlens_judge_prompts.json new file mode 100644 index 00000000..bf34840b --- /dev/null +++ b/plugins/memlens-judge/src/judge_memlens/_memlens_judge_prompts.json @@ -0,0 +1,29 @@ +{ + "source_revision": "77f3ab9a52fa2d6a17978e2dffe80438a4ecced2", + "license": "MIT License\n\nCopyright (c) 2026 Xiyu Ren, Zhaowei Wang, and the MemLens Authors\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n", + "template": "Now your role is a grading teacher. Your task is to review and score student answers based on reference standard answers for a question-answering benchmark. You need to notice the following key points:\n- First, extract the final answer from the student's solution, then analyze and judge whether the answer is correct.\n- Scoring should only refer to the final answer obtained by the student; there is no need to examine whether the intermediate problem-solving steps are correct.\n- If the response contains both hesitation and a clear answer, judge the answer itself.\n- If the response contains both a refusal and a guessed answer, judge the final committed answer.\n- If the response gives multiple inconsistent answers, assign 0 points.\n- If the student's response shows circular reasoning — repeatedly revisiting the same evidence, flip-flopping between answers (e.g., \"yes... wait no... actually yes... let me reconsider... no\"), with no clear final commitment — assign 0 points. A correct answer mentioned during reasoning does NOT count unless it is the clearly stated final conclusion.\n- Only the student's LAST clearly stated answer counts. If the student initially says \"yes\" but later concludes \"no\" (or vice versa), score based on the LAST position only. Intermediate answers during reasoning are not final answers.\n- If the student's output is entirely a reasoning trace (e.g., \"let me check... scanning evidence... re-reading session 5...\") with no clear concluding answer statement, assign 0 points — even if the correct answer appears somewhere within the reasoning.\n- If the student's output begins with \"[...truncated earlier reasoning...]\", this means only the final portion of a longer response is shown. Focus on extracting the answer from this final portion.\n\nBelow are examples of problematic student outputs that apply to ALL task types. Study these before grading.\n\n[Universal Example A — Circular Reasoning -> 0 points]\n: Is the cat in Mark and Jenny's 5-year anniversary post named Mittens?\n: Yes\n: let me check session 5 mark has cat named mittens session 7 mark and jenny anniversary post there is no link wait let me re-read session 5 says mittens session 7 says anniversary actually looking at this again wait is it possible let me reconsider mark has mittens in session 5 the post is about couple therefore answer is likely no wait let me double check session 5 mark has cat mittens actually maybe yes but text doesnt say\n[Scoring Rationale]: The output flip-flops between \"yes\" and \"no\" multiple times with no clear final commitment. The last position is ambiguous (\"maybe yes but text doesnt say\"). This is circular reasoning.\nIn summary, the student's answer deserves 0 points.\n[JSON]: {{\"answer_score\": 0}}\n\n[Universal Example B — Redundant But Committed -> Score Normally]\n: Is the trailing plant in the cafe safe for my kitten?\n: No\n: the plant looks like pothos pothos is toxic to cats therefore answer is no let me double check is pothos safe for cats no is answer no yes wait let me verify image 3 is pothos image 9 is pothos pothos is toxic therefore answer must be no\n[Scoring Rationale]: Despite excessive self-verification, the student consistently commits to \"no\" throughout and never wavers to a different answer. The final answer is clearly \"No\", which matches the standard answer.\nIn summary, the student's answer deserves 1 point.\n[JSON]: {{\"answer_score\": 1}}\n\n[Universal Example C — Reasoning Trace With No Answer -> 0 points]\n: How many sessions mentioned the coffee shop?\n: 3\n: scanning session 1 yes coffee mentioned session 2 no session 3 yes coffee again session 4 no session 5 maybe let me re-read session 5 it mentions cafe is that same as coffee shop need to check session 6 no session 7 unclear let me look at session 5 again\n[Scoring Rationale]: The student walks through sessions but never states a final count. The output is entirely a reasoning trace with no concluding answer.\nIn summary, the student's answer deserves 0 points.\n[JSON]: {{\"answer_score\": 0}}\n\nNow proceed with grading. Remember: the universal examples above apply regardless of task type.\n\n- When analyzing and judging whether the answer is correct, you need to write down the scoring rationale, organize it into clear statements that follow the logical flow. The summary of the scoring rationale should be placed at the end, using the following format: \"In summary, the student's answer deserves x points\" (where x represents the student's specific score).\n- Keep the whole process concise, within 150 words.\n- Provide the score based on your analysis and display it in a code block in \"JSON\" format.\n- An item is covered if it is strictly mentioned or unambiguously implied by a semantic equivalence. This includes numerical equivalence (e.g., 10% and 0.1), synonyms (e.g., UK and United Kingdom), plural/singular forms (e.g., \"apple\" and \"apples\"), and equivalent date formats (e.g., 2024-01-15 and January 15, 2024). However, do not accept loosely related concepts.\n- Ignore minor formatting differences, capitalization, punctuation, and equivalent wording when meaning is unchanged.\nYour output format is:\n[Scoring Rationale]:\n[Score]: x points\n[JSON]:\n{{\"answer_score\": }}\n\nBelow is the grading rubric:\n[Scores]:\nThe scoring scale consists of 2 levels in total, from highest to lowest: 1 point, 0 points (the minimum is 0 points).\n[Tier Details]:\n1 point: Assign 1 point if the student's final answer matches the standard answer under the task-specific criteria below.\n0 points: Assign 0 points if the student's final answer does not match the standard answer, the student refuses to answer, claims insufficient information (except for AR tasks), or does not clearly answer.\n\n[Task-Specific Criteria]:\n{task_criteria}\n\n{examples}\n", + "criteria": { + "IE_Entity": "[IE — Entity / Attribute Extraction]\nThis is an information extraction question about a concrete entity, count, attribute, text, object, location, or visual detail from the provided memory context.\n- Assign 1 point if the student's response contains the core information from the standard answer. Minor wording differences are acceptable, but the essential information must be present and correct.\n- Assign 0 points if the core information is missing, contradicted, too vague, refused, or incorrect.\n- For numeric or count answers, require the exact value unless the standard answer itself is approximate.", + "IE_PreviousInfo": "[IE — Previous Information Extraction]\nThis is an information extraction question about previously mentioned information, spatial relations, counts, or attributes from earlier conversation/image context.\n- Assign 1 point if the student's response recovers the same previous information as the standard answer, including the correct relation, location, count, or attribute.\n- Minor wording differences are acceptable if they preserve the same meaning.\n- Assign 0 points if the answer refers to the wrong previous item, gives the wrong relation/count/location, is too vague, refuses, or is unsupported.", + "KU_KnowledgeUpdate": "[KU — Knowledge Update]\nThis is a knowledge update question testing whether the model tracks the latest state.\nThe correct current answer and the outdated old answer are both provided.\n- Assign 1 point if the student gives the current (correct) answer.\n- Assign 0 points if the student gives the outdated old answer, any other incorrect answer, refuses, or does not clearly answer.", + "TR_DurationComparison": "[TR — Duration Comparison (A/B)]\nThis is a binary choice question. The student must select one of two options.\n- Assign 1 point if the student clearly selects the same option as the standard answer, regardless of phrasing.\n- Assign 0 points if the student selects the other option, refuses, or does not make a clear selection.", + "TR_OrderRanking": "[TR — Order Ranking]\nThis is an ordering question. The standard answer is an exact sequence of items.\n- Assign 1 point ONLY if the student gives the exact same sequence in the same order.\n- Assign 0 points if any item is missing, added, duplicated, misplaced, or if the order is unclear.", + "TR_DateExtraction": "[TR — Date Extraction]\nThis is a date extraction question. The standard answer is a specific date.\n- Assign 1 point if the student gives the same date in any unambiguous equivalent format (e.g., \"Jan 15, 2024\" and \"2024-01-15\" and \"15/01/2024\").\n- Assign 0 points if the date is different, incomplete when a full date is required, ambiguous, or missing.", + "MSR_YesNo": "[MSR — Yes/No (Entity Resolution)]\nThis is a yes/no question.\n- Assign 1 point if the student's final answer is semantically equivalent to the standard answer (e.g., \"Yes\" matches \"Yes\", \"No\" matches \"No\"), even if phrased differently (e.g., \"That's correct\" for \"Yes\").\n- Assign 0 points if the answer means the opposite, refuses, or does not clearly answer.", + "MSR_Counting": "[MSR — Counting]\nThis is a counting question. The standard answer is a specific number.\n- Assign 1 point if the student clearly states the exact number in any equivalent form (e.g., \"3\", \"three\", \"3.0\").\n- Assign 0 points if the student gives a different number, a range, an approximation, refuses, or does not clearly answer.", + "MSR_Arithmetic": "[MSR — Arithmetic]\nThis is an arithmetic question. The standard answer is a specific value.\n- Assign 1 point if the student gives the same value in an equivalent format (e.g., \"0.5\" and \"1/2\" and \"50%\").\n- Assign 0 points if the student gives a different value, an unsupported approximation, refuses, or does not clearly answer.", + "AR_AnswerRefusal": "[AR — Answer Refusal]\n- Assign 1 point if the student explicitly refuses, says there is not enough information, or indicates the question is unanswerable from the provided context.\n- Assign 0 points if the student gives a substantive answer, even if hedged." + }, + "examples": { + "IE_Entity": "\n[Example 1]\n: What specific word is painted in large white letters on the pavement in the middle lanes?\n: SCHOOL\n: SCHOOL\n\n[Scoring Rationale]: The student's answer is \"SCHOOL\", which exactly matches the standard answer.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: Where on the sign is the photo of the mountains located?\n: In the bottom right corner\n: The photo of the mountains is located in the upper right corner of the sign.\n\n[Scoring Rationale]: The student says \"upper right corner\" but the standard answer is \"bottom right corner\". The location is incorrect — \"upper\" vs \"bottom\" is a substantive difference.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "IE_PreviousInfo": "\n[Example 1]\n: Where on the sign is the photo of the mountains located?\n: In the bottom right corner\n: The photo is in the bottom-right corner of the sign.\n\n[Scoring Rationale]: The student's answer preserves the same location relation as the standard answer.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: Relative to the baking sheet, where is the salmon fillet placed?\n: The right side\n: It is on the left side of the baking sheet.\n\n[Scoring Rationale]: The student's answer gives the opposite spatial relation from the standard answer.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "KU_KnowledgeUpdate": "\n[Example 1]\n: What do I treat as my must-have now that I plan my meals around it?\n: brioche\n: sourdough\n: You now consider brioche your must-have bread that you plan meals around.\n\n[Scoring Rationale]: The student's answer is \"brioche\", which matches the current standard answer, not the outdated answer \"sourdough\".\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: What do I treat as my must-have now that I plan my meals around it?\n: brioche\n: sourdough\n: A versatile, high-quality protein source like chicken, beans, or tofu is a great starting point for building a week's worth of meals.\n\n[Scoring Rationale]: The student's answer mentions generic protein sources, which is neither the current answer \"brioche\" nor the outdated answer \"sourdough\". The answer is simply incorrect.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "TR_DurationComparison": "\n[Example 1]\n: Which of the following two durations is longer? Duration 1: The time I spent on Miller Hall. Duration 2: The time I spent on Music Festival Job. A. Duration 1 is longer. B. Duration 2 is longer.\n: B\n: B. Duration 2 is longer.\n\n[Scoring Rationale]: The student clearly selects \"B\", which matches the standard answer.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: Which of the following two durations is longer? Duration 1: The time I spent with the NutriPeak Nutrition. Duration 2: The time I spent with the BlendIt Pro Blender. A. Duration 1 is longer. B. Duration 2 is longer.\n: A\n: B. Duration 2 is longer.\n\n[Scoring Rationale]: The student selects \"B\", but the standard answer is \"A\". The selection is opposite.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "TR_OrderRanking": "\n[Example 1]\n: Rank the following cities by population from largest to smallest: Tokyo, Delhi, Shanghai.\n: Tokyo, Delhi, Shanghai\n: 1. Tokyo 2. Delhi 3. Shanghai\n\n[Scoring Rationale]: The student's sequence is Tokyo, Delhi, Shanghai, which exactly matches the standard answer.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: Rank events A, B, C in chronological order.\n: B, A, C\n: A, B, C\n\n[Scoring Rationale]: The student's sequence is A, B, C, but the standard answer is B, A, C. The order does not match.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "TR_DateExtraction": "\n[Example 1]\n: When did I complete the Scandinavian design course at Copenhagen Design Academy?\n: 2024/08/31\n: You completed the Scandinavian design course on August 31st, 2024.\n\n[Scoring Rationale]: The student's answer \"August 31st, 2024\" is an unambiguous equivalent of the standard answer \"2024/08/31\".\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: When did I visit Marutama Ra-men to taste their chicken broth?\n: 2023/01/20\n: January 20, 2021\n\n[Scoring Rationale]: The student's answer is \"January 20, 2021\", but the standard answer is \"2023/01/20\" — the year is wrong (2021 vs 2023).\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "MSR_YesNo": "\n[Example 1]\n: Is the cat in Mark and Jenny's 5-year anniversary post named Mittens?\n: Yes\n: Yes, the cat in Mark and Jenny's 5-year anniversary post is named Mittens.\n\n[Scoring Rationale]: The student's final answer is \"Yes\", which matches the standard answer \"Yes\". The additional explanation does not contradict the answer.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: Did the cat named after a crypto creator knock over the phone displaying the Binance alert?\n: Yes\n: The text does not explicitly state whether the cat knocked over the phone. It only mentions the cat was sitting next to the phone when it crashed.\n\n[Scoring Rationale]: The student does not give a clear \"Yes\". Instead, the student hedges and effectively refuses to commit to an answer. The standard answer is \"Yes\".\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "MSR_Counting": "", + "MSR_Arithmetic": "\n[Example 1]\n: How much total have I spent on coffee makers?\n: $260.00\n: Based on the conversation, you spent $110 on the first coffee maker and $150 on the second one, totaling $260.\n\n[Scoring Rationale]: The student's final answer is $260, which is numerically equivalent to the standard answer $260.00.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: How much total have I spent on coffee makers?\n: $260.00\n: I do not have access to your personal financial data, including your coffee maker purchases. I cannot provide you with a total amount spent on coffee makers.\n\n[Scoring Rationale]: The student refuses to answer and claims insufficient information, but this is not an AR task. The information was available in the conversation context.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "AR_AnswerRefusal": "\n[Example 1]\n: What is the model number printed on the power bank?\n: Insufficient information\n: I am sorry, but I cannot fulfill your request. The provided conversation history does not contain any information about a power bank or its model number.\n\n[Scoring Rationale]: The student correctly refuses to answer, stating the information is not available in the conversation. This is the expected behavior for an unanswerable question.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: How many vertical cables are visible on one side of the bridge?\n: Insufficient information\n: Three\n\n[Scoring Rationale]: The student provides a substantive answer (\"Three\") rather than refusing. For an unanswerable question, the correct response is to refuse or indicate insufficient information.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n" + } +} diff --git a/plugins/memlens-judge/src/judge_memlens/llm_judge.py b/plugins/memlens-judge/src/judge_memlens/llm_judge.py new file mode 100644 index 00000000..998e1685 --- /dev/null +++ b/plugins/memlens-judge/src/judge_memlens/llm_judge.py @@ -0,0 +1,72 @@ +"""Official MemLens task-specific judge as an isolated ReMe Job.""" + +from dataclasses import asdict + +from agentscope.message import UserMsg + +from reme.steps.base_step import BaseStep +from reme.steps.evolve.compressor import CompressorStep + +from .protocol import ( + SOURCE_REVISION, + judge_prompt, + normalize_for_judge, + parse_judge_response, + task_key, +) + + +class MemLensJudgeStep(BaseStep): + """Score saved answers without accessing the candidate's workspace or tools.""" + + async def execute(self): + """Apply the official judge prompt/parser to one saved prediction.""" + question = self.context["question"] + reference = self.context["reference"] + prediction = self.context["prediction"] + category = self.context["category"] + subtype = self.context.get("subtype", "") + normalized = normalize_for_judge(prediction) + prompt = judge_prompt( + question, + reference, + normalized["text"], + category, + subtype, + self.context.get("old_answer", ""), + ) + words = len(normalized["text"].split()) + result = { + "source_revision": SOURCE_REVISION, + "task_key": task_key(category, subtype, reference), + "original_len": normalized["original_len"], + "was_truncated": normalized["was_truncated"], + "is_degenerate": normalized["is_degenerate"], + **normalized["diagnostics"], + "parsed_words": words, + "auto_zero": words > 500, + } + raw, usage = "", None + if words > 500: + parsed = {"acc": 0.0, "parse_default_zero": False} + else: + response = await self.as_llm([UserMsg(name="user", content=prompt)]) + if hasattr(type(response), "__aiter__"): + last = None + async for chunk in response: + last = chunk + response = last + raw = await CompressorStep._response_text(response) # pylint: disable=protected-access + usage = asdict(response.usage) if response is not None and response.usage is not None else None + parsed = parse_judge_response(raw) + score = parsed["acc"] + self.context.response.success = True + self.context.response.answer = str(score) + self.context.response.metadata.update( + **result, + **parsed, + llm_judge_score=score, + raw_response=raw, + usage=usage, + ) + return self.context.response diff --git a/plugins/memlens-judge/src/judge_memlens/plugin.yaml b/plugins/memlens-judge/src/judge_memlens/plugin.yaml new file mode 100644 index 00000000..330566eb --- /dev/null +++ b/plugins/memlens-judge/src/judge_memlens/plugin.yaml @@ -0,0 +1,38 @@ +backends: + memlens_judge_step: judge_memlens.llm_judge:MemLensJudgeStep + +application_defaults: + jobs: + answer_judge: + backend: base + description: "memlens official judge protocol for a saved prediction" + parameters: + type: object + properties: + question: {type: string} + reference: {type: string} + prediction: {type: string} + category: {type: string, default: ""} + subtype: {type: string, default: ""} + old_answer: {type: string, default: ""} + required: [question, reference, prediction, category] + steps: + - backend: memlens_judge_step + as_llm: judge + components: + as_llm: + judge: + backend: openai + model: ${JUDGE_MODEL_NAME:-qwen3.8-max} + stream: false + max_retries: 3 + retry_delay: 60 + client_kwargs: {timeout: 600, max_retries: 0} + credential: + api_key: ${LLM_API_KEY:-} + base_url: ${LLM_BASE_URL:-} + extra_body: {enable_thinking: false} + parameters: + temperature: 0 + thinking_enable: false + max_tokens: 256 diff --git a/plugins/memlens-judge/src/judge_memlens/protocol.py b/plugins/memlens-judge/src/judge_memlens/protocol.py new file mode 100644 index 00000000..19ef3039 --- /dev/null +++ b/plugins/memlens-judge/src/judge_memlens/protocol.py @@ -0,0 +1,139 @@ +"""MemLens's task-specific binary judge for fixed saved predictions. + +Prompts, answer preprocessing and verdict parsing follow SOURCE_REVISION. +Copyright (c) 2026 Xiyu Ren, Zhaowei Wang, and the MemLens Authors. +The complete MIT notice accompanies the verbatim _memlens_judge_prompts.json. +Only exhausted API errors differ: they remain coverage failures, not wrong answers. +""" + +from functools import lru_cache +import json +from pathlib import Path +import re + +SOURCE_REVISION = "77f3ab9a52fa2d6a17978e2dffe80438a4ecced2" + +PROTOCOL_ID = "memlens-task-binary-api-qwen-v1" +_HEDGING_MARKERS = ( + "wait", + "let me re-read", + "let me re read", + "let me double check", + "let me double-check", + "let me reconsider", + "let me look again", + "let's reconsider", + "let s reconsider", + "re-evaluating", + "re evaluating", + "actually looking", + "actually let me", + "hold on", + "let me look at", + "let me check", + "let's re-read", + "let s re-read", + "i need to reconsider", + "wait let me", +) + + +@lru_cache(maxsize=1) +def _prompts() -> dict: + """Load the checked-in, verbatim official prompt strings, never a network resource.""" + return json.loads(Path(__file__).with_name("_memlens_judge_prompts.json").read_text(encoding="utf-8")) + + +def task_key(category: str, subtype: str, reference: str) -> str: + """Official question-type/subtype routing, including reference-based MSR choices.""" + if category == "multi_session_reasoning": + if subtype == "arithmetic": + return "MSR_Arithmetic" + if subtype == "entity_resolution" and reference.strip() in {"Yes", "No"}: + return "MSR_YesNo" + return "MSR_Counting" + if category == "temporal_reasoning": + return { + "duration_comparison": "TR_DurationComparison", + "order_ranking": "TR_OrderRanking", + }.get(subtype, "TR_DateExtraction") + fixed = {"knowledge_update": "KU_KnowledgeUpdate", "answer_refusal": "AR_AnswerRefusal"} + if category in fixed: + return fixed[category] + if category == "information_extraction": + return "IE_PreviousInfo" if subtype == "previnfo" else "IE_Entity" + raise ValueError(f"Unsupported MemLens judge category: {category}") + + +def judge_prompt( + question: str, + reference: str, + prediction: str, + category: str, + subtype: str = "", + old_answer: str = "", +) -> str: + """Render the exact official rubric/examples and current-case fields.""" + if any(not isinstance(value, str) for value in (question, reference, prediction, category, subtype, old_answer)): + raise ValueError("MemLens judge inputs must be strings") + key = task_key(category, subtype, reference) + prompts = _prompts() + prompt = prompts["template"].format(task_criteria=prompts["criteria"][key], examples=prompts["examples"][key]) + if key == "AR_AnswerRefusal": + reference = "Insufficient information" + old = f": {old_answer}\n" if key == "KU_KnowledgeUpdate" and old_answer else "" + return ( + f"{prompt}[Current Case]\n: {question}\n: {reference}\n" + f"{old}: {prediction}\n\n[Scoring Rationale]:" + ) + + +def normalize_for_judge(text: str) -> dict: + """Preserve native parsing, circularity diagnostics and final-30% tail selection.""" + if not text: + return {"text": "", "original_len": 0, "was_truncated": False, "is_degenerate": False, "diagnostics": {}} + original_len = len(text) + text = parse_model_output(text) + hedging, ratio = 0, 1.0 + if len(text) >= 500: + lower = text.lower() + hedging = sum(lower.count(marker) for marker in _HEDGING_MARKERS) + words = lower.split() + ngrams = [tuple(words[index : index + 5]) for index in range(len(words) - 4)] + ratio = len(set(ngrams)) / len(ngrams) if ngrams else 1.0 + truncated = len(text) > 3000 and len(text) > 6000 + if truncated: + text = "[...truncated earlier reasoning...]\n" + text[int(len(text) * (1 - 0.3)) :] + return { + "text": text, + "original_len": original_len, + "was_truncated": truncated, + "is_degenerate": hedging >= 6 or ratio < 0.40, + "diagnostics": {"hedging_count": hedging, "unique_5gram_ratio": round(ratio, 3)}, + } + + +def parse_judge_response(text: str) -> dict: + """Official first-match priority/clamping; even unparseable outputs natively score zero.""" + patterns = ( + (r'"answer_score"\s*:\s*(\d+)', 0), + (r"deserves\s+(\d+)\s+point", re.IGNORECASE), + (r"\[Score\]\s*:\s*(\d+)", 0), + ) + for pattern, flags in patterns: + match = re.search(pattern, text, flags) if text else None + if match: + return {"acc": float(min(int(match.group(1)), 1)), "parse_default_zero": False} + return {"acc": 0.0, "parse_default_zero": True} + + +def parse_model_output(text: str) -> str: + """Native special-token/think-trace parsing; retained raw outputs remain separate.""" + if not text: + return "" + if "" in text: + text = text.split("", 1)[1] + text = re.sub(r"<\|[^>]*\|>", "", text) + text = re.sub(r"<\|begin_of_box\|>(.*?)<\|end_of_box\|>", r"\1", text, flags=re.DOTALL) + text = text.replace("<|begin_of_box|>", "").replace("<|end_of_box|>", "") + return text.replace("[CONTENT_FILTER_REJECTED]", "").strip() diff --git a/plugins/memlens/README.md b/plugins/memlens/README.md new file mode 100644 index 00000000..81e982c4 --- /dev/null +++ b/plugins/memlens/README.md @@ -0,0 +1,48 @@ +# MemLens plugin + +Install with `pip install -e plugins/memlens` and load the `memlens-benchmark` configuration +with the `memlens` plugin. The executable ReMe Image evaluation lives +in [`benchmark/memlens`](../../benchmark/memlens). + +`auto_memory` extends native Auto Memory with a MemLens-specific system prompt. +It retains native source-session persistence and interleaved image inputs; `include_images` +selects whether the memory model sees historical images. No question or reference answer +is given during ingestion. + +`agentic_answer` accepts the official answer prompt as `query`, optional `query_time`, +and native image DataBlocks as `query_images`. It creates a fresh AgentScope Agent for +each question with BM25 `search` and scoped `read`. Question images remain visible in +both the text-memory and image-memory settings. +Its `max_iters` input defaults to 30 ReAct iterations and applies only +to the current answer Agent. The effective value is returned in answer metadata; +Auto Memory and context compression keep their existing settings. + +`include_source_sessions` defaults to `true`: memory prompts include exact physical +JSONL line references, and the answering Agent can search/read those source sessions. +When false, only generated daily/digest markdown is searchable and readable, regardless +of whether an old session index exists. Link expansion is disabled. JSONL persistence +is unchanged. The runner must also omit session directories from index updates for +the source-disabled setting. + +With `include_images=true`, native Auto Memory saves input images beneath +`session/images/` and supplies their paths to the memory Agent. The answer tool +`read_image(path)` calls ReMe's existing `read_image` Job and converts its successful +image response to an AgentScope DataBlock. The plugin supplies the native Job with +`max_bytes=52428800` (50 MiB). Image inputs retain their original bytes and resolution. +Native size errors and non-image results +stay textual tool errors; the generic ReMe Job is unchanged. + +Image access is independent of `include_source_sessions`: a source-disabled Agent +can still reopen images cited by its memory notes. ReMe Text has no historical-image +tool. All image paths are confined to the current scene's session image directory. +The runner retains official Image IDs as text beside the corresponding image blocks; +the IDs are answer evidence labels, not a separate registry or tool argument. +Do not put answer keys, judge results, another scene's data, or agent state under the +daily/digest/source evidence directories. Every scene starts from a separate workspace; +sessions within one scene accumulate in order. Query state and search deduplication +are independent even when questions execute concurrently. + +The plugin preserves the ReMe Image experiment's thinking, output budgets and context +settings. The runner selects the candidate model; scoring belongs to `memlens-judge`. +Answer metadata includes portable usage and the raw final-message usage when available; +missing provider cache/reasoning counters are not reported as zero. diff --git a/plugins/memlens/pyproject.toml b/plugins/memlens/pyproject.toml new file mode 100644 index 00000000..2bc97530 --- /dev/null +++ b/plugins/memlens/pyproject.toml @@ -0,0 +1,25 @@ +[project] +name = "reme-memlens" +version = "0.1.0" +description = "MemLens multimodal memory benchmark plugin for ReMe." +readme = "README.md" +license = "Apache-2.0" +requires-python = ">=3.11" +dependencies = ["reme-ai[as]>=0.4.1.11", "pillow>=10.0.0"] + +[project.entry-points."reme.plugins"] +memlens = "reme_memlens" + +[project.entry-points."reme.configs"] +memlens-benchmark = "reme_memlens:BENCHMARK_CONFIG" + +[tool.setuptools.packages.find] +where = ["src"] +include = ["reme_memlens*"] + +[tool.setuptools.package-data] +reme_memlens = ["plugin.yaml", "*.yaml"] + +[build-system] +requires = ["setuptools>=77", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/plugins/memlens/src/reme_memlens/__init__.py b/plugins/memlens/src/reme_memlens/__init__.py new file mode 100644 index 00000000..a81a2c73 --- /dev/null +++ b/plugins/memlens/src/reme_memlens/__init__.py @@ -0,0 +1,5 @@ +"""MemLens benchmark Jobs for ReMe.""" + +from pathlib import Path + +BENCHMARK_CONFIG = Path(__file__).with_name("benchmark.yaml") diff --git a/plugins/memlens/src/reme_memlens/agentic_answer.py b/plugins/memlens/src/reme_memlens/agentic_answer.py new file mode 100644 index 00000000..2115d95d --- /dev/null +++ b/plugins/memlens/src/reme_memlens/agentic_answer.py @@ -0,0 +1,136 @@ +"""Agentic BM25 answer with scene-local, optional multimodal source access.""" + +from pathlib import Path +from uuid import uuid4 + +from agentscope.message import Base64Source, DataBlock, TextBlock, ToolResultState, UserMsg +from agentscope.tool import FunctionTool, ToolChunk, Toolkit + +from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper +from reme.steps.base_step import BaseStep +from reme.steps.file_io._path import _check_path_permission, resolve_path + + +def evidence_paths(workspace: Path, daily_dir: str, digest_dir: str, session_dir: str, sources: bool) -> list[str]: + """List exact allowed evidence files, never agent state, settings, or another workspace.""" + roots = [(daily_dir, "*.md"), (digest_dir, "*.md")] + if sources: + roots.append((f"{session_dir}/dialog", "*.jsonl")) + workspace = workspace.resolve() + paths = [] + for directory, pattern in roots: + root = (workspace / directory).resolve() + if not root.is_relative_to(workspace): + raise ValueError("Evidence directory must stay inside this scene workspace") + for path in sorted(root.rglob(pattern)): + if path.is_file() and path.resolve().is_relative_to(root): + paths.append(path.relative_to(workspace).as_posix()) + return paths + + +def image_reader(step: BaseStep) -> FunctionTool: + """Bridge the native read_image Job response to an AgentScope image result.""" + workspace = step.workspace_path.resolve() + image_directory = f"{step.config_value('session_dir')}/images" + + async def read_image(path: str) -> ToolChunk: + """View an image using its session/images path cited in a memory note.""" + target, error = resolve_path(workspace, path) + if error or target is None or not _check_path_permission(workspace, target, [image_directory]): + return ToolChunk( + content=[TextBlock(text=error or "Image path must stay inside this scene's session image directory.")], + state=ToolResultState.ERROR, + ) + relative = target.relative_to(workspace).as_posix() + response = await step.run_job("read_image", path=relative) + metadata = response.metadata + media_type = metadata.get("mime", "") + if ( + not response.success + or metadata.get("oversized") + or metadata.get("non_image_warning") + or not media_type.startswith("image/") + ): + notice = str(response.answer) if not response.success or metadata.get("oversized") else "Not an image file." + return ToolChunk(content=[TextBlock(text=notice)], state=ToolResultState.ERROR) + image = DataBlock(source=Base64Source(media_type=media_type, data=response.answer)) + return ToolChunk(content=[TextBlock(text=f"Image source: {relative}"), image], state=ToolResultState.SUCCESS) + + return FunctionTool(func=read_image, name="read_image", is_concurrency_safe=True, is_read_only=True) + + +class MemLensAgenticAnswerStep(BaseStep): + """Run a fresh Agent for each question, with the official answer prompt unchanged.""" + + async def execute(self): + """Answer through a fresh, scene-scoped Agent and its allowed tools.""" + wrapper = self.agent_wrapper + if not isinstance(wrapper, AsAgentWrapper): + raise NotImplementedError("MemLens agentic_answer requires the AgentScope wrapper") + query = self.context.get("query", "") + if not query: + raise ValueError("query must contain the benchmark answer prompt") + max_iters = self.context.get("max_iters", 30) + if isinstance(max_iters, bool) or not isinstance(max_iters, int) or max_iters < 1: + raise ValueError("max_iters must be a positive integer") + sources = bool(self.context.get("include_source_sessions", True)) + images = bool(self.context.get("include_images", False)) + paths = evidence_paths( + self.workspace_path, + self.config_value("daily_dir"), + self.config_value("digest_dir"), + self.config_value("session_dir"), + sources, + ) + tool_context_id = f"memlens_answer_{uuid4().hex}" + tools = [] + if images: + tools.append(image_reader(self)) + injected = { + "_allowed_paths": paths, + "read_step_format_session": True, + "search_filter": {"paths": paths}, + "vector_weight": 0.0, + } + # Standard ReMe job tools keep their normal file handling and lifecycle. + tools.extend( + wrapper._make_tool(job, tool_context_id, injected) # pylint: disable=protected-access + for job in wrapper._resolve_job_tools(["search", "read"]) # pylint: disable=protected-access + ) + blocks = [TextBlock(text=query)] + blocks.extend(DataBlock.model_validate(block) for block in self.context.get("query_images", [])) + if any(not block.source.media_type.startswith("image/") for block in blocks[1:]): + raise ValueError("query_images accepts image DataBlocks only") + prompt = self.get_prompt("system_prompt") + if query_time := self.context.get("query_time"): + prompt += f"\nQuestion time: {query_time}" + if images: + prompt += ( + "\nUse read_image(path=...) for the session/images paths cited in memory, including source_images." + ) + else: + prompt += "\nHistory images cannot be reopened in this setting." + try: + result = await wrapper.reply( + UserMsg(name="user", content=blocks), + system_prompt=prompt, + toolkit=Toolkit(tools=tools), + job_tools=[], + builtin_tools=False, + session_id=str(uuid4()), + resume="", + react_config={"max_iters": max_iters}, + ) + self.context.response.answer = (result.get("result") or "").strip() + self.context.response.metadata.update( + usage=result.get("usage"), + agentscope_usage=(result.get("last_message") or {}).get("usage"), + agent_session_id=result.get("session_id"), + include_source_sessions=sources, + include_images=images, + max_iters=max_iters, + ) + finally: + if self.app_context is not None: + self.app_context.metadata.get("tool_contexts", {}).pop(tool_context_id, None) + return self.context.response diff --git a/plugins/memlens/src/reme_memlens/agentic_answer.yaml b/plugins/memlens/src/reme_memlens/agentic_answer.yaml new file mode 100644 index 00000000..42f862e3 --- /dev/null +++ b/plugins/memlens/src/reme_memlens/agentic_answer.yaml @@ -0,0 +1,10 @@ +system_prompt: | + Answer the supplied benchmark question by retrieving evidence from this scene's memory. + Use search to find relevant facts, then read exact files or line ranges when more context is needed. + Search again with different names, dates or distinctive visual attributes if the evidence is incomplete. + Keep different people, times and visually similar objects distinct. Source line references point to + the original conversation, when source access is enabled. If read_image is available, use the image + path cited in the memory note or its source_images field to verify the actual historical image. + Preserve Image IDs exactly when the question asks which image. Do not invent IDs or unseen details. + Follow the question's supplied answer-format instructions and answer only from supported evidence. + Do not include the retrieval process or tool transcript in the final answer. diff --git a/plugins/memlens/src/reme_memlens/auto_memory.py b/plugins/memlens/src/reme_memlens/auto_memory.py new file mode 100644 index 00000000..3f473536 --- /dev/null +++ b/plugins/memlens/src/reme_memlens/auto_memory.py @@ -0,0 +1,34 @@ +"""MemLens memory prompts with optional, exact source-line references.""" + +from agentscope.message import Msg + +from reme.steps.evolve.auto_memory import AutoMemoryStep + + +class MemLensAutoMemoryStep(AutoMemoryStep): + """Keep core persistence and multimodal input, adding benchmark instructions only.""" + + def _format_history(self, messages: list[Msg]) -> str: + if not self.context.get("include_source_sessions", True): + return "Source-session access is disabled; record self-contained facts.\n\n" + super()._format_history( + messages, + ) + + # Core execute has saved the session already. Read physical line numbers + # rather than assuming input order equals timestamp-sorted JSONL order. + session_id = self.context.get("session_id", "") + source = self._session_source_path(session_id) + # File iteration keeps Unicode line separators inside JSON strings intact. + with self._session_path(session_id).open(encoding="utf-8") as stream: + lines = {Msg.model_validate_json(line).id: number for number, line in enumerate(stream, 1) if line.strip()} + turns = [ + f"[L{lines[msg.id]} | {msg.name or msg.role} @ {msg.created_at}]\n{msg.get_text_content() or ''}" + for msg in messages + ] + return ( + f"Source file: {source}\n" + f"Cite useful evidence with the full source path, for example [[{source}#L1-L2]], " + "replacing the line range with the actual physical lines below. Never use bare links such as [[L5]]. " + "Preserve original Image IDs alongside their facts, and retain the supplied Image source paths.\n\n" + + "\n\n".join(turns) + ) diff --git a/plugins/memlens/src/reme_memlens/auto_memory.yaml b/plugins/memlens/src/reme_memlens/auto_memory.yaml new file mode 100644 index 00000000..f318b669 --- /dev/null +++ b/plugins/memlens/src/reme_memlens/auto_memory.yaml @@ -0,0 +1,25 @@ +system_prompt: | + You maintain long-horizon personal multimodal memory across dated conversations. Use tools to save useful evidence, + not a response to any future question. Retain concrete people, entities, events, dates, preferences, decisions, + quantities, places and relationships. Preserve changes over time and distinguish speakers and similar objects. + Track longitudinal updates to preferences, plans, possessions, health, work and relationships without erasing + earlier dated facts. Keep which person experienced each event explicit and resolve relative dates using the + session timestamp only when justified. Separate repeated similar events instead of conflating them. + + When images are present, jointly understand them and the surrounding conversation. Record distinctive visual + details (objects, attributes, layout, visible text, activities) that make an image identifiable. Keep each original + Image ID verbatim next to the corresponding visual facts; an ID is a reference, not evidence of image contents. + Do not invent visual details when no image is provided. Do not substitute generic summaries for specific facts. + + When physical source-line labels are supplied, cite the useful source ranges with the exact session path and + line numbers. Without source access, make the memory self-contained rather than adding source references. + Organize related facts under descriptive headings, use absolute dates when supported, preserve exact names + and numbers, and distinguish facts from uncertainty. Resolve pronouns only when the conversation supports it. + + Follow the normal daily-note contract: a stable filename stem in name, an informative description, and the + factual body. Never set status. Invoke the provided file tools as structured tool calls to actually save memory. + + Complete the task, then stop: create at most one note for a new session. After a successful daily_write, + return the requested one-sentence summary without calling daily_write again or creating a renamed duplicate. + Further writes are needed only to correct an explicit tool error or an identified omission, not to repeat a + successful save. For an existing note, finish after the necessary edits and frontmatter update succeed. diff --git a/plugins/memlens/src/reme_memlens/benchmark.yaml b/plugins/memlens/src/reme_memlens/benchmark.yaml new file mode 100644 index 00000000..1421f4fc --- /dev/null +++ b/plugins/memlens/src/reme_memlens/benchmark.yaml @@ -0,0 +1,30 @@ +extends: benchmark +plugins: [memlens] +components: + as_llm: + default: &memory_model + backend: openai + model: ${LLM_MODEL_NAME:-qwen3.8-max} + stream: false + context_size: 200000 + max_retries: 3 + retry_delay: 60 + client_kwargs: {timeout: 600, max_retries: 0} + extra_body: {enable_thinking: true} + parameters: {temperature: 0, thinking_enable: true, max_tokens: 65536} + bench: + <<: *memory_model + context_size: 400000 + parameters: {temperature: 0, thinking_enable: true, max_tokens: 16384} + compressor: + <<: *memory_model + agent_wrapper: + default: &memory_agent + react_config: {max_iters: 30} + context_config: {trigger_ratio: 0.8, reserve_ratio: 0.1, tool_result_limit: 50000, max_image_num: 4096} + model_config: {max_retries: 0} + bench: + <<: *memory_agent + file_store: + default: + embedding_store: null diff --git a/plugins/memlens/src/reme_memlens/plugin.yaml b/plugins/memlens/src/reme_memlens/plugin.yaml new file mode 100644 index 00000000..56a6f64e --- /dev/null +++ b/plugins/memlens/src/reme_memlens/plugin.yaml @@ -0,0 +1,75 @@ +backends: + memlens_auto_memory_step: reme_memlens.auto_memory:MemLensAutoMemoryStep + memlens_agentic_answer_step: reme_memlens.agentic_answer:MemLensAgenticAnswerStep + +application_defaults: + jobs: + read_image: + backend: base + description: "Read a saved session image through the native ReMe image Job." + parameters: + type: object + properties: + path: {type: string, description: "Workspace-relative image attachment path."} + required: [path] + steps: + - backend: read_image_step + max_bytes: 52428800 + auto_memory: + backend: base + description: "Record one MemLens session using native Auto Memory." + include_source_sessions: true + parameters: + type: object + properties: + messages: + type: array + items: {type: object} + session_id: {type: string} + date: {type: string, default: ""} + include_images: {type: boolean, default: false} + include_source_sessions: {type: boolean, default: true} + memory_hint: {type: string, default: ""} + required: [messages, session_id] + steps: + - backend: memlens_auto_memory_step + search: + backend: base + description: "Search this scene's allowed memory and source text using BM25." + parameters: + type: object + properties: + query: {type: string} + limit: {type: integer, description: "Number of search results."} + start_date: {type: string} + end_date: {type: string} + required: [query] + steps: + - backend: search_step + vector_weight: 0.0 + expand_links: false + agentic_answer: + backend: base + description: "Answer a MemLens query through a fresh multimodal retrieval Agent." + include_source_sessions: true + parameters: + type: object + properties: + query: {type: string, description: "Official benchmark answer prompt including the query."} + query_time: {type: string, default: ""} + query_images: + type: array + items: {type: object} + default: [] + scene_id: {type: string} + max_iters: {type: integer, minimum: 1, default: 30, description: "Maximum answer Agent ReAct iterations."} + include_images: {type: boolean, default: false} + include_source_sessions: {type: boolean, default: true} + required: [query, scene_id] + steps: + - backend: memlens_agentic_answer_step + agent_wrapper: bench + components: + file_store: + default: + embedding_store: null diff --git a/tests/unit/test_memlens_benchmark.py b/tests/unit/test_memlens_benchmark.py new file mode 100644 index 00000000..1a461ca5 --- /dev/null +++ b/tests/unit/test_memlens_benchmark.py @@ -0,0 +1,335 @@ +"""Focused offline benchmark contracts; no external services or experiment writes.""" + +# pylint: disable=missing-function-docstring,protected-access +import asyncio +import base64 +from importlib.metadata import EntryPoint +from io import BytesIO +import json +from pathlib import Path +import tomllib +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from agentscope.formatter import OpenAIChatFormatter +from agentscope.message import AssistantMsg, Base64Source, DataBlock, Msg, TextBlock, ToolResultBlock, ToolResultState +from PIL import Image +import pytest + +from reme.components import R +from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper +from reme.components.component_registry import create_application_registry +from reme.components.file_store import LocalFileStore +from reme.plugin import PluginManager, _load_plugin +from reme.schema import ApplicationConfig, Response +from reme.steps.evolve.auto_memory import _sanitize_msg_for_save +from reme.steps.file_io.read_image import ReadImageStep + +from benchmark.memlens import run as memlens + + +@pytest.fixture(autouse=True) +def memlens_judge_path(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins/memlens-judge/src")) + + +def test_candidate_declares_required_image_dependency(): + manifest = Path(__file__).resolve().parents[2] / "plugins/memlens/pyproject.toml" + project = tomllib.loads(manifest.read_text(encoding="utf-8"))["project"] + assert "pillow>=10.0.0" in project["dependencies"] + + +def test_memlens_subset_lengths_order_and_public_inputs(tmp_path): + image = {"file": "photo.png", "blip_caption": "HIDDEN_CAPTION"} + row = { + "question_type": "knowledge_update", + "question_subtype": "update", + "question": "What changed?", + "answer": "HIDDEN_GOLD", + "old_answer": "HIDDEN_OLD", + "question_date": "2024/05/31 (Fri) 07:58", + "haystack_dates": ["2024/05/01 (Wed) 08:00"], + "haystack_session_ids": ["original-session"], + "haystack_sessions": [ + [ + {"role": "user", "content": "Before after", "images": [image]}, + {"role": "assistant", "content": "Original reply", "images": [], "has_answer": True}, + ], + ], + "answer_session_ids": ["HIDDEN_EVIDENCE"], + } + rows = [{**row, "question_id": f"q_{index:03d}"} for index in range(196)] + ids = [item["question_id"] for item in rows[:195]] + (tmp_path / "agent_subset_195.json").write_text(json.dumps({"n_questions": 195, "question_ids": ids[::-1]})) + for length in ("128k", "256k"): + (tmp_path / f"dataset_{length}.json").write_text( + json.dumps([{**item, "haystack_session_ids": [f"session-{length}"]} for item in rows]), + ) + + image_dir = tmp_path / "release_images" + image_dir.mkdir() + Image.new("RGB", (1, 1), (12, 34, 56)).save(image_dir / "photo.png") + scenes = memlens.load_memlens(tmp_path, "256k") + assert [scene["id"] for scene in scenes] == ids + first = scenes[0] + assert first["sessions"][0]["id"] == "session-256k" + assert memlens.load_memlens(tmp_path, "128k")[0]["sessions"][0]["id"] == "session-128k" + assert len(first["questions"]) == 1 and first["questions"][0]["id"] == first["id"] + assert first["sessions"][0]["timestamp"] == "2024-05-01T08:00" + messages = memlens.session_messages(first["sessions"][0]) + assert [message["role"] for message in messages] == ["user", "assistant"] + assert [block["type"] for block in messages[0]["content"]] == ["text", "text", "data", "text"] + assert messages[0]["content"][0]["text"] == "Before " + assert messages[0]["content"][1]["text"] == "[Image ID: photo]" + assert messages[0]["content"][-1]["text"] == " after" + assert base64.b64decode(messages[0]["content"][2]["source"]["data"]) == (image_dir / "photo.png").read_bytes() + assert memlens.session_messages(scenes[1]["sessions"][0])[0]["id"] != messages[0]["id"] + public = json.dumps(messages) + memlens.query_prompt(first["questions"][0], "Evidence tools") + assert "HIDDEN" not in public and "has_answer" not in public + rows[0]["haystack_sessions"][0][0]["images"][0]["file"] = "../private.png" + (tmp_path / "dataset_256k.json").write_text(json.dumps(rows)) + with pytest.raises(ValueError, match="MemLens"): + memlens.load_memlens(tmp_path, "256k") + + +def test_official_metric_normalization(): + question = {"answer": "$460.00"} + assert memlens.score_answer(question, "ignore$460.00<|im_end|>") == {"sub_em": 1.0, "f1": 1.0} + + +@pytest.fixture(name="candidate") +def fixture_candidate(tmp_path, monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins/memlens/src")) + monkeypatch.chdir(tmp_path) + from reme_memlens import agentic_answer as answer, auto_memory as memory + + app = SimpleNamespace( + registry=R, + metadata={}, + app_config=ApplicationConfig(workspace_dir=str(tmp_path)), + jobs={}, + components={}, + ) + wrapper = AsAgentWrapper(backend="agentscope", as_llm="", app_context=app) + wrapper.reply = AsyncMock(return_value={"result": "answer", "session_id": "new", "usage": {"input_tokens": 7}}) + store = LocalFileStore(embedding_store="", app_context=app) + return memory, answer, app, wrapper, store, tmp_path + + +def test_candidate_registration_is_application_local(candidate): + memory, answer, *_ = candidate + plugin = _load_plugin("memlens", EntryPoint(name="memlens", value="reme_memlens", group="reme.plugins")) + registry = create_application_registry() + PluginManager([plugin]).register(registry) + assert registry.get("step", "memlens_auto_memory_step") is memory.MemLensAutoMemoryStep + assert registry.get("step", "memlens_agentic_answer_step") is answer.MemLensAgenticAnswerStep + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sources", [True, False]) +async def test_memory_saves_originals_and_native_image_links(candidate, sources): + memory, _, app, wrapper, store, workspace = candidate + step = memory.MemLensAutoMemoryStep(app_context=app, agent_wrapper=wrapper, file_store=store) + step._list_session_note = AsyncMock(return_value=None) + picture = BytesIO() + Image.new("RGB", (1, 1), (12, 34, 56)).save(picture, format="PNG") + image_bytes = picture.getvalue() + image_data = base64.b64encode(image_bytes).decode() + late = Msg.model_validate( + { + "id": "late", + "name": "Alice", + "role": "user", + "created_at": "2026-01-02T10:00:00", + "content": [ + TextBlock(text="Image ID: image-7"), + DataBlock(source=Base64Source(media_type="image/png", data=image_data)), + TextBlock(text="Before.\u2028\u2028After.\u2029End.\u0085Done."), + ], + }, + ) + early = late.model_copy(update={"id": "early", "created_at": "2026-01-01T10:00:00"}, deep=True) + before = late.model_dump() + await step(messages=[late, early], session_id="one", include_images=True, include_source_sessions=sources) + saved = workspace / "session/dialog/one.jsonl" + assert saved.read_text() == "".join(_sanitize_msg_for_save(msg).model_dump_json() + "\n" for msg in [early, late]) + inputs = wrapper.reply.call_args.args[0] + assert [block.type for block in inputs.content] == ["text", "data", "text", "data", "text"] + assert [block.source.data for block in inputs.content if isinstance(block, DataBlock)] == [image_data] * 2 + text = inputs.get_text_content() + assert "Image ID: image-7" in text and "Before.\u2028\u2028After.\u2029End.\u0085Done." in text + assert ("[L2 | Alice" in text) is sources + files = sorted((workspace / "session/images/one").glob("*")) + assert len(files) == 2 and all(path.read_bytes() == image_bytes for path in files) + assert all(f"[[{path.relative_to(workspace).as_posix()}]]" in text for path in files) + assert late.model_dump() == before + + +@pytest.mark.asyncio +async def test_read_image_returns_real_formatter_image_and_rejects_escape(candidate): + _, answer, app, wrapper, store, workspace = candidate + path = workspace / "session/images/one/photo.png" + path.parent.mkdir(parents=True) + Image.new("RGB", (2, 2), (12, 34, 56)).save(path) + + async def native_read(**kwargs): + step = ReadImageStep(app_context=app) + await step(**kwargs) + return step.context.response + + app.jobs["read_image"] = AsyncMock(side_effect=native_read) + step = answer.MemLensAgenticAnswerStep(app_context=app, agent_wrapper=wrapper, file_store=store) + tool = answer.image_reader(step) + result = await tool.call(path="session/images/one/photo.png") + assert result.state == ToolResultState.SUCCESS and isinstance(result.content[1], DataBlock) + assert base64.b64decode(result.content[1].source.data) == path.read_bytes() + formatted = await OpenAIChatFormatter().format( + [ + AssistantMsg(name="tool", content=[ToolResultBlock(id="call", name="read_image", output=result.content)]), + ], + ) + assert any( + block["type"] == "image_url" + for msg in formatted + if isinstance(msg.get("content"), list) + for block in msg["content"] + ) + assert (await tool.call(path="../other-scene.png")).state == ToolResultState.ERROR + assert (await tool.call(path="mem_session/private.png")).state == ToolResultState.ERROR + app.jobs["read_image"].assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sources", [True, False]) +async def test_answer_has_fresh_state_bm25_scope_and_no_gold(candidate, sources): + _, answer, app, wrapper, store, workspace = candidate + for relative in ["daily/note.md", "digest/note.md", "session/dialog/one.jsonl", "mem_session/private.md"]: + path = workspace / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("evidence") + assert answer.evidence_paths(workspace, "daily", "digest", "session", False) == ["daily/note.md", "digest/note.md"] + calls = [] + + async def search(**kwargs): + calls.append(kwargs) + return Response(answer="evidence") + + job = AsyncMock(side_effect=search) + job.name, job.description, job.parameters = ( + "search", + "search", + { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + ) + read = AsyncMock(return_value=Response(answer="evidence")) + read.name, read.description, read.parameters = "read", "read", {"type": "object", "properties": {}} + app.jobs = {"search": job, "read": read} + step = answer.MemLensAgenticAnswerStep(app_context=app, agent_wrapper=wrapper, file_store=store) + sessions, contexts = [], [] + for _ in range(2): + response = await step(query="Official question", include_source_sessions=sources) + options = wrapper.reply.call_args.kwargs + assert wrapper.reply.call_args.args[0].get_text_content() == "Official question" + assert options["resume"] == "" and options["builtin_tools"] is False and options["job_tools"] == [] + assert options["react_config"] == {"max_iters": 30} + sessions.append(options["session_id"]) + await (await options["toolkit"].get_tool("search")).call(query="evidence") + assert calls[-1]["vector_weight"] == 0.0 + assert ("session/dialog/one.jsonl" in calls[-1]["search_filter"]["paths"]) is sources + assert "mem_session/private.md" not in calls[-1]["_allowed_paths"] + contexts.append(calls[-1]["tool_context_id"]) + assert response.metadata["usage"] == {"input_tokens": 7} + assert len(set(sessions)) == len(set(contexts)) == 2 + + +@pytest.mark.asyncio +async def test_image_only_pipeline_preserves_sessions_and_official_query(candidate, monkeypatch): + import reme + from benchmark.memlens.run import normalized_usage, total_usage + + runner = memlens + *_, workspace = candidate + + query = { + "id": "q", + "text": "What changed?", + "category": "knowledge_update", + "question_date": "2024/05/31 (Fri) 07:58", + "answer": "HIDDEN_GOLD", + "subtype": "update", + "old_answer": "HIDDEN_OLD", + } + scene = { + "id": "one", + "questions": [query], + "sessions": [{"id": name, "timestamp": "2024-05-01T08:00"} for name in ("first", "second")], + } + calls = [] + + class Candidate: + """Offline native-Job boundary; no provider calls.""" + + def __init__(self, **_config): + self.context = SimpleNamespace( + components={ + "agent_wrapper": {name: SimpleNamespace(session_path=workspace) for name in ("default", "bench")}, + }, + ) + self.closed = False + + async def start(self): + pass + + async def close(self): + self.closed = True + + async def run_job(self, name, **kwargs): + calls.append((name, kwargs)) + return Response(answer="Retained answer", metadata={"agent_session_id": "new"}) + + app = Candidate() + judge = SimpleNamespace(run_job=AsyncMock(return_value=Response(metadata={"llm_judge_score": 1}))) + monkeypatch.setattr(reme, "Application", lambda **_: app) + monkeypatch.setattr(runner, "app_config", lambda *_, **__: {}) + monkeypatch.setattr(runner, "session_messages", lambda session: [{"text": session["id"]}]) + monkeypatch.setattr(runner, "observe_usage", lambda _: None) + rows, *_ = await runner.evaluate_scene(scene, 0, "candidate", workspace, judge, asyncio.Semaphore(2)) + assert [name for name, _ in calls] == ["auto_memory", "auto_memory", "index_update", "agentic_answer"] + assert [kw["session_id"] for name, kw in calls if name == "auto_memory"] == ["first", "second"] + for name, kwargs in calls: + assert "HIDDEN" not in json.dumps(kwargs) + if name != "index_update": + assert kwargs["include_images"] is kwargs["include_source_sessions"] is True + inputs = calls[-1][1] + assert inputs["max_iters"] == 30 and "category" not in inputs + assert inputs["query"] == memlens.query_prompt(query, "Use the available tools to retrieve relevant evidence.") + assert judge.run_job.call_args.kwargs["reference"] == "HIDDEN_GOLD" + assert rows[0]["metrics"]["llm_judge"] == 1 and app.closed + measured = normalized_usage( + { + "prompt_tokens": 100, + "completion_tokens": 20, + "prompt_tokens_details": {"cached_tokens": 30}, + }, + ) + assert measured["noncached_input_tokens"] == 70 + assert total_usage([measured])["total_tokens"] == 120 + assert normalized_usage({})["cached_input_tokens"] is None + + +@pytest.mark.asyncio +async def test_observation_preserves_native_reply_and_rejects_only_interrupted(): + from benchmark.memlens.run import observe_usage + + result = {"result": "", "last_message": {"finished_reason": "exceed_max_iters", "error": None}} + reply = AsyncMock(return_value=result) + wrapper = SimpleNamespace(reply=reply) + app = SimpleNamespace(context=SimpleNamespace(components={"agent_wrapper": {"default": wrapper}})) + observe_usage(app) + assert await wrapper.reply("Original input", react_config={"max_iters": 30}) is result + reply.assert_awaited_once_with("Original input", react_config={"max_iters": 30}) + reply.return_value = {"last_message": {"finished_reason": "interrupted"}} + with pytest.raises(RuntimeError, match="interrupted"): + await wrapper.reply("Original input") diff --git a/tests/unit/test_memlens_judge.py b/tests/unit/test_memlens_judge.py new file mode 100644 index 00000000..3d02c74e --- /dev/null +++ b/tests/unit/test_memlens_judge.py @@ -0,0 +1,101 @@ +"""Standalone judge Jobs preserve official prompts, parsers and generation defaults.""" + +# pylint: disable=missing-function-docstring,protected-access +import hashlib +from pathlib import Path +from unittest.mock import AsyncMock + +from agentscope.message import TextBlock +from agentscope.model import ChatResponse, ChatUsage +import pytest +import yaml + +from reme.components.runtime_context import RuntimeContext +from reme.plugin_manifest import parse_plugin_manifest + + +@pytest.fixture(name="judge") +def fixture_judge(monkeypatch): + root = Path(__file__).resolve().parents[2] / "plugins/memlens-judge/src" + monkeypatch.syspath_prepend(str(root)) + from judge_memlens import protocol + from judge_memlens.llm_judge import MemLensJudgeStep + + step = MemLensJudgeStep() + step.context = RuntimeContext.from_context( + None, + question="QUESTION", + reference="REFERENCE", + prediction="PREDICTION", + evidence=["first", "second"], + category="information_extraction", + ) + step.as_llm = AsyncMock( + return_value=ChatResponse( + content=[TextBlock(text='{"answer_score":1}')], + is_last=True, + usage=ChatUsage(100, 10, 0.01, cache_input_tokens=30), + ), + ) + return step, protocol, root + + +@pytest.mark.asyncio +async def test_judge_job_uses_only_official_prompt_and_separate_usage(judge): + step, protocol, root = judge + response = await step.execute() + assert response.success and response.metadata["llm_judge_score"] == 1 + assert response.metadata["usage"]["cache_input_tokens"] == 30 + kwargs = {"category": "information_extraction"} + expected = protocol.judge_prompt("QUESTION", "REFERENCE", "PREDICTION", **kwargs) + assert step.as_llm.call_args.args[0][0].get_text_content() == expected + assert len(step.as_llm.call_args.args[0]) == 1 and not step.as_llm.call_args.kwargs + text = (root / "judge_memlens/plugin.yaml").read_text() + parse_plugin_manifest(text, plugin_name="memlens-judge") + manifest = yaml.safe_load(text)["application_defaults"] + assert manifest["jobs"]["answer_judge"]["steps"] == [{"backend": "memlens_judge_step", "as_llm": "judge"}] + model = manifest["components"]["as_llm"]["judge"] + assert model["extra_body"] == {"enable_thinking": False} and model["parameters"]["temperature"] == 0 + assert model["parameters"]["max_tokens"] == 256 + step.as_llm.return_value = ChatResponse(content=[TextBlock(text="malformed")], is_last=True) + assert (await step.execute()).metadata["parse_default_zero"] is True + step.context.data["prediction"] = "word " * 501 + assert (await step.execute()).metadata["auto_zero"] is True + assert step.as_llm.await_count == 2 + step.context.data["prediction"] = "PREDICTION" + step.as_llm.side_effect = TimeoutError("provider") + with pytest.raises(TimeoutError): + await step.execute() + + +def test_memlens_official_registry_digests_and_routing(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins/memlens-judge/src")) + from judge_memlens import protocol + + # Independent upstream build_judge_prompt digests, not generated by this renderer. + digests = { + "IE_Entity": "83991403e50c7390f9bafd957618a9ad01bdde2ae6d26ada38ce9954c697ebb0", + "IE_PreviousInfo": "18082c484f096bef129b6007671f7cde3e9c44b9e1e5e781630b84a604c126f7", + "KU_KnowledgeUpdate": "ccd52d849c6eb132e819f1f2847faf8a3c1755e517fd7f3e17308d8b51798951", + "TR_DurationComparison": "b884e6688d2b195c63d25c40c9003d98a5c8915b8b008613b51b14f95cc1c733", + "TR_OrderRanking": "c57f6b289c986aef4a1c6032c0d5544a7980001354298aaa9b3546c0edb4f86b", + "TR_DateExtraction": "8c10197630ba720759a6e04a1d670e5af9781284c509d53f89310529ad8d17a5", + "MSR_YesNo": "79779f253bfee20cfa9666ff33c6cdcc37828b9d17f8c83f92af4ec7e2e557ff", + "MSR_Counting": "bde797036e6069ac10a649ae80fc6749facd4c1ef5e912f39418fcfe729752a7", + "MSR_Arithmetic": "089db6580a3a640a5cb925b490e9cd059a71b65d4d85b236b748ee17cb1766af", + "AR_AnswerRefusal": "c7a6702489c079cfa4c354febb916078825325a0396261627b2c0295aab0d018", + } + prompts = protocol._prompts() + assert prompts["source_revision"] == protocol.SOURCE_REVISION + assert set(prompts["criteria"]) == set(digests) + for key, digest in digests.items(): + base = prompts["template"].format(task_criteria=prompts["criteria"][key], examples=prompts["examples"][key]) + reference = "Insufficient information" if key == "AR_AnswerRefusal" else "REFERENCE" + old = ": OLD\n" if key == "KU_KnowledgeUpdate" else "" + rendered = ( + f"{base}[Current Case]\n: QUESTION\n: {reference}\n" + f"{old}: PREDICTION\n\n[Scoring Rationale]:" + ) + assert hashlib.sha256(rendered.encode()).hexdigest() == digest + assert protocol.task_key("multi_session_reasoning", "entity_resolution", "2") == "MSR_Counting" + assert protocol.parse_judge_response('{"answer_score":0} [Score]:1')["acc"] == 0 diff --git a/tests/unit/test_memlens_runner_failures.py b/tests/unit/test_memlens_runner_failures.py new file mode 100644 index 00000000..095278ff --- /dev/null +++ b/tests/unit/test_memlens_runner_failures.py @@ -0,0 +1,423 @@ +"""Offline failure, accounting and artifact contracts for the MemLens runner.""" + +# pylint: disable=missing-function-docstring,protected-access +import asyncio +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from benchmark.memlens import run as runner +from benchmark.memlens.artifacts import error_details, judge_settings, model_settings, response_record, safe_value +from reme.schema import Response + + +@pytest.fixture(autouse=True) +def memlens_judge_path(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins/memlens-judge/src")) + + +@pytest.fixture(name="scene_run") +def fixture_scene_run(tmp_path, monkeypatch): + import reme + + scene = { + "id": "scene-1", + "sessions": [{"id": "session-1", "timestamp": "2026-01-01"}], + "questions": [ + { + "id": "question-1", + "text": "Question", + "question_date": "2026-01-02", + "answer": "Answer", + "category": "information_extraction", + }, + ], + } + app = SimpleNamespace(start=AsyncMock(), close=AsyncMock(), run_job=AsyncMock()) + judge = SimpleNamespace(run_job=AsyncMock(return_value=Response(metadata={"llm_judge_score": 1}))) + monkeypatch.setattr(reme, "Application", lambda **_: app) + monkeypatch.setattr(runner, "app_config", lambda *_, **__: {}) + monkeypatch.setattr(runner, "observe_usage", lambda _: None) + monkeypatch.setattr(runner, "session_messages", lambda _: [{"role": "user", "content": "Text"}]) + case = tmp_path / "reme_image/scenes/0000-scene-1" + + async def evaluate(): + return await runner.evaluate_scene(scene, 0, "candidate", tmp_path, judge, asyncio.Semaphore(1)) + + return scene, app, judge, case, evaluate + + +def measured(tokens=120): + return runner.normalized_usage({"input_tokens": tokens - 20, "output_tokens": 20}) + + +@pytest.mark.asyncio +async def test_index_failure_retains_completed_session_usage(scene_run): + _, app, _, case, evaluate = scene_run + + async def job(name, **_): + if name == "auto_memory": + runner._CURRENT.get().append(measured()) + return Response(answer="Memory") + return Response(success=False, answer="Index unavailable") + + app.run_job.side_effect = job + results, memory, answers, status = await evaluate() + assert results == answers == [] + assert runner.total_usage(memory)["total_tokens"] == 120 + assert json.loads((case / "status.json").read_text()) == status + assert status["stage"] == "index_update" and status["scene_id"] == "scene-1" + assert "response detail: Index unavailable" in status["error"]["message"] + assert status["error_type"] == "RuntimeError" + assert status["memory_complete"] is False + assert len(list((case / "ingestion/attempt-001/sessions").glob("*.json"))) == 1 + artifact = json.loads((case / "ingestion/attempt-001/diagnostics/index-response.json").read_text()) + assert artifact["response"]["success"] is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct", [False, True]) +async def test_failed_memory_keeps_its_successful_requests_and_diagnostics(scene_run, direct): + scene, app, _, case, evaluate = scene_run + scene["sessions"].append({"id": "session-2", "timestamp": "2026-01-02"}) + + async def job(*_, session_id, **__): + runner._CURRENT.get().append(measured()) + if session_id == "session-1": + return Response(answer="Memory") + if direct: + try: + raise ValueError("Original provider failure") + except ValueError as exc: + raise RuntimeError("Memory failed") from exc + return Response(success=False, metadata={"error": "Memory failed"}) + + app.run_job.side_effect = job + _, memory, _, _ = await evaluate() + assert runner.total_usage(memory)["total_tokens"] == 240 + status = json.loads((case / "status.json").read_text()) + assert status["stage"] == "auto_memory" and status["session_id"] == "session-2" + assert status["memory_complete"] is False + assert (case / "ingestion/attempt-001/sessions/00000.json").exists() + assert not (case / "ingestion/attempt-001/sessions/00001.json").exists() + assert json.loads((case / "ingestion/attempt-001/diagnostics/00001-usage.json").read_text())["total_tokens"] == 120 + if direct: + assert status["error"]["cause"]["message"] == "Original provider failure" + assert status["error"]["traceback"] and "locals" not in json.dumps(status) + else: + artifact = json.loads((case / "ingestion/attempt-001/diagnostics/00001-response.json").read_text()) + assert artifact["response"]["success"] is False + app.close.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failed_job", ["agentic_answer", "judge"]) +async def test_failed_answer_or_judge_keeps_answer_usage_once(scene_run, failed_job): + _, app, judge, case, evaluate = scene_run + + async def job(name, **_): + if name in {"auto_memory", "agentic_answer"}: + runner._CURRENT.get().append(measured()) + return Response(answer="Answer", success=name != failed_job) + + app.run_job.side_effect = job + if failed_job == "judge": + judge.run_job.return_value = Response(success=False, metadata={"error": "Judge failed"}) + results, memory, answers, _ = await evaluate() + assert results == [] and len(memory) == len(answers) == 1 + assert runner.total_usage(answers)["total_tokens"] == 120 + status = json.loads((case / "status.json").read_text()) + assert status["question_id"] == "question-1" and status["stage"] == failed_job + assert status["memory_complete"] is True + path = "judge-response.json" if failed_job == "judge" else "answer-response.json" + assert json.loads((case / "questions/0000/diagnostics" / path).read_text())["response"]["success"] is False + assert (case / "questions/0000/answer.json").exists() is (failed_job == "judge") + assert json.loads((case / "questions/0000/diagnostics/answer-usage.json").read_text())["total_tokens"] == 120 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ["preprocess", "start", "cleanup", "cancel"]) +async def test_direct_lifecycle_failures_are_recorded(scene_run, monkeypatch, stage): + _, app, _, case, evaluate = scene_run + app.run_job.return_value = Response(answer="Answer") + if stage == "start": + app.start.side_effect = RuntimeError("startup failed") + elif stage == "cleanup": + app.close.side_effect = RuntimeError("cleanup failed") + elif stage == "cancel": + app.run_job.side_effect = asyncio.CancelledError() + else: + + def broken(_): + raise ValueError("Invalid image: release_images/missing.png") + + monkeypatch.setattr(runner, "session_messages", broken) + if stage == "cancel": + with pytest.raises(asyncio.CancelledError): + await evaluate() + else: + await evaluate() + status = json.loads((case / "status.json").read_text()) + assert status["complete"] is False + assert status["stage"] == ("auto_memory" if stage == "cancel" else stage) + assert status["error"]["type"] == ( + "CancelledError" if stage == "cancel" else "RuntimeError" if stage in {"start", "cleanup"} else "ValueError" + ) + if stage == "preprocess": + assert status["session_id"] == "session-1" and "missing.png" in status["error"]["message"] + assert status["error"]["traceback"] + + +@pytest.mark.asyncio +async def test_observer_separates_concurrent_scopes_including_compression_and_unknown_counters(): + def response(usage, status=200): + return SimpleNamespace(status_code=status, is_stream_consumed=True, json=lambda: {"usage": usage}) + + client = SimpleNamespace(send=AsyncMock(side_effect=response)) + compressor = SimpleNamespace(send=AsyncMock(side_effect=response)) + app = SimpleNamespace( + context=SimpleNamespace( + components={ + "as_llm": { + "default": SimpleNamespace(model=SimpleNamespace(client=SimpleNamespace(_client=client))), + "compressor": SimpleNamespace(model=SimpleNamespace(client=SimpleNamespace(_client=compressor))), + }, + }, + ), + ) + runner.observe_usage(app) + + async def request(tokens): + with runner.collect_usage() as records: + await client.send({"prompt_tokens": tokens, "completion_tokens": 20}) + await asyncio.sleep(0) + await compressor.send({}) + return records + + first, second = await asyncio.gather(request(100), request(200)) + assert runner.total_usage(first)["total_tokens"] == 120 + assert runner.total_usage(second)["total_tokens"] == 220 + assert runner.total_usage(first)["records"] == 2 + assert runner.total_usage(first)["coverage"]["total_tokens"] == 1 + assert runner.total_usage(first)["cached_input_tokens"] is None + assert runner._CURRENT.get() is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct", [False, True]) +async def test_failure_artifacts_redact_secrets_and_request_payloads(scene_run, monkeypatch, direct): + _, app, _, case, evaluate = scene_run + secret = "ENV_SECRET_SENTINEL_314159" + monkeypatch.setenv("LLM_API_KEY", secret) + detail = ( + f"provider failed {secret}; api_key=INLINE_SENTINEL " + "Bearer AUTH_SENTINEL https://USER_SENTINEL:PASS_SENTINEL@provider.test/v1?key=URL_SENTINEL " + f"data:image/png;base64,{'QmFzZTY0' * 20}" + ) + if direct: + app.run_job.side_effect = RuntimeError(detail) + else: + app.run_job.return_value = Response( + success=False, + metadata={ + "error": detail, + "request_headers": {"arbitrary": "HEADER_SENTINEL"}, + "config": {"arbitrary": "CONFIG_SENTINEL"}, + "source": {"type": "base64", "data": "SHORT_IMAGE_SENTINEL"}, + }, + ) + await evaluate() + artifacts = "".join(path.read_text() for path in case.rglob("*.json")) + for value in ( + secret, + "INLINE_SENTINEL", + "AUTH_SENTINEL", + "USER_SENTINEL", + "PASS_SENTINEL", + "URL_SENTINEL", + "HEADER_SENTINEL", + "CONFIG_SENTINEL", + "SHORT_IMAGE_SENTINEL", + "QmFzZTY0", + ): + assert value not in artifacts + assert "provider failed" in artifacts and "[REDACTED" in artifacts + + +@pytest.mark.asyncio +async def test_successful_predictions_and_references_remain_verbatim(scene_run): + scene, app, judge, case, evaluate = scene_run + answer = "token=ordinary-text https://example.test/path?a=b data:image/png;base64," + "QmFzZTY0" * 20 + scene["questions"][0]["answer"] = answer + app.run_job.return_value = Response(answer=answer) + results, *_ = await evaluate() + assert results[0]["prediction"] == results[0]["reference"] == answer + assert judge.run_job.call_args.kwargs["prediction"] == answer + for name in ("answer.json", "result.json"): + assert json.loads((case / "questions/0000" / name).read_text())["prediction"] == answer + assert json.loads((case / "status.json").read_text())["memory_complete"] is True + + +@pytest.mark.asyncio +async def test_cancellation_finalizes_observed_requests_before_propagating(scene_run): + _, app, _, case, evaluate = scene_run + + async def canceled(*_, **__): + runner._CURRENT.get().append(measured()) + raise asyncio.CancelledError() + + app.run_job.side_effect = canceled + with pytest.raises(asyncio.CancelledError): + await evaluate() + assert json.loads((case / "status.json").read_text())["memory_usage"]["total_tokens"] == 120 + assert json.loads((case / "ingestion/attempt-001/diagnostics/00000-usage.json").read_text())["total_tokens"] == 120 + app.close.assert_awaited_once() + + +@pytest.mark.parametrize("field", ["cookie", "set-cookie", "private_key", "id_token", "refresh_token", "token"]) +def test_diagnostic_credentials_are_redacted_without_removing_counters(field): + diagnostic = safe_value({field: "DIAGNOSTIC_SECRET_SENTINEL", "input_tokens": 120, "output_tokens": 20}) + assert diagnostic == {field: "[REDACTED]", "input_tokens": 120, "output_tokens": 20} + assert "DIAGNOSTIC_SECRET_SENTINEL" not in json.dumps( + error_details(RuntimeError(f"{field}: DIAGNOSTIC_SECRET_SENTINEL")), + ) + + +@pytest.mark.parametrize( + "secret", + ["response", "success", "answer", "metadata", "usage", "coverage", "records", "input_tokens"], +) +def test_failed_response_schema_survives_environment_name_collision(monkeypatch, secret): + monkeypatch.setenv("TEST_PRIVATE_KEY", secret) + usage = {"input_tokens": 100, "output_tokens": 20, "records": 1, "coverage": {"input_tokens": 1}} + raw = Response( + success=False, + answer=f"Provider unavailable: {secret}", + metadata={"usage": usage, "private_key": "PRIVATE_SENTINEL", f"debug {secret}": f"detail {secret}"}, + ) + response = json.loads(json.dumps(response_record(raw))) + assert Response.model_validate(response).success is False + assert response["answer"] == "Provider unavailable: [REDACTED]" + assert response["metadata"] == { + "usage": usage, + "private_key": "[REDACTED]", + "debug [REDACTED]": "detail [REDACTED]", + } + + +@pytest.mark.parametrize("secret", ["a", "model", "temperature"]) +def test_model_snapshot_preserves_whitelisted_schema(monkeypatch, secret): + monkeypatch.setenv("LLM_API_KEY", secret) + spec = {"backend": "agentscope", "model": "candidate", "parameters": {"temperature": 0, "max_tokens": 256}} + config = SimpleNamespace(components={"as_llm": {"default": SimpleNamespace(model_dump=lambda: spec)}}) + snapshot = model_settings(config)["default"] + assert snapshot["backend"] == "agentscope" and snapshot["model"] == "candidate" + assert snapshot["parameters"] == {"temperature": 0, "max_tokens": 256} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("secret", ["a", "stage"]) +@pytest.mark.parametrize("outcome", ["complete", "index_update", "cleanup"]) +async def test_run_summary_keeps_usage_and_independent_completion(scene_run, monkeypatch, secret, outcome): + import reme + from reme.schema import ApplicationConfig + + scene, app, judge, case, _ = scene_run + root = case.parents[2] + monkeypatch.setenv("LLM_API_KEY", secret) + monkeypatch.setattr(runner, "app_config", lambda workspace, *_, **__: {"workspace_dir": str(workspace)}) + monkeypatch.setattr(runner, "load_memlens", lambda *_: [scene]) + judge.start, judge.close = AsyncMock(), AsyncMock() + judge.config = ApplicationConfig(workspace_dir=root / "judge") + applications = iter([judge, app]) + monkeypatch.setattr(reme, "Application", lambda **_: next(applications)) + + async def job(name, **_): + if name != "index_update": + runner._CURRENT.get().append(measured()) + return Response(answer="Answer", success=name != outcome) + + app.run_job.side_effect = job + if outcome == "cleanup": + app.close.side_effect = RuntimeError("Close failed") + summary = await runner.run( + { + "dataset": str(root / "unused"), + "output": str(root / "output"), + "query_concurrency": 1, + "scenario_concurrency": 1, + "unexpected": "RAW_CONFIG_SENTINEL", + }, + "candidate", + "128k", + ) + group = root / "output/memlens-128k/candidate" + status = json.loads((group / "reme_image/scenes/0000-scene-1/status.json").read_text()) + complete = outcome == "complete" + memory_complete = outcome != "index_update" + assert status["stage"] == outcome and status["scene_id"] == "scene-1" + assert status["memory_complete"] is memory_complete and status["complete"] is complete + assert status["memory_usage"]["input_tokens"] == 100 + assert summary["scored_questions"] == summary["memory_complete_scenes"] == int(memory_complete) + assert summary["token_usage"]["total_tokens"] == (240 if memory_complete else 120) + assert summary["cleanup_error_count"] == int(outcome == "cleanup") + assert summary["completed_scenes"] == int(complete) + assert summary["usage_policy_version"] == "successful-http-requests-v2" + assert json.loads((group / "summary.json").read_text()) == summary + assert json.loads((group / "status.json").read_text())["complete"] is complete + settings = (group / "settings.json").read_text() + assert "RAW_CONFIG_SENTINEL" not in settings + assert json.loads(settings)["judge"]["protocol_id"] == "memlens-task-binary-api-qwen-v1" + if outcome == "cleanup": + assert status["error_type"] == status["error"]["type"] == "RuntimeError" + assert status["error"]["message"] == "Close failed" + + +@pytest.mark.asyncio +async def test_judge_snapshot_uses_validated_plugin_and_environment_config(tmp_path, monkeypatch): + from importlib.metadata import EntryPoint + + from reme import Application + from reme.plugin import PluginManager, _load_plugin + + plugin = _load_plugin( + "memlens-judge", + EntryPoint( + name="memlens-judge", + value="judge_memlens", + group="reme.plugins", + ), + ) + manager = PluginManager([plugin]) + monkeypatch.setattr(PluginManager, "discover", classmethod(lambda cls, _: manager)) + monkeypatch.setenv("JUDGE_MODEL_NAME", "judge-override-for-test") + monkeypatch.setenv("LLM_API_KEY", "SNAPSHOT_SECRET_SENTINEL") + monkeypatch.setenv("LLM_BASE_URL", "https://private.test/v1?key=URL_SENTINEL") + inputs = runner.app_config(tmp_path / "judge", "candidate-model", judge=True) + assert inputs["components"] == {} + app = Application(**inputs) + snapshot = judge_settings(app) + model = snapshot["models"]["judge"] + assert model["model"] == "judge-override-for-test" + assert model["parameters"] == { + "temperature": 0, + "thinking_enable": False, + "max_tokens": 256, + } + assert model["extra_body"] == {"enable_thinking": False} + assert model["client_kwargs"] == {"timeout": 600, "max_retries": 0} + assert snapshot["protocol_id"] == "memlens-task-binary-api-qwen-v1" + assert "SENTINEL" not in json.dumps(snapshot) and "credential" not in json.dumps(snapshot) + await app.close() + + +def test_cli_reports_sanitized_direct_failure(monkeypatch, capsys): + monkeypatch.setattr("sys.argv", ["memlens"]) + monkeypatch.setattr(runner, "run", AsyncMock(side_effect=RuntimeError("provider api_key=CLI_SECRET_SENTINEL"))) + assert runner.main() == 2 + diagnostic = capsys.readouterr().err + assert "CLI_SECRET_SENTINEL" not in diagnostic + assert json.loads(diagnostic)["error"]["message"] == "provider api_key=[REDACTED]"