Skip to content

Commit bcac9ba

Browse files
wuliang229copybara-github
authored andcommitted
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
1 parent 2f73cfd commit bcac9ba

File tree

3 files changed

+127
-17
lines changed

3 files changed

+127
-17
lines changed

src/google/adk/cli/cli_create.py

Lines changed: 68 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,20 @@
1414

1515
from __future__ import annotations
1616

17+
import enum
1718
import os
1819
import subprocess
1920
from typing import Optional
2021
from typing import Tuple
2122

2223
import click
2324

25+
26+
class Type(enum.Enum):
27+
CONFIG = "config"
28+
CODE = "code"
29+
30+
2431
_INIT_PY_TEMPLATE = """\
2532
from . import agent
2633
"""
@@ -36,6 +43,13 @@
3643
)
3744
"""
3845

46+
_AGENT_CONFIG_TEMPLATE = """\
47+
name: root_agent
48+
description: A helpful assistant for user questions.
49+
instruction: Answer user questions to the best of your knowledge
50+
model: {model_name}
51+
"""
52+
3953

4054
_GOOGLE_API_MSG = """
4155
Don't have API Key? Create one in AI Studio: https://aistudio.google.com/apikey
@@ -51,13 +65,20 @@
5165
https://google.github.io/adk-docs/agents/models
5266
"""
5367

54-
_SUCCESS_MSG = """
68+
_SUCCESS_MSG_CODE = """
5569
Agent created in {agent_folder}:
5670
- .env
5771
- __init__.py
5872
- agent.py
5973
"""
6074

75+
_SUCCESS_MSG_CONFIG = """
76+
Agent created in {agent_folder}:
77+
- .env
78+
- __init__.py
79+
- root_agent.yaml
80+
"""
81+
6182

6283
def _get_gcp_project_from_gcloud() -> str:
6384
"""Uses gcloud to get default project."""
@@ -158,13 +179,15 @@ def _generate_files(
158179
google_cloud_project: Optional[str] = None,
159180
google_cloud_region: Optional[str] = None,
160181
model: Optional[str] = None,
182+
type: Optional[Type] = None,
161183
):
162184
"""Generates a folder name for the agent."""
163185
os.makedirs(agent_folder, exist_ok=True)
164186

165187
dotenv_file_path = os.path.join(agent_folder, ".env")
166188
init_file_path = os.path.join(agent_folder, "__init__.py")
167-
agent_file_path = os.path.join(agent_folder, "agent.py")
189+
agent_py_file_path = os.path.join(agent_folder, "agent.py")
190+
agent_config_file_path = os.path.join(agent_folder, "root_agent.yaml")
168191

169192
with open(dotenv_file_path, "w", encoding="utf-8") as f:
170193
lines = []
@@ -180,29 +203,38 @@ def _generate_files(
180203
lines.append(f"GOOGLE_CLOUD_LOCATION={google_cloud_region}")
181204
f.write("\n".join(lines))
182205

183-
with open(init_file_path, "w", encoding="utf-8") as f:
184-
f.write(_INIT_PY_TEMPLATE)
185-
186-
with open(agent_file_path, "w", encoding="utf-8") as f:
187-
f.write(_AGENT_PY_TEMPLATE.format(model_name=model))
188-
189-
click.secho(
190-
_SUCCESS_MSG.format(agent_folder=agent_folder),
191-
fg="green",
192-
)
206+
if type == Type.CONFIG:
207+
with open(agent_config_file_path, "w", encoding="utf-8") as f:
208+
f.write(_AGENT_CONFIG_TEMPLATE.format(model_name=model))
209+
with open(init_file_path, "w", encoding="utf-8") as f:
210+
f.write("")
211+
click.secho(
212+
_SUCCESS_MSG_CONFIG.format(agent_folder=agent_folder),
213+
fg="green",
214+
)
215+
else:
216+
with open(init_file_path, "w", encoding="utf-8") as f:
217+
f.write(_INIT_PY_TEMPLATE)
218+
219+
with open(agent_py_file_path, "w", encoding="utf-8") as f:
220+
f.write(_AGENT_PY_TEMPLATE.format(model_name=model))
221+
click.secho(
222+
_SUCCESS_MSG_CODE.format(agent_folder=agent_folder),
223+
fg="green",
224+
)
193225

194226

195227
def _prompt_for_model() -> str:
196228
model_choice = click.prompt(
197229
"""\
198230
Choose a model for the root agent:
199-
1. gemini-2.0-flash-001
231+
1. gemini-2.5-flash
200232
2. Other models (fill later)
201233
Choose model""",
202234
type=click.Choice(["1", "2"]),
203235
)
204236
if model_choice == "1":
205-
return "gemini-2.0-flash-001"
237+
return "gemini-2.5-flash"
206238
else:
207239
click.secho(_OTHER_MODEL_MSG, fg="green")
208240
return "<FILL_IN_MODEL>"
@@ -231,13 +263,30 @@ def _prompt_to_choose_backend(
231263
return google_api_key, google_cloud_project, google_cloud_region
232264

233265

266+
def _prompt_to_choose_type() -> Type:
267+
"""Prompts user to choose type of agent to create."""
268+
type_choice = click.prompt(
269+
"""\
270+
Choose a type for the root agent:
271+
1. YAML config (experimental, may change without notice)
272+
2. Code
273+
Choose type""",
274+
type=click.Choice(["1", "2"]),
275+
)
276+
if type_choice == "1":
277+
return Type.CONFIG
278+
else:
279+
return Type.CODE
280+
281+
234282
def run_cmd(
235283
agent_name: str,
236284
*,
237285
model: Optional[str],
238286
google_api_key: Optional[str],
239287
google_cloud_project: Optional[str],
240288
google_cloud_region: Optional[str],
289+
type: Optional[Type],
241290
):
242291
"""Runs `adk create` command to create agent template.
243292
@@ -249,6 +298,7 @@ def run_cmd(
249298
VertexAI as backend.
250299
google_cloud_region: Optional[str], The Google Cloud region for using
251300
VertexAI as backend.
301+
type: Optional[Type], Whether to define agent with config file or code.
252302
"""
253303
agent_folder = os.path.join(os.getcwd(), agent_name)
254304
# check folder doesn't exist or it's empty. Otherwise, throw
@@ -272,10 +322,14 @@ def run_cmd(
272322
)
273323
)
274324

325+
if not type:
326+
type = _prompt_to_choose_type()
327+
275328
_generate_files(
276329
agent_folder,
277330
google_api_key=google_api_key,
278331
google_cloud_project=google_cloud_project,
279332
google_cloud_region=google_cloud_region,
280333
model=model,
334+
type=type,
281335
)

src/google/adk/cli/cli_tools_click.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,6 @@
3333
from . import cli_deploy
3434
from .. import version
3535
from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE
36-
from ..evaluation.gcs_eval_set_results_manager import GcsEvalSetResultsManager
37-
from ..evaluation.gcs_eval_sets_manager import GcsEvalSetsManager
3836
from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
3937
from ..sessions.in_memory_session_service import InMemorySessionService
4038
from .cli import run_cli
@@ -147,13 +145,26 @@ def deploy():
147145
type=str,
148146
help="Optional. The Google Cloud Region for using VertexAI as backend.",
149147
)
148+
@click.option(
149+
"--type",
150+
type=click.Choice([t.value for t in cli_create.Type]),
151+
help=(
152+
"EXPERIMENTAL Optional. Type of agent to create: 'config' or 'code'."
153+
" 'config' is not ready for use so it defaults to 'code'. It may change"
154+
" later once 'config' is ready for use."
155+
),
156+
default=cli_create.Type.CODE.value,
157+
show_default=True,
158+
hidden=True, # Won't show in --help output. Not ready for use.
159+
)
150160
@click.argument("app_name", type=str, required=True)
151161
def cli_create_cmd(
152162
app_name: str,
153163
model: Optional[str],
154164
api_key: Optional[str],
155165
project: Optional[str],
156166
region: Optional[str],
167+
type: Optional[cli_create.Type],
157168
):
158169
"""Creates a new app in the current folder with prepopulated agent template.
159170
@@ -169,6 +180,7 @@ def cli_create_cmd(
169180
google_api_key=api_key,
170181
google_cloud_project=project,
171182
google_cloud_region=region,
183+
type=type,
172184
)
173185

174186

tests/unittests/cli/utils/test_cli_create.py

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -147,9 +147,53 @@ def test_run_cmd_overwrite_reject(
147147
google_api_key=None,
148148
google_cloud_project=None,
149149
google_cloud_region=None,
150+
type=cli_create.Type.CODE,
150151
)
151152

152153

154+
def test_run_cmd_with_type_config(
155+
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
156+
) -> None:
157+
"""run_cmd with --type=config should generate YAML config file."""
158+
agent_name = "test_agent"
159+
160+
monkeypatch.setattr(os, "getcwd", lambda: str(tmp_path))
161+
monkeypatch.setattr(os.path, "exists", lambda _p: False)
162+
163+
cli_create.run_cmd(
164+
agent_name,
165+
model="gemini-2.0-flash-001",
166+
google_api_key="test-key",
167+
google_cloud_project=None,
168+
google_cloud_region=None,
169+
type=cli_create.Type.CONFIG,
170+
)
171+
172+
agent_dir = tmp_path / agent_name
173+
assert agent_dir.exists()
174+
175+
# Should create root_agent.yaml instead of agent.py
176+
yaml_file = agent_dir / "root_agent.yaml"
177+
assert yaml_file.exists()
178+
assert not (agent_dir / "agent.py").exists()
179+
180+
# Check YAML content
181+
yaml_content = yaml_file.read_text()
182+
assert "name: root_agent" in yaml_content
183+
assert "model: gemini-2.0-flash-001" in yaml_content
184+
assert "description: A helpful assistant for user questions." in yaml_content
185+
186+
# Should create empty __init__.py
187+
init_file = agent_dir / "__init__.py"
188+
assert init_file.exists()
189+
assert init_file.read_text().strip() == ""
190+
191+
# Should still create .env file
192+
env_file = agent_dir / ".env"
193+
assert env_file.exists()
194+
assert "GOOGLE_API_KEY=test-key" in env_file.read_text()
195+
196+
153197
# Prompt helpers
154198
def test_prompt_for_google_cloud(monkeypatch: pytest.MonkeyPatch) -> None:
155199
"""Prompt should return the project input."""
@@ -174,7 +218,7 @@ def test_prompt_for_google_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
174218
def test_prompt_for_model_gemini(monkeypatch: pytest.MonkeyPatch) -> None:
175219
"""Selecting option '1' should return the default Gemini model string."""
176220
monkeypatch.setattr(click, "prompt", lambda *a, **k: "1")
177-
assert cli_create._prompt_for_model() == "gemini-2.0-flash-001"
221+
assert cli_create._prompt_for_model() == "gemini-2.5-flash"
178222

179223

180224
def test_prompt_for_model_other(monkeypatch: pytest.MonkeyPatch) -> None:

0 commit comments

Comments
 (0)