Skip to content

Commit cff0f33

Browse files
viniciusdsmelloviniciusdsmello
committed
feat(closes OPEN-7119): add Google ADK tracing integration with Openlayer
- Add comprehensive Google ADK tracer with support for agents, LLM calls, and tools - Implement smart agent nesting to avoid excessive hierarchy depth - Add context variable tracking for agent transfers - Include interactive Jupyter notebook example - Add wrapt>=1.14.0 as dependency for method patching - Export trace_google_adk and unpatch_google_adk functions
1 parent cceb302 commit cff0f33

File tree

5 files changed

+1226
-1
lines changed

5 files changed

+1226
-1
lines changed

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

Lines changed: 291 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,291 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# Google ADK Tracing with Openlayer\n",
8+
"\n",
9+
"This notebook demonstrates how to trace Google Agent Development Kit (ADK) agents with Openlayer.\n",
10+
"\n",
11+
"## Prerequisites\n",
12+
"\n",
13+
"Install the required packages:\n",
14+
"```bash\n",
15+
"pip install openlayer google-adk wrapt\n",
16+
"```\n"
17+
]
18+
},
19+
{
20+
"cell_type": "markdown",
21+
"metadata": {},
22+
"source": [
23+
"## Setup\n",
24+
"\n",
25+
"First, configure your Openlayer credentials and Google Cloud credentials:\n"
26+
]
27+
},
28+
{
29+
"cell_type": "code",
30+
"execution_count": null,
31+
"metadata": {},
32+
"outputs": [],
33+
"source": [
34+
"import os\n",
35+
"\n",
36+
"# Openlayer configuration\n",
37+
"os.environ[\"OPENLAYER_API_KEY\"] = \"your-api-key-here\"\n",
38+
"os.environ[\"OPENLAYER_INFERENCE_PIPELINE_ID\"] = \"your-pipeline-id-here\"\n",
39+
"\n",
40+
"# Google AI API configuration (Option 1: Using Google AI Studio)\n",
41+
"# Get your API key from: https://aistudio.google.com/apikey\n",
42+
"os.environ[\"GOOGLE_API_KEY\"] = \"your-google-ai-api-key-here\"\n",
43+
"\n",
44+
"# Google Cloud Vertex AI configuration (Option 2: Using Google Cloud)\n",
45+
"# Uncomment these if you're using Vertex AI instead of Google AI\n",
46+
"# os.environ[\"GOOGLE_APPLICATION_CREDENTIALS\"] = \"path/to/your/service-account-key.json\"\n",
47+
"# os.environ[\"GOOGLE_CLOUD_PROJECT\"] = \"your-project-id\"\n",
48+
"# os.environ[\"GOOGLE_CLOUD_LOCATION\"] = \"us-central1\"\n"
49+
]
50+
},
51+
{
52+
"cell_type": "markdown",
53+
"metadata": {},
54+
"source": [
55+
"## Enable Google ADK Tracing\n",
56+
"\n",
57+
"Enable tracing before creating any agents. This patches Google ADK globally to send traces to Openlayer:\n"
58+
]
59+
},
60+
{
61+
"cell_type": "code",
62+
"execution_count": null,
63+
"metadata": {},
64+
"outputs": [],
65+
"source": [
66+
"from openlayer.lib.integrations import trace_google_adk\n",
67+
"\n",
68+
"# Enable tracing (must be called before creating agents)\n",
69+
"trace_google_adk()"
70+
]
71+
},
72+
{
73+
"cell_type": "markdown",
74+
"metadata": {},
75+
"source": [
76+
"## Example 1: Basic Agent with LLM Calls\n",
77+
"\n",
78+
"Create a simple agent that responds to user queries:\n"
79+
]
80+
},
81+
{
82+
"cell_type": "code",
83+
"execution_count": null,
84+
"metadata": {},
85+
"outputs": [],
86+
"source": [
87+
"\n",
88+
"from google.genai import types\n",
89+
"from google.adk.agents import LlmAgent\n",
90+
"from google.adk.runners import Runner\n",
91+
"from google.adk.sessions import InMemorySessionService\n",
92+
"\n",
93+
"# Setup constants\n",
94+
"APP_NAME = \"openlayer_demo\"\n",
95+
"USER_ID = \"user_123\"\n",
96+
"SESSION_ID = \"session_123\"\n",
97+
"\n",
98+
"# Create session service (shared across examples)\n",
99+
"session_service = InMemorySessionService()\n",
100+
"\n",
101+
"# Create a basic agent\n",
102+
"agent = LlmAgent(\n",
103+
" model=\"gemini-2.0-flash-exp\",\n",
104+
" name=\"Assistant\",\n",
105+
" instruction=\"You are a helpful assistant. Provide concise and accurate responses.\"\n",
106+
")\n",
107+
"\n",
108+
"# Create runner\n",
109+
"runner = Runner(\n",
110+
" agent=agent,\n",
111+
" app_name=APP_NAME,\n",
112+
" session_service=session_service\n",
113+
")\n",
114+
"\n",
115+
"# Define async function to run the agent\n",
116+
"async def run_basic_agent():\n",
117+
" # Create session\n",
118+
" await session_service.create_session(\n",
119+
" app_name=APP_NAME,\n",
120+
" user_id=USER_ID,\n",
121+
" session_id=SESSION_ID\n",
122+
" )\n",
123+
" \n",
124+
" # Run the agent\n",
125+
" query = \"What is the capital of France?\"\n",
126+
" content = types.Content(role='user', parts=[types.Part(text=query)])\n",
127+
" \n",
128+
" # Process events and get response\n",
129+
" async for event in runner.run_async(user_id=USER_ID, session_id=SESSION_ID, new_message=content):\n",
130+
" if event.is_final_response() and event.content:\n",
131+
" final_answer = event.content.parts[0].text.strip()\n",
132+
"\n",
133+
"# Run the async function\n",
134+
"await run_basic_agent()"
135+
]
136+
},
137+
{
138+
"cell_type": "markdown",
139+
"metadata": {},
140+
"source": [
141+
"## Example 2: Agent with Tools/Functions\n",
142+
"\n",
143+
"Create an agent with custom tools that can be called during execution:\n"
144+
]
145+
},
146+
{
147+
"cell_type": "code",
148+
"execution_count": null,
149+
"metadata": {},
150+
"outputs": [],
151+
"source": [
152+
"import os\n",
153+
"\n",
154+
"\n",
155+
"# Define custom tools as regular Python functions\n",
156+
"def get_weather(city: str) -> str:\n",
157+
" \"\"\"Retrieves the current weather report for a specified city.\n",
158+
" \n",
159+
" Args:\n",
160+
" city: The name of the city for which to retrieve the weather report.\n",
161+
" \n",
162+
" Returns:\n",
163+
" str: Weather report or error message.\n",
164+
" \"\"\"\n",
165+
" if city.lower() == \"san francisco\":\n",
166+
" return \"The weather in San Francisco is sunny with a temperature of 72°F (22°C).\"\n",
167+
" else:\n",
168+
" return f\"Sorry, weather information for '{city}' is not available.\"\n",
169+
"\n",
170+
"def calculate(expression: str) -> str:\n",
171+
" \"\"\"Evaluates a mathematical expression.\n",
172+
" \n",
173+
" Args:\n",
174+
" expression: A mathematical expression to evaluate.\n",
175+
" \n",
176+
" Returns:\n",
177+
" str: Calculation result or error message.\n",
178+
" \"\"\"\n",
179+
" try:\n",
180+
" result = eval(expression)\n",
181+
" return f\"The result is {result}\"\n",
182+
" except Exception as e:\n",
183+
" return f\"Error: {str(e)}\"\n",
184+
"\n",
185+
"# Use different session IDs for tool agent\n",
186+
"TOOL_USER_ID = \"user_456\"\n",
187+
"TOOL_SESSION_ID = \"session_456\"\n",
188+
"\n",
189+
"# Create agent with tools (pass functions directly)\n",
190+
"tool_agent = LlmAgent(\n",
191+
" model=\"gemini-2.0-flash-exp\",\n",
192+
" name=\"ToolAgent\",\n",
193+
" instruction=\"You are a helpful assistant with access to weather and calculation tools. Use them when appropriate.\",\n",
194+
" tools=[get_weather, calculate]\n",
195+
")\n",
196+
"\n",
197+
"# Create runner for tool agent (reuse the session_service)\n",
198+
"tool_runner = Runner(\n",
199+
" agent=tool_agent,\n",
200+
" app_name=APP_NAME,\n",
201+
" session_service=session_service\n",
202+
")\n",
203+
"\n",
204+
"# Define async function to run the tool agent\n",
205+
"async def run_tool_agent():\n",
206+
" # Create session\n",
207+
" await session_service.create_session(\n",
208+
" app_name=APP_NAME,\n",
209+
" user_id=TOOL_USER_ID,\n",
210+
" session_id=TOOL_SESSION_ID\n",
211+
" )\n",
212+
" \n",
213+
" # Run the agent with a query that requires tool use\n",
214+
" query = \"What's the weather in San Francisco? Also, what is 15 * 24?\"\n",
215+
" content = types.Content(role='user', parts=[types.Part(text=query)])\n",
216+
" \n",
217+
" # Process events and get response\n",
218+
" async for event in tool_runner.run_async(\n",
219+
" user_id=TOOL_USER_ID,\n",
220+
" session_id=TOOL_SESSION_ID,\n",
221+
" new_message=content\n",
222+
" ):\n",
223+
" if event.is_final_response() and event.content:\n",
224+
" final_answer = event.content.parts[0].text.strip()\n",
225+
"\n",
226+
"# Run the async function\n",
227+
"await run_tool_agent()\n"
228+
]
229+
},
230+
{
231+
"cell_type": "markdown",
232+
"metadata": {},
233+
"source": [
234+
"## View Traces in Openlayer\n",
235+
"\n",
236+
"After running these examples, you can view the traces in your Openlayer dashboard:\n",
237+
"\n",
238+
"1. Go to https://app.openlayer.com\n",
239+
"2. Navigate to your inference pipeline\n",
240+
"3. View the traces tab to see:\n",
241+
" - Agent execution steps\n",
242+
" - LLM calls with token counts\n",
243+
" - Tool executions with inputs and outputs\n",
244+
" - Latency for each operation\n",
245+
" - Complete execution hierarchy\n"
246+
]
247+
},
248+
{
249+
"cell_type": "markdown",
250+
"metadata": {},
251+
"source": [
252+
"## Disable Tracing\n",
253+
"\n",
254+
"When you're done, you can disable tracing to restore ADK's default behavior:\n"
255+
]
256+
},
257+
{
258+
"cell_type": "code",
259+
"execution_count": null,
260+
"metadata": {},
261+
"outputs": [],
262+
"source": [
263+
"from openlayer.lib.integrations import unpatch_google_adk\n",
264+
"\n",
265+
"# Disable tracing\n",
266+
"unpatch_google_adk()"
267+
]
268+
}
269+
],
270+
"metadata": {
271+
"kernelspec": {
272+
"display_name": ".venv",
273+
"language": "python",
274+
"name": "python3"
275+
},
276+
"language_info": {
277+
"codemirror_mode": {
278+
"name": "ipython",
279+
"version": 3
280+
},
281+
"file_extension": ".py",
282+
"mimetype": "text/x-python",
283+
"name": "python",
284+
"nbconvert_exporter": "python",
285+
"pygments_lexer": "ipython3",
286+
"version": "3.9.18"
287+
}
288+
},
289+
"nbformat": 4,
290+
"nbformat_minor": 2
291+
}

‎pyproject.toml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@ dependencies = [
1919
"pyarrow>=18.0.0; python_version >= '3.9'",
2020
"pyyaml>=6.0",
2121
"requests_toolbelt>=1.0.0",
22-
"tqdm"
22+
"tqdm",
23+
"wrapt>=1.14.0"
2324
]
2425
requires-python = ">= 3.8"
2526
classifiers = [

‎src/openlayer/lib/__init__.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
"trace_oci_genai",
1515
"trace_oci", # Alias for backward compatibility
1616
"trace_litellm",
17+
"trace_google_adk",
18+
"unpatch_google_adk",
1719
"update_current_trace",
1820
"update_current_step",
1921
# Offline buffer management functions
@@ -188,3 +190,48 @@ def trace_litellm():
188190
from .integrations import litellm_tracer
189191

190192
return litellm_tracer.trace_litellm()
193+
194+
195+
# ------------------------------ Google ADK ---------------------------------- #
196+
def trace_google_adk():
197+
"""Enable tracing for Google Agent Development Kit (ADK).
198+
199+
This function patches Google ADK to automatically trace agent execution,
200+
LLM calls, and tool calls made through the ADK framework.
201+
202+
Requirements:
203+
Google ADK and wrapt must be installed:
204+
pip install google-adk wrapt
205+
206+
Example:
207+
>>> import os
208+
>>> os.environ["OPENLAYER_API_KEY"] = "your-api-key"
209+
>>> os.environ["OPENLAYER_INFERENCE_PIPELINE_ID"] = "your-pipeline-id"
210+
>>> from openlayer.lib import trace_google_adk
211+
>>> # Enable tracing (must be called before creating agents)
212+
>>> trace_google_adk()
213+
>>> # Now create and run your ADK agents
214+
>>> from google.adk.agents import Agent
215+
>>> agent = Agent(name="Assistant", model="gemini-2.0-flash-exp")
216+
>>> result = await agent.run_async(...)
217+
"""
218+
# pylint: disable=import-outside-toplevel
219+
from .integrations import google_adk_tracer
220+
221+
return google_adk_tracer.trace_google_adk()
222+
223+
224+
def unpatch_google_adk():
225+
"""Remove Google ADK tracing patches.
226+
227+
This function restores Google ADK's original behavior by removing all
228+
Openlayer instrumentation.
229+
230+
Example:
231+
>>> from openlayer.lib import unpatch_google_adk
232+
>>> unpatch_google_adk()
233+
"""
234+
# pylint: disable=import-outside-toplevel
235+
from .integrations import google_adk_tracer
236+
237+
return google_adk_tracer.unpatch_google_adk()

‎src/openlayer/lib/integrations/__init__.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,3 +24,10 @@
2424
__all__.extend(["trace_oci_genai"])
2525
except ImportError:
2626
pass
27+
28+
try:
29+
from .google_adk_tracer import trace_google_adk, unpatch_google_adk
30+
31+
__all__.extend(["trace_google_adk", "unpatch_google_adk"])
32+
except ImportError:
33+
pass

0 commit comments

Comments
 (0)