mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
2f73cfde18
commit
bcac9ba44c
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user