engine/conversation.py

from __future__ import annotations

import argparse
import json
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Sequence

from engine.controller import Action, ControllerConfig, ControllerTick, WhenToSpeakController
from engine.live_brain import BrainClient, Dialogue, LiveBrainPanel, Persona


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_LOG_PATH = ROOT / "eval" / "conversation_log.json"


@dataclass(frozen=True)
class TranscriptChunk:
    text: str
    silence_flag: bool = False


@dataclass(frozen=True)
class ConversationResult:
    events: list[dict[str, Any]]
    dialogue: Dialogue
    personas: list[Persona]
    total_latency_ms: float
    model_name: str
    device_name: str
    generated_examples: list[dict[str, Any]]


def default_personas() -> list[Persona]:
    return [
        Persona(
            agent_id="numbers_vc",
            display_name="Numbers VC",
            system_prompt=(
                "You are a numbers-obsessed venture investor. Be blunt, specific, and quantitative. "
                "Ask for denominators, cohorts, margins, contract evidence, and arithmetic that actually closes."
            ),
        ),
        Persona(
            agent_id="vision_optimist",
            display_name="Vision Optimist",
            system_prompt=(
                "You are a big-vision optimist. You look for the huge version of the company, but your "
                "questions are crisp and founder-facing when the story needs a missing bridge."
            ),
        ),
        Persona(
            agent_id="ruthless_skeptic",
            display_name="Ruthless Skeptic",
            system_prompt=(
                "You are a ruthless startup skeptic. Interrupt bad claims in plain English. No pleasantries, "
                "no throat-clearing, no softening. Be sharp without being long."
            ),
        ),
    ]


def sample_pitch_stream() -> list[TranscriptChunk]:
    return [
        TranscriptChunk("so basically our startup helps small retailers manage inventory"),
        TranscriptChunk("we connect to their point of sale and purchase orders"),
        TranscriptChunk("we already have ten thousand stores and zero churn after launching last week"),
        TranscriptChunk("then we predict stockouts and write reorder suggestions automatically"),
        TranscriptChunk("we are converting pilots into paid contracts this month"),
        TranscriptChunk("so we think this becomes the operating system for local retail"),
        TranscriptChunk("that's the pitch", silence_flag=True),
    ]


def demo_controller_config() -> ControllerConfig:
    return ControllerConfig(
        tau=0.85,
        min_readiness=0.08,
        w_surprise=0.85,
        w_barge=0.85,
        w_readiness=0.75,
        w_end=1.05,
        backchannel_tau_fraction=0.72,
        barge_tau_fraction=0.50,
        turn_end_tau_discount=0.45,
    )


class Conversation:
    def __init__(
        self,
        personas: list[Persona],
        brain_panel: LiveBrainPanel,
        controller: WhenToSpeakController | None = None,
    ) -> None:
        self.personas = personas
        self.brain_panel = brain_panel
        self.controller = controller or WhenToSpeakController(
            brain_panel.agent_ids,
            config=demo_controller_config(),
        )

    def run(self, stream: Sequence[TranscriptChunk]) -> ConversationResult:
        started = time.perf_counter()
        dialogue: Dialogue = []
        current_user_text = ""
        events: list[dict[str, Any]] = []
        generated_examples: list[dict[str, Any]] = []

        for step_index, chunk in enumerate(stream, start=1):
            dialogue_before = _dialogue_with_current_user(dialogue, current_user_text)
            signals = self.brain_panel.step_all(dialogue_before, chunk.text, chunk.silence_flag)
            current_user_text = _join_text(current_user_text, chunk.text)
            tick = self.controller.tick(signals, floor_holder="human")
            event = self._event(step_index, chunk, tick)

            winner = tick.winner
            if winner is not None:
                if current_user_text:
                    dialogue.append({"role": "user", "speaker": "founder", "text": current_user_text})
                    current_user_text = ""

                generated = self.brain_panel.generate(winner, dialogue)
                reply_text = str(generated.get("reply_text", ""))
                if reply_text:
                    dialogue.append({"role": "assistant", "speaker": winner, "text": reply_text})
                event["generated"] = {
                    "agent_id": winner,
                    "reply_text": reply_text,
                    "reply_source": generated.get("reply_source"),
                    "raw_reply_text": generated.get("raw_reply_text"),
                    "latency_ms": generated.get("latency_ms"),
                    "model_name": generated.get("model_name"),
                }
                generated_examples.append(event["generated"])

            events.append(event)

        if current_user_text:
            dialogue.append({"role": "user", "speaker": "founder", "text": current_user_text})

        raw = self.brain_panel.last_raw or {}
        return ConversationResult(
            events=events,
            dialogue=dialogue,
            personas=self.personas,
            total_latency_ms=(time.perf_counter() - started) * 1000.0,
            model_name=str(raw.get("model_name", "")),
            device_name=str(raw.get("device_name", "")),
            generated_examples=generated_examples,
        )

    def _event(self, step_index: int, chunk: TranscriptChunk, tick: ControllerTick) -> dict[str, Any]:
        raw = self.brain_panel.last_raw or {}
        decisions = {}
        for agent_id, decision in tick.decisions.items():
            brain_raw = self.brain_panel.last_results.get(agent_id, {})
            decisions[agent_id] = {
                "action": decision.action.value,
                "urge": decision.urge,
                "z_surprise": decision.z_surprise,
                "change_score": decision.change_score,
                "readiness": decision.readiness,
                "p_end": decision.p_end,
                "hidden_delta": decision.hidden_delta,
                "map_run_length": decision.map_run_length,
                "brain_latency_ms": brain_raw.get("latency_ms"),
                "surprise": brain_raw.get("surprise"),
            }
        return {
            "step": step_index,
            "new_user_text": chunk.text,
            "silence_flag": chunk.silence_flag,
            "winner": tick.winner,
            "floor_holder": tick.floor_holder,
            "batch_latency_ms": raw.get("batch_latency_ms"),
            "model_name": raw.get("model_name"),
            "device_name": raw.get("device_name"),
            "decisions": decisions,
        }


def save_conversation_log(result: ConversationResult, path: str | Path = DEFAULT_LOG_PATH) -> Path:
    output = Path(path)
    output.parent.mkdir(parents=True, exist_ok=True)
    data = {
        "model_name": result.model_name,
        "device_name": result.device_name,
        "total_latency_ms": result.total_latency_ms,
        "personas": [asdict(persona) for persona in result.personas],
        "events": result.events,
        "dialogue": result.dialogue,
        "generated_examples": result.generated_examples,
    }
    output.write_text(json.dumps(data, indent=2), encoding="utf-8")
    return output


def readable_log(result: ConversationResult) -> str:
    persona_names = {persona.agent_id: persona.display_name for persona in result.personas}
    lines = [
        f"Model: {result.model_name or 'unknown'} on {result.device_name or 'unknown'}",
        f"Total wall latency: {result.total_latency_ms:.1f} ms",
    ]
    for event in result.events:
        lines.append(f"[{event['step']}] USER + {event['new_user_text']!r} silence={event['silence_flag']}")
        for agent_id, decision in event["decisions"].items():
            action = decision["action"]
            if action == Action.SILENT.value:
                continue
            lines.append(
                "    "
                f"{persona_names.get(agent_id, agent_id)} -> {action} "
                f"urge={decision['urge']:.2f} readiness={decision['readiness']:.2f} "
                f"p_end={decision['p_end']:.2f}"
            )
        if "generated" in event:
            generated = event["generated"]
            lines.append(f"    {persona_names.get(generated['agent_id'], generated['agent_id'])}: {generated['reply_text']}")
    return "\n".join(lines)


def _dialogue_with_current_user(dialogue: Dialogue, current_user_text: str) -> Dialogue:
    snapshot = [dict(turn) for turn in dialogue]
    snapshot.append({"role": "user", "speaker": "founder", "text": current_user_text})
    return snapshot


def _join_text(left: str, right: str) -> str:
    left = left.strip()
    right = right.strip()
    if not left:
        return right
    if not right:
        return left
    return f"{left} {right}"


def run_demo(log_path: str | Path = DEFAULT_LOG_PATH, client: BrainClient | None = None) -> ConversationResult:
    personas = default_personas()
    panel = LiveBrainPanel(personas, client=client)
    conversation = Conversation(personas, panel)
    result = conversation.run(sample_pitch_stream())
    save_conversation_log(result, log_path)
    print(readable_log(result))
    print(f"Wrote {log_path}")
    return result


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="Run the text-streamed WhenToSpeak conversation demo.")
    parser.add_argument("--log-path", default=str(DEFAULT_LOG_PATH))
    args = parser.parse_args(argv)
    run_demo(args.log_path)


if __name__ == "__main__":
    main()