Skip to content

Commit 1b1f8fc

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
docs: Update AdkApp.async_stream_query docstring for multimodal message dicts and add unit tests.
PiperOrigin-RevId: 1019370555
1 parent 74d17cc commit 1b1f8fc

6 files changed

Lines changed: 220 additions & 8 deletions

File tree

‎agentplatform/frameworks/adk.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1336,8 +1336,12 @@ async def async_stream_query(
13361336
"""Streams responses asynchronously from the ADK application.
13371337
13381338
Args:
1339-
message (str):
1340-
Required. The message to stream responses for.
1339+
message (Union[str, Dict[str, Any]]):
1340+
Required. The message to stream responses for. Either a text
1341+
string or a dictionary representing a
1342+
`google.genai.types.Content` object (for example,
1343+
`content.model_dump(exclude_none=True, mode="json")` for
1344+
multimodal inputs).
13411345
user_id (str):
13421346
Required. The ID of the user.
13431347
session_id (str):
@@ -1362,7 +1366,7 @@ async def async_stream_query(
13621366
13631367
Raises:
13641368
TypeError: If message is not a string or a dictionary representing
1365-
a Content object.
1369+
a Content object.
13661370
ValueError: If both session_id and session_events are specified.
13671371
"""
13681372
from agentplatform._genai import _runtimes_utils

‎tests/unit/agentplatform/frameworks/test_frameworks_adk.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -770,6 +770,64 @@ async def test_async_stream_query_with_content(
770770
events.append(event)
771771
assert len(events) == 1
772772

773+
@pytest.mark.asyncio
774+
@pytest.mark.usefixtures(
775+
"default_instrumentor_builder_mock", "get_project_id_mock"
776+
)
777+
async def test_async_stream_query_with_multimodal_content(self):
778+
"""Tests async_stream_query with a multimodal Content dictionary."""
779+
# pylint: disable=protected-access
780+
app = adk_template.AdkApp(agent=_TEST_AGENT)
781+
assert app._tmpl_attrs.get("runner") is None
782+
app.set_up()
783+
runner_kwargs = {}
784+
785+
async def run_async(**kwargs):
786+
runner_kwargs.update(kwargs)
787+
async for event in _MockRunner().run_async():
788+
yield event
789+
790+
app._tmpl_attrs["runner"] = mock.Mock(run_async=run_async)
791+
video_message = {
792+
"role": "user",
793+
"parts": [
794+
{
795+
"file_data": {
796+
"file_uri": (
797+
"gs://cloud-samples-data/generative-ai/video/pixel8.mp4"
798+
),
799+
"mime_type": "video/mp4",
800+
},
801+
},
802+
{
803+
"text": "Describe the attached media in 5 words!",
804+
},
805+
],
806+
}
807+
events = []
808+
async for event in app.async_stream_query(
809+
user_id=_TEST_USER_ID,
810+
message=video_message,
811+
):
812+
events.append(event)
813+
assert len(events) == 1
814+
assert runner_kwargs["new_message"] == types.Content(
815+
role="user",
816+
parts=[
817+
types.Part(
818+
file_data=types.FileData(
819+
file_uri=(
820+
"gs://cloud-samples-data/generative-ai/video/pixel8.mp4"
821+
),
822+
mime_type="video/mp4",
823+
),
824+
),
825+
types.Part(
826+
text="Describe the attached media in 5 words!",
827+
),
828+
],
829+
)
830+
773831
@pytest.mark.asyncio
774832
async def test_streaming_agent_run_with_events(
775833
self,

‎tests/unit/vertex_adk/test_agent_engine_templates_adk.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -773,6 +773,87 @@ async def test_async_stream_query_with_content(
773773
events.append(event)
774774
assert len(events) == 1
775775

776+
@pytest.mark.asyncio
777+
@pytest.mark.usefixtures(
778+
"default_instrumentor_builder_mock", "get_project_id_mock"
779+
)
780+
async def test_async_stream_query_with_multimodal_content(self):
781+
"""Tests async_stream_query with a multimodal Content dictionary."""
782+
# pylint: disable=protected-access
783+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
784+
assert app._tmpl_attrs.get("runner") is None
785+
app.set_up()
786+
runner_kwargs = {}
787+
788+
async def run_async(**kwargs):
789+
runner_kwargs.update(kwargs)
790+
async for event in _MockRunner().run_async():
791+
yield event
792+
793+
app._tmpl_attrs["runner"] = mock.Mock(run_async=run_async)
794+
video_message = {
795+
"role": "user",
796+
"parts": [
797+
{
798+
"file_data": {
799+
"file_uri": (
800+
"gs://cloud-samples-data/generative-ai/video/pixel8.mp4"
801+
),
802+
"mime_type": "video/mp4",
803+
},
804+
},
805+
{
806+
"text": "Describe the attached media in 5 words!",
807+
},
808+
],
809+
}
810+
events = []
811+
async for event in app.async_stream_query(
812+
user_id=_TEST_USER_ID,
813+
message=video_message,
814+
):
815+
events.append(event)
816+
assert len(events) == 1
817+
expected_content = types.Content(
818+
role="user",
819+
parts=[
820+
types.Part(
821+
file_data=types.FileData(
822+
file_uri=(
823+
"gs://cloud-samples-data/generative-ai/video/pixel8.mp4"
824+
),
825+
mime_type="video/mp4",
826+
),
827+
),
828+
types.Part(
829+
text="Describe the attached media in 5 words!",
830+
),
831+
],
832+
)
833+
assert runner_kwargs["new_message"] == expected_content
834+
835+
# Also verify the REST / requests JSON payload round-trips through
836+
# StreamQueryReasoningEngineRequest into async_stream_query.
837+
rest_request = aip_types.StreamQueryReasoningEngineRequest.from_json(
838+
json.dumps({
839+
"class_method": "async_stream_query",
840+
"input": {
841+
"user_id": _TEST_USER_ID,
842+
"message": video_message,
843+
},
844+
})
845+
)
846+
assert rest_request.class_method == "async_stream_query"
847+
runner_kwargs.clear()
848+
rest_events = [
849+
event
850+
async for event in getattr(app, rest_request.class_method)(
851+
**_utils.to_dict(rest_request)["input"]
852+
)
853+
]
854+
assert len(rest_events) == 1
855+
assert runner_kwargs["new_message"] == expected_content
856+
776857
@pytest.mark.asyncio
777858
async def test_streaming_agent_run_with_events(
778859
self,

‎tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -675,6 +675,63 @@ async def test_async_stream_query_with_content(self):
675675
events.append(event)
676676
assert len(events) == 1
677677

678+
@pytest.mark.asyncio
679+
async def test_async_stream_query_with_multimodal_content(self):
680+
"""Tests async_stream_query with a multimodal Content dictionary."""
681+
# pylint: disable=protected-access
682+
app = reasoning_engines.AdkApp(
683+
agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL)
684+
)
685+
assert app._tmpl_attrs.get("runner") is None
686+
app.set_up()
687+
runner_kwargs = {}
688+
689+
async def run_async(**kwargs):
690+
runner_kwargs.update(kwargs)
691+
async for event in _MockRunner().run_async():
692+
yield event
693+
694+
app._tmpl_attrs["runner"] = mock.Mock(run_async=run_async)
695+
video_message = {
696+
"role": "user",
697+
"parts": [
698+
{
699+
"file_data": {
700+
"file_uri": (
701+
"gs://cloud-samples-data/generative-ai/video/pixel8.mp4"
702+
),
703+
"mime_type": "video/mp4",
704+
},
705+
},
706+
{
707+
"text": "Describe the attached media in 5 words!",
708+
},
709+
],
710+
}
711+
events = []
712+
async for event in app.async_stream_query(
713+
user_id=_TEST_USER_ID,
714+
message=video_message,
715+
):
716+
events.append(event)
717+
assert len(events) == 1
718+
assert runner_kwargs["new_message"] == types.Content(
719+
role="user",
720+
parts=[
721+
types.Part(
722+
file_data=types.FileData(
723+
file_uri=(
724+
"gs://cloud-samples-data/generative-ai/video/pixel8.mp4"
725+
),
726+
mime_type="video/mp4",
727+
),
728+
),
729+
types.Part(
730+
text="Describe the attached media in 5 words!",
731+
),
732+
],
733+
)
734+
678735
def test_streaming_agent_run_with_events(self):
679736
app = reasoning_engines.AdkApp(
680737
agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL)

‎vertexai/agent_engines/templates/adk.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1307,8 +1307,12 @@ async def async_stream_query(
13071307
"""Streams responses asynchronously from the ADK application.
13081308
13091309
Args:
1310-
message (str):
1311-
Required. The message to stream responses for.
1310+
message (Union[str, Dict[str, Any]]):
1311+
Required. The message to stream responses for. Either a text
1312+
string or a dictionary representing a
1313+
`google.genai.types.Content` object (for example,
1314+
`content.model_dump(exclude_none=True, mode="json")` for
1315+
multimodal inputs).
13121316
user_id (str):
13131317
Required. The ID of the user.
13141318
session_id (str):
@@ -1333,7 +1337,7 @@ async def async_stream_query(
13331337
13341338
Raises:
13351339
TypeError: If message is not a string or a dictionary representing
1336-
a Content object.
1340+
a Content object.
13371341
ValueError: If both session_id and session_events are specified.
13381342
"""
13391343
from vertexai.agent_engines import _utils

‎vertexai/preview/reasoning_engines/templates/adk.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1243,8 +1243,12 @@ async def async_stream_query(
12431243
"""Streams responses asynchronously from the ADK application.
12441244
12451245
Args:
1246-
message (str):
1247-
Required. The message to stream responses for.
1246+
message (Union[str, Dict[str, Any]]):
1247+
Required. The message to stream responses for. Either a text
1248+
string or a dictionary representing a
1249+
`google.genai.types.Content` object (for example,
1250+
`content.model_dump(exclude_none=True, mode="json")` for
1251+
multimodal inputs).
12481252
user_id (str):
12491253
Required. The ID of the user.
12501254
session_id (str):
@@ -1260,6 +1264,10 @@ async def async_stream_query(
12601264
12611265
Yields:
12621266
Event dictionaries asynchronously.
1267+
1268+
Raises:
1269+
TypeError: If message is not a string or a dictionary representing
1270+
a Content object.
12631271
"""
12641272
from vertexai.agent_engines import _utils
12651273
from google.genai import types

0 commit comments

Comments
 (0)