feat: add soak corpus loaders and scenario generators

2026-07-21 01:17:46 -04:00
parent 5e01dc738f
commit cca7071dbd
4 changed files with 449 additions and 0 deletions
@@ -0,0 +1,190 @@
+"""P1-P6 场景请求生成器(findings §2 矩阵;async 生成器,产出调用参数)。
+
+每项产出 `(kind, kwargs)`: kind ∈ {"chat"};kwargs 直接喂
+`GatewayClient.chat(**kwargs)`。回放/组装场景一律掺 `cache_salt=run_id`
+破缓存(缓存行为归 P4);P4 子流量特意重复 messages 且不掺 salt。
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, AsyncIterator
+
+from pydantic import BaseModel
+
+from tools.soak.corpus import (
+    assemble_frame_messages,
+    chs_image_messages,
+    load_replay_payloads,
+    load_trace_chains,
+)
+
+Item = tuple[str, dict[str, Any]]
+
+
+class ChsExtraction(BaseModel):
+    """P3 结构化档的真实形态 schema(12 字段,仅存在于 tools/)。"""
+
+    modality: str
+    orientation: str
+    depth_cm: float | None = None
+    gain: str | None = None
+    focus_zone: str | None = None
+    vessel_visible: bool | None = None
+    lesion_present: bool | None = None
+    lesion_location: str | None = None
+    echo_pattern: str | None = None
+    doppler_signal: str | None = None
+    image_quality: str
+    notes: str | None = None
+
+
+@dataclass
+class SoakCorpus:
+    """一次 run 的语料句柄(装载一次,场景间共享)。"""
+
+    harness_db: Path
+    telemetry_db: Path
+    frames_root: Path
+    images_root: Path
+
+    def __post_init__(self) -> None:
+        self.chains = load_trace_chains(self.harness_db)
+        self.replays = load_replay_payloads(self.telemetry_db)
+        self.images = sorted(Path(self.images_root).glob("chs_*.jpg"))
+        if not (self.chains and self.replays and self.images):
+            raise ValueError("语料不完整: 请确认 data/soak/ 已按 findings §6 拉取")
+
+
+def weighted_mix(weights: dict[str, float], rng) -> str:
+    """按权重抽一个场景 key;权重和必须 > 0。"""
+    total = sum(weights.values())
+    if total <= 0:
+        raise ValueError("场景权重之和必须 > 0")
+    point = rng() * total
+    acc = 0.0
+    for key, w in weights.items():
+        acc += w
+        if point < acc:
+            return key
+    return next(reversed(weights))
+
+
+async def p1_trace_chains(corpus: SoakCorpus, run_id: str, rng=random.random) -> AsyncIterator[Item]:
+    """P1 文本长上下文回放: 单链串行,session/parent 链路照原样语义。"""
+    for chain_idx, chain in enumerate(corpus.chains):
+        session_id = f"{run_id}-p1-{chain_idx}"
+        parent: str | None = None
+        for snapshot in chain:
+            yield (
+                "chat",
+                {
+                    "messages": snapshot,
+                    "session_id": session_id,
+                    "parent_call_id": parent,
+                    "cache_salt": run_id,
+                },
+            )
+            parent = session_id  # 链内父子: 以 session 为锚(真实 call_id 由库生成)
+
+
+async def p2_multimodal_replay(
+    corpus: SoakCorpus, run_id: str, rng=random.random
+) -> AsyncIterator[Item]:
+    """P2 多模态重载: 真实 376 条原样回放 + 帧组装两档交错。"""
+    for i, messages in enumerate(corpus.replays):
+        yield ("chat", {"messages": messages, "cache_salt": run_id, "session_id": f"{run_id}-p2"})
+        if i % 3 == 0:
+            n = 4 if rng() < 0.5 else 6
+            yield (
+                "chat",
+                {
+                    "messages": assemble_frame_messages(corpus.frames_root, n_frames=n, rng=rng),
+                    "cache_salt": run_id,
+                    "session_id": f"{run_id}-p2",
+                },
+            )
+
+
+async def p3_single_image(
+    corpus: SoakCorpus, run_id: str, rng=random.random, *, structured_ratio: float = 0.5
+) -> AsyncIterator[Item]:
+    """P3 单图短指令高频主场景: 无状态,50% 走 structured 真实 schema。"""
+    while True:
+        image = corpus.images[int(rng() * len(corpus.images)) % len(corpus.images)]
+        kwargs: dict[str, Any] = {
+            "messages": chs_image_messages(image),
+            "cache_salt": run_id,
+            "session_id": f"{run_id}-p3",
+        }
+        if rng() < structured_ratio:
+            kwargs["structured"] = ChsExtraction
+        yield ("chat", kwargs)
+
+
+async def p4_cache_bidirectional(
+    corpus: SoakCorpus, run_id: str, rng=random.random
+) -> AsyncIterator[Item]:
+    """P4 缓存双向: 固定小图池重复 messages(命中侧,不掺 salt)+ salt 对照(强制 miss)。"""
+    pool = corpus.images[:5]
+    while True:
+        image = pool[int(rng() * len(pool)) % len(pool)]
+        messages = chs_image_messages(image)
+        if rng() < 0.5:
+            yield ("chat", {"messages": messages, "session_id": f"{run_id}-p4"})  # 可命中
+        else:
+            yield (
+                "chat",
+                {
+                    "messages": messages,
+                    "cache_salt": f"{run_id}-{rng()}",  # 强制 miss 对照
+                    "session_id": f"{run_id}-p4",
+                },
+            )
+
+
+async def p5_fault_mixed(corpus: SoakCorpus, run_id: str, rng=random.random) -> AsyncIterator[Item]:
+    """P5 故障源混编: 请求形态同 P3;故障性来自 SOAK scope 源池配置(findings §3),
+    生成器本身不造故障——真实故障由坏 key/黑洞/紧闸源在协议层自然发生。"""
+    async for item in p3_single_image(corpus, run_id, rng, structured_ratio=0.3):
+        yield item
+
+
+async def p6_mixed_soak(
+    corpus: SoakCorpus,
+    run_id: str,
+    rng=random.random,
+    *,
+    weights: dict[str, float] | None = None,
+) -> AsyncIterator[Item]:
+    """P6 混合浸泡: 按签字比例(设计 §8.1)加权混合;P4 已并入 P3 权重的 3/10。"""
+    weights = weights or {"P1": 0.10, "P2": 0.20, "P3": 0.35, "P4": 0.15, "P5": 0.20}
+    gens = {
+        "P1": p1_trace_chains(corpus, run_id, rng),
+        "P2": p2_multimodal_replay(corpus, run_id, rng),
+        "P3": p3_single_image(corpus, run_id, rng),
+        "P4": p4_cache_bidirectional(corpus, run_id, rng),
+        "P5": p5_fault_mixed(corpus, run_id, rng),
+    }
+    while True:
+        key = weighted_mix(weights, rng)
+        try:
+            yield await gens[key].__anext__()
+        except StopAsyncIteration:
+            # 有限语料场景(P1/P2)耗尽后重启一轮(浸泡语义: 语料循环使用)
+            gens[key] = {
+                "P1": p1_trace_chains,
+                "P2": p2_multimodal_replay,
+            }[key](corpus, f"{run_id}-r{rng()}", rng)
+
+
+SCENARIOS = {
+    "P1": p1_trace_chains,
+    "P2": p2_multimodal_replay,
+    "P3": p3_single_image,
+    "P4": p4_cache_bidirectional,
+    "P5": p5_fault_mixed,
+    "P6": p6_mixed_soak,
+}