Configure an ADK Agent with YAML

"""Check root_agent.yaml, its load failures and the Visual Builder save endpoint, offline.

Run: python -m unittest -v test_config.py
The model is replaced by its recorded responses; the YAML loader, Runner and dev server are real.
"""
import asyncio
import json
import logging
import tempfile
import unittest
import warnings
from pathlib import Path

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

from fastapi.testclient import TestClient  # noqa: E402
from google.adk.agents.config_agent_utils import from_config  # noqa: E402
from google.adk.cli.fast_api import get_fast_api_app  # 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 load  # noqa: E402
import record  # noqa: E402

HERE = Path(__file__).resolve().parent
RECORD = json.loads((HERE / "recorded-runs.json").read_text(encoding="utf-8"))


class ReplayLlm(BaseLlm):
    turns: list

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


def load_text(text: str):
    with tempfile.TemporaryDirectory() as tmp:
        path = Path(tmp) / "root_agent.yaml"
        path.write_text(text, encoding="utf-8")
        return from_config(str(path))


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

    def test_the_yaml_loads_as_an_llm_agent_with_its_tool(self):
        agent = load.load()
        self.assertEqual((type(agent).__name__, agent.name, agent.model), ("LlmAgent", "docs_assistant", RECORD["model"]))
        self.assertEqual([t.__name__ for t in agent.tools], ["read_page"])
        self.assertEqual(agent.generate_content_config.temperature, 0.2)

    def test_the_loaded_agent_runs_through_a_runner(self):
        agent = load.load().clone(update={"model": ReplayLlm(model="replay", turns=list(RECORD["load"]["model_turns"]))})
        events = asyncio.run(load.ask(agent, RECORD["question"]))
        self.assertEqual([load.describe(e) for e in events], RECORD["load"]["lines"])
        self.assertEqual(events[-1].content.parts[-1].text, RECORD["load"]["answer"])

    def test_a_misspelled_field_fails_validation(self):
        with self.assertRaisesRegex(Exception, "instructions\n  Extra inputs are not permitted"):
            load_text(record.BROKEN["misspelled_field"])

    def test_a_missing_tool_function_fails_to_load(self):
        with self.assertRaisesRegex(AttributeError, "has no attribute 'read_pages'"):
            load_text(record.BROKEN["missing_tool"])

    def test_the_loader_blocks_standard_library_references(self):
        with self.assertRaisesRegex(ValueError, "Blocked module reference: 'os.system'"):
            load_text(record.BROKEN["stdlib_tool"])


class VisualBuilderSave(unittest.TestCase):
    def setUp(self):
        warnings.simplefilter("ignore")
        # The builder writes into the app folder, so the test saves into this folder and restores the file.
        # (A temporary copy would not work here: docs_config also imports from this folder, and the
        # builder blocks an app whose name resolves to code outside its own folder.)
        self.yaml = HERE / "docs_config" / "root_agent.yaml"
        self.app = get_fast_api_app(agents_dir=str(HERE), web=True, session_service_uri="memory://", artifact_service_uri="memory://")
        self.local = TestClient(self.app, client=("127.0.0.1", 50000))

    def tearDown(self):
        self.yaml.write_text(record.YAML, encoding="utf-8")

    def save(self, name: str, http=None):
        files = [("files", ("docs_config/root_agent.yaml", record.UPLOADS[name].encode()))]
        return (http or self.local).post("/dev/apps/docs_config/builder/save", files=files)

    def test_a_builder_save_rewrites_root_agent_yaml(self):
        self.assertEqual(self.local.get("/dev/apps/docs_config/builder").text, record.YAML)
        self.assertEqual(self.save("edit").json(), True)
        self.assertIn("under 80 words", self.yaml.read_text(encoding="utf-8"))

    def test_a_builder_upload_with_an_args_key_is_rejected(self):
        r = self.save("args_key")
        self.assertEqual(r.status_code, 400)
        self.assertEqual(json.dumps(r.json(), separators=(",", ":")), RECORD["builder"]["args_key"].split("\nHTTP")[0])
        self.assertNotIn("args:", self.yaml.read_text(encoding="utf-8"))

    def test_a_builder_upload_naming_code_outside_the_app_is_rejected(self):
        r = self.save("outside_code")
        self.assertEqual(r.status_code, 400)
        self.assertIn("Blocked code reference 'subprocess.run'", r.json()["detail"])

    def test_a_remote_client_cannot_reach_the_builder(self):
        remote = TestClient(self.app, client=("10.0.0.7", 50000))
        r = self.save("edit", http=remote)
        self.assertEqual(r.status_code, 403)
        self.assertIn("ADK_ALLOW_REMOTE_AGENT_BUILDER", r.json()["detail"])
        self.assertNotIn("under 80 words", self.yaml.read_text(encoding="utf-8"))


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

"""Record the YAML docs assistant: adk run, load.py, three load failures and the Visual Builder endpoints.

Run from this folder with GOOGLE_API_KEY set: python record.py
Writes run-output.txt and recorded-runs.json. The builder calls run against a temporary copy of
this folder, because a save rewrites root_agent.yaml.
"""
import asyncio
import json
import os
import shutil
import subprocess
import sys
import tempfile
import time
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")

import httpx  # noqa: E402
from google.adk.agents.config_agent_utils import from_config  # noqa: E402

import load  # noqa: E402

HERE = Path(__file__).resolve().parent
ADK = str(Path(sys.executable).with_name("adk"))
QUESTION = "How do I roll back production?"
ENV = {**os.environ, "GOOGLE_GENAI_USE_ENTERPRISE": "0", "GOOGLE_GENAI_USE_VERTEXAI": "0", "PYTHONWARNINGS": "ignore"}
YAML = (HERE / "docs_config" / "root_agent.yaml").read_text(encoding="utf-8")
BROKEN = {
    "misspelled_field": YAML.replace("instruction:", "instructions:"),
    "missing_tool": YAML.replace("tools.read_page", "tools.read_pages"),
    "stdlib_tool": YAML.replace("docs_config.tools.read_page", "os.system"),
}
UPLOADS = {
    "edit": YAML.replace("under 100 words", "under 80 words"),
    "args_key": YAML.replace("  - name: docs_config.tools.read_page", "  - name: docs_config.tools.read_page\n    args:\n      page: tokens"),
    "outside_code": YAML.replace("docs_config.tools.read_page", "subprocess.run"),
}


def curl(*args: str) -> str:
    proc = subprocess.run(["curl", "-s", "-w", "\nHTTP %{http_code}", *args], capture_output=True, text=True, timeout=30)
    return proc.stdout


def main() -> None:
    rec, out = {"model": "gemini-3.8-flash", "question": QUESTION}, []
    with tempfile.TemporaryDirectory() as tmp:
        cmd = ["create", "--type=config", "--model", "gemini-3.8-flash", "--api_key", "your-gemini-api-key", "docs_config"]
        proc = subprocess.run([ADK, *cmd], cwd=tmp, env=ENV, capture_output=True, text=True, timeout=60)
        made = Path(tmp) / "docs_config"
        rec["adk_create"] = {"command": "adk " + " ".join(cmd), "files": sorted(p.name for p in made.iterdir()),
                             "root_agent.yaml": (made / "root_agent.yaml").read_text(encoding="utf-8")}
    out += [f"$ {rec['adk_create']['command']}", f"created docs_config/: {', '.join(rec['adk_create']['files'])}",
            rec["adk_create"]["root_agent.yaml"].rstrip(), ""]
    proc = subprocess.run([ADK, "run", "--in_memory", "docs_config", QUESTION], cwd=HERE, env=ENV, capture_output=True, text=True, timeout=300)
    rec["adk_run"] = [l for l in proc.stdout.splitlines() if l.strip() and not l.startswith(("Log setup", "To access latest log"))]
    out += [f"$ adk run --in_memory docs_config {QUESTION!r}", *rec["adk_run"], ""]

    agent = load.load()
    events = asyncio.run(load.ask(agent, QUESTION))
    rec["load"] = {"lines": [load.describe(e) for e in events], "answer": events[-1].content.parts[-1].text,
                   "model_turns": [{"content": {"role": "model", "parts": [
                       {k: v for k, v in p.model_dump(mode="json", exclude_none=True).items() if k != "thought_signature"}
                       for p in e.content.parts]}} for e in events if e.content and e.content.role == "model"]}
    cfg = agent.generate_content_config
    out += [f"$ python load.py {QUESTION!r}", f"LlmAgent {agent.name} model={agent.model} tools=['read_page']",
            f"temperature={cfg.temperature} thinking_level={cfg.thinking_config.thinking_level.value}",
            *rec["load"]["lines"], f"answer: {rec['load']['answer']}", ""]

    rec["load_failures"] = {}
    with tempfile.TemporaryDirectory() as tmp:
        for name, text in BROKEN.items():
            path = Path(tmp) / f"{name}.yaml"
            path.write_text(text, encoding="utf-8")
            try:
                from_config(str(path))
                rec["load_failures"][name] = "loaded"
            except Exception as error:  # the three failures raise three different types
                rec["load_failures"][name] = f"{type(error).__name__}: {error}"
        out += [f"from_config {name}: {msg.splitlines()[0]}" for name, msg in rec["load_failures"].items()] + [""]

        work = Path(tmp) / "agents"
        shutil.copytree(HERE / "docs_config", work / "docs_config", ignore=shutil.ignore_patterns("__pycache__", ".env"))
        shutil.copytree(HERE / "docs", work / "docs")
        for name, text in UPLOADS.items():
            (Path(tmp) / f"{name}.yaml").write_text(text, encoding="utf-8")
        server = subprocess.Popen([ADK, "web", "--port", "8766", "."], cwd=work, env=ENV, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
        try:
            for _ in range(60):
                try:
                    httpx.get("http://127.0.0.1:8766/list-apps")
                    break
                except httpx.ConnectError:
                    time.sleep(0.5)
            base = "http://127.0.0.1:8766/dev/apps/docs_config/builder"
            got = curl(base)
            rec["builder"] = {"get": got.splitlines()[-1], "get_matches_file": got.rsplit("\nHTTP", 1)[0] == YAML}
            for name in UPLOADS:
                rec["builder"][name] = curl("-X", "POST", f"{base}/save", "-F", f"files=@{tmp}/{name}.yaml;filename=docs_config/root_agent.yaml")
            rec["builder"]["saved_has_edit"] = "under 80 words" in (work / "docs_config" / "root_agent.yaml").read_text(encoding="utf-8")
        finally:
            server.terminate()
            server.wait(timeout=20)
    out += ["$ adk web --port 8766 .", "$ curl -s http://127.0.0.1:8766/dev/apps/docs_config/builder", rec["builder"]["get"]]
    for name in UPLOADS:
        out += [f"$ curl -s -X POST {'http://127.0.0.1:8766/dev/apps/docs_config/builder'}/save -F 'files=@{name}.yaml;filename=docs_config/root_agent.yaml'",
                rec["builder"][name]]
    (HERE / "run-output.txt").write_text("\n".join(out).rstrip() + "\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()