From 024a438920293e9de5f95b58c0669f858d39fa6f Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Sat, 22 Aug 2026 20:21:36 +0100 Subject: [PATCH] feat(cli): add max LLM calls option --- src/google/adk/cli/api_server.py | 23 ++++++++++--- src/google/adk/cli/cli_tools_click.py | 13 ++++++++ src/google/adk/cli/fast_api.py | 4 +++ .../cli/test_adk_web_server_run_live.py | 4 ++- .../cli/utils/test_cli_tools_click.py | 33 +++++++++++++++++++ 5 files changed, 71 insertions(+), 6 deletions(-) diff --git a/src/google/adk/cli/api_server.py b/src/google/adk/cli/api_server.py index 8285af05a9..542dcc96b6 100644 --- a/src/google/adk/cli/api_server.py +++ b/src/google/adk/cli/api_server.py @@ -839,6 +839,7 @@ def __init__( auto_create_session: bool = False, trigger_sources: Optional[list[str]] = None, default_llm_model: Optional[str] = None, + max_llm_calls: Optional[int] = None, ): self.agent_loader = agent_loader self.session_service = session_service @@ -859,6 +860,7 @@ def __init__( self.auto_create_session = auto_create_session self.trigger_sources = trigger_sources self.default_llm_model = default_llm_model + self.max_llm_calls = max_llm_calls self.default_app_name = os.getenv("ADK_DEFAULT_APP_NAME") async def get_runner_async(self, app_name: str) -> Runner: @@ -1807,11 +1809,12 @@ async def run_agent(req: RunAgentRequest, request: Request) -> list[Event]: self.current_app_name_ref.value = req.app_name runner = await self.get_runner_async(req.app_name) _set_telemetry_context_if_needed(runner) - run_config = ( - RunConfig(custom_metadata=req.custom_metadata) - if req.custom_metadata - else None - ) + run_config_kwargs: dict[str, Any] = {} + if req.custom_metadata: + run_config_kwargs["custom_metadata"] = req.custom_metadata + if self.max_llm_calls is not None: + run_config_kwargs["max_llm_calls"] = self.max_llm_calls + run_config = RunConfig(**run_config_kwargs) if run_config_kwargs else None async def worker(): try: @@ -1907,6 +1910,11 @@ async def event_generator(): run_config=RunConfig( streaming_mode=stream_mode, custom_metadata=req.custom_metadata, + **( + {"max_llm_calls": self.max_llm_calls} + if self.max_llm_calls is not None + else {} + ), ), invocation_id=req.invocation_id, ) @@ -2061,6 +2069,11 @@ async def forward_events(): ), save_live_blob=save_live_blob, explicit_vad_signal=explicit_vad_signal, + **( + {"max_llm_calls": self.max_llm_calls} + if self.max_llm_calls is not None + else {} + ), ) async with Aclosing( runner.run_live( diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index 49fb819388..1abebe216e 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -1927,6 +1927,15 @@ def decorator(func): ), default=None, ) + @click.option( + "--max_llm_calls", + type=int, + help=( + "Optional. Maximum number of LLM calls allowed for each agent" + " run. Values less than or equal to zero disable the limit." + ), + default=None, + ) # Parsed into list[str] by the wrapper below (server commands need a list). @click.option( "--trigger_sources", @@ -2012,6 +2021,7 @@ def cli_web( logo_text: str | None = None, logo_image_url: str | None = None, trigger_sources: list[str] | None = None, + max_llm_calls: int | None = None, ): """Starts a FastAPI server with Web UI for agents. @@ -2082,6 +2092,7 @@ async def _lifespan(app: FastAPI): logo_image_url=logo_image_url, trigger_sources=trigger_sources, default_llm_model=default_llm_model, + max_llm_calls=max_llm_calls, ) config = uvicorn.Config( app, @@ -2163,6 +2174,7 @@ def cli_api_server( with_ui: bool = False, gemini_enterprise_app_name: str | None = None, express_mode: bool = False, + max_llm_calls: int | None = None, ): """Starts a FastAPI server for agents. @@ -2223,6 +2235,7 @@ async def _lifespan(app: FastAPI) -> AsyncIterator[None]: trigger_sources=trigger_sources, gemini_enterprise_app_name=gemini_enterprise_app_name, express_mode=express_mode, + max_llm_calls=max_llm_calls, lifespan=_lifespan, ), host=host, diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index 3e437a7772..51c1d00654 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -122,6 +122,7 @@ def get_fast_api_app( default_llm_model: str | None = None, gemini_enterprise_app_name: str | None = None, express_mode: bool = False, + max_llm_calls: int | None = None, ) -> FastAPI: """Constructs and returns a FastAPI application for serving ADK agents. @@ -178,6 +179,8 @@ def get_fast_api_app( gemini_enterprise_app_name: The Gemini Enterprise app name to use for the agent. express_mode: Whether to enable express mode. + max_llm_calls: Maximum number of LLM calls allowed for each agent run. + When None, ``RunConfig`` resolves its normal default. Returns: The configured FastAPI application instance. @@ -300,6 +303,7 @@ def get_fast_api_app( auto_create_session=auto_create_session, trigger_sources=trigger_sources, default_llm_model=default_llm_model, + max_llm_calls=max_llm_calls, ) # In single agent mode, use that agent as the default app. diff --git a/tests/unittests/cli/test_adk_web_server_run_live.py b/tests/unittests/cli/test_adk_web_server_run_live.py index 316d6c1f21..b685acb1fd 100644 --- a/tests/unittests/cli/test_adk_web_server_run_live.py +++ b/tests/unittests/cli/test_adk_web_server_run_live.py @@ -61,7 +61,7 @@ async def run_live( yield Event(author="runner") -def test_run_live_applies_run_config_query_options(): +def test_run_live_applies_server_and_query_run_config_options(): session_service = InMemorySessionService() asyncio.run( session_service.create_session( @@ -82,6 +82,7 @@ def test_run_live_applies_run_config_query_options(): eval_sets_manager=types.SimpleNamespace(), eval_set_results_manager=types.SimpleNamespace(), agents_dir=".", + max_llm_calls=37, ) async def _get_runner_async(_self, _app_name: str): @@ -122,6 +123,7 @@ async def _get_runner_async(_self, _app_name: str): assert run_config.session_resumption.transparent is True assert run_config.save_live_blob is True assert run_config.explicit_vad_signal is True + assert run_config.max_llm_calls == 37 @pytest.mark.parametrize( diff --git a/tests/unittests/cli/utils/test_cli_tools_click.py b/tests/unittests/cli/utils/test_cli_tools_click.py index dddda4471f..f6a4845e15 100644 --- a/tests/unittests/cli/utils/test_cli_tools_click.py +++ b/tests/unittests/cli/utils/test_cli_tools_click.py @@ -1589,6 +1589,28 @@ def test_cli_api_server_invokes_uvicorn( assert _patch_uvicorn.calls, "uvicorn.Server.run must be called" +@pytest.mark.parametrize("command", ["web", "api_server"]) +def test_cli_server_passes_max_llm_calls( + tmp_path: Path, + _patch_uvicorn: _Recorder, + monkeypatch: pytest.MonkeyPatch, + command: str, +) -> None: + """Both server commands pass the requested LLM call limit to the app.""" + agents_dir = tmp_path / "agents" + agents_dir.mkdir() + mock_get_app = _Recorder() + monkeypatch.setattr("google.adk.cli.fast_api.get_fast_api_app", mock_get_app) + + result = CliRunner().invoke( + cli_tools_click.main, + [command, "--max_llm_calls", "37", str(agents_dir)], + ) + + assert result.exit_code == 0, (result.output, repr(result.exception)) + assert mock_get_app.calls[0][1]["max_llm_calls"] == 37 + + def test_cli_web_passes_service_uris( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, _patch_uvicorn: _Recorder ) -> None: @@ -2632,10 +2654,21 @@ def test_fast_api_common_options_documented_defaults() -> None: assert captured["a2a"] is False assert captured["allow_origins"] == () assert captured["log_level"] == "INFO" + assert captured["max_llm_calls"] is None # --verbose is consumed while folding it into log_level. assert "verbose" not in captured +def test_fast_api_common_options_parses_max_llm_calls() -> None: + """The common server option parses an explicit per-run LLM call limit.""" + command, captured = _fast_api_command() + + result = CliRunner().invoke(command, ["--max_llm_calls", "37"]) + + assert result.exit_code == 0, (result.output, repr(result.exception)) + assert captured["max_llm_calls"] == 37 + + # adk test @pytest.fixture def fake_pytest_run(monkeypatch: pytest.MonkeyPatch):