125 lines
4.6 KiB
Python
125 lines
4.6 KiB
Python
"""压测语料装载(纯函数): 真实库还原请求负载,不做任何网络调用。
|
|
|
|
数据源(data/soak/,不入 git):
|
|
- harness.db predictions.steps_json: VT agent 完整轨迹 → P1 长上下文累积链
|
|
- generate_questions_telemetry.db llm_calls.messages: 376 条真实重载 → P2 回放
|
|
- vt_frames/<video>/frames/*.jpg: 多模态组装;chs_images/chs_NNNN.jpg: 单图短指令
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
Messages = list[dict[str, Any]]
|
|
|
|
_TRACE_SYSTEM = (
|
|
"You are a video-understanding agent replayed for gateway soak testing. "
|
|
"Answer the tool-use conversation faithfully and concisely."
|
|
)
|
|
|
|
# CHS 形态的固定结构化短指令(业务 schema 只存在于 tools/,不入 src——零业务假设铁律)
|
|
CHS_INSTRUCTION = (
|
|
"Extract the following 12 fields from this ultrasound image and reply in strict JSON: "
|
|
"modality, orientation, depth_cm, gain, focus_zone, vessel_visible, lesion_present, "
|
|
"lesion_location, echo_pattern, doppler_signal, image_quality, notes. "
|
|
"Reply with a single JSON object only."
|
|
)
|
|
|
|
|
|
def load_trace_chains(harness_db: Path | str, *, max_steps: int = 40) -> list[list[Messages]]:
|
|
"""还原 P1 长上下文链: 每条 prediction → 逐步累积的 messages 快照序列。
|
|
|
|
快照 k = system + user(任务行)+ 前 k 步的 (assistant: thought+tool_call,
|
|
user: tool_output);steps_used=0 或 steps_json 坏行剔除。
|
|
"""
|
|
conn = sqlite3.connect(harness_db)
|
|
try:
|
|
rows = conn.execute(
|
|
"SELECT video_id, question_id, task_type, steps_json FROM predictions"
|
|
).fetchall()
|
|
finally:
|
|
conn.close()
|
|
chains: list[list[Messages]] = []
|
|
for video_id, question_id, task_type, steps_json in rows:
|
|
try:
|
|
steps = json.loads(steps_json or "[]")
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if not steps:
|
|
continue
|
|
base: Messages = [
|
|
{"role": "system", "content": _TRACE_SYSTEM},
|
|
{
|
|
"role": "user",
|
|
"content": f"[soak replay] video={video_id} question={question_id} task={task_type}",
|
|
},
|
|
]
|
|
snapshots: list[Messages] = []
|
|
for step in steps[:max_steps]:
|
|
base = base + [
|
|
{
|
|
"role": "assistant",
|
|
"content": f"{step.get('thought', '')}\n{json.dumps(step.get('tool_call', {}), ensure_ascii=False)}",
|
|
},
|
|
{"role": "user", "content": str(step.get("tool_output", ""))},
|
|
]
|
|
snapshots.append(base)
|
|
chains.append(snapshots)
|
|
return chains
|
|
|
|
|
|
def load_replay_payloads(telemetry_db: Path | str) -> list[Messages]:
|
|
"""P2 回放负载: 原样取真实 messages(内嵌帧的多模态重载),坏 JSON 行剔除。"""
|
|
conn = sqlite3.connect(telemetry_db)
|
|
try:
|
|
rows = conn.execute("SELECT messages FROM llm_calls").fetchall()
|
|
finally:
|
|
conn.close()
|
|
payloads: list[Messages] = []
|
|
for (raw,) in rows:
|
|
try:
|
|
messages = json.loads(raw)
|
|
except (json.JSONDecodeError, TypeError):
|
|
continue
|
|
if isinstance(messages, list) and messages:
|
|
payloads.append(messages)
|
|
return payloads
|
|
|
|
|
|
def _data_uri(path: Path) -> str:
|
|
return "data:image/jpeg;base64," + base64.b64encode(path.read_bytes()).decode("ascii")
|
|
|
|
|
|
def assemble_frame_messages(frames_root: Path | str, *, n_frames: int, rng) -> Messages:
|
|
"""VT 帧组装(P2 两档: 1-4 帧 / 5-6 帧): 随机选视频取前 n 帧 + 简短指令。"""
|
|
root = Path(frames_root)
|
|
videos = sorted(d for d in root.iterdir() if (d / "frames").is_dir())
|
|
if not videos:
|
|
raise ValueError(f"{root} 下没有 <video>/frames/ 目录")
|
|
video = videos[int(rng() * len(videos)) % len(videos)]
|
|
frames = sorted((video / "frames").glob("*.jpg"))[:n_frames]
|
|
if len(frames) < n_frames:
|
|
raise ValueError(f"{video} 帧数不足 {n_frames}")
|
|
parts: list[dict[str, Any]] = [
|
|
{"type": "text", "text": f"Describe the key visual content of these {n_frames} frames."}
|
|
]
|
|
parts.extend({"type": "image_url", "image_url": {"url": _data_uri(f)}} for f in frames)
|
|
return [{"role": "user", "content": parts}]
|
|
|
|
|
|
def chs_image_messages(image_path: Path | str) -> Messages:
|
|
"""CHS 形态(P3): 单图 + 固定 12 字段结构化短指令。"""
|
|
return [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": CHS_INSTRUCTION},
|
|
{"type": "image_url", "image_url": {"url": _data_uri(Path(image_path))}},
|
|
],
|
|
}
|
|
]
|