Skip to content

Commit 4406aa0

Browse files
authored
chore: add runner doc string (#265)
* add runner doc string * refine example code
1 parent 1ab8149 commit 4406aa0

File tree

1 file changed

+263
-1
lines changed

1 file changed

+263
-1
lines changed

veadk/runner.py

Lines changed: 263 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,28 @@
4848

4949

5050
async def pre_run_process(self, process_func, new_message, user_id, session_id):
51+
"""Pre-run hook invoked before agent execution.
52+
53+
Iterates over all ``parts`` of ``new_message`` and, when a ``part`` contains
54+
``inline_data`` and uploading is enabled, calls ``process_func`` to process
55+
the data (for example, upload to TOS and rewrite with an accessible URL).
56+
Typically used together with the ``intercept_new_message`` decorator.
57+
58+
Args:
59+
self: Runner instance.
60+
process_func: An async processing function with a signature like
61+
``(part, app_name, user_id, session_id)`` used to handle
62+
``inline_data`` in the message (e.g., upload to TOS).
63+
new_message (google.genai.types.Content): Incoming user message.
64+
user_id (str): User identifier.
65+
session_id (str): Session identifier.
66+
67+
Returns:
68+
None
69+
70+
Raises:
71+
Exception: Propagated if ``process_func`` raises and does not handle it.
72+
"""
5173
if new_message.parts:
5274
for part in new_message.parts:
5375
if part.inline_data and self.upload_inline_data_to_tos:
@@ -60,10 +82,42 @@ async def pre_run_process(self, process_func, new_message, user_id, session_id):
6082

6183

6284
def post_run_process(self):
85+
"""Post-run hook executed after agent run.
86+
87+
This is currently a no-op placeholder and can be extended to perform
88+
cleanup or finalize logic after a run.
89+
90+
Args:
91+
self: Runner instance.
92+
93+
Returns:
94+
None
95+
96+
Raises:
97+
None
98+
"""
6399
return
64100

65101

66102
def intercept_new_message(process_func):
103+
"""Create a decorator to insert pre/post hooks around ``run_async`` calls.
104+
105+
Internally it invokes :func:`pre_run_process` to preprocess the incoming
106+
message (e.g., upload image/video inline data to TOS), then iterates the
107+
underlying event stream and finally calls :func:`post_run_process`.
108+
109+
Args:
110+
process_func: Async function used to process ``inline_data`` (typically
111+
``_upload_image_to_tos``).
112+
113+
Returns:
114+
Callable: A decorator that can wrap ``run_async``.
115+
116+
Raises:
117+
Exception: May propagate exceptions raised by the wrapped function or
118+
the pre-processing step.
119+
"""
120+
67121
def decorator(func):
68122
@functools.wraps(func)
69123
async def wrapper(
@@ -97,7 +151,34 @@ def _convert_messages(
97151
user_id: str,
98152
session_id: str,
99153
) -> list:
100-
"""Convert VeADK formatted messages to Google ADK formatted messages."""
154+
"""Convert a VeADK ``RunnerMessage`` into a list of Google ADK messages.
155+
156+
Supported inputs:
157+
- ``str``: Single-turn text prompt.
158+
- :class:`veadk.types.MediaMessage`: Single-turn multimodal prompt (text + image/video).
159+
- ``list``: A list of the above types (multi-turn with mixed text and multimodal).
160+
161+
For multimodal inputs, this reads the local media file bytes and detects
162+
the MIME type via ``filetype``; only ``image/*`` and ``video/*`` are supported.
163+
164+
Args:
165+
messages (RunnerMessage): Input message or list of messages to convert.
166+
app_name (str): App name (not directly used; kept for consistency with upload path).
167+
user_id (str): User ID (not directly used; kept for consistency with upload path).
168+
session_id (str): Session ID (not directly used; kept for consistency with upload path).
169+
170+
Returns:
171+
list[google.genai.types.Content]: Converted ADK messages.
172+
173+
Raises:
174+
ValueError: If the message type is unknown or media type is unrecognized.
175+
AssertionError: If the media MIME type is not supported (only image/* and video/*).
176+
177+
Note:
178+
This function only performs structural conversion. To upload inline media
179+
to an object store and rewrite URLs, use it together with
180+
``intercept_new_message`` and ``_upload_image_to_tos``.
181+
"""
101182
if isinstance(messages, str):
102183
_messages = [types.Content(role="user", parts=[types.Part(text=messages)])]
103184
elif isinstance(messages, MediaMessage):
@@ -146,6 +227,25 @@ def _convert_messages(
146227
async def _upload_image_to_tos(
147228
part: genai.types.Part, app_name: str, user_id: str, session_id: str
148229
) -> None:
230+
"""Upload inline media data in a message part to TOS and rewrite its URL.
231+
232+
When ``part.inline_data`` has both ``display_name`` (original filename) and
233+
``data`` (bytes), it generates an object storage path based on
234+
``app_name``, ``user_id`` and ``session_id``. After upload, it replaces
235+
``display_name`` with a signed URL.
236+
237+
Args:
238+
part (google.genai.types.Part): Message part containing ``inline_data``.
239+
app_name (str): App name.
240+
user_id (str): User ID.
241+
session_id (str): Session ID.
242+
243+
Returns:
244+
None
245+
246+
Raises:
247+
None: All exceptions are caught and logged; nothing is propagated.
248+
"""
149249
try:
150250
if part.inline_data and part.inline_data.display_name and part.inline_data.data:
151251
from veadk.integrations.ve_tos.ve_tos import VeTOS
@@ -164,6 +264,51 @@ async def _upload_image_to_tos(
164264

165265

166266
class Runner(ADKRunner):
267+
"""VeADK Runner that augments ADK with session, memory, tracing, and media upload.
268+
269+
This class builds on Google ADK's ``Runner`` and adds:
270+
- Integration with short-term memory (ShortTermMemory) for auto session management.
271+
- Optional long-term memory integration and session persistence.
272+
- New message interception and media upload to TOS.
273+
- Tracing dump and Trace ID logging.
274+
- A simplified ``run`` entry that supports multi-turn text/multimodal inputs.
275+
276+
Attributes:
277+
user_id (str): Default user ID.
278+
long_term_memory: Long-term memory service instance, or ``None`` if not set.
279+
short_term_memory (veadk.memory.short_term_memory.ShortTermMemory | None):
280+
Short-term memory instance used to auto-create/manage sessions.
281+
upload_inline_data_to_tos (bool): Whether to upload inline media to TOS while running.
282+
session_service: Session service instance (may come from short-term memory).
283+
memory_service: Memory service instance (may come from agent's long-term memory).
284+
app_name (str): Application name used in session management and object pathing.
285+
286+
Examples:
287+
Create a runner and perform a text-only interaction:
288+
289+
```python
290+
from veadk.runner import Runner
291+
from veadk.agent import Agent # assume it's properly constructed
292+
runner = Runner(agent=my_agent, app_name="demo_app", user_id="u1")
293+
output = await runner.run("Hello")
294+
print(output)
295+
```
296+
297+
Send multimodal (text + image):
298+
299+
```python
300+
from veadk.types import MediaMessage
301+
msg = MediaMessage(text="Describe the image", media="/path/to/image.png")
302+
output = await runner.run(msg, upload_inline_data_to_tos=True)
303+
print(output)
304+
```
305+
306+
Note:
307+
This class wraps the parent ``run_async`` at initialization to insert media
308+
upload and post-run handling. If you override the underlying ``run_async``,
309+
ensure it remains compatible with this interception logic.
310+
"""
311+
167312
def __init__(
168313
self,
169314
agent: BaseAgent | Agent,
@@ -174,6 +319,33 @@ def __init__(
174319
*args,
175320
**kwargs,
176321
) -> None:
322+
"""Initialize a Runner instance.
323+
324+
Selects the session service based on provided short-term memory or an
325+
external ``session_service``. If long-term memory or an external
326+
``memory_service`` is provided, the passed service is preferred. After
327+
construction, it injects a message interception layer into the parent's
328+
``run_async`` to support inline media upload and post-run handling.
329+
330+
Args:
331+
agent (google.adk.agents.base_agent.BaseAgent | veadk.agent.Agent):
332+
The agent instance used to run interactions.
333+
short_term_memory (ShortTermMemory | None): Optional short-term memory; if
334+
not provided and no external ``session_service`` is supplied, an in-memory
335+
session service will be created.
336+
app_name (str): Application name. Defaults to ``"veadk_default_app"``.
337+
user_id (str): Default user ID. Defaults to ``"veadk_default_user"``.
338+
upload_inline_data_to_tos (bool): Whether to enable inline media upload. Defaults to ``False``.
339+
*args: Positional args passed through to ``ADKRunner``.
340+
**kwargs: Keyword args passed through to ``ADKRunner``; may include
341+
``session_service`` and ``memory_service`` to override defaults.
342+
343+
Returns:
344+
None
345+
346+
Raises:
347+
None
348+
"""
177349
self.user_id = user_id
178350
self.long_term_memory = None
179351
self.short_term_memory = short_term_memory
@@ -238,6 +410,30 @@ async def run(
238410
save_tracing_data: bool = False,
239411
upload_inline_data_to_tos: bool = False,
240412
):
413+
"""Run a conversation with multi-turn text and multimodal inputs.
414+
415+
When short-term memory is configured, a session is auto-created as needed.
416+
Inputs are converted into ADK message format. If ``upload_inline_data_to_tos``
417+
is ``True``, media upload is enabled temporarily for this run (does not change
418+
the Runner's global setting).
419+
420+
Args:
421+
messages (RunnerMessage): Input messages (``str``, ``MediaMessage`` or a list of them).
422+
user_id (str): Override default user ID; if empty, uses the constructed ``user_id``.
423+
session_id (str): Session ID. Defaults to a timestamp-based temporary ID.
424+
run_config (google.adk.agents.RunConfig | None): Run config; if ``None``, a default
425+
config is created using the environment var ``MODEL_AGENT_MAX_LLM_CALLS``.
426+
save_tracing_data (bool): Whether to dump tracing data to disk after the run. Defaults to ``False``.
427+
upload_inline_data_to_tos (bool): Whether to enable media upload only for this run. Defaults to ``False``.
428+
429+
Returns:
430+
str: The textual output from the last event, if present; otherwise an empty string.
431+
432+
Raises:
433+
ValueError: If an input contains an unsupported or unrecognized media type.
434+
AssertionError: If a media MIME type is not among ``image/*`` or ``video/*``.
435+
Exception: Exceptions from the underlying ADK/Agent execution may propagate.
436+
"""
241437
if upload_inline_data_to_tos:
242438
_upload_inline_data_to_tos = self.upload_inline_data_to_tos
243439
self.upload_inline_data_to_tos = upload_inline_data_to_tos
@@ -302,6 +498,17 @@ async def run(
302498
return final_output
303499

304500
def get_trace_id(self) -> str:
501+
"""Get the Trace ID from the current agent's tracer.
502+
503+
If the agent is not a :class:`veadk.agent.Agent` or no tracer is configured,
504+
returns ``"<unknown_trace_id>"``.
505+
506+
Returns:
507+
str: The Trace ID or ``"<unknown_trace_id>"``.
508+
509+
Raises:
510+
None
511+
"""
305512
if not isinstance(self.agent, Agent):
306513
logger.warning(
307514
("The agent is not an instance of VeADK Agent, no trace id provided.")
@@ -322,6 +529,17 @@ def get_trace_id(self) -> str:
322529
return "<unknown_trace_id>"
323530

324531
def _print_trace_id(self) -> None:
532+
"""Log the current tracer's Trace ID.
533+
534+
If the agent is not a :class:`veadk.agent.Agent` or no tracer is configured,
535+
nothing is printed.
536+
537+
Returns:
538+
None
539+
540+
Raises:
541+
None
542+
"""
325543
if not isinstance(self.agent, Agent):
326544
logger.warning(
327545
("The agent is not an instance of VeADK Agent, no trace id provided.")
@@ -342,6 +560,21 @@ def _print_trace_id(self) -> None:
342560
return
343561

344562
def save_tracing_file(self, session_id: str) -> str:
563+
"""Dump tracing data to disk and return the last written path.
564+
565+
Only effective when the agent is one of
566+
Agent/SequentialAgent/ParallelAgent/LoopAgent and a tracer is configured;
567+
otherwise returns an empty string.
568+
569+
Args:
570+
session_id (str): Session ID used to associate the tracing with a session.
571+
572+
Returns:
573+
str: The tracing file path; returns an empty string on failure or when no tracer.
574+
575+
Raises:
576+
None: All errors are logged and an empty string is returned.
577+
"""
345578
if not isinstance(
346579
self.agent, (Agent, SequentialAgent, ParallelAgent, LoopAgent)
347580
):
@@ -367,6 +600,18 @@ def save_tracing_file(self, session_id: str) -> str:
367600
return ""
368601

369602
async def save_eval_set(self, session_id: str, eval_set_id: str = "default") -> str:
603+
"""Save the current session as part of an evaluation set and return its path.
604+
605+
Args:
606+
session_id (str): Session ID.
607+
eval_set_id (str): Evaluation set identifier. Defaults to ``"default"``.
608+
609+
Returns:
610+
str: The exported evaluation set file path.
611+
612+
Raises:
613+
Exception: Propagated if the underlying export logic raises.
614+
"""
370615
eval_set_recorder = EvalSetRecorder(self.session_service, eval_set_id)
371616
eval_set_path = await eval_set_recorder.dump(
372617
self.app_name, self.user_id, session_id
@@ -376,6 +621,23 @@ async def save_eval_set(self, session_id: str, eval_set_id: str = "default") ->
376621
async def save_session_to_long_term_memory(
377622
self, session_id: str, user_id: str = "", app_name: str = ""
378623
) -> None:
624+
"""Save the specified session to long-term memory.
625+
626+
If ``long_term_memory`` is not configured, the function logs a warning and returns.
627+
It fetches the session from the session service and then calls the long-term memory's
628+
``add_session_to_memory`` for persistence.
629+
630+
Args:
631+
session_id (str): Session ID.
632+
user_id (str): Optional; override default user ID. If empty, uses ``self.user_id``.
633+
app_name (str): Optional; override default app name. If empty, uses ``self.app_name``.
634+
635+
Returns:
636+
None
637+
638+
Raises:
639+
Exception: May propagate if the underlying memory service raises during write.
640+
"""
379641
if not self.long_term_memory:
380642
logger.warning("Long-term memory is not enabled. Failed to save session.")
381643
return

0 commit comments

Comments
 (0)