Keep long sessions in budget with compaction and context caching

"""Replay the recorded Gemini turns offline and check what compaction does to the request.

Run: python -m unittest -v test_assistant.py
"""
import asyncio
import json
import logging
import unittest
import warnings
from pathlib import Path

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

from google.adk.agents.context_cache_config import ContextCacheConfig  # noqa: E402
from google.adk.apps.app import App, EventsCompactionConfig  # noqa: E402
from google.adk.models.base_llm import BaseLlm  # noqa: E402
from google.adk.models.llm_response import LlmResponse  # noqa: E402
from google.genai import types  # noqa: E402

import assistant  # noqa: E402

RECORD = json.loads(Path(__file__).with_name("recorded-runs.json").read_text(encoding="utf-8"))
BUDGETED, PLAIN = RECORD["runs"]


class ReplayLlm(BaseLlm):
    """Returns recorded responses in order and keeps every request it receives."""
    turns: list
    requests: list = []

    async def generate_content_async(self, llm_request, stream=False):
        self.requests.append(llm_request)
        yield LlmResponse(content=types.Content.model_validate(self.turns.pop(0)))


def replay(run: dict, compact: bool):
    """Run the six questions against recorded model turns; return the model requests and the stored session."""
    model = ReplayLlm(model="replay", turns=[c["turn"] for c in run["model_calls"]], requests=[])
    summaries = [{"role": "model", "parts": [{"text": c["summary"]}]} for c in run["compactions"]]
    summarizer = ReplayLlm(model="replay", turns=summaries, requests=[])
    app = assistant.make_app(assistant.root_agent.clone(update={"model": model}), cache=False,
                             compact=compact, summarizer_llm=summarizer)

    async def go():
        runner = assistant.make_runner(app)
        session = await runner.session_service.create_session(app_name="docs", user_id="dana")
        for question in [t["question"] for t in run["turns"]]:
            await assistant.ask(runner, session.id, question)
        stored = await runner.session_service.get_session(app_name="docs", user_id="dana", session_id=session.id)
        await runner.close()
        return model.requests, summarizer.requests, stored
    return asyncio.run(go())


class Compaction(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.budgeted = replay(BUDGETED, compact=True)
        cls.plain = replay(PLAIN, compact=False)

    def test_compaction_runs_after_every_second_turn(self):
        _, summarizer_requests, stored = self.budgeted
        self.assertEqual(len(summarizer_requests), len(BUDGETED["compactions"]))
        self.assertEqual(sum(1 for e in stored.events if e.actions.compaction), len(BUDGETED["compactions"]))

    def test_compaction_shrinks_the_request_history(self):
        budgeted = [len(r.contents) for r in self.budgeted[0]]
        plain = [len(r.contents) for r in self.plain[0]]
        self.assertEqual(budgeted, [c["contents"] for c in BUDGETED["model_calls"]])
        self.assertEqual(plain, [c["contents"] for c in PLAIN["model_calls"]])
        self.assertLess(budgeted[-1], plain[-1])

    def test_raw_events_stay_in_the_session(self):
        stored_budgeted, stored_plain = self.budgeted[2], self.plain[2]
        raw = [e for e in stored_budgeted.events if not e.actions.compaction]
        self.assertEqual(len(raw), len(stored_plain.events))

    def test_the_summary_replaces_old_tool_results(self):
        def logs_in(request):
            return [p.function_response.response["release"] for c in request.contents for p in c.parts
                    if p.function_response and "release" in p.function_response.response]
        self.assertEqual(logs_in(self.budgeted[0][-1]), [])  # both raw deploy logs are gone from the request
        self.assertEqual(logs_in(self.plain[0][-1]), ["r42", "r43"])


class Configuration(unittest.TestCase):
    def test_a_half_sliding_window_is_rejected(self):
        with self.assertRaisesRegex(ValueError, "compaction_interval and overlap_size must be set together"):
            EventsCompactionConfig(compaction_interval=2)

    def test_an_unknown_app_field_is_rejected(self):
        with self.assertRaisesRegex(ValueError, "context_cache"):
            App(name="docs", root_agent=assistant.root_agent, context_cache=ContextCacheConfig())

    def test_the_manual_is_the_static_system_instruction(self):
        self.assertIn("<manual>", assistant.root_agent.static_instruction)
        self.assertIn("docuctl rollback", assistant.root_agent.static_instruction)


class RecordedCacheUse(unittest.TestCase):
    """Checks on the recorded Gemini run: these read recorded-runs.json and make no model call."""

    def test_the_first_call_after_each_compaction_misses_the_cache(self):
        compacted_after = [max(c["turns"]) for c in BUDGETED["compactions"]]
        for turn in BUDGETED["turns"]:
            if turn["turn"] - 1 in compacted_after:
                self.assertEqual(turn["usage"][0]["cache"], "fingerprint only")
                self.assertEqual(turn["usage"][0]["cached"], 0)

    def test_a_short_prefix_never_gets_a_cache(self):
        short = RECORD["failures"]["short_prefix"]
        self.assertTrue(all(u["cache"] == "fingerprint only" for u in short["usage"]))
        self.assertIn("too small for caching", short["cache_log"][0])


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

"""Record one long session twice on Gemini: with caching and compaction, and without either.

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

warnings.filterwarnings("ignore")

from google.adk import Agent  # noqa: E402
from google.adk.agents.context_cache_config import ContextCacheConfig  # noqa: E402
from google.adk.apps.app import App, EventsCompactionConfig  # noqa: E402

import assistant  # noqa: E402

HERE = Path(__file__).resolve().parent
QUESTIONS = [
    "Deploy r42 to staging failed. What went wrong, according to its log?",
    "If I roll staging back, does that undo the migration r42 started?",
    "We fixed the migration and deployed r43, which failed differently. What happened?",
    "How long are staging logs kept?",
    "Which of the two failed deploys, r42 or r43, ran its migration to the end?",
    "What is the command to read the full pre_deploy output of r42?",
]


class CacheLog(logging.Handler):
    """Collect ADK's cache-manager messages: they say why a cache was or was not created."""
    def __init__(self):
        super().__init__(logging.INFO)
        self.lines = []

    def emit(self, record):
        if "gemini_context_cache_manager" in record.name:
            self.lines.append(record.getMessage())


def recording(agent: Agent, calls: list) -> Agent:
    def before(callback_context, llm_request):
        calls.append({"invocation": callback_context.invocation_id, "contents": len(llm_request.contents)})

    def after(callback_context, llm_response):
        content = llm_response.content.model_dump(mode="json", exclude_none=True) if llm_response.content else None
        for part in (content or {}).get("parts", []):
            part.pop("thought_signature", None)
        calls[-1]["turn"] = content
    return agent.clone(update={"before_model_callback": before, "after_model_callback": after})


async def session_run(label: str, cache: bool, compact: bool, agent: Agent = assistant.root_agent,
                      questions: list = QUESTIONS) -> dict:
    calls, log = [], CacheLog()
    logging.getLogger("google_adk").addHandler(log)
    logging.getLogger("google_adk").setLevel(logging.INFO)
    runner = assistant.make_runner(assistant.make_app(recording(agent, calls), cache=cache, compact=compact))
    session = await runner.session_service.create_session(app_name="docs", user_id="dana")
    turns = []
    for n, question in enumerate(questions, 1):
        events = await assistant.ask(runner, session.id, question)
        model_events = [e for e in events if e.usage_metadata]
        answer = next(("".join(p.text or "" for p in e.content.parts if not p.thought).strip()
                       for e in reversed(events) if e.is_final_response() and e.content), "")
        turns.append({"turn": n, "question": question,
                      "tool_calls": [f"{c.name}({c.args})" for e in events for c in e.get_function_calls()],
                      "usage": [assistant.usage(e) for e in model_events], "answer": answer})
    stored = await runner.session_service.get_session(app_name="docs", user_id="dana", session_id=session.id)
    invocations = list(dict.fromkeys(e.invocation_id for e in stored.events if not e.actions.compaction))
    compactions = []
    for e in stored.events:
        if e.actions.compaction:
            c = e.actions.compaction
            covered = sorted({invocations.index(x.invocation_id) + 1 for x in stored.events
                              if not x.actions.compaction and c.start_timestamp <= x.timestamp <= c.end_timestamp})
            compactions.append({"turns": covered, "summary": "".join(p.text or "" for p in c.compacted_content.parts)})
    for call in calls:
        call["turn_index"] = invocations.index(call.pop("invocation")) + 1
    await runner.close()
    logging.getLogger("google_adk").removeHandler(log)
    return {"label": label, "cache": cache, "compact": compact, "turns": turns, "model_calls": calls,
            "compactions": compactions, "stored_events": len(stored.events), "cache_log": log.lines}


async def failures() -> dict:
    out = {}
    # 1. A sliding window needs both numbers; one alone is rejected when the config is built.
    try:
        EventsCompactionConfig(compaction_interval=2)
    except ValueError as error:
        out["half_config"] = str(error).splitlines()[1].strip() if "\n" in str(error) else str(error)
    # 2. A short prefix: with no manual, the cacheable part is below Gemini's minimum cache size.
    short = assistant.root_agent.clone(update={"static_instruction": "You answer questions about docuctl."})
    run = await session_run("short prefix", cache=True, compact=False, agent=short,
                            questions=["What does docuctl rollback do?", "Does it undo migrations?"])
    out["short_prefix"] = {"usage": [u for t in run["turns"] for u in t["usage"]], "cache_log": run["cache_log"]}
    # 3. The App rejects a config object it does not know, so a typo cannot silently turn caching off.
    try:
        App(name="docs", root_agent=assistant.root_agent, context_cache=ContextCacheConfig())
    except ValueError as error:
        out["unknown_field"] = [line for line in str(error).splitlines() if line.strip()][:2]
    return out


def transcript(run: dict) -> str:
    lines = [f"# {run['label']}: cache={run['cache']} compaction={run['compact']}"]
    for t in run["turns"]:
        lines.append(f"> {t['turn']}. {t['question']}")
        lines += [f"call   {c}" for c in t["tool_calls"]]
        lines += [f"model  prompt={u['prompt']} cached={u['cached']} cache={u['cache']}" for u in t["usage"]]
        lines.append("answer " + t["answer"].splitlines()[0][:150])
    lines += [f"compaction of turns {c['turns']}: {len(c['summary'])} characters" for c in run["compactions"]]
    return "\n".join(lines)


async def main() -> None:
    budgeted = await session_run("budgeted", cache=True, compact=True)
    plain = await session_run("plain", cache=False, compact=False)
    record = {"model": assistant.MODEL, "manual_characters": len(assistant.MANUAL),
              "runs": [budgeted, plain], "failures": await failures()}
    (HERE / "recorded-runs.json").write_text(json.dumps(record, indent=2) + "\n", encoding="utf-8")
    (HERE / "run-output.txt").write_text(transcript(budgeted) + "\n\n" + transcript(plain) + "\n", encoding="utf-8")
    print(transcript(budgeted) + "\n\n" + transcript(plain))
    print(json.dumps(record["failures"], indent=2))


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