feat(harness): inference.py — async run_inference + DI

This commit is contained in:
2026-07-07 13:04:26 -04:00
parent 886a444d1d
commit d7f1bdeea6
2 changed files with 1100 additions and 0 deletions
+441
View File
@@ -0,0 +1,441 @@
"""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
+659
View File
@@ -0,0 +1,659 @@
"""app/harness/inference.py 单元测试。
测试覆盖:
- run_inference 基本流程(mock LLM + tool_dispatch
- 异常时 prediction 仍落库(stop_reason=error
- _to_text_field 归一化
- run_id 空串 → ValueError
- _aggregate_results 内存聚合
- 空 questions 列表零值返回
- 并发控制 Semaphore
- plugins_factory 调用
"""
from __future__ import annotations
import json
from typing import Any
from unittest.mock import AsyncMock
import pytest
from app.harness.inference import (
InferenceResult,
_aggregate_results,
_to_text_field,
run_inference,
)
from app.harness.log import HarnessLog
from core.types import GeneratedQuestion, LLMResponse
# ── 测试基础设施 ──────────────────────────────────────────────────
def _make_question(
question_id: str = "q1",
video_id: str = "v1",
task_type: str = "Action Reasoning",
answer: str = "B",
) -> GeneratedQuestion:
"""构造测试用题目。"""
return GeneratedQuestion(
question_id=question_id,
video_id=video_id,
task_type=task_type,
question="测试问题",
options=("A. 选项A", "B. 选项B", "C. 选项C", "D. 选项D"),
answer=answer,
source_nodes=("L1_001",),
difficulty="medium",
)
def _make_llm_response(answer: str = "B") -> LLMResponse:
"""构造测试用 LLMResponsesubmit_answer 场景)。"""
content = json.dumps(
{
"reflect": {"observation": "找到答案"},
"plan": {"next_step": "提交"},
"action": {
"tool": "submit_answer",
"args": {
"answer": answer,
"evidence": "证据文本",
"reasoning": "推理过程",
},
},
}
)
return LLMResponse(
content=content,
thinking="思考过程",
model="test-model",
provider="test",
prompt_tokens=100,
completion_tokens=50,
latency_ms=200,
ttft_ms=30.0,
max_inter_token_ms=5.0,
cache_hit=False,
call_id="test-call-001",
)
def _make_error_llm_response() -> LLMResponse:
"""构造触发解析失败的 LLMResponse。"""
return LLMResponse(
content="这不是JSON",
thinking="",
model="test-model",
provider="test",
prompt_tokens=10,
completion_tokens=5,
latency_ms=50,
ttft_ms=10.0,
max_inter_token_ms=2.0,
cache_hit=False,
call_id="test-call-err",
)
async def _stub_tool_dispatch(
tool_name: str, args: dict[str, Any], *, context: dict[str, Any]
) -> str:
"""测试用工具调度函数。"""
if tool_name == "submit_answer":
return "答案已提交"
raise ValueError(f"未知工具: {tool_name}")
def _stub_prompt_builder(qa: GeneratedQuestion) -> tuple[str, str]:
"""测试用 prompt 构建函数。"""
return "系统提示词", f"用户问题: {qa.question}"
@pytest.fixture
def harness_log(tmp_path: Any, request: Any) -> HarnessLog:
"""创建临时 HarnessLog 实例。
使用 test 节点名称的 hash 作为 db 文件名,避免冲突。
run_id 固定为 "test-run",实际 run_inference 中传入的 run_id
由 HarnessLog.insert 自动覆盖为 HarnessLog 构造时的值。
"""
db_name = f"harness_{id(request)}.db"
db_path = str(tmp_path / db_name)
log = HarnessLog(db_path, "test-run")
yield log
log.close()
# ── 测试用例 ──────────────────────────────────────────────────
class TestToTextField:
"""_to_text_field 归一化测试。"""
@pytest.mark.asyncio
async def test_string_passthrough(self) -> None:
"""字符串原样返回。"""
assert _to_text_field("hello") == "hello"
@pytest.mark.asyncio
async def test_empty_string(self) -> None:
"""空字符串原样返回。"""
assert _to_text_field("") == ""
@pytest.mark.asyncio
async def test_list_serialized(self) -> None:
"""list 被 JSON 序列化。"""
result = _to_text_field(["a", "b"])
assert result == '["a", "b"]'
@pytest.mark.asyncio
async def test_dict_serialized(self) -> None:
"""dict 被 JSON 序列化。"""
result = _to_text_field({"key": ""})
assert '"key"' in result
assert '""' in result
@pytest.mark.asyncio
async def test_int_serialized(self) -> None:
"""int 被 JSON 序列化。"""
assert _to_text_field(42) == "42"
@pytest.mark.asyncio
async def test_none_serialized(self) -> None:
"""None 被 JSON 序列化。"""
assert _to_text_field(None) == "null"
@pytest.mark.asyncio
async def test_unicode_preserved(self) -> None:
"""ensure_ascii=False 保留中文。"""
result = _to_text_field(["中文"])
assert "中文" in result
assert "\\u" not in result
class TestAggregateResults:
"""_aggregate_results 内存聚合测试。"""
@pytest.mark.asyncio
async def test_empty_records(self) -> None:
"""空列表返回零值 InferenceResult。"""
result = _aggregate_results([], "run-empty")
assert result.run_id == "run-empty"
assert result.accuracy == 0.0
assert result.total == 0
assert result.correct == 0
assert result.per_task_type == {}
assert result.steps_mean == 0.0
assert result.token_usage == {"prompt_tokens": 0, "completion_tokens": 0}
assert result.stop_reason_counts == {}
@pytest.mark.asyncio
async def test_single_correct(self) -> None:
"""单条正确记录 → accuracy=1.0。"""
records = [
{
"prediction": "B",
"answer": "B",
"task_type": "AR",
"steps_used": 3,
"prompt_tokens": 100,
"completion_tokens": 50,
"stop_reason": "finished",
}
]
result = _aggregate_results(records, "run-1")
assert result.accuracy == 1.0
assert result.total == 1
assert result.correct == 1
assert result.steps_mean == 3.0
@pytest.mark.asyncio
async def test_mixed_correct_wrong(self) -> None:
"""混合正确/错误 → 准确率与步数均正确聚合。"""
records = [
{
"prediction": "B",
"answer": "B",
"task_type": "AR",
"steps_used": 2,
"prompt_tokens": 100,
"completion_tokens": 50,
"stop_reason": "finished",
},
{
"prediction": "C",
"answer": "A",
"task_type": "AR",
"steps_used": 4,
"prompt_tokens": 200,
"completion_tokens": 100,
"stop_reason": "budget_exceeded",
},
{
"prediction": "D",
"answer": "D",
"task_type": "SP",
"steps_used": 1,
"prompt_tokens": 50,
"completion_tokens": 25,
"stop_reason": "finished",
},
]
result = _aggregate_results(records, "run-mix")
assert result.total == 3
assert result.correct == 2
assert abs(result.accuracy - 2 / 3) < 1e-9
assert abs(result.steps_mean - 7 / 3) < 1e-9
assert result.token_usage == {"prompt_tokens": 350, "completion_tokens": 175}
assert result.stop_reason_counts == {"finished": 2, "budget_exceeded": 1}
@pytest.mark.asyncio
async def test_per_task_type_grouping(self) -> None:
"""按 task_type 分组聚合。"""
records = [
{
"prediction": "B",
"answer": "B",
"task_type": "AR",
"steps_used": 1,
"prompt_tokens": 10,
"completion_tokens": 5,
"stop_reason": "finished",
},
{
"prediction": "A",
"answer": "C",
"task_type": "AR",
"steps_used": 2,
"prompt_tokens": 20,
"completion_tokens": 10,
"stop_reason": "finished",
},
{
"prediction": "D",
"answer": "D",
"task_type": "SP",
"steps_used": 3,
"prompt_tokens": 30,
"completion_tokens": 15,
"stop_reason": "finished",
},
]
result = _aggregate_results(records, "run-task")
assert "AR" in result.per_task_type
assert "SP" in result.per_task_type
assert result.per_task_type["AR"]["total"] == 2
assert result.per_task_type["AR"]["correct"] == 1
assert result.per_task_type["AR"]["accuracy"] == 0.5
assert result.per_task_type["SP"]["total"] == 1
assert result.per_task_type["SP"]["correct"] == 1
assert result.per_task_type["SP"]["accuracy"] == 1.0
class TestRunIdValidation:
"""run_id 校验测试。"""
@pytest.mark.asyncio
async def test_empty_string_raises(self, harness_log: HarnessLog) -> None:
"""空串 run_id → ValueError。"""
llm = AsyncMock()
with pytest.raises(ValueError, match="run_id 不得为空"):
await run_inference(
[],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
@pytest.mark.asyncio
async def test_whitespace_only_raises(self, harness_log: HarnessLog) -> None:
"""纯空白 run_id → ValueError。"""
llm = AsyncMock()
with pytest.raises(ValueError, match="run_id 不得为空"):
await run_inference(
[],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id=" ",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
class TestEmptyQuestions:
"""空题目列表测试。"""
@pytest.mark.asyncio
async def test_empty_questions_returns_zero(self, harness_log: HarnessLog) -> None:
"""空 questions 列表直接返回零值 InferenceResult。"""
llm = AsyncMock()
result = await run_inference(
[],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-empty",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert isinstance(result, InferenceResult)
assert result.run_id == "run-empty"
assert result.accuracy == 0.0
assert result.total == 0
assert result.correct == 0
# LLM 未被调用
llm.chat.assert_not_called()
class TestRunInferenceBasic:
"""run_inference 基本流程测试。"""
@pytest.mark.asyncio
async def test_single_question_correct(self, harness_log: HarnessLog) -> None:
"""单题正确推理 → accuracy=1.0, stop_reason=finished。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="B")
result = await run_inference(
[_make_question(answer="B")],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-basic",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert result.accuracy == 1.0
assert result.total == 1
assert result.correct == 1
assert result.stop_reason_counts.get("finished") == 1
@pytest.mark.asyncio
async def test_single_question_wrong(self, harness_log: HarnessLog) -> None:
"""单题错误推理 → accuracy=0.0。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="C")
result = await run_inference(
[_make_question(answer="B")],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-wrong",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert result.accuracy == 0.0
assert result.total == 1
assert result.correct == 0
@pytest.mark.asyncio
async def test_multiple_questions_concurrent(self, harness_log: HarnessLog) -> None:
"""3 题并发推理 → 结果正确聚合。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="B")
questions = [
_make_question(question_id="q1", answer="B"),
_make_question(question_id="q2", answer="B"),
_make_question(question_id="q3", answer="A"),
]
result = await run_inference(
questions,
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-multi",
concurrency=3,
max_steps=10,
skill_mode="auto",
)
assert result.total == 3
assert result.correct == 2
assert abs(result.accuracy - 2 / 3) < 1e-9
@pytest.mark.asyncio
async def test_token_usage_accumulated(self, harness_log: HarnessLog) -> None:
"""多题 token 累加验证。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="B")
questions = [
_make_question(question_id="q1"),
_make_question(question_id="q2"),
]
result = await run_inference(
questions,
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-token",
concurrency=2,
max_steps=10,
skill_mode="auto",
)
assert result.token_usage["prompt_tokens"] == 200
assert result.token_usage["completion_tokens"] == 100
class TestPredictionAlwaysWritten:
"""异常时 prediction 仍落库测试。"""
@pytest.mark.asyncio
async def test_error_still_persisted(self, harness_log: HarnessLog) -> None:
"""LLM 调用异常时,prediction 仍以 stop_reason=error 落库。"""
llm = AsyncMock()
llm.chat.side_effect = RuntimeError("LLM API 不可用")
result = await run_inference(
[_make_question()],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-error",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert result.total == 1
assert result.correct == 0
assert result.stop_reason_counts.get("error") == 1
# 验证 DB 中的记录(HarnessLog.insert 使用构造时的 run_id
rows = harness_log.query("SELECT * FROM predictions WHERE run_id = ?", ("test-run",))
assert len(rows) == 1
assert rows[0]["stop_reason"] == "error"
assert rows[0]["prediction"] is None
@pytest.mark.asyncio
async def test_parse_error_still_persisted(self, harness_log: HarnessLog) -> None:
"""LLM 返回非 JSON 内容,parse_error 后 prediction 仍落库。"""
llm = AsyncMock()
llm.chat.return_value = _make_error_llm_response()
result = await run_inference(
[_make_question()],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-parse-err",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert result.total == 1
# HarnessLog.insert 使用构造时的 run_id
rows = harness_log.query("SELECT * FROM predictions WHERE run_id = ?", ("test-run",))
assert len(rows) == 1
assert rows[0]["prediction"] is None
class TestPluginsFactory:
"""plugins_factory 调用测试。"""
@pytest.mark.asyncio
async def test_factory_called_per_question(self, harness_log: HarnessLog) -> None:
"""每题调用 plugins_factory,传入 (video_id, question_id)。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="B")
factory_calls: list[tuple[str, str]] = []
def _factory(video_id: str, question_id: str) -> list[object]:
factory_calls.append((video_id, question_id))
return []
questions = [
_make_question(question_id="q1", video_id="v1"),
_make_question(question_id="q2", video_id="v2"),
]
await run_inference(
questions,
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-factory",
concurrency=2,
max_steps=10,
skill_mode="auto",
plugins_factory=_factory,
)
assert len(factory_calls) == 2
call_set = set(factory_calls)
assert ("v1", "q1") in call_set
assert ("v2", "q2") in call_set
@pytest.mark.asyncio
async def test_no_factory_uses_empty_plugins(self, harness_log: HarnessLog) -> None:
"""plugins_factory=None 时使用空 plugins 列表。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="B")
result = await run_inference(
[_make_question()],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-no-factory",
concurrency=1,
max_steps=10,
skill_mode="auto",
plugins_factory=None,
)
assert result.total == 1
assert result.stop_reason_counts.get("finished") == 1
class TestConcurrencyControl:
"""并发控制 Semaphore 测试。"""
@pytest.mark.asyncio
async def test_concurrency_semaphore_limits(self, harness_log: HarnessLog) -> None:
"""Semaphore(1) 限制并发为 1 — 通过最大并发计数器验证。"""
import asyncio
llm = AsyncMock()
current_concurrent = 0
max_concurrent = 0
original_response = _make_llm_response(answer="B")
async def _slow_chat(
messages: Any,
*,
session_id: str | None = None,
parent_call_id: str | None = None,
) -> LLMResponse:
nonlocal current_concurrent, max_concurrent
current_concurrent += 1
max_concurrent = max(max_concurrent, current_concurrent)
await asyncio.sleep(0.01)
current_concurrent -= 1
return original_response
llm.chat.side_effect = _slow_chat
questions = [_make_question(question_id=f"q{i}") for i in range(5)]
await run_inference(
questions,
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-sem",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert max_concurrent == 1
class TestTablesCreated:
"""表创建测试。"""
@pytest.mark.asyncio
async def test_five_tables_created(self, harness_log: HarnessLog) -> None:
"""run_inference 启动时创建 5 张推理表。"""
llm = AsyncMock()
await run_inference(
[],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-tables",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
expected_tables = [
"predictions",
"traces",
"validation_flags",
"anchor_check",
"observe_frame_health",
]
for table_name in expected_tables:
rows = harness_log.query(
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
(table_name,),
)
assert len(rows) == 1, f"{table_name} 未创建"