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()