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
This commit is contained in:
Che Liu
2025-07-10 17:28:42 -07:00
committed by Copybara-Service
parent 16ba91cd01
commit 162228d208
11 changed files with 759 additions and 89 deletions
@@ -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__])
@@ -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__])