feat(config): add --type flag to adk create to allow starting with config

Updated the `adk create` default model version to gemini-2.5-flash.

PiperOrigin-RevId: 788589859
This commit is contained in:
Liang Wu
2025-07-29 13:45:11 -07:00
committed by Copybara-Service
parent 2f73cfde18
commit bcac9ba44c
3 changed files with 126 additions and 16 deletions
+67 -13
View File
@@ -14,6 +14,7 @@
from __future__ import annotations
import enum
import os
import subprocess
from typing import Optional
@@ -21,6 +22,12 @@ from typing import Tuple
import click
class Type(enum.Enum):
CONFIG = "config"
CODE = "code"
_INIT_PY_TEMPLATE = """\
from . import agent
"""
@@ -36,6 +43,13 @@ root_agent = Agent(
)
"""
_AGENT_CONFIG_TEMPLATE = """\
name: root_agent
description: A helpful assistant for user questions.
instruction: Answer user questions to the best of your knowledge
model: {model_name}
"""
_GOOGLE_API_MSG = """
Don't have API Key? Create one in AI Studio: https://aistudio.google.com/apikey
@@ -51,13 +65,20 @@ Please see below guide to configure other models:
https://google.github.io/adk-docs/agents/models
"""
_SUCCESS_MSG = """
_SUCCESS_MSG_CODE = """
Agent created in {agent_folder}:
- .env
- __init__.py
- agent.py
"""
_SUCCESS_MSG_CONFIG = """
Agent created in {agent_folder}:
- .env
- __init__.py
- root_agent.yaml
"""
def _get_gcp_project_from_gcloud() -> str:
"""Uses gcloud to get default project."""
@@ -158,13 +179,15 @@ def _generate_files(
google_cloud_project: Optional[str] = None,
google_cloud_region: Optional[str] = None,
model: Optional[str] = None,
type: Optional[Type] = None,
):
"""Generates a folder name for the agent."""
os.makedirs(agent_folder, exist_ok=True)
dotenv_file_path = os.path.join(agent_folder, ".env")
init_file_path = os.path.join(agent_folder, "__init__.py")
agent_file_path = os.path.join(agent_folder, "agent.py")
agent_py_file_path = os.path.join(agent_folder, "agent.py")
agent_config_file_path = os.path.join(agent_folder, "root_agent.yaml")
with open(dotenv_file_path, "w", encoding="utf-8") as f:
lines = []
@@ -180,29 +203,38 @@ def _generate_files(
lines.append(f"GOOGLE_CLOUD_LOCATION={google_cloud_region}")
f.write("\n".join(lines))
with open(init_file_path, "w", encoding="utf-8") as f:
f.write(_INIT_PY_TEMPLATE)
if type == Type.CONFIG:
with open(agent_config_file_path, "w", encoding="utf-8") as f:
f.write(_AGENT_CONFIG_TEMPLATE.format(model_name=model))
with open(init_file_path, "w", encoding="utf-8") as f:
f.write("")
click.secho(
_SUCCESS_MSG_CONFIG.format(agent_folder=agent_folder),
fg="green",
)
else:
with open(init_file_path, "w", encoding="utf-8") as f:
f.write(_INIT_PY_TEMPLATE)
with open(agent_file_path, "w", encoding="utf-8") as f:
f.write(_AGENT_PY_TEMPLATE.format(model_name=model))
click.secho(
_SUCCESS_MSG.format(agent_folder=agent_folder),
fg="green",
)
with open(agent_py_file_path, "w", encoding="utf-8") as f:
f.write(_AGENT_PY_TEMPLATE.format(model_name=model))
click.secho(
_SUCCESS_MSG_CODE.format(agent_folder=agent_folder),
fg="green",
)
def _prompt_for_model() -> str:
model_choice = click.prompt(
"""\
Choose a model for the root agent:
1. gemini-2.0-flash-001
1. gemini-2.5-flash
2. Other models (fill later)
Choose model""",
type=click.Choice(["1", "2"]),
)
if model_choice == "1":
return "gemini-2.0-flash-001"
return "gemini-2.5-flash"
else:
click.secho(_OTHER_MODEL_MSG, fg="green")
return "<FILL_IN_MODEL>"
@@ -231,6 +263,22 @@ def _prompt_to_choose_backend(
return google_api_key, google_cloud_project, google_cloud_region
def _prompt_to_choose_type() -> Type:
"""Prompts user to choose type of agent to create."""
type_choice = click.prompt(
"""\
Choose a type for the root agent:
1. YAML config (experimental, may change without notice)
2. Code
Choose type""",
type=click.Choice(["1", "2"]),
)
if type_choice == "1":
return Type.CONFIG
else:
return Type.CODE
def run_cmd(
agent_name: str,
*,
@@ -238,6 +286,7 @@ def run_cmd(
google_api_key: Optional[str],
google_cloud_project: Optional[str],
google_cloud_region: Optional[str],
type: Optional[Type],
):
"""Runs `adk create` command to create agent template.
@@ -249,6 +298,7 @@ def run_cmd(
VertexAI as backend.
google_cloud_region: Optional[str], The Google Cloud region for using
VertexAI as backend.
type: Optional[Type], Whether to define agent with config file or code.
"""
agent_folder = os.path.join(os.getcwd(), agent_name)
# check folder doesn't exist or it's empty. Otherwise, throw
@@ -272,10 +322,14 @@ def run_cmd(
)
)
if not type:
type = _prompt_to_choose_type()
_generate_files(
agent_folder,
google_api_key=google_api_key,
google_cloud_project=google_cloud_project,
google_cloud_region=google_cloud_region,
model=model,
type=type,
)
+14 -2
View File
@@ -33,8 +33,6 @@ from . import cli_create
from . import cli_deploy
from .. import version
from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE
from ..evaluation.gcs_eval_set_results_manager import GcsEvalSetResultsManager
from ..evaluation.gcs_eval_sets_manager import GcsEvalSetsManager
from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
from ..sessions.in_memory_session_service import InMemorySessionService
from .cli import run_cli
@@ -147,6 +145,18 @@ def deploy():
type=str,
help="Optional. The Google Cloud Region for using VertexAI as backend.",
)
@click.option(
"--type",
type=click.Choice([t.value for t in cli_create.Type]),
help=(
"EXPERIMENTAL Optional. Type of agent to create: 'config' or 'code'."
" 'config' is not ready for use so it defaults to 'code'. It may change"
" later once 'config' is ready for use."
),
default=cli_create.Type.CODE.value,
show_default=True,
hidden=True, # Won't show in --help output. Not ready for use.
)
@click.argument("app_name", type=str, required=True)
def cli_create_cmd(
app_name: str,
@@ -154,6 +164,7 @@ def cli_create_cmd(
api_key: Optional[str],
project: Optional[str],
region: Optional[str],
type: Optional[cli_create.Type],
):
"""Creates a new app in the current folder with prepopulated agent template.
@@ -169,6 +180,7 @@ def cli_create_cmd(
google_api_key=api_key,
google_cloud_project=project,
google_cloud_region=region,
type=type,
)
+45 -1
View File
@@ -147,9 +147,53 @@ def test_run_cmd_overwrite_reject(
google_api_key=None,
google_cloud_project=None,
google_cloud_region=None,
type=cli_create.Type.CODE,
)
def test_run_cmd_with_type_config(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
"""run_cmd with --type=config should generate YAML config file."""
agent_name = "test_agent"
monkeypatch.setattr(os, "getcwd", lambda: str(tmp_path))
monkeypatch.setattr(os.path, "exists", lambda _p: False)
cli_create.run_cmd(
agent_name,
model="gemini-2.0-flash-001",
google_api_key="test-key",
google_cloud_project=None,
google_cloud_region=None,
type=cli_create.Type.CONFIG,
)
agent_dir = tmp_path / agent_name
assert agent_dir.exists()
# Should create root_agent.yaml instead of agent.py
yaml_file = agent_dir / "root_agent.yaml"
assert yaml_file.exists()
assert not (agent_dir / "agent.py").exists()
# Check YAML content
yaml_content = yaml_file.read_text()
assert "name: root_agent" in yaml_content
assert "model: gemini-2.0-flash-001" in yaml_content
assert "description: A helpful assistant for user questions." in yaml_content
# Should create empty __init__.py
init_file = agent_dir / "__init__.py"
assert init_file.exists()
assert init_file.read_text().strip() == ""
# Should still create .env file
env_file = agent_dir / ".env"
assert env_file.exists()
assert "GOOGLE_API_KEY=test-key" in env_file.read_text()
# Prompt helpers
def test_prompt_for_google_cloud(monkeypatch: pytest.MonkeyPatch) -> None:
"""Prompt should return the project input."""
@@ -174,7 +218,7 @@ def test_prompt_for_google_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
def test_prompt_for_model_gemini(monkeypatch: pytest.MonkeyPatch) -> None:
"""Selecting option '1' should return the default Gemini model string."""
monkeypatch.setattr(click, "prompt", lambda *a, **k: "1")
assert cli_create._prompt_for_model() == "gemini-2.0-flash-001"
assert cli_create._prompt_for_model() == "gemini-2.5-flash"
def test_prompt_for_model_other(monkeypatch: pytest.MonkeyPatch) -> None: