from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path from typing import Sequence import numpy as np from engine.brain import BrainSignals, ReplayBrain ROOT = Path(__file__).resolve().parents[1] DEFAULT_MODAL_DUMP = ROOT / "eval" / "probe_dump_modal.npz" DEFAULT_AGENTS = ("investor_a", "investor_b") @dataclass(frozen=True) class TraceStep: signals_by_agent: dict[str, BrainSignals] floor_holder: str = "human" note: str = "" @dataclass(frozen=True) class SignalTrace: name: str steps: list[TraceStep] expected: dict[str, object] = field(default_factory=dict) @property def agent_ids(self) -> list[str]: return list(self.steps[0].signals_by_agent.keys()) if self.steps else [] def synthetic_traces() -> list[SignalTrace]: return [ clean_monologue_trace(), surprising_claim_trace(), topic_shift_trace(), rambling_pause_trace(), ] def load_modal_trace(path: str | Path = DEFAULT_MODAL_DUMP, agent_ids: Sequence[str] = DEFAULT_AGENTS) -> SignalTrace: path = Path(path) replay = ReplayBrain.from_npz(path, name="modal_probe_replay") samples = list(replay) steps: list[TraceStep] = [] for index, sample in enumerate(samples): signals: dict[str, BrainSignals] = {} for offset, agent_id in enumerate(agent_ids): readiness = max(0.0, sample.readiness - 0.08 * offset) signals[agent_id] = BrainSignals( surprise=sample.surprise, hidden=sample.hidden, readiness=readiness, p_end=sample.p_end, ) steps.append(TraceStep(signals_by_agent=signals, note=f"modal step {index + 1}")) return SignalTrace( name="modal_probe_replay", steps=steps, expected={"source": str(path), "n_steps": len(steps)}, ) def all_traces() -> list[SignalTrace]: traces = synthetic_traces() if DEFAULT_MODAL_DUMP.exists(): traces.append(load_modal_trace(DEFAULT_MODAL_DUMP)) return traces def clean_monologue_trace() -> SignalTrace: base = _unit([1, 0, 0, 0, 0, 0, 0, 0]) steps = _trace_from_series( name="clean_monologue_take_floor", hidden_series=[_nudge(base, index, 0.01) for index in range(6)], surprise=[2.0, 2.1, 2.0, 2.2, 2.1, 2.45], readiness_a=[0.25, 0.30, 0.35, 0.40, 0.45, 0.78], readiness_b=[0.20, 0.25, 0.30, 0.32, 0.35, 0.52], p_end=[0.02, 0.03, 0.04, 0.05, 0.10, 0.95], notes=["setup", "details", "still talking", "more context", "closing", "turn complete"], ) return SignalTrace(name="clean_monologue_take_floor", steps=steps, expected={"take_floor_step": 6}) def surprising_claim_trace() -> SignalTrace: base = _unit([1, 0, 0, 0, 0, 0, 0, 0]) steps = _trace_from_series( name="surprising_claim_interrupt", hidden_series=[_nudge(base, index, 0.01) for index in range(6)], surprise=[2.0, 2.1, 7.3, 3.0, 2.4, 2.2], readiness_a=[0.30, 0.42, 0.92, 0.82, 0.55, 0.45], readiness_b=[0.25, 0.32, 0.35, 0.36, 0.38, 0.40], p_end=[0.02, 0.04, 0.18, 0.20, 0.30, 0.55], notes=["setup", "build", "wild claim", "continues", "settles", "handoff"], ) return SignalTrace(name="surprising_claim_interrupt", steps=steps, expected={"interrupt_step": 3}) def topic_shift_trace() -> SignalTrace: topic_a = _unit([1, 0, 0, 0, 0, 0, 0, 0]) topic_b = _unit([0, 1, 0, 0, 0, 0, 0, 0]) steps = _trace_from_series( name="topic_shift_backchannel", hidden_series=[ _nudge(topic_a, 0, 0.01), _nudge(topic_a, 1, 0.01), _nudge(topic_a, 2, 0.01), _nudge(topic_a, 3, 0.01), _nudge(topic_b, 4, 0.01), _nudge(topic_b, 5, 0.01), _nudge(topic_b, 6, 0.01), ], surprise=[2.0, 2.1, 2.0, 2.1, 2.12, 2.08, 2.06], readiness_a=[0.25, 0.30, 0.35, 0.38, 0.50, 0.45, 0.40], readiness_b=[0.20, 0.25, 0.28, 0.30, 0.35, 0.34, 0.34], p_end=[0.02, 0.04, 0.06, 0.08, 0.22, 0.25, 0.28], notes=["topic a", "topic a", "topic a", "topic a", "topic b shift", "topic b", "topic b"], ) return SignalTrace(name="topic_shift_backchannel", steps=steps, expected={"backchannel_step": 5}) def rambling_pause_trace() -> SignalTrace: base = _unit([0, 0, 1, 0, 0, 0, 0, 0]) steps = _trace_from_series( name="rambling_pause_take_floor", hidden_series=[_nudge(base, index, 0.015) for index in range(7)], surprise=[2.1, 2.0, 2.2, 2.1, 2.0, 2.1, 2.45], readiness_a=[0.20, 0.25, 0.30, 0.38, 0.42, 0.42, 0.80], readiness_b=[0.18, 0.22, 0.25, 0.30, 0.32, 0.35, 0.50], p_end=[0.02, 0.05, 0.10, 0.38, 0.42, 0.44, 0.96], notes=["start", "ramble", "ramble", "awkward pause", "holds", "still unsure", "complete"], ) return SignalTrace(name="rambling_pause_take_floor", steps=steps, expected={"hold_step": 4, "take_floor_step": 7}) def _trace_from_series( *, name: str, hidden_series: Sequence[np.ndarray], surprise: Sequence[float], readiness_a: Sequence[float], readiness_b: Sequence[float], p_end: Sequence[float], notes: Sequence[str], ) -> list[TraceStep]: steps: list[TraceStep] = [] for index, hidden in enumerate(hidden_series): signals = { "investor_a": BrainSignals( surprise=surprise[index], hidden=hidden, readiness=readiness_a[index], p_end=p_end[index], ), "investor_b": BrainSignals( surprise=max(0.0, surprise[index] - 0.2), hidden=hidden, readiness=readiness_b[index], p_end=p_end[index], ), } steps.append(TraceStep(signals_by_agent=signals, note=notes[index])) return steps def _unit(values: Sequence[float]) -> np.ndarray: vector = np.asarray(values, dtype=np.float32) norm = np.linalg.norm(vector) if norm == 0: return vector return vector / norm def _nudge(base: np.ndarray, index: int, scale: float) -> np.ndarray: vector = base.astype(np.float32, copy=True) vector[(index % (len(vector) - 1)) + 1] += scale return _unit(vector)