iomgaa/Video-Tree-TRM5
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(harness): momentum.py — async 慢更新动量生成

Browse Source
This commit is contained in:
iomgaa
2026-07-07 13:01:04 -04:00
parent a550d39e1c
commit 886a444d1d
2 changed files with 439 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/harness/momentum.py
+178
View File
@@ -0,0 +1,178 @@
"""慢更新动量生成 — epoch 末为单个 skill 产出新的动量指导。
对标 SkillOpt 的 slow_update 机制：拿上一 epoch 末与当前 epoch 末两版 skill，
在固定样本上各跑一遍得到纵向对比（comparison_pairs），反思上一轮动量指导是否奏效、
本轮正文改动是改善还是漂移，据此重写动量指导。新指导经 patch 引擎的 replace_momentum
写回 skill 的 momentum 受保护区，作为下一轮进化的方向锚。
从 TRM4 core/harness/momentum.py（156 行）迁移 + async 化。
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from loguru import logger
from core.evolution.diagnose import extract_json_from_response
if TYPE_CHECKING:
from pathlib import Path
from core.protocols import LLMProvider
# =========================================================================
# 四类纵向对比类别名（单一真源）
# =========================================================================
IMPROVED = "improved" # 错→对
REGRESSED = "regressed" # 对→错
PERSISTENT_FAIL = "persistent_fail" # 错→错
STABLE_SUCCESS = "stable_success" # 对→对
# 类别名 → 展示标题，列表顺序即展示顺序。
# 回退（REGRESSED）刻意排在改善（IMPROVED）之前——它是最该警惕的伤害信号。
_CATEGORY_LABELS: tuple[tuple[str, str], ...] = (
(REGRESSED, "从对变错（回退，最高优先级）"),
(PERSISTENT_FAIL, "始终答错（持续失败）"),
(IMPROVED, "从错变对（改善）"),
(STABLE_SUCCESS, "始终答对（稳定成功）"),
)
# =========================================================================
# 辅助函数
# =========================================================================
def _categorize_pair(pair: dict[str, Any]) -> str:
"""按两版正误派生纵向对比类别。
用键值的真值（bool(...)）表示该题在两版上各自的正误：缺 correct_prev/
correct_curr 键时直接抛 KeyError 向上传播——这是上游数据损坏（不是裁判语义
歧义），静默当 False 会伪造 persistent_fail 证据、污染动量指导，故不掩盖。
参数:
pair: 单个纵向对比对，须含 correct_prev/correct_curr 两键。
返回:
四个类别命名常量之一：IMPROVED（错→对）/REGRESSED（对→错）/
PERSISTENT_FAIL（错→错）/STABLE_SUCCESS（对→对）。
异常:
KeyError: 缺 correct_prev 或 correct_curr 键时。
"""
correct_prev = bool(pair["correct_prev"])
correct_curr = bool(pair["correct_curr"])
if not correct_prev and correct_curr:
return IMPROVED
if correct_prev and not correct_curr:
return REGRESSED
if not correct_prev and not correct_curr:
return PERSISTENT_FAIL
return STABLE_SUCCESS
def _format_comparison_pairs(comparison_pairs: list[dict[str, Any]]) -> str:
"""将纵向对比对格式化为裁判可读文本，按 _CATEGORY_LABELS 分组与排序。
参数:
comparison_pairs: 每个 dict 含 question/prev_prediction/curr_prediction/
correct_prev/correct_curr 字段，描述一道固定样本上两版的成对结果。
返回:
可读的纵向对比文本；空列表返回占位说明。
异常:
KeyError: 任一 pair 缺 correct_prev/correct_curr 键时；不掩盖的理由见
_categorize_pair docstring。
"""
if not comparison_pairs:
return "（本轮无可用纵向对比样本）"
grouped: dict[str, list[dict[str, Any]]] = {key: [] for key, _ in _CATEGORY_LABELS}
for pair in comparison_pairs:
grouped[_categorize_pair(pair)].append(pair)
lines: list[str] = [f"固定样本总数：{len(comparison_pairs)}"]
for key, label in _CATEGORY_LABELS:
entries = grouped[key]
lines.append(f"\n### {label}（{len(entries)} 题）")
if not entries:
lines.append("（无）")
continue
for pair in entries:
lines.append(
f"- 题目：{pair.get('question', '')}\n"
f" 上版预测：{pair.get('prev_prediction', '')} | "
f"当前版预测：{pair.get('curr_prediction', '')}"
)
return "\n".join(lines)
# =========================================================================
# 入口
# =========================================================================
async def run_slow_momentum(
llm: LLMProvider,
diagnose_prompts_dir: Path,
skill_content: str,
prev_skill: str,
prev_guidance: str,
comparison_pairs: list[dict[str, Any]],
) -> str:
"""为单个 skill 生成新的慢更新动量指导。
参数:
llm: LLM 端口（async chat）。
diagnose_prompts_dir: 诊断 prompt 目录（根 prompts/，slow_momentum.md 在此）。
skill_content: 当前版 skill 正文。
prev_skill: 上一版 skill 正文。
prev_guidance: 上一轮写下的动量指导。
comparison_pairs: 固定样本上两版 rollout 的成对结果（含 question/
prev_prediction/curr_prediction/correct_prev/correct_curr）。
返回:
新的动量指导文本；解析失败时保留 prev_guidance。
关键实现细节:
- _format_comparison_pairs 刻意置于 try 块之外（prompt 构造阶段）：它对每个
pair 取 correct_prev/correct_curr，缺键抛 KeyError 直接向上传播，不被下方
针对裁判语义歧义的 except ValueError 吞掉。
- 解析失败保留上轮指导：extract_json_from_response 抛 ValueError、缺
slow_update_content 字段、或该字段非 str，均视为语义解析失败，返回
prev_guidance（判不准时保守保留上轮指导，对标 diagnose 的保护性 fallback）。
- P5 边界：仅捕 ValueError 这一语义歧义；llm.chat 的基础设施失败
（网络/API 异常）刻意不捕，向上传播，绝不用默认值掩盖。
"""
system_prompt = (diagnose_prompts_dir / "slow_momentum.md").read_text(encoding="utf-8")
# _format_comparison_pairs 刻意置于下方 try 块之外（prompt 构造阶段）：它对每个
# pair 取 correct_prev/correct_curr，缺键抛 KeyError 直接向上传播，不被下方针对
# 裁判语义歧义的 except ValueError 吞掉。异常类型选 KeyError（非 ValueError），
# 即便位置疏忽落入 try 也不会被误吞。
user_prompt = (
f"## 上一版 skill 正文\n{prev_skill}\n\n"
f"## 当前版 skill 正文\n{skill_content}\n\n"
f"## 上一轮的动量指导\n{prev_guidance}\n\n"
f"## 固定样本纵向对比（上版 vs 当前版）\n"
f"{_format_comparison_pairs(comparison_pairs)}"
)
response = await llm.chat(
[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
]
)
raw = response.content
try:
parsed = extract_json_from_response(raw)
new_guidance = parsed.get("slow_update_content")
if not isinstance(new_guidance, str):
raise ValueError("slow_update_content 字段缺失或非字符串")
except ValueError:
logger.warning("慢更新动量解析失败，保留上轮动量指导")
return prev_guidance
return new_guidance
tests/unit/test_harness_momentum.py
+261
View File
@@ -0,0 +1,261 @@
"""慢更新动量生成（app/harness/momentum.py）单元测试。
覆盖:
- _categorize_pair 四分类 + 缺键 KeyError
- _format_comparison_pairs 分组排序 + 空列表
- run_slow_momentum 正常解析 + 解析失败保留 prev_guidance
"""
from __future__ import annotations
import json
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from unittest.mock import AsyncMock
import pytest
from app.harness.momentum import (
IMPROVED,
PERSISTENT_FAIL,
REGRESSED,
STABLE_SUCCESS,
_categorize_pair,
_format_comparison_pairs,
run_slow_momentum,
)
if TYPE_CHECKING:
from pathlib import Path
# =========================================================================
# _categorize_pair
# =========================================================================
class TestCategorizePair:
"""_categorize_pair 四分类测试。"""
def test_categorize_pair_all_four(self) -> None:
"""四种 correct_prev/correct_curr 组合应返回对应类别常量。"""
assert _categorize_pair({"correct_prev": False, "correct_curr": True}) == IMPROVED
assert _categorize_pair({"correct_prev": True, "correct_curr": False}) == REGRESSED
assert _categorize_pair({"correct_prev": False, "correct_curr": False}) == PERSISTENT_FAIL
assert _categorize_pair({"correct_prev": True, "correct_curr": True}) == STABLE_SUCCESS
def test_categorize_pair_truthy_values(self) -> None:
"""非布尔真值（int / str）也能正确分类。"""
assert _categorize_pair({"correct_prev": 0, "correct_curr": 1}) == IMPROVED
assert _categorize_pair({"correct_prev": "yes", "correct_curr": ""}) == REGRESSED
def test_categorize_pair_missing_key(self) -> None:
"""缺 correct_prev 或 correct_curr 键时抛 KeyError。"""
with pytest.raises(KeyError):
_categorize_pair({"correct_prev": True})
with pytest.raises(KeyError):
_categorize_pair({"correct_curr": False})
with pytest.raises(KeyError):
_categorize_pair({})
# =========================================================================
# _format_comparison_pairs
# =========================================================================
class TestFormatComparisonPairs:
"""_format_comparison_pairs 分组排序测试。"""
def test_format_comparison_pairs_empty(self) -> None:
"""空列表返回占位说明。"""
result = _format_comparison_pairs([])
assert "无可用" in result
def test_format_comparison_pairs_order(self) -> None:
"""REGRESSED 标题应出现在 IMPROVED 标题之前（伤害信号优先）。"""
pairs = [
{
"question": "Q1",
"prev_prediction": "A",
"curr_prediction": "B",
"correct_prev": False,
"correct_curr": True,
},
{
"question": "Q2",
"prev_prediction": "C",
"curr_prediction": "D",
"correct_prev": True,
"correct_curr": False,
},
]
result = _format_comparison_pairs(pairs)
regressed_pos = result.index("回退")
improved_pos = result.index("改善")
assert regressed_pos < improved_pos, "REGRESSED 应排在 IMPROVED 之前"
def test_format_comparison_pairs_all_categories(self) -> None:
"""四种类别的 pair 都能被正确分组。"""
pairs = [
{"question": "Q1", "correct_prev": False, "correct_curr": True},
{"question": "Q2", "correct_prev": True, "correct_curr": False},
{"question": "Q3", "correct_prev": False, "correct_curr": False},
{"question": "Q4", "correct_prev": True, "correct_curr": True},
]
result = _format_comparison_pairs(pairs)
assert "固定样本总数：4" in result
# 每个类别都应标注 1 题
assert "1 题" in result
def test_format_comparison_pairs_missing_key(self) -> None:
"""pair 缺键时 KeyError 不被吞。"""
with pytest.raises(KeyError):
_format_comparison_pairs([{"question": "Q1"}])
# =========================================================================
# run_slow_momentum
# =========================================================================
def _make_mock_llm(response_content: str) -> Any:
"""构造一个返回指定 content 的 mock LLMProvider。"""
@dataclass(frozen=True)
class _FakeResponse:
content: str
thinking: str = ""
model: str = "mock"
provider: str = "mock"
prompt_tokens: int = 0
completion_tokens: int = 0
latency_ms: int = 0
ttft_ms: float | None = None
max_inter_token_ms: float | None = None
cache_hit: bool = False
call_id: str = "test-call-id"
mock_llm = AsyncMock()
mock_llm.chat.return_value = _FakeResponse(content=response_content)
return mock_llm
@pytest.mark.asyncio
async def test_run_slow_momentum_basic(tmp_path: Path) -> None:
"""正常解析时返回新的动量指导文本。"""
# 准备 prompt 文件
prompt_file = tmp_path / "slow_momentum.md"
prompt_file.write_text("你是一个慢更新动量裁判。", encoding="utf-8")
new_guidance_text = "新一轮的动量指导内容"
llm_response = json.dumps({"slow_update_content": new_guidance_text}, ensure_ascii=False)
mock_llm = _make_mock_llm(llm_response)
result = await run_slow_momentum(
llm=mock_llm,
diagnose_prompts_dir=tmp_path,
skill_content="当前 skill 正文",
prev_skill="上一版 skill 正文",
prev_guidance="旧的动量指导",
comparison_pairs=[
{
"question": "Q1",
"prev_prediction": "A",
"curr_prediction": "B",
"correct_prev": False,
"correct_curr": True,
}
],
)
assert result == new_guidance_text
mock_llm.chat.assert_awaited_once()
@pytest.mark.asyncio
async def test_run_slow_momentum_parse_failure(tmp_path: Path) -> None:
"""LLM 返回无法解析的内容时保留 prev_guidance。"""
prompt_file = tmp_path / "slow_momentum.md"
prompt_file.write_text("你是一个慢更新动量裁判。", encoding="utf-8")
mock_llm = _make_mock_llm("这不是 JSON，无法解析")
prev_guidance = "应该被保留的旧动量指导"
result = await run_slow_momentum(
llm=mock_llm,
diagnose_prompts_dir=tmp_path,
skill_content="当前 skill",
prev_skill="上一版 skill",
prev_guidance=prev_guidance,
comparison_pairs=[
{
"question": "Q1",
"prev_prediction": "A",
"curr_prediction": "B",
"correct_prev": True,
"correct_curr": True,
}
],
)
assert result == prev_guidance
@pytest.mark.asyncio
async def test_run_slow_momentum_missing_field(tmp_path: Path) -> None:
"""LLM 返回合法 JSON 但缺少 slow_update_content 字段时保留 prev_guidance。"""
prompt_file = tmp_path / "slow_momentum.md"
prompt_file.write_text("你是一个慢更新动量裁判。", encoding="utf-8")
llm_response = json.dumps({"other_field": "无关内容"}, ensure_ascii=False)
mock_llm = _make_mock_llm(llm_response)
prev_guidance = "应该被保留的旧动量指导"
result = await run_slow_momentum(
llm=mock_llm,
diagnose_prompts_dir=tmp_path,
skill_content="当前 skill",
prev_skill="上一版 skill",
prev_guidance=prev_guidance,
comparison_pairs=[],
)
assert result == prev_guidance
@pytest.mark.asyncio
async def test_run_slow_momentum_keyerror_not_swallowed(tmp_path: Path) -> None:
"""comparison_pairs 缺键时 KeyError 不被 ValueError 吞掉。"""
prompt_file = tmp_path / "slow_momentum.md"
prompt_file.write_text("你是一个慢更新动量裁判。", encoding="utf-8")
mock_llm = _make_mock_llm('{"slow_update_content": "ok"}')
with pytest.raises(KeyError):
await run_slow_momentum(
llm=mock_llm,
diagnose_prompts_dir=tmp_path,
skill_content="当前 skill",
prev_skill="上一版 skill",
prev_guidance="旧指导",
comparison_pairs=[{"question": "Q1"}], # 缺 correct_prev/correct_curr
)
@pytest.mark.asyncio
async def test_run_slow_momentum_fenced_json(tmp_path: Path) -> None:
"""LLM 返回 fenced code block 中的 JSON 也能正确解析。"""
prompt_file = tmp_path / "slow_momentum.md"
prompt_file.write_text("你是一个慢更新动量裁判。", encoding="utf-8")
new_guidance = "从 fenced block 中提取的指导"
llm_response = f'```json\n{{"slow_update_content": "{new_guidance}"}}\n```'
mock_llm = _make_mock_llm(llm_response)
result = await run_slow_momentum(
llm=mock_llm,
diagnose_prompts_dir=tmp_path,
skill_content="当前 skill",
prev_skill="上一版 skill",
prev_guidance="旧指导",
comparison_pairs=[],
)
assert result == new_guidance