Stream a voice conversation with run_live

"""Drive run_live, keep the event fields a reader needs, and record or replay the Live API traffic.

RecordingGemini wraps the real Live connection and logs what ADK sends and what the model
returns (audio is logged as a byte count). ReplayLive plays that log back through ADK's own
run_live loop, so the tests exercise ADK's tool execution, callbacks and session writes offline.
"""
import asyncio
import base64
import contextlib
import os
import wave
from importlib.metadata import version

from google.adk.agents import LiveRequestQueue
from google.adk.models import BaseLlm, Gemini, LlmResponse
from google.adk.models.base_llm_connection import BaseLlmConnection
from google.genai import types

# Current course Live model. Override with LIVE_MODEL.
MODEL = os.getenv("LIVE_MODEL", "gemini-3.8-live")
RATE_IN, CHUNK_MS = 16000, 100
PACE = 1.0  # seconds of wall time per second of audio; the tests set 0


def summarize(event) -> dict:
    """One event as the fields this course prints. Audio becomes a byte count."""
    row = {"author": event.author}
    for flag in ("partial", "turn_complete", "interrupted"):
        if getattr(event, flag):
            row[flag] = True
    for side in ("input_transcription", "output_transcription"):
        part = getattr(event, side)
        if part and (part.text or part.finished):
            row[side] = {"text": part.text or "", "finished": bool(part.finished)}
    for part in (event.content.parts or []) if event.content else []:
        if part.inline_data:
            row["audio_bytes"] = row.get("audio_bytes", 0) + len(part.inline_data.data or b"")
            row["mime_type"] = part.inline_data.mime_type
        elif part.text and not part.thought:
            row["text"] = row.get("text", "") + part.text
        if part.function_call:
            row.setdefault("calls", []).append({"name": part.function_call.name, "args": part.function_call.args})
        if part.function_response:
            row.setdefault("responses", []).append(
                {"name": part.function_response.name, "response": part.function_response.response})
    if event.usage_metadata and event.usage_metadata.total_token_count:
        row["total_tokens"] = event.usage_metadata.total_token_count
    if event.error_code:
        row["error"] = f"{event.error_code}: {event.error_message}"
    if event.live_session_resumption_update:
        row["resumption_update"] = bool(event.live_session_resumption_update.new_handle)
    if event.actions and event.actions.transfer_to_agent:
        row["transfer_to_agent"] = event.actions.transfer_to_agent
    return row


def read_pcm(path) -> bytes:
    """The PCM frames of a 16 kHz, 16-bit mono WAV file; any other format is refused."""
    with wave.open(str(path)) as clip:
        shape = (clip.getframerate(), clip.getsampwidth(), clip.getnchannels())
        if shape != (RATE_IN, 2, 1):
            raise ValueError(f"{path}: {shape} is not (16000 Hz, 16-bit, mono)")
        return clip.readframes(clip.getnframes())


def say_text(text):
    async def send(queue):
        queue.send_content(types.Content(role="user", parts=[types.Part(text=text)]))
    return send


def say_audio(pcm, mime_type="audio/pcm;rate=16000", silence_s=1.0):
    """Stream a clip in 100 ms chunks, then silence so voice activity detection ends the turn."""
    step = RATE_IN * 2 * CHUNK_MS // 1000
    data = pcm + bytes(int(RATE_IN * 2 * silence_s))

    async def stream(queue):
        for start in range(0, len(data), step):
            queue.send_realtime(types.Blob(mime_type=mime_type, data=data[start:start + step]))
            await asyncio.sleep(PACE * CHUNK_MS / 1000)

    async def send(queue):
        asyncio.get_running_loop().create_task(stream(queue))
    return send


async def talk(runner, session_id, run_config, turns, user_id="dev", timeout=90, queue=None):
    """Send each turn after the previous turn completes; return the summarized events."""
    queue = queue or LiveRequestQueue()
    rows, pending = [], list(turns)

    async def loop():
        if pending:
            await pending.pop(0)(queue)
        async for event in runner.run_live(user_id=user_id, session_id=session_id,
                                           live_request_queue=queue, run_config=run_config):
            rows.append(summarize(event))
            if event.turn_complete:
                if not pending:
                    break
                await pending.pop(0)(queue)
    try:
        await asyncio.wait_for(loop(), timeout)
    except asyncio.TimeoutError as error:
        error.rows = rows  # keep what arrived before the timeout
        raise
    finally:
        queue.close()
    return rows


def wait_for_turn():
    """A turn that sends nothing: wait for the model to speak on its own, for example a tool result."""
    async def send(queue):
        return None
    return send


def spoken(rows, side="output_transcription") -> str:
    """The transcript of one side: the finished fragments joined in order."""
    return "".join(r[side]["text"] for r in rows if side in r and not r.get("partial")).strip()


def environment() -> dict:
    return {"google-adk": version("google-adk"), "google-genai": version("google-genai"), "model": MODEL}


def _dump(response: LlmResponse) -> dict:
    data = response.model_dump(mode="json", exclude_none=True)
    for part in data.get("content", {}).get("parts", []):
        part.pop("thought_signature", None)
        if "inline_data" in part and isinstance(part["inline_data"].get("data"), str):
            part["inline_data"]["data"] = len(base64.urlsafe_b64decode(part["inline_data"]["data"] + "=="))
    data.pop("live_session_id", None)
    return data


def _load(data: dict) -> LlmResponse:
    data = {**data, "content": {**data["content"]}} if "content" in data else dict(data)
    parts = []
    for part in data.get("content", {}).get("parts", []):
        if isinstance(part.get("inline_data", {}).get("data"), int):
            part = {**part, "inline_data": {**part["inline_data"], "data": bytes(part["inline_data"]["data"])}}
        parts.append(part)
    if "content" in data:
        data["content"]["parts"] = parts
    return LlmResponse.model_validate(data)


class _Recorded(BaseLlmConnection):
    def __init__(self, inner, traffic):
        self.inner, self.traffic = inner, traffic

    def _sent(self, kind, **fields):
        last = self.traffic[-1] if self.traffic else {}
        if kind == "realtime" and last.get("sent") == "realtime" and last["mime_type"] == fields["mime_type"]:
            last["chunks"] += 1
            last["bytes"] += fields["bytes"]
        else:
            self.traffic.append({"sent": kind, **fields} | ({"chunks": 1} if kind == "realtime" else {}))

    async def send_history(self, history):
        self._sent("history", turns=[{"role": c.role, "text": "".join(p.text or "" for p in c.parts or [])[:120]}
                                     for c in history])
        await self.inner.send_history(history)

    async def send_content(self, content):
        await self._send_content(content)

    async def _send_content(self, content, *, partial=False):
        parts = content.parts or []
        self._sent("content", text="".join(p.text or "" for p in parts),
                   function_responses=[p.function_response.name for p in parts if p.function_response])
        await self.inner._send_content(content, partial=partial)

    async def send_realtime(self, blob):
        mime = getattr(blob, "mime_type", None) or type(blob).__name__
        self._sent("realtime", mime_type=mime, bytes=len(getattr(blob, "data", b"") or b""))
        await self.inner.send_realtime(blob)

    async def receive(self):
        async for response in self.inner.receive():
            self.traffic.append({"received": _dump(response)})
            yield response

    async def close(self):
        self.traffic.append({"sent": "close"})
        await self.inner.close()


class RecordingGemini(Gemini):
    """A Gemini model that logs each live connection: its config and its traffic."""

    connections: list = []

    @contextlib.asynccontextmanager
    async def connect(self, llm_request):
        declarations = [d for tool in llm_request.config.tools or [] for d in getattr(tool, "function_declarations", None) or []]
        log = {"model": llm_request.model, "traffic": [],
               "tools": [{"name": d.name, "behavior": d.behavior.value if d.behavior else None} for d in declarations]}
        self.connections.append(log)
        async with super().connect(llm_request) as connection:
            # Read the config after Gemini.connect has applied the agent-level voice and instruction.
            config = llm_request.live_connect_config.model_dump(mode="json", exclude_none=True)
            for private in ("system_instruction", "http_options", "tools"):
                config.pop(private, None)
            log["config"] = config
            yield _Recorded(connection, log["traffic"])


class _Replay(BaseLlmConnection):
    def __init__(self, traffic):
        self.traffic, self.index, self.sent, self.closed = traffic, 0, 0, False
        self.changed = asyncio.Event()
        self.received = []  # what ADK sent, for assertions

    def _note(self, item, weight=1):
        self.received.append(item)
        self.sent += weight
        self.changed.set()

    async def send_history(self, history):
        self._note(("history", len(history)))

    async def send_content(self, content):
        await self._send_content(content)

    async def _send_content(self, content, *, partial=False):
        self._note(("content", content))

    async def send_realtime(self, blob):
        self._note(("realtime", getattr(blob, "mime_type", type(blob).__name__)))

    async def receive(self):
        needed = sum(e.get("chunks", 1) for e in self.traffic[:self.index] if "sent" in e)
        while self.index < len(self.traffic):
            entry = self.traffic[self.index]
            self.index += 1
            if "sent" in entry:
                needed += entry.get("chunks", 1) if entry["sent"] != "close" else 0
                continue
            while self.sent < needed and not self.closed:
                self.changed.clear()
                await self.changed.wait()
            if self.closed:
                return
            if "raise" in entry:
                raise entry["raise"]
            yield _load(entry["received"])
        while not self.closed:  # an idle open connection
            self.changed.clear()
            await self.changed.wait()

    async def close(self):
        self.closed = True
        self.changed.set()


class ReplayLive(BaseLlm):
    """Plays recorded connections back in order; refuses a connection it has no recording for."""

    connections: list
    opened: list = []

    async def generate_content_async(self, llm_request, stream=False):
        raise NotImplementedError("ReplayLive only replays run_live connections")
        yield  # pragma: no cover

    @contextlib.asynccontextmanager
    async def connect(self, llm_request):
        if not self.connections:
            raise LookupError("no recorded live connection left")
        replay = _Replay(self.connections.pop(0)["traffic"])
        self.opened.append((llm_request, replay))
        yield replay

"""Record lesson runs on the Live API: two answered turns, ungrounded trials and three failures."""
import asyncio
import json
from pathlib import Path

from google.adk.agents.run_config import RunConfig

from live_harness import RecordingGemini, environment, MODEL, read_pcm, say_audio, say_text, spoken
from voice_docs_helper import ask, build_agent, mentions

HERE = Path(__file__).parent
QUESTION = "What does ADK stand for?"


def error_text(error) -> str:
    return f"{type(error).__name__}: {str(error).splitlines()[0][:240]}"


async def answered(name, turn, facts, grounded=True):
    model = RecordingGemini(model=MODEL, connections=[])
    rows, answer = await ask(build_agent(model, grounded=grounded), turn)
    run = {"name": name, "grounded": grounded, "answer": answer, "heard": spoken(rows, "input_transcription"),
           "expected": list(facts), "check_passed": mentions(answer, *facts), "events": rows,
           "connections": model.connections}
    print(f"[{name}] {answer!r} check_passed={run['check_passed']}")
    return run


async def failure(name, coro):
    try:
        await coro
        return {"name": name, "error": None}
    except Exception as error:  # recorded, not handled: these runs show what breaks
        print(f"[{name}] {error_text(error)}")
        return {"name": name, "error": error_text(error)}


async def missing_session():
    from google.adk.agents import LiveRequestQueue
    from google.adk.runners import Runner
    from google.adk.sessions import InMemorySessionService
    runner = Runner(app_name="voice_docs", agent=build_agent(), session_service=InMemorySessionService())
    async for _ in runner.run_live(user_id="dev", session_id="never-created",
                                   live_request_queue=LiveRequestQueue(), run_config=RunConfig()):
        break


async def main():
    runs = [await answered("text turn", say_text(QUESTION), ["Agent Development Kit"]),
            await answered("spoken turn", say_audio(read_pcm(HERE / "question.wav")), ["16"])]
    for trial in range(1, 4):
        run = await answered(f"ungrounded trial {trial}", say_text(QUESTION), ["Agent Development Kit"], grounded=False)
        run.pop("connections")
        runs.append(run)
    failures = [
        await failure("TEXT response modality", ask(build_agent(), say_text(QUESTION), RunConfig(response_modalities=["TEXT"]))),
        await failure("standard Gemini model", ask(build_agent("gemini-3.8-flash"), say_text(QUESTION))),
        await failure("session not created", missing_session()),
    ]
    record = {"environment": environment(), "question": QUESTION, "runs": runs, "failures": failures}
    (HERE / "stream-run.json").write_text(json.dumps(record, indent=1) + "\n")


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

"""Replay the recorded Live traffic through ADK's run_live loop; no network, no API key."""
import asyncio
import json
import logging
import tempfile
import unittest
import wave
from pathlib import Path

from google.adk.agents import LiveRequestQueue
from google.adk.agents.run_config import RunConfig
from google.adk.errors.session_not_found_error import SessionNotFoundError
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService

import live_harness
from live_harness import ReplayLive, read_pcm, say_audio, say_text
from voice_docs_helper import ask, build_agent, mentions

HERE = Path(__file__).parent
RECORD = json.loads((HERE / "stream-run.json").read_text())
RUNS = {run["name"]: run for run in RECORD["runs"]}
FAILURES = {f["name"]: f["error"] for f in RECORD["failures"]}
live_harness.PACE = 0
logging.disable(logging.CRITICAL)  # ADK logs the injected failures; the assertions check them  # stream recorded clips without waiting in real time


def replay(name):
    return ReplayLive(model="replay", connections=json.loads(json.dumps(RUNS[name]["connections"])), opened=[])


class VoiceDocsHelper(unittest.TestCase):
    def test_text_turn_replays_audio_and_a_transcript(self):
        model = replay("text turn")
        rows, answer = asyncio.run(ask(build_agent(model), say_text(RECORD["question"])))
        self.assertEqual(answer, RUNS["text turn"]["answer"])
        self.assertGreater(sum(r.get("audio_bytes", 0) for r in rows), 0)
        self.assertTrue(rows[-1].get("turn_complete"))

    def test_spoken_turn_streams_16khz_pcm_chunks(self):
        model = replay("spoken turn")
        asyncio.run(ask(build_agent(model), say_audio(read_pcm(HERE / "question.wav"))))
        sent = [item for _, connection in model.opened for item in connection.received]
        self.assertEqual({mime for kind, mime in sent if kind == "realtime"}, {"audio/pcm;rate=16000"})

    def test_a_44khz_clip_is_refused_before_it_is_sent(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "cd-quality.wav"
            with wave.open(str(path), "wb") as clip:
                clip.setnchannels(1), clip.setsampwidth(2), clip.setframerate(44100)
                clip.writeframes(bytes(4410))
            with self.assertRaisesRegex(ValueError, "not \\(16000 Hz"):
                read_pcm(path)

    def test_ungrounded_answers_fail_the_check(self):
        trials = [run for name, run in RUNS.items() if name.startswith("ungrounded")]
        self.assertTrue(trials)
        for run in trials:
            self.assertFalse(mentions(run["answer"], "Agent Development Kit"), run["answer"])

    def test_text_response_modality_is_rejected(self):
        self.assertIn("response modalities (TEXT) is not supported", FAILURES["TEXT response modality"])

    def test_a_standard_model_is_not_served_over_bidi(self):
        self.assertIn("not supported for bidiGenerateContent", FAILURES["standard Gemini model"])

    def test_run_live_needs_an_existing_session(self):
        runner = Runner(app_name="voice_docs", agent=build_agent(replay("text turn")),
                        session_service=InMemorySessionService())

        async def run():
            async for _ in runner.run_live(user_id="dev", session_id="never-created",
                                           live_request_queue=LiveRequestQueue(), run_config=RunConfig()):
                pass
        with self.assertRaises(SessionNotFoundError):
            asyncio.run(run())


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