Google ADK Tool Instrumentation

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"
     )

newrelic/hooks/mlmodel_googleadk.py

@@ -23,13 +23,15 @@
 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__)
 GOOGLEADK_VERSION = get_package_version("google-adk")
 
 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)

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)]

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)