refactor(ve_identity): introduce BaseRunProcessor to decouple Agent Identity dependency

loveyana · loveyana · commit 6e29ec01ce7e · 2025-10-31T00:28:03.000+08:00
## Problem
Agent class had a hard dependency on Agent Identity's AuthRequestProcessor,
causing all agents to depend on Agent Identity even when authentication is
not needed. This violates the Dependency Inversion Principle and could
break existing agents.

## Solution
Introduce an abstract processor layer following the Dependency Inversion
Principle:

1. Created `veadk/processors/` package with:
   - `BaseRunProcessor`: Abstract base class for runtime processors
   - `NoOpRunProcessor`: Default no-op implementation

2. Modified `veadk/agent.py`:
   - Added `run_processor: Optional[BaseRunProcessor]` field
   - Changed dependency from AuthRequestProcessor to BaseRunProcessor
   - Initialize with NoOpRunProcessor by default
   - Updated `run()` method parameter from `auth_request_processor` to `run_processor`

3. Updated `veadk/integrations/ve_identity/auth_processor.py`:
   - Made AuthRequestProcessor inherit from BaseRunProcessor
   - Renamed main method from `with_auth_loop` to `process_run`
diff --git a/veadk/agent.py b/veadk/agent.py
@@ -33,7 +33,7 @@
     DEFAULT_MODEL_EXTRA_CONFIG,
 )
 from veadk.evaluation import EvalSetRecorder
-from veadk.integrations.ve_identity import AuthRequestProcessor
+from veadk.processors import BaseRunProcessor, NoOpRunProcessor
 from veadk.knowledgebase import KnowledgeBase
 from veadk.memory.long_term_memory import LongTermMemory
 from veadk.memory.short_term_memory import ShortTermMemory
@@ -98,9 +98,33 @@ class Agent(LlmAgent):
     tracers: list[BaseTracer] = []
     """The tracers provided to agent."""
 
+    run_processor: Optional[BaseRunProcessor] = Field(default=None, exclude=True)
+    """Optional run processor for intercepting and processing agent execution flows.
+
+    The run processor can be used to implement cross-cutting concerns such as:
+    - Authentication flows (e.g., OAuth2 via VeIdentity)
+    - Request/response logging
+    - Error handling and retry logic
+    - Performance monitoring
+
+    If not provided, a NoOpRunProcessor will be used by default.
+
+    Example:
+        from veadk.integrations.ve_identity import AuthRequestProcessor
+
+        agent = Agent(
+            name="my-agent",
+            run_processor=AuthRequestProcessor()
+        )
+    """
+
     def model_post_init(self, __context: Any) -> None:
         super().model_post_init(None)  # for sub_agents init
 
+        # Initialize run_processor if not provided
+        if self.run_processor is None:
+            self.run_processor = NoOpRunProcessor()
+
         # combine user model config with VeADK defaults
         headers = DEFAULT_MODEL_EXTRA_CONFIG["extra_headers"].copy()
         body = DEFAULT_MODEL_EXTRA_CONFIG["extra_body"].copy()
@@ -168,11 +192,27 @@ async def _run(
         session_id: str,
         message: types.Content,
         stream: bool,
-        auth_request_processor: AuthRequestProcessor,
+        run_processor: Optional[BaseRunProcessor] = None,
     ):
+        """Internal run method with run processor support.
+
+        Args:
+            runner: The Runner instance.
+            user_id: User ID for the session.
+            session_id: Session ID.
+            message: The message to send.
+            stream: Whether to stream the output.
+            run_processor: Optional run processor to use. If not provided, uses self.run_processor.
+
+        Returns:
+            The final output string.
+        """
         stream_mode = StreamingMode.SSE if stream else StreamingMode.NONE
 
-        @auth_request_processor.with_auth_loop(runner=runner, message=message)
+        # Use provided run_processor or fall back to instance's run_processor
+        processor = run_processor or self.run_processor
+
+        @processor.process_run(runner=runner, message=message)
         async def event_generator():
             async for event in runner.run_async(
                 user_id=user_id,
@@ -248,7 +288,7 @@ async def run(
         collect_runtime_data: bool = False,
         eval_set_id: str = "",
         save_session_to_memory: bool = False,
-        auth_request_processor: AuthRequestProcessor = AuthRequestProcessor(),
+        run_processor: Optional[BaseRunProcessor] = None,
     ):
         """Running the agent. The runner and session service will be created automatically.
 
@@ -265,6 +305,8 @@ async def run(
             collect_runtime_data (bool, optional): Whether to collect runtime data. Defaults to False.
             eval_set_id (str, optional): The id of the eval set. Defaults to "".
             save_session_to_memory (bool, optional): Whether to save this turn session to memory. Defaults to False.
+            run_processor (Optional[BaseRunProcessor], optional): Optional run processor to use for this run.
+                If not provided, uses the agent's default run_processor. Defaults to None.
         """
 
         logger.warning(
@@ -299,7 +341,7 @@ async def run(
         for _prompt in prompt:
             message = types.Content(role="user", parts=[types.Part(text=_prompt)])
             final_output = await self._run(
-                runner, user_id, session_id, message, stream, auth_request_processor
+                runner, user_id, session_id, message, stream, run_processor
             )
 
         # VeADK features
diff --git a/veadk/integrations/ve_identity/auth_processor.py b/veadk/integrations/ve_identity/auth_processor.py
@@ -19,12 +19,13 @@
 import asyncio
 import json
 import time
-from typing import Awaitable, Callable, Optional, TYPE_CHECKING
+from typing import Any, Awaitable, Callable, Optional, TYPE_CHECKING
 
 from google.adk.auth import AuthConfig
 from google.genai import types
 from google.adk.auth.auth_credential import OAuth2Auth
 
+from veadk.processors.base_run_processor import BaseRunProcessor
 from veadk.integrations.ve_identity.identity_client import IdentityClient
 from veadk.integrations.ve_identity.models import AuthRequestConfig, OAuth2AuthPoller
 from veadk.integrations.ve_identity.utils import (
@@ -155,19 +156,7 @@ async def poll_for_auth(self) -> OAuth2Auth:
         )
 
 
-class _NoOpAuthProcessor:
-    """No-op auth processor that doesn't modify the event generator."""
-
-    def with_auth_loop(self, runner, message):
-        """Return a decorator that does nothing."""
-
-        def decorator(func):
-            return func
-
-        return decorator
-
-
-class AuthRequestProcessor:
+class AuthRequestProcessor(BaseRunProcessor):
     """Processor for handling authentication requests in agent conversations.
 
     This class manages the OAuth2 authentication flow when tools require user authorization.
@@ -277,19 +266,22 @@ async def process_auth_request(
         logger.info(f"Auth request {auth_request_event_id} processed successfully")
         return auth_content
 
-    def with_auth_loop(
+    def process_run(
         self,
         runner: Runner,
         message: types.Content,
-        task_updater: Optional[TaskUpdater] = None,
+        **kwargs: Any,
     ):
-        """Decorator to add authentication loop handling to event generators.
+        """Process the agent run by wrapping the event generator with authentication loop.
+
+        This method implements the BaseRunProcessor interface and adds authentication
+        loop handling to event generators.
 
         This decorator intercepts runner.run_async calls and automatically handles
         authentication loops. The event_generator code can remain completely unchanged!
 
         Usage example:
-            @auth_processor.with_auth_loop(
+            @auth_processor.process_run(
                 runner=runner,
                 message=message,
             )
@@ -319,10 +311,14 @@ async def event_generator():
         Args:
             runner: Runner instance (will be wrapped).
             message: Initial message to send.
+            **kwargs: Additional keyword arguments. Supports:
+                     - task_updater: Optional TaskUpdater for status updates.
 
         Returns:
             Decorated generator function.
         """
+        # Extract task_updater from kwargs
+        task_updater = kwargs.get("task_updater")
 
         def decorator(event_generator_func):
             async def wrapper():
@@ -338,13 +334,13 @@ async def wrapper():
                     # Create a wrapped runner to intercept run_async calls
                     original_run_async = runner.run_async
 
-                    async def wrapped_run_async(**kwargs):
+                    async def wrapped_run_async(**run_kwargs):
                         nonlocal auth_request_event_id, auth_config
 
                         # Override the message with the current message
-                        kwargs["new_message"] = current_message
+                        run_kwargs["new_message"] = current_message
 
-                        async for event in original_run_async(**kwargs):
+                        async for event in original_run_async(**run_kwargs):
                             # Detect authentication events
                             if is_pending_auth_event(event):
                                 auth_request_event_id = get_function_call_id(event)
diff --git a/veadk/processors/__init__.py b/veadk/processors/__init__.py
@@ -0,0 +1,27 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# 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.
+
+"""Runtime processors for VeADK agents.
+
+This module provides the abstract base classes for runtime processors that can
+intercept and process agent execution flows.
+"""
+
+from veadk.processors.base_run_processor import BaseRunProcessor, NoOpRunProcessor
+
+__all__ = [
+    "BaseRunProcessor",
+    "NoOpRunProcessor",
+]
+
diff --git a/veadk/processors/base_run_processor.py b/veadk/processors/base_run_processor.py
@@ -0,0 +1,120 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# 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.
+
+"""Base run processor for intercepting and processing agent execution."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable
+
+if TYPE_CHECKING:
+    from google.genai import types
+    from veadk.runner import Runner
+
+
+class BaseRunProcessor(ABC):
+    """Abstract base class for runtime processors.
+
+    A run processor can intercept and modify the agent execution flow by wrapping
+    the event generator function. This is useful for implementing cross-cutting
+    concerns such as:
+    - Authentication flows (e.g., OAuth2)
+    - Request/response logging
+    - Error handling and retry logic
+    - Performance monitoring
+    - Custom event filtering or transformation
+
+    The processor uses a decorator pattern to wrap the event generator, allowing
+    it to:
+    1. Intercept events from runner.run_async
+    2. Process or modify events
+    3. Inject additional events (e.g., authentication requests)
+    4. Control the execution flow (e.g., retry loops)
+
+    Example:
+        class MyProcessor(BaseRunProcessor):
+            def process_run(self, runner, message, **kwargs):
+                def decorator(event_generator_func):
+                    async def wrapper():
+                        # Pre-processing
+                        async for event in event_generator_func():
+                            # Process each event
+                            yield event
+                        # Post-processing
+                    return wrapper
+                return decorator
+    """
+
+    @abstractmethod
+    def process_run(
+        self,
+        runner: Runner,
+        message: types.Content,
+        **kwargs: Any,
+    ) -> Callable[[Callable[[], AsyncGenerator]], Callable[[], AsyncGenerator]]:
+        """Process the agent run by wrapping the event generator.
+
+        This method returns a decorator that wraps the event generator function.
+        The decorator can intercept events, modify them, or inject new events.
+
+        Args:
+            runner: The Runner instance executing the agent.
+            message: The initial message to send to the agent.
+            **kwargs: Additional keyword arguments that may be needed by specific
+                     implementations (e.g., task_updater for status updates).
+
+        Returns:
+            A decorator function that takes an event generator function and returns
+            a wrapped event generator function.
+
+        Example:
+            @processor.process_run(runner=runner, message=message)
+            async def event_generator():
+                async for event in runner.run_async(...):
+                    yield event
+        """
+        pass
+
+
+class NoOpRunProcessor(BaseRunProcessor):
+    """No-op run processor that doesn't modify the event generator.
+
+    This is the default processor used when no specific processing is needed.
+    It simply passes through all events without any modification.
+    """
+
+    def process_run(
+        self,
+        runner: Runner,
+        message: types.Content,
+        **kwargs: Any,
+    ) -> Callable[[Callable[[], AsyncGenerator]], Callable[[], AsyncGenerator]]:
+        """Return a decorator that does nothing.
+
+        Args:
+            runner: The Runner instance (unused).
+            message: The initial message (unused).
+            **kwargs: Additional keyword arguments (unused).
+
+        Returns:
+            A decorator that returns the original function unchanged.
+        """
+
+        def decorator(
+            event_generator_func: Callable[[], AsyncGenerator]
+        ) -> Callable[[], AsyncGenerator]:
+            return event_generator_func
+
+        return decorator
