engine/live_brain.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Protocol

import numpy as np

from engine.brain import Brain, BrainSignals


Dialogue = list[dict[str, str]]


class BrainClient(Protocol):
    def step(
        self,
        agent_id: str,
        system_prompt: str,
        dialogue_so_far: Dialogue,
        new_user_text: str,
        silence_flag: bool,
    ) -> dict[str, object]:
        ...

    def step_many(
        self,
        agent_payloads: list[dict[str, str]],
        dialogue_so_far: Dialogue,
        new_user_text: str,
        silence_flag: bool,
    ) -> dict[str, object]:
        ...

    def generate(self, agent_id: str, system_prompt: str, dialogue: Dialogue) -> dict[str, object]:
        ...


class ModalBrainClient:
    """Thin lazy wrapper around `modal_app.brain_modal` remote functions."""

    def __init__(self) -> None:
        from modal_app import brain_modal

        self._brain_modal = brain_modal

    def step(
        self,
        agent_id: str,
        system_prompt: str,
        dialogue_so_far: Dialogue,
        new_user_text: str,
        silence_flag: bool,
    ) -> dict[str, object]:
        return self._brain_modal.step.remote(agent_id, system_prompt, dialogue_so_far, new_user_text, silence_flag)

    def step_many(
        self,
        agent_payloads: list[dict[str, str]],
        dialogue_so_far: Dialogue,
        new_user_text: str,
        silence_flag: bool,
    ) -> dict[str, object]:
        return self._brain_modal.step_many.remote(agent_payloads, dialogue_so_far, new_user_text, silence_flag)

    def generate(self, agent_id: str, system_prompt: str, dialogue: Dialogue) -> dict[str, object]:
        return self._brain_modal.generate.remote(agent_id, system_prompt, dialogue)


@dataclass(frozen=True)
class Persona:
    agent_id: str
    display_name: str
    system_prompt: str


class LiveBrain(Brain):
    """One-agent Brain implementation backed by Modal."""

    def __init__(self, persona: Persona, client: BrainClient | None = None) -> None:
        self.persona = persona
        self.client = client or ModalBrainClient()
        self.last_raw: dict[str, object] | None = None
        self._latest: BrainSignals | None = None

    def step(self, dialogue_so_far: Dialogue, new_user_text: str, silence_flag: bool) -> BrainSignals:
        raw = self.client.step(
            self.persona.agent_id,
            self.persona.system_prompt,
            dialogue_so_far,
            new_user_text,
            silence_flag,
        )
        self.last_raw = raw
        self._latest = signals_from_raw(raw)
        return self._latest

    def next_signals(self) -> BrainSignals:
        if self._latest is None:
            raise RuntimeError("LiveBrain has no queued signals; call step() first")
        return self._latest

    def generate(self, dialogue: Dialogue) -> dict[str, object]:
        return self.client.generate(self.persona.agent_id, self.persona.system_prompt, dialogue)


class LiveBrainPanel:
    """Batched Modal brain calls for a multi-agent panel."""

    def __init__(self, personas: list[Persona], client: BrainClient | None = None) -> None:
        if not personas:
            raise ValueError("personas must not be empty")
        self.personas = personas
        self.client = client or ModalBrainClient()
        self.last_raw: dict[str, object] | None = None
        self.last_results: dict[str, dict[str, object]] = {}

    @property
    def agent_ids(self) -> list[str]:
        return [persona.agent_id for persona in self.personas]

    def step_all(self, dialogue_so_far: Dialogue, new_user_text: str, silence_flag: bool) -> dict[str, BrainSignals]:
        payloads = [
            {"agent_id": persona.agent_id, "system_prompt": persona.system_prompt}
            for persona in self.personas
        ]
        raw = self.client.step_many(payloads, dialogue_so_far, new_user_text, silence_flag)
        self.last_raw = raw
        results = raw.get("results", {}) if isinstance(raw, dict) else {}
        self.last_results = {agent_id: dict(result) for agent_id, result in results.items()}
        return {agent_id: signals_from_raw(result) for agent_id, result in self.last_results.items()}

    def generate(self, agent_id: str, dialogue: Dialogue) -> dict[str, object]:
        persona = self.persona(agent_id)
        return self.client.generate(agent_id, persona.system_prompt, dialogue)

    def persona(self, agent_id: str) -> Persona:
        for persona in self.personas:
            if persona.agent_id == agent_id:
                return persona
        raise KeyError(f"unknown persona {agent_id!r}")


def signals_from_raw(raw: dict[str, object]) -> BrainSignals:
    if not raw.get("ok", False):
        raise RuntimeError(str(raw.get("failure", "brain call failed")))
    return BrainSignals(
        surprise=float(raw["surprise"]),
        hidden=np.asarray(raw["hidden"], dtype=np.float32),
        readiness=float(raw["readiness"]),
        p_end=float(raw["p_end"]),
    )