encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat(config): support callbacks in BaseAgent and LlmAgent

Browse Source 
PiperOrigin-RevId: 783422160
This commit is contained in:
Liang Wu
2025-07-15 12:05:10 -07:00
committed by Copybara-Service
parent 3c6e18fd1c
commit 41f1888116
4 changed files with 255 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/agents/base_agent.py
+25
View File
@@ -43,6 +43,7 @@ from typing_extensions import TypeAlias
from ..events.event import Event
from ..utils.feature_decorator import working_in_progress
from .callback_context import CallbackContext
from .common_configs import CodeConfig
if TYPE_CHECKING:
from .invocation_context import InvocationContext
@@ -510,6 +511,7 @@ class BaseAgent(BaseModel):
The created agent.
"""
from .config_agent_utils import build_sub_agent
from .config_agent_utils import resolve_callbacks
kwargs: Dict[str, Any] = {
'name': config.name,
@@ -523,6 +525,15 @@ class BaseAgent(BaseModel):
)
sub_agents.append(sub_agent)
kwargs['sub_agents'] = sub_agents
if config.before_agent_callbacks:
kwargs['before_agent_callback'] = resolve_callbacks(
config.before_agent_callbacks
)
if config.after_agent_callbacks:
kwargs['after_agent_callback'] = resolve_callbacks(
config.after_agent_callbacks
)
return cls(**kwargs)
@@ -607,3 +618,17 @@ class BaseAgentConfig(BaseModel):
sub_agents: Optional[List[SubAgentConfig]] = None
"""Optional. The sub-agents of the agent."""
before_agent_callbacks: Optional[List[CodeConfig]] = None
"""Optional. The before_agent_callbacks of the agent.
Example:
```
before_agent_callbacks:
- name: my_library.security_callbacks.before_agent_callback
```
"""
after_agent_callbacks: Optional[List[CodeConfig]] = None
"""Optional. The after_agent_callbacks of the agent."""
src/google/adk/agents/config_agent_utils.py
+14 -1
View File
@@ -16,8 +16,8 @@ from __future__ import annotations
import importlib
import os
from pathlib import Path
from typing import Any
from typing import List
import yaml
@@ -169,3 +169,16 @@ def resolve_code_reference(code_config: CodeConfig) -> Any:
return obj(*positional_args, **kwargs)
else:
return obj
@working_in_progress("resolve_callbacks is not ready for use.")
def resolve_callbacks(callbacks_config: List[CodeConfig]) -> Any:
"""Resolve callbacks from configuration.
Args:
callbacks_config: List of callback configurations (CodeConfig objects).
Returns:
List of resolved callback objects.
"""
return [resolve_code_reference(config) for config in callbacks_config]
src/google/adk/agents/config_schemas/AgentConfig.json
+180
View File
@@ -93,6 +93,36 @@
"default": null,
"title": "Sub Agents"
},
"before_agent_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Before Agent Callbacks"
},
"after_agent_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "After Agent Callbacks"
},
"model": {
"anyOf": [
{
@@ -168,6 +198,66 @@
],
"default": null,
"title": "Tools"
},
"before_model_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Before Model Callbacks"
},
"after_model_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "After Model Callbacks"
},
"before_tool_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Before Tool Callbacks"
},
"after_tool_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "After Tool Callbacks"
}
},
"required": [
@@ -211,6 +301,36 @@
"default": null,
"title": "Sub Agents"
},
"before_agent_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Before Agent Callbacks"
},
"after_agent_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "After Agent Callbacks"
},
"max_iterations": {
"anyOf": [
{
@@ -263,6 +383,36 @@
],
"default": null,
"title": "Sub Agents"
},
"before_agent_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Before Agent Callbacks"
},
"after_agent_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "After Agent Callbacks"
}
},
"required": [
@@ -304,6 +454,36 @@
],
"default": null,
"title": "Sub Agents"
},
"before_agent_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Before Agent Callbacks"
},
"after_agent_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "After Agent Callbacks"
}
},
"required": [
src/google/adk/agents/llm_agent.py
+36 -1
View File
@@ -21,7 +21,6 @@ from typing import Any
from typing import AsyncGenerator
from typing import Awaitable
from typing import Callable
from typing import Dict
from typing import List
from typing import Literal
from typing import Optional
@@ -563,6 +562,8 @@ class LlmAgent(BaseAgent):
config: LlmAgentConfig,
config_abs_path: str,
) -> LlmAgent:
from .config_agent_utils import resolve_callbacks
agent = super().from_config(config, config_abs_path)
if config.model:
agent.model = config.model
@@ -578,6 +579,20 @@ class LlmAgent(BaseAgent):
agent.output_key = config.output_key
if config.tools:
agent.tools = cls._resolve_tools(config.tools)
if config.before_model_callbacks:
agent.before_model_callback = resolve_callbacks(
config.before_model_callbacks
)
if config.after_model_callbacks:
agent.after_model_callback = resolve_callbacks(
config.after_model_callbacks
)
if config.before_tool_callbacks:
agent.before_tool_callback = resolve_callbacks(
config.before_tool_callbacks
)
if config.after_tool_callbacks:
agent.after_tool_callback = resolve_callbacks(config.after_tool_callbacks)
return agent
@@ -664,3 +679,23 @@ class LlmAgentConfig(BaseAgentConfig):
- name: tools.my_mcp_toolset
```
"""
before_model_callbacks: Optional[List[CodeConfig]] = None
"""Optional. LlmAgent.before_model_callbacks.
Example:
```
before_model_callbacks:
- name: my_library.callbacks.before_model_callback
```
"""
after_model_callbacks: Optional[List[CodeConfig]] = None
"""Optional. LlmAgent.after_model_callbacks."""
before_tool_callbacks: Optional[List[CodeConfig]] = None
"""Optional. LlmAgent.before_tool_callbacks."""
after_tool_callbacks: Optional[List[CodeConfig]] = None
"""Optional. LlmAgent.after_tool_callbacks."""