Skip to content

Commit 67021d4

Browse files
viniciusdsmelloviniciusdsmello
committed
feat(CLOSES open-8261): enhance Google ADK tracing with callback support
- Add token usage capture (prompt_tokens, completion_tokens, total_tokens) - Support all 6 ADK callbacks (before/after agent, model, tool) - Enable Google Cloud Trace coexistence (ADK OTel remains active by default) - Fix callback hierarchy for correct chronological order - Add recursive step sorting by start_time - Update example notebook with comprehensive callback demo
1 parent 90b940c commit 67021d4

File tree

3 files changed

+1142
-85
lines changed

3 files changed

+1142
-85
lines changed

‎examples/tracing/google-adk/google_adk_tracing.ipynb‎

Lines changed: 316 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,13 @@
1010
"\n",
1111
"This notebook demonstrates how to trace Google Agent Development Kit (ADK) agents with Openlayer.\n",
1212
"\n",
13+
"## Features\n",
14+
"\n",
15+
"- **Full Agent Tracing**: Capture agent execution, LLM calls, and tool usage\n",
16+
"- **Token Usage Tracking**: Automatically captures prompt, completion, and total tokens\n",
17+
"- **All 6 ADK Callbacks**: Trace before_agent, after_agent, before_model, after_model, before_tool, after_tool\n",
18+
"- **Google Cloud Coexistence**: Use both Google Cloud telemetry (Cloud Trace) AND Openlayer simultaneously\n",
19+
"\n",
1320
"## Prerequisites\n",
1421
"\n",
1522
"Install the required packages:\n",
@@ -18,6 +25,15 @@
1825
"```\n"
1926
]
2027
},
28+
{
29+
"cell_type": "code",
30+
"execution_count": null,
31+
"metadata": {},
32+
"outputs": [],
33+
"source": [
34+
"!pip install google-adk wrapt"
35+
]
36+
},
2137
{
2238
"cell_type": "markdown",
2339
"metadata": {},
@@ -36,12 +52,12 @@
3652
"import os\n",
3753
"\n",
3854
"# Openlayer configuration\n",
39-
"os.environ[\"OPENLAYER_API_KEY\"] = \"your-api-key-here\"\n",
40-
"os.environ[\"OPENLAYER_INFERENCE_PIPELINE_ID\"] = \"your-pipeline-id-here\"\n",
55+
"os.environ[\"OPENLAYER_API_KEY\"] = \"your-api-key\"\n",
56+
"os.environ[\"OPENLAYER_INFERENCE_PIPELINE_ID\"] = \"your-pipeline-id\"\n",
4157
"\n",
4258
"# Google AI API configuration (Option 1: Using Google AI Studio)\n",
4359
"# Get your API key from: https://aistudio.google.com/apikey\n",
44-
"os.environ[\"GOOGLE_API_KEY\"] = \"your-google-ai-api-key-here\"\n",
60+
"os.environ[\"GOOGLE_API_KEY\"] = \"your-google-api-key\"\n",
4561
"\n",
4662
"# Google Cloud Vertex AI configuration (Option 2: Using Google Cloud)\n",
4763
"# Uncomment these if you're using Vertex AI instead of Google AI\n",
@@ -56,7 +72,9 @@
5672
"source": [
5773
"## Enable Google ADK Tracing\n",
5874
"\n",
59-
"Enable tracing before creating any agents. This patches Google ADK globally to send traces to Openlayer:\n"
75+
"Enable tracing before creating any agents. This patches Google ADK globally to send traces to Openlayer.\n",
76+
"\n",
77+
"**Note:** By default, ADK's built-in OpenTelemetry tracing remains active, allowing you to send data to both Google Cloud (Cloud Trace, Cloud Monitoring) AND Openlayer. If you only want Openlayer, use `trace_google_adk(disable_adk_otel=True)`.\n"
6078
]
6179
},
6280
{
@@ -102,7 +120,7 @@
102120
"\n",
103121
"# Create a basic agent\n",
104122
"agent = LlmAgent(\n",
105-
" model=\"gemini-2.0-flash-exp\",\n",
123+
" model=\"gemini-2.5-flash\",\n",
106124
" name=\"Assistant\",\n",
107125
" instruction=\"You are a helpful assistant. Provide concise and accurate responses.\"\n",
108126
")\n",
@@ -190,7 +208,7 @@
190208
"\n",
191209
"# Create agent with tools (pass functions directly)\n",
192210
"tool_agent = LlmAgent(\n",
193-
" model=\"gemini-2.0-flash-exp\",\n",
211+
" model=\"gemini-2.5-flash\",\n",
194212
" name=\"ToolAgent\",\n",
195213
" instruction=\"You are a helpful assistant with access to weather and calculation tools. Use them when appropriate.\",\n",
196214
" tools=[get_weather, calculate]\n",
@@ -229,6 +247,273 @@
229247
"await run_tool_agent()\n"
230248
]
231249
},
250+
{
251+
"cell_type": "markdown",
252+
"metadata": {},
253+
"source": [
254+
"## Example 3: Agent with All 6 Callbacks\n",
255+
"\n",
256+
"Google ADK supports 6 types of callbacks that allow you to observe, customize, and control agent behavior. Openlayer automatically traces all of them:\n",
257+
"\n",
258+
"| Callback | Description | When Called |\n",
259+
"|----------|-------------|-------------|\n",
260+
"| `before_agent_callback` | Agent pre-processing | Before the agent starts its main work |\n",
261+
"| `after_agent_callback` | Agent post-processing | After the agent finishes all its steps |\n",
262+
"| `before_model_callback` | LLM pre-call | Before sending a request to the LLM |\n",
263+
"| `after_model_callback` | LLM post-call | After receiving a response from the LLM |\n",
264+
"| `before_tool_callback` | Tool pre-execution | Before executing a tool |\n",
265+
"| `after_tool_callback` | Tool post-execution | After a tool finishes |\n",
266+
"\n",
267+
"Reference: https://google.github.io/adk-docs/callbacks/\n"
268+
]
269+
},
270+
{
271+
"cell_type": "code",
272+
"execution_count": null,
273+
"metadata": {},
274+
"outputs": [],
275+
"source": [
276+
"from typing import Any, Dict, Optional\n",
277+
"\n",
278+
"from google.adk.tools import ToolContext\n",
279+
"from google.adk.models import LlmRequest, LlmResponse\n",
280+
"from google.adk.tools.base_tool import BaseTool\n",
281+
"from google.adk.agents.callback_context import CallbackContext\n",
282+
"\n",
283+
"# ============================================================================\n",
284+
"# Define all 6 callback functions\n",
285+
"# ============================================================================\n",
286+
"\n",
287+
"# 1. Before Agent Callback\n",
288+
"# Called before the agent starts processing a request\n",
289+
"def before_agent_callback(callback_context: CallbackContext) -> Optional[Any]:\n",
290+
" \"\"\"\n",
291+
" Called before the agent starts its main work.\n",
292+
" \n",
293+
" Use cases:\n",
294+
" - Input validation\n",
295+
" - Session initialization\n",
296+
" - Logging request start\n",
297+
" - Adding default context\n",
298+
" \"\"\"\n",
299+
" print(f\"[before_agent] Agent '{callback_context.agent_name}' starting\") # noqa: T201\n",
300+
" print(f\"[before_agent] Invocation ID: {callback_context.invocation_id}\") # noqa: T201\n",
301+
" # Return None to allow the agent to proceed normally\n",
302+
" # Return a Content object to skip the agent and return that content directly\n",
303+
" return None\n",
304+
"\n",
305+
"\n",
306+
"# 2. After Agent Callback\n",
307+
"# Called after the agent finishes processing\n",
308+
"def after_agent_callback(callback_context: CallbackContext) -> Optional[Any]:\n",
309+
" \"\"\"\n",
310+
" Called after the agent has finished all its steps.\n",
311+
" \n",
312+
" Use cases:\n",
313+
" - Response post-processing\n",
314+
" - Logging request completion\n",
315+
" - Cleanup operations\n",
316+
" - Analytics\n",
317+
" \"\"\"\n",
318+
" print(f\"[after_agent] Agent '{callback_context.agent_name}' finished\") # noqa: T201\n",
319+
" # Return None to use the agent's original response\n",
320+
" # Return a Content object to replace the agent's response\n",
321+
" return None\n",
322+
"\n",
323+
"\n",
324+
"# 3. Before Model Callback\n",
325+
"# Called before each LLM call\n",
326+
"def before_model_callback(\n",
327+
" _callback_context: CallbackContext, \n",
328+
" llm_request: LlmRequest\n",
329+
") -> Optional[LlmResponse]:\n",
330+
" \"\"\"\n",
331+
" Called before sending a request to the LLM.\n",
332+
" \n",
333+
" Use cases:\n",
334+
" - Request modification (add system prompts)\n",
335+
" - Content filtering / guardrails\n",
336+
" - Caching (return cached response)\n",
337+
" - Rate limiting\n",
338+
" \"\"\"\n",
339+
" print(f\"[before_model] Calling model: {llm_request.model}\") # noqa: T201\n",
340+
" print(f\"[before_model] Request has {len(llm_request.contents)} content items\") # noqa: T201\n",
341+
" # Return None to proceed with the LLM call\n",
342+
" # Return an LlmResponse to skip the LLM and use that response instead\n",
343+
" return None\n",
344+
"\n",
345+
"\n",
346+
"# 4. After Model Callback\n",
347+
"# Called after receiving LLM response\n",
348+
"def after_model_callback(\n",
349+
" _callback_context: CallbackContext, \n",
350+
" llm_response: LlmResponse\n",
351+
") -> Optional[LlmResponse]:\n",
352+
" \"\"\"\n",
353+
" Called after receiving a response from the LLM.\n",
354+
" \n",
355+
" Use cases:\n",
356+
" - Response validation\n",
357+
" - Content filtering\n",
358+
" - Response transformation\n",
359+
" - Logging/analytics\n",
360+
" \"\"\"\n",
361+
" print(\"[after_model] Received response from LLM\") # noqa: T201\n",
362+
" if hasattr(llm_response, 'usage_metadata') and llm_response.usage_metadata:\n",
363+
" print(f\"[after_model] Tokens used: {llm_response.usage_metadata.total_token_count}\") # noqa: T201\n",
364+
" # Return None to use the original response\n",
365+
" # Return a modified LlmResponse to replace it\n",
366+
" return None\n",
367+
"\n",
368+
"\n",
369+
"# 5. Before Tool Callback\n",
370+
"# Called before tool execution\n",
371+
"def before_tool_callback(\n",
372+
" tool: BaseTool, \n",
373+
" args: Dict[str, Any], \n",
374+
" _tool_context: ToolContext\n",
375+
") -> Optional[Dict]:\n",
376+
" \"\"\"\n",
377+
" Called before executing a tool.\n",
378+
" \n",
379+
" Use cases:\n",
380+
" - Argument validation\n",
381+
" - Authorization checks\n",
382+
" - Tool call logging\n",
383+
" - Mocking tool responses for testing\n",
384+
" \"\"\"\n",
385+
" print(f\"[before_tool] Executing tool: {tool.name}\") # noqa: T201\n",
386+
" print(f\"[before_tool] Arguments: {args}\") # noqa: T201\n",
387+
" # Return None to proceed with the tool execution\n",
388+
" # Return a dict to skip the tool and use that as the response\n",
389+
" return None\n",
390+
"\n",
391+
"\n",
392+
"# 6. After Tool Callback\n",
393+
"# Called after tool execution\n",
394+
"def after_tool_callback(\n",
395+
" tool: BaseTool, \n",
396+
" _args: Dict[str, Any], \n",
397+
" _tool_context: ToolContext, \n",
398+
" tool_response: Dict\n",
399+
") -> Optional[Dict]:\n",
400+
" \"\"\"\n",
401+
" Called after a tool finishes execution.\n",
402+
" \n",
403+
" Use cases:\n",
404+
" - Response transformation\n",
405+
" - Error handling\n",
406+
" - Logging tool results\n",
407+
" - Caching responses\n",
408+
" \"\"\"\n",
409+
" print(f\"[after_tool] Tool '{tool.name}' completed\") # noqa: T201\n",
410+
" print(f\"[after_tool] Response: {tool_response}\") # noqa: T201\n",
411+
" # Return None to use the original tool response\n",
412+
" # Return a modified dict to replace the response\n",
413+
" return None\n"
414+
]
415+
},
416+
{
417+
"cell_type": "code",
418+
"execution_count": null,
419+
"metadata": {},
420+
"outputs": [],
421+
"source": [
422+
"# ============================================================================\n",
423+
"# Create agent with all callbacks\n",
424+
"# ============================================================================\n",
425+
"\n",
426+
"# Define a tool for the callback agent to use\n",
427+
"def get_current_time() -> str:\n",
428+
" \"\"\"Returns the current time.\n",
429+
" \n",
430+
" Returns:\n",
431+
" str: The current time as a formatted string.\n",
432+
" \"\"\"\n",
433+
" from datetime import datetime\n",
434+
" return f\"The current time is {datetime.now().strftime('%H:%M:%S')}\"\n",
435+
"\n",
436+
"\n",
437+
"# Use different session IDs for callback agent\n",
438+
"CALLBACK_USER_ID = \"user_789\"\n",
439+
"CALLBACK_SESSION_ID = \"session_789\"\n",
440+
"\n",
441+
"# Create agent with ALL 6 callbacks\n",
442+
"callback_agent = LlmAgent(\n",
443+
" model=\"gemini-2.5-flash\",\n",
444+
" name=\"CallbackDemoAgent\",\n",
445+
" instruction=\"You are a helpful assistant. Use the get_current_time tool when asked about time.\",\n",
446+
" tools=[get_current_time],\n",
447+
" # Register all 6 callbacks\n",
448+
" before_agent_callback=before_agent_callback,\n",
449+
" after_agent_callback=after_agent_callback,\n",
450+
" before_model_callback=before_model_callback,\n",
451+
" after_model_callback=after_model_callback,\n",
452+
" before_tool_callback=before_tool_callback,\n",
453+
" after_tool_callback=after_tool_callback,\n",
454+
")\n",
455+
"\n",
456+
"# Create runner for callback agent\n",
457+
"callback_runner = Runner(\n",
458+
" agent=callback_agent,\n",
459+
" app_name=APP_NAME,\n",
460+
" session_service=session_service\n",
461+
")\n",
462+
"\n",
463+
"# Define async function to run the callback agent\n",
464+
"async def run_callback_agent():\n",
465+
" # Create session\n",
466+
" await session_service.create_session(\n",
467+
" app_name=APP_NAME,\n",
468+
" user_id=CALLBACK_USER_ID,\n",
469+
" session_id=CALLBACK_SESSION_ID\n",
470+
" )\n",
471+
" \n",
472+
" # Run the agent with a query that will trigger tool use\n",
473+
" query = \"What time is it right now?\"\n",
474+
" content = types.Content(role='user', parts=[types.Part(text=query)])\n",
475+
" \n",
476+
" # Process events and get response\n",
477+
" async for event in callback_runner.run_async(\n",
478+
" user_id=CALLBACK_USER_ID,\n",
479+
" session_id=CALLBACK_SESSION_ID,\n",
480+
" new_message=content\n",
481+
" ):\n",
482+
" if event.is_final_response() and event.content:\n",
483+
" print(f\"Final Response: {event.content.parts[0].text.strip()}\") # noqa: T201\n",
484+
"\n",
485+
"# Run the async function\n",
486+
"await run_callback_agent()\n"
487+
]
488+
},
489+
{
490+
"cell_type": "markdown",
491+
"metadata": {},
492+
"source": [
493+
"### What You'll See in Openlayer\n",
494+
"\n",
495+
"After running the callback agent, you'll see in your Openlayer dashboard:\n",
496+
"\n",
497+
"1. **Agent Step** (`CallbackDemoAgent`):\n",
498+
" - Shows which callbacks are registered\n",
499+
" - Contains all nested steps in chronological order\n",
500+
"\n",
501+
"2. **All Callbacks as Siblings** (direct children of Agent):\n",
502+
" - `Callback: before_agent` - First, before any processing\n",
503+
" - `Callback: before_model` - Before each LLM call\n",
504+
" - `LLM Call: gemini-2.0-flash-exp` - The actual LLM invocation\n",
505+
" - `Callback: after_model` - After each LLM call (includes token counts)\n",
506+
" - `Callback: before_tool` - Before tool execution\n",
507+
" - `Tool: get_current_time` - The actual tool execution\n",
508+
" - `Callback: after_tool` - After tool completion\n",
509+
" - `Callback: after_agent` - Last, after all processing complete\n",
510+
"\n",
511+
"3. **Token Usage** (captured on LLM Call steps):\n",
512+
" - Prompt tokens\n",
513+
" - Completion tokens \n",
514+
" - Total tokens\n"
515+
]
516+
},
232517
{
233518
"cell_type": "markdown",
234519
"metadata": {},
@@ -240,11 +525,30 @@
240525
"1. Go to https://app.openlayer.com\n",
241526
"2. Navigate to your inference pipeline\n",
242527
"3. View the traces tab to see:\n",
243-
" - Agent execution steps\n",
244-
" - LLM calls with token counts\n",
245-
" - Tool executions with inputs and outputs\n",
246-
" - Latency for each operation\n",
247-
" - Complete execution hierarchy\n"
528+
" - **Agent execution steps** with nested hierarchy\n",
529+
" - **LLM calls** with token counts (prompt, completion, total)\n",
530+
" - **Tool executions** with inputs and outputs\n",
531+
" - **All 6 callback types** traced as separate steps\n",
532+
" - **Latency** for each operation\n",
533+
" - **Complete execution hierarchy** showing the flow\n",
534+
"\n",
535+
"The traces will show the hierarchy of operations in chronological order:\n",
536+
"```\n",
537+
"Agent: CallbackDemoAgent\n",
538+
"├── Callback: before_agent (CallbackDemoAgent)\n",
539+
"├── Callback: before_model (CallbackDemoAgent)\n",
540+
"├── LLM Call: gemini-2.0-flash-exp\n",
541+
"├── Callback: after_model (CallbackDemoAgent)\n",
542+
"├── Callback: before_tool (CallbackDemoAgent)\n",
543+
"├── Tool: get_current_time\n",
544+
"├── Callback: after_tool (CallbackDemoAgent)\n",
545+
"├── Callback: before_model (CallbackDemoAgent)\n",
546+
"├── LLM Call: gemini-2.0-flash-exp\n",
547+
"├── Callback: after_model (CallbackDemoAgent)\n",
548+
"└── Callback: after_agent (CallbackDemoAgent)\n",
549+
"```\n",
550+
"\n",
551+
"**Note:** All callbacks are direct children of the Agent step, appearing in chronological order alongside LLM calls and tool executions. This provides a clear timeline view of the agent execution flow.\n"
248552
]
249553
},
250554
{
@@ -285,7 +589,7 @@
285589
"name": "python",
286590
"nbconvert_exporter": "python",
287591
"pygments_lexer": "ipython3",
288-
"version": "3.9.18"
592+
"version": "3.12.8"
289593
}
290594
},
291595
"nbformat": 4,

0 commit comments

Comments
 (0)