diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index ce45e3a4..e6a1e69b 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -118,7 +118,7 @@ jobs: - name: Log in to GHCR if: env.PUBLISH_IMAGE == 'true' - uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4 + uses: docker/login-action@81c6f93d4bc991e510b342a00b419b3bb4efc5c6 # v4 with: registry: ghcr.io username: ${{ github.actor }} @@ -177,7 +177,7 @@ jobs: - uses: docker/setup-buildx-action@f87e5991a6d7451dcb8d9637bfbc97413f497069 # v4 - - uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4 + - uses: docker/login-action@81c6f93d4bc991e510b342a00b419b3bb4efc5c6 # v4 with: registry: ghcr.io username: ${{ github.actor }} diff --git a/.github/workflows/release-auto-fin.yml b/.github/workflows/release-auto-fin.yml index 31d294d1..ccb9f52e 100644 --- a/.github/workflows/release-auto-fin.yml +++ b/.github/workflows/release-auto-fin.yml @@ -133,7 +133,8 @@ jobs: assert set(manifest.backends) == { "auto_fin_data_step", "auto_fin_topic_step", - "auto_fin_merge_step", + "auto_fin_research_step", + "auto_fin_digest_step", } assert set(manifest.application_defaults["jobs"]) == { "auto_fin", @@ -165,6 +166,6 @@ jobs: path: dist/auto-fin - name: Publish reme-auto-fin - uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # release/v1 + uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # release/v1 with: packages-dir: dist/auto-fin diff --git a/.github/workflows/release-daily-paper.yml b/.github/workflows/release-daily-paper.yml index 7d51dc81..0d20c1f5 100644 --- a/.github/workflows/release-daily-paper.yml +++ b/.github/workflows/release-daily-paper.yml @@ -165,6 +165,6 @@ jobs: path: dist/daily-paper - name: Publish reme-daily-paper - uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # release/v1 + uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # release/v1 with: packages-dir: dist/daily-paper diff --git a/.github/workflows/release-python.yml b/.github/workflows/release-python.yml index d0769b91..5144b383 100644 --- a/.github/workflows/release-python.yml +++ b/.github/workflows/release-python.yml @@ -42,7 +42,7 @@ jobs: path: dist/reme - name: Publish ReMe - uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # release/v1 + uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # release/v1 with: packages-dir: dist/reme skip-existing: true diff --git a/.github/workflows/release-reme-studio.yml b/.github/workflows/release-reme-studio.yml index 01e8d328..fa6c539f 100644 --- a/.github/workflows/release-reme-studio.yml +++ b/.github/workflows/release-reme-studio.yml @@ -137,7 +137,7 @@ jobs: path: dist - name: Publish ReMe Studio to PyPI - uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # release/v1 + uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # release/v1 with: packages-dir: dist/studio-python skip-existing: true diff --git a/.github/workflows/security-codeql.yml b/.github/workflows/security-codeql.yml index 92c2d46d..e29819bf 100644 --- a/.github/workflows/security-codeql.yml +++ b/.github/workflows/security-codeql.yml @@ -34,12 +34,12 @@ jobs: persist-credentials: false - name: Initialize CodeQL - uses: github/codeql-action/init@6f2ac81f5cd81fe351a197ba0b00d91d03cd2aee # v4 + uses: github/codeql-action/init@9b3254085494b55a0f7c4c2b3d2d52cc712ce370 # v4 with: languages: ${{ matrix.language }} build-mode: none - name: Perform CodeQL analysis - uses: github/codeql-action/analyze@6f2ac81f5cd81fe351a197ba0b00d91d03cd2aee # v4 + uses: github/codeql-action/analyze@9b3254085494b55a0f7c4c2b3d2d52cc712ce370 # v4 with: category: /language:${{ matrix.language }} diff --git a/benchmark/memlens/README.md b/benchmark/memlens/README.md new file mode 100644 index 00000000..73dcebfa --- /dev/null +++ b/benchmark/memlens/README.md @@ -0,0 +1,107 @@ +# ReMe Image on MemLens-128K / 256K + +Evaluate ReMe Image on MemLens's official 195-question memory-agent subset. +Each length uses its own full histories; every question has an isolated workspace. + +## Run + +```bash +pip install -e ".[as]" -e plugins/memlens -e plugins/memlens-judge +export MEMLENS_DATA=/path/to/MEMLENS +export EXPERIMENT_DIR=/path/to/new-experiment +python -m benchmark.memlens.run --model qwen3.8-max --context-length 256k +``` + +Supply `LLM_API_KEY` and `LLM_BASE_URL` through the environment, outside Git. +The dataset directory contains `dataset_128k.json`, `dataset_256k.json`, +`agent_subset_195.json` and `release_images/`. Select 128K with +`--context-length 128k`; `--limit-scenes 1 --output /path/to/separate-experiment` +runs one complete history. Concurrency is configured in `config.yaml`. + +## Evaluation + +Original text, timestamps, Image IDs and interleaved image bytes go to native +Auto Memory. Supplied BLIP captions and answer annotations are not inputs. +BM25 indexes daily notes, digest notes and original session transcripts. +Each question starts a fresh 30-iteration Agent with `search`, `read` and +`read_image`; no embedding or auto-tagging is used. + +| Setting | Memory / compression | Answer | +| --- | --- | --- | +| Model | Selected candidate | Same candidate | +| Thinking / temperature / streaming | On / 0 / off | On / 0 / off | +| Output budget | 65,536 | 16,384 | +| Model context size | 200,000 | 400,000 | + +The [candidate preset](../../plugins/memlens/src/reme_memlens/benchmark.yaml) +retains 30 Agent iterations, context trigger/reserve ratios 0.8/0.1, a +50,000-token tool-result limit and 4,096-image retention. Native requests use +three retries, a 60-second delay and a 600-second HTTP timeout; nested retries +are disabled. The runner does not repair model outputs. + +The independent [MemLens Judge](../../plugins/memlens-judge/README.md) uses the +official category/subtype-specific binary accuracy protocol, Qwen3.8-Max by default, +thinking off, temperature 0 and 256 output tokens. SubEM and F1 follow the +official code. Reference answers, question categories and subtypes enter only scoring. +`JUDGE_MODEL_NAME` overrides the judge model independently of the candidate. +`settings.json` records the effective judge model, generation parameters and +protocol after plugin defaults and environment expansion, without credentials. + +## Results + +`/memlens-//` contains settings and summary JSON. +`reme_image/scenes//` retains ingestion responses, a native workspace, +answers, judgments, metrics and coverage status. The workspace's `mem_session/` +contains native Agent JSONL trajectories. Existing output groups are not overwritten. + +New runs use usage policy `successful-http-requests-v2`: usage includes observed +successful HTTP requests in memory/answer scopes, including compression and requests +completed before a later Job or history failure. Each scope is counted once. +Judge requests and unsuccessful HTTP attempts are excluded. This is measured request +usage, not a full-sample cost estimate or provider billing total. Unknown counters +remain `null`; `records` and per-counter `coverage` show how many requests were measured. +Expected, scored and unscored question counts make coverage gaps explicit. + +Technical failures are coverage gaps; valid empty or iteration-limit answers are scored. +Failure artifacts retain the failed Response before validation and record the stage, +scene/session/question IDs, sanitized error chain and traceback locations without locals. +These artifacts omit credentials, request headers, raw configuration and image data. +Failed responses use separate `diagnostics/*-response.json` files; `answer.json` continues to +indicate a successful answer Job and retains the original prediction. Per-scope +`diagnostics/*-usage.json` receipts survive direct exceptions and cancellation, while +`status.json.memory_complete` records successful ingestion and indexing independently. +Cleanup errors retain scores and usage, increment `cleanup_error_count`, and cause +a nonzero CLI exit even when all questions were scored. +Historical reports using selected successful Job/history usage retain their original +values and scope; they are not rewritten under the new policy. + +## Sources and license + +[Dataset](https://huggingface.co/datasets/xiyuRenBill/MEMLENS); +[official code](https://github.com/xrenaf/MEMLENS/tree/77f3ab9a52fa2d6a17978e2dffe80438a4ecced2). +The dataset retains its own data/image terms. Adapted query and metric code carries +the following upstream notice: + +```text +MIT License + +Copyright (c) 2026 Xiyu Ren, Zhaowei Wang, and the MemLens Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +``` diff --git a/benchmark/memlens/README_ZH.md b/benchmark/memlens/README_ZH.md new file mode 100644 index 00000000..5a384859 --- /dev/null +++ b/benchmark/memlens/README_ZH.md @@ -0,0 +1,69 @@ +# 在 MemLens-128K / 256K 上评测 ReMe Image + +使用官方固定的 195 题 memory-agent 子集。128K、256K 分别读取自己的完整历史， +每题使用独立 workspace。 + +## 运行 + +```bash +pip install -e ".[as]" -e plugins/memlens -e plugins/memlens-judge +export MEMLENS_DATA=/path/to/MEMLENS +export EXPERIMENT_DIR=/path/to/new-experiment +python -m benchmark.memlens.run --model qwen3.8-max --context-length 256k +``` + +通过环境变量提供 `LLM_API_KEY` 和 `LLM_BASE_URL`，不要提交凭据。 +数据目录包含 `dataset_128k.json`、`dataset_256k.json`、 +`agent_subset_195.json` 和 `release_images/`。使用 +`--context-length 128k` 切换长度；用 `--limit-scenes 1` 和独立的 +`--output` 路径小批量验证，不裁剪场景历史。并发配置在 `config.yaml`。 + +## 评测设置 + +原始文本、时间、Image ID 和交错的原图输入原生 Auto Memory； +不使用附带的 BLIP caption 或答案标注。BM25 索引 daily、digest 和原始 +session 正文。每题新建 30 轮回答 Agent，工具为 `search`、`read`、 +`read_image`，不启用 embedding 或 auto-tag。 + +| 设置 | 记忆／压缩 | 回答 | +| --- | --- | --- | +| 模型 | 选定的候选模型 | 同一候选模型 | +| Thinking／temperature／stream | 开启／0／关闭 | 开启／0／关闭 | +| 输出预算 | 65,536 | 16,384 | +| 模型 context size | 200,000 | 400,000 | + +[候选插件配置](../../plugins/memlens/src/reme_memlens/benchmark.yaml) 保留 +30 轮预算、0.8／0.1 上下文触发／预留比例、50,000 token 的工具结果限制和 +4,096 张图片保留上限。原生请求最多重试 3 次，间隔 60 秒，HTTP timeout +为 600 秒；不叠加 SDK／Agent 重试，也不修补模型输出。 + +独立的 [MemLens Judge](../../plugins/memlens-judge/README.md) 使用官方按 +category／subtype 区分的二元评分：默认 Qwen3.8-Max、thinking 关闭、temperature +0、输出预算 256。辅助指标为官方 SubEM、F1。答案与题型字段仅进入评分。 +`JUDGE_MODEL_NAME` 可独立覆盖 Judge 模型；`settings.json` 记录插件默认值和 +环境变量展开后实际生效的 Judge 模型、生成参数与协议，不保存凭据或服务端点。 + +## 结果 + +结果保存到 `/memlens-//`。该目录的 JSON 记录配置与汇总； +`reme_image/scenes//` 保存摄入响应、原生 workspace、回答、judge、 +指标及完成状态。workspace 下的 `mem_session/` 保留原生 Agent JSONL 轨迹。 +已有结果目录不会被覆盖。 + +新运行使用 `successful-http-requests-v2` 口径：统计记忆／回答作用域中观测到的 +成功 HTTP 请求，包括压缩，以及所属 Job 或历史处理随后失败前已成功的请求； +每个作用域只计一次。不计 Judge 请求及未成功的 HTTP 尝试。这是实测请求用量， +并非完整样本成本估计或供应商账单。未知计数保持 `null`，`records` 与各计数的 +`coverage` 标示观测覆盖；预期、已评分、未评分题数分别记录。 + +技术失败计为覆盖缺口，合法空答案与达到轮数上限的回答仍评分。失败诊断记录阶段、 +场景／session／题目 ID、脱敏错误链和不含局部变量的堆栈位置，不保存凭据、 +请求头、原始配置或图片数据。失败 Response 在检查前保存到独立的 +`diagnostics/*-response.json`；`answer.json` 仍只表示成功回答，保留原始预测。 +`diagnostics/*-usage.json` 在直接异常或取消时也保留已观测用量； +`status.json.memory_complete` 独立表示历史摄入和索引完成。 +清理失败保留评分及用量、增加 `cleanup_error_count`，并使 CLI 返回非零状态。 + +历史报告沿用其原有的选定成功 Job／历史统计口径，不按新口径重写数据。 + +数据来源、上游版本和完整授权文本见[英文 README](README.md#sources-and-license)。 diff --git a/benchmark/memlens/artifacts.py b/benchmark/memlens/artifacts.py new file mode 100644 index 00000000..18c48749 --- /dev/null +++ b/benchmark/memlens/artifacts.py @@ -0,0 +1,163 @@ +"""Small, credential-free snapshots and diagnostics for benchmark artifacts.""" + +import os +from pathlib import Path +import re +import traceback +from urllib.parse import urlsplit, urlunsplit + +_PRIVATE_KEY = re.compile( + r"credential|api[_-]?key|private[_-]?key|authorization|password|secret|cookie|access[_-]?token|refresh[_-]?token|" + r"(^|[_-])token$|headers|request[_-]?body|^environment$|^config(?:uration)?$|^messages$", + re.IGNORECASE, +) +_ASSIGNMENT = re.compile( + r"(?i)([\"']?(?:api[_-]?key|private[_-]?key|authorization|password|secret|(?:access|refresh|id)[_-]?token|token)" + r"[\"']?\s*[:=]\s*)(?:\"[^\"]*\"|'[^']*'|[^\s,;}]+)", +) +_MODEL_FIELDS = ( + "backend", + "model", + "stream", + "context_size", + "max_retries", + "retry_delay", +) +_MODEL_OPTIONS = { + "parameters": ( + "temperature", + "top_p", + "max_tokens", + "max_completion_tokens", + "thinking_enable", + ), + "extra_body": ("enable_thinking",), + "client_kwargs": ("timeout", "max_retries"), +} +_USAGE_FIELDS = { + "usage", + "coverage", + "records", + "input_tokens", + "output_tokens", + "total_tokens", + "cached_input_tokens", + "noncached_input_tokens", + "reasoning_tokens", +} + + +def _safe_text(value: str) -> str: + for key, secret in sorted(os.environ.items(), key=lambda item: len(item[1]), reverse=True): + if secret and _PRIVATE_KEY.search(key): + if len(secret) < 12: + value = re.sub(r"(?]+", safe_url, value) + value = re.sub(r"(?s)-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----", "[REDACTED_KEY]", value) + value = re.sub(r"(?im)(\b(?:set[-_])?cookie[\"']?\s*[:=]\s*)[^\r\n]+", r"\1[REDACTED]", value) + value = re.sub(r"(?i)\b(?:Bearer|Basic)\s+[A-Za-z0-9._~+/=-]+", "[REDACTED_AUTH]", value) + value = _ASSIGNMENT.sub(r"\1[REDACTED]", value) + value = re.sub(r"(?i)data:[^;,\s]+;base64,[A-Za-z0-9+/=\s]+", "[REDACTED_IMAGE]", value) + value = re.sub(r"\b[A-Za-z0-9+/]{80,}={0,2}", "[REDACTED_DATA]", value) + # Exceptions sometimes embed complete request/config dictionaries in their text. + return re.sub( + r"(?is)(\b(?:config|configuration|request_headers|headers|request_body|messages)\b[\"']?\s*[:=]).*", + r"\1[REDACTED]", + value, + ) + + +def safe_value(value, *, fixed_keys=()): + """Redact diagnostic data recursively without serializing arbitrary objects.""" + if isinstance(value, dict): + return { + (key if key in fixed_keys else _safe_text(str(key))): ( + "[REDACTED]" + if _PRIVATE_KEY.search(str(key)) or (key == "data" and value.get("type") == "base64") + else safe_value(item, fixed_keys=fixed_keys) + ) + for key, item in value.items() + } + if isinstance(value, (list, tuple)): + return [safe_value(item, fixed_keys=fixed_keys) for item in value] + if isinstance(value, str): + return _safe_text(value) + if value is None or isinstance(value, (bool, int, float)): + return value + return f"<{type(value).__name__}>" + + +def status_record(status: dict) -> dict: + """Keep runner-owned schema/enums/counters intact; errors are already sanitized.""" + return { + key: safe_value(value) if key in {"scene_id", "session_id", "question_id"} else value + for key, value in status.items() + } + + +def response_record(response) -> dict: + """Keep the Response schema while sanitizing its untrusted payload and extras.""" + raw = response.model_dump() + return { + **safe_value({key: value for key, value in raw.items() if key not in {"success", "answer", "metadata"}}), + "success": raw["success"], + "answer": safe_value(raw["answer"]), + "metadata": safe_value(raw["metadata"], fixed_keys=_USAGE_FIELDS), + } + + +def error_details(exc: BaseException, seen: set[int] | None = None) -> dict: + """Keep the error chain and frame locations, never frame locals or source text.""" + seen = set() if seen is None else seen + seen.add(id(exc)) + result = { + "type": type(exc).__name__, + "message": _safe_text(str(exc)), + "traceback": [ + { + "file": _safe_text(Path(frame.filename).name), + "line": frame.lineno, + "function": _safe_text(frame.name), + } + for frame in traceback.extract_tb(exc.__traceback__) + ], + } + cause = exc.__cause__ or (None if exc.__suppress_context__ else exc.__context__) + if cause is not None and id(cause) not in seen: + result["cause"] = error_details(cause, seen) + return result + + +def model_settings(config) -> dict: + """Whitelist reproducibility fields from a validated ApplicationConfig.""" + models = {} + for name, component in config.components.get("as_llm", {}).items(): + spec = component.model_dump() + snapshot = {key: safe_value(spec[key]) for key in _MODEL_FIELDS if key in spec} + for key, allowed in _MODEL_OPTIONS.items(): + options = spec.get(key) or {} + snapshot[key] = {name: safe_value(options[name]) for name in allowed if name in options} + models[_safe_text(name)] = snapshot + return models + + +def judge_settings(app) -> dict: + """Record the actual plugin/env-expanded judge model and official protocol.""" + from judge_memlens.protocol import PROTOCOL_ID, SOURCE_REVISION + + return { + "protocol_id": PROTOCOL_ID, + "source_revision": SOURCE_REVISION, + "models": model_settings(app.config), + } diff --git a/benchmark/memlens/config.yaml b/benchmark/memlens/config.yaml new file mode 100644 index 00000000..77d961c8 --- /dev/null +++ b/benchmark/memlens/config.yaml @@ -0,0 +1,4 @@ +dataset: ${MEMLENS_DATA} +output: ${EXPERIMENT_DIR} +scenario_concurrency: 4 +query_concurrency: 16 diff --git a/benchmark/memlens/run.py b/benchmark/memlens/run.py new file mode 100644 index 00000000..433270fc --- /dev/null +++ b/benchmark/memlens/run.py @@ -0,0 +1,621 @@ +"""Evaluate ReMe Image on the official MemLens memory-agent subset.""" + +import argparse +import base64 +from collections import Counter +from contextlib import contextmanager +from contextvars import ContextVar +from datetime import datetime +import io +import asyncio +import json +import os +from pathlib import Path +import re +import subprocess +import string +import sys + +import yaml + +from benchmark.memlens.artifacts import ( + error_details, + judge_settings, + model_settings, + response_record, + safe_value, + status_record, +) + +# Adapted query/metric code retains its upstream MIT notice in README.md. +SYSTEM_PROMPT = ( + "Provide answers based on the given conversation history. " + "If the question cannot be answered based on the given conversation, " + 'respond with "Insufficient information".' +) +INSTRUCTION = "Directly output the answer with no extra output." + + +def _image_path(root: Path, source: str) -> Path: + """Resolve only original images inside the released image directory.""" + path = (root / source).resolve() + if Path(source).is_absolute() or "\\" in source or not path.is_relative_to(root): + raise ValueError("MemLens image must stay inside release_images") + return path + + +def load_memlens(root: Path, context_length: str = "128k") -> list[dict]: + """Select the official 195 IDs in dataset order, with each question's own history.""" + root = root.expanduser().resolve() + source = root / f"dataset_{context_length}.json" + selected = json.loads((root / "agent_subset_195.json").read_text(encoding="utf-8"))["question_ids"] + if len(selected) != 195 or len(set(selected)) != 195: + raise ValueError("MemLens requires its official 195-question agent subset") + selected = set(selected) + images = (root / "release_images").resolve() + if not images.is_relative_to(root): + raise ValueError("MemLens images must stay inside the dataset directory") + result = [] + for row in json.loads(source.read_text(encoding="utf-8")): + if row["question_id"] not in selected: + continue + sessions = [] + for turns, date, session_id in zip( + row["haystack_sessions"], + row["haystack_dates"], + row["haystack_session_ids"], + strict=True, + ): + for turn in turns: + for image in turn.get("images", []): + _image_path(images, image["file"]) + sessions.append( + { + "id": session_id, + "timestamp": datetime.strptime(date, "%Y/%m/%d (%a) %H:%M").isoformat(timespec="minutes"), + "turns": turns, + "question_id": row["question_id"], + "image_root": images, + }, + ) + question = { + "id": row["question_id"], + "text": row["question"], + "question_date": row["question_date"], + "answer": row["answer"], + "category": row["question_type"], + "subtype": row.get("question_subtype", ""), + "old_answer": row.get("old_answer") or row.get("question_content", {}).get("old_answer", ""), + } + result.append({"id": row["question_id"], "sessions": sessions, "questions": [question]}) + if len(result) != 195 or {scene["id"] for scene in result} != selected: + raise ValueError("MemLens dataset must contain every selected question exactly once") + return result + + +def image_block(path: Path) -> dict: + """Encode untouched image bytes with their actual decoded MIME.""" + from PIL import Image + + if path.stat().st_size > 50 * 1024 * 1024: + raise ValueError(f"Oversized original image: {path.name}") + data = path.read_bytes() + if len(data) > 50 * 1024 * 1024: + raise ValueError(f"Oversized original image: {path.name}") + with Image.open(io.BytesIO(data)) as decoded: + if decoded.width * decoded.height > 50_000_000: + raise ValueError(f"Oversized original image dimensions: {path.name}") + media_type = decoded.get_format_mimetype() + if not media_type or not media_type.startswith("image/"): + raise ValueError(f"Missing image MIME: {path.name}") + decoded.load() + return { + "type": "data", + "source": { + "type": "base64", + "data": base64.b64encode(data).decode("ascii"), + "media_type": media_type, + }, + } + + +def session_messages(session: dict) -> list[dict]: + """Keep source roles, IDs, timestamps and the official placement.""" + messages = [] + for index, turn in enumerate(session["turns"]): + text, images = turn.get("content", ""), turn.get("images", []) + if not images: + stripped = text.replace("", "").strip() + parts = [stripped] if stripped else [] + elif "" not in text: + parts = [*images, *([text.strip()] if text.strip() else [])] + else: + parts, image_index = [], 0 + for part in re.split(r"()", text): + if part == "": + if image_index < len(images): + parts.append(images[image_index]) + image_index += 1 + elif part.strip(): + parts.append(part) + parts.extend(images[image_index:]) + content = [] + for part in parts: + if isinstance(part, str): + content.append({"type": "text", "text": part}) + else: + source = part["file"] + content.extend( + [ + {"type": "text", "text": f"[Image ID: {Path(source).stem}]"}, + image_block(_image_path(session["image_root"], source)), + ], + ) + messages.append( + { + "id": f"{session['question_id']}:{session['id']}:turn:{index}", + "name": turn["role"], + "role": turn["role"], + "content": content, + "created_at": session["timestamp"], + }, + ) + return messages + + +def query_prompt(question: dict, context: str) -> str: + """Keep the official dated query template unchanged.""" + return ( + f"{SYSTEM_PROMPT}\nConversation:\n{context}\n\n" + f"{INSTRUCTION}\nQuestion Date: {question['question_date']}\nQuestion: {question['text']}\n" + ) + + +def _normalize(text: str) -> str: + """Native SubEM/F1 normalization, preserving decimal points and date slashes.""" + text = text.lower() + text = "".join( + ( + char + if char not in string.punctuation + or (char in "./" and 0 < index < len(text) - 1 and text[index - 1].isdigit() and text[index + 1].isdigit()) + else " " + ) + for index, char in enumerate(text) + ) + text = re.sub(r"\ba(?=\s+[a-z])", " ", text) + text = re.sub(r"\b(an|the)\b", " ", text) + return " ".join(text.split()) + + +def score_answer(question: dict, prediction: str) -> dict: + """Official inline auxiliary SubEM/F1; primary accuracy comes from the LLM judge.""" + from judge_memlens.protocol import parse_model_output + + predicted, expected = _normalize(parse_model_output(prediction)), _normalize(question["answer"]) + prediction_tokens, answer_tokens = predicted.split(), expected.split() + overlap = sum((Counter(prediction_tokens) & Counter(answer_tokens)).values()) + return { + "sub_em": float(bool(predicted and expected) and expected in predicted), + "f1": 2 * overlap / (len(prediction_tokens) + len(answer_tokens)) if overlap else 0.0, + } + + +_CURRENT = ContextVar("memlens_usage", default=None) +_KEYS = ( + "input_tokens", + "output_tokens", + "total_tokens", + "cached_input_tokens", + "noncached_input_tokens", + "reasoning_tokens", +) + + +def normalized_usage(raw: dict) -> dict: + """Keep unknown counters unknown; cache/reasoning are subsets, not extra tokens.""" + incoming = raw.get("input_tokens", raw.get("prompt_tokens")) + outgoing = raw.get("output_tokens", raw.get("completion_tokens")) + total = raw.get("total_tokens") + if total is None and incoming is not None and outgoing is not None: + total = incoming + outgoing + cached = raw.get( + "cached_input_tokens", + raw.get("cached_tokens", (raw.get("prompt_tokens_details") or {}).get("cached_tokens")), + ) + reasoning = raw.get("reasoning_tokens", (raw.get("completion_tokens_details") or {}).get("reasoning_tokens")) + return { + "input_tokens": incoming, + "output_tokens": outgoing, + "total_tokens": total, + "cached_input_tokens": cached, + "noncached_input_tokens": incoming - cached if incoming is not None and cached is not None else None, + "reasoning_tokens": reasoning, + } + + +def total_usage(records: list[dict]) -> dict: + """Aggregate measured requests, recording coverage of optional counters.""" + return { + key: ( + sum(row[key] for row in records if row.get(key) is not None) + if any(row.get(key) is not None for row in records) + else None + ) + for key in _KEYS + } | { + "records": len(records), + "coverage": {key: sum(row.get(key) is not None for row in records) for key in _KEYS}, + } + + +@contextmanager +def collect_usage(target: list | None = None, receipt: Path | None = None): + """Keep concurrent calls' usage separate and include their compression calls.""" + records = [] + token = _CURRENT.set(records) + try: + yield records + finally: + _CURRENT.reset(token) + if target is not None: + target.extend(records) + if receipt is not None: + save(receipt, total_usage(records)) + + +def observe_usage(app) -> None: + """Observe owned clients and distinguish native interrupted replies from answers.""" + for wrapper in app.context.components.get("agent_wrapper", {}).values(): + original_reply = wrapper.reply + + async def reply(*args, _reply=original_reply, **kwargs): + result = await _reply(*args, **kwargs) + last = result.get("last_message") or {} + if last.get("finished_reason") == "interrupted" or last.get("error") is not None: + raise RuntimeError("Native Agent was interrupted; see the application log.") + return result + + wrapper.reply = reply + for component in app.context.components.get("as_llm", {}).values(): + client = getattr(getattr(component.model, "client", None), "_client", None) + if client is None or getattr(client, "_memlens_usage_observer", False): + continue + original = client.send + + async def send(request, *args, _send=original, **kwargs): + response = await _send(request, *args, **kwargs) + records = _CURRENT.get() + if records is not None and response.status_code == 200 and response.is_stream_consumed: + try: + payload = response.json() + except (ValueError, UnicodeError): + payload = {} + records.append(normalized_usage(payload.get("usage") or {})) + return response + + client.send = send + client._memlens_usage_observer = True # pylint: disable=protected-access + + +def save(path: Path, value) -> None: + """Save source/results unchanged; callers sanitize diagnostics and config snapshots.""" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, ensure_ascii=False, indent=2, default=str) + "\n", encoding="utf-8") + + +def app_config(workspace: Path, model: str, *, judge: bool = False) -> dict: + """Resolve standard plugin settings through ReMe's configuration parser.""" + from reme.config import resolve_app_config + + values = resolve_app_config( + config="benchmark" if judge else "memlens-benchmark", + plugins=["memlens-judge" if judge else "memlens"], + workspace_dir=str(workspace), + log_config=False, + enable_logo=False, + log_to_console=False, + log_to_file=True, + ) + # BM25-only: do not start unused embedding clients. + values["components"].pop("as_embedding", None) + values["components"].pop("embedding_store", None) + values["components"]["file_store"]["default"]["embedding_store"] = None + if not judge: + for name in ("default", "bench", "compressor"): + values["components"]["as_llm"][name]["model"] = model + else: + values["jobs"] = {"answer_judge": {}} + values["components"] = {} + return values + + +def check_response(response, diagnostic_path: Path) -> None: + """Save failed responses before validation; keep execution failures unscored.""" + if not response.success: + save(diagnostic_path, {"response": response_record(response)}) + if asyncio.current_task().cancelling(): + raise asyncio.CancelledError + if not response.success: + detail = response.answer or response.metadata.get("error") or "No failure detail provided." + raise RuntimeError(f"ReMe Job failed; response detail: {safe_value(detail)}") + + +async def close_application(app) -> None: + """Finish owned cleanup even when cancellation arrives during close.""" + closing = asyncio.create_task(app.close()) + cancellation = None + while True: + try: + await asyncio.shield(closing) + break + except asyncio.CancelledError as exc: + if closing.done(): + raise + cancellation = exc + except Exception as exc: + if cancellation is not None: + raise cancellation from exc + raise + if cancellation is not None: + raise cancellation + + +async def evaluate_scene(scene, number, model, group, judge, slots): + """Accumulate one isolated history, then answer and judge its question.""" + from reme import Application + + label = re.sub(r"[^\w.-]+", "_", scene["id"])[:100] + case = group / "reme_image" / "scenes" / f"{number:04d}-{label}" + app = None + memory, answers, results = [], [], [] + status = {"complete": False, "memory_complete": False, "stage": "initialize", "scene_id": scene["id"]} + try: + save(case / "source.json", {"id": scene["id"], "sessions": len(scene["sessions"])}) + app = Application(**app_config(case / "ingestion/attempt-001/workspace", model)) + status["stage"] = "start" + await app.start() + observe_usage(app) + for ordinal, session in enumerate(scene["sessions"]): + directory = case / "ingestion/attempt-001" + status.update(stage="preprocess", session_id=session["id"]) + async with slots: + with collect_usage(memory, directory / "diagnostics" / f"{ordinal:05d}-usage.json") as requests: + messages = await asyncio.to_thread(session_messages, session) + status["stage"] = "auto_memory" + response = await app.run_job( + "auto_memory", + messages=messages, + session_id=session["id"], + date=session["timestamp"][:10], + include_images=True, + include_source_sessions=True, + ) + check_response(response, directory / "diagnostics" / f"{ordinal:05d}-response.json") + save( + directory / "sessions" / f"{ordinal:05d}.json", + { + "session_id": session["id"], + "response": response.model_dump(), + "usage": total_usage(requests), + }, + ) + status.pop("session_id", None) + status["stage"] = "index_update" + response = await app.run_job("index_update") + check_response(response, case / "ingestion/attempt-001/diagnostics/index-response.json") + status["memory_complete"] = True + for index, question in enumerate(scene["questions"]): + directory = case / "questions" / f"{index:04d}" + status.update(stage="agentic_answer", question_id=question["id"]) + async with slots: + with collect_usage(answers, directory / "diagnostics/answer-usage.json") as requests: + response = await app.run_job( + "agentic_answer", + query=query_prompt(question, "Use the available tools to retrieve relevant evidence."), + query_time=question["question_date"], + query_images=[], + scene_id=scene["id"], + include_images=True, + include_source_sessions=True, + max_iters=30, + ) + check_response(response, directory / "diagnostics/answer-response.json") + save( + directory / "answer.json", + { + "question_id": question["id"], + "prediction": response.answer, + "usage": total_usage(requests), + "agent": response.metadata, + }, + ) + status["stage"] = "judge" + async with slots: + judgment = await judge.run_job( + "answer_judge", + question=question["text"], + reference=question["answer"], + prediction=response.answer, + category=question["category"], + subtype=question.get("subtype", ""), + old_answer=question.get("old_answer", ""), + ) + check_response(judgment, directory / "diagnostics/judge-response.json") + status["stage"] = "score" + metrics = score_answer(question, response.answer) + metrics["llm_judge"] = judgment.metadata["llm_judge_score"] + save(directory / "judge.json", judgment.metadata) + record = { + "question_id": question["id"], + "category": question["category"], + "question": question["text"], + "reference": question["answer"], + "prediction": response.answer, + "metrics": metrics, + } + save(directory / "result.json", record) + results.append(record) + status.update(complete=True, stage="complete") + except asyncio.CancelledError as exc: + status["error"] = error_details(exc) + raise + except Exception as exc: # pylint: disable=broad-exception-caught + status["error"] = error_details(exc) + finally: + try: + if app is not None: + await close_application(app) + except (Exception, asyncio.CancelledError) as exc: # pylint: disable=broad-exception-caught + if "error" in status: + status["cleanup_error"] = error_details(exc) + else: + status.update(complete=False, stage="cleanup", error=error_details(exc)) + if isinstance(exc, asyncio.CancelledError): + raise + finally: + status.update(memory_usage=total_usage(memory), answer_usage=total_usage(answers)) + if "error" in status: + status["error_type"] = status["error"]["type"] + save(case / "status.json", status_record(status)) + return results, memory, answers, status + + +async def run(settings: dict, model: str, length: str, limit: int | None = None) -> dict: + """Run one model in a new output group; existing runs are never overwritten.""" + from reme import Application + from reme.schema import ApplicationConfig + + group = Path(settings["output"]).expanduser().resolve() / f"memlens-{length}" / model + group.mkdir(parents=True, exist_ok=False) + judge = None + status = {"complete": False, "stage": "preprocess"} + try: + scenes = load_memlens(Path(settings["dataset"]), length) + if limit is not None: + scenes = scenes[:limit] + status["stage"] = "configure" + effective = ApplicationConfig(**app_config(group / "unused-workspace", model)) + judge = Application(**app_config(group / "judge-workspace", model, judge=True)) + save( + group / "settings.json", + { + "method": "reme_image", + "model": model, + "context_length": length, + "dataset": str(Path(settings["dataset"]).expanduser().resolve()), + "limit_scenes": limit, + "include_source_sessions": True, + "include_images": True, + "answer_max_iters": 30, + "concurrency": {key: settings[key] for key in ("scenario_concurrency", "query_concurrency")}, + "models": model_settings(effective), + "judge": judge_settings(judge), + "usage_policy_version": "successful-http-requests-v2", + "reme_commit": subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + }, + ) + slots = asyncio.Semaphore(settings["query_concurrency"]) + scene_slots = asyncio.Semaphore(settings["scenario_concurrency"]) + status["stage"] = "judge_start" + await judge.start() + status["stage"] = "evaluate" + + async def evaluate(number, scene): + async with scene_slots: + print(f"Starting {length} {model} {scene['id']}", flush=True) + result = await evaluate_scene(scene, number, model, group, judge, slots) + print(f"Finished {scene['id']}: {len(result[0])}/{len(scene['questions'])} scored", flush=True) + return result + + tasks = [asyncio.create_task(evaluate(index, scene)) for index, scene in enumerate(scenes)] + try: + outputs = await asyncio.gather(*tasks) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + status["stage"] = "summary" + results = [row for output in outputs for row in output[0]] + memory = [row for output in outputs for row in output[1]] + answers = [row for output in outputs for row in output[2]] + expected = sum(len(scene["questions"]) for scene in scenes) + statuses = [output[3] for output in outputs] + cleanup_errors = sum("cleanup_error" in row or row["stage"] == "cleanup" for row in statuses) + summary = { + "method": "reme_image", + "model": model, + "context_length": length, + "expected_scenes": len(scenes), + "completed_scenes": sum(row["complete"] for row in statuses), + "memory_complete_scenes": sum(row["memory_complete"] for row in statuses), + "cleanup_error_count": cleanup_errors, + "expected_questions": expected, + "scored_questions": len(results), + "unscored_questions": expected - len(results), + "metrics": { + key: (sum(row["metrics"][key] for row in results) / len(results) if results else None) + for key in ("llm_judge", "sub_em", "f1") + }, + "memory_usage": total_usage(memory), + "answer_usage": total_usage(answers), + "token_usage": total_usage(memory + answers), + "usage_policy_version": "successful-http-requests-v2", + "usage_policy": ( + "observed successful HTTP requests in memory/answer scopes, including compression and requests " + "before Job/history failure; judge excluded; not full-sample or billed cost" + ), + } + save(group / "summary.json", summary) + status.update(complete=len(results) == expected and not cleanup_errors, stage="complete") + except (Exception, asyncio.CancelledError) as exc: + status["error"] = error_details(exc) + raise + finally: + try: + if judge is not None: + await close_application(judge) + except (Exception, asyncio.CancelledError) as exc: + if "error" in status: + status["cleanup_error"] = error_details(exc) + if isinstance(exc, asyncio.CancelledError): + raise + else: + status.update(complete=False, stage="judge_cleanup", error=error_details(exc)) + raise + finally: + save(group / "status.json", status_record(status)) + return summary + + +def main(): + """Run one model at the requested MemLens history length.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, default=Path(__file__).with_name("config.yaml")) + parser.add_argument("--model", default="qwen3.8-max") + parser.add_argument("--context-length", choices=("256k", "128k"), default="256k") + parser.add_argument("--limit-scenes", type=int) + parser.add_argument("--output", type=Path) + arguments = parser.parse_args() + try: + settings = yaml.safe_load(os.path.expandvars(arguments.config.read_text(encoding="utf-8"))) + if arguments.output: + settings["output"] = str(arguments.output) + if arguments.limit_scenes is not None and arguments.limit_scenes < 1: + parser.error("--limit-scenes must be positive") + for key in ("query_concurrency", "scenario_concurrency"): + if not isinstance(settings[key], int) or settings[key] < 1: + parser.error(f"{key} must be positive") + summary = asyncio.run(run(settings, arguments.model, arguments.context_length, arguments.limit_scenes)) + except (Exception, asyncio.CancelledError, KeyboardInterrupt) as exc: # pylint: disable=broad-exception-caught + print(json.dumps({"complete": False, "error": error_details(exc)}, ensure_ascii=False), file=sys.stderr) + return 130 if isinstance(exc, (asyncio.CancelledError, KeyboardInterrupt)) else 2 + print(json.dumps(summary, ensure_ascii=False, indent=2)) + return ( + 0 if summary["scored_questions"] == summary["expected_questions"] and not summary["cleanup_error_count"] else 2 + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/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/docs/en/context_compaction.md b/docs/en/context_compaction.md index c6677668..48e705be 100644 --- a/docs/en/context_compaction.md +++ b/docs/en/context_compaction.md @@ -176,3 +176,11 @@ keep the instance alive until ReMe has finished its calls, then close it itself. replacing an internally constructed model or any model with active requests. Shutdown waits for active calls and owned-client cleanup even when cancelled, then propagates cancellation. Component shutdown detaches the instance; inject again before restarting a borrowed-model component. Standalone credential-based behavior is unchanged. + +Owned clients initialize in a worker thread so certificate loading does not block the event loop. Cancelling a request +during construction waits for initialization to finish and leaves the client owned by the component until shutdown; +the cancelled request is not sent. + +Hosts can import `CompactContextStep`, `ToolPair`, `collect_pairs`, and `preview_result` from `reme.steps.compaction`. +The pairing and preview helpers share the Step's rules for validating proposals. Archiving, resolving references, +and committing proposals remain host responsibilities. diff --git a/docs/en/quick_start.md b/docs/en/quick_start.md index 19784e8d..c338f557 100644 --- a/docs/en/quick_start.md +++ b/docs/en/quick_start.md @@ -179,11 +179,19 @@ Daily notes: ```bash reme write path=daily/2026-06-20/demo-session.md name=demo-session description="Demo session" content="Recorded content" reme daily_list +reme daily_list start_date=2026-06-01 end_date=2026-06-20 page=1 page_size=20 reme daily_reindex ``` `write` can create a daily note directly. Run `daily_reindex` when the day's index needs to be refreshed. +`daily_list` is read-only. Without new parameters, it still lists all notes for `date` (default: today) and returns +`{date, count, notes}`. Both range dates are required, inclusive, and cannot be combined with a non-empty `date`. +Ranges sort by date descending, then path ascending. Ranges or an explicit `page`/`page_size` enable pagination: +page defaults to 1, page size to 100 (maximum 1000), and metadata adds `page`, `page_size`, `total_notes`, `total_pages`. +An empty range has zero pages; a page past the end returns no notes. Totals count readable, valid, authorized notes. +Pagination limits output, not scanning; each request reflects the current files rather than a cross-page snapshot. + --- ## Automatic Memory diff --git a/docs/zh/context_compaction.md b/docs/zh/context_compaction.md index 33accb2e..78328411 100644 --- a/docs/zh/context_compaction.md +++ b/docs/zh/context_compaction.md @@ -162,3 +162,9 @@ YAML 中的 `model` 仍表示供应商模型名称；运行时 `model` 属性接 宿主应等待 ReMe 调用结束后自行关闭实例。替换 ReMe 自建模型或存在进行中请求的实例之前，须先关闭组件。 关闭过程即使被取消，也会等待活动请求和自建客户端清理完成，再传播取消信号。组件关闭时会解除实例引用； 重新启动使用外部模型的组件前需要再次注入。独立运行时按凭据创建模型的行为保持不变。 + +自建客户端在线程中初始化，证书加载不会阻塞事件循环。构造期间取消请求时，组件会等待构造完成并保留客户端所有权， +由后续组件关闭负责释放；此时不会发送该请求。 + +宿主可从 `reme.steps.compaction` 导入 `CompactContextStep`、`ToolPair`、`collect_pairs` 和 `preview_result`。 +配对识别与预览构造函数和压缩 Step 共享相同规则，供宿主校验提案使用；归档、引用解析和提交仍由宿主负责。 diff --git a/docs/zh/quick_start.md b/docs/zh/quick_start.md index 2df0a47f..c2e04bf6 100644 --- a/docs/zh/quick_start.md +++ b/docs/zh/quick_start.md @@ -173,11 +173,18 @@ Daily note： ```bash reme write path=daily/2026-06-20/demo-session.md name=demo-session description="Demo session" content="记录内容" reme daily_list +reme daily_list start_date=2026-06-01 end_date=2026-06-20 page=1 page_size=20 reme daily_reindex ``` `write` 可直接创建 daily note；需要刷新当天索引时运行 `daily_reindex`。 +`daily_list` 只读。未传新参数时，仍列出 `date`（默认今天）的全部笔记，返回 `{date, count, notes}`。 +日期范围必须同时提供起止日期，包含两端，且不能与非空的 `date` 混用；按日期降序、同日路径升序排列。 +日期范围或显式的 `page` / `page_size` 会启用分页：页码默认 1，每页默认 100 条、最多 1000 条， +metadata 增加 `page`、`page_size`、`total_notes`、`total_pages`。空范围页数为 0，越界页返回空列表； +总数只统计可读取、有效且有权限访问的笔记。分页限制返回量，不限制扫描量；每次查询读取当前文件，不提供跨页快照。 + --- ## 自动记忆 diff --git a/plugins/memlens-judge/LICENSE-CODE b/plugins/memlens-judge/LICENSE-CODE new file mode 100644 index 00000000..a6bed67f --- /dev/null +++ b/plugins/memlens-judge/LICENSE-CODE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Xiyu Ren, Zhaowei Wang, and the MemLens Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/plugins/memlens-judge/README.md b/plugins/memlens-judge/README.md new file mode 100644 index 00000000..da0b4e0b --- /dev/null +++ b/plugins/memlens-judge/README.md @@ -0,0 +1,30 @@ +# ReMe MemLens Judge + +Independent `answer_judge` Job for saved MemLens predictions. Install with +`pip install -e plugins/memlens-judge` and enable the `memlens-judge` plugin +in a separate evaluation application/workspace, not the candidate memory application. + +Inputs are `question`, `reference`, `prediction`, and `category`, with optional +`subtype` and `old_answer`. The result is available as +`response.metadata.llm_judge_score` and `acc`. Metadata includes the raw response, +task rubric, native answer-preprocessing diagnostics and judge token usage. +Judge usage is evaluation cost, not method cost. + +The official category/subtype rubrics, outdated-answer handling, verdict parser, +long-answer preprocessing and over-500-word zero rule are preserved. A successful +but unparseable judge output receives the official zero score; API errors instead +remain failed jobs for resumable judging of the same saved prediction. + +The default judge is `qwen3.8-max`, with temperature 0 and thinking disabled. +Set `JUDGE_MODEL_NAME` to override the judge model independently of the candidate. +The benchmark's `settings.json` records the effective judge model, generation +parameters and protocol from the validated application configuration, after plugin +defaults and environment expansion. The snapshot excludes credentials and endpoints. +The output +budget is 256 tokens from the released OpenAI-compatible API judge branch +(`max_completion_tokens`), not the separate vLLM branch's 2,048-token setting. +These judge settings are independent of the candidate reader's thinking mode. + +Source: [official MemLens judge](https://github.com/xrenaf/MEMLENS/blob/77f3ab9a52fa2d6a17978e2dffe80438a4ecced2/llm_judge.py). +Verbatim prompt strings and their MIT notice are packaged in `_memlens_judge_prompts.json`; +see also `LICENSE-CODE`. The plugin is independent of candidate code and other benchmarks. diff --git a/plugins/memlens-judge/pyproject.toml b/plugins/memlens-judge/pyproject.toml new file mode 100644 index 00000000..8782eb5b --- /dev/null +++ b/plugins/memlens-judge/pyproject.toml @@ -0,0 +1,22 @@ +[project] +name = "reme-memlens-judge" +version = "0.1.0" +description = "Official memlens judge protocol as a ReMe plugin." +readme = "README.md" +license = "Apache-2.0" +requires-python = ">=3.11" +dependencies = ["reme-ai[as]>=0.4.1.11"] + +[project.entry-points."reme.plugins"] +memlens-judge = "judge_memlens" + +[tool.setuptools.packages.find] +where = ["src"] +include = ["judge_memlens*"] + +[tool.setuptools.package-data] +judge_memlens = ["plugin.yaml", "*.json"] + +[build-system] +requires = ["setuptools>=77", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/plugins/memlens-judge/src/judge_memlens/__init__.py b/plugins/memlens-judge/src/judge_memlens/__init__.py new file mode 100644 index 00000000..39c8fd4a --- /dev/null +++ b/plugins/memlens-judge/src/judge_memlens/__init__.py @@ -0,0 +1,5 @@ +"""Independent memlens evaluation-only plugin.""" + +from .llm_judge import MemLensJudgeStep + +__all__ = ["MemLensJudgeStep"] diff --git a/plugins/memlens-judge/src/judge_memlens/_memlens_judge_prompts.json b/plugins/memlens-judge/src/judge_memlens/_memlens_judge_prompts.json new file mode 100644 index 00000000..bf34840b --- /dev/null +++ b/plugins/memlens-judge/src/judge_memlens/_memlens_judge_prompts.json @@ -0,0 +1,29 @@ +{ + "source_revision": "77f3ab9a52fa2d6a17978e2dffe80438a4ecced2", + "license": "MIT License\n\nCopyright (c) 2026 Xiyu Ren, Zhaowei Wang, and the MemLens Authors\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n", + "template": "Now your role is a grading teacher. Your task is to review and score student answers based on reference standard answers for a question-answering benchmark. You need to notice the following key points:\n- First, extract the final answer from the student's solution, then analyze and judge whether the answer is correct.\n- Scoring should only refer to the final answer obtained by the student; there is no need to examine whether the intermediate problem-solving steps are correct.\n- If the response contains both hesitation and a clear answer, judge the answer itself.\n- If the response contains both a refusal and a guessed answer, judge the final committed answer.\n- If the response gives multiple inconsistent answers, assign 0 points.\n- If the student's response shows circular reasoning — repeatedly revisiting the same evidence, flip-flopping between answers (e.g., \"yes... wait no... actually yes... let me reconsider... no\"), with no clear final commitment — assign 0 points. A correct answer mentioned during reasoning does NOT count unless it is the clearly stated final conclusion.\n- Only the student's LAST clearly stated answer counts. If the student initially says \"yes\" but later concludes \"no\" (or vice versa), score based on the LAST position only. Intermediate answers during reasoning are not final answers.\n- If the student's output is entirely a reasoning trace (e.g., \"let me check... scanning evidence... re-reading session 5...\") with no clear concluding answer statement, assign 0 points — even if the correct answer appears somewhere within the reasoning.\n- If the student's output begins with \"[...truncated earlier reasoning...]\", this means only the final portion of a longer response is shown. Focus on extracting the answer from this final portion.\n\nBelow are examples of problematic student outputs that apply to ALL task types. Study these before grading.\n\n[Universal Example A — Circular Reasoning -> 0 points]\n: Is the cat in Mark and Jenny's 5-year anniversary post named Mittens?\n: Yes\n: let me check session 5 mark has cat named mittens session 7 mark and jenny anniversary post there is no link wait let me re-read session 5 says mittens session 7 says anniversary actually looking at this again wait is it possible let me reconsider mark has mittens in session 5 the post is about couple therefore answer is likely no wait let me double check session 5 mark has cat mittens actually maybe yes but text doesnt say\n[Scoring Rationale]: The output flip-flops between \"yes\" and \"no\" multiple times with no clear final commitment. The last position is ambiguous (\"maybe yes but text doesnt say\"). This is circular reasoning.\nIn summary, the student's answer deserves 0 points.\n[JSON]: {{\"answer_score\": 0}}\n\n[Universal Example B — Redundant But Committed -> Score Normally]\n: Is the trailing plant in the cafe safe for my kitten?\n: No\n: the plant looks like pothos pothos is toxic to cats therefore answer is no let me double check is pothos safe for cats no is answer no yes wait let me verify image 3 is pothos image 9 is pothos pothos is toxic therefore answer must be no\n[Scoring Rationale]: Despite excessive self-verification, the student consistently commits to \"no\" throughout and never wavers to a different answer. The final answer is clearly \"No\", which matches the standard answer.\nIn summary, the student's answer deserves 1 point.\n[JSON]: {{\"answer_score\": 1}}\n\n[Universal Example C — Reasoning Trace With No Answer -> 0 points]\n: How many sessions mentioned the coffee shop?\n: 3\n: scanning session 1 yes coffee mentioned session 2 no session 3 yes coffee again session 4 no session 5 maybe let me re-read session 5 it mentions cafe is that same as coffee shop need to check session 6 no session 7 unclear let me look at session 5 again\n[Scoring Rationale]: The student walks through sessions but never states a final count. The output is entirely a reasoning trace with no concluding answer.\nIn summary, the student's answer deserves 0 points.\n[JSON]: {{\"answer_score\": 0}}\n\nNow proceed with grading. Remember: the universal examples above apply regardless of task type.\n\n- When analyzing and judging whether the answer is correct, you need to write down the scoring rationale, organize it into clear statements that follow the logical flow. The summary of the scoring rationale should be placed at the end, using the following format: \"In summary, the student's answer deserves x points\" (where x represents the student's specific score).\n- Keep the whole process concise, within 150 words.\n- Provide the score based on your analysis and display it in a code block in \"JSON\" format.\n- An item is covered if it is strictly mentioned or unambiguously implied by a semantic equivalence. This includes numerical equivalence (e.g., 10% and 0.1), synonyms (e.g., UK and United Kingdom), plural/singular forms (e.g., \"apple\" and \"apples\"), and equivalent date formats (e.g., 2024-01-15 and January 15, 2024). However, do not accept loosely related concepts.\n- Ignore minor formatting differences, capitalization, punctuation, and equivalent wording when meaning is unchanged.\nYour output format is:\n[Scoring Rationale]:\n[Score]: x points\n[JSON]:\n{{\"answer_score\": }}\n\nBelow is the grading rubric:\n[Scores]:\nThe scoring scale consists of 2 levels in total, from highest to lowest: 1 point, 0 points (the minimum is 0 points).\n[Tier Details]:\n1 point: Assign 1 point if the student's final answer matches the standard answer under the task-specific criteria below.\n0 points: Assign 0 points if the student's final answer does not match the standard answer, the student refuses to answer, claims insufficient information (except for AR tasks), or does not clearly answer.\n\n[Task-Specific Criteria]:\n{task_criteria}\n\n{examples}\n", + "criteria": { + "IE_Entity": "[IE — Entity / Attribute Extraction]\nThis is an information extraction question about a concrete entity, count, attribute, text, object, location, or visual detail from the provided memory context.\n- Assign 1 point if the student's response contains the core information from the standard answer. Minor wording differences are acceptable, but the essential information must be present and correct.\n- Assign 0 points if the core information is missing, contradicted, too vague, refused, or incorrect.\n- For numeric or count answers, require the exact value unless the standard answer itself is approximate.", + "IE_PreviousInfo": "[IE — Previous Information Extraction]\nThis is an information extraction question about previously mentioned information, spatial relations, counts, or attributes from earlier conversation/image context.\n- Assign 1 point if the student's response recovers the same previous information as the standard answer, including the correct relation, location, count, or attribute.\n- Minor wording differences are acceptable if they preserve the same meaning.\n- Assign 0 points if the answer refers to the wrong previous item, gives the wrong relation/count/location, is too vague, refuses, or is unsupported.", + "KU_KnowledgeUpdate": "[KU — Knowledge Update]\nThis is a knowledge update question testing whether the model tracks the latest state.\nThe correct current answer and the outdated old answer are both provided.\n- Assign 1 point if the student gives the current (correct) answer.\n- Assign 0 points if the student gives the outdated old answer, any other incorrect answer, refuses, or does not clearly answer.", + "TR_DurationComparison": "[TR — Duration Comparison (A/B)]\nThis is a binary choice question. The student must select one of two options.\n- Assign 1 point if the student clearly selects the same option as the standard answer, regardless of phrasing.\n- Assign 0 points if the student selects the other option, refuses, or does not make a clear selection.", + "TR_OrderRanking": "[TR — Order Ranking]\nThis is an ordering question. The standard answer is an exact sequence of items.\n- Assign 1 point ONLY if the student gives the exact same sequence in the same order.\n- Assign 0 points if any item is missing, added, duplicated, misplaced, or if the order is unclear.", + "TR_DateExtraction": "[TR — Date Extraction]\nThis is a date extraction question. The standard answer is a specific date.\n- Assign 1 point if the student gives the same date in any unambiguous equivalent format (e.g., \"Jan 15, 2024\" and \"2024-01-15\" and \"15/01/2024\").\n- Assign 0 points if the date is different, incomplete when a full date is required, ambiguous, or missing.", + "MSR_YesNo": "[MSR — Yes/No (Entity Resolution)]\nThis is a yes/no question.\n- Assign 1 point if the student's final answer is semantically equivalent to the standard answer (e.g., \"Yes\" matches \"Yes\", \"No\" matches \"No\"), even if phrased differently (e.g., \"That's correct\" for \"Yes\").\n- Assign 0 points if the answer means the opposite, refuses, or does not clearly answer.", + "MSR_Counting": "[MSR — Counting]\nThis is a counting question. The standard answer is a specific number.\n- Assign 1 point if the student clearly states the exact number in any equivalent form (e.g., \"3\", \"three\", \"3.0\").\n- Assign 0 points if the student gives a different number, a range, an approximation, refuses, or does not clearly answer.", + "MSR_Arithmetic": "[MSR — Arithmetic]\nThis is an arithmetic question. The standard answer is a specific value.\n- Assign 1 point if the student gives the same value in an equivalent format (e.g., \"0.5\" and \"1/2\" and \"50%\").\n- Assign 0 points if the student gives a different value, an unsupported approximation, refuses, or does not clearly answer.", + "AR_AnswerRefusal": "[AR — Answer Refusal]\n- Assign 1 point if the student explicitly refuses, says there is not enough information, or indicates the question is unanswerable from the provided context.\n- Assign 0 points if the student gives a substantive answer, even if hedged." + }, + "examples": { + "IE_Entity": "\n[Example 1]\n: What specific word is painted in large white letters on the pavement in the middle lanes?\n: SCHOOL\n: SCHOOL\n\n[Scoring Rationale]: The student's answer is \"SCHOOL\", which exactly matches the standard answer.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: Where on the sign is the photo of the mountains located?\n: In the bottom right corner\n: The photo of the mountains is located in the upper right corner of the sign.\n\n[Scoring Rationale]: The student says \"upper right corner\" but the standard answer is \"bottom right corner\". The location is incorrect — \"upper\" vs \"bottom\" is a substantive difference.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "IE_PreviousInfo": "\n[Example 1]\n: Where on the sign is the photo of the mountains located?\n: In the bottom right corner\n: The photo is in the bottom-right corner of the sign.\n\n[Scoring Rationale]: The student's answer preserves the same location relation as the standard answer.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: Relative to the baking sheet, where is the salmon fillet placed?\n: The right side\n: It is on the left side of the baking sheet.\n\n[Scoring Rationale]: The student's answer gives the opposite spatial relation from the standard answer.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "KU_KnowledgeUpdate": "\n[Example 1]\n: What do I treat as my must-have now that I plan my meals around it?\n: brioche\n: sourdough\n: You now consider brioche your must-have bread that you plan meals around.\n\n[Scoring Rationale]: The student's answer is \"brioche\", which matches the current standard answer, not the outdated answer \"sourdough\".\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: What do I treat as my must-have now that I plan my meals around it?\n: brioche\n: sourdough\n: A versatile, high-quality protein source like chicken, beans, or tofu is a great starting point for building a week's worth of meals.\n\n[Scoring Rationale]: The student's answer mentions generic protein sources, which is neither the current answer \"brioche\" nor the outdated answer \"sourdough\". The answer is simply incorrect.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "TR_DurationComparison": "\n[Example 1]\n: Which of the following two durations is longer? Duration 1: The time I spent on Miller Hall. Duration 2: The time I spent on Music Festival Job. A. Duration 1 is longer. B. Duration 2 is longer.\n: B\n: B. Duration 2 is longer.\n\n[Scoring Rationale]: The student clearly selects \"B\", which matches the standard answer.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: Which of the following two durations is longer? Duration 1: The time I spent with the NutriPeak Nutrition. Duration 2: The time I spent with the BlendIt Pro Blender. A. Duration 1 is longer. B. Duration 2 is longer.\n: A\n: B. Duration 2 is longer.\n\n[Scoring Rationale]: The student selects \"B\", but the standard answer is \"A\". The selection is opposite.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "TR_OrderRanking": "\n[Example 1]\n: Rank the following cities by population from largest to smallest: Tokyo, Delhi, Shanghai.\n: Tokyo, Delhi, Shanghai\n: 1. Tokyo 2. Delhi 3. Shanghai\n\n[Scoring Rationale]: The student's sequence is Tokyo, Delhi, Shanghai, which exactly matches the standard answer.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: Rank events A, B, C in chronological order.\n: B, A, C\n: A, B, C\n\n[Scoring Rationale]: The student's sequence is A, B, C, but the standard answer is B, A, C. The order does not match.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "TR_DateExtraction": "\n[Example 1]\n: When did I complete the Scandinavian design course at Copenhagen Design Academy?\n: 2024/08/31\n: You completed the Scandinavian design course on August 31st, 2024.\n\n[Scoring Rationale]: The student's answer \"August 31st, 2024\" is an unambiguous equivalent of the standard answer \"2024/08/31\".\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: When did I visit Marutama Ra-men to taste their chicken broth?\n: 2023/01/20\n: January 20, 2021\n\n[Scoring Rationale]: The student's answer is \"January 20, 2021\", but the standard answer is \"2023/01/20\" — the year is wrong (2021 vs 2023).\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "MSR_YesNo": "\n[Example 1]\n: Is the cat in Mark and Jenny's 5-year anniversary post named Mittens?\n: Yes\n: Yes, the cat in Mark and Jenny's 5-year anniversary post is named Mittens.\n\n[Scoring Rationale]: The student's final answer is \"Yes\", which matches the standard answer \"Yes\". The additional explanation does not contradict the answer.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: Did the cat named after a crypto creator knock over the phone displaying the Binance alert?\n: Yes\n: The text does not explicitly state whether the cat knocked over the phone. It only mentions the cat was sitting next to the phone when it crashed.\n\n[Scoring Rationale]: The student does not give a clear \"Yes\". Instead, the student hedges and effectively refuses to commit to an answer. The standard answer is \"Yes\".\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "MSR_Counting": "", + "MSR_Arithmetic": "\n[Example 1]\n: How much total have I spent on coffee makers?\n: $260.00\n: Based on the conversation, you spent $110 on the first coffee maker and $150 on the second one, totaling $260.\n\n[Scoring Rationale]: The student's final answer is $260, which is numerically equivalent to the standard answer $260.00.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: How much total have I spent on coffee makers?\n: $260.00\n: I do not have access to your personal financial data, including your coffee maker purchases. I cannot provide you with a total amount spent on coffee makers.\n\n[Scoring Rationale]: The student refuses to answer and claims insufficient information, but this is not an AR task. The information was available in the conversation context.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n", + "AR_AnswerRefusal": "\n[Example 1]\n: What is the model number printed on the power bank?\n: Insufficient information\n: I am sorry, but I cannot fulfill your request. The provided conversation history does not contain any information about a power bank or its model number.\n\n[Scoring Rationale]: The student correctly refuses to answer, stating the information is not available in the conversation. This is the expected behavior for an unanswerable question.\nIn summary, the student's answer deserves 1 point.\n[Score]: 1 point\n[JSON]:\n{\n \"answer_score\": 1\n}\n\n[Example 2]\n: How many vertical cables are visible on one side of the bridge?\n: Insufficient information\n: Three\n\n[Scoring Rationale]: The student provides a substantive answer (\"Three\") rather than refusing. For an unanswerable question, the correct response is to refuse or indicate insufficient information.\nIn summary, the student's answer deserves 0 points.\n[Score]: 0 points\n[JSON]:\n{\n \"answer_score\": 0\n}\n" + } +} diff --git a/plugins/memlens-judge/src/judge_memlens/llm_judge.py b/plugins/memlens-judge/src/judge_memlens/llm_judge.py new file mode 100644 index 00000000..998e1685 --- /dev/null +++ b/plugins/memlens-judge/src/judge_memlens/llm_judge.py @@ -0,0 +1,72 @@ +"""Official MemLens task-specific judge as an isolated ReMe Job.""" + +from dataclasses import asdict + +from agentscope.message import UserMsg + +from reme.steps.base_step import BaseStep +from reme.steps.evolve.compressor import CompressorStep + +from .protocol import ( + SOURCE_REVISION, + judge_prompt, + normalize_for_judge, + parse_judge_response, + task_key, +) + + +class MemLensJudgeStep(BaseStep): + """Score saved answers without accessing the candidate's workspace or tools.""" + + async def execute(self): + """Apply the official judge prompt/parser to one saved prediction.""" + question = self.context["question"] + reference = self.context["reference"] + prediction = self.context["prediction"] + category = self.context["category"] + subtype = self.context.get("subtype", "") + normalized = normalize_for_judge(prediction) + prompt = judge_prompt( + question, + reference, + normalized["text"], + category, + subtype, + self.context.get("old_answer", ""), + ) + words = len(normalized["text"].split()) + result = { + "source_revision": SOURCE_REVISION, + "task_key": task_key(category, subtype, reference), + "original_len": normalized["original_len"], + "was_truncated": normalized["was_truncated"], + "is_degenerate": normalized["is_degenerate"], + **normalized["diagnostics"], + "parsed_words": words, + "auto_zero": words > 500, + } + raw, usage = "", None + if words > 500: + parsed = {"acc": 0.0, "parse_default_zero": False} + else: + response = await self.as_llm([UserMsg(name="user", content=prompt)]) + if hasattr(type(response), "__aiter__"): + last = None + async for chunk in response: + last = chunk + response = last + raw = await CompressorStep._response_text(response) # pylint: disable=protected-access + usage = asdict(response.usage) if response is not None and response.usage is not None else None + parsed = parse_judge_response(raw) + score = parsed["acc"] + self.context.response.success = True + self.context.response.answer = str(score) + self.context.response.metadata.update( + **result, + **parsed, + llm_judge_score=score, + raw_response=raw, + usage=usage, + ) + return self.context.response diff --git a/plugins/memlens-judge/src/judge_memlens/plugin.yaml b/plugins/memlens-judge/src/judge_memlens/plugin.yaml new file mode 100644 index 00000000..330566eb --- /dev/null +++ b/plugins/memlens-judge/src/judge_memlens/plugin.yaml @@ -0,0 +1,38 @@ +backends: + memlens_judge_step: judge_memlens.llm_judge:MemLensJudgeStep + +application_defaults: + jobs: + answer_judge: + backend: base + description: "memlens official judge protocol for a saved prediction" + parameters: + type: object + properties: + question: {type: string} + reference: {type: string} + prediction: {type: string} + category: {type: string, default: ""} + subtype: {type: string, default: ""} + old_answer: {type: string, default: ""} + required: [question, reference, prediction, category] + steps: + - backend: memlens_judge_step + as_llm: judge + components: + as_llm: + judge: + backend: openai + model: ${JUDGE_MODEL_NAME:-qwen3.8-max} + stream: false + max_retries: 3 + retry_delay: 60 + client_kwargs: {timeout: 600, max_retries: 0} + credential: + api_key: ${LLM_API_KEY:-} + base_url: ${LLM_BASE_URL:-} + extra_body: {enable_thinking: false} + parameters: + temperature: 0 + thinking_enable: false + max_tokens: 256 diff --git a/plugins/memlens-judge/src/judge_memlens/protocol.py b/plugins/memlens-judge/src/judge_memlens/protocol.py new file mode 100644 index 00000000..19ef3039 --- /dev/null +++ b/plugins/memlens-judge/src/judge_memlens/protocol.py @@ -0,0 +1,139 @@ +"""MemLens's task-specific binary judge for fixed saved predictions. + +Prompts, answer preprocessing and verdict parsing follow SOURCE_REVISION. +Copyright (c) 2026 Xiyu Ren, Zhaowei Wang, and the MemLens Authors. +The complete MIT notice accompanies the verbatim _memlens_judge_prompts.json. +Only exhausted API errors differ: they remain coverage failures, not wrong answers. +""" + +from functools import lru_cache +import json +from pathlib import Path +import re + +SOURCE_REVISION = "77f3ab9a52fa2d6a17978e2dffe80438a4ecced2" + +PROTOCOL_ID = "memlens-task-binary-api-qwen-v1" +_HEDGING_MARKERS = ( + "wait", + "let me re-read", + "let me re read", + "let me double check", + "let me double-check", + "let me reconsider", + "let me look again", + "let's reconsider", + "let s reconsider", + "re-evaluating", + "re evaluating", + "actually looking", + "actually let me", + "hold on", + "let me look at", + "let me check", + "let's re-read", + "let s re-read", + "i need to reconsider", + "wait let me", +) + + +@lru_cache(maxsize=1) +def _prompts() -> dict: + """Load the checked-in, verbatim official prompt strings, never a network resource.""" + return json.loads(Path(__file__).with_name("_memlens_judge_prompts.json").read_text(encoding="utf-8")) + + +def task_key(category: str, subtype: str, reference: str) -> str: + """Official question-type/subtype routing, including reference-based MSR choices.""" + if category == "multi_session_reasoning": + if subtype == "arithmetic": + return "MSR_Arithmetic" + if subtype == "entity_resolution" and reference.strip() in {"Yes", "No"}: + return "MSR_YesNo" + return "MSR_Counting" + if category == "temporal_reasoning": + return { + "duration_comparison": "TR_DurationComparison", + "order_ranking": "TR_OrderRanking", + }.get(subtype, "TR_DateExtraction") + fixed = {"knowledge_update": "KU_KnowledgeUpdate", "answer_refusal": "AR_AnswerRefusal"} + if category in fixed: + return fixed[category] + if category == "information_extraction": + return "IE_PreviousInfo" if subtype == "previnfo" else "IE_Entity" + raise ValueError(f"Unsupported MemLens judge category: {category}") + + +def judge_prompt( + question: str, + reference: str, + prediction: str, + category: str, + subtype: str = "", + old_answer: str = "", +) -> str: + """Render the exact official rubric/examples and current-case fields.""" + if any(not isinstance(value, str) for value in (question, reference, prediction, category, subtype, old_answer)): + raise ValueError("MemLens judge inputs must be strings") + key = task_key(category, subtype, reference) + prompts = _prompts() + prompt = prompts["template"].format(task_criteria=prompts["criteria"][key], examples=prompts["examples"][key]) + if key == "AR_AnswerRefusal": + reference = "Insufficient information" + old = f": {old_answer}\n" if key == "KU_KnowledgeUpdate" and old_answer else "" + return ( + f"{prompt}[Current Case]\n: {question}\n: {reference}\n" + f"{old}: {prediction}\n\n[Scoring Rationale]:" + ) + + +def normalize_for_judge(text: str) -> dict: + """Preserve native parsing, circularity diagnostics and final-30% tail selection.""" + if not text: + return {"text": "", "original_len": 0, "was_truncated": False, "is_degenerate": False, "diagnostics": {}} + original_len = len(text) + text = parse_model_output(text) + hedging, ratio = 0, 1.0 + if len(text) >= 500: + lower = text.lower() + hedging = sum(lower.count(marker) for marker in _HEDGING_MARKERS) + words = lower.split() + ngrams = [tuple(words[index : index + 5]) for index in range(len(words) - 4)] + ratio = len(set(ngrams)) / len(ngrams) if ngrams else 1.0 + truncated = len(text) > 3000 and len(text) > 6000 + if truncated: + text = "[...truncated earlier reasoning...]\n" + text[int(len(text) * (1 - 0.3)) :] + return { + "text": text, + "original_len": original_len, + "was_truncated": truncated, + "is_degenerate": hedging >= 6 or ratio < 0.40, + "diagnostics": {"hedging_count": hedging, "unique_5gram_ratio": round(ratio, 3)}, + } + + +def parse_judge_response(text: str) -> dict: + """Official first-match priority/clamping; even unparseable outputs natively score zero.""" + patterns = ( + (r'"answer_score"\s*:\s*(\d+)', 0), + (r"deserves\s+(\d+)\s+point", re.IGNORECASE), + (r"\[Score\]\s*:\s*(\d+)", 0), + ) + for pattern, flags in patterns: + match = re.search(pattern, text, flags) if text else None + if match: + return {"acc": float(min(int(match.group(1)), 1)), "parse_default_zero": False} + return {"acc": 0.0, "parse_default_zero": True} + + +def parse_model_output(text: str) -> str: + """Native special-token/think-trace parsing; retained raw outputs remain separate.""" + if not text: + return "" + if "" in text: + text = text.split("", 1)[1] + text = re.sub(r"<\|[^>]*\|>", "", text) + text = re.sub(r"<\|begin_of_box\|>(.*?)<\|end_of_box\|>", r"\1", text, flags=re.DOTALL) + text = text.replace("<|begin_of_box|>", "").replace("<|end_of_box|>", "") + return text.replace("[CONTENT_FILTER_REJECTED]", "").strip() diff --git a/plugins/memlens/README.md b/plugins/memlens/README.md new file mode 100644 index 00000000..81e982c4 --- /dev/null +++ b/plugins/memlens/README.md @@ -0,0 +1,48 @@ +# MemLens plugin + +Install with `pip install -e plugins/memlens` and load the `memlens-benchmark` configuration +with the `memlens` plugin. The executable ReMe Image evaluation lives +in [`benchmark/memlens`](../../benchmark/memlens). + +`auto_memory` extends native Auto Memory with a MemLens-specific system prompt. +It retains native source-session persistence and interleaved image inputs; `include_images` +selects whether the memory model sees historical images. No question or reference answer +is given during ingestion. + +`agentic_answer` accepts the official answer prompt as `query`, optional `query_time`, +and native image DataBlocks as `query_images`. It creates a fresh AgentScope Agent for +each question with BM25 `search` and scoped `read`. Question images remain visible in +both the text-memory and image-memory settings. +Its `max_iters` input defaults to 30 ReAct iterations and applies only +to the current answer Agent. The effective value is returned in answer metadata; +Auto Memory and context compression keep their existing settings. + +`include_source_sessions` defaults to `true`: memory prompts include exact physical +JSONL line references, and the answering Agent can search/read those source sessions. +When false, only generated daily/digest markdown is searchable and readable, regardless +of whether an old session index exists. Link expansion is disabled. JSONL persistence +is unchanged. The runner must also omit session directories from index updates for +the source-disabled setting. + +With `include_images=true`, native Auto Memory saves input images beneath +`session/images/` and supplies their paths to the memory Agent. The answer tool +`read_image(path)` calls ReMe's existing `read_image` Job and converts its successful +image response to an AgentScope DataBlock. The plugin supplies the native Job with +`max_bytes=52428800` (50 MiB). Image inputs retain their original bytes and resolution. +Native size errors and non-image results +stay textual tool errors; the generic ReMe Job is unchanged. + +Image access is independent of `include_source_sessions`: a source-disabled Agent +can still reopen images cited by its memory notes. ReMe Text has no historical-image +tool. All image paths are confined to the current scene's session image directory. +The runner retains official Image IDs as text beside the corresponding image blocks; +the IDs are answer evidence labels, not a separate registry or tool argument. +Do not put answer keys, judge results, another scene's data, or agent state under the +daily/digest/source evidence directories. Every scene starts from a separate workspace; +sessions within one scene accumulate in order. Query state and search deduplication +are independent even when questions execute concurrently. + +The plugin preserves the ReMe Image experiment's thinking, output budgets and context +settings. The runner selects the candidate model; scoring belongs to `memlens-judge`. +Answer metadata includes portable usage and the raw final-message usage when available; +missing provider cache/reasoning counters are not reported as zero. diff --git a/plugins/memlens/pyproject.toml b/plugins/memlens/pyproject.toml new file mode 100644 index 00000000..2bc97530 --- /dev/null +++ b/plugins/memlens/pyproject.toml @@ -0,0 +1,25 @@ +[project] +name = "reme-memlens" +version = "0.1.0" +description = "MemLens multimodal memory benchmark plugin for ReMe." +readme = "README.md" +license = "Apache-2.0" +requires-python = ">=3.11" +dependencies = ["reme-ai[as]>=0.4.1.11", "pillow>=10.0.0"] + +[project.entry-points."reme.plugins"] +memlens = "reme_memlens" + +[project.entry-points."reme.configs"] +memlens-benchmark = "reme_memlens:BENCHMARK_CONFIG" + +[tool.setuptools.packages.find] +where = ["src"] +include = ["reme_memlens*"] + +[tool.setuptools.package-data] +reme_memlens = ["plugin.yaml", "*.yaml"] + +[build-system] +requires = ["setuptools>=77", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/plugins/memlens/src/reme_memlens/__init__.py b/plugins/memlens/src/reme_memlens/__init__.py new file mode 100644 index 00000000..a81a2c73 --- /dev/null +++ b/plugins/memlens/src/reme_memlens/__init__.py @@ -0,0 +1,5 @@ +"""MemLens benchmark Jobs for ReMe.""" + +from pathlib import Path + +BENCHMARK_CONFIG = Path(__file__).with_name("benchmark.yaml") diff --git a/plugins/memlens/src/reme_memlens/agentic_answer.py b/plugins/memlens/src/reme_memlens/agentic_answer.py new file mode 100644 index 00000000..2115d95d --- /dev/null +++ b/plugins/memlens/src/reme_memlens/agentic_answer.py @@ -0,0 +1,136 @@ +"""Agentic BM25 answer with scene-local, optional multimodal source access.""" + +from pathlib import Path +from uuid import uuid4 + +from agentscope.message import Base64Source, DataBlock, TextBlock, ToolResultState, UserMsg +from agentscope.tool import FunctionTool, ToolChunk, Toolkit + +from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper +from reme.steps.base_step import BaseStep +from reme.steps.file_io._path import _check_path_permission, resolve_path + + +def evidence_paths(workspace: Path, daily_dir: str, digest_dir: str, session_dir: str, sources: bool) -> list[str]: + """List exact allowed evidence files, never agent state, settings, or another workspace.""" + roots = [(daily_dir, "*.md"), (digest_dir, "*.md")] + if sources: + roots.append((f"{session_dir}/dialog", "*.jsonl")) + workspace = workspace.resolve() + paths = [] + for directory, pattern in roots: + root = (workspace / directory).resolve() + if not root.is_relative_to(workspace): + raise ValueError("Evidence directory must stay inside this scene workspace") + for path in sorted(root.rglob(pattern)): + if path.is_file() and path.resolve().is_relative_to(root): + paths.append(path.relative_to(workspace).as_posix()) + return paths + + +def image_reader(step: BaseStep) -> FunctionTool: + """Bridge the native read_image Job response to an AgentScope image result.""" + workspace = step.workspace_path.resolve() + image_directory = f"{step.config_value('session_dir')}/images" + + async def read_image(path: str) -> ToolChunk: + """View an image using its session/images path cited in a memory note.""" + target, error = resolve_path(workspace, path) + if error or target is None or not _check_path_permission(workspace, target, [image_directory]): + return ToolChunk( + content=[TextBlock(text=error or "Image path must stay inside this scene's session image directory.")], + state=ToolResultState.ERROR, + ) + relative = target.relative_to(workspace).as_posix() + response = await step.run_job("read_image", path=relative) + metadata = response.metadata + media_type = metadata.get("mime", "") + if ( + not response.success + or metadata.get("oversized") + or metadata.get("non_image_warning") + or not media_type.startswith("image/") + ): + notice = str(response.answer) if not response.success or metadata.get("oversized") else "Not an image file." + return ToolChunk(content=[TextBlock(text=notice)], state=ToolResultState.ERROR) + image = DataBlock(source=Base64Source(media_type=media_type, data=response.answer)) + return ToolChunk(content=[TextBlock(text=f"Image source: {relative}"), image], state=ToolResultState.SUCCESS) + + return FunctionTool(func=read_image, name="read_image", is_concurrency_safe=True, is_read_only=True) + + +class MemLensAgenticAnswerStep(BaseStep): + """Run a fresh Agent for each question, with the official answer prompt unchanged.""" + + async def execute(self): + """Answer through a fresh, scene-scoped Agent and its allowed tools.""" + wrapper = self.agent_wrapper + if not isinstance(wrapper, AsAgentWrapper): + raise NotImplementedError("MemLens agentic_answer requires the AgentScope wrapper") + query = self.context.get("query", "") + if not query: + raise ValueError("query must contain the benchmark answer prompt") + max_iters = self.context.get("max_iters", 30) + if isinstance(max_iters, bool) or not isinstance(max_iters, int) or max_iters < 1: + raise ValueError("max_iters must be a positive integer") + sources = bool(self.context.get("include_source_sessions", True)) + images = bool(self.context.get("include_images", False)) + paths = evidence_paths( + self.workspace_path, + self.config_value("daily_dir"), + self.config_value("digest_dir"), + self.config_value("session_dir"), + sources, + ) + tool_context_id = f"memlens_answer_{uuid4().hex}" + tools = [] + if images: + tools.append(image_reader(self)) + injected = { + "_allowed_paths": paths, + "read_step_format_session": True, + "search_filter": {"paths": paths}, + "vector_weight": 0.0, + } + # Standard ReMe job tools keep their normal file handling and lifecycle. + tools.extend( + wrapper._make_tool(job, tool_context_id, injected) # pylint: disable=protected-access + for job in wrapper._resolve_job_tools(["search", "read"]) # pylint: disable=protected-access + ) + blocks = [TextBlock(text=query)] + blocks.extend(DataBlock.model_validate(block) for block in self.context.get("query_images", [])) + if any(not block.source.media_type.startswith("image/") for block in blocks[1:]): + raise ValueError("query_images accepts image DataBlocks only") + prompt = self.get_prompt("system_prompt") + if query_time := self.context.get("query_time"): + prompt += f"\nQuestion time: {query_time}" + if images: + prompt += ( + "\nUse read_image(path=...) for the session/images paths cited in memory, including source_images." + ) + else: + prompt += "\nHistory images cannot be reopened in this setting." + try: + result = await wrapper.reply( + UserMsg(name="user", content=blocks), + system_prompt=prompt, + toolkit=Toolkit(tools=tools), + job_tools=[], + builtin_tools=False, + session_id=str(uuid4()), + resume="", + react_config={"max_iters": max_iters}, + ) + self.context.response.answer = (result.get("result") or "").strip() + self.context.response.metadata.update( + usage=result.get("usage"), + agentscope_usage=(result.get("last_message") or {}).get("usage"), + agent_session_id=result.get("session_id"), + include_source_sessions=sources, + include_images=images, + max_iters=max_iters, + ) + finally: + if self.app_context is not None: + self.app_context.metadata.get("tool_contexts", {}).pop(tool_context_id, None) + return self.context.response diff --git a/plugins/memlens/src/reme_memlens/agentic_answer.yaml b/plugins/memlens/src/reme_memlens/agentic_answer.yaml new file mode 100644 index 00000000..42f862e3 --- /dev/null +++ b/plugins/memlens/src/reme_memlens/agentic_answer.yaml @@ -0,0 +1,10 @@ +system_prompt: | + Answer the supplied benchmark question by retrieving evidence from this scene's memory. + Use search to find relevant facts, then read exact files or line ranges when more context is needed. + Search again with different names, dates or distinctive visual attributes if the evidence is incomplete. + Keep different people, times and visually similar objects distinct. Source line references point to + the original conversation, when source access is enabled. If read_image is available, use the image + path cited in the memory note or its source_images field to verify the actual historical image. + Preserve Image IDs exactly when the question asks which image. Do not invent IDs or unseen details. + Follow the question's supplied answer-format instructions and answer only from supported evidence. + Do not include the retrieval process or tool transcript in the final answer. diff --git a/plugins/memlens/src/reme_memlens/auto_memory.py b/plugins/memlens/src/reme_memlens/auto_memory.py new file mode 100644 index 00000000..3f473536 --- /dev/null +++ b/plugins/memlens/src/reme_memlens/auto_memory.py @@ -0,0 +1,34 @@ +"""MemLens memory prompts with optional, exact source-line references.""" + +from agentscope.message import Msg + +from reme.steps.evolve.auto_memory import AutoMemoryStep + + +class MemLensAutoMemoryStep(AutoMemoryStep): + """Keep core persistence and multimodal input, adding benchmark instructions only.""" + + def _format_history(self, messages: list[Msg]) -> str: + if not self.context.get("include_source_sessions", True): + return "Source-session access is disabled; record self-contained facts.\n\n" + super()._format_history( + messages, + ) + + # Core execute has saved the session already. Read physical line numbers + # rather than assuming input order equals timestamp-sorted JSONL order. + session_id = self.context.get("session_id", "") + source = self._session_source_path(session_id) + # File iteration keeps Unicode line separators inside JSON strings intact. + with self._session_path(session_id).open(encoding="utf-8") as stream: + lines = {Msg.model_validate_json(line).id: number for number, line in enumerate(stream, 1) if line.strip()} + turns = [ + f"[L{lines[msg.id]} | {msg.name or msg.role} @ {msg.created_at}]\n{msg.get_text_content() or ''}" + for msg in messages + ] + return ( + f"Source file: {source}\n" + f"Cite useful evidence with the full source path, for example [[{source}#L1-L2]], " + "replacing the line range with the actual physical lines below. Never use bare links such as [[L5]]. " + "Preserve original Image IDs alongside their facts, and retain the supplied Image source paths.\n\n" + + "\n\n".join(turns) + ) diff --git a/plugins/memlens/src/reme_memlens/auto_memory.yaml b/plugins/memlens/src/reme_memlens/auto_memory.yaml new file mode 100644 index 00000000..f318b669 --- /dev/null +++ b/plugins/memlens/src/reme_memlens/auto_memory.yaml @@ -0,0 +1,25 @@ +system_prompt: | + You maintain long-horizon personal multimodal memory across dated conversations. Use tools to save useful evidence, + not a response to any future question. Retain concrete people, entities, events, dates, preferences, decisions, + quantities, places and relationships. Preserve changes over time and distinguish speakers and similar objects. + Track longitudinal updates to preferences, plans, possessions, health, work and relationships without erasing + earlier dated facts. Keep which person experienced each event explicit and resolve relative dates using the + session timestamp only when justified. Separate repeated similar events instead of conflating them. + + When images are present, jointly understand them and the surrounding conversation. Record distinctive visual + details (objects, attributes, layout, visible text, activities) that make an image identifiable. Keep each original + Image ID verbatim next to the corresponding visual facts; an ID is a reference, not evidence of image contents. + Do not invent visual details when no image is provided. Do not substitute generic summaries for specific facts. + + When physical source-line labels are supplied, cite the useful source ranges with the exact session path and + line numbers. Without source access, make the memory self-contained rather than adding source references. + Organize related facts under descriptive headings, use absolute dates when supported, preserve exact names + and numbers, and distinguish facts from uncertainty. Resolve pronouns only when the conversation supports it. + + Follow the normal daily-note contract: a stable filename stem in name, an informative description, and the + factual body. Never set status. Invoke the provided file tools as structured tool calls to actually save memory. + + Complete the task, then stop: create at most one note for a new session. After a successful daily_write, + return the requested one-sentence summary without calling daily_write again or creating a renamed duplicate. + Further writes are needed only to correct an explicit tool error or an identified omission, not to repeat a + successful save. For an existing note, finish after the necessary edits and frontmatter update succeed. diff --git a/plugins/memlens/src/reme_memlens/benchmark.yaml b/plugins/memlens/src/reme_memlens/benchmark.yaml new file mode 100644 index 00000000..1421f4fc --- /dev/null +++ b/plugins/memlens/src/reme_memlens/benchmark.yaml @@ -0,0 +1,30 @@ +extends: benchmark +plugins: [memlens] +components: + as_llm: + default: &memory_model + backend: openai + model: ${LLM_MODEL_NAME:-qwen3.8-max} + stream: false + context_size: 200000 + max_retries: 3 + retry_delay: 60 + client_kwargs: {timeout: 600, max_retries: 0} + extra_body: {enable_thinking: true} + parameters: {temperature: 0, thinking_enable: true, max_tokens: 65536} + bench: + <<: *memory_model + context_size: 400000 + parameters: {temperature: 0, thinking_enable: true, max_tokens: 16384} + compressor: + <<: *memory_model + agent_wrapper: + default: &memory_agent + react_config: {max_iters: 30} + context_config: {trigger_ratio: 0.8, reserve_ratio: 0.1, tool_result_limit: 50000, max_image_num: 4096} + model_config: {max_retries: 0} + bench: + <<: *memory_agent + file_store: + default: + embedding_store: null diff --git a/plugins/memlens/src/reme_memlens/plugin.yaml b/plugins/memlens/src/reme_memlens/plugin.yaml new file mode 100644 index 00000000..56a6f64e --- /dev/null +++ b/plugins/memlens/src/reme_memlens/plugin.yaml @@ -0,0 +1,75 @@ +backends: + memlens_auto_memory_step: reme_memlens.auto_memory:MemLensAutoMemoryStep + memlens_agentic_answer_step: reme_memlens.agentic_answer:MemLensAgenticAnswerStep + +application_defaults: + jobs: + read_image: + backend: base + description: "Read a saved session image through the native ReMe image Job." + parameters: + type: object + properties: + path: {type: string, description: "Workspace-relative image attachment path."} + required: [path] + steps: + - backend: read_image_step + max_bytes: 52428800 + auto_memory: + backend: base + description: "Record one MemLens session using native Auto Memory." + include_source_sessions: true + parameters: + type: object + properties: + messages: + type: array + items: {type: object} + session_id: {type: string} + date: {type: string, default: ""} + include_images: {type: boolean, default: false} + include_source_sessions: {type: boolean, default: true} + memory_hint: {type: string, default: ""} + required: [messages, session_id] + steps: + - backend: memlens_auto_memory_step + search: + backend: base + description: "Search this scene's allowed memory and source text using BM25." + parameters: + type: object + properties: + query: {type: string} + limit: {type: integer, description: "Number of search results."} + start_date: {type: string} + end_date: {type: string} + required: [query] + steps: + - backend: search_step + vector_weight: 0.0 + expand_links: false + agentic_answer: + backend: base + description: "Answer a MemLens query through a fresh multimodal retrieval Agent." + include_source_sessions: true + parameters: + type: object + properties: + query: {type: string, description: "Official benchmark answer prompt including the query."} + query_time: {type: string, default: ""} + query_images: + type: array + items: {type: object} + default: [] + scene_id: {type: string} + max_iters: {type: integer, minimum: 1, default: 30, description: "Maximum answer Agent ReAct iterations."} + include_images: {type: boolean, default: false} + include_source_sessions: {type: boolean, default: true} + required: [query, scene_id] + steps: + - backend: memlens_agentic_answer_step + agent_wrapper: bench + components: + file_store: + default: + embedding_store: null diff --git a/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/reme/components/as_decision_model/agentscope.py b/reme/components/as_decision_model/agentscope.py index 640e3a93..a5699acf 100644 --- a/reme/components/as_decision_model/agentscope.py +++ b/reme/components/as_decision_model/agentscope.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter from .base_as_decision_model import BaseAsDecisionModel from ..component_registry import R +from ...utils.async_utils import complete_in_thread if TYPE_CHECKING: from agentscope.classifier import ClassifierModelBase, ClassifierResponse @@ -117,18 +118,19 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel): raise ValueError("Decision model inference requires DECISION_MODEL_API_KEY or credential.api_key.") if not credential.get("base_url"): credential["base_url"] = None + # Reserve ownership before yielding so host injection cannot race + # initialization. The worker publishes the model to this component + # before cancellation can release the lifecycle lock. + self._owns_model = True try: - credential_cls = import_module("agentscope.credential").TypeSafeCredential - self._model = model_cls( - credential=credential_cls(**credential), - **self.options.model_dump(exclude_none=True), - ) - self._owns_model = True + await complete_in_thread(self._create_model, model_cls, credential) except ImportError as exc: raise ImportError( "Decision model inference requires a compatible typesafe-sdk. " "Install it with `pip install 'reme-ai[decision]'`.", ) from exc + finally: + self._owns_model = self._model is not None model = self._model self._active_calls += 1 self._idle.clear() @@ -141,6 +143,14 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel): if self._active_calls == 0: self._idle.set() + def _create_model(self, model_cls, credential) -> None: + """Construct and retain the owned client, including synchronous TLS setup.""" + credential_cls = import_module("agentscope.credential").TypeSafeCredential + self._model = model_cls( + credential=credential_cls(**credential), + **self.options.model_dump(exclude_none=True), + ) + async def _close(self) -> None: """Finish draining and cleanup before propagating shutdown cancellation.""" cleanup = asyncio.create_task(self._release_model()) diff --git a/reme/config/default.yaml b/reme/config/default.yaml index c57c0b3c..19b4423d 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -16,6 +16,7 @@ jobs: dispatch_steps: [update_index_step] - backend: watch_changes_step watch_mode: awatch # Set to watch to isolate the blocking wait from AnyIO's shared thread limiter. + monitor_type: file_store dispatch_steps: - backend: update_index_step persist: False @@ -37,6 +38,8 @@ jobs: - auto_text_resource_step - backend: watch_changes_step watch_mode: awatch + monitor_type: file_catalog + monitor_name: resource dispatch_steps: - backend: update_catalog_step file_catalog: resource @@ -59,6 +62,8 @@ jobs: - backend: log_changes_step - backend: watch_changes_step watch_mode: awatch + monitor_type: file_catalog + monitor_name: digest dispatch_steps: - backend: update_catalog_step file_catalog: digest @@ -513,14 +518,33 @@ jobs: daily_list: backend: base - description: "List notes under a single day." + description: >- + List daily notes without changing files or indexes. Defaults to all notes for today; date selects one day. + Alternatively, start_date and end_date select an inclusive range, newest day first then path ascending. + Range queries or explicit page/page_size opt into pagination (defaults: page 1, page_size 100). + Metadata adds page, page_size, total_notes and total_pages; count and notes describe the returned page. parameters: type: object properties: date: type: string - description: "YYYY-MM-DD; empty = today" + description: "YYYY-MM-DD; empty = today when no range is given; cannot be combined with a range" default: "" + start_date: + type: string + description: "inclusive YYYY-MM-DD range start; requires end_date; cannot be combined with date" + end_date: + type: string + description: "inclusive YYYY-MM-DD range end; requires start_date" + page: + type: integer + minimum: 1 + description: "optional 1-based page; defaults to 1 when paginating; past the end returns no notes" + page_size: + type: integer + minimum: 1 + maximum: 1000 + description: "optional notes per page; defaults to 100 when paginating" steps: - backend: daily_list_step diff --git a/reme/steps/compaction/__init__.py b/reme/steps/compaction/__init__.py index 57aeeab5..eaeb79b4 100644 --- a/reme/steps/compaction/__init__.py +++ b/reme/steps/compaction/__init__.py @@ -1,5 +1,6 @@ -"""Tool-history compaction Steps.""" +"""Tool-history compaction and shared host validation helpers.""" from .compact_context import CompactContextStep +from .messages import ToolPair, collect_pairs, preview_result -__all__ = ["CompactContextStep"] +__all__ = ["CompactContextStep", "ToolPair", "collect_pairs", "preview_result"] diff --git a/reme/steps/evolve/auto_tag.py b/reme/steps/evolve/auto_tag.py index 05ea0604..823d59cb 100644 --- a/reme/steps/evolve/auto_tag.py +++ b/reme/steps/evolve/auto_tag.py @@ -123,6 +123,8 @@ class AutoTagStep(BaseStep): path = self.workspace_path / target.path metadata = dict(frontmatter.loads(path.read_text(encoding="utf-8")).metadata or {}) + if tag_key not in metadata: + raise RuntimeError(f"Auto-tag agent did not write the required {tag_key!r} field") normalized = normalize_memory_tags( metadata.get(tag_key), max_tags_per_file=tag_index.max_tags_per_file, diff --git a/reme/steps/file_io/_daily_index.py b/reme/steps/file_io/_daily_index.py index 72a8863e..49d66fde 100644 --- a/reme/steps/file_io/_daily_index.py +++ b/reme/steps/file_io/_daily_index.py @@ -6,7 +6,7 @@ from pathlib import Path import frontmatter -from ._path import resolve_path, validate_filename_component +from ._path import _check_path_permission, resolve_path, validate_filename_component from ...utils import get_logger logger = get_logger(log_to_file=False) @@ -64,7 +64,7 @@ def _rebuild_body(body: str, notes_content: str) -> str: return f"{body.rstrip()}\n\n{block}\n" if body.strip() else f"{block}\n" -def scan_notes(workspace_dir: Path, date: str, daily_dir: str) -> list[dict]: +def scan_notes(workspace_dir: Path, date: str, daily_dir: str, *, allowed_paths=None) -> list[dict]: """Walk ``//*.md`` and pull each note's frontmatter. Files under the date folder are named by note ``name``. Conversation @@ -99,6 +99,8 @@ def scan_notes(workspace_dir: Path, date: str, daily_dir: str) -> list[dict]: target, err = resolve_path(workspace_dir, str(md_path)) if err or target is None: continue + if not _check_path_permission(workspace_dir, target, allowed_paths): + continue if not target.is_file(): continue except RuntimeError: diff --git a/reme/steps/file_io/_file_io.py b/reme/steps/file_io/_file_io.py index 82c5264d..22e945b0 100644 --- a/reme/steps/file_io/_file_io.py +++ b/reme/steps/file_io/_file_io.py @@ -1,7 +1,10 @@ """Encoding-aware file IO, output truncation, and per-path write locks.""" import asyncio +import os from pathlib import Path +from secrets import token_hex +from stat import S_IMODE from typing import Iterable import aiofiles @@ -10,6 +13,7 @@ import aiofiles.os from ...constants import DEFAULT_MAX_BYTES, MAX_FILE_READ_BYTES, TRUNCATION_NOTICE_MARKER from ...schema import FileChunk from ...utils import get_logger +from ...utils.async_utils import complete_in_thread logger = get_logger(log_to_file=False) @@ -152,7 +156,7 @@ async def detect_file_encoding(file_path, sniff_bytes: int = 8192) -> str: async def write_file_safe(file_path: Path, content: str | bytes, encoding: str = "utf-8") -> None: - """Write ``content`` to ``file_path`` in binary mode; creates parent dirs.""" + """Publish complete bytes atomically while retaining normal file permissions.""" file_path.parent.mkdir(parents=True, exist_ok=True) if isinstance(content, str): try: @@ -165,8 +169,32 @@ async def write_file_safe(file_path: Path, content: str | bytes, encoding: str = payload = content.encode("utf-8") else: payload = content - async with aiofiles.open(str(file_path), "wb") as f: - await f.write(payload) + + def publish(): + flags = os.O_WRONLY | getattr(os, "O_BINARY", 0) + # Check OS write access without truncating the existing bytes. + try: + existing_fd = os.open(file_path, flags) + except FileNotFoundError: + mode = None + else: + try: + mode = S_IMODE(os.fstat(existing_fd).st_mode) + finally: + os.close(existing_fd) + + temporary = file_path.parent / f".reme-{token_hex(16)}.tmp" + fd = os.open(temporary, flags | os.O_CREAT | os.O_EXCL, 0o600 if mode is not None else 0o666) + try: + with os.fdopen(fd, "wb") as handle: + handle.write(payload) + if mode is not None: + os.chmod(temporary, mode) + os.replace(temporary, file_path) + finally: + Path(temporary).unlink(missing_ok=True) + + await complete_in_thread(publish) # --------------------------------------------------------------------------- diff --git a/reme/steps/file_io/daily_list.py b/reme/steps/file_io/daily_list.py index 643a2f15..53058119 100644 --- a/reme/steps/file_io/daily_list.py +++ b/reme/steps/file_io/daily_list.py @@ -1,4 +1,4 @@ -"""``daily_list`` — list the notes under a single day (pure read, no side effects). +"""``daily_list`` — list daily notes with optional date ranges and pagination (pure read). Returns one row per ``daily//.md`` note file. The human answer contains the workspace-relative ``path`` plus every frontmatter key/value. @@ -10,13 +10,15 @@ Response metadata includes ``notes`` as the original list of flat dicts: explicitly when the index page needs to be rebuilt. Decoupling read from write keeps each step's effect predictable. -Input is a single optional ``date`` (``YYYY-MM-DD``); falls back -to today. +Without range or pagination parameters, the optional ``date`` defaults to today +and the original response shape is preserved. Ranges are inclusive and ordered +newest day first, then by path. Pagination limits output, not scan I/O. """ +import asyncio from pathlib import Path -from ._daily_index import scan_notes +from ._daily_index import parse_daily_date, scan_notes from ._path import resolve_path from ..base_step import BaseStep from ...components import R @@ -25,16 +27,58 @@ from ...steps.evolve import now @R.register("daily_list_step") class DailyListStep(BaseStep): - """List the notes under a single day. Pure read — no index refresh.""" + """Read daily notes without refreshing indexes or changing workspace files.""" - def _collect_params(self) -> tuple[str, str, Path]: - """Read ``date`` (default today, ``YYYY-MM-DD``), ``daily_dir``, and the workspace root.""" + def _collect_params(self) -> tuple[dict, int | None, int | None]: + """Keep legacy calls unpaged; range queries opt into bounded pages.""" assert self.context is not None - tz = self.app_context.app_config.timezone if self.app_context is not None else None - day = self.context.get("date", "") or now(tz).strftime("%Y-%m-%d") - daily_dir = self.config_value("daily_dir") - workspace_dir = Path(self.file_store.workspace_path or ".").resolve() - return day, daily_dir, workspace_dir + ranged = any(key in self.context for key in ("start_date", "end_date")) + if ranged: + if self.context.get("date"): + raise ValueError("date cannot be combined with start_date/end_date") + scope = {key: self.context.get(key) for key in ("start_date", "end_date")} + else: + tz = self.app_context.app_config.timezone if self.app_context is not None else None + scope = {"date": self.context.get("date") or now(tz).strftime("%Y-%m-%d")} + for key, value in scope.items(): + if not isinstance(value, str) or parse_daily_date(value) != value: + raise ValueError(f"{key} must be YYYY-MM-DD") + if ranged and scope["start_date"] > scope["end_date"]: + raise ValueError("start_date must not be later than end_date") + if not ranged and not any(key in self.context for key in ("page", "page_size")): + return scope, None, None + page, page_size = self.context.get("page", 1), self.context.get("page_size", 100) + if type(page) is not int or page < 1: + raise ValueError("page must be an integer greater than or equal to 1") + if type(page_size) is not int or not 1 <= page_size <= 1000: + raise ValueError("page_size must be an integer between 1 and 1000") + return scope, page, page_size + + @classmethod + def _scan(cls, workspace_dir: Path, daily_dir: str, scope: dict, allowed_paths) -> list[dict]: + """Inspect existing date directories only; permission checks precede file reads.""" + root, err = resolve_path(workspace_dir, daily_dir) + if err or root is None: + raise ValueError(err or "invalid daily directory") + if "date" in scope: + days = [scope["date"]] + elif root.is_dir(): + days = sorted( + ( + entry.name + for entry in root.iterdir() + if parse_daily_date(entry.name) == entry.name + and scope["start_date"] <= entry.name <= scope["end_date"] + ), + reverse=True, + ) + else: + days = [] + return [ + cls._project(note) + for day in days + for note in scan_notes(workspace_dir, day, daily_dir, allowed_paths=allowed_paths) + ] @staticmethod def _project(note: dict) -> dict: @@ -56,19 +100,46 @@ class DailyListStep(BaseStep): return " ".join(parts) async def execute(self): - """Scan ``//`` and emit one projected record per note.""" + """List authorized notes, preserving the legacy single-day response by default.""" assert self.context is not None - day, daily_dir, workspace_dir = self._collect_params() - _target_dir, err = resolve_path(workspace_dir, f"{daily_dir}/{day}") - if err: + try: + scope, page, page_size = self._collect_params() + workspace_dir = Path(self.file_store.workspace_path or ".").resolve() + notes = await asyncio.to_thread( + self._scan, + workspace_dir, + self.config_value("daily_dir"), + scope, + self.context.get("_allowed_paths"), + ) + except (ValueError, OSError, RuntimeError) as exc: self.context.response.success = False - self.context.response.answer = f"Error: {err}" - self.context.response.metadata.update({"date": day, "error": err}) - self.logger.info(f"[{self.name}] date={day} error={err!r}") + self.context.response.answer = f"Error: {exc}" + self.context.response.metadata["error"] = str(exc) return - notes = [self._project(n) for n in scan_notes(workspace_dir, day, daily_dir)] + metadata = dict(scope) + if page is not None and page_size is not None: + total = len(notes) + metadata.update( + { + "page": page, + "page_size": page_size, + "total_notes": total, + "total_pages": (total + page_size - 1) // page_size, + }, + ) + notes = notes[(page - 1) * page_size : page * page_size] + metadata.update({"count": len(notes), "notes": notes}) self.context.response.success = True - lines = [self._format_note_line(n) for n in notes] - self.context.response.answer = "\n".join(lines) if lines else f"No notes found for {day}" - self.context.response.metadata.update({"date": day, "count": len(notes), "notes": notes}) - self.logger.info(f"[{self.name}] date={day} notes={len(notes)}") + lines = [self._format_note_line(note) for note in notes] + label = scope.get("date") or f"{scope['start_date']} to {scope['end_date']}" + empty = f"No notes found for {label}" + if metadata.get("total_notes") and not notes: + empty = f"No notes on page {page} (total notes: {metadata['total_notes']})" + answer = "\n".join(lines) if lines else empty + if page is not None: + header = " ".join(f"{key}: {metadata[key]}" for key in ("page", "page_size", "total_notes", "total_pages")) + answer = f"{header}\n{answer}" + self.context.response.answer = answer + self.context.response.metadata.update(metadata) + self.logger.info(f"[{self.name}] scope={scope} notes={len(notes)}") diff --git a/reme/steps/file_io/edit.py b/reme/steps/file_io/edit.py index b9a457d5..31a0fd76 100644 --- a/reme/steps/file_io/edit.py +++ b/reme/steps/file_io/edit.py @@ -1,5 +1,7 @@ """Find-and-replace text in a markdown file body (front matter is preserved).""" +import re + import frontmatter import yaml @@ -9,6 +11,44 @@ from ._path import _check_path_permission, NON_MD_WARNING, gate_md, resolve_path from ...components import R +def _replace_normalized_newlines(body: str, old: str, new: str) -> tuple[str, int]: + """Find literal LF text, mapping its spans back without normalizing untouched text.""" + normalized = body.replace("\r\n", "\n").replace("\r", "\n") + position = normalized.find(old) + if position < 0: + return body, 0 + + # Only CRLF pairs change offsets. Consume them once as match boundaries advance; + # no per-character offset array or regex search for the user-supplied needle. + crlf_starts = (match.start() for match in re.finditer("\r\n", body)) + next_crlf = next(crlf_starts, len(body)) + removed = 0 + source_cursor = 0 + count = 0 + parts = [] + while position >= 0: + source_start = position + removed + while next_crlf < source_start: + removed += 1 + source_start += 1 + next_crlf = next(crlf_starts, len(body)) + + end = position + len(old) + source_end = end + removed + while next_crlf < source_end: + removed += 1 + source_end += 1 + next_crlf = next(crlf_starts, len(body)) + + parts.extend((body[source_cursor:source_start], new)) + source_cursor = source_end + count += 1 + position = normalized.find(old, end) + + parts.append(body[source_cursor:]) + return "".join(parts), count + + @R.register("edit_step") class EditStep(_FileIOStep): """Replace every occurrence of ``old`` with ``new`` inside the file body. @@ -82,13 +122,16 @@ class EditStep(_FileIOStep): body = raw_text not_found_msg = f"text to replace was not found in {target}" - if old_str not in body: + if "\r" in body and "\n" in old_str and "\r" not in old_str: + new_body, count = _replace_normalized_newlines(body, old_str, new_str) + else: + count = body.count(old_str) + new_body = body.replace(old_str, new_str) + + if count == 0: self._fail(not_found_msg, path=str(target)) return None - count = body.count(old_str) - new_body = body.replace(old_str, new_str) - if is_md and post is not None: post.content = new_body # Re-serialize: keep front matter when present, otherwise emit body alone diff --git a/reme/steps/file_io/read.py b/reme/steps/file_io/read.py index 268f5455..f8093487 100644 --- a/reme/steps/file_io/read.py +++ b/reme/steps/file_io/read.py @@ -152,6 +152,8 @@ class ReadStep(_FileIOStep): if content is None: return None + # Match text-mode streaming reads and the chunker's source line numbers. + content = content.replace("\r\n", "\n").replace("\r", "\n") all_lines = content.split("\n") if content.endswith("\n"): # split() yields a trailing empty element for the final newline; diff --git a/reme/steps/index/watch_changes.py b/reme/steps/index/watch_changes.py index a01d205f..a39eb467 100644 --- a/reme/steps/index/watch_changes.py +++ b/reme/steps/index/watch_changes.py @@ -60,6 +60,8 @@ class WatchChangesStep(BaseStep): step: int = DEFAULT_WATCH_STEP_MS, poll_delay_ms: int = DEFAULT_LOW_POWER_POLL_MS, watch_mode: Literal["awatch", "watch"] = "awatch", + monitor_type: Literal["file_store", "file_catalog"] | None = None, + monitor_name: str = "default", **kwargs, ): if watch_mode not in ("awatch", "watch"): @@ -71,6 +73,8 @@ class WatchChangesStep(BaseStep): self.debounce: int = debounce self.step: int = step self.poll_delay_ms: int = poll_delay_ms + self.monitor_type = monitor_type + self.monitor_name = monitor_name self._rules: list[WatchRule] = [] def _get_watch_rules(self) -> list[WatchRule]: @@ -112,7 +116,9 @@ class WatchChangesStep(BaseStep): step=self.step, poll_delay_ms=self.poll_delay_ms, stop_event=stop_event, + yield_on_timeout=self.monitor_type is not None, ) + needs_initial_scan = self.monitor_type is not None async with aclosing(watcher): async for raw_changes in watcher: if stop_event.is_set(): @@ -122,6 +128,23 @@ class WatchChangesStep(BaseStep): for c, p in raw_changes if c in (Change.added, Change.modified, Change.deleted) ] + if needs_initial_scan: + # Subscribe first, then recover writes made after the + # preceding scan but before watcher initialization. + await self.dispatch_steps( + [ + { + "backend": "init_changes_step", + "monitor_type": self.monitor_type, + "monitor_name": self.monitor_name, + "recursive": self.recursive, + }, + ], + ) + # The scan only produces changes. Merge overlapping events + # before running handlers, which may have model side effects. + changes.extend(self.context.get("changes") or []) + needs_initial_scan = False changes = coalesce_changes(changes) if changes: self.logger.info(f"Detected {len(changes)} change(s)") diff --git a/reme/utils/async_utils.py b/reme/utils/async_utils.py index 290a4456..8f54fbcb 100644 --- a/reme/utils/async_utils.py +++ b/reme/utils/async_utils.py @@ -7,11 +7,21 @@ from typing import Any async def complete_in_thread(func: Callable[..., Any], /, *args) -> Any: - """Finish a side-effecting thread call before propagating cancellation.""" + """Finish a side-effecting thread call, including through repeated cancellation.""" task = asyncio.create_task(asyncio.to_thread(func, *args)) - try: - return await asyncio.shield(task) - except asyncio.CancelledError: + cancellation: asyncio.CancelledError | None = None + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + except Exception: + if cancellation is None: + raise + break + if cancellation is not None: with suppress(Exception): - await task - raise + task.result() + raise cancellation + return task.result() diff --git a/tests/unit/test_async_utils.py b/tests/unit/test_async_utils.py new file mode 100644 index 00000000..70b90adc --- /dev/null +++ b/tests/unit/test_async_utils.py @@ -0,0 +1,100 @@ +"""Side-effecting thread calls retain ownership through cancellation.""" + +import asyncio +import threading +import time + +import anyio +import pytest + +from reme.utils.async_utils import complete_in_thread + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancellations", [0, 1, 3]) +@pytest.mark.parametrize("fail", [False, True]) +async def test_complete_in_thread_finishes_before_propagating_result(cancellations, fail): + """Repeated cancellation waits for success or failure without detaching work.""" + started, release, finished = threading.Event(), threading.Event(), threading.Event() + + def operation(): + started.set() + try: + assert release.wait(5) + if fail: + raise RuntimeError("thread failed") + return "complete" + finally: + finished.set() + + task = asyncio.create_task(complete_in_thread(operation)) + try: + assert await asyncio.to_thread(started.wait, 5) + for _ in range(cancellations): + task.cancel() + await asyncio.sleep(0) + assert not task.done() + assert not finished.is_set() + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + + assert finished.is_set() + if cancellations: + with pytest.raises(asyncio.CancelledError): + task.result() + elif fail: + with pytest.raises(RuntimeError, match="thread failed"): + task.result() + else: + assert task.result() == "complete" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope_name", ["fail_after", "move_on_after"]) +@pytest.mark.parametrize("fail", [False, True]) +async def test_complete_in_thread_preserves_anyio_cancel_scope(scope_name, fail): + """AnyIO must recognize its original cancellation after the thread finishes.""" + finished = threading.Event() + + def operation(): + time.sleep(0.05) + finished.set() + if fail: + raise RuntimeError("thread failed after cancellation") + + async def run_scope(): + with getattr(anyio, scope_name)(0.01) as scope: + await complete_in_thread(operation) + assert scope.cancelled_caught + + if scope_name == "fail_after": + with pytest.raises(TimeoutError): + await run_scope() + else: + await run_scope() + assert finished.is_set() + + +@pytest.mark.asyncio +async def test_complete_in_thread_retains_first_cancellation_message(): + """Repeated cancellation must retain the original caller's reason.""" + started, release = threading.Event(), threading.Event() + + def operation(): + started.set() + assert release.wait(5) + + task = asyncio.create_task(complete_in_thread(operation)) + try: + assert await asyncio.to_thread(started.wait, 5) + for message in ("original cancellation", "cancel again", "cancel once more"): + task.cancel(message) + await asyncio.sleep(0) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError, match="original cancellation"): + await task + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) diff --git a/tests/unit/test_atomic_file_write.py b/tests/unit/test_atomic_file_write.py new file mode 100644 index 00000000..b2a68dbd --- /dev/null +++ b/tests/unit/test_atomic_file_write.py @@ -0,0 +1,184 @@ +"""Memory writers publish complete files to concurrent index readers.""" + +import asyncio +import os +from pathlib import Path +import stat +import subprocess +import sys +import threading + +import pytest + +from reme.steps.file_io._file_io import write_file_safe +from reme.components.application_context import ApplicationContext +from reme.steps.file_io.save import SaveStep + + +@pytest.mark.asyncio +async def test_overwrite_preserves_complete_old_content_until_commit(tmp_path, monkeypatch): + """Readers retain the prior complete revision until atomic publication.""" + target = tmp_path / "memory.md" + target.write_text("old complete memory", encoding="utf-8") + target.chmod(0o640) + replace = os.replace + + def observe_commit(source, destination): + assert target.read_text(encoding="utf-8") == "old complete memory" + assert source != destination + replace(source, destination) + + monkeypatch.setattr(os, "replace", observe_commit) + await write_file_safe(target, "new complete memory") + + assert target.read_text(encoding="utf-8") == "new complete memory" + assert stat.S_IMODE(target.stat().st_mode) == 0o640 + assert sorted(path.name for path in tmp_path.iterdir()) == ["memory.md"] + + +@pytest.mark.asyncio +async def test_failed_commit_preserves_original_and_cleans_temporary_file(tmp_path, monkeypatch): + """Failed publication leaves source bytes intact and no temporary assets.""" + target = tmp_path / "memory.md" + target.write_bytes(b"original memory") + + def fail_commit(_source, _destination): + raise OSError("controlled commit failure") + + monkeypatch.setattr(os, "replace", fail_commit) + with pytest.raises(OSError, match="controlled commit failure"): + await write_file_safe(target, "replacement") + + assert target.read_bytes() == b"original memory" + assert sorted(path.name for path in tmp_path.iterdir()) == ["memory.md"] + + +@pytest.mark.asyncio +async def test_new_encoded_file_and_binary_payload(tmp_path): + """Atomic writes preserve the existing encoding and binary contracts.""" + target = tmp_path / "nested" / "memory.csv" + await write_file_safe(target, "中文", encoding="gbk") + assert target.read_bytes() == "中文".encode("gbk") + await write_file_safe(target, b"\x00\xff") + assert target.read_bytes() == b"\x00\xff" + + +@pytest.mark.asyncio +async def test_atomic_write_accepts_long_valid_filename(tmp_path): + """Temporary names must not grow with the destination's filename.""" + target = tmp_path / ("a" * 241 + ".md") + target.write_text("original", encoding="utf-8") + + await write_file_safe(target, "replacement") + + assert target.read_text(encoding="utf-8") == "replacement" + assert list(tmp_path.iterdir()) == [target] + + +@pytest.mark.asyncio +async def test_cancelled_save_keeps_path_lock_until_publication(tmp_path, monkeypatch): + """A repeatedly cancelled writer cannot overwrite a later successful save.""" + target = tmp_path / "memory.md" + target.write_text("original", encoding="utf-8") + context = ApplicationContext(workspace_dir=str(tmp_path)) + started, release = threading.Event(), threading.Event() + replace = os.replace + + def delayed_replace(source, destination): + if not started.is_set(): + started.set() + assert release.wait(5) + replace(source, destination) + + monkeypatch.setattr(os, "replace", delayed_replace) + first = asyncio.create_task(SaveStep(app_context=context)(path="memory.md", content="cancelled write")) + second = None + try: + assert await asyncio.to_thread(started.wait, 5) + for _ in range(3): + first.cancel() + await asyncio.sleep(0) + second = asyncio.create_task(SaveStep(app_context=context)(path="memory.md", content="successful write")) + await asyncio.sleep(0) + assert not first.done() + assert not second.done() + finally: + release.set() + results = await asyncio.gather(first, *([second] if second is not None else []), return_exceptions=True) + + assert isinstance(results[0], asyncio.CancelledError) + assert results[1].success + assert target.read_text(encoding="utf-8") == "successful write" + assert list(tmp_path.iterdir()) == [target] + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name != "posix", reason="POSIX file permissions") +async def test_save_rejects_read_only_target_without_replacing_it(tmp_path): + """A writable parent must not let saving bypass the target's access rights.""" + if os.geteuid() == 0: + pytest.skip("Root bypasses ordinary file write permissions") + target = tmp_path / "readonly.md" + target.write_bytes(b"original") + target.chmod(0o444) + try: + response = await SaveStep(app_context=ApplicationContext(workspace_dir=str(tmp_path)))( + path=target.name, + content="replacement", + ) + assert response.success is False + assert target.read_bytes() == b"original" + assert stat.S_IMODE(target.stat().st_mode) == 0o444 + assert list(tmp_path.iterdir()) == [target] + finally: + target.chmod(0o644) + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name != "posix", reason="POSIX file permissions") +async def test_atomic_write_does_not_require_target_read_access(tmp_path): + """Permission checks retain the existing write-only file contract.""" + target = tmp_path / "writeonly.md" + target.write_bytes(b"original") + target.chmod(0o200) + try: + await write_file_safe(target, b"replacement") + assert stat.S_IMODE(target.stat().st_mode) == 0o200 + finally: + target.chmod(0o600) + assert target.read_bytes() == b"replacement" + + +@pytest.mark.skipif(os.name != "posix", reason="POSIX umask semantics") +@pytest.mark.parametrize("mask, expected", [(0o002, 0o664), (0o077, 0o600)]) +def test_new_save_respects_process_umask(tmp_path, mask, expected): + """Check real creation modes without changing the test process's umask.""" + script = """ +import asyncio +import os +from pathlib import Path +import stat +import sys +from reme.components.application_context import ApplicationContext +from reme.steps.file_io.save import SaveStep + +workspace = Path(sys.argv[1]) +os.umask(int(sys.argv[2])) +response = asyncio.run(SaveStep(app_context=ApplicationContext(workspace_dir=str(workspace)))( + path="new.md", content="memory" +)) +assert response.success, response.answer +target = workspace / "new.md" +assert target.read_text() == "memory" +assert list(workspace.iterdir()) == [target] +print(stat.S_IMODE(target.stat().st_mode)) +""" + result = subprocess.run( + [sys.executable, "-c", script, str(tmp_path), str(mask)], + cwd=Path(__file__).resolve().parents[2], + capture_output=True, + text=True, + check=True, + timeout=15, + ) + assert int(result.stdout.splitlines()[-1]) == expected diff --git a/tests/unit/test_auto_tag.py b/tests/unit/test_auto_tag.py index 3336a381..69d6a6f1 100644 --- a/tests/unit/test_auto_tag.py +++ b/tests/unit/test_auto_tag.py @@ -3,6 +3,7 @@ # pylint: disable=missing-function-docstring from pathlib import Path +from unittest.mock import AsyncMock import frontmatter import pytest @@ -55,6 +56,39 @@ def _file_store(*, name: str = "default", **tag_index_kwargs) -> LocalFileStore: return store +@pytest.mark.asyncio +async def test_auto_tag_reports_missing_agent_write_without_inventing_empty_tags(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + note = tmp_path / "memory/note.md" + _write_note(note) + original = note.read_bytes() + wrapper = _TaggingWrapper(tmp_path) + monkeypatch.setattr(wrapper, "reply", AsyncMock(return_value={"result": "", "last_message": {}})) + step = AutoTagStep(agent_wrapper=wrapper, file_store=_file_store()) + + response = await step(RuntimeContext(changes=[{"change": "added", "path": "memory/note.md"}])) + + assert response.metadata["auto_tag"]["failed"] == 1 + assert response.metadata["auto_tag"]["succeeded"] == 0 + assert "did not write" in response.metadata["auto_tag"]["results"][0]["error"] + assert note.read_bytes() == original + + +@pytest.mark.asyncio +async def test_auto_tag_accepts_explicitly_written_empty_tags(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + note = tmp_path / "memory/note.md" + _write_note(note) + wrapper = _TaggingWrapper(tmp_path, tags=[]) + step = AutoTagStep(agent_wrapper=wrapper, file_store=_file_store()) + + response = await step(RuntimeContext(changes=[{"change": "added", "path": "memory/note.md"}])) + + assert response.metadata["auto_tag"]["succeeded"] == 1 + assert response.metadata["auto_tag"]["failed"] == 0 + assert frontmatter.loads(note.read_text(encoding="utf-8")).metadata["memory_tags"] == [] + + @pytest.mark.asyncio async def test_auto_tag_handles_noop_invalid_changes_and_no_tag_index(tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) diff --git a/tests/unit/test_compact_context.py b/tests/unit/test_compact_context.py index 5d10fae0..3f48f2a4 100644 --- a/tests/unit/test_compact_context.py +++ b/tests/unit/test_compact_context.py @@ -14,8 +14,7 @@ from reme.components.as_decision_model import BaseAsDecisionModel from reme.components.job import BaseJob from reme.enumeration import ComponentEnum from reme.schema.compaction import CompactionOptions, CompactionRequest -from reme.steps.compaction import CompactContextStep -from reme.steps.compaction.messages import collect_pairs +from reme.steps.compaction import CompactContextStep, collect_pairs from reme.steps.compaction.state import build_batches, estimate_decision_tokens, serialize diff --git a/tests/unit/test_daily_list_range.py b/tests/unit/test_daily_list_range.py new file mode 100644 index 00000000..11c67687 --- /dev/null +++ b/tests/unit/test_daily_list_range.py @@ -0,0 +1,318 @@ +"""Date-range and pagination contracts for the read-only daily_list job.""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import pytest_asyncio + +from reme.components import ApplicationContext +from reme.components.file_store import LocalFileStore +from reme.steps.file_io.daily_list import DailyListStep + + +@pytest_asyncio.fixture(name="step") +async def daily_step(tmp_path, monkeypatch): + """Use a real file store in an isolated workspace, without model services.""" + monkeypatch.chdir(tmp_path) + store = LocalFileStore(name="daily", embedding_store="") + await store.start() + context = ApplicationContext(workspace_dir=str(tmp_path), timezone="UTC") + try: + yield DailyListStep(file_store=store, app_context=context) + finally: + await store.close() + + +def note(root: Path, day: str, name: str, *, daily_dir="daily", metadata="") -> Path: + """Create a note whose full frontmatter must survive projection.""" + path = root / daily_dir / day / f"{name}.md" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(f"---\nname: {name}\n{metadata}---\nbody\n", encoding="utf-8") + return path + + +@pytest.mark.asyncio +async def test_range_is_inclusive_newest_first_and_keeps_full_frontmatter(step, tmp_path): + """Range boundaries and page order must not overwrite user metadata.""" + for day, name in [("2024-02-28", "old"), ("2024-02-29", "b"), ("2024-02-29", "a"), ("2024-03-01", "new")]: + note(tmp_path, day, name, metadata="date: custom\nlabels: [one, two]\n") + await step(start_date="2024-02-29", end_date="2024-03-01", page_size=2) + response = step.context.response + assert response.success + assert response.metadata == { + "start_date": "2024-02-29", + "end_date": "2024-03-01", + "page": 1, + "page_size": 2, + "total_notes": 3, + "total_pages": 2, + "count": 2, + "notes": [ + {"path": "daily/2024-03-01/new.md", "name": "new", "date": "custom", "labels": ["one", "two"]}, + {"path": "daily/2024-02-29/a.md", "name": "a", "date": "custom", "labels": ["one", "two"]}, + ], + } + assert "daily/2024-03-01/new.md" in response.answer + assert "daily/2024-02-29/b.md" not in response.answer + + +@pytest.mark.asyncio +async def test_opt_in_pagination_preserves_unpaged_single_day_shape(step, tmp_path): + """Pagination is optional, stable and explicit about pages beyond the end.""" + for name in ("a", "b", "c"): + note(tmp_path, "2024-03-01", name) + await step(date="2024-03-01") + assert set(step.context.response.metadata) == {"date", "count", "notes"} + assert step.context.response.metadata["count"] == 3 + await step(date="2024-03-01", page=2, page_size=2) + result = step.context.response.metadata + assert result["count"] == 1 and result["total_notes"] == 3 + assert result["notes"][0]["name"] == "c" + await step(date="2024-03-01", page=3, page_size=2) + assert step.context.response.success + assert step.context.response.metadata["notes"] == [] + assert step.context.response.metadata["total_pages"] == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "params", + [ + {"start_date": "2024-03-01"}, + {"end_date": "2024-03-01"}, + {"start_date": None, "end_date": "2024-03-01"}, + {"start_date": "", "end_date": "2024-03-01"}, + {"start_date": "2024-W09-4", "end_date": "2024-03-01"}, + {"start_date": "2024-03-02", "end_date": "2024-03-01"}, + {"start_date": "2024-02-30", "end_date": "2024-03-01"}, + {"start_date": "2024-03-01T00:00:00", "end_date": "2024-03-02"}, + {"date": "2024-03-01", "start_date": "2024-03-01", "end_date": "2024-03-02"}, + {"page": 0}, + {"page_size": 1001}, + {"page_size": 0}, + {"page": "2"}, + {"page": None}, + {"page": True}, + {"page_size": False}, + {"page_size": 1.5}, + ], +) +async def test_invalid_requests_fail_instead_of_falling_back_to_today(step, params): + """Malformed or conflicting queries must not quietly list today's notes.""" + await step(**params) + assert not step.context.response.success + assert "Error:" in step.context.response.answer + + +@pytest.mark.asyncio +async def test_default_range_page_is_bounded_but_legacy_calls_return_every_note(step, tmp_path): + """Omitted pagination must not truncate existing single-day callers.""" + for index in range(101): + note(tmp_path, "2024-03-01", f"note-{index:03}") + await step(date="2024-03-01") + assert step.context.response.metadata["count"] == 101 + assert set(step.context.response.metadata) == {"date", "count", "notes"} + await step(start_date="2024-03-01", end_date="2024-03-01") + result = step.context.response.metadata + assert result["count"] == result["page_size"] == 100 + assert result["total_notes"] == 101 and result["total_pages"] == 2 + await step(start_date="2024-03-01", end_date="2024-03-01", page=2) + assert step.context.response.metadata["notes"][0]["name"] == "note-100" + + +@pytest.mark.asyncio +async def test_empty_range_is_read_only_and_has_no_pages(step, tmp_path): + """A query must not create missing date directories or indexes.""" + await step(start_date="0001-01-01", end_date="9999-12-31") + assert step.context.response.success + assert step.context.response.metadata["total_notes"] == 0 + assert step.context.response.metadata["total_pages"] == 0 + assert step.context.response.metadata["notes"] == [] + assert not (tmp_path / "daily").exists() + + +@pytest.mark.asyncio +async def test_range_uses_custom_directory_and_skips_indexes_invalid_dates_and_nested_files(step, tmp_path): + """Scan existing valid date directories, never every date or nested content.""" + step.app_context.app_config.daily_dir = "journal/days" + note(tmp_path, "2024-03-01", "kept", daily_dir="journal/days") + note(tmp_path, "not-a-date", "ignored", daily_dir="journal/days") + note(tmp_path, "2024-W09-5", "week-date", daily_dir="journal/days") + note(tmp_path, "2024-03-01/nested", "ignored", daily_dir="journal/days") + note(tmp_path, "2024-03-01", "other-root") + index = tmp_path / "journal/days/2024-03-01.md" + index.write_text("day index", encoding="utf-8") + (tmp_path / "journal/days/2024-03-02").write_text("not a directory", encoding="utf-8") + await step(start_date="0001-01-01", end_date="9999-12-31") + assert step.context.response.metadata["notes"] == [{"name": "kept", "path": "journal/days/2024-03-01/kept.md"}] + + +@pytest.mark.asyncio +async def test_permissions_filter_before_reading_and_before_counting(step, tmp_path, monkeypatch): + """An exact-file scope must hide other dates and sibling frontmatter.""" + allowed = note(tmp_path, "2024-03-01", "allowed") + denied = {note(tmp_path, "2024-03-01", "sibling"), note(tmp_path, "2024-02-29", "other-day")} + original = Path.read_text + + def guarded_read(path, *args, **kwargs): + if path in denied: + pytest.fail("permission-denied note was read") + return original(path, *args, **kwargs) + + monkeypatch.setattr(Path, "read_text", guarded_read) + for params in ({"date": "2024-03-01"}, {"start_date": "2024-02-29", "end_date": "2024-03-01"}): + await step(**params, _allowed_paths=[str(allowed.relative_to(tmp_path))]) + result = step.context.response.metadata + assert result["count"] == 1 + assert result["notes"] == [{"name": "allowed", "path": "daily/2024-03-01/allowed.md"}] + assert result.get("total_notes", 1) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", [[], {}, 42, ["../outside"], ["daily", "../outside"]]) +async def test_invalid_permission_constraints_fail_closed(step, tmp_path, scope): + """An invalid scope cannot become an unrestricted listing.""" + note(tmp_path, "2024-03-01", "hidden") + await step(start_date="2024-03-01", end_date="2024-03-01", _allowed_paths=scope) + assert step.context.response.metadata["notes"] == [] + assert step.context.response.metadata["total_notes"] == 0 + + +@pytest.mark.asyncio +async def test_directory_scope_counts_only_authorized_notes(step, tmp_path): + """Filtering happens before pagination and totals, including directory scopes.""" + note(tmp_path, "2024-03-01", "a") + note(tmp_path, "2024-03-01", "b") + note(tmp_path, "2024-02-29", "hidden") + await step(start_date="2024-02-29", end_date="2024-03-01", page_size=1, _allowed_paths=["daily/2024-03-01"]) + result = step.context.response.metadata + assert result["total_notes"] == result["total_pages"] == 2 + assert result["notes"][0]["name"] == "a" + + +@pytest.mark.asyncio +async def test_symlink_escapes_and_loops_do_not_leak_notes(step, tmp_path): + """Both date-directory links and note links retain workspace containment.""" + outside = tmp_path.parent / f"{tmp_path.name}-outside" + outside.mkdir() + (outside / "secret.md").write_text("---\nname: secret\n---\nprivate", encoding="utf-8") + kept = note(tmp_path, "2024-03-01", "kept") + try: + (kept.parent / "escape.md").symlink_to(outside / "secret.md") + (tmp_path / "daily/2024-02-29").symlink_to(outside, target_is_directory=True) + (tmp_path / "daily/2024-03-02").symlink_to("2024-03-02", target_is_directory=True) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + await step(start_date="2024-02-29", end_date="2024-03-02") + assert step.context.response.metadata["notes"] == [{"name": "kept", "path": "daily/2024-03-01/kept.md"}] + + +@pytest.mark.asyncio +async def test_queries_preserve_note_and_index_contents_and_mtimes(step, tmp_path): + """Listing never reindexes, rewrites frontmatter, or touches user files.""" + note(tmp_path, "2024-03-01", "a") + note(tmp_path, "2024-02-29", "b") + index = tmp_path / "daily/2024-03-01.md" + index.write_text("handwritten index\n", encoding="utf-8") + before = {p: (p.read_bytes(), p.stat().st_mtime_ns) for p in (tmp_path / "daily").rglob("*") if p.is_file()} + await step(start_date="2024-02-29", end_date="2024-03-01", page_size=1) + await step(date="2024-03-01", page=2, page_size=1) + after = {p: (p.read_bytes(), p.stat().st_mtime_ns) for p in (tmp_path / "daily").rglob("*") if p.is_file()} + assert before == after + + +@pytest.mark.asyncio +async def test_invalid_or_unreadable_notes_do_not_count_toward_pages(step, tmp_path, monkeypatch): + """Match the existing scanner's skip behavior for malformed and unreadable notes.""" + good = note(tmp_path, "2024-03-01", "good") + malformed = note(tmp_path, "2024-03-01", "bad") + malformed.write_text("---\nlabels: [\n---\n", encoding="utf-8") + unreadable = note(tmp_path, "2024-03-01", "unreadable") + original = Path.read_text + + def guarded_read(path, *args, **kwargs): + if path == unreadable: + raise PermissionError("unreadable note") + return original(path, *args, **kwargs) + + monkeypatch.setattr(Path, "read_text", guarded_read) + await step(start_date="2024-03-01", end_date="2024-03-01", page_size=1) + assert step.context.response.metadata["total_notes"] == 1 + assert step.context.response.metadata["notes"][0]["path"] == good.relative_to(tmp_path).as_posix() + + +@pytest.mark.asyncio +async def test_scanning_stays_off_the_event_loop(step, tmp_path, monkeypatch): + """File parsing uses the event loop's managed executor, not the caller's thread.""" + from threading import current_thread + import frontmatter + + note(tmp_path, "2024-03-01", "a") + original = frontmatter.loads + caller = current_thread().ident + threads = [] + + def capture_thread(*args, **kwargs): + threads.append(current_thread().ident) + return original(*args, **kwargs) + + monkeypatch.setattr(frontmatter, "loads", capture_thread) + await step(start_date="2024-03-01", end_date="2024-03-01") + assert threads and all(thread != caller for thread in threads) + + +@pytest.mark.asyncio +async def test_registered_job_and_http_endpoint_use_the_updated_schema(step, tmp_path): + """Exercise the configured Job and HTTP boundary rather than only a Step helper.""" + import httpx + import yaml + + from reme.components.job import BaseJob + from reme.components.service import HttpService, MCPService + from reme.enumeration import ComponentEnum + + default = yaml.safe_load((Path(__file__).parents[2] / "reme/config/default.yaml").read_text(encoding="utf-8")) + spec = default["jobs"]["daily_list"] + properties = spec["parameters"]["properties"] + assert {"date", "start_date", "end_date", "page", "page_size"} == set(properties) + assert "default" not in properties["page"] and "default" not in properties["page_size"] + assert properties["page_size"]["maximum"] == 1000 + note(tmp_path, "2024-02-29", "old") + note(tmp_path, "2024-03-01", "new") + context = step.app_context + context.components[ComponentEnum.FILE_STORE] = {"default": step.file_store} + job = BaseJob(name="daily_list", app_context=context, **spec) + await job.start() + try: + response = await job(start_date="2024-02-29", end_date="2024-03-01", page_size=1) + assert response.success and response.metadata["notes"][0]["name"] == "new" + service = HttpService(mcp_enabled=False, web_enabled=False) + service.build_service(SimpleNamespace(config=context.app_config)) + service.add_job(job) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=service.service), + base_url="http://test", + ) as client: + response = await client.post( + "/daily_list", + json={"start_date": "2024-02-29", "end_date": "2024-03-01", "page_size": 1}, + ) + assert response.status_code == 200 + assert response.json()["metadata"]["total_notes"] == 2 + response = await client.post("/daily_list", json={"date": "2024-03-01"}) + assert set(response.json()["metadata"]) == {"date", "count", "notes"} + response = await client.post("/daily_list", json={"start_date": "2024-03-01"}) + assert not response.json()["success"] + mcp = MCPService() + mcp.build_service(SimpleNamespace(config=context.app_config)) + mcp.add_job(job) + tool = await mcp.service.get_tool("daily_list") + result = await tool.run({"start_date": "2024-02-29", "end_date": "2024-03-01", "page_size": 1}) + assert "total_notes: 2" in str(result.content) + assert "daily/2024-03-01/new.md" in str(result.content) + assert "daily/2024-02-29/old.md" not in str(result.content) + legacy = await tool.run({"date": "2024-03-01"}) + assert "total_notes:" not in str(legacy.content) + finally: + await job.close() diff --git a/tests/unit/test_decision_model.py b/tests/unit/test_decision_model.py index 1a4c3beb..b7416127 100644 --- a/tests/unit/test_decision_model.py +++ b/tests/unit/test_decision_model.py @@ -4,6 +4,7 @@ import asyncio import sys +import threading from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, Mock @@ -101,6 +102,13 @@ def test_valid_questions_reuse_model_and_restart_closes_clients(provider, provid } async def run(): + loop_thread = threading.get_ident() + + def construct(**_kwargs): + assert threading.get_ident() != loop_thread + return model + + factory.side_effect = construct component = AgentScopeAsDecisionModel( credential={"api_key": "test-key", "base_url": ""}, model="decision-test", @@ -138,6 +146,53 @@ def test_valid_questions_reuse_model_and_restart_closes_clients(provider, provid asyncio.run(run()) +@pytest.mark.parametrize("construction_fails", [False, True]) +def test_cancelled_construction_keeps_ownership_until_close(provider, construction_fails): + """Cancellation cannot abandon a client or let shutdown overtake its worker.""" + factory, model = provider + + async def run(): + loop = asyncio.get_running_loop() + entered, release = asyncio.Event(), threading.Event() + + def construct(**_kwargs): + loop.call_soon_threadsafe(entered.set) + assert release.wait(5), "Construction was not released" + if construction_fails: + raise RuntimeError("construction failed") + return model + + factory.side_effect = construct + component = AgentScopeAsDecisionModel(credential={"api_key": "test-key"}) + call = asyncio.create_task(component(state="hello", questions={"q": {"type": "binary"}})) + closing = None + try: + await asyncio.wait_for(entered.wait(), 2) + with pytest.raises(RuntimeError, match="Close the decision component"): + component.model = AsyncMock() + for _ in range(2): + call.cancel() + await asyncio.sleep(0) + assert not call.done() + closing = asyncio.create_task(component.close()) + await asyncio.sleep(0) + assert not closing.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(call, 2) + await asyncio.wait_for(closing, 2) + assert component.model is None + assert not component._owns_model + model.assert_not_awaited() + assert model.client.aclose.await_count == (0 if construction_fails else 1) + finally: + release.set() + await asyncio.gather(call, *([closing] if closing is not None else []), return_exceptions=True) + await component.close() + + asyncio.run(run()) + + @pytest.mark.parametrize( "question", [ diff --git a/tests/unit/test_markdown_file_chunker.py b/tests/unit/test_markdown_file_chunker.py index 3b18601b..73926bf5 100644 --- a/tests/unit/test_markdown_file_chunker.py +++ b/tests/unit/test_markdown_file_chunker.py @@ -28,6 +28,28 @@ def _write_md(tmpdir: str, name: str, body: str) -> str: return name +@pytest.mark.asyncio +async def test_chunk_keeps_pre_read_mtime_when_source_changes(tmp_path, monkeypatch): + """A concurrent write must remain visible to the next filesystem scan.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / "changing.md" + path.write_text("", encoding="utf-8") + os.utime(path, (1000, 1000)) + chunker = MarkdownFileChunker() + + async def read_then_finish_write(_path): + path.write_text("Completed memory note.", encoding="utf-8") + os.utime(path, (2000, 2000)) + return "" + + monkeypatch.setattr(chunker, "_read_text_for_indexing", read_then_finish_write) + node, chunks = await chunker.chunk("changing.md") + + assert chunks == [] + assert node.st_mtime == 1000 + assert path.stat().st_mtime != node.st_mtime + + @pytest.mark.asyncio async def test_parse_empty_file(): """An empty .md → FileNode, no chunks, no links.""" diff --git a/tests/unit/test_memlens_benchmark.py b/tests/unit/test_memlens_benchmark.py new file mode 100644 index 00000000..1a461ca5 --- /dev/null +++ b/tests/unit/test_memlens_benchmark.py @@ -0,0 +1,335 @@ +"""Focused offline benchmark contracts; no external services or experiment writes.""" + +# pylint: disable=missing-function-docstring,protected-access +import asyncio +import base64 +from importlib.metadata import EntryPoint +from io import BytesIO +import json +from pathlib import Path +import tomllib +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from agentscope.formatter import OpenAIChatFormatter +from agentscope.message import AssistantMsg, Base64Source, DataBlock, Msg, TextBlock, ToolResultBlock, ToolResultState +from PIL import Image +import pytest + +from reme.components import R +from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper +from reme.components.component_registry import create_application_registry +from reme.components.file_store import LocalFileStore +from reme.plugin import PluginManager, _load_plugin +from reme.schema import ApplicationConfig, Response +from reme.steps.evolve.auto_memory import _sanitize_msg_for_save +from reme.steps.file_io.read_image import ReadImageStep + +from benchmark.memlens import run as memlens + + +@pytest.fixture(autouse=True) +def memlens_judge_path(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins/memlens-judge/src")) + + +def test_candidate_declares_required_image_dependency(): + manifest = Path(__file__).resolve().parents[2] / "plugins/memlens/pyproject.toml" + project = tomllib.loads(manifest.read_text(encoding="utf-8"))["project"] + assert "pillow>=10.0.0" in project["dependencies"] + + +def test_memlens_subset_lengths_order_and_public_inputs(tmp_path): + image = {"file": "photo.png", "blip_caption": "HIDDEN_CAPTION"} + row = { + "question_type": "knowledge_update", + "question_subtype": "update", + "question": "What changed?", + "answer": "HIDDEN_GOLD", + "old_answer": "HIDDEN_OLD", + "question_date": "2024/05/31 (Fri) 07:58", + "haystack_dates": ["2024/05/01 (Wed) 08:00"], + "haystack_session_ids": ["original-session"], + "haystack_sessions": [ + [ + {"role": "user", "content": "Before after", "images": [image]}, + {"role": "assistant", "content": "Original reply", "images": [], "has_answer": True}, + ], + ], + "answer_session_ids": ["HIDDEN_EVIDENCE"], + } + rows = [{**row, "question_id": f"q_{index:03d}"} for index in range(196)] + ids = [item["question_id"] for item in rows[:195]] + (tmp_path / "agent_subset_195.json").write_text(json.dumps({"n_questions": 195, "question_ids": ids[::-1]})) + for length in ("128k", "256k"): + (tmp_path / f"dataset_{length}.json").write_text( + json.dumps([{**item, "haystack_session_ids": [f"session-{length}"]} for item in rows]), + ) + + image_dir = tmp_path / "release_images" + image_dir.mkdir() + Image.new("RGB", (1, 1), (12, 34, 56)).save(image_dir / "photo.png") + scenes = memlens.load_memlens(tmp_path, "256k") + assert [scene["id"] for scene in scenes] == ids + first = scenes[0] + assert first["sessions"][0]["id"] == "session-256k" + assert memlens.load_memlens(tmp_path, "128k")[0]["sessions"][0]["id"] == "session-128k" + assert len(first["questions"]) == 1 and first["questions"][0]["id"] == first["id"] + assert first["sessions"][0]["timestamp"] == "2024-05-01T08:00" + messages = memlens.session_messages(first["sessions"][0]) + assert [message["role"] for message in messages] == ["user", "assistant"] + assert [block["type"] for block in messages[0]["content"]] == ["text", "text", "data", "text"] + assert messages[0]["content"][0]["text"] == "Before " + assert messages[0]["content"][1]["text"] == "[Image ID: photo]" + assert messages[0]["content"][-1]["text"] == " after" + assert base64.b64decode(messages[0]["content"][2]["source"]["data"]) == (image_dir / "photo.png").read_bytes() + assert memlens.session_messages(scenes[1]["sessions"][0])[0]["id"] != messages[0]["id"] + public = json.dumps(messages) + memlens.query_prompt(first["questions"][0], "Evidence tools") + assert "HIDDEN" not in public and "has_answer" not in public + rows[0]["haystack_sessions"][0][0]["images"][0]["file"] = "../private.png" + (tmp_path / "dataset_256k.json").write_text(json.dumps(rows)) + with pytest.raises(ValueError, match="MemLens"): + memlens.load_memlens(tmp_path, "256k") + + +def test_official_metric_normalization(): + question = {"answer": "$460.00"} + assert memlens.score_answer(question, "ignore$460.00<|im_end|>") == {"sub_em": 1.0, "f1": 1.0} + + +@pytest.fixture(name="candidate") +def fixture_candidate(tmp_path, monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins/memlens/src")) + monkeypatch.chdir(tmp_path) + from reme_memlens import agentic_answer as answer, auto_memory as memory + + app = SimpleNamespace( + registry=R, + metadata={}, + app_config=ApplicationConfig(workspace_dir=str(tmp_path)), + jobs={}, + components={}, + ) + wrapper = AsAgentWrapper(backend="agentscope", as_llm="", app_context=app) + wrapper.reply = AsyncMock(return_value={"result": "answer", "session_id": "new", "usage": {"input_tokens": 7}}) + store = LocalFileStore(embedding_store="", app_context=app) + return memory, answer, app, wrapper, store, tmp_path + + +def test_candidate_registration_is_application_local(candidate): + memory, answer, *_ = candidate + plugin = _load_plugin("memlens", EntryPoint(name="memlens", value="reme_memlens", group="reme.plugins")) + registry = create_application_registry() + PluginManager([plugin]).register(registry) + assert registry.get("step", "memlens_auto_memory_step") is memory.MemLensAutoMemoryStep + assert registry.get("step", "memlens_agentic_answer_step") is answer.MemLensAgenticAnswerStep + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sources", [True, False]) +async def test_memory_saves_originals_and_native_image_links(candidate, sources): + memory, _, app, wrapper, store, workspace = candidate + step = memory.MemLensAutoMemoryStep(app_context=app, agent_wrapper=wrapper, file_store=store) + step._list_session_note = AsyncMock(return_value=None) + picture = BytesIO() + Image.new("RGB", (1, 1), (12, 34, 56)).save(picture, format="PNG") + image_bytes = picture.getvalue() + image_data = base64.b64encode(image_bytes).decode() + late = Msg.model_validate( + { + "id": "late", + "name": "Alice", + "role": "user", + "created_at": "2026-01-02T10:00:00", + "content": [ + TextBlock(text="Image ID: image-7"), + DataBlock(source=Base64Source(media_type="image/png", data=image_data)), + TextBlock(text="Before.\u2028\u2028After.\u2029End.\u0085Done."), + ], + }, + ) + early = late.model_copy(update={"id": "early", "created_at": "2026-01-01T10:00:00"}, deep=True) + before = late.model_dump() + await step(messages=[late, early], session_id="one", include_images=True, include_source_sessions=sources) + saved = workspace / "session/dialog/one.jsonl" + assert saved.read_text() == "".join(_sanitize_msg_for_save(msg).model_dump_json() + "\n" for msg in [early, late]) + inputs = wrapper.reply.call_args.args[0] + assert [block.type for block in inputs.content] == ["text", "data", "text", "data", "text"] + assert [block.source.data for block in inputs.content if isinstance(block, DataBlock)] == [image_data] * 2 + text = inputs.get_text_content() + assert "Image ID: image-7" in text and "Before.\u2028\u2028After.\u2029End.\u0085Done." in text + assert ("[L2 | Alice" in text) is sources + files = sorted((workspace / "session/images/one").glob("*")) + assert len(files) == 2 and all(path.read_bytes() == image_bytes for path in files) + assert all(f"[[{path.relative_to(workspace).as_posix()}]]" in text for path in files) + assert late.model_dump() == before + + +@pytest.mark.asyncio +async def test_read_image_returns_real_formatter_image_and_rejects_escape(candidate): + _, answer, app, wrapper, store, workspace = candidate + path = workspace / "session/images/one/photo.png" + path.parent.mkdir(parents=True) + Image.new("RGB", (2, 2), (12, 34, 56)).save(path) + + async def native_read(**kwargs): + step = ReadImageStep(app_context=app) + await step(**kwargs) + return step.context.response + + app.jobs["read_image"] = AsyncMock(side_effect=native_read) + step = answer.MemLensAgenticAnswerStep(app_context=app, agent_wrapper=wrapper, file_store=store) + tool = answer.image_reader(step) + result = await tool.call(path="session/images/one/photo.png") + assert result.state == ToolResultState.SUCCESS and isinstance(result.content[1], DataBlock) + assert base64.b64decode(result.content[1].source.data) == path.read_bytes() + formatted = await OpenAIChatFormatter().format( + [ + AssistantMsg(name="tool", content=[ToolResultBlock(id="call", name="read_image", output=result.content)]), + ], + ) + assert any( + block["type"] == "image_url" + for msg in formatted + if isinstance(msg.get("content"), list) + for block in msg["content"] + ) + assert (await tool.call(path="../other-scene.png")).state == ToolResultState.ERROR + assert (await tool.call(path="mem_session/private.png")).state == ToolResultState.ERROR + app.jobs["read_image"].assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sources", [True, False]) +async def test_answer_has_fresh_state_bm25_scope_and_no_gold(candidate, sources): + _, answer, app, wrapper, store, workspace = candidate + for relative in ["daily/note.md", "digest/note.md", "session/dialog/one.jsonl", "mem_session/private.md"]: + path = workspace / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("evidence") + assert answer.evidence_paths(workspace, "daily", "digest", "session", False) == ["daily/note.md", "digest/note.md"] + calls = [] + + async def search(**kwargs): + calls.append(kwargs) + return Response(answer="evidence") + + job = AsyncMock(side_effect=search) + job.name, job.description, job.parameters = ( + "search", + "search", + { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + ) + read = AsyncMock(return_value=Response(answer="evidence")) + read.name, read.description, read.parameters = "read", "read", {"type": "object", "properties": {}} + app.jobs = {"search": job, "read": read} + step = answer.MemLensAgenticAnswerStep(app_context=app, agent_wrapper=wrapper, file_store=store) + sessions, contexts = [], [] + for _ in range(2): + response = await step(query="Official question", include_source_sessions=sources) + options = wrapper.reply.call_args.kwargs + assert wrapper.reply.call_args.args[0].get_text_content() == "Official question" + assert options["resume"] == "" and options["builtin_tools"] is False and options["job_tools"] == [] + assert options["react_config"] == {"max_iters": 30} + sessions.append(options["session_id"]) + await (await options["toolkit"].get_tool("search")).call(query="evidence") + assert calls[-1]["vector_weight"] == 0.0 + assert ("session/dialog/one.jsonl" in calls[-1]["search_filter"]["paths"]) is sources + assert "mem_session/private.md" not in calls[-1]["_allowed_paths"] + contexts.append(calls[-1]["tool_context_id"]) + assert response.metadata["usage"] == {"input_tokens": 7} + assert len(set(sessions)) == len(set(contexts)) == 2 + + +@pytest.mark.asyncio +async def test_image_only_pipeline_preserves_sessions_and_official_query(candidate, monkeypatch): + import reme + from benchmark.memlens.run import normalized_usage, total_usage + + runner = memlens + *_, workspace = candidate + + query = { + "id": "q", + "text": "What changed?", + "category": "knowledge_update", + "question_date": "2024/05/31 (Fri) 07:58", + "answer": "HIDDEN_GOLD", + "subtype": "update", + "old_answer": "HIDDEN_OLD", + } + scene = { + "id": "one", + "questions": [query], + "sessions": [{"id": name, "timestamp": "2024-05-01T08:00"} for name in ("first", "second")], + } + calls = [] + + class Candidate: + """Offline native-Job boundary; no provider calls.""" + + def __init__(self, **_config): + self.context = SimpleNamespace( + components={ + "agent_wrapper": {name: SimpleNamespace(session_path=workspace) for name in ("default", "bench")}, + }, + ) + self.closed = False + + async def start(self): + pass + + async def close(self): + self.closed = True + + async def run_job(self, name, **kwargs): + calls.append((name, kwargs)) + return Response(answer="Retained answer", metadata={"agent_session_id": "new"}) + + app = Candidate() + judge = SimpleNamespace(run_job=AsyncMock(return_value=Response(metadata={"llm_judge_score": 1}))) + monkeypatch.setattr(reme, "Application", lambda **_: app) + monkeypatch.setattr(runner, "app_config", lambda *_, **__: {}) + monkeypatch.setattr(runner, "session_messages", lambda session: [{"text": session["id"]}]) + monkeypatch.setattr(runner, "observe_usage", lambda _: None) + rows, *_ = await runner.evaluate_scene(scene, 0, "candidate", workspace, judge, asyncio.Semaphore(2)) + assert [name for name, _ in calls] == ["auto_memory", "auto_memory", "index_update", "agentic_answer"] + assert [kw["session_id"] for name, kw in calls if name == "auto_memory"] == ["first", "second"] + for name, kwargs in calls: + assert "HIDDEN" not in json.dumps(kwargs) + if name != "index_update": + assert kwargs["include_images"] is kwargs["include_source_sessions"] is True + inputs = calls[-1][1] + assert inputs["max_iters"] == 30 and "category" not in inputs + assert inputs["query"] == memlens.query_prompt(query, "Use the available tools to retrieve relevant evidence.") + assert judge.run_job.call_args.kwargs["reference"] == "HIDDEN_GOLD" + assert rows[0]["metrics"]["llm_judge"] == 1 and app.closed + measured = normalized_usage( + { + "prompt_tokens": 100, + "completion_tokens": 20, + "prompt_tokens_details": {"cached_tokens": 30}, + }, + ) + assert measured["noncached_input_tokens"] == 70 + assert total_usage([measured])["total_tokens"] == 120 + assert normalized_usage({})["cached_input_tokens"] is None + + +@pytest.mark.asyncio +async def test_observation_preserves_native_reply_and_rejects_only_interrupted(): + from benchmark.memlens.run import observe_usage + + result = {"result": "", "last_message": {"finished_reason": "exceed_max_iters", "error": None}} + reply = AsyncMock(return_value=result) + wrapper = SimpleNamespace(reply=reply) + app = SimpleNamespace(context=SimpleNamespace(components={"agent_wrapper": {"default": wrapper}})) + observe_usage(app) + assert await wrapper.reply("Original input", react_config={"max_iters": 30}) is result + reply.assert_awaited_once_with("Original input", react_config={"max_iters": 30}) + reply.return_value = {"last_message": {"finished_reason": "interrupted"}} + with pytest.raises(RuntimeError, match="interrupted"): + await wrapper.reply("Original input") diff --git a/tests/unit/test_memlens_judge.py b/tests/unit/test_memlens_judge.py new file mode 100644 index 00000000..3d02c74e --- /dev/null +++ b/tests/unit/test_memlens_judge.py @@ -0,0 +1,101 @@ +"""Standalone judge Jobs preserve official prompts, parsers and generation defaults.""" + +# pylint: disable=missing-function-docstring,protected-access +import hashlib +from pathlib import Path +from unittest.mock import AsyncMock + +from agentscope.message import TextBlock +from agentscope.model import ChatResponse, ChatUsage +import pytest +import yaml + +from reme.components.runtime_context import RuntimeContext +from reme.plugin_manifest import parse_plugin_manifest + + +@pytest.fixture(name="judge") +def fixture_judge(monkeypatch): + root = Path(__file__).resolve().parents[2] / "plugins/memlens-judge/src" + monkeypatch.syspath_prepend(str(root)) + from judge_memlens import protocol + from judge_memlens.llm_judge import MemLensJudgeStep + + step = MemLensJudgeStep() + step.context = RuntimeContext.from_context( + None, + question="QUESTION", + reference="REFERENCE", + prediction="PREDICTION", + evidence=["first", "second"], + category="information_extraction", + ) + step.as_llm = AsyncMock( + return_value=ChatResponse( + content=[TextBlock(text='{"answer_score":1}')], + is_last=True, + usage=ChatUsage(100, 10, 0.01, cache_input_tokens=30), + ), + ) + return step, protocol, root + + +@pytest.mark.asyncio +async def test_judge_job_uses_only_official_prompt_and_separate_usage(judge): + step, protocol, root = judge + response = await step.execute() + assert response.success and response.metadata["llm_judge_score"] == 1 + assert response.metadata["usage"]["cache_input_tokens"] == 30 + kwargs = {"category": "information_extraction"} + expected = protocol.judge_prompt("QUESTION", "REFERENCE", "PREDICTION", **kwargs) + assert step.as_llm.call_args.args[0][0].get_text_content() == expected + assert len(step.as_llm.call_args.args[0]) == 1 and not step.as_llm.call_args.kwargs + text = (root / "judge_memlens/plugin.yaml").read_text() + parse_plugin_manifest(text, plugin_name="memlens-judge") + manifest = yaml.safe_load(text)["application_defaults"] + assert manifest["jobs"]["answer_judge"]["steps"] == [{"backend": "memlens_judge_step", "as_llm": "judge"}] + model = manifest["components"]["as_llm"]["judge"] + assert model["extra_body"] == {"enable_thinking": False} and model["parameters"]["temperature"] == 0 + assert model["parameters"]["max_tokens"] == 256 + step.as_llm.return_value = ChatResponse(content=[TextBlock(text="malformed")], is_last=True) + assert (await step.execute()).metadata["parse_default_zero"] is True + step.context.data["prediction"] = "word " * 501 + assert (await step.execute()).metadata["auto_zero"] is True + assert step.as_llm.await_count == 2 + step.context.data["prediction"] = "PREDICTION" + step.as_llm.side_effect = TimeoutError("provider") + with pytest.raises(TimeoutError): + await step.execute() + + +def test_memlens_official_registry_digests_and_routing(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins/memlens-judge/src")) + from judge_memlens import protocol + + # Independent upstream build_judge_prompt digests, not generated by this renderer. + digests = { + "IE_Entity": "83991403e50c7390f9bafd957618a9ad01bdde2ae6d26ada38ce9954c697ebb0", + "IE_PreviousInfo": "18082c484f096bef129b6007671f7cde3e9c44b9e1e5e781630b84a604c126f7", + "KU_KnowledgeUpdate": "ccd52d849c6eb132e819f1f2847faf8a3c1755e517fd7f3e17308d8b51798951", + "TR_DurationComparison": "b884e6688d2b195c63d25c40c9003d98a5c8915b8b008613b51b14f95cc1c733", + "TR_OrderRanking": "c57f6b289c986aef4a1c6032c0d5544a7980001354298aaa9b3546c0edb4f86b", + "TR_DateExtraction": "8c10197630ba720759a6e04a1d670e5af9781284c509d53f89310529ad8d17a5", + "MSR_YesNo": "79779f253bfee20cfa9666ff33c6cdcc37828b9d17f8c83f92af4ec7e2e557ff", + "MSR_Counting": "bde797036e6069ac10a649ae80fc6749facd4c1ef5e912f39418fcfe729752a7", + "MSR_Arithmetic": "089db6580a3a640a5cb925b490e9cd059a71b65d4d85b236b748ee17cb1766af", + "AR_AnswerRefusal": "c7a6702489c079cfa4c354febb916078825325a0396261627b2c0295aab0d018", + } + prompts = protocol._prompts() + assert prompts["source_revision"] == protocol.SOURCE_REVISION + assert set(prompts["criteria"]) == set(digests) + for key, digest in digests.items(): + base = prompts["template"].format(task_criteria=prompts["criteria"][key], examples=prompts["examples"][key]) + reference = "Insufficient information" if key == "AR_AnswerRefusal" else "REFERENCE" + old = ": OLD\n" if key == "KU_KnowledgeUpdate" else "" + rendered = ( + f"{base}[Current Case]\n: QUESTION\n: {reference}\n" + f"{old}: PREDICTION\n\n[Scoring Rationale]:" + ) + assert hashlib.sha256(rendered.encode()).hexdigest() == digest + assert protocol.task_key("multi_session_reasoning", "entity_resolution", "2") == "MSR_Counting" + assert protocol.parse_judge_response('{"answer_score":0} [Score]:1')["acc"] == 0 diff --git a/tests/unit/test_memlens_runner_failures.py b/tests/unit/test_memlens_runner_failures.py new file mode 100644 index 00000000..095278ff --- /dev/null +++ b/tests/unit/test_memlens_runner_failures.py @@ -0,0 +1,423 @@ +"""Offline failure, accounting and artifact contracts for the MemLens runner.""" + +# pylint: disable=missing-function-docstring,protected-access +import asyncio +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from benchmark.memlens import run as runner +from benchmark.memlens.artifacts import error_details, judge_settings, model_settings, response_record, safe_value +from reme.schema import Response + + +@pytest.fixture(autouse=True) +def memlens_judge_path(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "plugins/memlens-judge/src")) + + +@pytest.fixture(name="scene_run") +def fixture_scene_run(tmp_path, monkeypatch): + import reme + + scene = { + "id": "scene-1", + "sessions": [{"id": "session-1", "timestamp": "2026-01-01"}], + "questions": [ + { + "id": "question-1", + "text": "Question", + "question_date": "2026-01-02", + "answer": "Answer", + "category": "information_extraction", + }, + ], + } + app = SimpleNamespace(start=AsyncMock(), close=AsyncMock(), run_job=AsyncMock()) + judge = SimpleNamespace(run_job=AsyncMock(return_value=Response(metadata={"llm_judge_score": 1}))) + monkeypatch.setattr(reme, "Application", lambda **_: app) + monkeypatch.setattr(runner, "app_config", lambda *_, **__: {}) + monkeypatch.setattr(runner, "observe_usage", lambda _: None) + monkeypatch.setattr(runner, "session_messages", lambda _: [{"role": "user", "content": "Text"}]) + case = tmp_path / "reme_image/scenes/0000-scene-1" + + async def evaluate(): + return await runner.evaluate_scene(scene, 0, "candidate", tmp_path, judge, asyncio.Semaphore(1)) + + return scene, app, judge, case, evaluate + + +def measured(tokens=120): + return runner.normalized_usage({"input_tokens": tokens - 20, "output_tokens": 20}) + + +@pytest.mark.asyncio +async def test_index_failure_retains_completed_session_usage(scene_run): + _, app, _, case, evaluate = scene_run + + async def job(name, **_): + if name == "auto_memory": + runner._CURRENT.get().append(measured()) + return Response(answer="Memory") + return Response(success=False, answer="Index unavailable") + + app.run_job.side_effect = job + results, memory, answers, status = await evaluate() + assert results == answers == [] + assert runner.total_usage(memory)["total_tokens"] == 120 + assert json.loads((case / "status.json").read_text()) == status + assert status["stage"] == "index_update" and status["scene_id"] == "scene-1" + assert "response detail: Index unavailable" in status["error"]["message"] + assert status["error_type"] == "RuntimeError" + assert status["memory_complete"] is False + assert len(list((case / "ingestion/attempt-001/sessions").glob("*.json"))) == 1 + artifact = json.loads((case / "ingestion/attempt-001/diagnostics/index-response.json").read_text()) + assert artifact["response"]["success"] is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct", [False, True]) +async def test_failed_memory_keeps_its_successful_requests_and_diagnostics(scene_run, direct): + scene, app, _, case, evaluate = scene_run + scene["sessions"].append({"id": "session-2", "timestamp": "2026-01-02"}) + + async def job(*_, session_id, **__): + runner._CURRENT.get().append(measured()) + if session_id == "session-1": + return Response(answer="Memory") + if direct: + try: + raise ValueError("Original provider failure") + except ValueError as exc: + raise RuntimeError("Memory failed") from exc + return Response(success=False, metadata={"error": "Memory failed"}) + + app.run_job.side_effect = job + _, memory, _, _ = await evaluate() + assert runner.total_usage(memory)["total_tokens"] == 240 + status = json.loads((case / "status.json").read_text()) + assert status["stage"] == "auto_memory" and status["session_id"] == "session-2" + assert status["memory_complete"] is False + assert (case / "ingestion/attempt-001/sessions/00000.json").exists() + assert not (case / "ingestion/attempt-001/sessions/00001.json").exists() + assert json.loads((case / "ingestion/attempt-001/diagnostics/00001-usage.json").read_text())["total_tokens"] == 120 + if direct: + assert status["error"]["cause"]["message"] == "Original provider failure" + assert status["error"]["traceback"] and "locals" not in json.dumps(status) + else: + artifact = json.loads((case / "ingestion/attempt-001/diagnostics/00001-response.json").read_text()) + assert artifact["response"]["success"] is False + app.close.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failed_job", ["agentic_answer", "judge"]) +async def test_failed_answer_or_judge_keeps_answer_usage_once(scene_run, failed_job): + _, app, judge, case, evaluate = scene_run + + async def job(name, **_): + if name in {"auto_memory", "agentic_answer"}: + runner._CURRENT.get().append(measured()) + return Response(answer="Answer", success=name != failed_job) + + app.run_job.side_effect = job + if failed_job == "judge": + judge.run_job.return_value = Response(success=False, metadata={"error": "Judge failed"}) + results, memory, answers, _ = await evaluate() + assert results == [] and len(memory) == len(answers) == 1 + assert runner.total_usage(answers)["total_tokens"] == 120 + status = json.loads((case / "status.json").read_text()) + assert status["question_id"] == "question-1" and status["stage"] == failed_job + assert status["memory_complete"] is True + path = "judge-response.json" if failed_job == "judge" else "answer-response.json" + assert json.loads((case / "questions/0000/diagnostics" / path).read_text())["response"]["success"] is False + assert (case / "questions/0000/answer.json").exists() is (failed_job == "judge") + assert json.loads((case / "questions/0000/diagnostics/answer-usage.json").read_text())["total_tokens"] == 120 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ["preprocess", "start", "cleanup", "cancel"]) +async def test_direct_lifecycle_failures_are_recorded(scene_run, monkeypatch, stage): + _, app, _, case, evaluate = scene_run + app.run_job.return_value = Response(answer="Answer") + if stage == "start": + app.start.side_effect = RuntimeError("startup failed") + elif stage == "cleanup": + app.close.side_effect = RuntimeError("cleanup failed") + elif stage == "cancel": + app.run_job.side_effect = asyncio.CancelledError() + else: + + def broken(_): + raise ValueError("Invalid image: release_images/missing.png") + + monkeypatch.setattr(runner, "session_messages", broken) + if stage == "cancel": + with pytest.raises(asyncio.CancelledError): + await evaluate() + else: + await evaluate() + status = json.loads((case / "status.json").read_text()) + assert status["complete"] is False + assert status["stage"] == ("auto_memory" if stage == "cancel" else stage) + assert status["error"]["type"] == ( + "CancelledError" if stage == "cancel" else "RuntimeError" if stage in {"start", "cleanup"} else "ValueError" + ) + if stage == "preprocess": + assert status["session_id"] == "session-1" and "missing.png" in status["error"]["message"] + assert status["error"]["traceback"] + + +@pytest.mark.asyncio +async def test_observer_separates_concurrent_scopes_including_compression_and_unknown_counters(): + def response(usage, status=200): + return SimpleNamespace(status_code=status, is_stream_consumed=True, json=lambda: {"usage": usage}) + + client = SimpleNamespace(send=AsyncMock(side_effect=response)) + compressor = SimpleNamespace(send=AsyncMock(side_effect=response)) + app = SimpleNamespace( + context=SimpleNamespace( + components={ + "as_llm": { + "default": SimpleNamespace(model=SimpleNamespace(client=SimpleNamespace(_client=client))), + "compressor": SimpleNamespace(model=SimpleNamespace(client=SimpleNamespace(_client=compressor))), + }, + }, + ), + ) + runner.observe_usage(app) + + async def request(tokens): + with runner.collect_usage() as records: + await client.send({"prompt_tokens": tokens, "completion_tokens": 20}) + await asyncio.sleep(0) + await compressor.send({}) + return records + + first, second = await asyncio.gather(request(100), request(200)) + assert runner.total_usage(first)["total_tokens"] == 120 + assert runner.total_usage(second)["total_tokens"] == 220 + assert runner.total_usage(first)["records"] == 2 + assert runner.total_usage(first)["coverage"]["total_tokens"] == 1 + assert runner.total_usage(first)["cached_input_tokens"] is None + assert runner._CURRENT.get() is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct", [False, True]) +async def test_failure_artifacts_redact_secrets_and_request_payloads(scene_run, monkeypatch, direct): + _, app, _, case, evaluate = scene_run + secret = "ENV_SECRET_SENTINEL_314159" + monkeypatch.setenv("LLM_API_KEY", secret) + detail = ( + f"provider failed {secret}; api_key=INLINE_SENTINEL " + "Bearer AUTH_SENTINEL https://USER_SENTINEL:PASS_SENTINEL@provider.test/v1?key=URL_SENTINEL " + f"data:image/png;base64,{'QmFzZTY0' * 20}" + ) + if direct: + app.run_job.side_effect = RuntimeError(detail) + else: + app.run_job.return_value = Response( + success=False, + metadata={ + "error": detail, + "request_headers": {"arbitrary": "HEADER_SENTINEL"}, + "config": {"arbitrary": "CONFIG_SENTINEL"}, + "source": {"type": "base64", "data": "SHORT_IMAGE_SENTINEL"}, + }, + ) + await evaluate() + artifacts = "".join(path.read_text() for path in case.rglob("*.json")) + for value in ( + secret, + "INLINE_SENTINEL", + "AUTH_SENTINEL", + "USER_SENTINEL", + "PASS_SENTINEL", + "URL_SENTINEL", + "HEADER_SENTINEL", + "CONFIG_SENTINEL", + "SHORT_IMAGE_SENTINEL", + "QmFzZTY0", + ): + assert value not in artifacts + assert "provider failed" in artifacts and "[REDACTED" in artifacts + + +@pytest.mark.asyncio +async def test_successful_predictions_and_references_remain_verbatim(scene_run): + scene, app, judge, case, evaluate = scene_run + answer = "token=ordinary-text https://example.test/path?a=b data:image/png;base64," + "QmFzZTY0" * 20 + scene["questions"][0]["answer"] = answer + app.run_job.return_value = Response(answer=answer) + results, *_ = await evaluate() + assert results[0]["prediction"] == results[0]["reference"] == answer + assert judge.run_job.call_args.kwargs["prediction"] == answer + for name in ("answer.json", "result.json"): + assert json.loads((case / "questions/0000" / name).read_text())["prediction"] == answer + assert json.loads((case / "status.json").read_text())["memory_complete"] is True + + +@pytest.mark.asyncio +async def test_cancellation_finalizes_observed_requests_before_propagating(scene_run): + _, app, _, case, evaluate = scene_run + + async def canceled(*_, **__): + runner._CURRENT.get().append(measured()) + raise asyncio.CancelledError() + + app.run_job.side_effect = canceled + with pytest.raises(asyncio.CancelledError): + await evaluate() + assert json.loads((case / "status.json").read_text())["memory_usage"]["total_tokens"] == 120 + assert json.loads((case / "ingestion/attempt-001/diagnostics/00000-usage.json").read_text())["total_tokens"] == 120 + app.close.assert_awaited_once() + + +@pytest.mark.parametrize("field", ["cookie", "set-cookie", "private_key", "id_token", "refresh_token", "token"]) +def test_diagnostic_credentials_are_redacted_without_removing_counters(field): + diagnostic = safe_value({field: "DIAGNOSTIC_SECRET_SENTINEL", "input_tokens": 120, "output_tokens": 20}) + assert diagnostic == {field: "[REDACTED]", "input_tokens": 120, "output_tokens": 20} + assert "DIAGNOSTIC_SECRET_SENTINEL" not in json.dumps( + error_details(RuntimeError(f"{field}: DIAGNOSTIC_SECRET_SENTINEL")), + ) + + +@pytest.mark.parametrize( + "secret", + ["response", "success", "answer", "metadata", "usage", "coverage", "records", "input_tokens"], +) +def test_failed_response_schema_survives_environment_name_collision(monkeypatch, secret): + monkeypatch.setenv("TEST_PRIVATE_KEY", secret) + usage = {"input_tokens": 100, "output_tokens": 20, "records": 1, "coverage": {"input_tokens": 1}} + raw = Response( + success=False, + answer=f"Provider unavailable: {secret}", + metadata={"usage": usage, "private_key": "PRIVATE_SENTINEL", f"debug {secret}": f"detail {secret}"}, + ) + response = json.loads(json.dumps(response_record(raw))) + assert Response.model_validate(response).success is False + assert response["answer"] == "Provider unavailable: [REDACTED]" + assert response["metadata"] == { + "usage": usage, + "private_key": "[REDACTED]", + "debug [REDACTED]": "detail [REDACTED]", + } + + +@pytest.mark.parametrize("secret", ["a", "model", "temperature"]) +def test_model_snapshot_preserves_whitelisted_schema(monkeypatch, secret): + monkeypatch.setenv("LLM_API_KEY", secret) + spec = {"backend": "agentscope", "model": "candidate", "parameters": {"temperature": 0, "max_tokens": 256}} + config = SimpleNamespace(components={"as_llm": {"default": SimpleNamespace(model_dump=lambda: spec)}}) + snapshot = model_settings(config)["default"] + assert snapshot["backend"] == "agentscope" and snapshot["model"] == "candidate" + assert snapshot["parameters"] == {"temperature": 0, "max_tokens": 256} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("secret", ["a", "stage"]) +@pytest.mark.parametrize("outcome", ["complete", "index_update", "cleanup"]) +async def test_run_summary_keeps_usage_and_independent_completion(scene_run, monkeypatch, secret, outcome): + import reme + from reme.schema import ApplicationConfig + + scene, app, judge, case, _ = scene_run + root = case.parents[2] + monkeypatch.setenv("LLM_API_KEY", secret) + monkeypatch.setattr(runner, "app_config", lambda workspace, *_, **__: {"workspace_dir": str(workspace)}) + monkeypatch.setattr(runner, "load_memlens", lambda *_: [scene]) + judge.start, judge.close = AsyncMock(), AsyncMock() + judge.config = ApplicationConfig(workspace_dir=root / "judge") + applications = iter([judge, app]) + monkeypatch.setattr(reme, "Application", lambda **_: next(applications)) + + async def job(name, **_): + if name != "index_update": + runner._CURRENT.get().append(measured()) + return Response(answer="Answer", success=name != outcome) + + app.run_job.side_effect = job + if outcome == "cleanup": + app.close.side_effect = RuntimeError("Close failed") + summary = await runner.run( + { + "dataset": str(root / "unused"), + "output": str(root / "output"), + "query_concurrency": 1, + "scenario_concurrency": 1, + "unexpected": "RAW_CONFIG_SENTINEL", + }, + "candidate", + "128k", + ) + group = root / "output/memlens-128k/candidate" + status = json.loads((group / "reme_image/scenes/0000-scene-1/status.json").read_text()) + complete = outcome == "complete" + memory_complete = outcome != "index_update" + assert status["stage"] == outcome and status["scene_id"] == "scene-1" + assert status["memory_complete"] is memory_complete and status["complete"] is complete + assert status["memory_usage"]["input_tokens"] == 100 + assert summary["scored_questions"] == summary["memory_complete_scenes"] == int(memory_complete) + assert summary["token_usage"]["total_tokens"] == (240 if memory_complete else 120) + assert summary["cleanup_error_count"] == int(outcome == "cleanup") + assert summary["completed_scenes"] == int(complete) + assert summary["usage_policy_version"] == "successful-http-requests-v2" + assert json.loads((group / "summary.json").read_text()) == summary + assert json.loads((group / "status.json").read_text())["complete"] is complete + settings = (group / "settings.json").read_text() + assert "RAW_CONFIG_SENTINEL" not in settings + assert json.loads(settings)["judge"]["protocol_id"] == "memlens-task-binary-api-qwen-v1" + if outcome == "cleanup": + assert status["error_type"] == status["error"]["type"] == "RuntimeError" + assert status["error"]["message"] == "Close failed" + + +@pytest.mark.asyncio +async def test_judge_snapshot_uses_validated_plugin_and_environment_config(tmp_path, monkeypatch): + from importlib.metadata import EntryPoint + + from reme import Application + from reme.plugin import PluginManager, _load_plugin + + plugin = _load_plugin( + "memlens-judge", + EntryPoint( + name="memlens-judge", + value="judge_memlens", + group="reme.plugins", + ), + ) + manager = PluginManager([plugin]) + monkeypatch.setattr(PluginManager, "discover", classmethod(lambda cls, _: manager)) + monkeypatch.setenv("JUDGE_MODEL_NAME", "judge-override-for-test") + monkeypatch.setenv("LLM_API_KEY", "SNAPSHOT_SECRET_SENTINEL") + monkeypatch.setenv("LLM_BASE_URL", "https://private.test/v1?key=URL_SENTINEL") + inputs = runner.app_config(tmp_path / "judge", "candidate-model", judge=True) + assert inputs["components"] == {} + app = Application(**inputs) + snapshot = judge_settings(app) + model = snapshot["models"]["judge"] + assert model["model"] == "judge-override-for-test" + assert model["parameters"] == { + "temperature": 0, + "thinking_enable": False, + "max_tokens": 256, + } + assert model["extra_body"] == {"enable_thinking": False} + assert model["client_kwargs"] == {"timeout": 600, "max_retries": 0} + assert snapshot["protocol_id"] == "memlens-task-binary-api-qwen-v1" + assert "SENTINEL" not in json.dumps(snapshot) and "credential" not in json.dumps(snapshot) + await app.close() + + +def test_cli_reports_sanitized_direct_failure(monkeypatch, capsys): + monkeypatch.setattr("sys.argv", ["memlens"]) + monkeypatch.setattr(runner, "run", AsyncMock(side_effect=RuntimeError("provider api_key=CLI_SECRET_SENTINEL"))) + assert runner.main() == 2 + diagnostic = capsys.readouterr().err + assert "CLI_SECRET_SENTINEL" not in diagnostic + assert json.loads(diagnostic)["error"]["message"] == "provider api_key=[REDACTED]" 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) diff --git a/tests/unit/test_read_newlines.py b/tests/unit/test_read_newlines.py new file mode 100644 index 00000000..081f31b5 --- /dev/null +++ b/tests/unit/test_read_newlines.py @@ -0,0 +1,245 @@ +"""Keep ranged reads aligned with universal-newline chunking and large-file reads.""" + +import asyncio +import time + +import pytest + +from reme.components import ApplicationContext +from reme.components.file_chunker import DefaultFileChunker +from reme.steps.file_io import edit as edit_module +from reme.steps.file_io import read as read_module + + +@pytest.mark.parametrize("newline,suffix", [("\n", ".md"), ("\r\n", ".txt")]) +@pytest.mark.parametrize("found", [False, True]) +def test_edit_repeated_prefix_search(tmp_path, newline, suffix, found): + """A small stale edit must not spend seconds retrying a long shared prefix.""" + path = tmp_path / ("note" + suffix) + source = ("x" + newline) * 100_000 + if found: + source += "z" + path.write_bytes(source.encode("utf-8")) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + editor = edit_module.EditStep(app_context=app_context) + started = time.process_time() + await editor(path=path.name, old="x\n" * 2_000 + "z", new="replacement") + cpu_seconds = time.process_time() - started + assert editor.context.response.success is found + if found: + expected = ("x" + newline) * 98_000 + "replacement" + ("\n" if suffix == ".md" else "") + assert "Replaced 1 occurrence(s)" in editor.context.response.answer + assert path.read_bytes() == expected.encode("utf-8") + else: + assert "not found" in editor.context.response.answer + assert path.read_bytes() == source.encode("utf-8") + # CPU time excludes delays from other processes on shared CI runners. + # The literal-search implementation is orders of magnitude below this bound. + assert cpu_seconds < 1.0, f"repeated-prefix search used {cpu_seconds:.3f}s CPU" + + asyncio.run(run()) + + +@pytest.mark.parametrize( + "source,old,expected", + [ + ("\r\na\r\nb\r\n", "a\nb", "\r\nR\r\n"), + ("\r\na\r\nb", "\na", "R\r\nb"), + ("a\r\nb\r\n", "b\n", "a\r\nR"), + ("\r\n\r\n", "\n", "RR"), + ("a\r\nb\rc\nd", "\n", "aRbRcRd"), + ("a\r\n\r\nb", "\n\n", "aRb"), + ("a\r\na\na\ra", "a\na", "R\nR"), + ("\r\na\u2028b\x0bc\r\n", "a\u2028b\x0bc\n", "\r\nR"), + ], +) +def test_edit_normalized_match_boundaries(tmp_path, source, old, expected): + """Mapped spans may start/end at a newline and must remain non-overlapping.""" + path = tmp_path / "note.txt" + path.write_bytes(source.encode("utf-8")) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + editor = edit_module.EditStep(app_context=app_context) + await editor(path=path.name, old=old, new="R") + assert editor.context.response.success is True + assert path.read_bytes() == expected.encode("utf-8") + + asyncio.run(run()) + + +@pytest.mark.parametrize("encoding", ["utf-8", "utf-16"]) +@pytest.mark.parametrize("trailing_newline", [False, True]) +@pytest.mark.parametrize("newline", ["\n", "\r\n", "\r", "mixed"]) +def test_read_line_ranges_match_chunker_and_streaming_path(tmp_path, monkeypatch, encoding, trailing_newline, newline): + """The same bytes must produce the same line ranges on both read paths.""" + logical_lines = ["first\u2028inline", "second 中文", "third\x0binline", "fourth"] + separators = ["\r\n", "\r", "\n"] if newline == "mixed" else [newline] * 3 + text = logical_lines[0] + "".join(separator + line for separator, line in zip(separators, logical_lines[1:])) + if trailing_newline: + text += "\r" if newline == "mixed" else newline + path = tmp_path / "note.md" + source = text.encode(encoding) + path.write_bytes(source) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + # UTF-8 is the normal indexer configuration. UTF-16 still exercises the + # read helper's BOM detection without asserting byte-chunker geometry. + if encoding == "utf-8": + _, chunks = await DefaultFileChunker(encoding=encoding, app_context=app_context).chunk(path) + assert chunks[0].start_line == 1 + assert chunks[0].end_line == len(logical_lines) + + for threshold in (len(source) + 1, 1): + monkeypatch.setattr(read_module, "MAX_FILE_READ_BYTES", threshold) + step = read_module.ReadStep(app_context=app_context) + await step(path="note.md", start_line=2, end_line=3) + assert step.context.response.success is True + assert step.context.response.answer == "\n".join(logical_lines[1:3]) + + complete = read_module.ReadStep(app_context=app_context) + await complete(path="note.md") + assert complete.context.response.success is True + assert complete.context.response.answer == "\n".join(logical_lines) + + past_end = read_module.ReadStep(app_context=app_context) + await past_end(path="note.md", start_line=5) + assert past_end.context.response.success is False + assert "exceeds file length (4 lines)" in past_end.context.response.answer + + assert path.read_bytes() == source + + asyncio.run(run()) + + +@pytest.mark.parametrize("encoding", ["utf-8", "utf-16"]) +@pytest.mark.parametrize("trailing_newline", [False, True]) +@pytest.mark.parametrize("newline", ["\n", "\r\n", "\r"]) +@pytest.mark.parametrize("suffix", [".md", ".txt"]) +def test_read_then_edit_multiline(tmp_path, encoding, trailing_newline, newline, suffix): + """Read output remains usable as an edit needle without changing write rules.""" + text = f"first\u2028inline{newline}second\x0binline" + if trailing_newline: + text += newline + path = tmp_path / ("note" + suffix) + source = text.encode(encoding) + path.write_bytes(source) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + reader = read_module.ReadStep(app_context=app_context) + await reader(path=path.name) + assert reader.context.response.success is True + assert reader.context.response.answer == "first\u2028inline\nsecond\x0binline" + assert path.read_bytes() == source + + editor = edit_module.EditStep(app_context=app_context) + await editor(path=path.name, old=reader.context.response.answer, new="replacement") + assert editor.context.response.success is True + assert "Replaced 1 occurrence(s)" in editor.context.response.answer + ending = "\n" if suffix == ".md" else newline if trailing_newline else "" + assert path.read_bytes() == ("replacement" + ending).encode(encoding) + + asyncio.run(run()) + + +@pytest.mark.parametrize("newline", ["\r\n", "\r"]) +@pytest.mark.parametrize("filename,encoding", [("note.txt", "utf-8"), ("note.txt", "utf-16"), ("note.csv", "gbk")]) +def test_ranged_read_then_edit_preserves_unmatched_bytes(tmp_path, newline, filename, encoding): + """Only the matched span changes; encoding, other newlines and replacement literals survive.""" + prefix = "开头" + newline + target = "first [1].*" + newline + r"second (2)\path" + suffix = newline + "结尾" + newline + path = tmp_path / filename + path.write_bytes((prefix + target + suffix).encode(encoding)) + replacement = r"\1 replacement" + "\n" + r"next\g<0>" + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + reader = read_module.ReadStep(app_context=app_context) + await reader(path=filename, start_line=2, end_line=3) + assert reader.context.response.success is True + editor = edit_module.EditStep(app_context=app_context) + await editor(path=filename, old=reader.context.response.answer, new=replacement) + assert editor.context.response.success is True + assert path.read_bytes() == (prefix + replacement + suffix).encode(encoding) + + asyncio.run(run()) + + +def test_edit_normalized_needle_replaces_all_newline_styles(tmp_path): + """A mixed-newline body must not leave non-LF matches behind.""" + path = tmp_path / "note.txt" + path.write_bytes(b"prefix\r\nfirst\r\nsecond\nmiddle\rfirst\nsecond\r\nfirst\rsecond\r\nsuffix") + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + editor = edit_module.EditStep(app_context=app_context) + await editor(path=path.name, old="first\nsecond", new="done") + assert editor.context.response.success is True + assert "Replaced 3 occurrence(s)" in editor.context.response.answer + assert path.read_bytes() == b"prefix\r\ndone\nmiddle\rdone\r\ndone\r\nsuffix" + + asyncio.run(run()) + + +@pytest.mark.parametrize("separator", ["\r\n", "\r\n\r\n", "\r\r", "\n\n"]) +def test_edit_does_not_treat_crlf_as_two_line_breaks(tmp_path, separator): + """A normalized blank line cannot match a single CRLF separator.""" + path = tmp_path / "note.txt" + source = ("first" + separator + "second").encode("utf-8") + path.write_bytes(source) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + editor = edit_module.EditStep(app_context=app_context) + await editor(path=path.name, old="first\n\nsecond", new="wrong") + if separator == "\r\n": + assert editor.context.response.success is False + assert "not found" in editor.context.response.answer + assert path.read_bytes() == source + else: + assert editor.context.response.success is True + assert path.read_bytes() == b"wrong" + + asyncio.run(run()) + + +def test_edit_normalized_needle_preserves_frontmatter(tmp_path): + """Universal-newline matching stays inside the Markdown body.""" + path = tmp_path / "note.md" + path.write_bytes(b"---\nname: first\ndescription: second\n---\nfirst\rsecond\rsuffix") + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + editor = edit_module.EditStep(app_context=app_context) + await editor(path=path.name, old="first\nsecond", new="done") + assert editor.context.response.success is True + assert path.read_bytes() == b"---\ndescription: second\nname: first\n---\n\ndone\rsuffix\n" + + before = path.read_bytes() + await editor(path=path.name, old="description: second\nname: first", new="wrong") + assert editor.context.response.success is False + assert path.read_bytes() == before + + asyncio.run(run()) + + +@pytest.mark.parametrize("old", ["first", "first\rsecond", "first\r\nsecond"]) +def test_edit_literal_needles_preserve_existing_matching(tmp_path, old): + """Single-line and explicitly raw-newline needles retain exact replacement behavior.""" + path = tmp_path / "note.txt" + source = "prefix\r\nfirst\rsecond / first\r\nsecond / first\nsecond\r\nsuffix" + path.write_bytes(source.encode("utf-8")) + app_context = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + editor = edit_module.EditStep(app_context=app_context) + await editor(path=path.name, old=old, new="done") + assert editor.context.response.success is True + assert path.read_bytes() == source.replace(old, "done").encode("utf-8") + + asyncio.run(run()) diff --git a/tests/unit/test_threaded_watch.py b/tests/unit/test_threaded_watch.py index d93e7c02..5a56bf48 100644 --- a/tests/unit/test_threaded_watch.py +++ b/tests/unit/test_threaded_watch.py @@ -10,7 +10,10 @@ import pytest from watchfiles import Change from reme.components.application_context import ApplicationContext +from reme.components.file_catalog import LocalFileCatalog from reme.components.runtime_context import RuntimeContext +from reme.schema import FileNode +from reme.steps.base_step import BaseStep from reme.steps.index import _threaded_watch from reme.steps.index._threaded_watch import watch_in_thread from reme.steps.index.watch_changes import WatchChangesStep @@ -49,6 +52,78 @@ def test_watch_mode_defaults_and_validation(): WatchChangesStep(watch_mode=value) +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["watch", "awatch"]) +async def test_watcher_reconciles_startup_gap_on_first_idle_batch(tmp_path, monkeypatch, mode): + """Both watchers reconcile once even when no subsequent change arrives.""" + + async def idle_batches(*_paths, **kwargs): + assert kwargs["yield_on_timeout"] is True + yield set() + yield set() + + monkeypatch.setattr("reme.steps.index.watch_changes.awatch", idle_batches) + monkeypatch.setattr("reme.steps.index.watch_changes.watch_in_thread", idle_batches) + step = _step(tmp_path, asyncio.Event(), mode=mode) + step.monitor_type = "file_store" + step.dispatch_step_specs = ["update_index_step"] + + await step.execute() + + step.dispatch_steps.assert_awaited_once_with( + [ + { + "backend": "init_changes_step", + "monitor_type": "file_store", + "monitor_name": "default", + "recursive": True, + }, + ], + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["watch", "awatch"]) +async def test_watcher_merges_startup_scan_and_first_change_batch(tmp_path, monkeypatch, mode): + """Overlapping events and missed additions/deletions dispatch once together.""" + step = _step(tmp_path, asyncio.Event(), mode=mode) + catalog = LocalFileCatalog(app_context=step.app_context) + step.app_context.components["file_catalog"] = {"default": catalog} + await catalog.upsert([FileNode(path="daily/deleted.md", st_mtime=1)]) + step.monitor_type = "file_catalog" + step.dispatch_step_specs = ["update_catalog_step"] + changed = tmp_path / "daily/changed.md" + missed = tmp_path / "daily/missed.md" + + async def first_batch(*_paths, **_kwargs): + changed.write_text("new", encoding="utf-8") + missed.write_text("missed", encoding="utf-8") + yield {(Change.added, str(changed))} + yield set() + + batches = [] + + async def dispatch(specs, **kwargs): + if specs[0] == "update_catalog_step": + batches.append(kwargs["changes"]) + return [] + return await BaseStep.dispatch_steps(step, specs, **kwargs) + + monkeypatch.setattr("reme.steps.index.watch_changes.awatch", first_batch) + monkeypatch.setattr("reme.steps.index.watch_changes.watch_in_thread", first_batch) + step.dispatch_steps = dispatch + await step.execute() + + assert len(batches) == 1 + assert sorted((item["path"], item["change"]) for item in batches[0]) == sorted( + [ + (str(changed), "added"), + (str(missed), "added"), + (str(tmp_path / "daily/deleted.md"), "deleted"), + ], + ) + + @pytest.mark.asyncio async def test_default_mode_uses_awatch_and_closes_on_dispatch_failure(tmp_path, monkeypatch): """An exception in downstream processing closes the default generator."""