feat: Integrating Plugin with ADK

This change integrates the plugin system with ADK. PluginManager is attached to the invocation context similar to session/artifact/memory. It includes integrations with following ADK internal callbacks: * App callbacks: Integrated in the BaseRunner class, in run_async and run_live * On Message callbacks: Integrated in the BaseRunner class, triggers on run_async. * Agent callbacks: Integrated in the BaseAgent class. Leveraging the existing *callback functions * Model callbacks: Integrating in the base_llm_flow. * Tool callbacks: Integrated in functions.py, wrapped around the code for agent tool_callbacks Sample code to use plugins: ```python # Add plugins to Runner runner = Runner( app_name="my-app", agent=root_agent, artifact_service=artifact_service, session_service=session_service, memory_service=memory_service, plugins=[ MySamplePlugin(), LoggingPlugin(), ], ) ``` PiperOrigin-RevId: 781746586

Author: pwwpche

src/google/adk/agents/base_agent.py

@@ -227,11 +227,18 @@ async def run_live(
  """
  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,82 +342,117 @@ async def __handle_before_agent_callback(
  ) -> 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 callback_context.state.has_delta():
+ # 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():
+ 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,
  branch=invocation_context.branch,
  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:

src/google/adk/agents/invocation_context.py

@@ -24,6 +24,7 @@
 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.

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -565,21 +565,32 @@ async def _handle_before_model_callback(
  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 @@ async def _handle_after_model_callback(
  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,

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.