Skip to content

Commit c4fff3f

Browse files
TimPansinoTimPansino
committed
Merge remote-tracking branch 'origin/develop-google-adk' into feat-google-adk-tool-instrumentation
2 parents 0088aea + fef8ce4 commit c4fff3f

5 files changed

Lines changed: 296 additions & 4 deletions

File tree

‎newrelic/config.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3185,6 +3185,9 @@ def _process_module_builtin_defaults():
31853185
"newrelic.hooks.mlmodel_autogen",
31863186
"instrument_autogen_agentchat_agents__assistant_agent",
31873187
)
3188+
_process_module_definition(
3189+
"google.adk.agents.llm_agent", "newrelic.hooks.mlmodel_googleadk", "instrument_googleadk_agents_llm_agent"
3190+
)
31883191
_process_module_definition(
31893192
"google.adk.flows.llm_flows.functions",
31903193
"newrelic.hooks.mlmodel_googleadk",

‎newrelic/hooks/mlmodel_googleadk.py‎

Lines changed: 117 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from newrelic.api.function_trace import FunctionTrace
2121
from newrelic.api.time_trace import get_trace_linking_metadata
2222
from newrelic.api.transaction import current_transaction
23-
from newrelic.common.llm_utils import _get_llm_metadata
23+
from newrelic.common.llm_utils import AsyncLLMStreamProxy, _get_llm_metadata
2424
from newrelic.common.object_names import callable_name
2525
from newrelic.common.object_wrapper import wrap_function_wrapper
2626
from newrelic.common.package_version_utils import get_package_version
@@ -31,9 +31,120 @@
3131
GOOGLEADK_VERSION = get_package_version("google-adk")
3232

3333
RECORD_EVENTS_FAILURE_LOG_MESSAGE = "Exception occurred in Google ADK instrumentation: Failed to record LLM events. Please report this issue to New Relic Support."
34+
AGENT_EVENT_FAILURE_LOG_MESSAGE = "Exception occurred in Google ADK instrumentation: Failed to record agent data. Please report this issue to New Relic Support."
3435
TOOL_EXTRACTOR_FAILURE_LOG_MESSAGE = "Exception occurred in Google ADK instrumentation: Failed to extract tool information. If the issue persists, report this issue to New Relic support.\n"
3536

3637

38+
def wrap_llm_agent__run_async_impl(wrapped, instance, args, kwargs):
39+
transaction = current_transaction()
40+
if not transaction:
41+
return wrapped(*args, **kwargs)
42+
43+
settings = transaction.settings or global_settings()
44+
if not settings.ai_monitoring.enabled:
45+
return wrapped(*args, **kwargs)
46+
47+
agent_name = getattr(instance, "name", "agent")
48+
function_trace_name = f"run_async/{agent_name}"
49+
agentic_subcomponent_data = {"type": "APM-AI_AGENT", "name": agent_name}
50+
51+
ft = FunctionTrace(name=function_trace_name, group="Llm/agent/GoogleADK")
52+
ft.__enter__()
53+
ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data))
54+
55+
linking_metadata = get_trace_linking_metadata()
56+
agent_id = str(uuid.uuid4())
57+
58+
try:
59+
return_val = wrapped(*args, **kwargs)
60+
except Exception:
61+
ft.__exit__(*sys.exc_info())
62+
raise
63+
64+
try:
65+
proxied_return_val = AsyncLLMStreamProxy(
66+
wrapped=return_val,
67+
on_stop_iteration=_record_agent_event_on_stop_iteration,
68+
on_error=_handle_agent_streaming_completion_error,
69+
)
70+
proxied_return_val._nr_ft = ft
71+
proxied_return_val._nr_metadata = linking_metadata
72+
proxied_return_val._nr_adk_attrs = {"agent_name": agent_name, "agent_id": agent_id}
73+
return proxied_return_val
74+
except Exception:
75+
ft.__exit__(*sys.exc_info())
76+
return return_val
77+
78+
79+
def _record_agent_event_on_stop_iteration(self, transaction):
80+
if hasattr(self, "_nr_ft"):
81+
linking_metadata = self._nr_metadata or get_trace_linking_metadata()
82+
self._nr_ft.__exit__(None, None, None)
83+
try:
84+
adk_attrs = getattr(self, "_nr_adk_attrs", {})
85+
if not adk_attrs:
86+
return
87+
88+
agent_name = adk_attrs.get("agent_name", "agent")
89+
agent_id = adk_attrs.get("agent_id")
90+
agent_event_dict = _construct_base_agent_event_dict(
91+
agent_name=agent_name, agent_id=agent_id, transaction=transaction, linking_metadata=linking_metadata
92+
)
93+
agent_event_dict["duration"] = self._nr_ft.duration * 1000
94+
transaction.record_custom_event("LlmAgent", agent_event_dict)
95+
except Exception:
96+
_logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE, exc_info=True)
97+
finally:
98+
if hasattr(self, "_nr_adk_attrs"):
99+
self._nr_adk_attrs.clear()
100+
101+
102+
def _handle_agent_streaming_completion_error(self, transaction):
103+
if hasattr(self, "_nr_ft"):
104+
adk_attrs = getattr(self, "_nr_adk_attrs", {})
105+
if not adk_attrs:
106+
self._nr_ft.__exit__(*sys.exc_info())
107+
return
108+
109+
linking_metadata = self._nr_metadata or get_trace_linking_metadata()
110+
111+
try:
112+
agent_name = adk_attrs.get("agent_name", "agent")
113+
agent_id = adk_attrs.get("agent_id")
114+
115+
self._nr_ft.notice_error(attributes={"agent_id": agent_id})
116+
self._nr_ft.__exit__(*sys.exc_info())
117+
118+
agent_event_dict = _construct_base_agent_event_dict(
119+
agent_name=agent_name, agent_id=agent_id, transaction=transaction, linking_metadata=linking_metadata
120+
)
121+
agent_event_dict.update({"duration": self._nr_ft.duration * 1000, "error": True})
122+
transaction.record_custom_event("LlmAgent", agent_event_dict)
123+
except Exception:
124+
_logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE, exc_info=True)
125+
finally:
126+
if hasattr(self, "_nr_adk_attrs"):
127+
self._nr_adk_attrs.clear()
128+
129+
130+
def _construct_base_agent_event_dict(agent_name, agent_id, transaction, linking_metadata):
131+
try:
132+
agent_event_dict = {
133+
"id": agent_id,
134+
"name": agent_name,
135+
"span_id": linking_metadata.get("span.id"),
136+
"trace_id": linking_metadata.get("trace.id"),
137+
"vendor": "google_adk",
138+
"ingest_source": "Python",
139+
}
140+
agent_event_dict.update(_get_llm_metadata(transaction))
141+
except Exception:
142+
_logger.warning(AGENT_EVENT_FAILURE_LOG_MESSAGE, exc_info=True)
143+
agent_event_dict = {}
144+
145+
return agent_event_dict
146+
147+
37148
async def wrap__execute_single_function_call_async(wrapped, instance, args, kwargs):
38149
transaction = current_transaction()
39150
if not transaction:
@@ -171,6 +282,11 @@ def _construct_base_tool_event_dict(
171282
return tool_event_dict
172283

173284

285+
def instrument_googleadk_agents_llm_agent(module):
286+
if hasattr(module, "LlmAgent") and hasattr(module.LlmAgent, "_run_async_impl"):
287+
wrap_function_wrapper(module, "LlmAgent._run_async_impl", wrap_llm_agent__run_async_impl)
288+
289+
174290
def instrument_googleadk_flows_llm_flows_functions(module):
175291
if hasattr(module, "_execute_single_function_call_async"):
176292
wrap_function_wrapper(module, "_execute_single_function_call_async", wrap__execute_single_function_call_async)

‎tests/mlmodel_googleadk/_test_agent.py‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414

1515
from google.adk.agents import LlmAgent
1616

17-
# Model is hit live during recording; replay mode never reaches the network.
1817
MODEL = "gemini-3.5-flash"
1918
AGENT_NAME = "my_agent"
2019
AGENT_INSTRUCTION = "Answer the user's question in one word."
@@ -24,3 +23,35 @@
2423
def build_agent(tools=None):
2524
"""Return an LlmAgent. tools defaults to none (pure-LLM path)."""
2625
return LlmAgent(name=AGENT_NAME, model=MODEL, instruction=AGENT_INSTRUCTION, tools=tools or [])
26+
27+
28+
agent_recorded_event = [
29+
(
30+
{"type": "LlmAgent"},
31+
{
32+
"id": None,
33+
"name": AGENT_NAME,
34+
"span_id": None,
35+
"trace_id": "trace-id",
36+
"vendor": "google_adk",
37+
"ingest_source": "Python",
38+
"duration": None,
39+
},
40+
)
41+
]
42+
43+
agent_recorded_event_error = [
44+
(
45+
{"type": "LlmAgent"},
46+
{
47+
"id": None,
48+
"name": AGENT_NAME,
49+
"span_id": None,
50+
"trace_id": "trace-id",
51+
"vendor": "google_adk",
52+
"ingest_source": "Python",
53+
"duration": None,
54+
"error": True,
55+
},
56+
)
57+
]

‎tests/mlmodel_googleadk/test_agent.py‎

Lines changed: 83 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,91 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from _test_agent import PROMPT, build_agent
15+
from _test_agent import AGENT_NAME, PROMPT, agent_recorded_event, build_agent
16+
from conftest import GOOGLEADK_VERSION
17+
from testing_support.fixtures import dt_enabled, reset_core_stats_engine, validate_attributes
18+
from testing_support.ml_testing_utils import (
19+
disabled_ai_monitoring_record_content_settings,
20+
disabled_ai_monitoring_settings,
21+
events_with_context_attrs,
22+
)
23+
from testing_support.validators.validate_custom_event import validate_custom_event_count
24+
from testing_support.validators.validate_custom_events import validate_custom_events
25+
from testing_support.validators.validate_span_events import validate_span_events
26+
from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
1627

28+
from newrelic.api.background_task import background_task
29+
from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes
1730

18-
def test_agent_no_harm(exercise_agent, set_trace_info):
31+
EXPECTED_AGENT_METRIC = (f"Llm/agent/GoogleADK/run_async/{AGENT_NAME}", 1)
32+
EXPECTED_SUPPORTABILITY_METRIC = (f"Supportability/Python/ML/GoogleADK/{GOOGLEADK_VERSION}", 1)
33+
34+
35+
@dt_enabled
36+
@reset_core_stats_engine()
37+
@validate_custom_events(events_with_context_attrs(agent_recorded_event))
38+
@validate_custom_event_count(count=4) # Agent, Input, Output, Summary.
39+
@validate_transaction_metrics(
40+
"test_agent:test_agent",
41+
scoped_metrics=[EXPECTED_AGENT_METRIC],
42+
rollup_metrics=[EXPECTED_AGENT_METRIC],
43+
custom_metrics=[EXPECTED_SUPPORTABILITY_METRIC],
44+
background_task=True,
45+
)
46+
@validate_attributes("agent", ["llm"])
47+
@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "my_agent"}'})
48+
@background_task()
49+
def test_agent(exercise_agent, set_trace_info):
50+
set_trace_info()
51+
agent = build_agent()
52+
with WithLlmCustomAttributes({"context": "attr"}):
53+
events = exercise_agent(agent, PROMPT)
54+
55+
assert len(events) == 1
56+
assert events[0].content.parts[0].text == "Paris"
57+
58+
59+
@dt_enabled
60+
@reset_core_stats_engine()
61+
@disabled_ai_monitoring_record_content_settings
62+
@validate_custom_events(agent_recorded_event)
63+
@validate_custom_event_count(count=4) # Agent, Input, Output, Summary.
64+
@validate_transaction_metrics(
65+
"test_agent:test_agent_no_content",
66+
scoped_metrics=[EXPECTED_AGENT_METRIC],
67+
rollup_metrics=[EXPECTED_AGENT_METRIC],
68+
background_task=True,
69+
)
70+
@validate_attributes("agent", ["llm"])
71+
@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "my_agent"}'})
72+
@background_task()
73+
def test_agent_no_content(exercise_agent, set_trace_info):
74+
set_trace_info()
75+
agent = build_agent()
76+
events = exercise_agent(agent, PROMPT)
77+
78+
assert len(events) == 1
79+
assert events[0].content.parts[0].text == "Paris"
80+
81+
82+
@dt_enabled
83+
@reset_core_stats_engine()
84+
@disabled_ai_monitoring_settings
85+
@validate_custom_event_count(count=0)
86+
@validate_transaction_metrics("test_agent:test_agent_disabled_ai_monitoring", background_task=True)
87+
@background_task()
88+
def test_agent_disabled_ai_monitoring(exercise_agent, set_trace_info):
89+
set_trace_info()
90+
agent = build_agent()
91+
events = exercise_agent(agent, PROMPT)
92+
93+
assert len(events) == 1
94+
assert events[0].content.parts[0].text == "Paris"
95+
96+
97+
@reset_core_stats_engine()
98+
@validate_custom_event_count(count=0)
99+
def test_agent_outside_transaction(exercise_agent, set_trace_info):
19100
set_trace_info()
20101
agent = build_agent()
21102
events = exercise_agent(agent, PROMPT)

‎tests/mlmodel_googleadk/test_agent_error.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
# Copyright 2010 New Relic, Inc.
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+
import pytest
16+
from _test_agent import AGENT_NAME, agent_recorded_event_error, build_agent
17+
from testing_support.fixtures import dt_enabled, reset_core_stats_engine, validate_attributes
18+
from testing_support.validators.validate_custom_event import validate_custom_event_count
19+
from testing_support.validators.validate_custom_events import validate_custom_events
20+
from testing_support.validators.validate_error_trace_attributes import validate_error_trace_attributes
21+
from testing_support.validators.validate_span_events import validate_span_events
22+
from testing_support.validators.validate_transaction_error_event_count import validate_transaction_error_event_count
23+
from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
24+
25+
from newrelic.api.background_task import background_task
26+
from newrelic.common.object_names import callable_name
27+
from newrelic.common.object_wrapper import transient_function_wrapper
28+
29+
EXPECTED_AGENT_METRIC = (f"Llm/agent/GoogleADK/run_async/{AGENT_NAME}", 1)
30+
31+
32+
@dt_enabled
33+
@reset_core_stats_engine()
34+
@validate_transaction_error_event_count(1)
35+
@validate_error_trace_attributes(callable_name(ValueError), exact_attrs={"agent": {}, "intrinsic": {}, "user": {}})
36+
@validate_custom_events(agent_recorded_event_error)
37+
@validate_custom_event_count(count=1)
38+
@validate_transaction_metrics(
39+
"test_agent_error:test_agent_error",
40+
scoped_metrics=[EXPECTED_AGENT_METRIC],
41+
rollup_metrics=[EXPECTED_AGENT_METRIC],
42+
background_task=True,
43+
)
44+
@validate_attributes("agent", ["llm"])
45+
@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "my_agent"}'})
46+
@background_task()
47+
def test_agent_error(exercise_agent, set_trace_info):
48+
# Inject a ValueError inside _run_async_impl's async iteration
49+
# so the exception flows through the async generator's athrow.
50+
@transient_function_wrapper("google.adk.flows.llm_flows.base_llm_flow", "BaseLlmFlow.run_async")
51+
def inject_exception(wrapped, instance, args, kwargs):
52+
raise ValueError("Oops")
53+
54+
@inject_exception
55+
def _test():
56+
set_trace_info()
57+
agent = build_agent()
58+
with pytest.raises(ValueError):
59+
exercise_agent(agent, "trigger error")
60+
61+
_test()

0 commit comments

Comments
 (0)