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.