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