Skip to content

Commit e50f05a

Browse files
jawoszekcopybara-github
authored andcommitted
feat(otel): env variable for disabling llm_request and llm_response in spans
The default without the variable set is to keep the content in spans to keep backward compatible behavior for existing users. This allows to enable tracing without potential PII data from request and response. Google GenAI instrumentation lib requires an explicit opt-in to enable request and response content - see https://github.com/open-telemetry/opentelemetry-python-contrib/tree/main/instrumentation-genai/opentelemetry-instrumentation-google-genai#enabling-message-content. PiperOrigin-RevId: 820351154
1 parent 9dc0036 commit e50f05a

File tree

2 files changed

+208
-29
lines changed

2 files changed

+208
-29
lines changed

src/google/adk/telemetry/tracing.py

Lines changed: 62 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from __future__ import annotations
2525

2626
import json
27+
import os
2728
from typing import Any
2829
from typing import TYPE_CHECKING
2930

@@ -33,6 +34,9 @@
3334
from .. import version
3435
from ..events.event import Event
3536

37+
# By default some ADK spans include attributes with potential PII data.
38+
# This env, when set to false, allows to disable populating those attributes.
39+
ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS = 'ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS'
3640
# TODO: Replace with constant from opentelemetry.semconv when it reaches version 1.37 in g3.
3741
GEN_AI_AGENT_DESCRIPTION = 'gen_ai.agent.description'
3842
GEN_AI_AGENT_NAME = 'gen_ai.agent.name'
@@ -138,10 +142,13 @@ def trace_tool_call(
138142
span.set_attribute('gcp.vertex.agent.llm_request', '{}')
139143
span.set_attribute('gcp.vertex.agent.llm_response', '{}')
140144

141-
span.set_attribute(
142-
'gcp.vertex.agent.tool_call_args',
143-
_safe_json_serialize(args),
144-
)
145+
if _should_add_request_response_to_spans():
146+
span.set_attribute(
147+
'gcp.vertex.agent.tool_call_args',
148+
_safe_json_serialize(args),
149+
)
150+
else:
151+
span.set_attribute('gcp.vertex.agent.tool_call_args', {})
145152

146153
# Tracing tool response
147154
tool_call_id = '<not specified>'
@@ -163,10 +170,13 @@ def trace_tool_call(
163170
if not isinstance(tool_response, dict):
164171
tool_response = {'result': tool_response}
165172
span.set_attribute('gcp.vertex.agent.event_id', function_response_event.id)
166-
span.set_attribute(
167-
'gcp.vertex.agent.tool_response',
168-
_safe_json_serialize(tool_response),
169-
)
173+
if _should_add_request_response_to_spans():
174+
span.set_attribute(
175+
'gcp.vertex.agent.tool_response',
176+
_safe_json_serialize(tool_response),
177+
)
178+
else:
179+
span.set_attribute('gcp.vertex.agent.tool_response', {})
170180

171181

172182
def trace_merged_tool_calls(
@@ -200,10 +210,13 @@ def trace_merged_tool_calls(
200210
except Exception: # pylint: disable=broad-exception-caught
201211
function_response_event_json = '<not serializable>'
202212

203-
span.set_attribute(
204-
'gcp.vertex.agent.tool_response',
205-
function_response_event_json,
206-
)
213+
if _should_add_request_response_to_spans():
214+
span.set_attribute(
215+
'gcp.vertex.agent.tool_response',
216+
function_response_event_json,
217+
)
218+
else:
219+
span.set_attribute('gcp.vertex.agent.tool_response', {})
207220
# Setting empty llm request and response (as UI expect these) while not
208221
# applicable for tool_response.
209222
span.set_attribute('gcp.vertex.agent.llm_request', '{}')
@@ -243,10 +256,13 @@ def trace_call_llm(
243256
)
244257
span.set_attribute('gcp.vertex.agent.event_id', event_id)
245258
# Consider removing once GenAI SDK provides a way to record this info.
246-
span.set_attribute(
247-
'gcp.vertex.agent.llm_request',
248-
_safe_json_serialize(_build_llm_request_for_trace(llm_request)),
249-
)
259+
if _should_add_request_response_to_spans():
260+
span.set_attribute(
261+
'gcp.vertex.agent.llm_request',
262+
_safe_json_serialize(_build_llm_request_for_trace(llm_request)),
263+
)
264+
else:
265+
span.set_attribute('gcp.vertex.agent.llm_request', {})
250266
# Consider removing once GenAI SDK provides a way to record this info.
251267
if llm_request.config:
252268
if llm_request.config.top_p:
@@ -265,10 +281,13 @@ def trace_call_llm(
265281
except Exception: # pylint: disable=broad-exception-caught
266282
llm_response_json = '<not serializable>'
267283

268-
span.set_attribute(
269-
'gcp.vertex.agent.llm_response',
270-
llm_response_json,
271-
)
284+
if _should_add_request_response_to_spans():
285+
span.set_attribute(
286+
'gcp.vertex.agent.llm_response',
287+
llm_response_json,
288+
)
289+
else:
290+
span.set_attribute('gcp.vertex.agent.llm_response', {})
272291

273292
if llm_response.usage_metadata is not None:
274293
span.set_attribute(
@@ -309,15 +328,18 @@ def trace_send_data(
309328
span.set_attribute('gcp.vertex.agent.event_id', event_id)
310329
# Once instrumentation is added to the GenAI SDK, consider whether this
311330
# information still needs to be recorded by the Agent Development Kit.
312-
span.set_attribute(
313-
'gcp.vertex.agent.data',
314-
_safe_json_serialize([
315-
types.Content(role=content.role, parts=content.parts).model_dump(
316-
exclude_none=True
317-
)
318-
for content in data
319-
]),
320-
)
331+
if _should_add_request_response_to_spans():
332+
span.set_attribute(
333+
'gcp.vertex.agent.data',
334+
_safe_json_serialize([
335+
types.Content(role=content.role, parts=content.parts).model_dump(
336+
exclude_none=True
337+
)
338+
for content in data
339+
]),
340+
)
341+
else:
342+
span.set_attribute('gcp.vertex.agent.data', {})
321343

322344

323345
def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]:
@@ -350,3 +372,14 @@ def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]:
350372
)
351373
)
352374
return result
375+
376+
377+
# Defaults to true for now to preserve backward compatibility.
378+
# Once prompt and response logging is well established in ADK, we might start
379+
# a deprecation of request/response content in spans by switching the default
380+
# to false.
381+
def _should_add_request_response_to_spans() -> bool:
382+
disabled_via_env_var = os.getenv(
383+
ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'true'
384+
).lower() in ('false', '0')
385+
return not disabled_via_env_var

tests/unittests/telemetry/test_spans.py

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
import json
16+
import os
1617
from typing import Any
1718
from typing import Dict
1819
from typing import Optional
@@ -23,6 +24,7 @@
2324
from google.adk.models.llm_request import LlmRequest
2425
from google.adk.models.llm_response import LlmResponse
2526
from google.adk.sessions.in_memory_session_service import InMemorySessionService
27+
from google.adk.telemetry.tracing import ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS
2628
from google.adk.telemetry.tracing import trace_agent_invocation
2729
from google.adk.telemetry.tracing import trace_call_llm
2830
from google.adk.telemetry.tracing import trace_merged_tool_calls
@@ -371,3 +373,147 @@ def test_trace_merged_tool_calls_sets_correct_attributes(
371373
expected_calls, any_order=True
372374
)
373375
mock_event_fixture.model_dumps_json.assert_called_once_with(exclude_none=True)
376+
377+
378+
@pytest.mark.asyncio
379+
async def test_call_llm_disabling_request_response_content(
380+
monkeypatch, mock_span_fixture
381+
):
382+
"""Test trace_call_llm doesn't set request and response attributes if env is set to false"""
383+
# Arrange
384+
monkeypatch.setenv(ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'false')
385+
monkeypatch.setattr(
386+
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
387+
)
388+
389+
agent = LlmAgent(name='test_agent')
390+
invocation_context = await _create_invocation_context(agent)
391+
llm_request = LlmRequest(
392+
model='gemini-pro',
393+
contents=[
394+
types.Content(
395+
role='user',
396+
parts=[types.Part(text='Hello, how are you?')],
397+
),
398+
],
399+
)
400+
llm_response = LlmResponse(
401+
turn_complete=True,
402+
finish_reason=types.FinishReason.STOP,
403+
)
404+
405+
# Act
406+
trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)
407+
408+
# Assert
409+
assert not any(
410+
call_obj.args[0] == 'gcp.vertex.agent.llm_request'
411+
and call_obj.args[1] != {}
412+
for call_obj in mock_span_fixture.set_attribute.call_args_list
413+
), "Attribute 'gcp.vertex.agent.llm_request' was incorrectly set on the span."
414+
415+
assert not any(
416+
call_obj.args[0] == 'gcp.vertex.agent.llm_response'
417+
and call_obj.args[1] != {}
418+
for call_obj in mock_span_fixture.set_attribute.call_args_list
419+
), (
420+
"Attribute 'gcp.vertex.agent.llm_response' was incorrectly set on the"
421+
' span.'
422+
)
423+
424+
425+
def test_trace_tool_call_disabling_request_response_content(
426+
monkeypatch,
427+
mock_span_fixture,
428+
mock_tool_fixture,
429+
mock_event_fixture,
430+
):
431+
"""Test trace_tool_call doesn't set request and response attributes if env is set to false"""
432+
# Arrange
433+
monkeypatch.setenv(ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'false')
434+
monkeypatch.setattr(
435+
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
436+
)
437+
438+
test_args: Dict[str, Any] = {'query': 'details', 'id_list': [1, 2, 3]}
439+
test_tool_call_id: str = 'tool_call_id_002'
440+
test_event_id: str = 'event_id_dict_002'
441+
dict_function_response: Dict[str, Any] = {
442+
'data': 'structured_data',
443+
'count': 5,
444+
}
445+
446+
mock_event_fixture.id = test_event_id
447+
mock_event_fixture.content = types.Content(
448+
role='user',
449+
parts=[
450+
types.Part(
451+
function_response=types.FunctionResponse(
452+
id=test_tool_call_id,
453+
name='test_function_1',
454+
response=dict_function_response,
455+
)
456+
),
457+
],
458+
)
459+
460+
# Act
461+
trace_tool_call(
462+
tool=mock_tool_fixture,
463+
args=test_args,
464+
function_response_event=mock_event_fixture,
465+
)
466+
467+
# Assert
468+
assert not any(
469+
call_obj.args[0] == 'gcp.vertex.agent.tool_call_args'
470+
and call_obj.args[1] != {}
471+
for call_obj in mock_span_fixture.set_attribute.call_args_list
472+
), (
473+
"Attribute 'gcp.vertex.agent.tool_call_args' was incorrectly set on the"
474+
' span.'
475+
)
476+
477+
assert not any(
478+
call_obj.args[0] == 'gcp.vertex.agent.tool_response'
479+
and call_obj.args[1] != {}
480+
for call_obj in mock_span_fixture.set_attribute.call_args_list
481+
), (
482+
"Attribute 'gcp.vertex.agent.tool_response' was incorrectly set on the"
483+
' span.'
484+
)
485+
486+
487+
def test_trace_merged_tool_disabling_request_response_content(
488+
monkeypatch,
489+
mock_span_fixture,
490+
mock_event_fixture,
491+
):
492+
"""Test trace_merged_tool doesn't set request and response attributes if env is set to false"""
493+
# Arrange
494+
monkeypatch.setenv(ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'false')
495+
monkeypatch.setattr(
496+
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
497+
)
498+
499+
test_response_event_id = 'merged_evt_id_001'
500+
custom_event_json_output = (
501+
'{"custom_event_payload": true, "details": "merged_details"}'
502+
)
503+
mock_event_fixture.model_dumps_json.return_value = custom_event_json_output
504+
505+
# Act
506+
trace_merged_tool_calls(
507+
response_event_id=test_response_event_id,
508+
function_response_event=mock_event_fixture,
509+
)
510+
511+
# Assert
512+
assert not any(
513+
call_obj.args[0] == 'gcp.vertex.agent.tool_response'
514+
and call_obj.args[1] != {}
515+
for call_obj in mock_span_fixture.set_attribute.call_args_list
516+
), (
517+
"Attribute 'gcp.vertex.agent.tool_response' was incorrectly set on the"
518+
' span.'
519+
)

0 commit comments

Comments
 (0)