Tune a run with RunConfig and read the event loop

"""Replay the recorded model output through tune.py's runner, offline.

Run: python -m unittest -v test_tune.py
The model is replaced by its recorded responses; the Runner, RunConfig handling, session service and tool are real.
"""
import asyncio
import json
import logging
import unittest
import warnings
from pathlib import Path
from unittest import mock

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)

from google.adk import Runner  # noqa: E402
from google.adk.models.base_llm import BaseLlm  # noqa: E402
from google.adk.models.llm_response import LlmResponse  # noqa: E402
from google.adk.sessions import InMemorySessionService  # noqa: E402
from google.genai import types  # noqa: E402

import tune  # noqa: E402

RECORD = json.loads(Path(__file__).with_name("recorded-runs.json").read_text(encoding="utf-8"))


class ReplayLlm(BaseLlm):
    """Yields one recorded response per call and keeps each request; raises when the recording ends."""
    turns: list
    requests: list = []

    async def generate_content_async(self, llm_request, stream=False):
        self.requests.append(llm_request)
        if not self.turns:
            raise RuntimeError("no recorded response for this call")
        yield LlmResponse(content=types.Content.model_validate(self.turns.pop(0)["content"]))


def replay(name: str):
    rec = RECORD[name]
    model = ReplayLlm(model="replay", turns=list(rec["model_turns"]), requests=[])
    runner = Runner(app_name=tune.APP, agent=tune.root_agent.clone(update={"model": model}), session_service=InMemorySessionService())

    async def go():
        tune.LOG.clear()
        session = await runner.session_service.create_session(app_name=tune.APP, user_id=tune.USER)
        error = None
        try:
            for request_id, text, recent, timeout_ms in rec["turns"]:
                await tune.ask(runner, session.id, text, tune.run_config(request_id, recent, timeout_ms))
        except RuntimeError as exc:
            error = exc
        finally:
            await runner.close()
        stored = await runner.session_service.get_session(app_name=tune.APP, user_id=tune.USER, session_id=session.id)
        return model.requests, stored, error

    return asyncio.run(go())


def recorded(name: str, prefix: str) -> list:
    return [line for line in RECORD[name]["log"] if line.startswith(prefix)]


class RunConfigAndEventLoop(unittest.TestCase):
    def setUp(self):
        warnings.simplefilter("ignore")

    def test_a_page_path_is_rejected_before_state_or_file_access(self):
        with mock.patch.object(Path, "is_file", return_value=True), mock.patch.object(Path, "read_text", return_value="private"):
            self.assertIn("error", asyncio.run(tune.read_page("../secret", None)))

    def test_state_commits_when_the_runner_appends_the_event(self):
        replay("tuned")
        self.assertEqual([l for l in tune.LOG if l.startswith(("  tool", "  event"))], recorded("tuned", ("  tool", "  event")))
        tool_line = next(l for l in tune.LOG if l.startswith("  tool"))
        self.assertIn("state[last_page]='rollback' stored[last_page]=None", tool_line)

    def test_custom_metadata_is_stamped_on_every_stored_event(self):
        _, stored, _ = replay("tuned")
        tagged = {e.invocation_id: e.custom_metadata["request_id"] for e in stored.events if e.author != "user"}
        self.assertEqual(sorted(tagged.values()), ["req-1", "req-2", "req-3"])

    def test_num_recent_events_limits_the_history_the_model_receives(self):
        requests, _, _ = replay("tuned")
        self.assertEqual([f"  model  call with {len(r.contents)} contents" for r in requests], recorded("tuned", "  model"))
        self.assertEqual(requests[-1].contents[0].parts[0].text, "Which page covers CI tokens?")

    def test_a_cap_of_3_starts_the_history_at_a_function_call(self):
        requests, _, error = replay("cap_at_function_call")
        self.assertIsNotNone(error)  # the recording has no response: Gemini rejected this request with 400
        first = requests[-1].contents[0]
        self.assertEqual(first.role, "model")
        self.assertIsNotNone(first.parts[0].function_call)
        self.assertIn("INVALID_ARGUMENT", RECORD["cap_at_function_call"]["error"])

    def test_the_http_timeout_reaches_the_model_request(self):
        requests, _, _ = replay("timeout")
        self.assertEqual(requests[0].config.http_options.timeout, 1)
        self.assertEqual(requests[0].config.http_options.retry_options.attempts, 3)
        self.assertTrue(RECORD["timeout"]["error"].startswith("builtins.TimeoutError"))


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

"""Record tune.py on Gemini, plus two RunConfig failures: a history cap that starts at a
function call, and a model-call timeout that is too short.

Run from this folder with GOOGLE_API_KEY set: python record.py
Writes run-output.txt and recorded-runs.json. The tests replay it offline.
"""
import asyncio
import json
import logging
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)  # ADK logs each expected failure with a traceback

from google.adk import Runner  # noqa: E402
from google.adk.sessions import InMemorySessionService  # noqa: E402

import tune  # noqa: E402

HERE = Path(__file__).resolve().parent
MODEL_TURNS: list = []


def keep_turn(callback_context, llm_response):
    """after_model_callback: keep each model response (thought signatures dropped) for the replay tests."""
    if llm_response.content:
        parts = [{k: v for k, v in p.model_dump(mode="json", exclude_none=True).items() if k != "thought_signature"}
                 for p in llm_response.content.parts]
        MODEL_TURNS.append({"content": {"role": "model", "parts": parts}})


async def scenario(turns) -> dict:
    tune.LOG.clear()
    MODEL_TURNS.clear()
    agent = tune.root_agent.clone(update={"after_model_callback": keep_turn})
    runner = Runner(app_name=tune.APP, agent=agent, session_service=InMemorySessionService())
    session = await runner.session_service.create_session(app_name=tune.APP, user_id=tune.USER)
    error = None
    try:
        for request_id, text, recent, timeout_ms in turns:
            tune.LOG.append(f"{request_id}: {text}" + (f"  (num_recent_events={recent})" if recent else "")
                            + (f"  (timeout={timeout_ms} ms)" if timeout_ms != 60_000 else ""))
            events = await tune.ask(runner, session.id, text, tune.run_config(request_id, recent, timeout_ms))
            tune.LOG.append(f"  answer: {events[-1].content.parts[-1].text}")
    except Exception as exc:  # the failures are what this scenario records
        error = f"{type(exc).__module__}.{type(exc).__name__}: {exc}"
        tune.LOG.append(f"  raised {error}")
    finally:
        await runner.close()
    return {"turns": [list(t) for t in turns], "log": list(tune.LOG), "model_turns": list(MODEL_TURNS), "error": error}


BASE = [("req-1", "How do I roll back production?", None, 60_000), ("req-2", "Which page covers CI tokens?", None, 60_000)]
SCENARIOS = {
    "tuned": BASE + [("req-3", "What was my first question?", 4, 60_000)],
    "cap_at_function_call": BASE + [("req-3", "What was my first question?", 3, 60_000)],
    "timeout": [("req-1", "How do I roll back production?", None, 1)],
}


def main() -> None:
    rec = {"model": tune.MODEL, **{name: asyncio.run(scenario(turns)) for name, turns in SCENARIOS.items()}}
    out = [f"$ python tune.py   # {tune.MODEL}", *rec["tuned"]["log"], "",
           "# the same session, capped at 3 events", *rec["cap_at_function_call"]["log"][-3:], "",
           "# a 1 ms model-call timeout", *rec["timeout"]["log"]]
    (HERE / "run-output.txt").write_text("\n".join(out) + "\n", encoding="utf-8")
    (HERE / "recorded-runs.json").write_text(json.dumps(rec, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    print("\n".join(out))


if __name__ == "__main__":
    main()