mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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
This commit is contained in:
committed by
Copybara-Service
parent
9dc00360e3
commit
e50f05a9fc
@@ -24,6 +24,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
@@ -33,6 +34,9 @@ from opentelemetry import trace
|
||||
from .. import version
|
||||
from ..events.event import Event
|
||||
|
||||
# By default some ADK spans include attributes with potential PII data.
|
||||
# This env, when set to false, allows to disable populating those attributes.
|
||||
ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS = 'ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS'
|
||||
# TODO: Replace with constant from opentelemetry.semconv when it reaches version 1.37 in g3.
|
||||
GEN_AI_AGENT_DESCRIPTION = 'gen_ai.agent.description'
|
||||
GEN_AI_AGENT_NAME = 'gen_ai.agent.name'
|
||||
@@ -138,10 +142,13 @@ def trace_tool_call(
|
||||
span.set_attribute('gcp.vertex.agent.llm_request', '{}')
|
||||
span.set_attribute('gcp.vertex.agent.llm_response', '{}')
|
||||
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.tool_call_args',
|
||||
_safe_json_serialize(args),
|
||||
)
|
||||
if _should_add_request_response_to_spans():
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.tool_call_args',
|
||||
_safe_json_serialize(args),
|
||||
)
|
||||
else:
|
||||
span.set_attribute('gcp.vertex.agent.tool_call_args', {})
|
||||
|
||||
# Tracing tool response
|
||||
tool_call_id = '<not specified>'
|
||||
@@ -163,10 +170,13 @@ def trace_tool_call(
|
||||
if not isinstance(tool_response, dict):
|
||||
tool_response = {'result': tool_response}
|
||||
span.set_attribute('gcp.vertex.agent.event_id', function_response_event.id)
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.tool_response',
|
||||
_safe_json_serialize(tool_response),
|
||||
)
|
||||
if _should_add_request_response_to_spans():
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.tool_response',
|
||||
_safe_json_serialize(tool_response),
|
||||
)
|
||||
else:
|
||||
span.set_attribute('gcp.vertex.agent.tool_response', {})
|
||||
|
||||
|
||||
def trace_merged_tool_calls(
|
||||
@@ -200,10 +210,13 @@ def trace_merged_tool_calls(
|
||||
except Exception: # pylint: disable=broad-exception-caught
|
||||
function_response_event_json = '<not serializable>'
|
||||
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.tool_response',
|
||||
function_response_event_json,
|
||||
)
|
||||
if _should_add_request_response_to_spans():
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.tool_response',
|
||||
function_response_event_json,
|
||||
)
|
||||
else:
|
||||
span.set_attribute('gcp.vertex.agent.tool_response', {})
|
||||
# Setting empty llm request and response (as UI expect these) while not
|
||||
# applicable for tool_response.
|
||||
span.set_attribute('gcp.vertex.agent.llm_request', '{}')
|
||||
@@ -243,10 +256,13 @@ def trace_call_llm(
|
||||
)
|
||||
span.set_attribute('gcp.vertex.agent.event_id', event_id)
|
||||
# Consider removing once GenAI SDK provides a way to record this info.
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.llm_request',
|
||||
_safe_json_serialize(_build_llm_request_for_trace(llm_request)),
|
||||
)
|
||||
if _should_add_request_response_to_spans():
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.llm_request',
|
||||
_safe_json_serialize(_build_llm_request_for_trace(llm_request)),
|
||||
)
|
||||
else:
|
||||
span.set_attribute('gcp.vertex.agent.llm_request', {})
|
||||
# Consider removing once GenAI SDK provides a way to record this info.
|
||||
if llm_request.config:
|
||||
if llm_request.config.top_p:
|
||||
@@ -265,10 +281,13 @@ def trace_call_llm(
|
||||
except Exception: # pylint: disable=broad-exception-caught
|
||||
llm_response_json = '<not serializable>'
|
||||
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.llm_response',
|
||||
llm_response_json,
|
||||
)
|
||||
if _should_add_request_response_to_spans():
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.llm_response',
|
||||
llm_response_json,
|
||||
)
|
||||
else:
|
||||
span.set_attribute('gcp.vertex.agent.llm_response', {})
|
||||
|
||||
if llm_response.usage_metadata is not None:
|
||||
span.set_attribute(
|
||||
@@ -309,15 +328,18 @@ def trace_send_data(
|
||||
span.set_attribute('gcp.vertex.agent.event_id', event_id)
|
||||
# Once instrumentation is added to the GenAI SDK, consider whether this
|
||||
# information still needs to be recorded by the Agent Development Kit.
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.data',
|
||||
_safe_json_serialize([
|
||||
types.Content(role=content.role, parts=content.parts).model_dump(
|
||||
exclude_none=True
|
||||
)
|
||||
for content in data
|
||||
]),
|
||||
)
|
||||
if _should_add_request_response_to_spans():
|
||||
span.set_attribute(
|
||||
'gcp.vertex.agent.data',
|
||||
_safe_json_serialize([
|
||||
types.Content(role=content.role, parts=content.parts).model_dump(
|
||||
exclude_none=True
|
||||
)
|
||||
for content in data
|
||||
]),
|
||||
)
|
||||
else:
|
||||
span.set_attribute('gcp.vertex.agent.data', {})
|
||||
|
||||
|
||||
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]:
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
# Defaults to true for now to preserve backward compatibility.
|
||||
# Once prompt and response logging is well established in ADK, we might start
|
||||
# a deprecation of request/response content in spans by switching the default
|
||||
# to false.
|
||||
def _should_add_request_response_to_spans() -> bool:
|
||||
disabled_via_env_var = os.getenv(
|
||||
ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'true'
|
||||
).lower() in ('false', '0')
|
||||
return not disabled_via_env_var
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
# limitations under the License.
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any
|
||||
from typing import Dict
|
||||
from typing import Optional
|
||||
@@ -23,6 +24,7 @@ from google.adk.agents.llm_agent import LlmAgent
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.models.llm_response import LlmResponse
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.adk.telemetry.tracing import ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS
|
||||
from google.adk.telemetry.tracing import trace_agent_invocation
|
||||
from google.adk.telemetry.tracing import trace_call_llm
|
||||
from google.adk.telemetry.tracing import trace_merged_tool_calls
|
||||
@@ -371,3 +373,147 @@ def test_trace_merged_tool_calls_sets_correct_attributes(
|
||||
expected_calls, any_order=True
|
||||
)
|
||||
mock_event_fixture.model_dumps_json.assert_called_once_with(exclude_none=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_llm_disabling_request_response_content(
|
||||
monkeypatch, mock_span_fixture
|
||||
):
|
||||
"""Test trace_call_llm doesn't set request and response attributes if env is set to false"""
|
||||
# Arrange
|
||||
monkeypatch.setenv(ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'false')
|
||||
monkeypatch.setattr(
|
||||
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
|
||||
)
|
||||
|
||||
agent = LlmAgent(name='test_agent')
|
||||
invocation_context = await _create_invocation_context(agent)
|
||||
llm_request = LlmRequest(
|
||||
model='gemini-pro',
|
||||
contents=[
|
||||
types.Content(
|
||||
role='user',
|
||||
parts=[types.Part(text='Hello, how are you?')],
|
||||
),
|
||||
],
|
||||
)
|
||||
llm_response = LlmResponse(
|
||||
turn_complete=True,
|
||||
finish_reason=types.FinishReason.STOP,
|
||||
)
|
||||
|
||||
# Act
|
||||
trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)
|
||||
|
||||
# Assert
|
||||
assert not any(
|
||||
call_obj.args[0] == 'gcp.vertex.agent.llm_request'
|
||||
and call_obj.args[1] != {}
|
||||
for call_obj in mock_span_fixture.set_attribute.call_args_list
|
||||
), "Attribute 'gcp.vertex.agent.llm_request' was incorrectly set on the span."
|
||||
|
||||
assert not any(
|
||||
call_obj.args[0] == 'gcp.vertex.agent.llm_response'
|
||||
and call_obj.args[1] != {}
|
||||
for call_obj in mock_span_fixture.set_attribute.call_args_list
|
||||
), (
|
||||
"Attribute 'gcp.vertex.agent.llm_response' was incorrectly set on the"
|
||||
' span.'
|
||||
)
|
||||
|
||||
|
||||
def test_trace_tool_call_disabling_request_response_content(
|
||||
monkeypatch,
|
||||
mock_span_fixture,
|
||||
mock_tool_fixture,
|
||||
mock_event_fixture,
|
||||
):
|
||||
"""Test trace_tool_call doesn't set request and response attributes if env is set to false"""
|
||||
# Arrange
|
||||
monkeypatch.setenv(ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'false')
|
||||
monkeypatch.setattr(
|
||||
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
|
||||
)
|
||||
|
||||
test_args: Dict[str, Any] = {'query': 'details', 'id_list': [1, 2, 3]}
|
||||
test_tool_call_id: str = 'tool_call_id_002'
|
||||
test_event_id: str = 'event_id_dict_002'
|
||||
dict_function_response: Dict[str, Any] = {
|
||||
'data': 'structured_data',
|
||||
'count': 5,
|
||||
}
|
||||
|
||||
mock_event_fixture.id = test_event_id
|
||||
mock_event_fixture.content = types.Content(
|
||||
role='user',
|
||||
parts=[
|
||||
types.Part(
|
||||
function_response=types.FunctionResponse(
|
||||
id=test_tool_call_id,
|
||||
name='test_function_1',
|
||||
response=dict_function_response,
|
||||
)
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
# Act
|
||||
trace_tool_call(
|
||||
tool=mock_tool_fixture,
|
||||
args=test_args,
|
||||
function_response_event=mock_event_fixture,
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert not any(
|
||||
call_obj.args[0] == 'gcp.vertex.agent.tool_call_args'
|
||||
and call_obj.args[1] != {}
|
||||
for call_obj in mock_span_fixture.set_attribute.call_args_list
|
||||
), (
|
||||
"Attribute 'gcp.vertex.agent.tool_call_args' was incorrectly set on the"
|
||||
' span.'
|
||||
)
|
||||
|
||||
assert not any(
|
||||
call_obj.args[0] == 'gcp.vertex.agent.tool_response'
|
||||
and call_obj.args[1] != {}
|
||||
for call_obj in mock_span_fixture.set_attribute.call_args_list
|
||||
), (
|
||||
"Attribute 'gcp.vertex.agent.tool_response' was incorrectly set on the"
|
||||
' span.'
|
||||
)
|
||||
|
||||
|
||||
def test_trace_merged_tool_disabling_request_response_content(
|
||||
monkeypatch,
|
||||
mock_span_fixture,
|
||||
mock_event_fixture,
|
||||
):
|
||||
"""Test trace_merged_tool doesn't set request and response attributes if env is set to false"""
|
||||
# Arrange
|
||||
monkeypatch.setenv(ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'false')
|
||||
monkeypatch.setattr(
|
||||
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
|
||||
)
|
||||
|
||||
test_response_event_id = 'merged_evt_id_001'
|
||||
custom_event_json_output = (
|
||||
'{"custom_event_payload": true, "details": "merged_details"}'
|
||||
)
|
||||
mock_event_fixture.model_dumps_json.return_value = custom_event_json_output
|
||||
|
||||
# Act
|
||||
trace_merged_tool_calls(
|
||||
response_event_id=test_response_event_id,
|
||||
function_response_event=mock_event_fixture,
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert not any(
|
||||
call_obj.args[0] == 'gcp.vertex.agent.tool_response'
|
||||
and call_obj.args[1] != {}
|
||||
for call_obj in mock_span_fixture.set_attribute.call_args_list
|
||||
), (
|
||||
"Attribute 'gcp.vertex.agent.tool_response' was incorrectly set on the"
|
||||
' span.'
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user