From 1990895a9e23ac19451d89253aacd93c242813da Mon Sep 17 00:00:00 2001 From: Bicheng Wang Date: Thu, 8 Oct 2026 13:13:13 -0700 Subject: [PATCH] feat(cli): add adk graph command Renders an agent's structure with the existing agent_graph builder, so the tree can be inspected without starting the dev UI. Prints DOT by default and can write DOT, PNG or SVG to a file (PNG and SVG need Graphviz). --ascii-only drops the emoji label prefix for consoles that cannot encode it. --- src/google/adk/cli/cli_tools_click.py | 104 ++++++++++ .../cli/utils/test_cli_tools_click.py | 178 ++++++++++++++++++ 2 files changed, 282 insertions(+) diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index 8b08978f988..806cc42d75e 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -24,6 +24,7 @@ import logging import os from pathlib import Path +import re import sys import tempfile import textwrap @@ -1171,6 +1172,109 @@ def cli_run( ) +@main.command("graph", cls=HelpfulCommand) +@click.option( + "--format", + "output_format", + type=click.Choice(["dot", "png", "svg"]), + default="dot", + show_default=True, + help=( + "Optional. Output format. png and svg are rendered by Graphviz, which" + " must be installed." + ), +) +@click.option( + "--output", + type=click.Path(dir_okay=False, file_okay=True, resolve_path=True), + help=( + "Optional. File to write the graph to. Defaults to stdout for dot;" + " required for png and svg." + ), +) +@click.option( + "--ascii-only", + is_flag=True, + help=( + "Optional. Drop the emoji prefix from node labels in the dot output," + " for consoles that cannot encode them. Non-ASCII agent names are kept." + ), +) +@click.argument( + "agent", + type=click.Path( + exists=True, dir_okay=True, file_okay=False, resolve_path=True + ), +) +def cli_graph( + agent: str, + output_format: str, + output: Optional[str], + ascii_only: bool, +): + """Renders an agent's structure as a Graphviz graph. + + AGENT: The path to the agent source code folder. + + Example: + + adk graph path/to/my_agent + adk graph path/to/my_agent --format svg --output my_agent.svg + """ + import graphviz + + from ..apps.app import App + from .agent_graph import get_agent_graph + from .utils.agent_loader import AgentLoader + + if output_format != "dot" and not output: + raise click.UsageError( + f"--output is required for --format {output_format}." + ) + if ascii_only and output_format != "dot": + raise click.UsageError("--ascii-only only applies to --format dot.") + + agent_dir = Path(agent) + try: + loaded = AgentLoader(agents_dir=str(agent_dir.parent)).load_agent( + agent_dir.name + ) + except Exception as e: # pylint: disable=broad-exception-caught + raise click.ClickException(f"Failed to load agent '{agent}': {e}") from e + root_agent = loaded.root_agent if isinstance(loaded, App) else loaded + + try: + graph = asyncio.run(get_agent_graph(root_agent, [])) + except Exception as e: # pylint: disable=broad-exception-caught + raise click.ClickException(f"Failed to build the agent graph: {e}") from e + + if output_format == "dot": + dot = graph.source + if ascii_only: + dot = re.sub(r'label="[^\x00-\x7f]+ ', 'label="', dot) + if not output: + click.echo(dot) + return + content = dot.encode("utf-8") + else: + try: + content = graph.pipe(format=output_format) + except graphviz.ExecutableNotFound as e: + raise click.ClickException( + "Graphviz 'dot' executable not found; install Graphviz to render" + f" {output_format}, or use --format dot. ({e})" + ) from e + except graphviz.CalledProcessError as e: + raise click.ClickException( + f"Graphviz failed to render {output_format}: {e}" + ) from e + + try: + Path(output).write_bytes(content) + except OSError as e: + raise click.ClickException(f"Failed to write '{output}': {e}") from e + + @main.command( "test", cls=HelpfulCommand, diff --git a/tests/unittests/cli/utils/test_cli_tools_click.py b/tests/unittests/cli/utils/test_cli_tools_click.py index bd31ceedbc8..a47d3a08aeb 100644 --- a/tests/unittests/cli/utils/test_cli_tools_click.py +++ b/tests/unittests/cli/utils/test_cli_tools_click.py @@ -22,6 +22,7 @@ import logging import os from pathlib import Path +import re import sys from types import SimpleNamespace from typing import Any @@ -3624,3 +3625,180 @@ def explode(*args: Any, **kwargs: Any) -> None: assert result.exit_code == 1 assert "Migration failed: destination schema is newer" in result.output + + +@pytest.fixture +def graph_agent_dir(tmp_path: Path) -> Path: + # A name unique to the test, so sys.modules never serves a stale package. + agent_dir = tmp_path / f"graph_agent_{tmp_path.name}" + agent_dir.mkdir() + (agent_dir / "__init__.py").write_text("from . import agent\n") + (agent_dir / "agent.py").write_text( + "from google.adk.agents import LlmAgent\n" + "def lookup(q: str) -> str:\n" + ' """Looks something up."""\n' + " return q\n" + 'helper = LlmAgent(name="helper", model="gemini-2.0-flash",' + " tools=[lookup])\n" + 'root_agent = LlmAgent(name="root", model="gemini-2.0-flash",' + " sub_agents=[helper])\n" + ) + return agent_dir + + +@pytest.mark.unmute_click +def test_cli_graph_prints_dot_to_stdout(graph_agent_dir: Path) -> None: + result = CliRunner().invoke( + cli_tools_click.main, ["graph", str(graph_agent_dir)] + ) + + assert result.exit_code == 0, result.output + assert "root -> helper" in result.output + assert "helper -> lookup" in result.output + assert "🤖 root" in result.output + + +@pytest.mark.unmute_click +def test_cli_graph_ascii_only_strips_non_ascii( + graph_agent_dir: Path, tmp_path: Path +) -> None: + out = tmp_path / "graph.dot" + + result = CliRunner().invoke( + cli_tools_click.main, + ["graph", str(graph_agent_dir), "--ascii-only", "--output", str(out)], + ) + + assert result.exit_code == 0, result.output + dot = out.read_text(encoding="utf-8") + assert dot.isascii() + assert 'label="root"' in dot + assert 'label="lookup"' in dot + + +@pytest.mark.unmute_click +def test_cli_graph_image_format_requires_output(graph_agent_dir: Path) -> None: + result = CliRunner().invoke( + cli_tools_click.main, ["graph", str(graph_agent_dir), "--format", "svg"] + ) + + assert result.exit_code == 2 + assert "--output is required" in result.output + + +@pytest.mark.unmute_click +def test_cli_graph_ascii_only_rejects_image_format( + graph_agent_dir: Path, tmp_path: Path +) -> None: + result = CliRunner().invoke( + cli_tools_click.main, + [ + "graph", + str(graph_agent_dir), + "--format", + "png", + "--ascii-only", + "--output", + str(tmp_path / "g.png"), + ], + ) + + assert result.exit_code == 2 + assert "--ascii-only only applies" in result.output + + +@pytest.mark.unmute_click +def test_cli_graph_writes_rendered_image( + graph_agent_dir: Path, tmp_path: Path +) -> None: + out = tmp_path / "graph.svg" + + with mock.patch("graphviz.Digraph.pipe", return_value=b"") as pipe: + result = CliRunner().invoke( + cli_tools_click.main, + [ + "graph", + str(graph_agent_dir), + "--format", + "svg", + "--output", + str(out), + ], + ) + + assert result.exit_code == 0, result.output + pipe.assert_called_once_with(format="svg") + assert out.read_bytes() == b"" + + +@pytest.mark.unmute_click +def test_cli_graph_reports_missing_graphviz( + graph_agent_dir: Path, tmp_path: Path +) -> None: + import graphviz + + with mock.patch( + "graphviz.Digraph.pipe", side_effect=graphviz.ExecutableNotFound(["dot"]) + ): + result = CliRunner().invoke( + cli_tools_click.main, + [ + "graph", + str(graph_agent_dir), + "--format", + "png", + "--output", + str(tmp_path / "g.png"), + ], + ) + + assert result.exit_code == 1 + assert "install Graphviz" in result.output + assert not (tmp_path / "g.png").exists() + + +@pytest.mark.unmute_click +def test_cli_graph_ascii_only_keeps_non_ascii_names_distinct( + tmp_path: Path, +) -> None: + agent_dir = tmp_path / f"graph_names_{tmp_path.name}" + agent_dir.mkdir() + (agent_dir / "__init__.py").write_text("from . import agent\n") + (agent_dir / "agent.py").write_text( + "from google.adk.agents import LlmAgent\n" + 'a = LlmAgent(name="助手", model="gemini-2.0-flash")\n' + 'b = LlmAgent(name="café", model="gemini-2.0-flash")\n' + 'root_agent = LlmAgent(name="root", model="gemini-2.0-flash",' + " sub_agents=[a, b])\n", + encoding="utf-8", + ) + out = tmp_path / "graph.dot" + + result = CliRunner().invoke( + cli_tools_click.main, + ["graph", str(agent_dir), "--ascii-only", "--output", str(out)], + ) + + assert result.exit_code == 0, result.output + dot = out.read_text(encoding="utf-8") + assert re.search(r'root -> "?助手"? ', dot) + assert re.search(r'root -> "?café"? ', dot) + assert 'label="助手"' in dot + + +@pytest.mark.unmute_click +def test_cli_graph_reports_unwritable_output( + graph_agent_dir: Path, tmp_path: Path +) -> None: + result = CliRunner().invoke( + cli_tools_click.main, + [ + "graph", + str(graph_agent_dir), + "--output", + str(tmp_path / "missing_dir" / "graph.dot"), + ], + ) + + assert result.exit_code == 1 + assert "Failed to write" in result.output