iomgaa
730caa7e9a
Reuse build_units/unit_correctness (pair contract single entry) in the inference aggregation step: single questions count as one unit, AR pairs collapse original+mirror into one unit scored by bidirectional AND. total/ correct/per_task_type are unit-grained; orphan pairs (missing one side) are warned and dropped, not counted. Per-question predictions still land row by row (traceability unchanged).
511 lines
17 KiB
Python
511 lines
17 KiB
Python
"""async 推理编排 — 训练循环的 forward()。
|
||
|
||
从 TRM4 core/harness/inference.py (~560 行) 迁移，重大重构:
|
||
- 同步 ThreadPoolExecutor → asyncio.Semaphore + asyncio.gather
|
||
- LLMClient.from_env() 每题构造 → llm: LLMProvider 注入共享
|
||
- SentenceTransformer/OCR 内部构造 → 调用方通过 tool_dispatch_fn 注入
|
||
- run_id 必传，空串 → ValueError
|
||
- _aggregate_results 从内存 results 聚合（非 DB 回读）
|
||
- record_run 由调用方（Runner）负责
|
||
- prompt 构建由调用方注入 prompt_builder
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
from collections import defaultdict
|
||
from dataclasses import dataclass
|
||
from typing import TYPE_CHECKING, Any
|
||
|
||
from loguru import logger
|
||
|
||
from app.harness.question_units import build_units, unit_correctness
|
||
from core.agent.loop import AgentLoop
|
||
|
||
if TYPE_CHECKING:
|
||
from collections.abc import Callable
|
||
|
||
from app.harness.log import HarnessLog
|
||
from core.agent.types import LoopResult
|
||
from core.protocols import LLMProvider
|
||
from core.types import GeneratedQuestion, QuestionUnit
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class InferenceResult:
|
||
"""推理聚合结果。
|
||
|
||
属性:
|
||
run_id: 运行标识。
|
||
accuracy: 总正确率。
|
||
total: 总题数。
|
||
correct: 正确题数。
|
||
per_task_type: 按题型分组的指标 {task_type: {accuracy, total, correct}}。
|
||
steps_mean: 平均步数。
|
||
token_usage: token 总用量 {prompt_tokens, completion_tokens}。
|
||
stop_reason_counts: 终止原因计数 {reason: count}。
|
||
"""
|
||
|
||
run_id: str
|
||
accuracy: float
|
||
total: int
|
||
correct: int
|
||
per_task_type: dict[str, dict]
|
||
steps_mean: float
|
||
token_usage: dict[str, int]
|
||
stop_reason_counts: dict[str, int]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 表 Schema 定义（5 张表，保留 TRM4 全部 schema）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
PREDICTIONS_SCHEMA: dict[str, str] = {
|
||
"video_id": "TEXT",
|
||
"question_id": "TEXT",
|
||
"task_type": "TEXT",
|
||
"prediction": "TEXT",
|
||
"answer": "TEXT",
|
||
"evidence": "TEXT",
|
||
"reasoning": "TEXT",
|
||
"steps_used": "INTEGER",
|
||
"prompt_tokens": "INTEGER",
|
||
"completion_tokens": "INTEGER",
|
||
"stop_reason": "TEXT",
|
||
"steps_json": "JSON",
|
||
}
|
||
|
||
TRACES_SCHEMA: dict[str, str] = {
|
||
"video_id": "TEXT",
|
||
"question_id": "TEXT",
|
||
"step": "INTEGER",
|
||
"tool_name": "TEXT",
|
||
"tool_args": "JSON",
|
||
"tool_output": "TEXT",
|
||
"thought": "TEXT",
|
||
}
|
||
|
||
VALIDATION_FLAGS_SCHEMA: dict[str, str] = {
|
||
"video_id": "TEXT",
|
||
"question_id": "TEXT",
|
||
"has_l3_visit": "INTEGER",
|
||
"l1_count": "INTEGER",
|
||
"l2_count": "INTEGER",
|
||
"l3_count": "INTEGER",
|
||
}
|
||
|
||
ANCHOR_CHECK_SCHEMA: dict[str, str] = {
|
||
"video_id": "TEXT",
|
||
"question_id": "TEXT",
|
||
"step": "INTEGER",
|
||
"n_assertions": "INTEGER",
|
||
"n_anchored": "INTEGER",
|
||
"n_illegal": "INTEGER",
|
||
"n_expanded": "INTEGER",
|
||
"n_trunc": "INTEGER",
|
||
"output_chars": "INTEGER",
|
||
}
|
||
|
||
OF_HEALTH_SCHEMA: dict[str, str] = {
|
||
"video_id": "TEXT",
|
||
"question_id": "TEXT",
|
||
"step": "INTEGER",
|
||
"ocr_injected": "INTEGER",
|
||
"ocr_chars": "INTEGER",
|
||
"ocr_failed": "INTEGER",
|
||
"discrepancy": "INTEGER",
|
||
"abstain": "INTEGER",
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 内部工具
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class _DispatcherAdapter:
|
||
"""将裸 async callable 包装为 ToolDispatcher Protocol 实例。
|
||
|
||
AgentLoop 要求 ToolDispatcher（有 dispatch 方法），而 run_inference
|
||
接收的 tool_dispatch_fn 是裸 async callable。此适配器桥接两者。
|
||
|
||
参数:
|
||
fn: async def (tool_name, args, *, context) -> str。
|
||
"""
|
||
|
||
def __init__(self, fn: Callable[..., Any]) -> None:
|
||
self._fn = fn
|
||
|
||
async def dispatch(
|
||
self, tool_name: str, args: dict[str, Any], *, context: dict[str, Any]
|
||
) -> str:
|
||
"""转发工具调用给被包装的 callable。"""
|
||
return await self._fn(tool_name, args, context=context)
|
||
|
||
|
||
def _to_text_field(value: Any) -> str:
|
||
"""把 prediction 的 evidence/reasoning 归一为可入库的文本。
|
||
|
||
LLM 有时把这些字段返回成 list 或 dict（而非字符串）。sqlite 无法绑定
|
||
非标量类型，直接入库会抛 ProgrammingError 致该题丢失预测行、进而触发
|
||
rollout 完整性护栏中止整轮。凡非 str 一律 JSON 序列化为文本。
|
||
|
||
参数:
|
||
value: evidence/reasoning 原始值（可能是 str/list/dict）。
|
||
|
||
返回:
|
||
可直接入库的字符串。
|
||
"""
|
||
if isinstance(value, str):
|
||
return value
|
||
return json.dumps(value, ensure_ascii=False)
|
||
|
||
|
||
def _zero_result(run_id: str) -> InferenceResult:
|
||
"""空记录时的零值 InferenceResult。
|
||
|
||
参数:
|
||
run_id: 运行标识。
|
||
|
||
返回:
|
||
全零的 InferenceResult。
|
||
"""
|
||
return InferenceResult(
|
||
run_id=run_id,
|
||
accuracy=0.0,
|
||
total=0,
|
||
correct=0,
|
||
per_task_type={},
|
||
steps_mean=0.0,
|
||
token_usage={"prompt_tokens": 0, "completion_tokens": 0},
|
||
stop_reason_counts={},
|
||
)
|
||
|
||
|
||
def _group_by_task_type(graded: list[tuple[QuestionUnit, bool]]) -> dict[str, dict[str, Any]]:
|
||
"""按 task_type 分组聚合 unit 级正确率指标。
|
||
|
||
pair 单元整体计 1 个 unit，归入其 task_type；single 单元计 1 个 unit。
|
||
|
||
参数:
|
||
graded: (单元, 该单元是否整体正确) 元组列表。
|
||
|
||
返回:
|
||
{task_type: {accuracy, total, correct}} 映射（unit 粒度）。
|
||
"""
|
||
task_groups: dict[str, list[bool]] = defaultdict(list)
|
||
for unit, is_correct in graded:
|
||
task_groups[unit.task_type].append(is_correct)
|
||
|
||
per_task_type: dict[str, dict[str, Any]] = {}
|
||
for task_type, verdicts in task_groups.items():
|
||
t_total = len(verdicts)
|
||
t_correct = sum(verdicts)
|
||
per_task_type[task_type] = {
|
||
"accuracy": t_correct / t_total,
|
||
"total": t_total,
|
||
"correct": t_correct,
|
||
}
|
||
return per_task_type
|
||
|
||
|
||
def _drop_orphan_pairs(questions: list[GeneratedQuestion]) -> list[GeneratedQuestion]:
|
||
"""剔除收不齐 2 条 / 角色非法的孤儿 pair，告警不静默。
|
||
|
||
每条题目均会各答一次并逐题落库；能否合成 pair 单元仅取决于 questions
|
||
是否同时含该 pair_id 的 original + mirror。收不齐者告警并整对剔除，使
|
||
后续 build_units 只面对合法孪生对（不触发 fail-fast），孤儿 unit 不计入
|
||
total（对齐设计 §8 聚合入口的"告警 + 剔除"）。
|
||
|
||
参数:
|
||
questions: 待聚合的题目列表（可混含 single 与孪生对成员）。
|
||
|
||
返回:
|
||
可安全交给 build_units 的题目列表（single 全保留，pair 仅保留成对者）。
|
||
"""
|
||
by_pair: dict[str, list[GeneratedQuestion]] = defaultdict(list)
|
||
singles: list[GeneratedQuestion] = []
|
||
for q in questions:
|
||
if q.pair_id:
|
||
by_pair[q.pair_id].append(q)
|
||
else:
|
||
singles.append(q)
|
||
|
||
kept_pairs: list[GeneratedQuestion] = []
|
||
for pair_id, group in by_pair.items():
|
||
originals = [q for q in group if q.question_role == "pair_original"]
|
||
mirrors = [q for q in group if q.question_role == "pair_mirror"]
|
||
if len(originals) == 1 and len(mirrors) == 1:
|
||
kept_pairs.extend(group)
|
||
else:
|
||
logger.warning(
|
||
"孤儿 pair {}：收不齐 2 条（original={} mirror={}），剔除该 unit 不计入 total",
|
||
pair_id,
|
||
len(originals),
|
||
len(mirrors),
|
||
)
|
||
return singles + kept_pairs
|
||
|
||
|
||
def _per_question_correctness(records: list[dict[str, Any]]) -> dict[str, bool]:
|
||
"""由逐题 record 构造 question_id → 该题作答是否正确 的映射。
|
||
|
||
prediction 为 None（作答异常）时与 answer 不相等 → False，天然计错。
|
||
|
||
参数:
|
||
records: _run_single_question 返回的逐题 record 列表。
|
||
|
||
返回:
|
||
{question_id: prediction == answer} 映射，供 unit_correctness 取值。
|
||
"""
|
||
return {r["question_id"]: r["prediction"] == r["answer"] for r in records}
|
||
|
||
|
||
def _aggregate_results(
|
||
records: list[dict[str, Any]],
|
||
questions: list[GeneratedQuestion],
|
||
run_id: str,
|
||
) -> InferenceResult:
|
||
"""从内存 records + 题目列表按 unit 粒度聚合推理指标。
|
||
|
||
逐题 record 保留逐题溯源（token/steps/stop_reason 诊断仍按 record 汇总）；
|
||
正确率则按 unit 粒度计：single 计 1，AR pair 经 build_units 收齐 original +
|
||
mirror 后走 unit_correctness 的双向 AND 判定，整对计 1 个 unit。孤儿 pair
|
||
在 _drop_orphan_pairs 中告警 + 剔除，不计入 total。
|
||
|
||
参数:
|
||
records: _run_single_question 返回的逐题 record 列表。
|
||
questions: 与 records 对应的题目列表（提供 pair_id/question_role 元数据）。
|
||
run_id: 当前运行标识。
|
||
|
||
返回:
|
||
InferenceResult 冻结实例（total/correct/per_task_type 为 unit 粒度）。
|
||
"""
|
||
if not records:
|
||
return _zero_result(run_id)
|
||
|
||
per_q = _per_question_correctness(records)
|
||
units = build_units(_drop_orphan_pairs(questions))
|
||
# unit 内任一题缺 record → unit_correctness 抛 KeyError（防静默兜底/读回校验）。
|
||
graded = [(unit, unit_correctness(unit, per_q)) for unit in units]
|
||
|
||
total = len(graded)
|
||
correct = sum(1 for _, is_correct in graded if is_correct)
|
||
|
||
stop_counts: dict[str, int] = defaultdict(int)
|
||
for r in records:
|
||
stop_counts[r["stop_reason"]] += 1
|
||
|
||
n_records = len(records)
|
||
return InferenceResult(
|
||
run_id=run_id,
|
||
accuracy=correct / total if total else 0.0,
|
||
total=total,
|
||
correct=correct,
|
||
per_task_type=_group_by_task_type(graded),
|
||
steps_mean=sum(r["steps_used"] for r in records) / n_records,
|
||
token_usage={
|
||
"prompt_tokens": sum(r["prompt_tokens"] for r in records),
|
||
"completion_tokens": sum(r["completion_tokens"] for r in records),
|
||
},
|
||
stop_reason_counts=dict(stop_counts),
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 单题推理
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def _run_single_question(
|
||
qa: GeneratedQuestion,
|
||
*,
|
||
llm: LLMProvider,
|
||
tool_dispatch_fn: Callable[..., Any],
|
||
prompt_builder: Callable[[GeneratedQuestion], tuple[str, str]],
|
||
log: HarnessLog,
|
||
max_steps: int,
|
||
plugins: list[object],
|
||
) -> dict[str, Any]:
|
||
"""执行单道题目的 Agent 推理。
|
||
|
||
悲观默认值：record 初始 stop_reason="error"，成功后覆盖。
|
||
prediction 必落库：log.insert 在 try/except 之后（无论成败）。
|
||
|
||
参数:
|
||
qa: 待推理的题目。
|
||
llm: LLMProvider 共享实例。
|
||
tool_dispatch_fn: async 工具调度函数 (tool_name, args, *, context) -> str。
|
||
prompt_builder: (GeneratedQuestion) -> (system_prompt, user_prompt)。
|
||
log: HarnessLog 实例（线程安全）。
|
||
max_steps: AgentLoop 最大步数。
|
||
plugins: pluggy 插件列表。
|
||
|
||
返回:
|
||
预测结果字典（含 video_id, question_id, prediction, answer 等）。
|
||
"""
|
||
record: dict[str, Any] = {
|
||
"video_id": qa.video_id,
|
||
"question_id": qa.question_id,
|
||
"task_type": qa.task_type,
|
||
"prediction": None,
|
||
"answer": qa.answer,
|
||
"evidence": "",
|
||
"reasoning": "",
|
||
"steps_used": 0,
|
||
"prompt_tokens": 0,
|
||
"completion_tokens": 0,
|
||
"stop_reason": "error", # 悲观默认
|
||
"steps_json": "[]",
|
||
}
|
||
|
||
try:
|
||
system_prompt, user_prompt = prompt_builder(qa)
|
||
dispatcher = _DispatcherAdapter(tool_dispatch_fn)
|
||
loop = AgentLoop(llm, max_steps=max_steps)
|
||
loop_result: LoopResult = await loop.run(
|
||
system_prompt,
|
||
user_prompt,
|
||
dispatcher,
|
||
plugins=plugins,
|
||
session_id=qa.question_id,
|
||
)
|
||
|
||
result_dict = loop_result.result if isinstance(loop_result.result, dict) else {}
|
||
evidence = _to_text_field(result_dict.get("evidence", ""))
|
||
reasoning = _to_text_field(result_dict.get("reasoning", ""))
|
||
record.update(
|
||
{
|
||
"prediction": result_dict.get("answer"),
|
||
"evidence": evidence,
|
||
"reasoning": reasoning,
|
||
"steps_used": loop_result.steps_used,
|
||
"prompt_tokens": loop_result.token_usage["prompt_tokens"],
|
||
"completion_tokens": loop_result.token_usage["completion_tokens"],
|
||
"stop_reason": loop_result.stop_reason,
|
||
"steps_json": json.dumps(
|
||
[
|
||
{
|
||
"thought": s.thought,
|
||
"tool_call": s.tool_call,
|
||
"tool_output": s.tool_output,
|
||
}
|
||
for s in loop_result.steps
|
||
],
|
||
ensure_ascii=False,
|
||
),
|
||
}
|
||
)
|
||
except Exception:
|
||
logger.exception("[{}] QA {} 执行异常", qa.video_id, qa.question_id)
|
||
|
||
# prediction 必落库（try 外，无论成败）
|
||
await asyncio.to_thread(log.insert, "predictions", record)
|
||
return record
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 建表
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _ensure_tables(log: HarnessLog) -> None:
|
||
"""创建推理所需的 5 张表。
|
||
|
||
参数:
|
||
log: HarnessLog 实例。
|
||
"""
|
||
log.create_table("predictions", PREDICTIONS_SCHEMA)
|
||
log.create_table("traces", TRACES_SCHEMA)
|
||
log.create_table("validation_flags", VALIDATION_FLAGS_SCHEMA)
|
||
log.create_table("anchor_check", ANCHOR_CHECK_SCHEMA)
|
||
log.create_table("observe_frame_health", OF_HEALTH_SCHEMA)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 公共入口
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def run_inference(
|
||
questions: list[GeneratedQuestion],
|
||
*,
|
||
llm: LLMProvider,
|
||
tool_dispatch_fn: Callable[..., Any],
|
||
prompt_builder: Callable[[GeneratedQuestion], tuple[str, str]],
|
||
log: HarnessLog,
|
||
run_id: str,
|
||
concurrency: int,
|
||
max_steps: int,
|
||
skill_mode: str,
|
||
plugins_factory: Callable[[str, str], list[object]] | None = None,
|
||
) -> InferenceResult:
|
||
"""在视频树上执行 Agent 推理，对应训练循环的 forward()。
|
||
|
||
参数:
|
||
questions: 待推理的题目列表。
|
||
llm: LLMProvider 共享实例（依赖注入）。
|
||
tool_dispatch_fn: async 工具调度函数 (tool_name, args, *, context) -> str。
|
||
prompt_builder: prompt 构建函数 (GeneratedQuestion) -> (system_prompt, user_prompt)。
|
||
log: HarnessLog 实例（由调用方管理生命周期）。
|
||
run_id: 运行标识（必传，空串 → ValueError）。
|
||
concurrency: 最大并发数（asyncio.Semaphore 控制）。
|
||
max_steps: AgentLoop 单题最大步数。
|
||
skill_mode: "auto" / "manual" / "none"（传递给调用方的 prompt/plugin 构建逻辑）。
|
||
plugins_factory: 可选的插件工厂 (video_id, question_id) -> plugins 列表。
|
||
|
||
返回:
|
||
InferenceResult（含 accuracy、per_task_type 等聚合指标）。
|
||
|
||
异常:
|
||
ValueError: run_id 为空串或纯空白。
|
||
"""
|
||
if not run_id or not run_id.strip():
|
||
raise ValueError("run_id 不得为空串或纯空白")
|
||
|
||
_ensure_tables(log)
|
||
|
||
if not questions:
|
||
logger.info("题目列表为空，返回零值 InferenceResult")
|
||
return _aggregate_results([], [], run_id)
|
||
|
||
sem = asyncio.Semaphore(concurrency)
|
||
total_count = len(questions)
|
||
|
||
async def _bounded(index: int, qa: GeneratedQuestion) -> dict[str, Any]:
|
||
"""信号量限流的单题推理包装。"""
|
||
async with sem:
|
||
plugins = (
|
||
plugins_factory(qa.video_id, qa.question_id) if plugins_factory is not None else []
|
||
)
|
||
result = await _run_single_question(
|
||
qa,
|
||
llm=llm,
|
||
tool_dispatch_fn=tool_dispatch_fn,
|
||
prompt_builder=prompt_builder,
|
||
log=log,
|
||
max_steps=max_steps,
|
||
plugins=plugins,
|
||
)
|
||
logger.info(
|
||
"[{}/{}] {} QA {} 完成 (stop={})",
|
||
index + 1,
|
||
total_count,
|
||
qa.video_id,
|
||
qa.question_id,
|
||
result["stop_reason"],
|
||
)
|
||
return result
|
||
|
||
results = await asyncio.gather(*[_bounded(i, qa) for i, qa in enumerate(questions)])
|
||
|
||
inference_result = _aggregate_results(list(results), questions, run_id)
|
||
logger.info(
|
||
"推理完成: accuracy={:.2%} ({}/{})",
|
||
inference_result.accuracy,
|
||
inference_result.correct,
|
||
inference_result.total,
|
||
)
|
||
return inference_result
|