Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion .mega-linter.yml
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,15 @@ CLEAR_REPORT_FOLDER: true
VALIDATE_ALL_CODEBASE: true
IGNORE_GITIGNORED_FILES: true
FAIL_IF_MISSING_LINTER_IN_FLAVOR: true
FILTER_REGEX_EXCLUDE: "(.*/?packages/.*)" # Ignore packages directories
FILTER_REGEX_EXCLUDE: "(.*cassette.*\\.ya?ml)" # Ignore recorded cassette.yaml files
ADDITIONAL_EXCLUDED_DIRECTORIES:
- .pytest_cache
- .ruff_cache
- .tox
- .venv
- megalinter-reports
- newrelic.egg-info
- newrelic/packages

ENABLE_LINTERS: # If you use ENABLE_LINTERS variable, all other linters will be disabled by default
- ACTION_ACTIONLINT
Expand Down
3 changes: 3 additions & 0 deletions newrelic/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -3185,6 +3185,9 @@ def _process_module_builtin_defaults():
"newrelic.hooks.mlmodel_autogen",
"instrument_autogen_agentchat_agents__assistant_agent",
)
_process_module_definition(
"google.adk.agents.llm_agent", "newrelic.hooks.mlmodel_googleadk", "instrument_googleadk_agents_llm_agent"
)
_process_module_definition(
"strands.agent.agent", "newrelic.hooks.mlmodel_strands", "instrument_strands_agent_agent"
)
Expand Down
150 changes: 150 additions & 0 deletions newrelic/hooks/mlmodel_googleadk.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,150 @@
# 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 json
import logging
import sys
import uuid

from newrelic.api.function_trace import FunctionTrace
from newrelic.api.time_trace import get_trace_linking_metadata
from newrelic.api.transaction import current_transaction
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.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."


def wrap_llm_agent__run_async_impl(wrapped, instance, args, kwargs):
transaction = current_transaction()
if not transaction:
return wrapped(*args, **kwargs)

settings = transaction.settings or global_settings()
if not settings.ai_monitoring.enabled:
return wrapped(*args, **kwargs)

transaction.add_ml_model_info("GoogleADK", GOOGLEADK_VERSION)
transaction._add_agent_attribute("llm", True)

agent_name = getattr(instance, "name", "agent")
function_trace_name = f"run_async/{agent_name}"
agentic_subcomponent_data = {"type": "APM-AI_AGENT", "name": agent_name}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Going forward, we will ideally only want to send the subcomponent attr if we have a local agent executing (rather than an invocation by a local or remote agent). Is there a way in the ADK repo to distinguish between remote and local agents and executions vs invocations?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remote agents are handled by a separate class entirely, and it looks like there's a parent_agent field on all agents that can determine if this is a sub agent and not a top level agent.


ft = FunctionTrace(name=function_trace_name, group="Llm/agent/GoogleADK")
ft.__enter__()
ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data))

linking_metadata = get_trace_linking_metadata()
agent_id = str(uuid.uuid4())

try:
return_val = wrapped(*args, **kwargs)
except Exception:
ft.__exit__(*sys.exc_info())
raise

try:
proxied_return_val = AsyncLLMStreamProxy(
wrapped=return_val,
on_stop_iteration=_record_agent_event_on_stop_iteration,
on_error=_handle_agent_streaming_completion_error,
)
proxied_return_val._nr_ft = ft
proxied_return_val._nr_metadata = linking_metadata
proxied_return_val._nr_adk_attrs = {"agent_name": agent_name, "agent_id": agent_id}
return proxied_return_val
except Exception:
ft.__exit__(*sys.exc_info())
return return_val


def _record_agent_event_on_stop_iteration(self, transaction):
if hasattr(self, "_nr_ft"):
linking_metadata = self._nr_metadata or get_trace_linking_metadata()
self._nr_ft.__exit__(None, None, None)
try:
adk_attrs = getattr(self, "_nr_adk_attrs", {})
if not adk_attrs:
return

agent_name = adk_attrs.get("agent_name", "agent")
agent_id = adk_attrs.get("agent_id")
agent_event_dict = _construct_base_agent_event_dict(
agent_name=agent_name, agent_id=agent_id, transaction=transaction, linking_metadata=linking_metadata
)
agent_event_dict["duration"] = self._nr_ft.duration * 1000
transaction.record_custom_event("LlmAgent", agent_event_dict)
except Exception:
_logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE, exc_info=True)
finally:
if hasattr(self, "_nr_adk_attrs"):
self._nr_adk_attrs.clear()


def _handle_agent_streaming_completion_error(self, transaction):
if hasattr(self, "_nr_ft"):
adk_attrs = getattr(self, "_nr_adk_attrs", {})
if not adk_attrs:
self._nr_ft.__exit__(*sys.exc_info())
return

linking_metadata = self._nr_metadata or get_trace_linking_metadata()

try:
agent_name = adk_attrs.get("agent_name", "agent")
agent_id = adk_attrs.get("agent_id")

self._nr_ft.notice_error(attributes={"agent_id": agent_id})
self._nr_ft.__exit__(*sys.exc_info())

agent_event_dict = _construct_base_agent_event_dict(
agent_name=agent_name, agent_id=agent_id, transaction=transaction, linking_metadata=linking_metadata
)
agent_event_dict.update({"duration": self._nr_ft.duration * 1000, "error": True})
transaction.record_custom_event("LlmAgent", agent_event_dict)
except Exception:
_logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE, exc_info=True)
finally:
if hasattr(self, "_nr_adk_attrs"):
self._nr_adk_attrs.clear()


def _construct_base_agent_event_dict(agent_name, agent_id, transaction, linking_metadata):
try:
agent_event_dict = {
"id": agent_id,
"name": agent_name,
"span_id": linking_metadata.get("span.id"),
"trace_id": linking_metadata.get("trace.id"),
"vendor": "google_adk",
"ingest_source": "Python",
}
agent_event_dict.update(_get_llm_metadata(transaction))
except Exception:
_logger.warning(AGENT_EVENT_FAILURE_LOG_MESSAGE, exc_info=True)
agent_event_dict = {}

return agent_event_dict


def instrument_googleadk_agents_llm_agent(module):
if hasattr(module, "LlmAgent") and hasattr(module.LlmAgent, "_run_async_impl"):
Comment thread
umaannamalai marked this conversation as resolved.
wrap_function_wrapper(module, "LlmAgent._run_async_impl", wrap_llm_agent__run_async_impl)
33 changes: 32 additions & 1 deletion tests/mlmodel_googleadk/_test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@

from google.adk.agents import LlmAgent

# Model is hit live during recording; replay mode never reaches the network.
MODEL = "gemini-3.5-flash"
AGENT_NAME = "my_agent"
AGENT_INSTRUCTION = "Answer the user's question in one word."
Expand All @@ -24,3 +23,35 @@
def build_agent(tools=None):
"""Return an LlmAgent. tools defaults to none (pure-LLM path)."""
return LlmAgent(name=AGENT_NAME, model=MODEL, instruction=AGENT_INSTRUCTION, tools=tools or [])


agent_recorded_event = [
(
{"type": "LlmAgent"},
{
"id": None,
"name": AGENT_NAME,
"span_id": None,
"trace_id": "trace-id",
"vendor": "google_adk",
"ingest_source": "Python",
"duration": None,
},
)
]

agent_recorded_event_error = [
(
{"type": "LlmAgent"},
{
"id": None,
"name": AGENT_NAME,
"span_id": None,
"trace_id": "trace-id",
"vendor": "google_adk",
"ingest_source": "Python",
"duration": None,
"error": True,
},
)
]
Loading
Loading