Read live events and manage live sessions

"""Record: a barge-in, the saved session, a second connection on the same session, a reused queue."""
import asyncio
import json
from pathlib import Path

from google.adk.agents import LiveRequestQueue

from captions import RUN_CONFIG, Captions, build_agent, make_runner, open_session, remember_questions, stored
from live_harness import RecordingGemini, environment, MODEL, read_pcm, say_audio, say_text, spoken, summarize

HERE = Path(__file__).parent
LONG = "Read me every limit on the page, one at a time, with a sentence about each."


async def barge_in(runner, session_id, queue):
    """Ask a long question, then talk over the answer once its audio starts."""
    rows, interrupted = [], False
    await say_text(LONG)(queue)
    async for event in runner.run_live(user_id="dev", session_id=session_id,
                                       live_request_queue=queue, run_config=RUN_CONFIG):
        row = summarize(event)
        rows.append(row)
        if row.get("audio_bytes") and not interrupted:
            interrupted = True
            await say_audio(read_pcm(HERE / "interrupt.wav"))(queue)
        if sum(1 for r in rows if r.get("turn_complete")) == 2:
            break
    queue.close()
    return rows


async def one_turn(runner, session_id, queue, turn, timeout=60):
    rows = []

    async def loop():
        await turn(queue)
        async for event in runner.run_live(user_id="dev", session_id=session_id,
                                           live_request_queue=queue, run_config=RUN_CONFIG):
            rows.append(summarize(event))
            if event.turn_complete:
                break
    try:
        await asyncio.wait_for(loop(), timeout)
        outcome = "turn complete"
    except asyncio.TimeoutError:
        outcome = f"no event for {timeout} s"
    queue.close()
    return rows, outcome


async def main():
    model = RecordingGemini(model=MODEL, connections=[])
    runner = make_runner(build_agent(model))
    session = await open_session(runner, "review-call")
    first = await asyncio.wait_for(barge_in(runner, session.id, LiveRequestQueue()), 120)
    captions = Captions()
    for row in first:
        captions.apply(row)
    saved_after_first = stored(await open_session(runner, "review-call"))
    second, _ = await one_turn(runner, session.id, LiveRequestQueue(), say_audio(read_pcm(HERE / "followup.wav")))
    asked = await remember_questions(runner, await open_session(runner, "review-call"))
    third, _ = await one_turn(runner, session.id, LiveRequestQueue(), say_text("What did I ask you a moment ago?"))
    closed = LiveRequestQueue()
    closed.close()
    reused, outcome = await one_turn(runner, session.id, closed, say_text("Are you still there?"), timeout=20)
    print(f"[first connection] interrupted={any(r.get('interrupted') for r in first)} "
          f"resumption updates={sum(1 for r in first if 'resumption_update' in r)}")
    for who, line in captions.lines:
        print(f"  {who}: {line}")
    print(f"[saved] {len(saved_after_first)} events, audio parts saved: "
          f"{sum('inline_data' in r['kinds'] for r in saved_after_first)}")
    print(f"[second connection] history sent: {model.connections[1]["traffic"][0]["turns"]}")
    print(f"  user: {spoken(second, 'input_transcription')}\n  agent: {spoken(second)}")
    print(f"[third connection, state asked={asked}]\n  agent: {spoken(third)}")
    print(f"[reused closed queue] {len(reused)} events, {outcome}")
    record = {"environment": environment(), "first": first, "captions": captions.lines,
              "saved_after_first": saved_after_first, "second": second, "third": third, "asked": asked,
              "reused_queue": {"events": reused, "outcome": outcome}, "connections": model.connections}
    (HERE / "events-run.json").write_text(json.dumps(record, indent=1) + "\n")


if __name__ == "__main__":
    asyncio.run(main())

"""Replay the recorded live session through ADK; inject a dropped connection; no network."""
import asyncio
import copy
import json
import logging
import unittest
from pathlib import Path

from google.adk.agents import LiveRequestQueue
from google.adk.agents.run_config import RunConfig
from websockets.exceptions import ConnectionClosedError

import live_harness
from captions import RUN_CONFIG, Captions, build_agent, make_runner, open_session
from live_harness import ReplayLive, read_pcm, say_audio, say_text, spoken, summarize
from record import barge_in, one_turn

HERE = Path(__file__).parent
RECORD = json.loads((HERE / "events-run.json").read_text())
live_harness.PACE = 0
logging.disable(logging.CRITICAL)  # ADK logs the injected failures; the assertions check them


def replay(*indexes, drop_after_handle=False, strip_handles=False):
    connections = [copy.deepcopy(RECORD["connections"][i]) for i in indexes]
    traffic = connections[0]["traffic"]
    if strip_handles:
        traffic[:] = [e for e in traffic if "live_session_resumption_update" not in e.get("received", {})]
    if drop_after_handle or strip_handles:
        mark = "live_session_resumption_update" if drop_after_handle else "output_transcription"
        at = next(i for i, e in enumerate(traffic) if e.get("received", {}).get(mark, {}).get(
            "new_handle" if drop_after_handle else "text"))
        traffic.insert(at + 1, {"raise": ConnectionClosedError(None, None)})
    return ReplayLive(model="replay", connections=connections, opened=[])


class LiveEvents(unittest.TestCase):
    def test_captions_keep_each_finished_line_once(self):
        captions = Captions()
        for row in RECORD["first"]:
            captions.apply(row)
        self.assertEqual([list(line) for line in captions.lines], RECORD["captions"])
        self.assertEqual(captions.state, "listening")

    def test_joining_every_fragment_doubles_the_transcript(self):
        naive = "".join(r["output_transcription"]["text"] for r in RECORD["first"] if "output_transcription" in r)
        first_line = RECORD["captions"][0][1]
        self.assertEqual(naive.count(first_line), 2)
        self.assertEqual(spoken(RECORD["first"]).count(first_line), 1)

    def test_barge_in_sets_interrupted_before_the_turn_completes(self):
        flags = [("interrupted" if r.get("interrupted") else "turn_complete")
                 for r in RECORD["first"] if r.get("interrupted") or r.get("turn_complete")]
        self.assertEqual(flags[:2], ["interrupted", "turn_complete"])

    def test_audio_is_streamed_but_not_saved(self):
        self.assertGreater(sum(r.get("audio_bytes", 0) for r in RECORD["first"]), 0)
        self.assertFalse(any("inline_data" in r["kinds"] for r in RECORD["saved_after_first"]))

    def test_a_new_connection_replays_the_saved_history(self):
        model = replay(0, 1)
        runner = make_runner(build_agent(model))

        async def run():
            session = await open_session(runner, "review-call")
            await barge_in(runner, session.id, LiveRequestQueue())
            await one_turn(runner, session.id, LiveRequestQueue(), say_audio(read_pcm(HERE / "followup.wav")))
        asyncio.run(run())
        self.assertEqual(model.opened[1][1].received[0], ("history", 4))

    def test_replayed_history_is_not_recall(self):
        self.assertNotIn("sample rate", spoken(RECORD["second"]))
        self.assertIn("sample rate", spoken(RECORD["third"]))

    def test_a_closed_queue_ends_run_live_with_no_events(self):
        runner = make_runner(build_agent(replay(0)))

        async def run():
            session = await open_session(runner, "review-call")
            queue = LiveRequestQueue()
            queue.close()
            return [summarize(e) async for e in runner.run_live(
                user_id="dev", session_id=session.id, live_request_queue=queue, run_config=RUN_CONFIG)]
        self.assertEqual(asyncio.run(run()), [])

    def test_a_drop_with_a_cached_handle_reconnects_without_history(self):
        model = replay(0, 1, drop_after_handle=True)
        runner = make_runner(build_agent(model))

        async def run():
            session = await open_session(runner, "review-call")
            await asyncio.wait_for(barge_in(runner, session.id, LiveRequestQueue()), 3)
        try:
            asyncio.run(run())
        except asyncio.TimeoutError:
            pass  # the second recording expects other input; reconnecting is what is under test
        request, connection = model.opened[1]
        self.assertTrue(request.live_connect_config.session_resumption.handle)
        self.assertNotIn("history", [kind for kind, _ in connection.received])

    def test_a_drop_without_a_handle_reaches_the_caller(self):
        runner = make_runner(build_agent(replay(0, strip_handles=True)))

        async def run():
            session = await open_session(runner, "review-call")
            await barge_in(runner, session.id, LiveRequestQueue())
        with self.assertRaises(ConnectionClosedError):
            asyncio.run(run())


if __name__ == "__main__":
    unittest.main()