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/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 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/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/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) 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