diff --git a/newrelic/config.py b/newrelic/config.py index 45743eb43..62d7aba01 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -3188,6 +3188,11 @@ def _process_module_builtin_defaults(): _process_module_definition( "google.adk.agents.llm_agent", "newrelic.hooks.mlmodel_googleadk", "instrument_googleadk_agents_llm_agent" ) + _process_module_definition( + "google.adk.flows.llm_flows.functions", + "newrelic.hooks.mlmodel_googleadk", + "instrument_googleadk_flows_llm_flows_functions", + ) _process_module_definition( "strands.agent.agent", "newrelic.hooks.mlmodel_strands", "instrument_strands_agent_agent" ) diff --git a/newrelic/hooks/mlmodel_googleadk.py b/newrelic/hooks/mlmodel_googleadk.py index 49e17ef00..eaeae49a0 100644 --- a/newrelic/hooks/mlmodel_googleadk.py +++ b/newrelic/hooks/mlmodel_googleadk.py @@ -23,6 +23,7 @@ from newrelic.common.llm_utils import AsyncLLMStreamProxy, _get_llm_metadata from newrelic.common.object_wrapper import wrap_function_wrapper from newrelic.common.package_version_utils import get_package_version +from newrelic.common.signature import bind_args from newrelic.core.config import global_settings _logger = logging.getLogger(__name__) @@ -30,6 +31,7 @@ RECORD_EVENTS_FAILURE_LOG_MESSAGE = "Exception occurred in Google ADK instrumentation: Failed to record LLM events. Please report this issue to New Relic Support." AGENT_EVENT_FAILURE_LOG_MESSAGE = "Exception occurred in Google ADK instrumentation: Failed to record agent data. Please report this issue to New Relic Support." +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" def wrap_llm_agent__run_async_impl(wrapped, instance, args, kwargs): @@ -145,6 +147,147 @@ def _construct_base_agent_event_dict(agent_name, agent_id, transaction, linking_ return agent_event_dict +async def wrap__execute_single_function_call_async(wrapped, instance, args, kwargs): + transaction = current_transaction() + if not transaction: + return await wrapped(*args, **kwargs) + + settings = transaction.settings or global_settings() + if not settings.ai_monitoring.enabled: + return await wrapped(*args, **kwargs) + + transaction.add_ml_model_info("GoogleADK", GOOGLEADK_VERSION) + transaction._add_agent_attribute("llm", True) + + tool_name = "tool" + run_id = "" + tool_input = None + agent_name = "agent" + is_local_tool = False + try: + bound_args = bind_args(wrapped, args, kwargs) + function_call = bound_args.get("function_call") + agent = bound_args.get("agent") + tools_dict = bound_args.get("tools_dict") + if function_call is not None: + tool_name = getattr(function_call, "name", "tool") or "tool" + run_id = getattr(function_call, "id", "") or "" + tool_input = getattr(function_call, "args", None) + if tools_dict is not None: + from google.adk.tools.function_tool import FunctionTool + + is_local_tool = isinstance(tools_dict.get(tool_name), FunctionTool) + if agent is not None: + agent_name = getattr(agent, "name", "agent") or "agent" + except Exception: + _logger.warning(TOOL_EXTRACTOR_FAILURE_LOG_MESSAGE, exc_info=True) + + function_trace_name = f"execute_single_function_call_async/{tool_name}" + + ft = FunctionTrace(name=function_trace_name, group="Llm/tool/GoogleADK") + ft.__enter__() + if is_local_tool: + agentic_subcomponent_data = {"type": "APM-AI_TOOL", "name": tool_name} + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) + linking_metadata = get_trace_linking_metadata() + tool_id = str(uuid.uuid4()) + + try: + tool_output = await wrapped(*args, **kwargs) + except Exception: + ft.notice_error(attributes={"tool_id": tool_id}) + ft.__exit__(*sys.exc_info()) + try: + tool_event_dict = _construct_base_tool_event_dict( + tool_name=tool_name, + tool_id=tool_id, + run_id=run_id, + tool_input=tool_input, + tool_output=None, + agent_name=agent_name, + error=True, + transaction=transaction, + linking_metadata=linking_metadata, + ) + if tool_event_dict: + tool_event_dict["duration"] = ft.duration * 1000 + transaction.record_custom_event("LlmTool", tool_event_dict) + except Exception: + _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE, exc_info=True) + raise + + ft.__exit__(None, None, None) + try: + response_dict = _extract_tool_response_dict(tool_output) + tool_event_dict = _construct_base_tool_event_dict( + tool_name=tool_name, + tool_id=tool_id, + run_id=run_id, + tool_input=tool_input, + tool_output=response_dict, + agent_name=agent_name, + error=False, + transaction=transaction, + linking_metadata=linking_metadata, + ) + if tool_event_dict: + tool_event_dict["duration"] = ft.duration * 1000 + transaction.record_custom_event("LlmTool", tool_event_dict) + except Exception: + _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE, exc_info=True) + + return tool_output + + +def _extract_tool_response_dict(tool_output): + """Return the dict at content.parts[*].function_response.response, or None.""" + try: + parts = tool_output.content.parts + for part in parts: + function_response = getattr(part, "function_response", None) + if function_response is not None: + return getattr(function_response, "response", None) + except (AttributeError, TypeError): + pass + return None + + +def _construct_base_tool_event_dict( + tool_name, tool_id, run_id, tool_input, tool_output, agent_name, error, transaction, linking_metadata +): + try: + settings = transaction.settings or global_settings() + + tool_event_dict = { + "id": tool_id, + "run_id": run_id, + "name": tool_name, + "span_id": linking_metadata.get("span.id"), + "trace_id": linking_metadata.get("trace.id"), + "agent_name": agent_name, + "vendor": "google_adk", + "ingest_source": "Python", + } + if error: + tool_event_dict["error"] = True + + if settings.ai_monitoring.record_content.enabled: + tool_event_dict["input"] = str(tool_input) if tool_input else None + tool_event_dict["output"] = str(tool_output) if tool_output else None + + tool_event_dict.update(_get_llm_metadata(transaction)) + except Exception: + tool_event_dict = {} + _logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE, exc_info=True) + + return tool_event_dict + + def instrument_googleadk_agents_llm_agent(module): if hasattr(module, "LlmAgent") and hasattr(module.LlmAgent, "_run_async_impl"): wrap_function_wrapper(module, "LlmAgent._run_async_impl", wrap_llm_agent__run_async_impl) + + +def instrument_googleadk_flows_llm_flows_functions(module): + if hasattr(module, "_execute_single_function_call_async"): + wrap_function_wrapper(module, "_execute_single_function_call_async", wrap__execute_single_function_call_async) diff --git a/tests/mlmodel_googleadk/_test_tools.py b/tests/mlmodel_googleadk/_test_tools.py new file mode 100644 index 000000000..f736d4ac4 --- /dev/null +++ b/tests/mlmodel_googleadk/_test_tools.py @@ -0,0 +1,75 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import functools + +from _test_agent import AGENT_NAME +from google.adk.tools import FunctionTool + +TOOL_NAME = "get_capital" + + +def get_capital(country: str) -> dict: + """Return the capital of a country.""" + capitals = {"France": "Paris", "Japan": "Tokyo"} + return {"output": capitals.get(country, "Unknown")} + + +@functools.wraps(get_capital) # Make function names match +def _raising_capital(country: str) -> dict: + """Return the capital of a country.""" + raise ValueError("intentional tool failure") + + +get_capital_tool = FunctionTool(func=get_capital) +raising_tool = FunctionTool(func=_raising_capital) + +EXPECTED_TOOL_INPUT_STR = "{'country': 'France'}" +EXPECTED_TOOL_OUTPUT_STR = "{'output': 'Paris'}" + + +def tool_recorded_event(record_content: bool): + base = { + "id": None, + "run_id": None, + "name": TOOL_NAME, + "span_id": None, + "trace_id": "trace-id", + "agent_name": AGENT_NAME, + "vendor": "google_adk", + "ingest_source": "Python", + "duration": None, + } + if record_content: + base["input"] = EXPECTED_TOOL_INPUT_STR + base["output"] = EXPECTED_TOOL_OUTPUT_STR + return [({"type": "LlmTool"}, base)] + + +def tool_recorded_event_error(record_content: bool): + base = { + "id": None, + "run_id": None, + "name": TOOL_NAME, + "span_id": None, + "trace_id": "trace-id", + "agent_name": AGENT_NAME, + "vendor": "google_adk", + "ingest_source": "Python", + "duration": None, + "error": True, + } + if record_content: + base["input"] = EXPECTED_TOOL_INPUT_STR + return [({"type": "LlmTool"}, base)] diff --git a/tests/mlmodel_googleadk/test_tools.py b/tests/mlmodel_googleadk/test_tools.py new file mode 100644 index 000000000..07edfa899 --- /dev/null +++ b/tests/mlmodel_googleadk/test_tools.py @@ -0,0 +1,110 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from _test_agent import AGENT_NAME, PROMPT, agent_recorded_event, build_agent +from _test_tools import TOOL_NAME, get_capital_tool, tool_recorded_event +from testing_support.fixtures import dt_enabled, reset_core_stats_engine, validate_attributes +from testing_support.ml_testing_utils import ( + disabled_ai_monitoring_record_content_settings, + disabled_ai_monitoring_settings, + events_with_context_attrs, +) +from testing_support.validators.validate_custom_event import validate_custom_event_count +from testing_support.validators.validate_custom_events import validate_custom_events +from testing_support.validators.validate_span_events import validate_span_events +from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics + +from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes +from newrelic.hooks.mlmodel_googleadk import GOOGLEADK_VERSION + +EXPECTED_AGENT_METRIC = (f"Llm/agent/GoogleADK/run_async/{AGENT_NAME}", 1) +EXPECTED_TOOL_METRIC = (f"Llm/tool/GoogleADK/execute_single_function_call_async/{TOOL_NAME}", 1) +EXPECTED_SUPPORTABILITY_METRIC = (f"Supportability/Python/ML/GoogleADK/{GOOGLEADK_VERSION}", 1) + + +def _validate_events(events): + assert len(events) == 3 + assert events[0].content.parts[0].function_call.name == "get_capital" + assert events[1].content.parts[0].function_response.response["output"] == "Paris" + assert events[2].content.parts[0].text == "Paris" + + +@dt_enabled +@reset_core_stats_engine() +@validate_custom_events(events_with_context_attrs(agent_recorded_event + tool_recorded_event(record_content=True))) +@validate_custom_event_count(count=8) # LlmAgent, LlmTool, 2x (Summary, Input, Output) for the two LLM calls. +@validate_transaction_metrics( + "test_tools:test_tool", + scoped_metrics=[EXPECTED_AGENT_METRIC, EXPECTED_TOOL_METRIC], + rollup_metrics=[EXPECTED_AGENT_METRIC, EXPECTED_TOOL_METRIC], + custom_metrics=[EXPECTED_SUPPORTABILITY_METRIC], + background_task=True, +) +@validate_attributes("agent", ["llm"]) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "my_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": f'{{"type": "APM-AI_TOOL", "name": "{TOOL_NAME}"}}'}) +@background_task() +def test_tool(exercise_agent, set_trace_info): + set_trace_info() + agent = build_agent(tools=[get_capital_tool]) + with WithLlmCustomAttributes({"context": "attr"}): + events = exercise_agent(agent, PROMPT) + + _validate_events(events) + + +@dt_enabled +@reset_core_stats_engine() +@disabled_ai_monitoring_record_content_settings +@validate_custom_events(agent_recorded_event + tool_recorded_event(record_content=False)) +@validate_custom_event_count(count=8) # LlmAgent, LlmTool, 2x (Summary, Input, Output) for the two LLM calls. +@validate_transaction_metrics( + "test_tools:test_tool_no_content", + scoped_metrics=[EXPECTED_AGENT_METRIC, EXPECTED_TOOL_METRIC], + rollup_metrics=[EXPECTED_AGENT_METRIC, EXPECTED_TOOL_METRIC], + background_task=True, +) +@validate_attributes("agent", ["llm"]) +@background_task() +def test_tool_no_content(exercise_agent, set_trace_info): + set_trace_info() + agent = build_agent(tools=[get_capital_tool]) + events = exercise_agent(agent, PROMPT) + + _validate_events(events) + + +@dt_enabled +@reset_core_stats_engine() +@disabled_ai_monitoring_settings +@validate_custom_event_count(count=0) +@validate_transaction_metrics("test_tools:test_tool_disabled_ai_monitoring", background_task=True) +@background_task() +def test_tool_disabled_ai_monitoring(exercise_agent, set_trace_info): + set_trace_info() + agent = build_agent(tools=[get_capital_tool]) + events = exercise_agent(agent, PROMPT) + + _validate_events(events) + + +@reset_core_stats_engine() +@validate_custom_event_count(count=0) +def test_tool_outside_transaction(exercise_agent, set_trace_info): + set_trace_info() + agent = build_agent(tools=[get_capital_tool]) + events = exercise_agent(agent, PROMPT) + + _validate_events(events) diff --git a/tests/mlmodel_googleadk/test_tools_error.py b/tests/mlmodel_googleadk/test_tools_error.py new file mode 100644 index 000000000..510a5e201 --- /dev/null +++ b/tests/mlmodel_googleadk/test_tools_error.py @@ -0,0 +1,53 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from _test_agent import AGENT_NAME, PROMPT, agent_recorded_event_error, build_agent +from _test_tools import TOOL_NAME, raising_tool, tool_recorded_event_error +from testing_support.fixtures import dt_enabled, reset_core_stats_engine, validate_attributes +from testing_support.validators.validate_custom_event import validate_custom_event_count +from testing_support.validators.validate_custom_events import validate_custom_events +from testing_support.validators.validate_error_trace_attributes import validate_error_trace_attributes +from testing_support.validators.validate_span_events import validate_span_events +from testing_support.validators.validate_transaction_error_event_count import validate_transaction_error_event_count +from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics + +from newrelic.api.background_task import background_task +from newrelic.common.object_names import callable_name + +EXPECTED_AGENT_METRIC = (f"Llm/agent/GoogleADK/run_async/{AGENT_NAME}", 1) +EXPECTED_TOOL_METRIC = (f"Llm/tool/GoogleADK/execute_single_function_call_async/{TOOL_NAME}", 1) + + +@dt_enabled +@reset_core_stats_engine() +@validate_transaction_error_event_count(1) +@validate_error_trace_attributes(callable_name(ValueError), exact_attrs={"agent": {}, "intrinsic": {}, "user": {}}) +@validate_custom_events(agent_recorded_event_error + tool_recorded_event_error(record_content=True)) +@validate_custom_event_count(count=5) # LlmAgent, LlmTool, Summary, Input, Output +@validate_transaction_metrics( + "test_tools_error:test_tool_error", + scoped_metrics=[EXPECTED_AGENT_METRIC, EXPECTED_TOOL_METRIC], + rollup_metrics=[EXPECTED_AGENT_METRIC, EXPECTED_TOOL_METRIC], + background_task=True, +) +@validate_attributes("agent", ["llm"]) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "my_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": f'{{"type": "APM-AI_TOOL", "name": "{TOOL_NAME}"}}'}) +@background_task() +def test_tool_error(exercise_agent, set_trace_info): + set_trace_info() + agent = build_agent(tools=[raising_tool]) + with pytest.raises(ValueError): + exercise_agent(agent, PROMPT)