mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-11-04 07:30:21 +00:00
Merge current upstream main into PR #599
This commit is contained in:
commit
b431ac2ffd
79 changed files with 6458 additions and 60 deletions
4
.github/workflows/docker.yml
vendored
4
.github/workflows/docker.yml
vendored
|
|
@ -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 }}
|
||||
|
|
|
|||
5
.github/workflows/release-auto-fin.yml
vendored
5
.github/workflows/release-auto-fin.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
2
.github/workflows/release-daily-paper.yml
vendored
2
.github/workflows/release-daily-paper.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
2
.github/workflows/release-python.yml
vendored
2
.github/workflows/release-python.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
2
.github/workflows/release-reme-studio.yml
vendored
2
.github/workflows/release-reme-studio.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
4
.github/workflows/security-codeql.yml
vendored
4
.github/workflows/security-codeql.yml
vendored
|
|
@ -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 }}
|
||||
|
|
|
|||
107
benchmark/memlens/README.md
Normal file
107
benchmark/memlens/README.md
Normal file
|
|
@ -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
|
||||
|
||||
`<output>/memlens-<length>/<model>/` contains settings and summary JSON.
|
||||
`reme_image/scenes/<scene>/` 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.
|
||||
```
|
||||
69
benchmark/memlens/README_ZH.md
Normal file
69
benchmark/memlens/README_ZH.md
Normal file
|
|
@ -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 模型、生成参数与协议,不保存凭据或服务端点。
|
||||
|
||||
## 结果
|
||||
|
||||
结果保存到 `<output>/memlens-<length>/<model>/`。该目录的 JSON 记录配置与汇总;
|
||||
`reme_image/scenes/<scene>/` 保存摄入响应、原生 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)。
|
||||
163
benchmark/memlens/artifacts.py
Normal file
163
benchmark/memlens/artifacts.py
Normal file
|
|
@ -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"(?<![\w-])" + re.escape(secret) + r"(?![\w-])", "[REDACTED]", value)
|
||||
else:
|
||||
value = value.replace(secret, "[REDACTED]")
|
||||
|
||||
def safe_url(match):
|
||||
try:
|
||||
url = urlsplit(match.group())
|
||||
host = url.netloc.rsplit("@", 1)[-1]
|
||||
return urlunsplit((url.scheme, host, url.path, "[REDACTED]" if url.query else "", ""))
|
||||
except ValueError:
|
||||
return "[REDACTED_URL]"
|
||||
|
||||
value = re.sub(r"https?://[^\s\"'<>]+", 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),
|
||||
}
|
||||
4
benchmark/memlens/config.yaml
Normal file
4
benchmark/memlens/config.yaml
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
dataset: ${MEMLENS_DATA}
|
||||
output: ${EXPERIMENT_DIR}
|
||||
scenario_concurrency: 4
|
||||
query_concurrency: 16
|
||||
621
benchmark/memlens/run.py
Normal file
621
benchmark/memlens/run.py
Normal file
|
|
@ -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 <image> placement."""
|
||||
messages = []
|
||||
for index, turn in enumerate(session["turns"]):
|
||||
text, images = turn.get("content", ""), turn.get("images", [])
|
||||
if not images:
|
||||
stripped = text.replace("<image>", "").strip()
|
||||
parts = [stripped] if stripped else []
|
||||
elif "<image>" not in text:
|
||||
parts = [*images, *([text.strip()] if text.strip() else [])]
|
||||
else:
|
||||
parts, image_index = [], 0
|
||||
for part in re.split(r"(<image>)", text):
|
||||
if part == "<image>":
|
||||
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())
|
||||
96
benchmark/mobilemem/README.md
Normal file
96
benchmark/mobilemem/README.md
Normal file
|
|
@ -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 `<output>/mobilemem/<model>` 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.
|
||||
61
benchmark/mobilemem/README_ZH.md
Normal file
61
benchmark/mobilemem/README_ZH.md
Normal file
|
|
@ -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` 覆盖。
|
||||
每次运行创建 `<output>/mobilemem/<model>`,已有目录不会被覆盖。
|
||||
|
||||
## 评测设置
|
||||
|
||||
每个用户使用独立的新工作空间,按时间顺序通过多模态 `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)。
|
||||
182
benchmark/mobilemem/_audit.py
Normal file
182
benchmark/mobilemem/_audit.py
Normal file
|
|
@ -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"(?<![\w-])" + re.escape(secret) + r"(?![\w-])", "[redacted]", text)
|
||||
else:
|
||||
text = text.replace(secret, "[redacted]")
|
||||
text = re.sub(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
"[private key redacted]",
|
||||
text,
|
||||
flags=re.S,
|
||||
)
|
||||
text = re.sub(r"(?i)\b(?:request[_ -]?)?headers\s*[:=].*", "headers=[redacted]", text)
|
||||
text = re.sub(r"(?i)\b(?:set-cookie|cookie)\s*:[^\r\n]*", "cookie: [redacted]", text)
|
||||
text = re.sub(r"data:[^\s,;]+(?:;[^\s,]*)?,[^\s\"']+", "[image data redacted]", text, flags=re.I)
|
||||
text = re.sub(r"https?://[^\s\"'<>]+", "[url redacted]", text, flags=re.I)
|
||||
text = re.sub(r"(?<![\w])(?:[A-Za-z]:\\|/)[^\s\"']+", "[path redacted]", text)
|
||||
text = re.sub(r"\b(?:Bearer|Basic)\s+[^\s,;\"'}]+", "[authorization redacted]", text, flags=re.I)
|
||||
text = re.sub(
|
||||
r"(?i)([\"']?(?:api[_-]?key|private[_-]?key|authorization|password|secret|access[_-]?token|"
|
||||
r"refresh[_-]?token|id[_-]?token|token|cookie)[\"']?\s*[:=]\s*)"
|
||||
r"(?:[\"'][^\"']*[\"']|[^\s,;}]+)",
|
||||
r"\1[redacted]",
|
||||
text,
|
||||
)
|
||||
text = re.sub(r"\b(?:sk-|sk_|ghp_|github_pat_)[A-Za-z0-9_-]+", "[key redacted]", text)
|
||||
text = re.sub(r"[A-Za-z0-9+/=_-]{80,}", "[payload redacted]", text)
|
||||
return text[:limit]
|
||||
|
||||
|
||||
def safe_value(value, depth: int = 0, *, fixed_keys=()):
|
||||
"""Sanitize failed response envelopes, never serialize arbitrary objects/locals."""
|
||||
if depth > 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,
|
||||
}
|
||||
3
benchmark/mobilemem/config.yaml
Normal file
3
benchmark/mobilemem/config.yaml
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
model: qwen3.8-max
|
||||
user_concurrency: 4
|
||||
query_concurrency: 16
|
||||
621
benchmark/mobilemem/run.py
Normal file
621
benchmark/mobilemem/run.py
Normal file
|
|
@ -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())
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -162,3 +162,9 @@ YAML 中的 `model` 仍表示供应商模型名称;运行时 `model` 属性接
|
|||
宿主应等待 ReMe 调用结束后自行关闭实例。替换 ReMe 自建模型或存在进行中请求的实例之前,须先关闭组件。
|
||||
关闭过程即使被取消,也会等待活动请求和自建客户端清理完成,再传播取消信号。组件关闭时会解除实例引用;
|
||||
重新启动使用外部模型的组件前需要再次注入。独立运行时按凭据创建模型的行为保持不变。
|
||||
|
||||
自建客户端在线程中初始化,证书加载不会阻塞事件循环。构造期间取消请求时,组件会等待构造完成并保留客户端所有权,
|
||||
由后续组件关闭负责释放;此时不会发送该请求。
|
||||
|
||||
宿主可从 `reme.steps.compaction` 导入 `CompactContextStep`、`ToolPair`、`collect_pairs` 和 `preview_result`。
|
||||
配对识别与预览构造函数和压缩 Step 共享相同规则,供宿主校验提案使用;归档、引用解析和提交仍由宿主负责。
|
||||
|
|
|
|||
|
|
@ -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,越界页返回空列表;
|
||||
总数只统计可读取、有效且有权限访问的笔记。分页限制返回量,不限制扫描量;每次查询读取当前文件,不提供跨页快照。
|
||||
|
||||
---
|
||||
|
||||
## 自动记忆
|
||||
|
|
|
|||
21
plugins/memlens-judge/LICENSE-CODE
Normal file
21
plugins/memlens-judge/LICENSE-CODE
Normal file
|
|
@ -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.
|
||||
30
plugins/memlens-judge/README.md
Normal file
30
plugins/memlens-judge/README.md
Normal file
|
|
@ -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.
|
||||
22
plugins/memlens-judge/pyproject.toml
Normal file
22
plugins/memlens-judge/pyproject.toml
Normal file
|
|
@ -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"
|
||||
5
plugins/memlens-judge/src/judge_memlens/__init__.py
Normal file
5
plugins/memlens-judge/src/judge_memlens/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
"""Independent memlens evaluation-only plugin."""
|
||||
|
||||
from .llm_judge import MemLensJudgeStep
|
||||
|
||||
__all__ = ["MemLensJudgeStep"]
|
||||
File diff suppressed because one or more lines are too long
72
plugins/memlens-judge/src/judge_memlens/llm_judge.py
Normal file
72
plugins/memlens-judge/src/judge_memlens/llm_judge.py
Normal file
|
|
@ -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
|
||||
38
plugins/memlens-judge/src/judge_memlens/plugin.yaml
Normal file
38
plugins/memlens-judge/src/judge_memlens/plugin.yaml
Normal file
|
|
@ -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
|
||||
139
plugins/memlens-judge/src/judge_memlens/protocol.py
Normal file
139
plugins/memlens-judge/src/judge_memlens/protocol.py
Normal file
|
|
@ -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 (Outdated) Answer>: {old_answer}\n" if key == "KU_KnowledgeUpdate" and old_answer else ""
|
||||
return (
|
||||
f"{prompt}[Current Case]\n<Question>: {question}\n<Standard Answer>: {reference}\n"
|
||||
f"{old}<Student Answer>: {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 "</think>" in text:
|
||||
text = text.split("</think>", 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()
|
||||
48
plugins/memlens/README.md
Normal file
48
plugins/memlens/README.md
Normal file
|
|
@ -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.
|
||||
25
plugins/memlens/pyproject.toml
Normal file
25
plugins/memlens/pyproject.toml
Normal file
|
|
@ -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"
|
||||
5
plugins/memlens/src/reme_memlens/__init__.py
Normal file
5
plugins/memlens/src/reme_memlens/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
"""MemLens benchmark Jobs for ReMe."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
BENCHMARK_CONFIG = Path(__file__).with_name("benchmark.yaml")
|
||||
136
plugins/memlens/src/reme_memlens/agentic_answer.py
Normal file
136
plugins/memlens/src/reme_memlens/agentic_answer.py
Normal file
|
|
@ -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
|
||||
10
plugins/memlens/src/reme_memlens/agentic_answer.yaml
Normal file
10
plugins/memlens/src/reme_memlens/agentic_answer.yaml
Normal file
|
|
@ -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.
|
||||
34
plugins/memlens/src/reme_memlens/auto_memory.py
Normal file
34
plugins/memlens/src/reme_memlens/auto_memory.py
Normal file
|
|
@ -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)
|
||||
)
|
||||
25
plugins/memlens/src/reme_memlens/auto_memory.yaml
Normal file
25
plugins/memlens/src/reme_memlens/auto_memory.yaml
Normal file
|
|
@ -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.
|
||||
30
plugins/memlens/src/reme_memlens/benchmark.yaml
Normal file
30
plugins/memlens/src/reme_memlens/benchmark.yaml
Normal file
|
|
@ -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
|
||||
75
plugins/memlens/src/reme_memlens/plugin.yaml
Normal file
75
plugins/memlens/src/reme_memlens/plugin.yaml
Normal file
|
|
@ -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
|
||||
21
plugins/mobilemem-judge/LICENSE-CODE
Normal file
21
plugins/mobilemem-judge/LICENSE-CODE
Normal file
|
|
@ -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.
|
||||
29
plugins/mobilemem-judge/README.md
Normal file
29
plugins/mobilemem-judge/README.md
Normal file
|
|
@ -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`.
|
||||
22
plugins/mobilemem-judge/pyproject.toml
Normal file
22
plugins/mobilemem-judge/pyproject.toml
Normal file
|
|
@ -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"
|
||||
5
plugins/mobilemem-judge/src/judge_mobilemem/__init__.py
Normal file
5
plugins/mobilemem-judge/src/judge_mobilemem/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
"""Independent MobileMem evaluation-only plugin."""
|
||||
|
||||
from .llm_judge import MobileMemJudgeStep
|
||||
|
||||
__all__ = ["MobileMemJudgeStep"]
|
||||
43
plugins/mobilemem-judge/src/judge_mobilemem/llm_judge.py
Normal file
43
plugins/mobilemem-judge/src/judge_mobilemem/llm_judge.py
Normal file
|
|
@ -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
|
||||
40
plugins/mobilemem-judge/src/judge_mobilemem/plugin.yaml
Normal file
40
plugins/mobilemem-judge/src/judge_mobilemem/plugin.yaml
Normal file
|
|
@ -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
|
||||
118
plugins/mobilemem-judge/src/judge_mobilemem/protocol.py
Normal file
118
plugins/mobilemem-judge/src/judge_mobilemem/protocol.py
Normal file
|
|
@ -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,
|
||||
}
|
||||
17
plugins/mobilemem/README.md
Normal file
17
plugins/mobilemem/README.md
Normal file
|
|
@ -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.
|
||||
25
plugins/mobilemem/pyproject.toml
Normal file
25
plugins/mobilemem/pyproject.toml
Normal file
|
|
@ -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"
|
||||
5
plugins/mobilemem/src/reme_mobilemem/__init__.py
Normal file
5
plugins/mobilemem/src/reme_mobilemem/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
"""MobileMem-Omni benchmark Jobs for ReMe."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
BENCHMARK_CONFIG = Path(__file__).with_name("benchmark.yaml")
|
||||
136
plugins/mobilemem/src/reme_mobilemem/agentic_answer.py
Normal file
136
plugins/mobilemem/src/reme_mobilemem/agentic_answer.py
Normal file
|
|
@ -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
|
||||
8
plugins/mobilemem/src/reme_mobilemem/agentic_answer.yaml
Normal file
8
plugins/mobilemem/src/reme_mobilemem/agentic_answer.yaml
Normal file
|
|
@ -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.
|
||||
34
plugins/mobilemem/src/reme_mobilemem/auto_memory.py
Normal file
34
plugins/mobilemem/src/reme_mobilemem/auto_memory.py
Normal file
|
|
@ -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)
|
||||
)
|
||||
17
plugins/mobilemem/src/reme_mobilemem/auto_memory.yaml
Normal file
17
plugins/mobilemem/src/reme_mobilemem/auto_memory.yaml
Normal file
|
|
@ -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.
|
||||
32
plugins/mobilemem/src/reme_mobilemem/benchmark.yaml
Normal file
32
plugins/mobilemem/src/reme_mobilemem/benchmark.yaml
Normal file
|
|
@ -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}
|
||||
71
plugins/mobilemem/src/reme_mobilemem/plugin.yaml
Normal file
71
plugins/mobilemem/src/reme_mobilemem/plugin.yaml
Normal file
|
|
@ -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
|
||||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 ``<daily_dir>/<date>/*.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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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/<date>/<name>.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 ``<daily_dir>/<date>/`` 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)}")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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)")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
100
tests/unit/test_async_utils.py
Normal file
100
tests/unit/test_async_utils.py
Normal file
|
|
@ -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)
|
||||
184
tests/unit/test_atomic_file_write.py
Normal file
184
tests/unit/test_atomic_file_write.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
318
tests/unit/test_daily_list_range.py
Normal file
318
tests/unit/test_daily_list_range.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
335
tests/unit/test_memlens_benchmark.py
Normal file
335
tests/unit/test_memlens_benchmark.py
Normal file
|
|
@ -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 <image> 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, "<think>ignore</think>$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")
|
||||
101
tests/unit/test_memlens_judge.py
Normal file
101
tests/unit/test_memlens_judge.py
Normal file
|
|
@ -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 (Outdated) Answer>: OLD\n" if key == "KU_KnowledgeUpdate" else ""
|
||||
rendered = (
|
||||
f"{base}[Current Case]\n<Question>: QUESTION\n<Standard Answer>: {reference}\n"
|
||||
f"{old}<Student Answer>: 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
|
||||
423
tests/unit/test_memlens_runner_failures.py
Normal file
423
tests/unit/test_memlens_runner_failures.py
Normal file
|
|
@ -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]"
|
||||
475
tests/unit/test_mobilemem_audit.py
Normal file
475
tests/unit/test_mobilemem_audit.py
Normal file
|
|
@ -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
|
||||
407
tests/unit/test_mobilemem_benchmark.py
Normal file
407
tests/unit/test_mobilemem_benchmark.py
Normal file
|
|
@ -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")
|
||||
85
tests/unit/test_mobilemem_judge.py
Normal file
85
tests/unit/test_mobilemem_judge.py
Normal file
|
|
@ -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)
|
||||
245
tests/unit/test_read_newlines.py
Normal file
245
tests/unit/test_read_newlines.py
Normal file
|
|
@ -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())
|
||||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue