mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
16ba91cd01
commit
162228d208
@@ -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__])
|
||||
Reference in New Issue
Block a user