scripts/walk_smoke.py

from __future__ import annotations

import os
import sys
from time import perf_counter
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
ROOT_ENV = ROOT / ".env"
LOCAL_ENV = ROOT / "data" / "local.env"

if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))


def main() -> int:
    _load_local_env(ROOT_ENV, override=True)
    _load_local_env(LOCAL_ENV, override=True)

    from time_machine.application.container import create_container
    from time_machine.domain.events import (
        AudioEvent,
        ConversationTextEvent,
        DestinationEvent,
        PersonaEvent,
        SouvenirEvent,
    )

    if not (os.getenv("TIME_MACHINE_LLM_API_KEY") or os.getenv("TOGETHER_API_KEY")):
        print("Missing TIME_MACHINE_LLM_API_KEY or TOGETHER_API_KEY. Put it in data/local.env or the shell env.")
        return 2

    container = create_container(adapter_profile="dev")
    launch_started = perf_counter()
    launch_events = list(
        container.encounter_service.start_encounter(
            mode=os.getenv("TIME_MACHINE_WALK_MODE", "surprise"),
            coordinate_prompt=os.getenv("TIME_MACHINE_WALK_PROMPT", "a rain-soaked workshop before electricity"),
            session_id="walk-smoke",
        )
    )
    launch_elapsed = perf_counter() - launch_started

    destination = _single_event(launch_events, DestinationEvent).destination
    persona = _single_event(launch_events, PersonaEvent).persona
    encounter_id = launch_events[0].metadata["encounter_id"]

    turn_started = perf_counter()
    turn_events = list(
        container.encounter_service.handle_text_turn(
            encounter_id=encounter_id,
            text=os.getenv("TIME_MACHINE_WALK_TEXT", "Can you hear the machine through the wall?"),
        )
    )
    turn_elapsed = perf_counter() - turn_started
    response = _single_event(turn_events, ConversationTextEvent)
    audio = _single_event(turn_events, AudioEvent).audio
    if not audio.path or not Path(audio.path).exists():
        raise RuntimeError(f"Audio path was not created: {audio.path!r}")

    souvenir_events = list(container.encounter_service.generate_souvenir(encounter_id))
    souvenir = _single_event(souvenir_events, SouvenirEvent).souvenir

    print("Walk smoke passed.")
    print(f"Profile: {container.adapter_profile}")
    print(f"Destination: {destination.place}, {destination.year}")
    print(f"Contact: {persona.name}, {persona.role}")
    print(f"Reply: {response.text[:240]}")
    print(f"Audio: {audio.path} ({audio.duration_seconds}s)")
    print(f"Launch wall time: {launch_elapsed:.2f}s")
    print(f"Turn wall time: {turn_elapsed:.2f}s")
    print("Turn stage timings:")
    for timing in _event_timings(turn_events):
        print(f"  - {timing['stage']}: {timing['elapsed_ms'] / 1000:.2f}s")
    if audio.description and "fallback" in audio.description.lower():
        print(f"Audio note: {audio.description}")
    print(f"Souvenir: {souvenir.stamp_name}")
    return 0


def _load_local_env(path: Path, *, override: bool = False) -> None:
    if not path.exists():
        return
    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        if key and (override or key not in os.environ):
            os.environ[key] = value


def _single_event(events: list[object], event_type: type) -> object:
    matches = [event for event in events if isinstance(event, event_type)]
    if not matches:
        raise RuntimeError(f"Expected event {event_type.__name__}, got {[type(e).__name__ for e in events]}")
    return matches[-1]


def _event_timings(events: list[object]) -> list[dict[str, object]]:
    timings: list[dict[str, object]] = []
    seen: set[tuple[str, float]] = set()
    for event in events:
        metadata = getattr(event, "metadata", {})
        for timing in metadata.get("timings", []):
            stage = timing.get("stage")
            elapsed_ms = timing.get("elapsed_ms")
            key = (stage, elapsed_ms)
            if key in seen:
                continue
            seen.add(key)
            timings.append(timing)
    return timings


if __name__ == "__main__":
    raise SystemExit(main())