| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Iterable, Protocol, Sequence |
|
|
| import numpy as np |
|
|
|
|
| @dataclass(frozen=True) |
| class BrainSignals: |
| """Signals emitted by one agent-context brain for one transcript update. |
| |
| `readiness` is the controller-facing score for a speculative short reply: |
| |
| readiness = 1 / (1 + mean_token_entropy) |
| |
| A low-entropy draft means this agent has a confident next move. `p_end` is a |
| turn-completion heuristic: in the live loop it should combine trailing |
| silence, sentence-final punctuation, and/or high EOS probability. |
| """ |
|
|
| surprise: float |
| hidden: np.ndarray |
| readiness: float |
| p_end: float |
|
|
| def __post_init__(self) -> None: |
| hidden = np.asarray(self.hidden, dtype=np.float32) |
| if hidden.ndim != 1: |
| raise ValueError("hidden must be a 1-D float32 vector") |
| if not 0.0 <= float(self.readiness) <= 1.0: |
| raise ValueError("readiness must be in [0, 1]") |
| if not 0.0 <= float(self.p_end) <= 1.0: |
| raise ValueError("p_end must be in [0, 1]") |
| object.__setattr__(self, "hidden", hidden) |
| object.__setattr__(self, "surprise", float(self.surprise)) |
| object.__setattr__(self, "readiness", float(self.readiness)) |
| object.__setattr__(self, "p_end", float(self.p_end)) |
|
|
|
|
| class Brain(Protocol): |
| """Interface the live instrumented LLM must satisfy for one agent context.""" |
|
|
| def next_signals(self) -> BrainSignals: |
| """Return signals for the newest incremental transcript update.""" |
|
|
|
|
| class ReplayBrain: |
| """Deterministic `Brain` backed by precomputed signal samples.""" |
|
|
| def __init__(self, samples: Sequence[BrainSignals | dict[str, object]], *, name: str = "replay") -> None: |
| self.name = name |
| self._samples = [coerce_signals(sample) for sample in samples] |
| self._index = 0 |
|
|
| def __len__(self) -> int: |
| return len(self._samples) |
|
|
| def __iter__(self) -> Iterable[BrainSignals]: |
| return iter(self._samples) |
|
|
| def reset(self) -> None: |
| self._index = 0 |
|
|
| def next_signals(self) -> BrainSignals: |
| if self._index >= len(self._samples): |
| raise StopIteration(f"ReplayBrain {self.name!r} is exhausted") |
| sample = self._samples[self._index] |
| self._index += 1 |
| return sample |
|
|
| @classmethod |
| def from_npz( |
| cls, |
| path: str | Path, |
| *, |
| readiness: Sequence[float] | None = None, |
| p_end: Sequence[float] | None = None, |
| name: str | None = None, |
| ) -> "ReplayBrain": |
| dump = np.load(path, allow_pickle=True) |
| surprises = np.asarray(dump["nll_series"], dtype=np.float32) |
| hidden = np.asarray(dump["hidden_states"], dtype=np.float32) |
| if hidden.ndim != 2: |
| raise ValueError("hidden_states in npz must be a 2-D matrix") |
|
|
| n_steps = int(surprises.shape[0]) |
| readiness_values = ( |
| np.asarray(readiness, dtype=np.float32) |
| if readiness is not None |
| else np.linspace(0.35, 0.75, n_steps, dtype=np.float32) |
| ) |
| p_end_values = np.asarray(p_end, dtype=np.float32) if p_end is not None else np.zeros(n_steps, dtype=np.float32) |
| if n_steps: |
| p_end_values[-1] = max(float(p_end_values[-1]), 0.95) |
|
|
| samples = [ |
| BrainSignals( |
| surprise=float(surprises[index]), |
| hidden=hidden[index], |
| readiness=float(readiness_values[index]), |
| p_end=float(p_end_values[index]), |
| ) |
| for index in range(n_steps) |
| ] |
| return cls(samples, name=name or Path(path).stem) |
|
|
|
|
| def coerce_signals(sample: BrainSignals | dict[str, object]) -> BrainSignals: |
| if isinstance(sample, BrainSignals): |
| return sample |
| return BrainSignals( |
| surprise=float(sample["surprise"]), |
| hidden=np.asarray(sample["hidden"], dtype=np.float32), |
| readiness=float(sample["readiness"]), |
| p_end=float(sample["p_end"]), |
| ) |
|
|
|
|
