442 lines
14 KiB
Python
442 lines
14 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 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
|
||
|
||
|
||
@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(records: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
|
||
"""按 task_type 分组聚合正确率指标。
|
||
|
||
参数:
|
||
records: 预测记录列表。
|
||
|
||
返回:
|
||
{task_type: {accuracy, total, correct}} 映射。
|
||
"""
|
||
task_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||
for r in records:
|
||
task_groups[r["task_type"]].append(r)
|
||
|
||
per_task_type: dict[str, dict[str, Any]] = {}
|
||
for task_type, group in task_groups.items():
|
||
t_total = len(group)
|
||
t_correct = sum(1 for r in group if r["prediction"] == r["answer"])
|
||
per_task_type[task_type] = {
|
||
"accuracy": t_correct / t_total,
|
||
"total": t_total,
|
||
"correct": t_correct,
|
||
}
|
||
return per_task_type
|
||
|
||
|
||
def _aggregate_results(records: list[dict[str, Any]], run_id: str) -> InferenceResult:
|
||
"""从内存 records 聚合推理指标。
|
||
|
||
TRM4 从 DB 回读 predictions 表聚合；TRM5 改为从内存直接聚合，
|
||
避免 DB 回读的同步开销和额外依赖。
|
||
|
||
参数:
|
||
records: _run_single_question 返回的 record 列表。
|
||
run_id: 当前运行标识。
|
||
|
||
返回:
|
||
InferenceResult 冻结实例。
|
||
"""
|
||
total = len(records)
|
||
if total == 0:
|
||
return _zero_result(run_id)
|
||
|
||
correct = sum(1 for r in records if r["prediction"] == r["answer"])
|
||
stop_counts: dict[str, int] = defaultdict(int)
|
||
for r in records:
|
||
stop_counts[r["stop_reason"]] += 1
|
||
|
||
return InferenceResult(
|
||
run_id=run_id,
|
||
accuracy=correct / total,
|
||
total=total,
|
||
correct=correct,
|
||
per_task_type=_group_by_task_type(records),
|
||
steps_mean=sum(r["steps_used"] for r in records) / total,
|
||
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), run_id)
|
||
logger.info(
|
||
"推理完成: accuracy={:.2%} ({}/{})",
|
||
inference_result.accuracy,
|
||
inference_result.correct,
|
||
inference_result.total,
|
||
)
|
||
return inference_result
|