14
14
15
15
from __future__ import annotations
16
16
17
+ import enum
17
18
import os
18
19
import subprocess
19
20
from typing import Optional
20
21
from typing import Tuple
21
22
22
23
import click
23
24
25
+
26
+ class Type (enum .Enum ):
27
+ CONFIG = "config"
28
+ CODE = "code"
29
+
30
+
24
31
_INIT_PY_TEMPLATE = """\
25
32
from . import agent
26
33
"""
36
43
)
37
44
"""
38
45
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
+
39
53
40
54
_GOOGLE_API_MSG = """
41
55
Don't have API Key? Create one in AI Studio: https://aistudio.google.com/apikey
51
65
https://google.github.io/adk-docs/agents/models
52
66
"""
53
67
54
- _SUCCESS_MSG = """
68
+ _SUCCESS_MSG_CODE = """
55
69
Agent created in {agent_folder}:
56
70
- .env
57
71
- __init__.py
58
72
- agent.py
59
73
"""
60
74
75
+ _SUCCESS_MSG_CONFIG = """
76
+ Agent created in {agent_folder}:
77
+ - .env
78
+ - __init__.py
79
+ - root_agent.yaml
80
+ """
81
+
61
82
62
83
def _get_gcp_project_from_gcloud () -> str :
63
84
"""Uses gcloud to get default project."""
@@ -158,13 +179,15 @@ def _generate_files(
158
179
google_cloud_project : Optional [str ] = None ,
159
180
google_cloud_region : Optional [str ] = None ,
160
181
model : Optional [str ] = None ,
182
+ type : Optional [Type ] = None ,
161
183
):
162
184
"""Generates a folder name for the agent."""
163
185
os .makedirs (agent_folder , exist_ok = True )
164
186
165
187
dotenv_file_path = os .path .join (agent_folder , ".env" )
166
188
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" )
168
191
169
192
with open (dotenv_file_path , "w" , encoding = "utf-8" ) as f :
170
193
lines = []
@@ -180,29 +203,38 @@ def _generate_files(
180
203
lines .append (f"GOOGLE_CLOUD_LOCATION={ google_cloud_region } " )
181
204
f .write ("\n " .join (lines ))
182
205
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
+ )
193
225
194
226
195
227
def _prompt_for_model () -> str :
196
228
model_choice = click .prompt (
197
229
"""\
198
230
Choose a model for the root agent:
199
- 1. gemini-2.0 -flash-001
231
+ 1. gemini-2.5 -flash
200
232
2. Other models (fill later)
201
233
Choose model""" ,
202
234
type = click .Choice (["1" , "2" ]),
203
235
)
204
236
if model_choice == "1" :
205
- return "gemini-2.0 -flash-001 "
237
+ return "gemini-2.5 -flash"
206
238
else :
207
239
click .secho (_OTHER_MODEL_MSG , fg = "green" )
208
240
return "<FILL_IN_MODEL>"
@@ -231,13 +263,30 @@ def _prompt_to_choose_backend(
231
263
return google_api_key , google_cloud_project , google_cloud_region
232
264
233
265
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
+
234
282
def run_cmd (
235
283
agent_name : str ,
236
284
* ,
237
285
model : Optional [str ],
238
286
google_api_key : Optional [str ],
239
287
google_cloud_project : Optional [str ],
240
288
google_cloud_region : Optional [str ],
289
+ type : Optional [Type ],
241
290
):
242
291
"""Runs `adk create` command to create agent template.
243
292
@@ -249,6 +298,7 @@ def run_cmd(
249
298
VertexAI as backend.
250
299
google_cloud_region: Optional[str], The Google Cloud region for using
251
300
VertexAI as backend.
301
+ type: Optional[Type], Whether to define agent with config file or code.
252
302
"""
253
303
agent_folder = os .path .join (os .getcwd (), agent_name )
254
304
# check folder doesn't exist or it's empty. Otherwise, throw
@@ -272,10 +322,14 @@ def run_cmd(
272
322
)
273
323
)
274
324
325
+ if not type :
326
+ type = _prompt_to_choose_type ()
327
+
275
328
_generate_files (
276
329
agent_folder ,
277
330
google_api_key = google_api_key ,
278
331
google_cloud_project = google_cloud_project ,
279
332
google_cloud_region = google_cloud_region ,
280
333
model = model ,
334
+ type = type ,
281
335
)
0 commit comments