google-adk 1.6.1__py3-none-any.whl → 1.7.0__py3-none-any.whl

google/adk/plugins/base_plugin.py

@@ -0,0 +1,317 @@
+# 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")
+
+
+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

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