Skip to content

Commit 21f63f6

Browse files
wikaaaaacopybara-github
authored andcommitted
feat(cli): add otel_to_cloud flag to adk deploy agent_engine command
Co-authored-by: Wiktoria Walczak <wwalczak@google.com> PiperOrigin-RevId: 858546581
1 parent 69ad605 commit 21f63f6

File tree

3 files changed

+53
-0
lines changed

3 files changed

+53
-0
lines changed

src/google/adk/cli/cli_deploy.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -689,6 +689,7 @@ def to_agent_engine(
689689
adk_app: str,
690690
staging_bucket: Optional[str] = None,
691691
trace_to_cloud: Optional[bool] = None,
692+
otel_to_cloud: Optional[bool] = None,
692693
api_key: Optional[str] = None,
693694
adk_app_object: Optional[str] = None,
694695
agent_engine_id: Optional[str] = None,
@@ -733,6 +734,8 @@ def to_agent_engine(
733734
staging_bucket (str): Deprecated. This argument is no longer required or
734735
used.
735736
trace_to_cloud (bool): Whether to enable Cloud Trace.
737+
otel_to_cloud (bool): Whether to enable exporting OpenTelemetry signals
738+
to Google Cloud.
736739
api_key (str): Optional. The API key to use for Express Mode.
737740
If not provided, the API key from the GOOGLE_API_KEY environment variable
738741
will be used. It will only be used if GOOGLE_GENAI_USE_VERTEXAI is true.
@@ -910,6 +913,14 @@ def to_agent_engine(
910913
if 'GOOGLE_API_KEY' in env_vars:
911914
api_key = env_vars['GOOGLE_API_KEY']
912915
click.echo(f'api_key set by GOOGLE_API_KEY in {env_file}')
916+
if otel_to_cloud:
917+
if 'GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY' in env_vars:
918+
click.secho(
919+
'Ignoring GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in .env'
920+
' as `--otel_to_cloud` was explicitly passed and takes precedence',
921+
fg='yellow',
922+
)
923+
env_vars['GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY'] = 'true'
913924
if env_vars:
914925
if 'env_vars' in agent_config:
915926
click.echo(

src/google/adk/cli/cli_tools_click.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1751,6 +1751,14 @@ def cli_migrate_session(
17511751
default=None,
17521752
help="Optional. Whether to enable Cloud Trace for Agent Engine.",
17531753
)
1754+
@click.option(
1755+
"--otel_to_cloud",
1756+
type=bool,
1757+
is_flag=True,
1758+
show_default=True,
1759+
default=None,
1760+
help="Optional. Whether to enable OpenTelemetry for Agent Engine.",
1761+
)
17541762
@click.option(
17551763
"--display_name",
17561764
type=str,
@@ -1842,6 +1850,7 @@ def cli_deploy_agent_engine(
18421850
staging_bucket: Optional[str],
18431851
agent_engine_id: Optional[str],
18441852
trace_to_cloud: Optional[bool],
1853+
otel_to_cloud: Optional[bool],
18451854
api_key: Optional[str],
18461855
display_name: str,
18471856
description: str,
@@ -1872,6 +1881,7 @@ def cli_deploy_agent_engine(
18721881
region=region,
18731882
agent_engine_id=agent_engine_id,
18741883
trace_to_cloud=trace_to_cloud,
1884+
otel_to_cloud=otel_to_cloud,
18751885
api_key=api_key,
18761886
adk_app_object=adk_app_object,
18771887
display_name=display_name,

tests/unittests/cli/utils/test_cli_tools_click.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -410,6 +410,38 @@ def test_cli_deploy_agent_engine_success(
410410
assert called_kwargs.get("region") == "us-central1"
411411

412412

413+
# cli deploy agent_engine with --otel_to_cloud
414+
def test_cli_deploy_agent_engine_otel_to_cloud_success(
415+
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
416+
) -> None:
417+
"""Successful path should call cli_deploy.to_agent_engine with --otel_to_cloud."""
418+
rec = _Recorder()
419+
monkeypatch.setattr(cli_tools_click.cli_deploy, "to_agent_engine", rec)
420+
421+
agent_dir = tmp_path / "agent_ae"
422+
agent_dir.mkdir()
423+
runner = CliRunner()
424+
result = runner.invoke(
425+
cli_tools_click.main,
426+
[
427+
"deploy",
428+
"agent_engine",
429+
"--project",
430+
"test-proj",
431+
"--region",
432+
"us-central1",
433+
"--otel_to_cloud",
434+
str(agent_dir),
435+
],
436+
)
437+
assert result.exit_code == 0
438+
assert rec.calls, "cli_deploy.to_agent_engine must be invoked"
439+
called_kwargs = rec.calls[0][1]
440+
assert called_kwargs.get("project") == "test-proj"
441+
assert called_kwargs.get("region") == "us-central1"
442+
assert called_kwargs.get("otel_to_cloud")
443+
444+
413445
# cli deploy gke
414446
def test_cli_deploy_gke_success(
415447
tmp_path: Path, monkeypatch: pytest.MonkeyPatch

0 commit comments

Comments
 (0)