-
Notifications
You must be signed in to change notification settings - Fork 3.5k
Expand file tree
/
Copy pathlogging_plugin.py
More file actions
323 lines (281 loc) · 10.8 KB
/
logging_plugin.py
File metadata and controls
323 lines (281 loc) · 10.8 KB
Edit and raw actions
OlderNewer
1
# Copyright 2026 Google LLC
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
# http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14
15
from __future__ import annotations
16
17
from typing import Any
18
from typing import Optional
19
from typing import TYPE_CHECKING
20
21
from google.genai import types
22
from typing_extensions import override
23
24
from ..agents.base_agent import BaseAgent
25
from ..agents.callback_context import CallbackContext
26
from ..events.event import Event
27
from ..models.llm_request import LlmRequest
28
from ..models.llm_response import LlmResponse
29
from ..tools.base_tool import BaseTool
30
from ..tools.tool_context import ToolContext
31
from .base_plugin import BasePlugin
32
33
if TYPE_CHECKING:
34
from ..agents.invocation_context import InvocationContext
35
36
37
class LoggingPlugin(BasePlugin):
38
"""A plugin that logs important information at each callback point.
39
40
This plugin helps print all critical events in the console. It is not a
41
replacement of existing logging in ADK. It rather helps terminal based
42
debugging by showing all logs in the console, and serves as a simple demo for
43
everyone to leverage when developing new plugins.
44
45
This plugin helps users track the invocation status by logging:
46
- User messages and invocation context
47
- Agent execution flow
48
- LLM requests and responses
49
- Tool calls with arguments and results
50
- Events and final responses
51
- Errors during model and tool execution
52
53
Example:
54
>>> logging_plugin = LoggingPlugin()
55
>>> runner = Runner(
56
... agents=[my_agent],
57
... # ...
58
... plugins=[logging_plugin],
59
... )
60
"""
61
62
def __init__(self, name: str = "logging_plugin"):
63
"""Initialize the logging plugin.
64
65
Args:
66
name: The name of the plugin instance.
67
"""
68
super().__init__(name)
69
70
@override
71
async def on_user_message_callback(
72
self,
73
*,
74
invocation_context: InvocationContext,
75
user_message: types.Content,
76
) -> Optional[types.Content]:
77
"""Log user message and invocation start."""
78
self._log(f"🚀 USER MESSAGE RECEIVED")
79
self._log(f" Invocation ID: {invocation_context.invocation_id}")
80
self._log(f" Session ID: {invocation_context.session.id}")
81
self._log(f" User ID: {invocation_context.user_id}")
82
self._log(f" App Name: {invocation_context.app_name}")
83
self._log(
84
" Root Agent:"
85
f" {invocation_context.agent.name if hasattr(invocation_context.agent, 'name') else 'Unknown'}"
86
)
87
self._log(f" User Content: {self._format_content(user_message)}")
88
if invocation_context.branch:
89
self._log(f" Branch: {invocation_context.branch}")
90
return None
91
92
@override
93
async def before_run_callback(
94
self, *, invocation_context: InvocationContext
95
) -> Optional[types.Content]:
96
"""Log invocation start."""
97
self._log(f"🏃 INVOCATION STARTING")
98
self._log(f" Invocation ID: {invocation_context.invocation_id}")
99
self._log(
100
" Starting Agent:"
101
f" {invocation_context.agent.name if hasattr(invocation_context.agent, 'name') else 'Unknown'}"
102
)
103
return None
104
105
@override
106
async def on_event_callback(
107
self, *, invocation_context: InvocationContext, event: Event
108
) -> Optional[Event]:
109
"""Log events yielded from the runner."""
110
self._log(f"📢 EVENT YIELDED")
111
self._log(f" Event ID: {event.id}")
112
self._log(f" Author: {event.author}")
113
self._log(f" Content: {self._format_content(event.content)}")
114
self._log(f" Final Response: {event.is_final_response()}")
115
116
if event.get_function_calls():
117
func_calls = [fc.name for fc in event.get_function_calls()]
118
self._log(f" Function Calls: {func_calls}")
119
120
if event.get_function_responses():
121
func_responses = [fr.name for fr in event.get_function_responses()]
122
self._log(f" Function Responses: {func_responses}")
123
124
if event.long_running_tool_ids:
125
self._log(f" Long Running Tools: {list(event.long_running_tool_ids)}")
126
127
return None
128
129
@override
130
async def after_run_callback(
131
self, *, invocation_context: InvocationContext
132
) -> Optional[None]:
133
"""Log invocation completion."""
134
self._log(f"✅ INVOCATION COMPLETED")
135
self._log(f" Invocation ID: {invocation_context.invocation_id}")
136
self._log(
137
" Final Agent:"
138
f" {invocation_context.agent.name if hasattr(invocation_context.agent, 'name') else 'Unknown'}"
139
)
140
return None
141
142
@override
143
async def before_agent_callback(
144
self, *, agent: BaseAgent, callback_context: CallbackContext
145
) -> Optional[types.Content]:
146
"""Log agent execution start."""
147
self._log(f"🤖 AGENT STARTING")
148
self._log(f" Agent Name: {callback_context.agent_name}")
149
self._log(f" Invocation ID: {callback_context.invocation_id}")
150
if callback_context._invocation_context.branch:
151
self._log(f" Branch: {callback_context._invocation_context.branch}")
152
return None
153
154
@override
155
async def after_agent_callback(
156
self, *, agent: BaseAgent, callback_context: CallbackContext
157
) -> Optional[types.Content]:
158
"""Log agent execution completion."""
159
self._log(f"🤖 AGENT COMPLETED")
160
self._log(f" Agent Name: {callback_context.agent_name}")
161
self._log(f" Invocation ID: {callback_context.invocation_id}")
162
return None
163
164
@override
165
async def before_model_callback(
166
self, *, callback_context: CallbackContext, llm_request: LlmRequest
167
) -> Optional[LlmResponse]:
168
"""Log LLM request before sending to model."""
169
self._log(f"🧠 LLM REQUEST")
170
self._log(f" Model: {llm_request.model or 'default'}")
171
self._log(f" Agent: {callback_context.agent_name}")
172
173
# Log system instruction if present
174
if llm_request.config and llm_request.config.system_instruction:
175
sys_instruction = llm_request.config.system_instruction[:200]
176
if len(llm_request.config.system_instruction) > 200:
177
sys_instruction += "..."
178
self._log(f" System Instruction: '{sys_instruction}'")
179
180
# Note: Content logging removed due to type compatibility issues
181
# Users can still see content in the LLM response
182
183
# Log available tools
184
if llm_request.tools_dict:
185
tool_names = list(llm_request.tools_dict.keys())
186
self._log(f" Available Tools: {tool_names}")
187
188
return None
189
190
@override
191
async def after_model_callback(
192
self, *, callback_context: CallbackContext, llm_response: LlmResponse
193
) -> Optional[LlmResponse]:
194
"""Log LLM response after receiving from model."""
195
self._log(f"🧠 LLM RESPONSE")
196
self._log(f" Agent: {callback_context.agent_name}")
197
198
if llm_response.error_code:
199
self._log(f" ❌ ERROR - Code: {llm_response.error_code}")
200
self._log(f" Error Message: {llm_response.error_message}")
201
else:
202
self._log(f" Content: {self._format_content(llm_response.content)}")
203
if llm_response.partial:
204
self._log(f" Partial: {llm_response.partial}")
205
if llm_response.turn_complete is not None:
206
self._log(f" Turn Complete: {llm_response.turn_complete}")
207
208
# Log usage metadata if available
209
if llm_response.usage_metadata:
210
self._log(
211
" Token Usage - Input:"
212
f" {llm_response.usage_metadata.prompt_token_count}, Output:"
213
f" {llm_response.usage_metadata.candidates_token_count}"
214
)
215
216
return None
217
218
@override
219
async def before_tool_callback(
220
self,
221
*,
222
tool: BaseTool,
223
tool_args: dict[str, Any],
224
tool_context: ToolContext,
225
) -> Optional[dict]:
226
"""Log tool execution start."""
227
self._log(f"🔧 TOOL STARTING")
228
self._log(f" Tool Name: {tool.name}")
229
self._log(f" Agent: {tool_context.agent_name}")
230
self._log(f" Function Call ID: {tool_context.function_call_id}")
231
self._log(f" Arguments: {self._format_args(tool_args)}")
232
return None
233
234
@override
235
async def after_tool_callback(
236
self,
237
*,
238
tool: BaseTool,
239
tool_args: dict[str, Any],
240
tool_context: ToolContext,
241
result: dict,
242
) -> Optional[dict]:
243
"""Log tool execution completion."""
244
self._log(f"🔧 TOOL COMPLETED")
245
self._log(f" Tool Name: {tool.name}")
246
self._log(f" Agent: {tool_context.agent_name}")
247
self._log(f" Function Call ID: {tool_context.function_call_id}")
248
self._log(f" Result: {self._format_args(result)}")
249
return None
250
251
@override
252
async def on_model_error_callback(
253
self,
254
*,
255
callback_context: CallbackContext,
256
llm_request: LlmRequest,
257
error: Exception,
258
) -> Optional[LlmResponse]:
259
"""Log LLM error."""
260
self._log(f"🧠 LLM ERROR")
261
self._log(f" Agent: {callback_context.agent_name}")
262
self._log(f" Error: {error}")
263
264
return None
265
266
@override
267
async def on_tool_error_callback(
268
self,
269
*,
270
tool: BaseTool,
271
tool_args: dict[str, Any],
272
tool_context: ToolContext,
273
error: Exception,
274
) -> Optional[dict]:
275
"""Log tool error."""
276
self._log(f"🔧 TOOL ERROR")
277
self._log(f" Tool Name: {tool.name}")
278
self._log(f" Agent: {tool_context.agent_name}")
279
self._log(f" Function Call ID: {tool_context.function_call_id}")
280
self._log(f" Arguments: {self._format_args(tool_args)}")
281
self._log(f" Error: {error}")
282
return None
283
284
def _log(self, message: str) -> None:
285
"""Internal method to format and print log messages."""
286
# ANSI color codes: \033[90m for grey, \033[0m to reset
287
formatted_message: str = f"\033[90m[{self.name}] {message}\033[0m"
288
print(formatted_message)
289
290
def _format_content(
291
self, content: Optional[types.Content], max_length: int = 200
292
) -> str:
293
"""Format content for logging, truncating if too long."""
294
if not content or not content.parts:
295
return "None"
296
297
parts = []
298
for part in content.parts:
299
if part.text:
300
text = part.text.strip()
301
if len(text) > max_length:
302
text = text[:max_length] + "..."
303
parts.append(f"text: '{text}'")
304
elif part.function_call:
305
parts.append(f"function_call: {part.function_call.name}")
306
elif part.function_response:
307
parts.append(f"function_response: {part.function_response.name}")
308
elif part.code_execution_result:
309
parts.append("code_execution_result")
310
else:
311
parts.append("other_part")
312
313
return " | ".join(parts)
314
315
def _format_args(self, args: dict[str, Any], max_length: int = 300) -> str:
316
"""Format arguments dictionary for logging."""
317
if not args:
318
return "{}"
319
320
formatted = str(args)
321
if len(formatted) > max_length:
322
formatted = formatted[:max_length] + "...}"
323
return formatted