From d776f22c8e3e2ee570c8ac84d83ecdca3c77f6e9 Mon Sep 17 00:00:00 2001 From: George Weale Date: Thu, 30 Jul 2026 15:14:39 -0700 Subject: [PATCH] refactor: define StreamingMode in a leaf module the CLI can import Co-authored-by: George Weale PiperOrigin-RevId: 956760523 --- src/google/adk/agents/_streaming_mode.py | 147 ++++++++++++++++++ src/google/adk/agents/run_config.py | 132 +--------------- src/google/adk/cli/cli_tools_click.py | 8 +- .../cli/utils/test_cli_tools_click.py | 10 -- 4 files changed, 151 insertions(+), 146 deletions(-) create mode 100644 src/google/adk/agents/_streaming_mode.py diff --git a/src/google/adk/agents/_streaming_mode.py b/src/google/adk/agents/_streaming_mode.py new file mode 100644 index 00000000..2fc032d1 --- /dev/null +++ b/src/google/adk/agents/_streaming_mode.py @@ -0,0 +1,147 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from enum import Enum + + +class StreamingMode(Enum): + """Streaming modes for agent execution. + + This enum defines different streaming behaviors for how the agent returns + events as model response. + """ + + NONE = None + """Non-streaming mode (default). + + In this mode: + - The runner returns one single content in a turn (one user / model + interaction). + - No partial/intermediate events are produced + - Suitable for: CLI tools, batch processing, synchronous workflows + + Example: + ```python + config = RunConfig(streaming_mode=StreamingMode.NONE) + async for event in runner.run_async(..., run_config=config): + # event.partial is always False + # Only final responses are yielded + if event.content: + print(event.content.parts[0].text) + ``` + """ + + SSE = 'sse' + """Server-Sent Events (SSE) streaming mode. + + In this mode: + - The runner yields events progressively as the LLM generates responses + - Both partial events (streaming chunks) and aggregated events are yielded + - Suitable for: real-time display with typewriter effects in Web UIs, chat + applications, interactive displays + + Event Types in SSE Mode: + - **Partial text events** (event.partial=True, contains text): + Streaming text chunks for typewriter effect. These should typically be + displayed to users in real-time. + + - **Partial function call events** (event.partial=True, contains function_call): + Internal streaming chunks used to progressively build function call + arguments. These are typically NOT displayed to end users. + + - **Aggregated events** (event.partial=False): + The complete, aggregated response after all streaming chunks. Contains + the full text or complete function call with all arguments. + + Important Considerations: + 1. **Duplicate text issue**: With Progressive SSE Streaming enabled + (default), you will receive both partial text chunks AND a final + aggregated text event. To avoid displaying text twice: + - Option A: Only display partial text events, skip final text events + - Option B: Only display final events, skip all partial events + - Option C: Track what's been displayed and skip duplicates + + 2. **Event filtering**: Applications should filter events based on their + needs. Common patterns: + + # Pattern 1: Display only partial text + final function calls + async for event in runner.run_async(...): + if event.partial and event.content and event.content.parts: + # Check if it's text (not function call) + if any(part.text for part in event.content.parts): + if not any(part.function_call for part in event.content.parts): + # Display partial text for typewriter effect + text = ''.join(p.text or '' for p in event.content.parts) + print(text, end='', flush=True) + elif not event.partial and event.get_function_calls(): + # Display final function calls + for fc in event.get_function_calls(): + print(f"Calling {fc.name}({fc.args})") + + # Pattern 2: Display only final events (no streaming effect) + async for event in runner.run_async(...): + if not event.partial: + # Only process final responses + if event.content: + text = ''.join(p.text or '' for p in event.content.parts) + print(text) + + 3. **Progressive SSE Streaming feature**: Controlled by the + ADK_ENABLE_PROGRESSIVE_SSE_STREAMING environment variable (default: ON). + - When ON: Preserves original part ordering, supports function call + argument streaming, produces partial events + final aggregated event + - When OFF: Simple text accumulation, may lose some information + + Example: + ```python + config = RunConfig(streaming_mode=StreamingMode.SSE) + displayed_text = "" + + async for event in runner.run_async(..., run_config=config): + if event.partial: + # Partial streaming event + if event.content and event.content.parts: + # Check if this is text (not a function call) + has_text = any(part.text for part in event.content.parts) + has_fc = any(part.function_call for part in event.content.parts) + + if has_text and not has_fc: + # Display partial text chunks for typewriter effect + text = ''.join(p.text or '' for p in event.content.parts) + print(text, end='', flush=True) + displayed_text += text + else: + # Final event - check if we already displayed this content + if event.content: + final_text = ''.join(p.text or '' for p in event.content.parts) + if final_text != displayed_text: + # New content not yet displayed + print(final_text) + ``` + + See Also: + - Event.is_final_response() for identifying final responses + """ + + BIDI = 'bidi' + """Bidirectional streaming mode. + + So far this mode is not used in the standard execution path. The actual + bidirectional streaming behavior via runner.run_live() uses a completely + different code path that doesn't rely on streaming_mode. + + For bidirectional streaming, use runner.run_live() instead of run_async(). + """ diff --git a/src/google/adk/agents/run_config.py b/src/google/adk/agents/run_config.py index 2b4f1bd3..0d4133a3 100644 --- a/src/google/adk/agents/run_config.py +++ b/src/google/adk/agents/run_config.py @@ -14,7 +14,6 @@ from __future__ import annotations -from enum import Enum import logging import sys from typing import Any @@ -30,6 +29,7 @@ from pydantic import model_validator from ..sessions.base_session_service import GetSessionConfig from ..telemetry.context import TelemetryConfig +from ._streaming_mode import StreamingMode logger = logging.getLogger('google_adk.' + __name__) @@ -52,136 +52,6 @@ class ToolThreadPoolConfig(BaseModel): ) -class StreamingMode(Enum): - """Streaming modes for agent execution. - - This enum defines different streaming behaviors for how the agent returns - events as model response. - """ - - NONE = None - """Non-streaming mode (default). - - In this mode: - - The runner returns one single content in a turn (one user / model - interaction). - - No partial/intermediate events are produced - - Suitable for: CLI tools, batch processing, synchronous workflows - - Example: - ```python - config = RunConfig(streaming_mode=StreamingMode.NONE) - async for event in runner.run_async(..., run_config=config): - # event.partial is always False - # Only final responses are yielded - if event.content: - print(event.content.parts[0].text) - ``` - """ - - SSE = 'sse' - """Server-Sent Events (SSE) streaming mode. - - In this mode: - - The runner yields events progressively as the LLM generates responses - - Both partial events (streaming chunks) and aggregated events are yielded - - Suitable for: real-time display with typewriter effects in Web UIs, chat - applications, interactive displays - - Event Types in SSE Mode: - - **Partial text events** (event.partial=True, contains text): - Streaming text chunks for typewriter effect. These should typically be - displayed to users in real-time. - - - **Partial function call events** (event.partial=True, contains function_call): - Internal streaming chunks used to progressively build function call - arguments. These are typically NOT displayed to end users. - - - **Aggregated events** (event.partial=False): - The complete, aggregated response after all streaming chunks. Contains - the full text or complete function call with all arguments. - - Important Considerations: - 1. **Duplicate text issue**: With Progressive SSE Streaming enabled - (default), you will receive both partial text chunks AND a final - aggregated text event. To avoid displaying text twice: - - Option A: Only display partial text events, skip final text events - - Option B: Only display final events, skip all partial events - - Option C: Track what's been displayed and skip duplicates - - 2. **Event filtering**: Applications should filter events based on their - needs. Common patterns: - - # Pattern 1: Display only partial text + final function calls - async for event in runner.run_async(...): - if event.partial and event.content and event.content.parts: - # Check if it's text (not function call) - if any(part.text for part in event.content.parts): - if not any(part.function_call for part in event.content.parts): - # Display partial text for typewriter effect - text = ''.join(p.text or '' for p in event.content.parts) - print(text, end='', flush=True) - elif not event.partial and event.get_function_calls(): - # Display final function calls - for fc in event.get_function_calls(): - print(f"Calling {fc.name}({fc.args})") - - # Pattern 2: Display only final events (no streaming effect) - async for event in runner.run_async(...): - if not event.partial: - # Only process final responses - if event.content: - text = ''.join(p.text or '' for p in event.content.parts) - print(text) - - 3. **Progressive SSE Streaming feature**: Controlled by the - ADK_ENABLE_PROGRESSIVE_SSE_STREAMING environment variable (default: ON). - - When ON: Preserves original part ordering, supports function call - argument streaming, produces partial events + final aggregated event - - When OFF: Simple text accumulation, may lose some information - - Example: - ```python - config = RunConfig(streaming_mode=StreamingMode.SSE) - displayed_text = "" - - async for event in runner.run_async(..., run_config=config): - if event.partial: - # Partial streaming event - if event.content and event.content.parts: - # Check if this is text (not a function call) - has_text = any(part.text for part in event.content.parts) - has_fc = any(part.function_call for part in event.content.parts) - - if has_text and not has_fc: - # Display partial text chunks for typewriter effect - text = ''.join(p.text or '' for p in event.content.parts) - print(text, end='', flush=True) - displayed_text += text - else: - # Final event - check if we already displayed this content - if event.content: - final_text = ''.join(p.text or '' for p in event.content.parts) - if final_text != displayed_text: - # New content not yet displayed - print(final_text) - ``` - - See Also: - - Event.is_final_response() for identifying final responses - """ - - BIDI = 'bidi' - """Bidirectional streaming mode. - - So far this mode is not used in the standard execution path. The actual - bidirectional streaming behavior via runner.run_live() uses a completely - different code path that doesn't rely on streaming_mode. - - For bidirectional streaming, use runner.run_live() instead of run_async(). - """ - - class RunConfig(BaseModel): """Configs for runtime behavior of agents. diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index b2c070bb..d0b863d5 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -37,6 +37,7 @@ import click from click.core import ParameterSource from .. import version +from ..agents._streaming_mode import StreamingMode from ..features import FeatureName from ..features import override_feature_enabled from ..utils._telemetry_config import read_telemetry_consent @@ -49,7 +50,6 @@ if TYPE_CHECKING: from fastapi import FastAPI from ..agents.llm_agent import LlmAgent - from ..agents.run_config import StreamingMode LOG_LEVELS = click.Choice( @@ -57,7 +57,7 @@ LOG_LEVELS = click.Choice( case_sensitive=False, ) -_STREAMING_MODE_CHOICES = ("None", "sse", "bidi") +_STREAMING_MODE_CHOICES = tuple(str(mode.value) for mode in StreamingMode) def _missing_eval_dependencies_message() -> str: @@ -72,12 +72,10 @@ def _parse_streaming_mode( param: click.Parameter, value: str | None, ) -> StreamingMode | None: - """Converts a validated CLI value without importing the runtime for help.""" + """Converts a validated CLI value to its streaming mode.""" if value is None: return None - from ..agents.run_config import StreamingMode - mode = next( (m for m in StreamingMode if str(m.value).lower() == value.lower()), None ) diff --git a/tests/unittests/cli/utils/test_cli_tools_click.py b/tests/unittests/cli/utils/test_cli_tools_click.py index 4dc56da4..d1589240 100644 --- a/tests/unittests/cli/utils/test_cli_tools_click.py +++ b/tests/unittests/cli/utils/test_cli_tools_click.py @@ -90,16 +90,6 @@ def _mute_click(request, monkeypatch: pytest.MonkeyPatch) -> None: # monkeypatch.setattr(click, "secho", lambda *a, **k: None) -# streaming mode choices -def test_streaming_mode_choices_match_enum() -> None: - """The CLI choices are hardcoded to defer the runtime import; pin them.""" - from google.adk.agents.run_config import StreamingMode - - assert set(cli_tools_click._STREAMING_MODE_CHOICES) == { - str(mode.value) for mode in StreamingMode - } - - # validate_exclusive def test_validate_exclusive_allows_single() -> None: """Providing exactly one exclusive option should pass."""