From 4dce9ef5196c1b174e81da106b7b0b002f0fbc16 Mon Sep 17 00:00:00 2001 From: Che Liu Date: Thu, 10 Jul 2025 17:22:12 -0700 Subject: [PATCH 1/3] feat: ADK Plugin Base Class Plugin is a collection of callbacks (lifecycle hooks), including callbacks for Agent, Tools, Model and App. It helps developers to add customizations and change behaviors of their agent application easily This is the first change that introduces only the base interfaces based on the doc. plugins folder is generated with `create_symlink.sh plugins` PiperOrigin-RevId: 781745044 --- src/google/adk/plugins/__init__.py | 17 ++ src/google/adk/plugins/base_plugin.py | 320 ++++++++++++++++++++++++++ 2 files changed, 337 insertions(+) create mode 100644 src/google/adk/plugins/__init__.py create mode 100644 src/google/adk/plugins/base_plugin.py diff --git a/src/google/adk/plugins/__init__.py b/src/google/adk/plugins/__init__.py new file mode 100644 index 00000000..b0c771ed --- /dev/null +++ b/src/google/adk/plugins/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2025 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 in 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. + +from .base_plugin import BasePlugin + +__all__ = ['BasePlugin'] diff --git a/src/google/adk/plugins/base_plugin.py b/src/google/adk/plugins/base_plugin.py new file mode 100644 index 00000000..cc7d85f3 --- /dev/null +++ b/src/google/adk/plugins/base_plugin.py @@ -0,0 +1,320 @@ +# Copyright 2025 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 in 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. + +from __future__ import annotations + +from abc import ABC +from typing import Any +from typing import Optional +from typing import TYPE_CHECKING +from typing import TypeVar + +from google.genai import types + +from ..agents.base_agent import BaseAgent +from ..agents.callback_context import CallbackContext +from ..events.event import Event +from ..models.llm_request import LlmRequest +from ..models.llm_response import LlmResponse +from ..tools.base_tool import BaseTool +from ..utils.feature_decorator import working_in_progress + +if TYPE_CHECKING: + from ..agents.invocation_context import InvocationContext + from ..tools.tool_context import ToolContext + + +# Type alias: The value may or may not be awaitable, and value is optional. +T = TypeVar("T") + + +@working_in_progress( + "Plugin is under development now. Check again around Jul. 2025" +) +class BasePlugin(ABC): + """Base class for creating plugins. + + Plugins provide a structured way to intercept and modify agent, tool, and + LLM behaviors at critical execution points in a callback manner. While agent + callbacks apply to a particular agent, plugins applies globally to all + agents added in the runner. Plugins are best used for adding custom behaviors + like logging, monitoring, caching, or modifying requests and responses at key + stages. + + A plugin can implement one or more methods of callbacks, but should not + implement the same method of callback for multiple times. + + Relation with [Agent callbacks](https://google.github.io/adk-docs/callbacks/): + + **Execution Order** + Similar to Agent callbacks, Plugins are executed in the order they are + registered. However, Plugin and Agent Callbacks are executed sequentially, + with Plugins takes precedence over agent callbacks. When the callback in a + plugin returns a value, it will short circuit all remaining plugins and + agent callbacks, causing all remaining plugins and agent callbacks + to be skipped. + + **Change Propagation** + Plugins and agent callbacks can both modify the value of the input parameters, + including agent input, tool input, and LLM request/response, etc. They work in + the exactly same way. The modifications will be visible and passed to the next + callback in the chain. For example, if a plugin modifies the tool input with + before_tool_callback, the modified tool input will be passed to the + before_tool_callback of the next plugin, and further passed to the agent + callbacks if not short circuited. + + To use a plugin, implement the desired callback methods and pass an instance + of your custom plugin class to the ADK Runner. + + Examples: + A simple plugin that logs every tool call. + + >>> class ToolLoggerPlugin(BasePlugin): + .. def __init__(self): + .. super().__init__(name="tool_logger") + .. + .. async def before_tool_callback( + .. self, *, tool: BaseTool, tool_args: dict[str, Any], + tool_context: + ToolContext + .. ): + .. print(f"[{self.name}] Calling tool '{tool.name}' with args: + {tool_args}") + .. + .. async def after_tool_callback( + .. self, *, tool: BaseTool, tool_args: dict, tool_context: + ToolContext, result: dict + .. ): + .. print(f"[{self.name}] Tool '{tool.name}' finished with result: + {result}") + .. + >>> # Add the plugin to ADK Runner + >>> # runner = Runner( + >>> # ... + >>> # plugins=[ToolLoggerPlugin(), AgentPolicyPlugin()], + >>> # ) + """ + + def __init__(self, name: str): + """Initializes the plugin. + + Args: + name: A unique identifier for this plugin instance. + """ + super().__init__() + self.name = name + + async def on_user_message_callback( + self, + *, + invocation_context: InvocationContext, + user_message: types.Content, + ) -> Optional[types.Content]: + """Callback executed when a user message is received before an invocation starts. + + This callback helps logging and modifying the user message before the + runner starts the invocation. + + Args: + invocation_context: The context for the entire invocation. + user_message: The message content input by user. + + Returns: + An optional `types.Content` to be returned to the ADK. Returning a + value to replace the user message. Returning `None` to proceed + normally. + """ + pass + + async def before_run_callback( + self, *, invocation_context: InvocationContext + ) -> Optional[types.Content]: + """Callback executed before the ADK runner runs. + + This is the first callback to be called in the lifecycle, ideal for global + setup or initialization tasks. + + Args: + invocation_context: The context for the entire invocation, containing + session information, the root agent, etc. + + Returns: + An optional `Event` to be returned to the ADK. Returning a value to + halt execution of the runner and ends the runner with that event. Return + `None` to proceed normally. + """ + pass + + async def on_event_callback( + self, *, invocation_context: InvocationContext, event: Event + ) -> Optional[Event]: + """Callback executed after an event is yielded from runner. + + This is the ideal place to make modification to the event before the event + is handled by the underlying agent app. + + Args: + invocation_context: The context for the entire invocation. + event: The event raised by the runner. + + Returns: + An optional value. A non-`None` return may be used by the framework to + modify or replace the response. Returning `None` allows the original + response to be used. + """ + pass + + async def after_run_callback( + self, *, invocation_context: InvocationContext + ) -> Optional[None]: + """Callback executed after an ADK runner run has completed. + + This is the final callback in the ADK lifecycle, suitable for cleanup, final + logging, or reporting tasks. + + Args: + invocation_context: The context for the entire invocation. + + Returns: + None + """ + pass + + async def before_agent_callback( + self, *, agent: BaseAgent, callback_context: CallbackContext + ) -> Optional[types.Content]: + """Callback executed before an agent's primary logic is invoked. + + This callback can be used for logging, setup, or to short-circuit the + agent's execution by returning a value. + + Args: + agent: The agent that is about to run. + callback_context: The context for the agent invocation. + + Returns: + An optional `types.Content` object. If a value is returned, it will bypass + the agent's callbacks and its execution, and return this value directly. + Returning `None` allows the agent to proceed normally. + """ + pass + + async def after_agent_callback( + self, *, agent: BaseAgent, callback_context: CallbackContext + ) -> Optional[types.Content]: + """Callback executed after an agent's primary logic has completed. + + This callback can be used to inspect, log, or modify the agent's final + result before it is returned. + + Args: + agent: The agent that has just run. + callback_context: The context for the agent invocation. + + Returns: + An optional `types.Content` object. If a value is returned, it will + replace the agent's original result. Returning `None` uses the original, + unmodified result. + """ + pass + + async def before_model_callback( + self, *, callback_context: CallbackContext, llm_request: LlmRequest + ) -> Optional[LlmResponse]: + """Callback executed before a request is sent to the model. + + This provides an opportunity to inspect, log, or modify the `LlmRequest` + object. It can also be used to implement caching by returning a cached + `LlmResponse`, which would skip the actual model call. + + Args: + callback_context: The context for the current agent call. + llm_request: The prepared request object to be sent to the model. + + Returns: + An optional value. The interpretation of a non-`None` trigger an early + exit and returns the response immediately. Returning `None` allows the LLM + request to proceed normally. + """ + pass + + async def after_model_callback( + self, *, callback_context: CallbackContext, llm_response: LlmResponse + ) -> Optional[LlmResponse]: + """Callback executed after a response is received from the model. + + This is the ideal place to log model responses, collect metrics on token + usage, or perform post-processing on the raw `LlmResponse`. + + Args: + callback_context: The context for the current agent call. + llm_response: The response object received from the model. + + Returns: + An optional value. A non-`None` return may be used by the framework to + modify or replace the response. Returning `None` allows the original + response to be used. + """ + pass + + async def before_tool_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + ) -> Optional[dict]: + """Callback executed before a tool is called. + + This callback is useful for logging tool usage, input validation, or + modifying the arguments before they are passed to the tool. + + Args: + tool: The tool instance that is about to be executed. + tool_args: The dictionary of arguments to be used for invoking the tool. + tool_context: The context specific to the tool execution. + + Returns: + An optional dictionary. If a dictionary is returned, it will stop the tool + execution and return this response immediately. Returning `None` uses the + original, unmodified arguments. + """ + pass + + async def after_tool_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + result: dict, + ) -> Optional[dict]: + """Callback executed after a tool has been called. + + This callback allows for inspecting, logging, or modifying the result + returned by a tool. + + Args: + tool: The tool instance that has just been executed. + tool_args: The original arguments that were passed to the tool. + tool_context: The context specific to the tool execution. + result: The dictionary returned by the tool invocation. + + Returns: + An optional dictionary. If a dictionary is returned, it will **replace** + the original result from the tool. This allows for post-processing or + altering tool outputs. Returning `None` uses the original, unmodified + result. + """ + pass From 16ba91cd01eede0d7411ae85be12134e449d4c51 Mon Sep 17 00:00:00 2001 From: Che Liu Date: Thu, 10 Jul 2025 17:24:59 -0700 Subject: [PATCH 2/3] feat: Implement PluginService for registering and executing plugins PluginService takes the registration of plugins, and provide the wrapper utilities to execute all plugins. PiperOrigin-RevId: 781745769 --- src/google/adk/plugins/plugin_manager.py | 265 ++++++++++++++++++ tests/unittests/plugins/test_base_plugin.py | 239 ++++++++++++++++ .../unittests/plugins/test_plugin_manager.py | 250 +++++++++++++++++ 3 files changed, 754 insertions(+) create mode 100644 src/google/adk/plugins/plugin_manager.py create mode 100644 tests/unittests/plugins/test_base_plugin.py create mode 100644 tests/unittests/plugins/test_plugin_manager.py diff --git a/src/google/adk/plugins/plugin_manager.py b/src/google/adk/plugins/plugin_manager.py new file mode 100644 index 00000000..3680c351 --- /dev/null +++ b/src/google/adk/plugins/plugin_manager.py @@ -0,0 +1,265 @@ +# Copyright 2025 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. + +from __future__ import annotations + +import logging +from typing import Any +from typing import List +from typing import Literal +from typing import Optional +from typing import TYPE_CHECKING + +from google.genai import types + +from .base_plugin import BasePlugin + +if TYPE_CHECKING: + from ..agents.base_agent import BaseAgent + from ..agents.callback_context import CallbackContext + from ..agents.invocation_context import InvocationContext + from ..events.event import Event + from ..models.llm_request import LlmRequest + from ..models.llm_response import LlmResponse + from ..tools.base_tool import BaseTool + from ..tools.tool_context import ToolContext + +# A type alias for the names of the available plugin callbacks. +# This helps with static analysis and prevents typos when calling run_callbacks. +PluginCallbackName = Literal[ + "on_user_message_callback", + "before_run_callback", + "after_run_callback", + "on_event_callback", + "before_agent_callback", + "after_agent_callback", + "before_tool_callback", + "after_tool_callback", + "before_model_callback", + "after_model_callback", +] + +logger = logging.getLogger("google_adk." + __name__) + + +class PluginManager: + """Manages the registration and execution of plugins. + + The PluginManager is an internal class that orchestrates the invocation of + plugin callbacks at key points in the SDK's execution lifecycle. It maintains + a list of registered plugins and ensures they are called in the order they + were registered. + + The core execution logic implements an "early exit" strategy: if any plugin + callback returns a non-`None` value, the execution of subsequent plugins for + that specific event is halted, and the returned value is propagated up the + call stack. This allows plugins to short-circuit operations like agent runs, + tool calls, or model requests. + """ + + def __init__(self, plugins: Optional[List[BasePlugin]] = None): + """Initializes the plugin service. + + Args: + plugins: An optional list of plugins to register upon initialization. + """ + self.plugins: List[BasePlugin] = [] + if plugins: + for plugin in plugins: + self.register_plugin(plugin) + + def register_plugin(self, plugin: BasePlugin) -> None: + """Registers a new plugin. + + Args: + plugin: The plugin instance to register. + + Raises: + ValueError: If a plugin with the same name is already registered. + """ + if any(p.name == plugin.name for p in self.plugins): + raise ValueError(f"Plugin with name '{plugin.name}' already registered.") + self.plugins.append(plugin) + logger.info("Plugin '%s' registered.", plugin.name) + + def get_plugin(self, plugin_name: str) -> Optional[BasePlugin]: + """Retrieves a registered plugin by its name. + + Args: + plugin_name: The name of the plugin to retrieve. + + Returns: + The plugin instance if found, otherwise `None`. + """ + return next((p for p in self.plugins if p.name == plugin_name), None) + + async def run_on_user_message_callback( + self, + *, + user_message: types.Content, + invocation_context: InvocationContext, + ) -> Optional[types.Content]: + """Runs the `on_user_message_callback` for all plugins.""" + return await self._run_callbacks( + "on_user_message_callback", + user_message=user_message, + invocation_context=invocation_context, + ) + + async def run_before_run_callback( + self, *, invocation_context: InvocationContext + ) -> Optional[types.Content]: + """Runs the `before_run_callback` for all plugins.""" + return await self._run_callbacks( + "before_run_callback", invocation_context=invocation_context + ) + + async def run_after_run_callback( + self, *, invocation_context: InvocationContext + ) -> Optional[None]: + """Runs the `after_run_callback` for all plugins.""" + return await self._run_callbacks( + "after_run_callback", invocation_context=invocation_context + ) + + async def run_on_event_callback( + self, *, invocation_context: InvocationContext, event: Event + ) -> Optional[Event]: + """Runs the `on_event_callback` for all plugins.""" + return await self._run_callbacks( + "on_event_callback", + invocation_context=invocation_context, + event=event, + ) + + async def run_before_agent_callback( + self, *, agent: BaseAgent, callback_context: CallbackContext + ) -> Optional[types.Content]: + """Runs the `before_agent_callback` for all plugins.""" + return await self._run_callbacks( + "before_agent_callback", + agent=agent, + callback_context=callback_context, + ) + + async def run_after_agent_callback( + self, *, agent: BaseAgent, callback_context: CallbackContext + ) -> Optional[types.Content]: + """Runs the `after_agent_callback` for all plugins.""" + return await self._run_callbacks( + "after_agent_callback", + agent=agent, + callback_context=callback_context, + ) + + async def run_before_tool_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + ) -> Optional[dict]: + """Runs the `before_tool_callback` for all plugins.""" + return await self._run_callbacks( + "before_tool_callback", + tool=tool, + tool_args=tool_args, + tool_context=tool_context, + ) + + async def run_after_tool_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + result: dict, + ) -> Optional[dict]: + """Runs the `after_tool_callback` for all plugins.""" + return await self._run_callbacks( + "after_tool_callback", + tool=tool, + tool_args=tool_args, + tool_context=tool_context, + result=result, + ) + + async def run_before_model_callback( + self, *, callback_context: CallbackContext, llm_request: LlmRequest + ) -> Optional[LlmResponse]: + """Runs the `before_model_callback` for all plugins.""" + return await self._run_callbacks( + "before_model_callback", + callback_context=callback_context, + llm_request=llm_request, + ) + + async def run_after_model_callback( + self, *, callback_context: CallbackContext, llm_response: LlmResponse + ) -> Optional[LlmResponse]: + """Runs the `after_model_callback` for all plugins.""" + return await self._run_callbacks( + "after_model_callback", + callback_context=callback_context, + llm_response=llm_response, + ) + + async def _run_callbacks( + self, callback_name: PluginCallbackName, **kwargs: Any + ) -> Optional[Any]: + """Executes a specific callback for all registered plugins. + + This private method iterates through the plugins and calls the specified + callback method on each one, passing the provided keyword arguments. + + The execution stops as soon as a plugin's callback returns a non-`None` + value. This "early exit" value is then returned by this method. If all + plugins are executed and all return `None`, this method also returns `None`. + + Args: + callback_name: The name of the callback method to execute. + **kwargs: Keyword arguments to be passed to the callback method. + + Returns: + The first non-`None` value returned by a plugin callback, or `None` if + all callbacks return `None`. + + Raises: + RuntimeError: If a plugin encounters an unhandled exception during + execution. The original exception is chained. + """ + for plugin in self.plugins: + # Each plugin might not implement all callbacks. The base class provides + # default `pass` implementations, so `getattr` will always succeed. + callback_method = getattr(plugin, callback_name) + try: + result = await callback_method(**kwargs) + if result is not None: + # Early exit: A plugin has returned a value. We stop + # processing further plugins and return this value immediately. + logger.debug( + "Plugin '%s' returned a value for callback '%s', exiting early.", + plugin.name, + callback_name, + ) + return result + except Exception as e: + error_message = ( + f"Error in plugin '{plugin.name}' during '{callback_name}'" + f" callback: {e}" + ) + logger.error(error_message, exc_info=True) + raise RuntimeError(error_message) from e + + return None diff --git a/tests/unittests/plugins/test_base_plugin.py b/tests/unittests/plugins/test_base_plugin.py new file mode 100644 index 00000000..04b1c3e9 --- /dev/null +++ b/tests/unittests/plugins/test_base_plugin.py @@ -0,0 +1,239 @@ +# Copyright 2025 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. + +from __future__ import annotations + +from unittest.mock import Mock + +from google.adk.agents.base_agent import BaseAgent +from google.adk.agents.callback_context import CallbackContext +from google.adk.agents.invocation_context import InvocationContext +from google.adk.events.event import Event +from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse +from google.adk.plugins.base_plugin import BasePlugin +from google.adk.tools.base_tool import BaseTool +from google.adk.tools.tool_context import ToolContext +from google.genai import types +import pytest + + +class TestablePlugin(BasePlugin): + __test__ = False + """A concrete implementation of BasePlugin for testing purposes.""" + pass + + +class FullOverridePlugin(BasePlugin): + __test__ = False + + """A plugin that overrides every single callback method for testing.""" + + def __init__(self, name: str = "full_override"): + super().__init__(name) + + async def on_user_message_callback(self, **kwargs) -> str: + return "overridden_on_user_message" + + async def before_run_callback(self, **kwargs) -> str: + return "overridden_before_run" + + async def after_run_callback(self, **kwargs) -> str: + return "overridden_after_run" + + async def on_event_callback(self, **kwargs) -> str: + return "overridden_on_event" + + async def before_agent_callback(self, **kwargs) -> str: + return "overridden_before_agent" + + async def after_agent_callback(self, **kwargs) -> str: + return "overridden_after_agent" + + async def before_tool_callback(self, **kwargs) -> str: + return "overridden_before_tool" + + async def after_tool_callback(self, **kwargs) -> str: + return "overridden_after_tool" + + async def before_model_callback(self, **kwargs) -> str: + return "overridden_before_model" + + async def after_model_callback(self, **kwargs) -> str: + return "overridden_after_model" + + +def test_base_plugin_initialization(): + """Tests that a plugin is initialized with the correct name.""" + plugin_name = "my_test_plugin" + plugin = TestablePlugin(name=plugin_name) + assert plugin.name == plugin_name + + +@pytest.mark.asyncio +async def test_base_plugin_default_callbacks_return_none(): + """Tests that the default (non-overridden) callbacks in BasePlugin exist + + and return None as expected. + """ + plugin = TestablePlugin(name="default_plugin") + + # Mocking all necessary context objects + mock_context = Mock() + mock_user_message = Mock() + + # The default implementations should do nothing and return None. + assert ( + await plugin.on_user_message_callback( + user_message=mock_user_message, + invocation_context=mock_context, + ) + is None + ) + assert ( + await plugin.before_run_callback(invocation_context=mock_context) is None + ) + assert ( + await plugin.after_run_callback(invocation_context=mock_context) is None + ) + assert ( + await plugin.on_event_callback( + invocation_context=mock_context, event=mock_context + ) + is None + ) + assert ( + await plugin.before_agent_callback( + agent=mock_context, callback_context=mock_context + ) + is None + ) + assert ( + await plugin.after_agent_callback( + agent=mock_context, callback_context=mock_context + ) + is None + ) + assert ( + await plugin.before_tool_callback( + tool=mock_context, tool_args={}, tool_context=mock_context + ) + is None + ) + assert ( + await plugin.after_tool_callback( + tool=mock_context, tool_args={}, tool_context=mock_context, result={} + ) + is None + ) + assert ( + await plugin.before_model_callback( + callback_context=mock_context, llm_request=mock_context + ) + is None + ) + assert ( + await plugin.after_model_callback( + callback_context=mock_context, llm_response=mock_context + ) + is None + ) + + +@pytest.mark.asyncio +async def test_base_plugin_all_callbacks_can_be_overridden(): + """Verifies that a user can create a subclass of BasePlugin and that all + + overridden methods are correctly called. + """ + plugin = FullOverridePlugin() + + # Create mock objects for all required arguments. We don't need real + # objects, just placeholders to satisfy the method signatures. + mock_user_message = Mock(spec=types.Content) + mock_invocation_context = Mock(spec=InvocationContext) + mock_callback_context = Mock(spec=CallbackContext) + mock_agent = Mock(spec=BaseAgent) + mock_tool = Mock(spec=BaseTool) + mock_tool_context = Mock(spec=ToolContext) + mock_llm_request = Mock(spec=LlmRequest) + mock_llm_response = Mock(spec=LlmResponse) + mock_event = Mock(spec=Event) + + # Call each method and assert it returns the unique string from the override. + # This proves that the subclass's method was executed. + assert ( + await plugin.on_user_message_callback( + user_message=mock_user_message, + invocation_context=mock_invocation_context, + ) + == "overridden_on_user_message" + ) + assert ( + await plugin.before_run_callback( + invocation_context=mock_invocation_context + ) + == "overridden_before_run" + ) + assert ( + await plugin.after_run_callback( + invocation_context=mock_invocation_context + ) + == "overridden_after_run" + ) + assert ( + await plugin.on_event_callback( + invocation_context=mock_invocation_context, event=mock_event + ) + == "overridden_on_event" + ) + assert ( + await plugin.before_agent_callback( + agent=mock_agent, callback_context=mock_callback_context + ) + == "overridden_before_agent" + ) + assert ( + await plugin.after_agent_callback( + agent=mock_agent, callback_context=mock_callback_context + ) + == "overridden_after_agent" + ) + assert ( + await plugin.before_model_callback( + callback_context=mock_callback_context, llm_request=mock_llm_request + ) + == "overridden_before_model" + ) + assert ( + await plugin.after_model_callback( + callback_context=mock_callback_context, llm_response=mock_llm_response + ) + == "overridden_after_model" + ) + assert ( + await plugin.before_tool_callback( + tool=mock_tool, tool_args={}, tool_context=mock_tool_context + ) + == "overridden_before_tool" + ) + assert ( + await plugin.after_tool_callback( + tool=mock_tool, + tool_args={}, + tool_context=mock_tool_context, + result={}, + ) + == "overridden_after_tool" + ) diff --git a/tests/unittests/plugins/test_plugin_manager.py b/tests/unittests/plugins/test_plugin_manager.py new file mode 100644 index 00000000..76d32a61 --- /dev/null +++ b/tests/unittests/plugins/test_plugin_manager.py @@ -0,0 +1,250 @@ +# Copyright 2025 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. + +"""Unit tests for the PluginManager.""" + +from __future__ import annotations + +from unittest.mock import Mock + +from google.adk.models.llm_response import LlmResponse +from google.adk.plugins.base_plugin import BasePlugin +# Assume the following path to your modules +# You might need to adjust this based on your project structure. +from google.adk.plugins.plugin_manager import PluginCallbackName +from google.adk.plugins.plugin_manager import PluginManager +import pytest + + +# A helper class to use in tests instead of mocks. +# This makes tests more explicit and easier to debug. +class TestPlugin(BasePlugin): + __test__ = False + """ + A test plugin that can be configured to return specific values or raise + exceptions for any callback, and it logs which callbacks were invoked. + """ + + def __init__(self, name: str): + super().__init__(name) + # A log to track the names of callbacks that have been called. + self.call_log: list[PluginCallbackName] = [] + # A map to configure return values for specific callbacks. + self.return_values: dict[PluginCallbackName, any] = {} + # A map to configure exceptions to be raised by specific callbacks. + self.exceptions_to_raise: dict[PluginCallbackName, Exception] = {} + + async def _handle_callback(self, name: PluginCallbackName): + """Generic handler for all callback methods.""" + self.call_log.append(name) + if name in self.exceptions_to_raise: + raise self.exceptions_to_raise[name] + return self.return_values.get(name) + + # Implement all callback methods from the BasePlugin interface. + async def on_user_message_callback(self, **kwargs): + return await self._handle_callback("on_user_message_callback") + + async def before_run_callback(self, **kwargs): + return await self._handle_callback("before_run_callback") + + async def after_run_callback(self, **kwargs): + return await self._handle_callback("after_run_callback") + + async def on_event_callback(self, **kwargs): + return await self._handle_callback("on_event_callback") + + async def before_agent_callback(self, **kwargs): + return await self._handle_callback("before_agent_callback") + + async def after_agent_callback(self, **kwargs): + return await self._handle_callback("after_agent_callback") + + async def before_tool_callback(self, **kwargs): + return await self._handle_callback("before_tool_callback") + + async def after_tool_callback(self, **kwargs): + return await self._handle_callback("after_tool_callback") + + async def before_model_callback(self, **kwargs): + return await self._handle_callback("before_model_callback") + + async def after_model_callback(self, **kwargs): + return await self._handle_callback("after_model_callback") + + +@pytest.fixture +def service() -> PluginManager: + """Provides a clean PluginManager instance for each test.""" + return PluginManager() + + +@pytest.fixture +def plugin1() -> TestPlugin: + """Provides a clean instance of our test plugin named 'plugin1'.""" + return TestPlugin(name="plugin1") + + +@pytest.fixture +def plugin2() -> TestPlugin: + """Provides a clean instance of our test plugin named 'plugin2'.""" + return TestPlugin(name="plugin2") + + +def test_register_and_get_plugin(service: PluginManager, plugin1: TestPlugin): + """Tests successful registration and retrieval of a plugin.""" + service.register_plugin(plugin1) + + assert len(service.plugins) == 1 + assert service.plugins[0] is plugin1 + assert service.get_plugin("plugin1") is plugin1 + + +def test_register_duplicate_plugin_name_raises_value_error( + service: PluginManager, plugin1: TestPlugin +): + """Tests that registering a plugin with a duplicate name raises an error.""" + plugin1_duplicate = TestPlugin(name="plugin1") + service.register_plugin(plugin1) + + with pytest.raises( + ValueError, match="Plugin with name 'plugin1' already registered." + ): + service.register_plugin(plugin1_duplicate) + + +@pytest.mark.asyncio +async def test_early_exit_stops_subsequent_plugins( + service: PluginManager, plugin1: TestPlugin, plugin2: TestPlugin +): + """Tests the core "early exit" logic: if a plugin returns a value, + + subsequent plugins for that callback should not be executed. + """ + # Configure plugin1 to return a value, simulating a cache hit. + mock_response = Mock(spec=LlmResponse) + plugin1.return_values["before_run_callback"] = mock_response + + service.register_plugin(plugin1) + service.register_plugin(plugin2) + + # Execute the callback chain. + result = await service.run_before_run_callback(invocation_context=Mock()) + + # Assert that the final result is the one returned by the first plugin. + assert result is mock_response + # Assert that the first plugin was called. + assert "before_run_callback" in plugin1.call_log + # CRITICAL: Assert that the second plugin was never called. + assert "before_run_callback" not in plugin2.call_log + + +@pytest.mark.asyncio +async def test_normal_flow_all_plugins_are_called( + service: PluginManager, plugin1: TestPlugin, plugin2: TestPlugin +): + """Tests that if no plugin returns a value, all plugins in the chain + + are executed in order. + """ + # By default, plugins are configured to return None. + service.register_plugin(plugin1) + service.register_plugin(plugin2) + + result = await service.run_before_run_callback(invocation_context=Mock()) + + # The final result should be None as no plugin interrupted the flow. + assert result is None + # Both plugins must have been called. + assert "before_run_callback" in plugin1.call_log + assert "before_run_callback" in plugin2.call_log + + +@pytest.mark.asyncio +async def test_plugin_exception_is_wrapped_in_runtime_error( + service: PluginManager, plugin1: TestPlugin +): + """Tests that if a plugin callback raises an exception, the PluginManager + + catches it and raises a descriptive RuntimeError. + """ + # Configure the plugin to raise an error during a specific callback. + original_exception = ValueError("Something went wrong inside the plugin!") + plugin1.exceptions_to_raise["before_run_callback"] = original_exception + service.register_plugin(plugin1) + + with pytest.raises(RuntimeError) as excinfo: + await service.run_before_run_callback(invocation_context=Mock()) + + # Check that the error message is informative. + assert "Error in plugin 'plugin1'" in str(excinfo.value) + assert "before_run_callback" in str(excinfo.value) + # Check that the original exception is chained for better tracebacks. + assert excinfo.value.__cause__ is original_exception + + +@pytest.mark.asyncio +async def test_all_callbacks_are_supported( + service: PluginManager, plugin1: TestPlugin +): + """Tests that all callbacks defined in the BasePlugin interface are supported + + by the PluginManager. + """ + service.register_plugin(plugin1) + mock_context = Mock() + mock_user_message = Mock() + + # Test all callbacks + await service.run_on_user_message_callback( + user_message=mock_user_message, invocation_context=mock_context + ) + await service.run_before_run_callback(invocation_context=mock_context) + await service.run_after_run_callback(invocation_context=mock_context) + await service.run_on_event_callback( + invocation_context=mock_context, event=mock_context + ) + await service.run_before_agent_callback( + agent=mock_context, callback_context=mock_context + ) + await service.run_after_agent_callback( + agent=mock_context, callback_context=mock_context + ) + await service.run_before_tool_callback( + tool=mock_context, tool_args={}, tool_context=mock_context + ) + await service.run_after_tool_callback( + tool=mock_context, tool_args={}, tool_context=mock_context, result={} + ) + await service.run_before_model_callback( + callback_context=mock_context, llm_request=mock_context + ) + await service.run_after_model_callback( + callback_context=mock_context, llm_response=mock_context + ) + + # Verify all callbacks were logged + expected_callbacks = [ + "on_user_message_callback", + "before_run_callback", + "after_run_callback", + "on_event_callback", + "before_agent_callback", + "after_agent_callback", + "before_tool_callback", + "after_tool_callback", + "before_model_callback", + "after_model_callback", + ] + assert set(plugin1.call_log) == set(expected_callbacks) From 162228d208dca39550a75221030edf9876bf8e3a Mon Sep 17 00:00:00 2001 From: Che Liu Date: Thu, 10 Jul 2025 17:27:52 -0700 Subject: [PATCH 3/3] feat: Integrating Plugin with ADK This change integrates the plugin system with ADK. PluginManager is attached to the invocation context similar to session/artifact/memory. It includes integrations with following ADK internal callbacks: * App callbacks: Integrated in the BaseRunner class, in run_async and run_live * On Message callbacks: Integrated in the BaseRunner class, triggers on run_async. * Agent callbacks: Integrated in the BaseAgent class. Leveraging the existing *callback functions * Model callbacks: Integrating in the base_llm_flow. * Tool callbacks: Integrated in functions.py, wrapped around the code for agent tool_callbacks Sample code to use plugins: ```python # Add plugins to Runner runner = Runner( app_name="my-app", agent=root_agent, artifact_service=artifact_service, session_service=session_service, memory_service=memory_service, plugins=[ MySamplePlugin(), LoggingPlugin(), ], ) ``` PiperOrigin-RevId: 781746586 --- src/google/adk/agents/base_agent.py | 124 +++++++++++------ src/google/adk/agents/invocation_context.py | 4 + .../adk/flows/llm_flows/base_llm_flow.py | 54 +++++--- src/google/adk/flows/llm_flows/functions.py | 81 +++++++---- src/google/adk/runners.py | 96 ++++++++++++- tests/unittests/agents/test_base_agent.py | 99 +++++++++++++- .../unittests/agents/test_langgraph_agent.py | 2 + .../llm_flows/test_plugin_model_callbacks.py | 128 ++++++++++++++++++ .../llm_flows/test_plugin_tool_callbacks.py | 128 ++++++++++++++++++ tests/unittests/test_runners.py | 122 +++++++++++++++++ tests/unittests/testing_utils.py | 10 +- 11 files changed, 759 insertions(+), 89 deletions(-) create mode 100644 tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py create mode 100644 tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py index c474c178..645256fd 100644 --- a/src/google/adk/agents/base_agent.py +++ b/src/google/adk/agents/base_agent.py @@ -227,11 +227,18 @@ class BaseAgent(BaseModel): """ with tracer.start_as_current_span(f'agent_run [{self.name}]'): ctx = self._create_invocation_context(parent_context) - # TODO(hangfei): support before/after_agent_callback + + if event := await self.__handle_before_agent_callback(ctx): + yield event + if ctx.end_invocation: + return async for event in self._run_live_impl(ctx): yield event + if event := await self.__handle_after_agent_callback(ctx): + yield event + async def _run_async_impl( self, ctx: InvocationContext ) -> AsyncGenerator[Event, None]: @@ -335,73 +342,99 @@ class BaseAgent(BaseModel): ) -> Optional[Event]: """Runs the before_agent_callback if it exists. + Args: + ctx: InvocationContext, the invocation context for this agent. + Returns: Optional[Event]: an event if callback provides content or changed state. """ - ret_event = None - - if not self.canonical_before_agent_callbacks: - return ret_event - callback_context = CallbackContext(ctx) - for callback in self.canonical_before_agent_callbacks: - before_agent_callback_content = callback( - callback_context=callback_context - ) - if inspect.isawaitable(before_agent_callback_content): - before_agent_callback_content = await before_agent_callback_content - if before_agent_callback_content: - ret_event = Event( - invocation_id=ctx.invocation_id, - author=self.name, - branch=ctx.branch, - content=before_agent_callback_content, - actions=callback_context._event_actions, + # Run callbacks from the plugins. + before_agent_callback_content = ( + await ctx.plugin_manager.run_before_agent_callback( + agent=self, callback_context=callback_context ) - ctx.end_invocation = True - return ret_event + ) + + # If no overrides are provided from the plugins, further run the canonical + # callbacks. + if ( + not before_agent_callback_content + and self.canonical_before_agent_callbacks + ): + for callback in self.canonical_before_agent_callbacks: + before_agent_callback_content = callback( + callback_context=callback_context + ) + if inspect.isawaitable(before_agent_callback_content): + before_agent_callback_content = await before_agent_callback_content + if before_agent_callback_content: + break + + # Process the override content if exists, and further process the state + # change if exists. + if before_agent_callback_content: + ret_event = Event( + invocation_id=ctx.invocation_id, + author=self.name, + branch=ctx.branch, + content=before_agent_callback_content, + actions=callback_context._event_actions, + ) + ctx.end_invocation = True + return ret_event if callback_context.state.has_delta(): - ret_event = Event( + return Event( invocation_id=ctx.invocation_id, author=self.name, branch=ctx.branch, actions=callback_context._event_actions, ) - return ret_event + return None async def __handle_after_agent_callback( self, invocation_context: InvocationContext ) -> Optional[Event]: """Runs the after_agent_callback if it exists. + Args: + invocation_context: InvocationContext, the invocation context for this + agent. + Returns: Optional[Event]: an event if callback provides content or changed state. """ - ret_event = None - - if not self.canonical_after_agent_callbacks: - return ret_event callback_context = CallbackContext(invocation_context) - for callback in self.canonical_after_agent_callbacks: - after_agent_callback_content = callback(callback_context=callback_context) - if inspect.isawaitable(after_agent_callback_content): - after_agent_callback_content = await after_agent_callback_content - if after_agent_callback_content: - ret_event = Event( - invocation_id=invocation_context.invocation_id, - author=self.name, - branch=invocation_context.branch, - content=after_agent_callback_content, - actions=callback_context._event_actions, + # Run callbacks from the plugins. + after_agent_callback_content = ( + await invocation_context.plugin_manager.run_after_agent_callback( + agent=self, callback_context=callback_context ) - return ret_event + ) - if callback_context.state.has_delta(): + # If no overrides are provided from the plugins, further run the canonical + # callbacks. + if ( + not after_agent_callback_content + and self.canonical_after_agent_callbacks + ): + for callback in self.canonical_after_agent_callbacks: + after_agent_callback_content = callback( + callback_context=callback_context + ) + if inspect.isawaitable(after_agent_callback_content): + after_agent_callback_content = await after_agent_callback_content + if after_agent_callback_content: + break + + # Process the override content if exists, and further process the state + # change if exists. + if after_agent_callback_content: ret_event = Event( invocation_id=invocation_context.invocation_id, author=self.name, @@ -409,8 +442,17 @@ class BaseAgent(BaseModel): content=after_agent_callback_content, actions=callback_context._event_actions, ) + return ret_event - return ret_event + if callback_context.state.has_delta(): + return Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + branch=invocation_context.branch, + content=after_agent_callback_content, + actions=callback_context._event_actions, + ) + return None @override def model_post_init(self, __context: Any) -> None: diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 765f22a2..8af54743 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -24,6 +24,7 @@ from pydantic import ConfigDict from ..artifacts.base_artifact_service import BaseArtifactService from ..auth.credential_service.base_credential_service import BaseCredentialService from ..memory.base_memory_service import BaseMemoryService +from ..plugins.plugin_manager import PluginManager from ..sessions.base_session_service import BaseSessionService from ..sessions.session import Session from .active_streaming_tool import ActiveStreamingTool @@ -153,6 +154,9 @@ class InvocationContext(BaseModel): run_config: Optional[RunConfig] = None """Configurations for live agents under this invocation.""" + plugin_manager: PluginManager = PluginManager() + """The manager for keeping track of plugins in this invocation.""" + _invocation_cost_manager: _InvocationCostManager = _InvocationCostManager() """A container to keep track of different kinds of costs incurred as a part of this invocation. diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 16927063..bb2759f6 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -565,21 +565,32 @@ class BaseLlmFlow(ABC): if not isinstance(agent, LlmAgent): return - if not agent.canonical_before_model_callbacks: - return - callback_context = CallbackContext( invocation_context, event_actions=model_response_event.actions ) + # First run callbacks from the plugins. + callback_response = ( + await invocation_context.plugin_manager.run_before_model_callback( + callback_context=callback_context, + llm_request=llm_request, + ) + ) + if callback_response: + return callback_response + + # If no overrides are provided from the plugins, further run the canonical + # callbacks. + if not agent.canonical_before_model_callbacks: + return for callback in agent.canonical_before_model_callbacks: - before_model_callback_content = callback( + callback_response = callback( callback_context=callback_context, llm_request=llm_request ) - if inspect.isawaitable(before_model_callback_content): - before_model_callback_content = await before_model_callback_content - if before_model_callback_content: - return before_model_callback_content + if inspect.isawaitable(callback_response): + callback_response = await callback_response + if callback_response: + return callback_response async def _handle_after_model_callback( self, @@ -593,21 +604,32 @@ class BaseLlmFlow(ABC): if not isinstance(agent, LlmAgent): return - if not agent.canonical_after_model_callbacks: - return - callback_context = CallbackContext( invocation_context, event_actions=model_response_event.actions ) + # First run callbacks from the plugins. + callback_response = ( + await invocation_context.plugin_manager.run_after_model_callback( + callback_context=CallbackContext(invocation_context), + llm_response=llm_response, + ) + ) + if callback_response: + return callback_response + + # If no overrides are provided from the plugins, further run the canonical + # callbacks. + if not agent.canonical_after_model_callbacks: + return for callback in agent.canonical_after_model_callbacks: - after_model_callback_content = callback( + callback_response = callback( callback_context=callback_context, llm_response=llm_response ) - if inspect.isawaitable(after_model_callback_content): - after_model_callback_content = await after_model_callback_content - if after_model_callback_content: - return after_model_callback_content + if inspect.isawaitable(callback_response): + callback_response = await callback_response + if callback_response: + return callback_response def _finalize_model_response_event( self, diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 5c690f1f..c64b7cec 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -153,37 +153,67 @@ async def handle_function_calls_async( # do not use "args" as the variable name, because it is a reserved keyword # in python debugger. function_args = function_call.args or {} - function_response: Optional[dict] = None - for callback in agent.canonical_before_tool_callbacks: - function_response = callback( - tool=tool, args=function_args, tool_context=tool_context - ) - if inspect.isawaitable(function_response): - function_response = await function_response - if function_response: - break + # Step 1: Check if plugin before_tool_callback overrides the function + # response. + function_response = ( + await invocation_context.plugin_manager.run_before_tool_callback( + tool=tool, tool_args=function_args, tool_context=tool_context + ) + ) - if not function_response: + # Step 2: If no overrides are provided from the plugins, further run the + # canonical callback. + if function_response is None: + for callback in agent.canonical_before_tool_callbacks: + function_response = callback( + tool=tool, args=function_args, tool_context=tool_context + ) + if inspect.isawaitable(function_response): + function_response = await function_response + if function_response: + break + + # Step 3: Otherwise, proceed calling the tool normally. + if function_response is None: function_response = await __call_tool_async( tool, args=function_args, tool_context=tool_context ) - for callback in agent.canonical_after_tool_callbacks: - altered_function_response = callback( - tool=tool, - args=function_args, - tool_context=tool_context, - tool_response=function_response, - ) - if inspect.isawaitable(altered_function_response): - altered_function_response = await altered_function_response - if altered_function_response is not None: - function_response = altered_function_response - break + # Step 4: Check if plugin after_tool_callback overrides the function + # response. + altered_function_response = ( + await invocation_context.plugin_manager.run_after_tool_callback( + tool=tool, + tool_args=function_args, + tool_context=tool_context, + result=function_response, + ) + ) + + # Step 5: If no overrides are provided from the plugins, further run the + # canonical after_tool_callbacks. + if altered_function_response is None: + for callback in agent.canonical_after_tool_callbacks: + altered_function_response = callback( + tool=tool, + args=function_args, + tool_context=tool_context, + tool_response=function_response, + ) + if inspect.isawaitable(altered_function_response): + altered_function_response = await altered_function_response + if altered_function_response: + break + + # Step 6: If alternative response exists from after_tool_callback, use it + # instead of the original function response. + if altered_function_response is not None: + function_response = altered_function_response if tool.is_long_running: - # Allow long running function to return None to not provide function response. + # Allow long running function to return None to not provide function + # response. if not function_response: continue @@ -264,6 +294,7 @@ async def handle_function_calls_live( # ) # if new_response: # function_response = new_response + altered_function_response = None if agent.after_tool_callback: altered_function_response = agent.after_tool_callback( tool=tool, @@ -273,8 +304,8 @@ async def handle_function_calls_live( ) if inspect.isawaitable(altered_function_response): altered_function_response = await altered_function_response - if altered_function_response is not None: - function_response = altered_function_response + if altered_function_response is not None: + function_response = altered_function_response if tool.is_long_running: # Allow async function to return None to not provide function response. diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index a95f1081..ccca19e0 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -18,7 +18,9 @@ import asyncio import logging import queue from typing import AsyncGenerator +from typing import Callable from typing import Generator +from typing import List from typing import Optional import warnings @@ -40,6 +42,8 @@ from .flows.llm_flows.functions import find_matching_function_call from .memory.base_memory_service import BaseMemoryService from .memory.in_memory_memory_service import InMemoryMemoryService from .platform.thread import create_thread +from .plugins.base_plugin import BasePlugin +from .plugins.plugin_manager import PluginManager from .sessions.base_session_service import BaseSessionService from .sessions.in_memory_session_service import InMemorySessionService from .sessions.session import Session @@ -60,6 +64,7 @@ class Runner: app_name: The application name of the runner. agent: The root agent to run. artifact_service: The artifact service for the runner. + plugin_manager: The plugin manager for the runner. session_service: The session service for the runner. memory_service: The memory service for the runner. """ @@ -70,6 +75,8 @@ class Runner: """The root agent to run.""" artifact_service: Optional[BaseArtifactService] = None """The artifact service for the runner.""" + plugin_manager: PluginManager + """The plugin manager for the runner.""" session_service: BaseSessionService """The session service for the runner.""" memory_service: Optional[BaseMemoryService] = None @@ -82,6 +89,7 @@ class Runner: *, app_name: str, agent: BaseAgent, + plugins: Optional[List[BasePlugin]] = None, artifact_service: Optional[BaseArtifactService] = None, session_service: BaseSessionService, memory_service: Optional[BaseMemoryService] = None, @@ -102,6 +110,7 @@ class Runner: self.session_service = session_service self.memory_service = memory_service self.credential_service = credential_service + self.plugin_manager = PluginManager(plugins=plugins) def run( self, @@ -191,6 +200,15 @@ class Runner: ) root_agent = self.agent + # Modify user message before execution. + modified_user_message = ( + await invocation_context.plugin_manager.run_on_user_message_callback( + invocation_context=invocation_context, user_message=new_message + ) + ) + if modified_user_message is not None: + new_message = modified_user_message + if new_message: await self._append_new_message_to_session( session, @@ -200,10 +218,64 @@ class Runner: ) invocation_context.agent = self._find_agent_to_run(session, root_agent) - async for event in invocation_context.agent.run_async(invocation_context): + + async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]: + async for event in ctx.agent.run_async(ctx): + yield event + + async for event in self._exec_with_plugin( + invocation_context, session, execute + ): + yield event + + async def _exec_with_plugin( + self, + invocation_context: InvocationContext, + session: Session, + execute_fn: Callable[[InvocationContext], AsyncGenerator[Event, None]], + ) -> AsyncGenerator[Event, None]: + """Wraps execution with plugin callbacks. + + Args: + invocation_context: The invocation context + session: The current session + execute_fn: A callable that returns an AsyncGenerator of Events + + Yields: + Events from the execution, including any generated by plugins + """ + + plugin_manager = invocation_context.plugin_manager + + # Step 1: Run the before_run callbacks to see if we should early exit. + early_exit_result = await plugin_manager.run_before_run_callback( + invocation_context=invocation_context + ) + if isinstance(early_exit_result, Event): + await self.session_service.append_event( + session=session, + event=Event( + invocation_id=invocation_context.invocation_id, + author='model', + content=early_exit_result, + ), + ) + yield early_exit_result + else: + # Step 2: Otherwise continue with normal execution + async for event in execute_fn(invocation_context): if not event.partial: await self.session_service.append_event(session=session, event=event) - yield event + # Step 3: Run the on_event callbacks to optionally modify the event. + modified_event = await plugin_manager.run_on_event_callback( + invocation_context=invocation_context, event=event + ) + yield (modified_event if modified_event else event) + + # Step 4: Run the after_run callbacks to optionally modify the context. + await plugin_manager.run_after_run_callback( + invocation_context=invocation_context + ) async def _append_new_message_to_session( self, @@ -345,8 +417,14 @@ class Runner: invocation_context.active_streaming_tools[tool.__name__] = ( active_streaming_tool ) - async for event in invocation_context.agent.run_live(invocation_context): - await self.session_service.append_event(session=session, event=event) + + async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]: + async for event in ctx.agent.run_live(ctx): + yield event + + async for event in self._exec_with_plugin( + invocation_context, session, execute + ): yield event def _find_agent_to_run( @@ -450,6 +528,7 @@ class Runner: session_service=self.session_service, memory_service=self.memory_service, credential_service=self.credential_service, + plugin_manager=self.plugin_manager, invocation_id=invocation_id, agent=self.agent, session=session, @@ -538,7 +617,13 @@ class InMemoryRunner(Runner): session service for the runner. """ - def __init__(self, agent: BaseAgent, *, app_name: str = 'InMemoryRunner'): + def __init__( + self, + agent: BaseAgent, + *, + app_name: str = 'InMemoryRunner', + plugins: Optional[list[BasePlugin]] = None, + ): """Initializes the InMemoryRunner. Args: @@ -551,6 +636,7 @@ class InMemoryRunner(Runner): app_name=app_name, agent=agent, artifact_service=InMemoryArtifactService(), + plugins=plugins, session_service=self._in_memory_session_service, memory_service=InMemoryMemoryService(), ) diff --git a/tests/unittests/agents/test_base_agent.py b/tests/unittests/agents/test_base_agent.py index 25aca8ff..4f8bd770 100644 --- a/tests/unittests/agents/test_base_agent.py +++ b/tests/unittests/agents/test_base_agent.py @@ -26,6 +26,8 @@ from google.adk.agents.base_agent import BaseAgent from google.adk.agents.callback_context import CallbackContext from google.adk.agents.invocation_context import InvocationContext from google.adk.events import Event +from google.adk.plugins.base_plugin import BasePlugin +from google.adk.plugins.plugin_manager import PluginManager from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.genai import types import pytest @@ -83,6 +85,35 @@ async def _async_after_agent_callback_append_agent_reply( ) +class MockPlugin(BasePlugin): + before_agent_text = 'before_agent_text from MockPlugin' + after_agent_text = 'after_agent_text from MockPlugin' + + def __init__(self, name='mock_plugin'): + self.name = name + self.enable_before_agent_callback = False + self.enable_after_agent_callback = False + + async def before_agent_callback( + self, *, agent: BaseAgent, callback_context: CallbackContext + ) -> Optional[types.Content]: + if not self.enable_before_agent_callback: + return None + return types.Content(parts=[types.Part(text=self.before_agent_text)]) + + async def after_agent_callback( + self, *, agent: BaseAgent, callback_context: CallbackContext + ) -> Optional[types.Content]: + if not self.enable_after_agent_callback: + return None + return types.Content(parts=[types.Part(text=self.after_agent_text)]) + + +@pytest.fixture +def mock_plugin(): + return MockPlugin() + + class _IncompleteAgent(BaseAgent): pass @@ -113,7 +144,10 @@ class _TestingAgent(BaseAgent): async def _create_parent_invocation_context( - test_name: str, agent: BaseAgent, branch: Optional[str] = None + test_name: str, + agent: BaseAgent, + branch: Optional[str] = None, + plugins: list[BasePlugin] = [], ) -> InvocationContext: session_service = InMemorySessionService() session = await session_service.create_session( @@ -125,6 +159,7 @@ async def _create_parent_invocation_context( agent=agent, session=session, session_service=session_service, + plugin_manager=PluginManager(plugins=plugins), ) @@ -190,6 +225,36 @@ async def test_run_async_before_agent_callback_noop( spy_run_async_impl.assert_called_once() +@pytest.mark.asyncio +async def test_run_async_before_agent_callback_use_plugin( + request: pytest.FixtureRequest, + mocker: pytest_mock.MockerFixture, + mock_plugin: MockPlugin, +): + """Test that the before agent callback uses the plugin response if both plugin callback and canonical agent callbacks are present.""" + # Arrange + agent = _TestingAgent( + name=f'{request.function.__name__}_test_agent', + before_agent_callback=_before_agent_callback_bypass_agent, + ) + parent_ctx = await _create_parent_invocation_context( + request.function.__name__, agent, plugins=[mock_plugin] + ) + mock_plugin.enable_before_agent_callback = True + spy_run_async_impl = mocker.spy(agent, BaseAgent._run_async_impl.__name__) + spy_before_agent_callback = mocker.spy(agent, 'before_agent_callback') + + # Act + events = [e async for e in agent.run_async(parent_ctx)] + + # Assert + spy_before_agent_callback.assert_not_called() + spy_run_async_impl.assert_not_called() + + assert len(events) == 1 + assert events[0].content.parts[0].text == MockPlugin.before_agent_text + + @pytest.mark.asyncio async def test_run_async_with_async_before_agent_callback_noop( request: pytest.FixtureRequest, @@ -486,6 +551,34 @@ async def test_after_agent_callbacks_chain( mock_cb.assert_called(expected_calls_count) +@pytest.mark.asyncio +async def test_run_async_after_agent_callback_use_plugin( + request: pytest.FixtureRequest, + mocker: pytest_mock.MockerFixture, + mock_plugin: MockPlugin, +): + # Arrange + agent = _TestingAgent( + name=f'{request.function.__name__}_test_agent', + after_agent_callback=_after_agent_callback_noop, + ) + mock_plugin.enable_after_agent_callback = True + parent_ctx = await _create_parent_invocation_context( + request.function.__name__, agent, plugins=[mock_plugin] + ) + spy_after_agent_callback = mocker.spy(agent, 'after_agent_callback') + + # Act + events = [e async for e in agent.run_async(parent_ctx)] + + # Assert + spy_after_agent_callback.assert_not_called() + # The first event is regular model response, the second event is + # after_agent_callback response. + assert len(events) == 2 + assert events[1].content.parts[0].text == mock_plugin.after_agent_text + + @pytest.mark.asyncio async def test_run_async_after_agent_callback_noop( request: pytest.FixtureRequest, @@ -757,3 +850,7 @@ def test_set_parent_agent_for_sub_agent_twice( name=f'{request.function.__name__}_parent_2', sub_agents=[sub_agent], ) + + +if __name__ == '__main__': + pytest.main([__file__]) diff --git a/tests/unittests/agents/test_langgraph_agent.py b/tests/unittests/agents/test_langgraph_agent.py index aa6cb6ac..4e5d3481 100644 --- a/tests/unittests/agents/test_langgraph_agent.py +++ b/tests/unittests/agents/test_langgraph_agent.py @@ -17,6 +17,7 @@ from unittest.mock import MagicMock from google.adk.agents.invocation_context import InvocationContext from google.adk.agents.langgraph_agent import LangGraphAgent from google.adk.events import Event +from google.adk.plugins.plugin_manager import PluginManager from google.genai import types from langchain_core.messages import AIMessage from langchain_core.messages import HumanMessage @@ -169,6 +170,7 @@ async def test_langgraph_agent( mock_session.events = events_list mock_parent_context.invocation_id = "test_invocation_id" mock_parent_context.model_copy.return_value = mock_parent_context + mock_parent_context.plugin_manager = PluginManager(plugins=[]) weather_agent = LangGraphAgent( name="weather_agent", diff --git a/tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py b/tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py new file mode 100644 index 00000000..b9b2dec3 --- /dev/null +++ b/tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py @@ -0,0 +1,128 @@ +# Copyright 2025 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. + +from typing import Optional + +from google.adk.agents import Agent +from google.adk.agents.callback_context import CallbackContext +from google.adk.models import LlmRequest +from google.adk.models import LlmResponse +from google.adk.plugins.base_plugin import BasePlugin +from google.genai import types +import pytest + +from ... import testing_utils + + +class MockPlugin(BasePlugin): + before_model_text = 'before_model_text from MockPlugin' + after_model_text = 'after_model_text from MockPlugin' + + def __init__(self, name='mock_plugin'): + self.name = name + self.enable_before_model_callback = False + self.enable_after_model_callback = False + self.before_model_response = LlmResponse( + content=testing_utils.ModelContent( + [types.Part.from_text(text=self.before_model_text)] + ) + ) + self.after_model_response = LlmResponse( + content=testing_utils.ModelContent( + [types.Part.from_text(text=self.after_model_text)] + ) + ) + + async def before_model_callback( + self, *, callback_context: CallbackContext, llm_request: LlmRequest + ) -> Optional[LlmResponse]: + if not self.enable_before_model_callback: + return None + return self.before_model_response + + async def after_model_callback( + self, *, callback_context: CallbackContext, llm_response: LlmResponse + ) -> Optional[LlmResponse]: + if not self.enable_after_model_callback: + return None + return self.after_model_response + + +CANONICAL_MODEL_CALLBACK_CONTENT = 'canonical_model_callback_content' + + +def canonical_agent_model_callback(**kwargs) -> Optional[LlmResponse]: + return LlmResponse( + content=testing_utils.ModelContent( + [types.Part.from_text(text=CANONICAL_MODEL_CALLBACK_CONTENT)] + ) + ) + + +@pytest.fixture +def mock_plugin(): + return MockPlugin() + + +def test_before_model_callback_with_plugin(mock_plugin): + """Tests that the model response is overridden by before_model_callback from the plugin.""" + responses = ['model_response'] + mock_model = testing_utils.MockModel.create(responses=responses) + mock_plugin.enable_before_model_callback = True + agent = Agent( + name='root_agent', + model=mock_model, + ) + + runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin]) + assert testing_utils.simplify_events(runner.run('test')) == [ + ('root_agent', mock_plugin.before_model_text), + ] + + +def test_before_model_fallback_canonical_callback(mock_plugin): + """Tests that when plugin returns empty response, the model response is overridden by the canonical agent model callback.""" + responses = ['model_response'] + mock_plugin.enable_before_model_callback = False + mock_model = testing_utils.MockModel.create(responses=responses) + agent = Agent( + name='root_agent', + model=mock_model, + before_model_callback=canonical_agent_model_callback, + ) + + runner = testing_utils.InMemoryRunner(agent) + assert testing_utils.simplify_events(runner.run('test')) == [ + ('root_agent', CANONICAL_MODEL_CALLBACK_CONTENT), + ] + + +def test_before_model_callback_fallback_model(mock_plugin): + """Tests that the model response is executed normally when both plugin and canonical agent model callback return empty response.""" + responses = ['model_response'] + mock_plugin.enable_before_model_callback = False + mock_model = testing_utils.MockModel.create(responses=responses) + agent = Agent( + name='root_agent', + model=mock_model, + ) + + runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin]) + assert testing_utils.simplify_events(runner.run('test')) == [ + ('root_agent', 'model_response'), + ] + + +if __name__ == '__main__': + pytest.main([__file__]) diff --git a/tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py b/tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py new file mode 100644 index 00000000..a79e562a --- /dev/null +++ b/tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py @@ -0,0 +1,128 @@ +# Copyright 2025 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. + +from typing import Any +from typing import Dict +from typing import Optional + +from google.adk.agents import Agent +from google.adk.events.event import Event +from google.adk.flows.llm_flows.functions import handle_function_calls_async +from google.adk.plugins.base_plugin import BasePlugin +from google.adk.tools.base_tool import BaseTool +from google.adk.tools.function_tool import FunctionTool +from google.adk.tools.tool_context import ToolContext +from google.genai import types +import pytest + +from ... import testing_utils + + +class MockPlugin(BasePlugin): + before_tool_response = {"MockPlugin": "before_tool_response from MockPlugin"} + after_tool_response = {"MockPlugin": "after_tool_response from MockPlugin"} + + def __init__(self, name="mock_plugin"): + self.name = name + self.enable_before_tool_callback = False + self.enable_after_tool_callback = False + + async def before_tool_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + ) -> Optional[dict]: + if not self.enable_before_tool_callback: + return None + return self.before_tool_response + + async def after_tool_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + result: dict, + ) -> Optional[dict]: + if not self.enable_after_tool_callback: + return None + return self.after_tool_response + + +@pytest.fixture +def mock_tool(): + def simple_fn(**kwargs) -> Dict[str, Any]: + return {"initial": "response"} + + return FunctionTool(simple_fn) + + +@pytest.fixture +def mock_plugin(): + return MockPlugin() + + +async def invoke_tool_with_plugin(mock_tool, mock_plugin) -> Optional[Event]: + """Invokes a tool with a plugin.""" + model = testing_utils.MockModel.create(responses=[]) + agent = Agent( + name="agent", + model=model, + tools=[mock_tool], + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content="", plugins=[mock_plugin] + ) + # Build function call event + function_call = types.FunctionCall(name=mock_tool.name, args={}) + content = types.Content(parts=[types.Part(function_call=function_call)]) + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=content, + ) + tools_dict = {mock_tool.name: mock_tool} + return await handle_function_calls_async( + invocation_context, + event, + tools_dict, + ) + + +@pytest.mark.asyncio +async def test_async_before_tool_callback(mock_tool, mock_plugin): + mock_plugin.enable_before_tool_callback = True + + result_event = await invoke_tool_with_plugin(mock_tool, mock_plugin) + + assert result_event is not None + part = result_event.content.parts[0] + assert part.function_response.response == mock_plugin.before_tool_response + + +@pytest.mark.asyncio +async def test_async_after_tool_callback(mock_tool, mock_plugin): + mock_plugin.enable_after_tool_callback = True + + result_event = await invoke_tool_with_plugin(mock_tool, mock_plugin) + + assert result_event is not None + part = result_event.content.parts[0] + assert part.function_response.response == mock_plugin.after_tool_response + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index 8d5bd241..61ee4b01 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -15,13 +15,20 @@ from typing import Optional from google.adk.agents.base_agent import BaseAgent +from google.adk.agents.invocation_context import InvocationContext from google.adk.agents.llm_agent import LlmAgent from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService from google.adk.events.event import Event +from google.adk.plugins.base_plugin import BasePlugin from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.sessions.session import Session from google.genai import types +import pytest + +TEST_APP_ID = "test_app" +TEST_USER_ID = "test_user" +TEST_SESSION_ID = "test_session" class MockAgent(BaseAgent): @@ -72,6 +79,51 @@ class MockLlmAgent(LlmAgent): ) +class MockPlugin(BasePlugin): + """Mock plugin for unit testing.""" + + ON_USER_CALLBACK_MSG = ( + "Modified user message ON_USER_CALLBACK_MSG from MockPlugin" + ) + ON_EVENT_CALLBACK_MSG = "Modified event ON_EVENT_CALLBACK_MSG from MockPlugin" + + def __init__(self): + super().__init__(name="mock_plugin") + self.enable_user_message_callback = False + self.enable_event_callback = False + + async def on_user_message_callback( + self, + *, + invocation_context: InvocationContext, + user_message: types.Content, + ) -> Optional[types.Content]: + if not self.enable_user_message_callback: + return None + return types.Content( + role="model", + parts=[types.Part(text=self.ON_USER_CALLBACK_MSG)], + ) + + async def on_event_callback( + self, *, invocation_context: InvocationContext, event: Event + ) -> Optional[Event]: + if not self.enable_event_callback: + return None + return Event( + invocation_id="", + author="", + content=types.Content( + parts=[ + types.Part( + text=self.ON_EVENT_CALLBACK_MSG, + ) + ], + role=event.content.role, + ), + ) + + class TestRunnerFindAgentToRun: """Tests for Runner._find_agent_to_run method.""" @@ -308,3 +360,73 @@ class TestRunnerFindAgentToRun: # MockAgent inherits from BaseAgent, not LlmAgent, so it should return False result = self.runner._is_transferable_across_agent_tree(non_llm_agent) assert result is False + + +class TestRunnerWithPlugins: + """Tests for Runner with plugins.""" + + def setup_method(self): + self.plugin = MockPlugin() + self.session_service = InMemorySessionService() + self.artifact_service = InMemoryArtifactService() + self.root_agent = MockLlmAgent("root_agent") + self.runner = Runner( + app_name="test_app", + agent=MockLlmAgent("test_agent"), + session_service=self.session_service, + artifact_service=self.artifact_service, + plugins=[self.plugin], + ) + + async def run_test(self, original_user_input="Hello") -> list[Event]: + """Prepares the test by creating a session and running the runner.""" + await self.session_service.create_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + events = [] + async for event in self.runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=types.Content( + role="user", parts=[types.Part(text=original_user_input)] + ), + ): + events.append(event) + return events + + @pytest.mark.asyncio + async def test_runner_is_initialized_with_plugins(self): + """Test that the runner is initialized with plugins.""" + await self.run_test() + + assert self.runner.plugin_manager is not None + + @pytest.mark.asyncio + async def test_runner_modifies_user_message_before_execution(self): + """Test that the runner modifies the user message before execution.""" + original_user_input = "original_input" + self.plugin.enable_user_message_callback = True + + await self.run_test(original_user_input=original_user_input) + session = await self.session_service.get_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + generated_event = session.events[0] + modified_user_message = generated_event.content.parts[0].text + + assert modified_user_message == MockPlugin.ON_USER_CALLBACK_MSG + + @pytest.mark.asyncio + async def test_runner_modifies_event_after_execution(self): + """Test that the runner modifies the event after execution.""" + self.plugin.enable_event_callback = True + + events = await self.run_test() + generated_event = events[0] + modified_event_message = generated_event.content.parts[0].text + + assert modified_event_message == MockPlugin.ON_EVENT_CALLBACK_MSG + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/unittests/testing_utils.py b/tests/unittests/testing_utils.py index b1d5ff82..810a6c44 100644 --- a/tests/unittests/testing_utils.py +++ b/tests/unittests/testing_utils.py @@ -30,6 +30,8 @@ from google.adk.models.base_llm import BaseLlm from google.adk.models.base_llm_connection import BaseLlmConnection from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse +from google.adk.plugins.base_plugin import BasePlugin +from google.adk.plugins.plugin_manager import PluginManager from google.adk.runners import InMemoryRunner as AfInMemoryRunner from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService @@ -57,7 +59,10 @@ class ModelContent(types.Content): async def create_invocation_context( - agent: Agent, user_content: str = '', run_config: RunConfig = None + agent: Agent, + user_content: str = '', + run_config: RunConfig = None, + plugins: list[BasePlugin] = [], ): invocation_id = 'test_id' artifact_service = InMemoryArtifactService() @@ -67,6 +72,7 @@ async def create_invocation_context( artifact_service=artifact_service, session_service=session_service, memory_service=memory_service, + plugin_manager=PluginManager(plugins=plugins), invocation_id=invocation_id, agent=agent, session=await session_service.create_session( @@ -165,6 +171,7 @@ class InMemoryRunner: self, root_agent: Union[Agent, LlmAgent], response_modalities: list[str] = None, + plugins: list[BasePlugin] = [], ): self.root_agent = root_agent self.runner = Runner( @@ -173,6 +180,7 @@ class InMemoryRunner: artifact_service=InMemoryArtifactService(), session_service=InMemorySessionService(), memory_service=InMemoryMemoryService(), + plugins=plugins, ) self.session_id = None