diff --git a/benchmark/mobilemem/README.md b/benchmark/mobilemem/README.md new file mode 100644 index 00000000..ca10e49f --- /dev/null +++ b/benchmark/mobilemem/README.md @@ -0,0 +1,96 @@ +# MobileMem-Omni Visual Reasoning + +Evaluate ReMe Image on the published `visual_reasoning` questions using each user's complete chronological +dialogue history. Original dialogue text, including any descriptions already in it, is retained alongside images. +Separate summaries/captions, reference answers, evidence annotations and question types are not added to candidate Jobs. +Native `multi_visual_reasoning` questions are not relabelled or included. + +## Run + +Install both plugins and extract the official archive without changing its images: + +```bash +pip install -e ".[as]" -e plugins/mobilemem -e plugins/mobilemem-judge +mkdir -p /path/to/omni/image +python -m zipfile -e /path/to/omni/image.zip /path/to/omni/image +export LLM_API_KEY="..." +export LLM_BASE_URL="..." +python -m benchmark.mobilemem.run --dataset /path/to/omni --output /path/to/exp-xxxx +``` + +[config.yaml](config.yaml) supplies defaults: `qwen3.8-max`, four concurrent users and 16 concurrent questions. +Override them with `--model`, `--user-concurrency` and `--query-concurrency`. +Every run creates `/mobilemem/` and rejects an existing directory. + +## Evaluation + +Each user has a fresh workspace. All source sessions are ingested chronologically through multimodal +`auto_memory`, then indexed once. Every question starts a fresh 30-iteration `agentic_answer` Agent with +BM25 `search`, `read` and `read_image`. Search covers daily notes, digest notes and original session text. +Original Image IDs accompany source images; saved native attachments can be reopened by the answer Agent. + +The [candidate preset](../../plugins/mobilemem/src/reme_mobilemem/benchmark.yaml) uses the same model for memory, +answering and compression: temperature 0, thinking enabled, 65,536 output tokens, 1,000,000-token context and +at most 4,096 images. Native context ratios remain 0.8/0.1. Input images are not resized; MPO inputs retain +the existing primary-picture PNG conversion. Published path aliases are resolved by unique image basename. + +The official QA template wraps the question with evidence-tool instructions. Only the isolated +[MobileMem Judge](../../plugins/mobilemem-judge/README.md) receives references and evidence. Its default is the +MemBase executable-default `mobilemem-omni-judge` rubric, `qwen3.8-max`, temperature 0, thinking disabled and +4,096 output tokens. Scores report binary LLM-Judge accuracy and the official token-overlap F1 and BLEU-1. + +## Results + +Each user's `workspace/` retains notes, source JSONL, images and native AgentScope trajectories. +`sessions/*.json` and `questions/*/` retain responses, judgments, metrics and technical errors. +`summary.json` reports completed/expected questions and scores over successfully scored results only. +Native interrupted replies and failed Jobs remain unscored technical gaps, not zero-score answers. +The runner adds no fallback or retry queue. + +Usage policy version 2 counts observed successful HTTP requests, including compression and requests made before +a candidate history or answer later fails. Memory cost comes from `sessions/*.json` and `index-response.json`; +answer cost comes only from `answer-response.json`, including failed responses or `response: null` when a call raises. +Successful and failed answers use this single response file; its `response.success` determines whether scoring proceeds. +Earlier runs' `answer.json` copies, result files and response metadata are not counted again. `memory_usage.json` remains a +completion marker: it is written only after the full history and index update succeed. +Judge usage is excluded; absent counters stay unknown and each counter includes its observation coverage. +The `mean_tokens_per_user` / `mean_tokens_per_question` fields amortize this run's observed cost over scheduled users / +expected questions, with explicit denominators and attempt/completion counts. They are not successful-sample averages. +Earlier selected-success-only reports retain their original accounting policy. + +Errors retain the failing stage, user/session/question identifiers, safe image paths, sanitized reasons and stack +locations; credentials, request headers, image payloads and frame locals are omitted. Cleanup errors are separate; +`cleanup_error_count` makes the run exit with status 2 without discarding already scored results. +`judge-settings.json` records a credential-free whitelist of the validated Judge model, parameters and protocol, +including an optional `JUDGE_MODEL_NAME` override, so scoring can be traced to the effective configuration. + +Sources: [MobileMem evaluation code](https://github.com/zjunlp/MobileMem/tree/fa5cd13c9c659699df9caf650987ba2bea8ca085/omni/eval/eval) +and [MemBase judge](https://github.com/zjunlp/MemBase/blob/a9ba9495d48541d3e68588ff097294cc7a1201bc/membase/inference_utils/prompts.py). +Dataset revision: `989ba00900bddf4e0d4b0c7dd52f0840a7ee3a54`. + +## Upstream code license + +The QA template and overlap metrics in `run.py` are adapted from ZJUNLP MobileMem's evaluation code at +`fa5cd13c9c659699df9caf650987ba2bea8ca085`, under the following license. + +MIT License + +Copyright (c) 2026 ZJUNLP + +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/mobilemem/README_ZH.md b/benchmark/mobilemem/README_ZH.md new file mode 100644 index 00000000..5b1c01a7 --- /dev/null +++ b/benchmark/mobilemem/README_ZH.md @@ -0,0 +1,61 @@ +# MobileMem-Omni 视觉推理评测 + +使用每个用户完整、按时间排序的原始对话,评测 ReMe Image 在官方 `visual_reasoning` 题目上的表现。 +记忆与回答 Job 仅接收原始文本和图像,不接收数据集摘要、caption、标准答案、证据标注或题型。 +官方 `multi_visual_reasoning` 不被重命名,也不纳入这个评测范围。 + +## 运行 + +安装两个插件,并按原目录结构解压官方图片: + +```bash +pip install -e ".[as]" -e plugins/mobilemem -e plugins/mobilemem-judge +mkdir -p /path/to/omni/image +python -m zipfile -e /path/to/omni/image.zip /path/to/omni/image +export LLM_API_KEY="..." +export LLM_BASE_URL="..." +python -m benchmark.mobilemem.run --dataset /path/to/omni --output /path/to/exp-xxxx +``` + +[config.yaml](config.yaml) 默认使用 `qwen3.8-max`、4 个用户并发、16 个题目并发; +可通过 `--model`、`--user-concurrency`、`--query-concurrency` 覆盖。 +每次运行创建 `/mobilemem/`,已有目录不会被覆盖。 + +## 评测设置 + +每个用户使用独立的新工作空间,按时间顺序通过多模态 `auto_memory` 摄入全部 session,然后更新一次索引。 +每题使用新的 `agentic_answer` Agent,最多 30 轮,工具为 BM25 `search`、`read` 和 `read_image`。 +检索范围包含 daily、digest 笔记及原始 session 正文;保存的原生图像附件与 Image ID 可供回答 Agent 重新读图。 + +[候选模型配置](../../plugins/mobilemem/src/reme_mobilemem/benchmark.yaml) 在记忆、回答和压缩阶段使用同一模型: +temperature 为 0,开启 thinking,输出预算 65,536,上下文为 1,000,000 token,最多 4,096 张图。 +原生上下文比例保持 0.8/0.1。输入图片不缩放;MPO 保持既有的主图 PNG 转换。 +数据集图片路径别名通过唯一 basename 对应原件。 + +问题使用官方 QA 模板,并提示通过证据工具作答。标准答案与证据仅交给独立的 +[MobileMem Judge](../../plugins/mobilemem-judge/README.md)。默认采用 MemBase 可执行默认 +`mobilemem-omni-judge` 评分模板,模型为 `qwen3.8-max`,temperature 为 0,关闭 thinking,输出预算 4,096。 +报告包含二元 LLM-Judge 准确率及官方 token-overlap F1、BLEU-1。 + +## 结果 + +每个用户的 `workspace/` 保存笔记、原始 JSONL、图像及原生 AgentScope 轨迹。 +`sessions/*.json` 与 `questions/*/` 保存响应、评分、指标和技术错误。 +`summary.json` 汇总完成/应测题数,指标只对成功评分的结果求均值。 +原生中断与失败 Job 保留为未评分的技术缺口,不作为零分答案;runner 不额外加入 fallback 或重试队列。 + +usage policy 第 2 版统计已观测的成功 HTTP 请求,包括压缩调用及摄入/回答最终失败前已经产生的请求。 +记忆成本只汇总 `sessions/*.json` 和 `index-response.json`;回答成本只汇总 `answer-response.json`, +包括失败响应及直接异常时的 `response: null`。成功与失败的回答统一保存在这个文件中, +根据 `response.success` 决定是否进入评分;旧运行的 `answer.json` 副本、result 和 response metadata 不重复计费。 +`memory_usage.json` 仍仅在完整历史摄入及索引更新成功后生成,不能将它误作部分摄入标记。 +Judge 用量不计入方法成本;缺失字段保持未知,并记录每个字段的观测覆盖数。 +`mean_tokens_per_user` / `mean_tokens_per_question` 是本次已观测成本按计划用户数/应测题数均摊, +summary 明确记录分母和尝试/完成数,不代表成功样本平均成本。旧的 selected-success-only 报告不改写口径。 + +错误记录包含失败阶段、用户/会话/题号、安全的图片相对路径、脱敏原因和调用栈位置,不保存凭据、请求 headers、 +图片载荷或 frame locals;清理错误单独保留,`cleanup_error_count` 使本轮以状态码 2 退出,但不丢弃已评分结果。 +`judge-settings.json` 从实际校验后的配置中保存 Judge 模型、参数和协议白名单, +包括可选的 `JUDGE_MODEL_NAME` 覆盖值,不导出凭据。 + +协议来源、数据版本和改编代码的完整 MIT 许可证见 [英文说明](README.md)。 diff --git a/benchmark/mobilemem/_audit.py b/benchmark/mobilemem/_audit.py new file mode 100644 index 00000000..5bb1ba2f --- /dev/null +++ b/benchmark/mobilemem/_audit.py @@ -0,0 +1,182 @@ +"""Small, credential-safe diagnostics and accounting helpers for this runner.""" + +import os +from pathlib import Path, PurePosixPath +import re +import traceback + +USAGE_POLICY = { + "version": 2, + "scope": "observed successful HTTP requests, including failed candidate flows and compression; judge excluded", + "memory_sources": ["users/*/sessions/*.json", "users/*/index-response.json"], + "answer_source": "users/*/questions/*/answer-response.json", + "means": "observed run cost amortized over scheduled users / expected questions, not successful samples", +} +_SENSITIVE = re.compile( + r"authorization|credential|secret|password|api.?key|private[_-]?key|cookie|headers|" + r"(?:^|[_-])(?:access[_-]?token|refresh[_-]?token|id[_-]?token|token)$", + re.I, +) +_USAGE_FIELDS = { + "usage", + "coverage", + "records", + "input_tokens", + "output_tokens", + "total_tokens", + "cached_input_tokens", + "noncached_input_tokens", + "reasoning_tokens", +} + + +def safe_text(value, limit: int = 1200) -> str: + """Retain short diagnostic text, without credentials, URLs or binary payloads.""" + text = str(value) + for key, secret in os.environ.items(): + if secret and (_SENSITIVE.search(key) or re.search(r"(?:^|_)TOKEN(?:$|_)", key.upper())): + if len(secret) < 12: + text = re.sub(r"(?]+", "[url redacted]", text, flags=re.I) + text = re.sub(r"(? 6: + return "[nested data omitted]" + if isinstance(value, dict): + return { + (key if key in fixed_keys else safe_text(key, 100)): ( + "[redacted]" + if _SENSITIVE.search(str(key)) or str(key).lower() in {"base64", "data", "body", "request", "traceback"} + else safe_value(item, depth + 1, fixed_keys=fixed_keys) + ) + for key, item in list(value.items())[:50] + } + if isinstance(value, (list, tuple)): + return [safe_value(item, depth + 1, fixed_keys=fixed_keys) for item in value[:20]] + if value is None or isinstance(value, (bool, int, float)): + return value + return safe_text(value) if isinstance(value, str) else f"[{type(value).__name__} omitted]" + + +def response_record(response): + """Keep successful response artifacts compatible; sanitize failed envelopes.""" + if response is None: + return None + value = response.model_dump() + if response.success: + return value + return { + **safe_value({key: item for key, item in value.items() if key not in {"success", "answer", "metadata"}}), + "success": value["success"], + "answer": safe_value(value["answer"]), + "metadata": safe_value(value["metadata"], fixed_keys=_USAGE_FIELDS), + } + + +def safe_image_path(source) -> str: + """Only retain bounded dataset-relative image identifiers in diagnostics.""" + source = str(source) + path = PurePosixPath(source) + if not source or len(source) > 300 or path.is_absolute() or ".." in path.parts: + return "[invalid image path]" + if any(character in source for character in ("\\", ":", "\n", "\r")): + return "[invalid image path]" + return safe_text(source, 300) + + +class ImageInputError(ValueError): + """Attach only the dataset-relative image identifier to a preprocessing error.""" + + def __init__(self, source, error): + super().__init__(safe_text(error)) + self.image_path = safe_image_path(source) + + +def error_record(error, *, stage: str, user_id=None, session_id=None, question_id=None, response=None) -> dict: + """Preserve stage, a bounded cause chain and stack locations, never frame locals.""" + chain, seen = [], set() + current = error + image_path = None + while current is not None and id(current) not in seen and len(chain) < 5: + seen.add(id(current)) + image_path = image_path or getattr(current, "image_path", None) + chain.append( + { + "type": type(current).__name__, + "reason": safe_text(current), + "traceback": [ + {"file": Path(frame.filename).name, "line": frame.lineno, "function": frame.name} + for frame in traceback.extract_tb(current.__traceback__)[-12:] + ], + }, + ) + current = current.__cause__ or (None if current.__suppress_context__ else current.__context__) + reason = response.answer if response is not None and not response.success else str(error) + return { + "stage": stage, + "user_id": safe_text(user_id, 120) if user_id is not None else None, + "session_id": safe_text(session_id, 120) if session_id is not None else None, + "question_id": safe_text(question_id, 120) if question_id is not None else None, + "image_path": image_path, + "error_type": type(error).__name__, + "reason": safe_value(reason), + "chain": chain, + } + + +def mean_tokens(total: dict, denominator: int) -> dict: + """Amortize observed counters over the explicit scheduled population.""" + return { + key: total[key] / denominator if denominator and total[key] is not None else None + for key in ("input_tokens", "output_tokens", "total_tokens", "cached_input_tokens", "noncached_input_tokens") + } + + +def judge_settings(app) -> dict: + """Export a whitelist from the validated application, not unresolved YAML/env text.""" + from judge_mobilemem.protocol import MEMBASE_REVISION, SOURCE_REVISION, protocol_id + + model = app.config.components["as_llm"]["judge"].model_dump() + mode = app.config.jobs["answer_judge"].model_dump().get("judge_protocol", "membase_default") + fields = ("backend", "model", "stream", "context_size", "max_retries", "retry_delay") + return { + **{key: safe_value(model[key]) for key in fields if key in model}, + "parameters": { + key: safe_value(value) + for key, value in (model.get("parameters") or {}).items() + if key in {"temperature", "max_tokens", "thinking_enable", "top_p"} + }, + "extra_body": {"enable_thinking": (model.get("extra_body") or {}).get("enable_thinking")}, + "client_kwargs": { + key: value for key, value in (model.get("client_kwargs") or {}).items() if key in {"timeout", "max_retries"} + }, + "protocol": protocol_id(mode), + "protocol_mode": mode, + "mobilemem_revision": SOURCE_REVISION, + "membase_revision": MEMBASE_REVISION, + } diff --git a/benchmark/mobilemem/config.yaml b/benchmark/mobilemem/config.yaml new file mode 100644 index 00000000..f6ff1e6d --- /dev/null +++ b/benchmark/mobilemem/config.yaml @@ -0,0 +1,3 @@ +model: qwen3.8-max +user_concurrency: 4 +query_concurrency: 16 diff --git a/benchmark/mobilemem/run.py b/benchmark/mobilemem/run.py new file mode 100644 index 00000000..28962ab4 --- /dev/null +++ b/benchmark/mobilemem/run.py @@ -0,0 +1,621 @@ +"""Evaluate ReMe Image on MobileMem-Omni Visual Reasoning. + +Official QA and overlap metrics: Copyright (c) 2026 ZJUNLP, MIT License. +The complete upstream notice is retained in README.md. +""" + +import argparse +import asyncio +import base64 +from collections import Counter +from contextlib import contextmanager +from contextvars import ContextVar +from datetime import datetime +import io +import json +import math +import mimetypes +from pathlib import Path, PurePosixPath +import re + +from PIL import Image, ImageOps, JpegImagePlugin +import yaml + +from reme import Application +from reme.config import resolve_app_config + +from benchmark.mobilemem._audit import ( + ImageInputError, + USAGE_POLICY, + error_record, + judge_settings, + mean_tokens, + response_record, +) + + +def timestamp(value: str) -> str: + """Preserve published start-time precision in AgentScope ISO timestamps.""" + return datetime.fromisoformat(value.replace("Z", "+00:00")).isoformat() + + +def load_mobilemem(root: Path) -> list[dict]: + """Keep complete raw trajectories, selecting only the native visual category.""" + root = root.expanduser().resolve() + records = [] + for name in ("data.jsonl", "questions.jsonl"): + rows = {} + with (root / name).open(encoding="utf-8") as stream: + for line in stream: + if line.strip(): + row = json.loads(line) + identity = str(row["uuid"]) + if identity in rows: + raise ValueError(f"Duplicate MobileMem user: {identity}") + rows[identity] = row + records.append(rows) + histories, questions = records + if set(histories) != set(questions): + raise ValueError("Every MobileMem user must have a trajectory and question record") + return [ + { + "id": f"uuid_{identity}", + "image_root": root / "image", + "sessions": sorted( + row["sessions"], + key=lambda session: datetime.fromisoformat(timestamp(session["event_start_time"])), + ), + "questions": [ + question + for question in questions[identity]["questions"] + if question["question_type"] == "visual_reasoning" + ], + } + for identity, row in histories.items() + ] + + +def resolve_image(root: Path, source: str, aliases: dict) -> Path: + """Resolve exact paths or unique published basename aliases inside image/.""" + root = root.resolve() + relative = PurePosixPath(source) + if relative.is_absolute() or ".." in relative.parts or "\\" in source or ":" in source: + raise ValueError(f"Invalid MobileMem image path: {source}") + parts = relative.parts[1:] if relative.parts[0] == "image" else relative.parts + target = (root / Path(*parts)).resolve() + if not target.is_relative_to(root): + raise ValueError(f"MobileMem image escapes image/: {source}") + if target.is_file(): + return target + if not aliases: + for path in root.rglob("*"): + if path.is_file() and path.resolve().is_relative_to(root): + if (mimetypes.guess_type(path)[0] or "").startswith("image/"): + aliases.setdefault(path.name, []).append(path.resolve()) + matches = aliases.get(relative.name, []) + if len(matches) != 1: + raise ValueError(f"Missing or ambiguous MobileMem image: {source}") + return matches[0] + + +def image_block(path: Path) -> dict: + """Use original bytes, or the existing lossless MPO primary-picture conversion.""" + data = path.read_bytes() + if len(data) > 50 * 1024 * 1024: + raise ValueError(f"Oversized MobileMem image: {path.name}") + opener = JpegImagePlugin.JpegImageFile if data.startswith(b"\xff\xd8\xff") else Image.open + with opener(io.BytesIO(data)) as image: + mime, mpo = image.get_format_mimetype(), "mp" in image.info + if not mime or not mime.startswith("image/") or image.width * image.height > 50_000_000: + raise ValueError(f"Invalid or oversized MobileMem image: {path.name}") + image.verify() + with opener(io.BytesIO(data)) as image: + image.load() + if mpo: + with ImageOps.exif_transpose(image) as oriented, oriented.convert("RGB") as picture: + buffer = io.BytesIO() + picture.save(buffer, format="PNG") + data, mime = buffer.getvalue(), "image/png" + return { + "type": "data", + "source": { + "type": "base64", + "data": base64.b64encode(data).decode("ascii"), + "media_type": mime, + }, + } + + +def session_messages(session: dict, image_root: Path, aliases=None) -> list[dict]: + """Convert original dialogue directly; annotations never enter the candidate.""" + aliases = {} if aliases is None else aliases + messages = [] + for turn in session["dialogue"]: + content = [{"type": "text", "text": turn.get("content", "")}] + refs = turn.get("image_inline") or [] + refs = [refs] if isinstance(refs, str) else refs + if turn["content_type"] == "image" and not refs: + raise ValueError("MobileMem image turn has no image reference") + for source in refs: + try: + image = image_block(resolve_image(image_root, source, aliases)) + except Exception as error: # pylint: disable=broad-exception-caught + raise ImageInputError(source, error) from error + content.extend( + [ + {"type": "text", "text": f"[Image ID: {source}]"}, + image, + ], + ) + messages.append( + { + "id": f"{session['session_id']}:{turn['turn']}", + "name": turn["role"], + "role": turn["role"], + "content": content, + "created_at": timestamp(session["event_start_time"]), + }, + ) + return messages + + +def evidence_values(question: dict) -> list[str]: + """Keep the official Raw2Locomo/MemBase evidence-string conversion.""" + return [ + str(item["explanation"]) if isinstance(item, dict) else item + for item in question.get("evidence", []) + if (isinstance(item, dict) and item.get("explanation")) or isinstance(item, str) + ] + + +ANSWER_RULES = ( + "Answer Rules:\n" + "1. If options are provided in the question, use the original option text as your answer.\n" + "2. For multiple-choice questions (multiple-select), select all correct options.\n" + "3. For open-ended questions, respond freely." +) + + +def query_prompt(question: str, context: str) -> str: + """Use the published Text Memory QA template verbatim.""" + return ( + f"Question:\n{question}\nPlease answer the question based on the following memories:\n" + f"{context}\n{ANSWER_RULES}" + ) + + +def tokenize(text: str) -> list[str]: + """Official lowercase ASCII word and individual CJK-character tokens.""" + text = re.sub(r"[\u0000-\u002F\u003A-\u0040\u005B-\u0060\u007B-\u007E]", " ", str(text).lower()) + return re.findall(r"[\u4e00-\u9fff]|[a-z0-9]+", text) + + +def score_answer(reference: str, prediction: str) -> dict[str, float]: + """Official token-overlap F1 and BLEU-1.""" + pt, rt = tokenize(prediction), tokenize(reference) + overlap = sum((Counter(pt) & Counter(rt)).values()) + f1 = 1.0 if not pt and not rt else 0.0 + if pt and rt and overlap: + precision, recall = overlap / len(pt), overlap / len(rt) + f1 = 2 * precision * recall / (precision + recall) + bleu = overlap / len(pt) * (math.exp(1 - len(rt) / len(pt)) if len(pt) < len(rt) else 1) if pt else 0.0 + return {"f1": f1, "bleu_1": bleu} + + +_CURRENT = ContextVar("mobilemem_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 absent cache fields unknown and avoid double-counting cache subsets.""" + 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")), + ) + 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": raw.get( + "reasoning_tokens", + (raw.get("completion_tokens_details") or {}).get("reasoning_tokens"), + ), + } + + +def aggregate_usage(records: list[dict]) -> dict: + """Return totals and coverage; unavailable counters remain unknown.""" + 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": sum(row.get("records", 1) for row in records), + "coverage": { + key: sum(row.get("coverage", {}).get(key, int(row.get(key) is not None)) for row in records) + for key in _KEYS + }, + } + + +@contextmanager +def usage_scope(): + """Attribute each successful request to its invoking async task.""" + records = [] + token = _CURRENT.set(records) + try: + yield records + finally: + _CURRENT.reset(token) + + +def observe_usage(app) -> None: + """Observe lifecycle-owned non-streaming clients, including the compressor.""" + 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, "_mobilemem_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: + raw = response.json() + except (ValueError, UnicodeError): + raw = {} + usage = raw.get("usage") if isinstance(raw, dict) else None + records.append(normalized_usage(usage if isinstance(usage, dict) else {})) + return response + + client.send = send + client._mobilemem_usage_observer = True # pylint: disable=protected-access + + +def observe_completion(app) -> None: + """Keep native interrupted replies unscored; successful replies are unchanged.""" + for wrapper in app.context.components.get("agent_wrapper", {}).values(): + original = wrapper.reply + + async def reply(inputs, *args, _reply=original, **kwargs): + result = await _reply(inputs, *args, **kwargs) + last = result.get("last_message") or {} + if str(last.get("finished_reason")) == "interrupted" or last.get("error") is not None: + raise RuntimeError("Native Agent was interrupted; its state remains saved.") + return result + + wrapper.reply = reply + + +def save(path: Path, value) -> None: + """Write an evaluation artifact under the newly created output directory.""" + 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 candidate_config(workspace: Path, model: str) -> dict: + """Select one candidate model; behavioral defaults belong to the plugin preset.""" + config = resolve_app_config( + config="mobilemem-benchmark", + plugins=["mobilemem"], + workspace_dir=str(workspace), + log_config=False, + enable_logo=False, + log_to_console=False, + log_to_file=True, + components={"as_llm": {name: {"model": model} for name in ("default", "bench", "compressor")}}, + ) + config["components"].pop("as_embedding", None) + config["components"].pop("embedding_store", None) + return config + + +def judge_config(workspace: Path) -> dict: + """Keep scoring separate from candidate memory, model and evidence tools.""" + return { + "plugins": ["mobilemem-judge"], + "service": {"backend": "http"}, + "workspace_dir": str(workspace), + "enable_logo": False, + "log_to_console": False, + "log_to_file": True, + } + + +def check_response(response): + """Keep failed Jobs as technical errors rather than incorrect predictions.""" + if response.metadata.get("cancelled"): + raise asyncio.CancelledError("ReMe Job cancelled") + if not response.success: + raise RuntimeError("ReMe Job failed; see the saved response artifact.") + return response + + +async def close_application(app, directory, *, user_id=None): + """Finish owned cleanup; keep its diagnostics separate from the primary failure.""" + if app is None: + return + closing = asyncio.create_task(app.close()) + try: + await asyncio.shield(closing) + except asyncio.CancelledError as cancellation: + if closing.cancelled(): + save(directory / "cleanup-error.json", error_record(cancellation, stage="cleanup", user_id=user_id)) + return + try: + await closing + except (Exception, asyncio.CancelledError) as error: # pylint: disable=broad-exception-caught + save(directory / "cleanup-error.json", error_record(error, stage="cleanup", user_id=user_id)) + raise + except Exception as error: # pylint: disable=broad-exception-caught + save(directory / "cleanup-error.json", error_record(error, stage="cleanup", user_id=user_id)) + + +async def gather_tasks(coroutines) -> None: + """Wait for owned tasks, cancelling and joining siblings on failure or cancellation.""" + tasks = [asyncio.create_task(coroutine) for coroutine in coroutines] + try: + await asyncio.gather(*tasks) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +async def evaluate_user(scene, directory: Path, model: str, judge, query_slots) -> None: + """Ingest the complete chronological history into one fresh user workspace.""" + app, response, session_id = None, None, None + stage = "startup" + completed = 0 + memory_usage, aliases = [], {} + try: + app = Application(**candidate_config(directory / "workspace", model)) + await app.start() + observe_usage(app) + observe_completion(app) + for ordinal, session in enumerate(scene["sessions"]): + session_id, response = session["session_id"], None + stage = "preprocessing" + with usage_scope() as measured: + try: + messages = await asyncio.to_thread(session_messages, session, scene["image_root"], aliases) + stage = "memory" + response = await app.run_job( + "auto_memory", + messages=messages, + session_id=session_id, + date=timestamp(session["event_start_time"])[:10], + include_images=True, + include_source_sessions=True, + ) + check_response(response) + finally: + save( + directory / "sessions" / f"{ordinal:05d}.json", + { + "session_id": session_id, + "response": response_record(response), + "usage": aggregate_usage(measured), + }, + ) + memory_usage.extend(measured) + print(f"{scene['id']}: session {ordinal + 1}/{len(scene['sessions'])}", flush=True) + stage, session_id, response = "index", None, None + with usage_scope() as measured: + try: + response = await app.run_job("index_update") + check_response(response) + finally: + save( + directory / "index-response.json", + {"response": response_record(response), "usage": aggregate_usage(measured)}, + ) + memory_usage.extend(measured) + # Compatibility marker: never write this for a partial history or failed index. + save(directory / "memory_usage.json", aggregate_usage(memory_usage)) + + async def answer_question(ordinal, question): + nonlocal completed + question_dir = directory / "questions" / f"{ordinal:04d}" + answer_stage, answer_response = "answer", None + try: + async with query_slots: + with usage_scope() as measured: + try: + answer_response = await app.run_job( + "agentic_answer", + query=query_prompt(question["question"], "Use the available evidence tools."), + include_images=True, + include_source_sessions=True, + max_iters=30, + ) + check_response(answer_response) + finally: + answer_record = { + "question_id": question["question_id"], + "response": response_record(answer_response), + "usage": aggregate_usage(measured), + } + save(question_dir / "answer-response.json", answer_record) + prediction = answer_response.answer + answer_stage, answer_response = "judge", None + try: + answer_response = await judge.run_job( + "answer_judge", + question=question["question"], + reference=question["answer"], + prediction=prediction, + evidence=evidence_values(question), + ) + check_response(answer_response) + finally: + save(question_dir / "judge.json", response_record(answer_response)) + answer_stage = "metrics" + metrics = score_answer(question["answer"], prediction) + metrics["llm_judge"] = answer_response.metadata["llm_judge_score"] + result = { + "question_id": question["question_id"], + "category": question["question_type"], + "question": question["question"], + "reference": question["answer"], + "prediction": prediction, + "metrics": metrics, + "usage": aggregate_usage(measured), + "judge_parse_default_zero": answer_response.metadata["parse_default_zero"], + } + save(question_dir / "result.json", result) + completed += 1 + print(f"{scene['id']}: {question['question_id']} completed", flush=True) + except (Exception, asyncio.CancelledError) as error: # pylint: disable=broad-exception-caught + save( + question_dir / "error.json", + error_record( + error, + stage=answer_stage, + user_id=scene["id"], + question_id=question["question_id"], + response=answer_response, + ), + ) + if isinstance(error, asyncio.CancelledError): + raise + print(f"{scene['id']}: {question['question_id']} failed ({type(error).__name__})", flush=True) + + stage, response = "answers", None + await gather_tasks(answer_question(i, question) for i, question in enumerate(scene["questions"])) + except (Exception, asyncio.CancelledError) as error: # pylint: disable=broad-exception-caught + save( + directory / "error.json", + error_record(error, stage=stage, user_id=scene["id"], session_id=session_id, response=response), + ) + if isinstance(error, asyncio.CancelledError): + raise + print(f"{scene['id']}: {stage} failed ({type(error).__name__})", flush=True) + finally: + try: + await close_application(app, directory, user_id=scene["id"]) + finally: + save(directory / "summary.json", {"completed": completed, "expected": len(scene["questions"])}) + + +def summarize_run(scenes, output: Path, model: str) -> dict: + """Read each canonical receipt/result once, including a run interrupted after partial progress.""" + expected = sum(len(scene["questions"]) for scene in scenes) + rows = [json.loads(path.read_text(encoding="utf-8")) for path in output.glob("users/*/questions/*/result.json")] + sessions = [json.loads(path.read_text(encoding="utf-8")) for path in output.glob("users/*/sessions/*.json")] + indexes = [json.loads(path.read_text(encoding="utf-8")) for path in output.glob("users/*/index-response.json")] + memory = [row["usage"] for row in sessions + indexes] + answers = [ + json.loads(path.read_text(encoding="utf-8")) for path in output.glob("users/*/questions/*/answer-response.json") + ] + answer_usage = [row["usage"] for row in answers] + total = aggregate_usage(memory + answer_usage) + summary = { + "model": model, + "completed": len(rows), + "expected": expected, + "metrics": { + key: sum(row["metrics"][key] for row in rows) / len(rows) if rows else None + for key in ("f1", "bleu_1", "llm_judge") + }, + "memory_usage": aggregate_usage(memory), + "answer_usage": aggregate_usage(answer_usage), + "token_usage": total, + "usage_policy": USAGE_POLICY, + "cleanup_error_count": len(list(output.glob("users/*/cleanup-error.json"))) + + int((output / "cleanup-error.json").exists()), + "usage_denominators": {"scheduled_users": len(scenes), "expected_questions": expected}, + "usage_counts": { + "attempted_users": len(list(output.glob("users/*/summary.json"))), + "attempted_sessions": len(sessions), + "successful_sessions": sum(bool((row.get("response") or {}).get("success")) for row in sessions), + "completed_ingestion_users": len(list(output.glob("users/*/memory_usage.json"))), + "attempted_answers": len(answers), + "successful_answers": sum(bool((row.get("response") or {}).get("success")) for row in answers), + "scored_questions": len(rows), + }, + "mean_tokens_per_user": mean_tokens(total, len(scenes)), + "mean_tokens_per_question": mean_tokens(total, expected), + "judge_parse_default_zero": sum(row["judge_parse_default_zero"] for row in rows), + } + save(output / "summary.json", summary) + return summary + + +async def run(scenes, output: Path, model: str, user_concurrency: int, query_concurrency: int) -> bool: + """Run independent users concurrently and save official metrics plus method cost.""" + judge = None + user_slots, query_slots = asyncio.Semaphore(user_concurrency), asyncio.Semaphore(query_concurrency) + try: + judge = Application(**judge_config(output / "judge-workspace")) + save(output / "judge-settings.json", judge_settings(judge)) + await judge.start() + + async def user_job(scene): + async with user_slots: + await evaluate_user(scene, output / "users" / scene["id"], model, judge, query_slots) + + await gather_tasks(user_job(scene) for scene in scenes) + except (Exception, asyncio.CancelledError) as error: # pylint: disable=broad-exception-caught + save(output / "error.json", error_record(error, stage="run")) + raise + finally: + try: + await close_application(judge, output) + finally: + summary = summarize_run(scenes, output, model) + print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True) + return summary["completed"] == summary["expected"] and summary["cleanup_error_count"] == 0 + + +def main() -> int: + """Create a fresh evaluation; the PR reproduces ReMe Image only.""" + defaults = yaml.safe_load(Path(__file__).with_name("config.yaml").read_text(encoding="utf-8")) + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--model", default=defaults["model"]) + parser.add_argument("--user-concurrency", type=int, default=defaults["user_concurrency"]) + parser.add_argument("--query-concurrency", type=int, default=defaults["query_concurrency"]) + args = parser.parse_args() + if min(args.user_concurrency, args.query_concurrency) < 1: + parser.error("Concurrency must be positive") + scenes = load_mobilemem(args.dataset) + output = args.output.expanduser().resolve() / "mobilemem" / args.model + output.mkdir(parents=True, exist_ok=False) + save( + output / "settings.json", + { + "dataset": str(args.dataset.expanduser().resolve()), + "model": args.model, + "candidate": "mobilemem-benchmark", + "judge": "mobilemem-judge", + "question_type": "visual_reasoning", + "user_concurrency": args.user_concurrency, + "query_concurrency": args.query_concurrency, + }, + ) + return 0 if asyncio.run(run(scenes, output, args.model, args.user_concurrency, args.query_concurrency)) else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/plugins/mobilemem-judge/LICENSE-CODE b/plugins/mobilemem-judge/LICENSE-CODE new file mode 100644 index 00000000..e213f97f --- /dev/null +++ b/plugins/mobilemem-judge/LICENSE-CODE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 ZJUNLP + +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/mobilemem-judge/README.md b/plugins/mobilemem-judge/README.md new file mode 100644 index 00000000..65faee79 --- /dev/null +++ b/plugins/mobilemem-judge/README.md @@ -0,0 +1,29 @@ +# ReMe MobileMem Judge + +Independent `answer_judge` Job for saved MobileMem-Omni predictions. Install with +`pip install -e plugins/mobilemem-judge` and enable `mobilemem-judge` in a separate +evaluation application/workspace, not the candidate memory application. + +Inputs are `question`, `reference`, `prediction`, and optional `evidence` (a list +of original evidence values). `judge_protocol` defaults to `membase_default`, the +executable MemBase rubric used by this benchmark. `native_strict` selects the +separately published strict rubric. Their protocol identifiers remain distinct. + +Both modes retain the official input interpolation, WRONG-first verdict parsing, +and default-zero handling of successful responses without a valid label. +`response.metadata` contains `llm_judge_score`, `label`, `parse_default_zero`, +`default_judge_protocol`, `prompt`, `raw_response`, and `usage`. Provider errors +remain failed Jobs. Judge usage is evaluation cost, not method cost. + +Defaults are `qwen3.8-max`, temperature 0, thinking disabled and 4,096 output tokens. +`JUDGE_MODEL_NAME` can override the Judge model independently of the candidate. +Candidate methods and their workspaces are not dependencies of this plugin. +The benchmark runner invokes this Job in a separate application and saves its +prompt, provider output and verdict alongside the candidate prediction. Its +`judge-settings.json` records the validated model, generation/retry settings and +protocol identifier through a credential-free whitelist, rather than recording +only the default model name. + +Sources: [MemBase executable judge](https://github.com/zjunlp/MemBase/blob/a9ba9495d48541d3e68588ff097294cc7a1201bc/membase/inference_utils/prompts.py) +and [MobileMem published strict rubric](https://github.com/zjunlp/MobileMem/blob/fa5cd13c9c659699df9caf650987ba2bea8ca085/omni/eval/eval/question_answering_and_judge_prompts.txt). +The upstream prompt/parser code retains its MIT notice in `LICENSE-CODE`. diff --git a/plugins/mobilemem-judge/pyproject.toml b/plugins/mobilemem-judge/pyproject.toml new file mode 100644 index 00000000..c9b03c05 --- /dev/null +++ b/plugins/mobilemem-judge/pyproject.toml @@ -0,0 +1,22 @@ +[project] +name = "reme-mobilemem-judge" +version = "0.1.0" +description = "Official MobileMem-Omni judge protocols as an independent 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"] +mobilemem-judge = "judge_mobilemem" + +[tool.setuptools.packages.find] +where = ["src"] +include = ["judge_mobilemem*"] + +[tool.setuptools.package-data] +judge_mobilemem = ["plugin.yaml"] + +[build-system] +requires = ["setuptools>=77", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/plugins/mobilemem-judge/src/judge_mobilemem/__init__.py b/plugins/mobilemem-judge/src/judge_mobilemem/__init__.py new file mode 100644 index 00000000..f243bdb5 --- /dev/null +++ b/plugins/mobilemem-judge/src/judge_mobilemem/__init__.py @@ -0,0 +1,5 @@ +"""Independent MobileMem evaluation-only plugin.""" + +from .llm_judge import MobileMemJudgeStep + +__all__ = ["MobileMemJudgeStep"] diff --git a/plugins/mobilemem-judge/src/judge_mobilemem/llm_judge.py b/plugins/mobilemem-judge/src/judge_mobilemem/llm_judge.py new file mode 100644 index 00000000..8b11ca7d --- /dev/null +++ b/plugins/mobilemem-judge/src/judge_mobilemem/llm_judge.py @@ -0,0 +1,43 @@ +"""Official MobileMem binary 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 judge_prompt, parse_judge, protocol_id + + +class MobileMemJudgeStep(BaseStep): + """Score saved predictions without accessing the candidate memory or tools.""" + + async def execute(self): + """Apply the unchanged benchmark prompt and verdict parser.""" + mode = self.context.get("judge_protocol", "membase_default") + prompt = judge_prompt( + self.context["question"], + self.context["reference"], + self.context["prediction"], + self.context.get("evidence", ()), + judge_protocol=mode, + ) + 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 + parsed = parse_judge(raw) + self.context.response.success = True + self.context.response.answer = str(parsed["llm_judge_score"]) + self.context.response.metadata.update( + **parsed, + default_judge_protocol=protocol_id(mode), + prompt=prompt, + raw_response=raw, + usage=asdict(response.usage) if response is not None and response.usage is not None else None, + ) + return self.context.response diff --git a/plugins/mobilemem-judge/src/judge_mobilemem/plugin.yaml b/plugins/mobilemem-judge/src/judge_mobilemem/plugin.yaml new file mode 100644 index 00000000..73b17a10 --- /dev/null +++ b/plugins/mobilemem-judge/src/judge_mobilemem/plugin.yaml @@ -0,0 +1,40 @@ +backends: + mobilemem_judge_step: judge_mobilemem.llm_judge:MobileMemJudgeStep + +application_defaults: + jobs: + answer_judge: + backend: base + description: "MobileMem-Omni official judge protocol for a saved prediction" + parameters: + type: object + properties: + question: {type: string} + reference: {type: string} + prediction: {type: string} + evidence: {type: array, default: []} + judge_protocol: + type: string + enum: [membase_default, native_strict] + default: membase_default + required: [question, reference, prediction] + steps: + - backend: mobilemem_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: 4096 diff --git a/plugins/mobilemem-judge/src/judge_mobilemem/protocol.py b/plugins/mobilemem-judge/src/judge_mobilemem/protocol.py new file mode 100644 index 00000000..c2ba851f --- /dev/null +++ b/plugins/mobilemem-judge/src/judge_mobilemem/protocol.py @@ -0,0 +1,118 @@ +"""Unchanged MobileMem/MemBase judge prompts and binary parser. + +Adapted from ZJUNLP's MobileMem and MemBase revisions below. +MIT License, Copyright (c) 2026 ZJUNLP; see LICENSE-CODE. +""" + +import json +import re +from string import Template + +SOURCE_REVISION = "fa5cd13c9c659699df9caf650987ba2bea8ca085" +MEMBASE_REVISION = "a9ba9495d48541d3e68588ff097294cc7a1201bc" +PROTOCOL_ID = "mobilemem-omni-membase-default-judge-v1" +STRICT_PROTOCOL_ID = "mobilemem-omni-final-memory-native-v1" +TEMPLATE_NAME = "mobilemem-omni-judge" +SOURCE_URL = f"https://github.com/zjunlp/MemBase/blob/{MEMBASE_REVISION}/membase/inference_utils/prompts.py" +JUDGE_TEMPLATE = ( + "Your task is to classify an answer to a question as CORRECT or WRONG. " + "You are given a question, a gold answer, a generated answer, and optional " + "evidence from the user's prior conversations.\n\n" + "GRADING RULES:\n" + "1. The generated answer must include every key element in the gold answer.\n" + "2. Unsupported or irrelevant content that changes the meaning makes the answer WRONG.\n" + "3. Equivalent date formats are acceptable when they refer to the same time.\n" + "4. Use the evidence only to verify the answer; do not infer unsupported facts.\n\n" + "Question: $question\n" + "Gold answer: $golden_answers\n" + "Evidence: $evidence\n" + "Generated answer: $prediction\n\n" + 'Return exactly one JSON object: {"label": "CORRECT"} or ' + '{"label": "WRONG"}.' +) +STRICT_JUDGE_TEMPLATE = ( + "Your task is to classify an answer to a question as CORRECT or WRONG. " + "You will be provided with the following information:\n" + "(1) A question (asked by a user to an AI assistant);\n" + "(2) A gold (ground truth) answer;\n" + "(3) A generated answer;\n" + "(4) A set of evidence sentences relevant to the question, which may be used to verify the generated answer.\n" + "The question focuses on information that an AI assistant should recall about the user " + "based on prior conversations. " + "The gold answer is typically concise and includes all key facts.\n\n" + "GRADING RULES (Must Be Strictly Followed)\n" + "1. COMPLETE COVERAGE REQUIRED: The generated answer must include every key element present in the gold answer. " + "Missing any required fact results in a WRONG classification.\n" + "2. NO REDUNDANT CONTENT: Including excessive or irrelevant details not supported by the evidence or gold answer " + "results in a WRONG classification.\n" + "3. TIME-RELATED QUESTIONS: Different date formats (e.g., May 7th vs. 7 May) are acceptable " + "as long as the same date or time period is referenced.\n\n" + "INSTRUCTIONS FOR THIS TASK\n" + "Use the provided evidence only as supporting context for verification when available. " + "Do not infer facts that are not present in the question, gold answer, or evidence.\n" + "First, provide a brief (one-sentence) explanation of your reasoning. Then, output the label CORRECT or WRONG.\n" + "Do not include both labels in your response, as this will break the evaluation script.\n" + "Return the label in JSON format with the key label.\n\n" + "QUESTION TO EVALUATE\n" + "Question: [Question]\n" + "Gold answer: [Gold Answer]\n" + "Evidence: [Evidences]\n" + "Generated answer: [Generated Answer]" +) + + +def protocol_id(judge_protocol: str) -> str: + """Identify the selected rubric without silently substituting another.""" + if judge_protocol == "membase_default": + return PROTOCOL_ID + if judge_protocol == "native_strict": + return STRICT_PROTOCOL_ID + raise ValueError(f"Unsupported MobileMem judge protocol: {judge_protocol!r}") + + +def judge_prompt( + question: str, + reference: str, + prediction: str, + evidence=(), + *, + judge_protocol="membase_default", +) -> str: + """Preserve native single-pass substitution and untrimmed evidence.""" + protocol_id(judge_protocol) + evidence_text = "\n".join(str(item) for item in evidence) or "Not provided." + if judge_protocol == "membase_default": + return Template(JUDGE_TEMPLATE).substitute( + question=question, + golden_answers=reference, + evidence=evidence_text, + prediction=prediction, + ) + replacements = { + "[Question]": question, + "[Gold Answer]": reference, + "[Evidences]": evidence_text, + "[Generated Answer]": prediction, + } + return re.sub( + r"\[Question\]|\[Gold Answer\]|\[Evidences\]|\[Generated Answer\]", + lambda match: replacements[match[0]], + STRICT_JUDGE_TEMPLATE, + ) + + +def parse_judge(text: str) -> dict: + """Preserve MemBase's WRONG-first binary parsing and native default zero.""" + text = text.strip() + try: + parsed = json.loads(text) + label = str(parsed.get("label", "")) + except (json.JSONDecodeError, AttributeError): + label = text + normalised = label.strip().upper() + score = 0.0 if "WRONG" in normalised else float("CORRECT" in normalised) + return { + "llm_judge_score": score, + "label": "CORRECT" if score else "WRONG", + "parse_default_zero": "WRONG" not in normalised and "CORRECT" not in normalised, + } diff --git a/plugins/mobilemem/README.md b/plugins/mobilemem/README.md new file mode 100644 index 00000000..268e4de7 --- /dev/null +++ b/plugins/mobilemem/README.md @@ -0,0 +1,17 @@ +# MobileMem plugin + +Install this plugin from the repository: + +```bash +pip install -e plugins/mobilemem +``` + +Enable the `mobilemem` ReMe plugin for the MobileMem-Omni benchmark. Its Steps use native ReMe Auto Memory, +source-line citations, BM25 retrieval, workspace containment, and an AgentScope bridge for the native image +reader. The plugin installs independently of MemLens. Evaluation uses the separate +[`mobilemem-judge`](../mobilemem-judge/README.md) plugin. + +The runner supplies only original dialogue messages and image attachments. Memory extraction retains dates, +updates, personal relationships, app details and visual evidence; each question uses a fresh answer Agent. Source +session reading and `read_image` remain scoped to that user's workspace. Supplied captions and benchmark gold +annotations must never be passed to either Job. diff --git a/plugins/mobilemem/pyproject.toml b/plugins/mobilemem/pyproject.toml new file mode 100644 index 00000000..5e5461ea --- /dev/null +++ b/plugins/mobilemem/pyproject.toml @@ -0,0 +1,25 @@ +[project] +name = "reme-mobilemem" +version = "0.1.0" +description = "MobileMem-Omni native 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"] +mobilemem = "reme_mobilemem" + +[project.entry-points."reme.configs"] +mobilemem-benchmark = "reme_mobilemem:BENCHMARK_CONFIG" + +[tool.setuptools.packages.find] +where = ["src"] +include = ["reme_mobilemem*"] + +[tool.setuptools.package-data] +reme_mobilemem = ["plugin.yaml", "*.yaml"] + +[build-system] +requires = ["setuptools>=77", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/plugins/mobilemem/src/reme_mobilemem/__init__.py b/plugins/mobilemem/src/reme_mobilemem/__init__.py new file mode 100644 index 00000000..e9027ef5 --- /dev/null +++ b/plugins/mobilemem/src/reme_mobilemem/__init__.py @@ -0,0 +1,5 @@ +"""MobileMem-Omni benchmark Jobs for ReMe.""" + +from pathlib import Path + +BENCHMARK_CONFIG = Path(__file__).with_name("benchmark.yaml") diff --git a/plugins/mobilemem/src/reme_mobilemem/agentic_answer.py b/plugins/mobilemem/src/reme_mobilemem/agentic_answer.py new file mode 100644 index 00000000..90f19510 --- /dev/null +++ b/plugins/mobilemem/src/reme_mobilemem/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 MobileMemAgenticAnswerStep(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("MobileMem 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"mobilemem_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/mobilemem/src/reme_mobilemem/agentic_answer.yaml b/plugins/mobilemem/src/reme_mobilemem/agentic_answer.yaml new file mode 100644 index 00000000..661c9f20 --- /dev/null +++ b/plugins/mobilemem/src/reme_mobilemem/agentic_answer.yaml @@ -0,0 +1,8 @@ +system_prompt: | + Answer the supplied MobileMem question from this user's saved memories and original source conversations. + Search for the relevant people, events, dates and distinctive details, then read the source files or exact line + ranges when additional context is needed. Distinguish changes over time and similar events. When read_image is + available, reopen cited session/images paths to verify visual details rather than guessing them from an ID. + Use another evidence search when the first result is incomplete. Follow the supplied benchmark answer rules. + Include all facts needed to answer the question, without unsupported or irrelevant details. Do not include the + retrieval process or tool transcript in the final answer. diff --git a/plugins/mobilemem/src/reme_mobilemem/auto_memory.py b/plugins/mobilemem/src/reme_mobilemem/auto_memory.py new file mode 100644 index 00000000..3c3c5dcc --- /dev/null +++ b/plugins/mobilemem/src/reme_mobilemem/auto_memory.py @@ -0,0 +1,34 @@ +"""MobileMem memory prompts with optional, exact source-line references.""" + +from agentscope.message import Msg + +from reme.steps.evolve.auto_memory import AutoMemoryStep + + +class MobileMemAutoMemoryStep(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/mobilemem/src/reme_mobilemem/auto_memory.yaml b/plugins/mobilemem/src/reme_mobilemem/auto_memory.yaml new file mode 100644 index 00000000..7afcac59 --- /dev/null +++ b/plugins/mobilemem/src/reme_mobilemem/auto_memory.yaml @@ -0,0 +1,17 @@ +system_prompt: | + Maintain the user's long-term mobile memory from this original dated conversation. Save concrete personal facts, + names, relationships, dates, quantities, locations, plans, decisions and preferences. Distinguish separate events + across the year and preserve earlier dated facts when newer information updates them. Keep whose experience + each fact describes explicit. Use the supplied session timestamp for relative dates only when justified. + + When original images are present, understand them jointly with surrounding text. Retain distinctive visual + details, visible app text, amounts, dates, people, objects and activities. Preserve each supplied Image ID and + image source path with its facts so the original can be reopened. An identifier is not evidence of its contents. + Do not infer unavailable image details or introduce outside information. + + Cite useful physical source-line ranges with the full supplied source path when source access is enabled. + Otherwise retain self-contained facts. Follow the normal daily-note contract with a stable filename stem, + informative description and factual body. Never set status. Use structured file tool calls to save the memory. + + Create at most one new note for this session. After a successful daily_write, finish with the requested + one-sentence summary. Further writes are needed only to correct a tool error or a concrete omission. diff --git a/plugins/mobilemem/src/reme_mobilemem/benchmark.yaml b/plugins/mobilemem/src/reme_mobilemem/benchmark.yaml new file mode 100644 index 00000000..0fba4066 --- /dev/null +++ b/plugins/mobilemem/src/reme_mobilemem/benchmark.yaml @@ -0,0 +1,32 @@ +extends: benchmark +plugins: [mobilemem] +components: + file_store: + default: + embedding_store: null + as_llm: + default: &candidate_model + backend: openai + model: ${LLM_MODEL_NAME:-qwen3.8-max} + stream: false + context_size: 1000000 + max_retries: 3 + retry_delay: 60 + credential: + api_key: ${LLM_API_KEY:-} + base_url: ${LLM_BASE_URL:-} + extra_body: {enable_thinking: true} + parameters: + temperature: 0 + thinking_enable: true + max_tokens: 65536 + client_kwargs: {timeout: 600, max_retries: 0} + bench: *candidate_model + compressor: *candidate_model + agent_wrapper: + default: + context_config: {max_image_num: 4096} + model_config: {max_retries: 0} + bench: + context_config: {max_image_num: 4096} + model_config: {max_retries: 0} diff --git a/plugins/mobilemem/src/reme_mobilemem/plugin.yaml b/plugins/mobilemem/src/reme_mobilemem/plugin.yaml new file mode 100644 index 00000000..c4e2bdc5 --- /dev/null +++ b/plugins/mobilemem/src/reme_mobilemem/plugin.yaml @@ -0,0 +1,71 @@ +backends: + mobilemem_auto_memory_step: reme_mobilemem.auto_memory:MobileMemAutoMemoryStep + mobilemem_agentic_answer_step: reme_mobilemem.agentic_answer:MobileMemAgenticAnswerStep + +application_defaults: + jobs: + read_image: + backend: base + description: "Read an original saved session image through ReMe's native 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 original MobileMem-Omni session using native Auto Memory." + include_images: true + 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: true} + include_source_sessions: {type: boolean, default: true} + memory_hint: {type: string, default: ""} + required: [messages, session_id] + steps: + - backend: mobilemem_auto_memory_step + search: + backend: base + description: "Search this user'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 MobileMem question through a fresh evidence retrieval Agent." + include_images: true + 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: []} + max_iters: {type: integer, minimum: 1, default: 30} + include_images: {type: boolean, default: true} + include_source_sessions: {type: boolean, default: true} + required: [query] + steps: + - backend: mobilemem_agentic_answer_step + agent_wrapper: bench + components: + file_store: + default: + embedding_store: null diff --git a/tests/unit/test_mobilemem_audit.py b/tests/unit/test_mobilemem_audit.py new file mode 100644 index 00000000..e9a78def --- /dev/null +++ b/tests/unit/test_mobilemem_audit.py @@ -0,0 +1,475 @@ +"""Failure-path receipts and cost accounting, without provider or server calls.""" + +# pylint: disable=missing-function-docstring,protected-access +import asyncio +import json +from types import SimpleNamespace + +import pytest + +from benchmark.mobilemem import run +from benchmark.mobilemem._audit import ImageInputError, error_record, mean_tokens, response_record +from reme.schema import Response + + +def read(path): + return json.loads(path.read_text(encoding="utf-8")) + + +@pytest.fixture(name="flow") +def fixture_flow(tmp_path, monkeypatch): + scenes = [ + { + "id": "uuid_0", + "image_root": tmp_path, + "sessions": [{"session_id": "s0", "event_start_time": "2025-01-01T09:00:00"}], + "questions": [ + {"question_id": "q0", "question": "What?", "answer": "answer", "question_type": "visual_reasoning"}, + ], + }, + ] + state = SimpleNamespace( + scenes=scenes, + apps=[], + fail=None, + direct=False, + cancel=False, + cleanup_fail=False, + startup_fail=False, + constructor_fail=False, + calls=[], + ) + + class FakeApp: + """Mock one owned candidate or Judge application.""" + + def __init__(self, *, judge): + if state.constructor_fail and not judge: + raise ValueError("candidate construction failed") + self.judge, self.closed = judge, False + state.apps.append(self) + + async def start(self): + if state.startup_fail and not self.judge: + raise ValueError("candidate startup failed") + + async def close(self): + self.closed = True + if state.cleanup_fail and not self.judge: + raise ValueError("candidate cleanup failed") + + async def run_job(self, name, **kwargs): + identity = kwargs.get("session_id", kwargs.get("query", "")) + state.calls.append((name, kwargs)) + if name == "answer_judge": + assert run._CURRENT.get() is None # Judge traffic cannot join candidate cost. + elif name != "index_update": + run._CURRENT.get().append(run.normalized_usage({"prompt_tokens": 100, "completion_tokens": 20})) + await asyncio.sleep(0) # Force interleaving of independently scoped requests. + if state.fail in (name, identity): + if state.cancel: + raise asyncio.CancelledError("cancelled offline call") + if state.direct: + raise ValueError("direct offline failure") + return Response(success=False, answer="HTTP 400 input length exceeds limit") + if name == "answer_judge": + return Response(answer="1", metadata={"llm_judge_score": 1.0, "parse_default_zero": False}) + # This portable counter is deliberately different: never sum it with HTTP receipts. + return Response(answer="answer", metadata={"usage": {"total_tokens": 999999}}) + + monkeypatch.setattr(run, "Application", FakeApp) + monkeypatch.setattr(run, "candidate_config", lambda *_: {"judge": False}) + monkeypatch.setattr(run, "judge_config", lambda *_: {"judge": True}) + monkeypatch.setattr(run, "judge_settings", lambda _: {"model": "offline", "protocol": "offline"}) + monkeypatch.setattr(run, "session_messages", lambda *_: [{"original": "unchanged"}]) + monkeypatch.setattr(run, "observe_usage", lambda _: None) + monkeypatch.setattr(run, "observe_completion", lambda _: None) + + async def execute(): + complete = await run.run(scenes, tmp_path / "run", "candidate", 4, 16) + return complete, read(tmp_path / "run/summary.json") + + state.execute = execute + state.directory = tmp_path / "run/users/uuid_0" + return state + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct", [False, True]) +async def test_failed_second_memory_keeps_both_requests(flow, direct): + flow.scenes[0]["sessions"].append({"session_id": "s1", "event_start_time": "2025-01-02T09:00:00"}) + flow.fail, flow.direct = "s1", direct + complete, summary = await flow.execute() + assert not complete and summary["completed"] == 0 + assert summary["memory_usage"]["total_tokens"] == summary["token_usage"]["total_tokens"] == 240 + assert summary["token_usage"]["records"] == 2 + assert not (flow.directory / "memory_usage.json").exists() + response = read(flow.directory / "sessions/00001.json")["response"] + assert response is None if direct else response["success"] is False + error = read(flow.directory / "error.json") + assert (error["stage"], error["user_id"], error["session_id"]) == ("memory", "uuid_0", "s1") + assert all(app.closed for app in flow.apps) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct", [False, True]) +async def test_failed_answer_keeps_memory_and_answer_cost_once(flow, direct): + flow.fail, flow.direct = "agentic_answer", direct + complete, summary = await flow.execute() + assert not complete and summary["completed"] == 0 and summary["metrics"]["llm_judge"] is None + assert summary["token_usage"]["total_tokens"] == 240 + assert summary["memory_usage"]["total_tokens"] == summary["answer_usage"]["total_tokens"] == 120 + assert summary["usage_counts"]["attempted_answers"] == 1 + assert summary["usage_counts"]["successful_answers"] == 0 + assert summary["mean_tokens_per_question"]["total_tokens"] == 240 + assert summary["usage_denominators"] == {"scheduled_users": 1, "expected_questions": 1} + assert summary["usage_policy"]["version"] == 2 + assert (flow.directory / "memory_usage.json").exists() + question = flow.directory / "questions/0000" + receipt = read(question / "answer-response.json") + assert receipt["usage"]["total_tokens"] == 120 + assert (receipt["response"] is None) is direct + assert not (question / "answer.json").exists() and not (question / "result.json").exists() + assert read(question / "error.json")["stage"] == "answer" + + +@pytest.mark.asyncio +async def test_preprocessing_failure_keeps_earlier_memory_and_image_identity(flow, monkeypatch): + flow.scenes[0]["sessions"].append({"session_id": "s1", "event_start_time": "2025-01-02T09:00:00"}) + + def messages(session, *_): + if session["session_id"] == "s1": + try: + raise FileNotFoundError("missing original image") + except FileNotFoundError as error: + raise ImageInputError("image/uid0/photo.png", error) from error + return [] + + monkeypatch.setattr(run, "session_messages", messages) + _, summary = await flow.execute() + assert summary["memory_usage"]["total_tokens"] == 120 + receipt = read(flow.directory / "sessions/00001.json") + assert receipt["response"] is None and receipt["usage"]["records"] == 0 + error = read(flow.directory / "error.json") + assert error["stage"] == "preprocessing" and error["image_path"] == "image/uid0/photo.png" + assert [cause["type"] for cause in error["chain"]] == ["ImageInputError", "FileNotFoundError"] + assert not (flow.directory / "memory_usage.json").exists() + + +@pytest.mark.asyncio +async def test_index_failure_is_saved_without_complete_memory_marker(flow): + flow.fail = "index_update" + _, summary = await flow.execute() + assert summary["memory_usage"]["total_tokens"] == 120 + assert read(flow.directory / "index-response.json")["response"]["success"] is False + assert read(flow.directory / "error.json")["stage"] == "index" + assert not (flow.directory / "memory_usage.json").exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct", [False, True]) +async def test_judge_failure_preserves_answer_without_double_counting(flow, direct): + flow.fail, flow.direct = "answer_judge", direct + complete, summary = await flow.execute() + assert not complete and summary["token_usage"]["total_tokens"] == 240 + assert summary["answer_usage"]["records"] == 1 + question = flow.directory / "questions/0000" + assert read(question / "answer-response.json")["response"]["success"] is True + assert not (question / "answer.json").exists() + assert not (question / "result.json").exists() + assert read(question / "error.json")["stage"] == "judge" + assert (read(question / "judge.json") is None) is direct + + +@pytest.mark.asyncio +async def test_all_success_old_payloads_metrics_and_usage_stay_compatible(flow): + flow.scenes[0]["sessions"].append({"session_id": "s1", "event_start_time": "2025-01-02T09:00:00"}) + complete, summary = await flow.execute() + assert complete and summary["completed"] == summary["expected"] == 1 + assert summary["metrics"] == {"f1": 1.0, "bleu_1": 1.0, "llm_judge": 1.0} + assert summary["token_usage"]["total_tokens"] == 360 + assert summary["memory_usage"]["total_tokens"] == 240 + assert summary["answer_usage"]["total_tokens"] == 120 + assert summary["mean_tokens_per_user"]["total_tokens"] == summary["mean_tokens_per_question"]["total_tokens"] == 360 + assert summary["token_usage"]["cached_input_tokens"] is None + assert summary["token_usage"]["coverage"]["cached_input_tokens"] == 0 + assert read(flow.directory / "memory_usage.json") == summary["memory_usage"] + assert [name for name, _ in flow.calls] == [ + "auto_memory", + "auto_memory", + "index_update", + "agentic_answer", + "answer_judge", + ] + question = flow.directory / "questions/0000" + assert read(question / "answer-response.json")["response"]["success"] is True + assert not (question / "answer.json").exists() + result = read(question / "result.json") + assert result == { + "question_id": "q0", + "category": "visual_reasoning", + "question": "What?", + "reference": "answer", + "prediction": "answer", + "metrics": summary["metrics"], + "usage": summary["answer_usage"], + "judge_parse_default_zero": False, + } + + +@pytest.mark.asyncio +async def test_concurrent_questions_keep_cost_scopes_and_scheduled_denominator(flow): + flow.scenes[0]["questions"].append({**flow.scenes[0]["questions"][0], "question_id": "q1"}) + flow.scenes.append({**flow.scenes[0], "id": "uuid_1"}) + complete, summary = await flow.execute() + assert complete and summary["completed"] == 4 + assert summary["token_usage"]["records"] == 6 and summary["token_usage"]["total_tokens"] == 720 + assert summary["mean_tokens_per_user"]["total_tokens"] == 360 + assert summary["mean_tokens_per_question"]["total_tokens"] == 180 + + +@pytest.mark.asyncio +async def test_failure_means_use_scheduled_not_successful_denominators(flow): + flow.scenes[0]["questions"].append({**flow.scenes[0]["questions"][0], "question_id": "q1"}) + flow.fail = "agentic_answer" + _, summary = await flow.execute() + assert summary["completed"] == 0 and summary["expected"] == 2 + assert summary["token_usage"]["total_tokens"] == 360 + assert summary["mean_tokens_per_question"]["total_tokens"] == 180 + + +@pytest.mark.asyncio +async def test_empty_scheduled_population_preserves_unknown_means(flow): + flow.scenes.clear() + _, summary = await flow.execute() + assert summary["usage_denominators"] == {"scheduled_users": 0, "expected_questions": 0} + assert all(value is None for value in summary["mean_tokens_per_user"].values()) + assert all(value is None for value in summary["mean_tokens_per_question"].values()) + assert summary["token_usage"]["records"] == 0 + + +@pytest.mark.asyncio +async def test_summary_ignores_legacy_answer_and_result_usage_copies(flow, monkeypatch): + original = run.close_application + + async def close(app, directory, **kwargs): + await original(app, directory, **kwargs) + if app is not None and app.judge: + for name in ("answer.json", "result.json"): + path = flow.directory / "questions/0000" / name + data = read(path) if path.exists() else {"response": {"success": True}} + data["usage"] = {"total_tokens": 999999} + run.save(path, data) + + monkeypatch.setattr(run, "close_application", close) + _, summary = await flow.execute() + assert summary["token_usage"]["total_tokens"] == 240 + assert summary["answer_usage"]["records"] == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("constructor", [False, True]) +async def test_startup_and_cleanup_failures_are_separate(flow, constructor): + flow.constructor_fail, flow.startup_fail = constructor, not constructor + flow.cleanup_fail = True + complete, summary = await flow.execute() + assert not complete and summary["token_usage"]["total_tokens"] is None + assert read(flow.directory / "error.json")["stage"] == "startup" + assert (flow.directory / "cleanup-error.json").exists() is not constructor + assert all(app.closed for app in flow.apps) + + +@pytest.mark.asyncio +async def test_cancelled_request_finalizes_once_and_closes_owned_apps(flow): + flow.fail, flow.cancel = "auto_memory", True + with pytest.raises(asyncio.CancelledError): + await flow.execute() + receipt = read(flow.directory / "sessions/00000.json") + assert receipt["response"] is None and receipt["usage"]["total_tokens"] == 120 + assert read(flow.directory / "error.json")["error_type"] == "CancelledError" + assert all(app.closed for app in flow.apps) + assert read(flow.directory.parents[1] / "summary.json")["memory_usage"]["total_tokens"] == 120 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel", [False, True]) +async def test_failed_task_joins_sibling_cleanup_before_propagating(cancel): + started, closed = asyncio.Event(), asyncio.Event() + + async def waiting(): + try: + started.set() + await asyncio.Event().wait() + finally: + await asyncio.sleep(0) + closed.set() + + async def interrupted(): + await started.wait() + if cancel: + raise asyncio.CancelledError() + raise ValueError("offline task failed") + + with pytest.raises(asyncio.CancelledError if cancel else ValueError): + await run.gather_tasks([waiting(), interrupted()]) + assert closed.is_set() + + +def test_error_artifacts_remove_secrets_payloads_and_locals(monkeypatch): + secret = "offline-sensitive-value" + monkeypatch.setenv("LLM_API_KEY", secret) + binary = "YWJj" * 100 + response = Response( + success=False, + answer=f"input length error api_key={secret} data:image/png;base64,{binary}", + metadata={ + "headers": {"X-Private": "hidden-header"}, + "request": {"body": "hidden-body"}, + "credential": {"token": "hidden-credential"}, + "traceback": "hidden-local", + "status_code": 400, + }, + ) + try: + raise ValueError(f"https://user:password@example.invalid/?key={secret} Bearer token-value sk-private-key") + except ValueError as error: + record = error_record(error, stage="answer", user_id="uuid_0", question_id="q0", response=response) + serialized = json.dumps([record, response_record(response)]) + for value in ( + secret, + binary, + "hidden-header", + "hidden-body", + "hidden-credential", + "hidden-local", + "token-value", + "sk-private-key", + "user:password", + ): + assert value not in serialized + assert "input length" in record["reason"] + assert set(record["chain"][0]["traceback"][0]) == {"file", "line", "function"} + assert ImageInputError("../../private.png", ValueError("bad")).image_path == "[invalid image path]" + successful = Response(answer="URL https://example.test/ and literal sk-example-text", metadata={"raw": binary}) + assert response_record(successful) == successful.model_dump() + + +@pytest.mark.parametrize( + "name", + ["token", "refresh_token", "id_token", "access_token", "auth_token", "private_key", "private-key", "privateKey"], +) +def test_credentials_are_redacted_in_failed_responses_and_error_chains(name): + sentinel = "0123456789abcdef" * 4 + usage = {"input_tokens": 100, "output_tokens": 20, "total_tokens": 120, "cached_input_tokens": 30} + response = Response( + success=False, + answer=f'provider rejection {name}="{sentinel}"', + metadata={name: sentinel, "usage": usage}, + ) + try: + try: + raise ValueError(f"provider {name}: {sentinel}") + except ValueError as cause: + raise RuntimeError(f"job {name}='{sentinel}'") from cause + except RuntimeError as error: + diagnostic = error_record(error, stage="answer", user_id="uuid_0", question_id="q0", response=response) + record = response_record(response) + assert sentinel not in json.dumps([record, diagnostic]) + assert record["metadata"] == {name: "[redacted]", "usage": usage} + assert record["answer"] == f"provider rejection {name}=[redacted]" + assert diagnostic["reason"] == record["answer"] + assert [item["reason"] for item in diagnostic["chain"]] == [ + f"job {name}=[redacted]", + f"provider {name}: [redacted]", + ] + assert diagnostic["stage"] == "answer" and diagnostic["question_id"] == "q0" + successful = response.model_copy(update={"success": True}) + assert response_record(successful) == successful.model_dump() + + +@pytest.mark.parametrize( + "sentinel", + ["0123456789abcdef" * 4, "a", "success", "answer", "metadata", "usage", "coverage", "records", "input_tokens"], +) +def test_private_key_environment_values_are_redacted_without_changing_usage_schema(monkeypatch, sentinel): + monkeypatch.setenv("TEST_PRIVATE_KEY", sentinel) + usage = {"input_tokens": 100, "output_tokens": 20, "records": 1, "coverage": {"input_tokens": 1}} + response = Response( + success=False, + answer=f"provider returned {sentinel}", + metadata={"usage": usage, "private_key": "PRIVATE_SENTINEL", f"debug {sentinel}": f"detail {sentinel}"}, + ) + record = response_record(response) + assert set(record) == {"success", "answer", "metadata"} + assert Response.model_validate(record).success is False + assert record["answer"] == "provider returned [redacted]" + assert record["metadata"] == { + "usage": usage, + "private_key": "[redacted]", + "debug [redacted]": "detail [redacted]", + } + diagnostic = error_record(RuntimeError(f"provider returned {sentinel}"), stage="startup") + assert diagnostic["reason"] == diagnostic["chain"][0]["reason"] == "provider returned [redacted]" + + +@pytest.mark.parametrize("header", ["Cookie", "Set-Cookie"]) +def test_failed_response_redacts_whole_cookie_header_without_changing_ordinary_prose(header): + response = Response( + success=False, + answer=f"{header}: session=opaque; other=secondsecret\nThe cookie parser failed with status 400.", + ) + record = response_record(response) + assert "opaque" not in record["answer"] and "secondsecret" not in record["answer"] + assert "The cookie parser failed with status 400." in record["answer"] + + +@pytest.mark.asyncio +async def test_http_observer_includes_compressor_subsets_and_null_coverage_without_duplicate_wrapping(): + async def send(request): + await asyncio.sleep(0) + return SimpleNamespace(status_code=200, is_stream_consumed=True, json=lambda: {"usage": request}) + + clients = [SimpleNamespace(send=send) for _ in range(3)] + app = SimpleNamespace( + context=SimpleNamespace( + components={ + "as_llm": { + name: SimpleNamespace(model=SimpleNamespace(client=SimpleNamespace(_client=client))) + for name, client in zip(("default", "bench", "compressor"), clients) + }, + }, + ), + ) + run.observe_usage(app) + run.observe_usage(app) + + async def scoped(amount): + with run.usage_scope() as records: + await clients[0].send( + {"prompt_tokens": amount, "completion_tokens": 20, "prompt_tokens_details": {"cached_tokens": 30}}, + ) + await clients[2].send( + {"prompt_tokens": 50, "completion_tokens": 10, "completion_tokens_details": {"reasoning_tokens": 4}}, + ) + return run.aggregate_usage(records) + + first, second = await asyncio.gather(scoped(100), scoped(200)) + assert first["records"] == second["records"] == 2 + assert (first["total_tokens"], second["total_tokens"]) == (180, 280) + assert first["cached_input_tokens"] == 30 and first["noncached_input_tokens"] == 70 + assert first["coverage"]["cached_input_tokens"] == 1 and first["reasoning_tokens"] == 4 + assert mean_tokens(first, 0)["total_tokens"] is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failed_job,completed,cost", [("auto_memory", 0, 120), (None, 1, 240)]) +async def test_cleanup_failure_preserves_primary_outcome(flow, failed_job, completed, cost): + flow.fail, flow.cleanup_fail = failed_job, True + complete, summary = await flow.execute() + assert not complete and summary["completed"] == completed + assert summary["cleanup_error_count"] == 1 and summary["token_usage"]["total_tokens"] == cost + assert read(flow.directory / "cleanup-error.json")["stage"] == "cleanup" + if failed_job: + assert read(flow.directory / "error.json")["stage"] == "memory" + else: + assert summary["metrics"]["llm_judge"] == 1.0 diff --git a/tests/unit/test_mobilemem_benchmark.py b/tests/unit/test_mobilemem_benchmark.py new file mode 100644 index 00000000..54fc256f --- /dev/null +++ b/tests/unit/test_mobilemem_benchmark.py @@ -0,0 +1,407 @@ +"""Focused offline benchmark contracts; no external services or experiment writes.""" + +# pylint: disable=missing-function-docstring,protected-access +import base64 +from importlib import import_module +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.mobilemem import run as mobilemem + + +def test_candidate_declares_required_image_dependency(): + manifest = Path(__file__).resolve().parents[2] / "plugins/mobilemem/pyproject.toml" + project = tomllib.loads(manifest.read_text(encoding="utf-8"))["project"] + assert "pillow>=10.0.0" in project["dependencies"] + + +@pytest.fixture(name="mobile_scene") +def fixture_mobile_scene(tmp_path): + folder = tmp_path / "image/uid0/camera_photos" + folder.mkdir(parents=True) + Image.new("RGB", (3, 2), (12, 34, 56)).save(folder / "photo.png") + sessions = [ + { + "session_id": f"0_{index}", + "event_start_time": f"2025-0{index}-01 09:00:00", + "event_end_time": f"2025-0{index}-01 10:00:00", + "dialogue": [ + {"turn": 1, "role": "user", "content_type": "text", "content": f"Original turn {index}"}, + {"turn": 2, "role": "user", "content_type": "image", "image_inline": "image/uid0/event/photo.png"}, + {"turn": 3, "role": "assistant", "content_type": "text", "content": "Original reply"}, + ], + "dialogue_summary": "HIDDEN_SUMMARY", + "memory_points": [{"memory_content": "HIDDEN_MEMORY"}], + } + for index in (2, 1) + ] + question = { + "question_id": "0_q_1", + "question": "What happened?", + "answer": "HIDDEN_GOLD", + "question_format": "open_ended", + "question_type": "visual_reasoning", + "source_session_ids": ["0_1"], + "image_refs": ["HIDDEN_IMAGE"], + "evidence": [{"explanation": "HIDDEN_EVIDENCE"}], + } + (tmp_path / "data.jsonl").write_text(json.dumps({"uuid": 0, "language": "en", "sessions": sessions}) + "\n") + (tmp_path / "questions.jsonl").write_text(json.dumps({"uuid": 0, "questions": [question]}) + "\n") + return tmp_path, mobilemem.load_mobilemem(tmp_path)[0] + + +def test_mobilemem_complete_history_ids_and_no_gold(mobile_scene): + root, scene = mobile_scene + assert [session["session_id"] for session in scene["sessions"]] == ["0_1", "0_2"] + assert scene["questions"][0]["question_type"] == "visual_reasoning" + converted = mobilemem.session_messages(scene["sessions"][0], scene["image_root"]) + assert "HIDDEN" not in json.dumps(converted) + original = root / "image/uid0/camera_photos/photo.png" + before = original.read_bytes() + assert converted[1]["content"][1]["text"] == "[Image ID: image/uid0/event/photo.png]" + assert converted[1]["content"][2]["type"] == "data" + assert base64.b64decode(converted[1]["content"][2]["source"]["data"]) == before + assert original.read_bytes() == before + with pytest.raises(ValueError, match="Invalid MobileMem image path"): + mobilemem.resolve_image(root / "image", "../outside.png", {}) + original.unlink() + with pytest.raises(ValueError, match="Missing or ambiguous MobileMem image"): + mobilemem.session_messages(scene["sessions"][0], scene["image_root"]) + + +def test_official_query_and_overlap_metrics(): + assert mobilemem.query_prompt("QUESTION", "CONTEXT") == ( + "Question:\nQUESTION\nPlease answer the question based on the following memories:\nCONTEXT\n" + + mobilemem.ANSWER_RULES + ) + score = mobilemem.score_answer("北京 北京", "北京") + assert score["f1"] == pytest.approx(2 / 3) + assert score["bleu_1"] == pytest.approx(0.36787944117144233) + + +@pytest.fixture(name="candidate") +def fixture_candidate(tmp_path, monkeypatch): + name = "mobilemem" + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins" / name / "src")) + monkeypatch.chdir(tmp_path) + memory = import_module(f"reme_{name}.auto_memory") + answer = import_module(f"reme_{name}.agentic_answer") + prefix = "MobileMem" + 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 name, memory, answer, prefix, app, wrapper, store, tmp_path + + +def test_candidate_registration_is_application_local(candidate): + name, memory, answer, prefix, *_ = candidate + plugin = _load_plugin(name, EntryPoint(name=name, value=f"reme_{name}", group="reme.plugins")) + registry = create_application_registry() + PluginManager([plugin]).register(registry) + assert registry.get("step", f"{name}_auto_memory_step") is getattr(memory, f"{prefix}AutoMemoryStep") + assert registry.get("step", f"{name}_agentic_answer_step") is getattr(answer, f"{prefix}AgenticAnswerStep") + + +@pytest.mark.asyncio +async def test_candidate_and_judge_start_with_native_settings(tmp_path, monkeypatch, mobile_scene): + from reme import Application + + root = Path(__file__).resolve().parents[2] / "plugins" + for package in ("mobilemem", "mobilemem-judge"): + monkeypatch.syspath_prepend(str(root / package / "src")) + entries = [ + EntryPoint(name="mobilemem", value="reme_mobilemem", group="reme.plugins"), + EntryPoint(name="mobilemem-judge", value="judge_mobilemem", group="reme.plugins"), + EntryPoint(name="mobilemem-benchmark", value="reme_mobilemem:BENCHMARK_CONFIG", group="reme.configs"), + ] + + def installed(group, name): + return [entry for entry in entries if entry.group == group and entry.name == name] + + monkeypatch.setattr("reme.plugin.find_entry_points", installed) + monkeypatch.setattr("reme.config.config_parser.find_entry_points", installed) + monkeypatch.setenv("LLM_API_KEY", "offline") + monkeypatch.setenv("LLM_BASE_URL", "https://provider.invalid/v1") + monkeypatch.setenv("JUDGE_MODEL_NAME", "offline-judge-override") + candidate_app = Application(**mobilemem.candidate_config(tmp_path / "candidate", "candidate-model")) + judge_app = Application(**mobilemem.judge_config(tmp_path / "judge")) + try: + await candidate_app.start() + await judge_app.start() + assert set(judge_app.context.jobs) == {"answer_judge"} + assert set(judge_app.context.components) == {"as_llm"} + audit = mobilemem.judge_settings(judge_app) + assert audit["model"] == "offline-judge-override" + assert audit["parameters"] == {"max_tokens": 4096, "temperature": 0, "thinking_enable": False} + with monkeypatch.context() as environment: + for field in ("temperature", "max_tokens", "thinking_enable"): + environment.setenv("TEST_PRIVATE_KEY", field) + assert mobilemem.judge_settings(judge_app)["parameters"] == audit["parameters"] + assert audit["protocol"] == "mobilemem-omni-membase-default-judge-v1" + assert "credential" not in json.dumps(audit) and "provider.invalid" not in json.dumps(audit) + models = candidate_app.context.components["as_llm"] + for component in models.values(): + assert component.kwargs["model"] == "candidate-model" + assert component.kwargs["stream"] is False and component.kwargs["context_size"] == 1000000 + assert component.kwargs["parameters"] == {"max_tokens": 65536, "thinking_enable": True, "temperature": 0} + assert component.kwargs["max_retries"] == 3 and component.kwargs["retry_delay"] == 60 + assert component.kwargs["client_kwargs"] == {"timeout": 600, "max_retries": 0} + assert candidate_app.config.components["file_store"]["default"].model_dump()["embedding_store"] is None + assert candidate_app.config.jobs["index_update"].model_dump()["watch_dirs"] == [ + "daily_dir", + "digest_dir", + "session_dir/dialog", + ] + from reme_mobilemem.auto_memory import MobileMemAutoMemoryStep + + monkeypatch.setattr(MobileMemAutoMemoryStep, "_list_session_note", AsyncMock(return_value=None)) + wrappers = candidate_app.context.components["agent_wrapper"] + for wrapper in wrappers.values(): + wrapper.reply = AsyncMock(return_value={"result": "answer"}) + _, scene = mobile_scene + messages = mobilemem.session_messages(scene["sessions"][0], scene["image_root"]) + for images, kwargs in ((True, {}), (False, {"include_images": False})): + answer = await candidate_app.run_job("agentic_answer", query="Question", **kwargs) + assert answer.success and answer.metadata["include_images"] is images + toolkit = wrappers["bench"].reply.call_args.kwargs["toolkit"] + assert ((await toolkit.get_tool("read_image")) is not None) is images + memory = await candidate_app.run_job("auto_memory", messages=messages, session_id=f"s-{images}", **kwargs) + assert memory.success + assert isinstance(wrappers["default"].reply.call_args.args[0], Msg) is images + assert (tmp_path / "candidate/session/images" / f"s-{images}").exists() is images + finally: + await judge_app.close() + await candidate_app.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sources", [True, False]) +async def test_memory_saves_originals_and_native_image_links(candidate, sources): + _, memory, _, prefix, app, wrapper, store, workspace = candidate + step = getattr(memory, f"{prefix}AutoMemoryStep")(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, prefix, 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 = getattr(answer, f"{prefix}AgenticAnswerStep")(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, prefix, 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 = getattr(answer, f"{prefix}AgenticAnswerStep")(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_keeps_all_sessions_and_gold_is_judge_only(mobile_scene, tmp_path, monkeypatch): + from benchmark.mobilemem.run import normalized_usage, observe_completion + + _, scene = mobile_scene + calls, apps = [], [] + + class Candidate: + """Offline Application stand-in preserving the native Job call boundary.""" + + def __init__(self, **config): + self.config = config + self.closed = False + self.context = SimpleNamespace( + components={ + "agent_wrapper": {name: SimpleNamespace(session_path=tmp_path) for name in ("default", "bench")}, + }, + ) + apps.append(self) + + 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-agent"}) + + judge = SimpleNamespace( + run_job=AsyncMock( + return_value=Response( + metadata={"llm_judge_score": 1, "parse_default_zero": False}, + ), + ), + ) + monkeypatch.setattr(mobilemem, "Application", Candidate) + + def configured(workspace, model): + return {"workspace": workspace, "model": model} + + monkeypatch.setattr(mobilemem, "candidate_config", configured) + monkeypatch.setattr(mobilemem, "observe_usage", lambda _: None) + monkeypatch.setattr(mobilemem, "observe_completion", lambda _: None) + await mobilemem.evaluate_user( + scene, + tmp_path / "new-user", + "candidate", + judge, + __import__("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"] == ["0_1", "0_2"] + for name, kwargs in calls: + assert "HIDDEN" not in json.dumps(kwargs) + if name in {"auto_memory", "agentic_answer"}: + assert kwargs["include_images"] is kwargs["include_source_sessions"] is True + answer = calls[-1][1] + assert answer["max_iters"] == 30 + assert answer["query"] == mobilemem.query_prompt( + scene["questions"][0]["question"], + "Use the available evidence tools.", + ) + assert "category" not in answer and "question_type" not in answer + judge_inputs = judge.run_job.call_args.kwargs + assert judge_inputs["reference"] == "HIDDEN_GOLD" + assert judge_inputs["evidence"] == ["HIDDEN_EVIDENCE"] + result = json.loads((tmp_path / "new-user/questions/0000/result.json").read_text(encoding="utf-8")) + assert result["metrics"]["llm_judge"] == 1 and apps[0].closed + assert normalized_usage({"prompt_tokens": 100, "completion_tokens": 20})["cached_input_tokens"] is None + assert ( + normalized_usage({"prompt_tokens": 100, "prompt_tokens_details": {"cached_tokens": 30}})[ + "noncached_input_tokens" + ] + == 70 + ) + raw = {"last_message": {"finished_reason": "exceed_max_iters"}, "result": "native answer"} + native = AsyncMock(return_value=raw) + wrapper = SimpleNamespace(reply=native) + observe_completion(SimpleNamespace(context=SimpleNamespace(components={"agent_wrapper": {"bench": wrapper}}))) + assert await wrapper.reply("unchanged input", max_iters=30) is raw + native.assert_awaited_once_with("unchanged input", max_iters=30) + native.return_value = {"last_message": {"finished_reason": "interrupted"}} + with pytest.raises(RuntimeError, match="interrupted"): + await wrapper.reply("unchanged input") diff --git a/tests/unit/test_mobilemem_judge.py b/tests/unit/test_mobilemem_judge.py new file mode 100644 index 00000000..4f8630c8 --- /dev/null +++ b/tests/unit/test_mobilemem_judge.py @@ -0,0 +1,85 @@ +"""Standalone judge Jobs preserve official prompts, parsers and generation defaults.""" + +# pylint: disable=missing-function-docstring,protected-access +import hashlib +from importlib import import_module +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): + name = "mobilemem" + root = Path(__file__).resolve().parents[2] / "plugins" / f"{name}-judge" / "src" + monkeypatch.syspath_prepend(str(root)) + module = import_module(f"judge_{name}.llm_judge") + prefix = "MobileMem" + step = getattr(module, f"{prefix}JudgeStep")() + 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='{"label":"CORRECT"}')], + is_last=True, + usage=ChatUsage(100, 10, 0.01, cache_input_tokens=30), + ), + ) + return name, step, import_module(f"judge_{name}.protocol"), root + + +@pytest.mark.asyncio +async def test_judge_job_uses_only_official_prompt_and_separate_usage(judge): + name, 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 = {"evidence": ["first", "second"]} + 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 / f"judge_{name}" / "plugin.yaml").read_text() + parse_plugin_manifest(text, plugin_name=f"{name}-judge") + manifest = yaml.safe_load(text)["application_defaults"] + assert manifest["jobs"]["answer_judge"]["steps"] == [{"backend": f"{name}_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"] == 4096 + step.as_llm.return_value = ChatResponse(content=[TextBlock(text="malformed")], is_last=True) + assert (await step.execute()).metadata["parse_default_zero"] is True + step.as_llm.side_effect = TimeoutError("provider") + with pytest.raises(TimeoutError): + await step.execute() + + +def test_mobilemem_official_prompt_and_parser_golden_contracts(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins/mobilemem-judge/src")) + protocol = import_module("judge_mobilemem.protocol") + for mode, digest in [ + ("membase_default", "7880f17ac2335307094dc8b02fa2c0093e1bec7b335215e3362e95a785a0f495"), + ("native_strict", "8710fcee94ffaa1b40911073b1f6404b62b93d36a4cc8236f3a10fef73622ca7"), + ]: + prompt = protocol.judge_prompt( + "QUESTION $prediction", + "REFERENCE", + "PREDICTION $question", + ["first", "second"], + judge_protocol=mode, + ) + assert hashlib.sha256(prompt.encode()).hexdigest() == digest + for raw, score, default in [('{"label":"CORRECT"}', 1, False), ("CORRECT WRONG", 0, False), ("bad", 0, True)]: + parsed = protocol.parse_judge(raw) + assert (parsed["llm_judge_score"], parsed["parse_default_zero"]) == (score, default)