Skip to content

Commit 973925b

Browse files
authored
feat: add run_processor support for Agent Runner (#296)
1. **Runner run_processor support** - Add `run_processor` parameter to Runner.__init__() and Runner.run() - Support multiple sources with priority: run() arg > __init__ arg > agent.run_processor > NoOpRunProcessor 2. **Fix IdentityClient region initialization** - Change default region from hardcoded "cn-beijing" to None in WorkloadTokenManager - Auto-detect region using _get_default_region() when not specified - Import _get_default_region from auth_config module Changes: - veadk/runner.py: Add run_processor support with priority chain - veadk/integrations/ve_identity/token_manager.py: Fix region initialization - veadk/integrations/ve_identity/auth_mixins.py: Minor formatting fix
1 parent fe6474a commit 973925b

File tree

3 files changed

+41
-7
lines changed

3 files changed

+41
-7
lines changed

veadk/integrations/ve_identity/auth_processor.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
from google.genai import types
2626
from google.adk.auth.auth_credential import OAuth2Auth
2727

28+
from veadk.integrations.ve_identity.auth_config import _get_default_region
2829
from veadk.processors.base_run_processor import BaseRunProcessor
2930
from veadk.integrations.ve_identity.identity_client import IdentityClient
3031
from veadk.integrations.ve_identity.models import AuthRequestConfig, OAuth2AuthPoller
@@ -178,6 +179,10 @@ def __init__(self, *, config: Optional[AuthRequestConfig] = None):
178179
f"Please open this URL in your browser to authorize: {url}"
179180
)
180181
)
182+
# Use provided region or get from config
183+
if self.config.region is None:
184+
self.config.region = _get_default_region()
185+
181186
self._identity_client = self.config.identity_client or IdentityClient(
182187
region=self.config.region
183188
)

veadk/integrations/ve_identity/token_manager.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from google.adk.tools.tool_context import ToolContext
2323
from google.adk.agents.readonly_context import ReadonlyContext
2424

25+
from veadk.integrations.ve_identity.auth_config import _get_default_region
2526
from veadk.utils.logger import get_logger
2627

2728
from veadk.integrations.ve_identity.identity_client import IdentityClient
@@ -49,13 +50,16 @@ class WorkloadTokenManager:
4950
def __init__(
5051
self,
5152
identity_client: IdentityClient = None,
52-
region: Optional[str] = "cn-beijing",
53+
region: Optional[str] = None,
5354
):
5455
"""Initialize the token manager.
5556
5657
Args:
5758
identity_client: The IdentityClient instance to use for token requests.
5859
"""
60+
if region is None:
61+
region = _get_default_region()
62+
5963
self._identity_client = identity_client or IdentityClient(region=region)
6064

6165
def _build_cache_key(

veadk/runner.py

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
from veadk.config import getenv
3333
from veadk.evaluation import EvalSetRecorder
3434
from veadk.memory.short_term_memory import ShortTermMemory
35+
from veadk.processors.base_run_processor import BaseRunProcessor
3536
from veadk.types import MediaMessage
3637
from veadk.utils.logger import get_logger
3738
from veadk.utils.misc import formatted_timestamp, read_file_to_bytes
@@ -418,6 +419,7 @@ def __init__(
418419
app_name: str = "veadk_default_app",
419420
user_id: str = "veadk_default_user",
420421
upload_inline_data_to_tos: bool = False,
422+
run_processor: "BaseRunProcessor | None" = None,
421423
*args,
422424
**kwargs,
423425
) -> None:
@@ -438,6 +440,8 @@ def __init__(
438440
app_name (str): Application name. Defaults to ``"veadk_default_app"``.
439441
user_id (str): Default user ID. Defaults to ``"veadk_default_user"``.
440442
upload_inline_data_to_tos (bool): Whether to enable inline media upload. Defaults to ``False``.
443+
run_processor (BaseRunProcessor | None): Optional run processor for intercepting agent execution.
444+
If not provided, will try to get from agent. If agent doesn't have one, uses NoOpRunProcessor.
441445
*args: Positional args passed through to ``ADKRunner``.
442446
**kwargs: Keyword args passed through to ``ADKRunner``; may include
443447
``session_service`` and ``memory_service`` to override defaults.
@@ -456,6 +460,16 @@ def __init__(
456460
session_service = kwargs.pop("session_service", None)
457461
memory_service = kwargs.pop("memory_service", None)
458462

463+
# Handle run_processor: priority is runner arg > agent.run_processor > NoOpRunProcessor
464+
if run_processor is not None:
465+
self.run_processor = run_processor
466+
elif hasattr(agent, "run_processor") and agent.run_processor is not None: # type: ignore
467+
self.run_processor = agent.run_processor # type: ignore
468+
else:
469+
from veadk.processors import NoOpRunProcessor
470+
471+
self.run_processor = NoOpRunProcessor()
472+
459473
if session_service:
460474
if short_term_memory:
461475
logger.warning(
@@ -511,6 +525,7 @@ async def run(
511525
run_config: RunConfig | None = None,
512526
save_tracing_data: bool = False,
513527
upload_inline_data_to_tos: bool = False,
528+
run_processor: "BaseRunProcessor | None" = None,
514529
):
515530
"""Run a conversation with multi-turn text and multimodal inputs.
516531
@@ -527,6 +542,8 @@ async def run(
527542
config is created using the environment var ``MODEL_AGENT_MAX_LLM_CALLS``.
528543
save_tracing_data (bool): Whether to dump tracing data to disk after the run. Defaults to ``False``.
529544
upload_inline_data_to_tos (bool): Whether to enable media upload only for this run. Defaults to ``False``.
545+
run_processor (BaseRunProcessor | None): Optional run processor to use for this run.
546+
If not provided, uses the runner's default run_processor. Defaults to None.
530547
531548
Returns:
532549
str: The textual output from the last event, if present; otherwise an empty string.
@@ -567,12 +584,20 @@ async def run(
567584
final_output = ""
568585
for converted_message in converted_messages:
569586
try:
570-
async for event in self.run_async(
571-
user_id=user_id,
572-
session_id=session_id,
573-
new_message=converted_message,
574-
run_config=run_config,
575-
):
587+
588+
@(run_processor or self.run_processor).process_run(
589+
runner=self, message=converted_message
590+
)
591+
async def event_generator():
592+
async for event in self.run_async(
593+
user_id=user_id,
594+
session_id=session_id,
595+
new_message=converted_message,
596+
run_config=run_config,
597+
):
598+
yield event
599+
600+
async for event in event_generator():
576601
if event.get_function_calls():
577602
for function_call in event.get_function_calls():
578603
logger.debug(f"Function call: {function_call}")

0 commit comments

Comments
 (0)