George Weale
a308848969
Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 962282749
109 lines
3.1 KiB
Python
109 lines
3.1 KiB
Python
# 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.
|
|
|
|
import asyncio
|
|
import base64
|
|
import os
|
|
from typing import Any
|
|
|
|
from google.adk import Event
|
|
from google.adk import Workflow
|
|
from google.genai import types
|
|
|
|
|
|
async def sleep_if_not_pytest(seconds: float):
|
|
if "PYTEST_CURRENT_TEST" not in os.environ:
|
|
await asyncio.sleep(seconds)
|
|
|
|
|
|
def send_string(node_input: Any = None):
|
|
"""Sends a single string message."""
|
|
yield Event(message="#1 This is a simple string message.")
|
|
|
|
|
|
def send_multimodal(node_input: Any = None):
|
|
"""Sends a multi-modal message containing a string and an inline image."""
|
|
# A 16x16 solid red PNG base64 encoded
|
|
red_square_png = base64.b64decode(
|
|
"iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAXElEQVR4nO2TSQ7AIAwD"
|
|
"7fz/z+ZQtapwmrJc8QklmjBIgZJgIZMiAIl9KYbhjx4fgwosbNxgMrF0+4uhgHnYDM6"
|
|
"AzQHJeg5HYtyHFfgy2AztN/5tZWfrBtVzkl4DzfQkEPd+cEkAAAAASUVORK5CYII="
|
|
)
|
|
yield Event(
|
|
message=[
|
|
types.Part.from_text(
|
|
text=(
|
|
"#2 Here is a multi-modal message with an inline image (red"
|
|
" circle):"
|
|
)
|
|
),
|
|
types.Part.from_bytes(data=red_square_png, mime_type="image/png"),
|
|
]
|
|
)
|
|
|
|
|
|
async def multiple_messages(node_input: Any = None):
|
|
"""Sends multiple complete messages from the same node with an interval."""
|
|
yield Event(message="#3 Multiple messages")
|
|
await sleep_if_not_pytest(1.0)
|
|
|
|
yield Event(message="Processing step 1...")
|
|
await sleep_if_not_pytest(1.0)
|
|
|
|
yield Event(message="Processing step 2...")
|
|
await sleep_if_not_pytest(1.0)
|
|
|
|
yield Event(message="Done processing.")
|
|
|
|
|
|
async def stream_sentence(node_input: Any = None):
|
|
"""
|
|
Demonstrates streaming by sending a sentence in chunks.
|
|
The `partial=True` flag tells the UI that this is part of an ongoing message.
|
|
Partial events are not written to the session, so the node ends by yielding
|
|
the assembled sentence once as a non-partial event.
|
|
"""
|
|
yield Event(message="#4 Starting to stream...")
|
|
sentence = """\
|
|
This is a streaming message sent in chunks.
|
|
|
|
You can stream in markdown as well. For example, the table below:
|
|
|
|
| Header 1 | Header 2 |
|
|
|----------|----------|
|
|
| Cell 1 | Cell 2 |
|
|
| Cell 3 | Cell 4 |
|
|
"""
|
|
|
|
for i in range(0, len(sentence), 5):
|
|
chunk = sentence[i : i + 5]
|
|
yield Event(message=chunk, partial=True)
|
|
await sleep_if_not_pytest(0.2)
|
|
|
|
yield Event(message=sentence)
|
|
|
|
|
|
root_agent = Workflow(
|
|
name="message",
|
|
edges=[
|
|
(
|
|
"START",
|
|
send_string,
|
|
send_multimodal,
|
|
multiple_messages,
|
|
stream_sentence,
|
|
),
|
|
],
|
|
)
|