diff --git a/newrelic/core/attribute.py b/newrelic/core/attribute.py index 79b9a56cb2..ed3e5bffa6 100644 --- a/newrelic/core/attribute.py +++ b/newrelic/core/attribute.py @@ -100,6 +100,7 @@ "response.headers.contentType", "response.status", "server.address", + "subcomponent", "zeebe.client.bpmnProcessId", "zeebe.client.messageName", "zeebe.client.correlationKey", diff --git a/newrelic/hooks/adapter_mcp.py b/newrelic/hooks/adapter_mcp.py index bcc8ae0a39..e891df0325 100644 --- a/newrelic/hooks/adapter_mcp.py +++ b/newrelic/hooks/adapter_mcp.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import logging from newrelic.api.function_trace import FunctionTrace @@ -37,8 +38,10 @@ async def wrap_call_tool(wrapped, instance, args, kwargs): bound_args = bind_args(wrapped, args, kwargs) tool_name = bound_args.get("name") or "tool" function_trace_name = f"{func_name}/{tool_name}" + agentic_subcomponent_data = {"type": "APM-AI_TOOL", "name": tool_name} - with FunctionTrace(name=function_trace_name, group="Llm/tool/MCP", source=wrapped): + with FunctionTrace(name=function_trace_name, group="Llm/tool/MCP", source=wrapped) as ft: + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) return await wrapped(*args, **kwargs) diff --git a/newrelic/hooks/mlmodel_langchain.py b/newrelic/hooks/mlmodel_langchain.py index e682f1bff3..a1c22e331f 100644 --- a/newrelic/hooks/mlmodel_langchain.py +++ b/newrelic/hooks/mlmodel_langchain.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import logging import sys import time @@ -161,9 +162,11 @@ def invoke(self, *args, **kwargs): agent_id = str(uuid.uuid4()) agent_event_dict = _construct_base_agent_event_dict(agent_name, agent_id, transaction) function_trace_name = f"invoke/{agent_name}" + agentic_subcomponent_data = {"type": "APM-AI_AGENT", "name": agent_name} ft = FunctionTrace(name=function_trace_name, group="Llm/agent/LangChain") ft.__enter__() + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) try: return_val = self.__wrapped__.invoke(*args, **kwargs) except Exception: @@ -189,9 +192,11 @@ async def ainvoke(self, *args, **kwargs): agent_id = str(uuid.uuid4()) agent_event_dict = _construct_base_agent_event_dict(agent_name, agent_id, transaction) function_trace_name = f"ainvoke/{agent_name}" + agentic_subcomponent_data = {"type": "APM-AI_AGENT", "name": agent_name} ft = FunctionTrace(name=function_trace_name, group="Llm/agent/LangChain") ft.__enter__() + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) try: return_val = await self.__wrapped__.ainvoke(*args, **kwargs) except Exception: @@ -217,9 +222,11 @@ def stream(self, *args, **kwargs): agent_id = str(uuid.uuid4()) agent_event_dict = _construct_base_agent_event_dict(agent_name, agent_id, transaction) function_trace_name = f"stream/{agent_name}" + agentic_subcomponent_data = {"type": "APM-AI_AGENT", "name": agent_name} ft = FunctionTrace(name=function_trace_name, group="Llm/agent/LangChain") ft.__enter__() + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) try: return_val = self.__wrapped__.stream(*args, **kwargs) return_val = GeneratorProxy( @@ -242,9 +249,11 @@ def astream(self, *args, **kwargs): agent_id = str(uuid.uuid4()) agent_event_dict = _construct_base_agent_event_dict(agent_name, agent_id, transaction) function_trace_name = f"astream/{agent_name}" + agentic_subcomponent_data = {"type": "APM-AI_AGENT", "name": agent_name} ft = FunctionTrace(name=function_trace_name, group="Llm/agent/LangChain") ft.__enter__() + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) try: return_val = self.__wrapped__.astream(*args, **kwargs) return_val = AsyncGeneratorProxy( @@ -267,9 +276,11 @@ def transform(self, *args, **kwargs): agent_id = str(uuid.uuid4()) agent_event_dict = _construct_base_agent_event_dict(agent_name, agent_id, transaction) function_trace_name = f"stream/{agent_name}" + agentic_subcomponent_data = {"type": "APM-AI_AGENT", "name": agent_name} ft = FunctionTrace(name=function_trace_name, group="Llm/agent/LangChain") ft.__enter__() + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) try: return_val = self.__wrapped__.transform(*args, **kwargs) return_val = GeneratorProxy( @@ -292,9 +303,11 @@ def atransform(self, *args, **kwargs): agent_id = str(uuid.uuid4()) agent_event_dict = _construct_base_agent_event_dict(agent_name, agent_id, transaction) function_trace_name = f"astream/{agent_name}" + agentic_subcomponent_data = {"type": "APM-AI_AGENT", "name": agent_name} ft = FunctionTrace(name=function_trace_name, group="Llm/agent/LangChain") ft.__enter__() + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) try: return_val = self.__wrapped__.atransform(*args, **kwargs) return_val = AsyncGeneratorProxy( @@ -512,8 +525,11 @@ def wrap_tool_sync_run(wrapped, instance, args, kwargs): except Exception: filtered_tool_input = tool_input + agentic_subcomponent_data = {"type": "APM-AI_TOOL", "name": tool_name} + ft = FunctionTrace(name=f"{wrapped.__name__}/{tool_name}", group="Llm/tool/LangChain") ft.__enter__() + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) linking_metadata = get_trace_linking_metadata() try: return_val = wrapped(**run_args) @@ -573,8 +589,11 @@ async def wrap_tool_async_run(wrapped, instance, args, kwargs): except Exception: filtered_tool_input = tool_input + agentic_subcomponent_data = {"type": "APM-AI_TOOL", "name": tool_name} + ft = FunctionTrace(name=f"{wrapped.__name__}/{tool_name}", group="Llm/tool/LangChain") ft.__enter__() + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) linking_metadata = get_trace_linking_metadata() try: return_val = await wrapped(**run_args) diff --git a/newrelic/hooks/mlmodel_strands.py b/newrelic/hooks/mlmodel_strands.py index a4ac6e5d72..bc045df190 100644 --- a/newrelic/hooks/mlmodel_strands.py +++ b/newrelic/hooks/mlmodel_strands.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import logging import sys import uuid @@ -94,9 +95,12 @@ def wrap_stream_async(wrapped, instance, args, kwargs): func_name = callable_name(wrapped) agent_name = getattr(instance, "name", "agent") function_trace_name = f"{func_name}/{agent_name}" + agentic_subcomponent_data = {"type": "APM-AI_AGENT", "name": agent_name} ft = FunctionTrace(name=function_trace_name, group="Llm/agent/Strands") ft.__enter__() + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) + linking_metadata = get_trace_linking_metadata() agent_id = str(uuid.uuid4()) @@ -105,7 +109,6 @@ def wrap_stream_async(wrapped, instance, args, kwargs): except Exception: raise - # For streaming responses, wrap with proxy and attach metadata try: # For streaming responses, wrap with proxy and attach metadata proxied_return_val = AsyncGeneratorProxy( @@ -126,7 +129,6 @@ def _record_agent_event_on_stop_iteration(self, transaction): # Use saved linking metadata to maintain correct span association linking_metadata = self._nr_metadata or get_trace_linking_metadata() self._nr_ft.__exit__(None, None, None) - try: strands_attrs = getattr(self, "_nr_strands_attrs", {}) @@ -352,9 +354,11 @@ def wrap_tool_executor__stream(wrapped, instance, args, kwargs): func_name = callable_name(wrapped) function_trace_name = f"{func_name}/{tool_name}" + agentic_subcomponent_data = {"type": "APM-AI_TOOL", "name": tool_name} ft = FunctionTrace(name=function_trace_name, group="Llm/tool/Strands") ft.__enter__() + ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data)) linking_metadata = get_trace_linking_metadata() tool_id = str(uuid.uuid4()) diff --git a/tests/adapter_mcp/test_mcp.py b/tests/adapter_mcp/test_mcp.py index 5ba6a81074..5424b57ca7 100644 --- a/tests/adapter_mcp/test_mcp.py +++ b/tests/adapter_mcp/test_mcp.py @@ -19,6 +19,7 @@ from mcp.server.fastmcp.tools import ToolManager from testing_support.ml_testing_utils import disabled_ai_monitoring_settings from testing_support.validators.validate_function_not_called import validate_function_not_called +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 @@ -57,6 +58,7 @@ def echo_prompt(message: str): rollup_metrics=[("Llm/tool/MCP/mcp.client.session:ClientSession.call_tool/add_exclamation", 1)], background_task=True, ) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task() def test_tool_tracing_via_client_session(loop, fastmcp_server): async def _test(): @@ -75,6 +77,7 @@ async def _test(): rollup_metrics=[("Llm/tool/MCP/mcp.server.fastmcp.tools.tool_manager:ToolManager.call_tool/add_exclamation", 1)], background_task=True, ) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task() def test_tool_tracing_via_tool_manager(loop): async def _test(): diff --git a/tests/mlmodel_langchain/test_agents.py b/tests/mlmodel_langchain/test_agents.py index 9ec7b20dff..6a1c471ecd 100644 --- a/tests/mlmodel_langchain/test_agents.py +++ b/tests/mlmodel_langchain/test_agents.py @@ -15,7 +15,7 @@ import pytest from langchain.messages import HumanMessage from langchain.tools import tool -from testing_support.fixtures import reset_core_stats_engine, validate_attributes +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, @@ -24,6 +24,7 @@ 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 @@ -76,6 +77,7 @@ def add_exclamation(message: str) -> str: return f"{message}!" +@dt_enabled @reset_core_stats_engine() def test_agent(exercise_agent, create_agent_runnable, set_trace_info, method_name): @validate_custom_events(events_with_context_attrs(agent_recorded_event)) @@ -87,6 +89,8 @@ def test_agent(exercise_agent, create_agent_runnable, set_trace_info, method_nam 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task(name="test_agent") def _test(): set_trace_info() @@ -100,6 +104,7 @@ def _test(): _test() +@dt_enabled @reset_core_stats_engine() @disabled_ai_monitoring_record_content_settings def test_agent_no_content(exercise_agent, create_agent_runnable, set_trace_info, method_name): @@ -112,6 +117,8 @@ def test_agent_no_content(exercise_agent, create_agent_runnable, set_trace_info, 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task(name="test_agent_no_content") def _test(): set_trace_info() @@ -123,6 +130,7 @@ def _test(): _test() +@dt_enabled @reset_core_stats_engine() @validate_custom_event_count(count=0) def test_agent_outside_txn(exercise_agent, create_agent_runnable): @@ -130,6 +138,7 @@ def test_agent_outside_txn(exercise_agent, create_agent_runnable): exercise_agent(my_agent, PROMPT) +@dt_enabled @disabled_ai_monitoring_settings @reset_core_stats_engine() @validate_custom_event_count(count=0) @@ -140,6 +149,7 @@ def test_agent_disabled_ai_monitoring_events(exercise_agent, create_agent_runnab exercise_agent(my_agent, PROMPT) +@dt_enabled @reset_core_stats_engine() def test_agent_execution_error(exercise_agent, create_agent_runnable, set_trace_info, method_name, agent_runnable_type): # Add a wrapper to intentionally force an error in the Agent code @@ -159,6 +169,8 @@ def inject_exception(wrapped, instance, args, kwargs): background_task=True, ) @validate_attributes("agent", ["llm"]) + # Only an agent span is expected here and not a tool because the error is injected before the tool is called + @validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "my_agent"}'}) @background_task(name="test_agent_execution_error") def _test(): set_trace_info() diff --git a/tests/mlmodel_langchain/test_tools.py b/tests/mlmodel_langchain/test_tools.py index 19778997db..3ad250fb45 100644 --- a/tests/mlmodel_langchain/test_tools.py +++ b/tests/mlmodel_langchain/test_tools.py @@ -14,7 +14,7 @@ import pytest from langchain.messages import HumanMessage -from testing_support.fixtures import reset_core_stats_engine, validate_attributes +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, events_with_context_attrs, @@ -23,6 +23,7 @@ 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 @@ -95,6 +96,7 @@ ] +@dt_enabled @reset_core_stats_engine() def test_tool(exercise_agent, set_trace_info, create_agent_runnable, add_exclamation, tool_method_name): @validate_custom_events(events_with_context_attrs(tool_recorded_event)) @@ -106,6 +108,8 @@ def test_tool(exercise_agent, set_trace_info, create_agent_runnable, add_exclama 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task(name="test_tool") def _test(): set_trace_info() @@ -119,6 +123,7 @@ def _test(): _test() +@dt_enabled @reset_core_stats_engine() @disabled_ai_monitoring_record_content_settings def test_tool_no_content(exercise_agent, set_trace_info, create_agent_runnable, add_exclamation, tool_method_name): @@ -131,6 +136,8 @@ def test_tool_no_content(exercise_agent, set_trace_info, create_agent_runnable, 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task(name="test_tool_no_content") def _test(): set_trace_info() @@ -142,6 +149,7 @@ def _test(): _test() +@dt_enabled @reset_core_stats_engine() def test_tool_execution_error(exercise_agent, set_trace_info, create_agent_runnable, add_exclamation, tool_method_name): @validate_transaction_error_event_count(1) @@ -157,6 +165,8 @@ def test_tool_execution_error(exercise_agent, set_trace_info, create_agent_runna 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task(name="test_tool_execution_error") def _test(): set_trace_info() @@ -169,6 +179,7 @@ def _test(): _test() +@dt_enabled @reset_core_stats_engine() def test_tool_pre_execution_exception( exercise_agent, set_trace_info, create_agent_runnable, add_exclamation, tool_method_name @@ -190,6 +201,8 @@ def inject_exception(wrapped, instance, args, kwargs): 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task(name="test_tool_pre_execution_exception") def _test(): set_trace_info() diff --git a/tests/mlmodel_strands/test_agents.py b/tests/mlmodel_strands/test_agents.py index b0a1965eea..93d635a716 100644 --- a/tests/mlmodel_strands/test_agents.py +++ b/tests/mlmodel_strands/test_agents.py @@ -14,7 +14,7 @@ import pytest from strands import Agent -from testing_support.fixtures import reset_core_stats_engine, validate_attributes +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, @@ -23,6 +23,7 @@ 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 @@ -65,6 +66,7 @@ ] +@dt_enabled @reset_core_stats_engine() @validate_custom_events(events_with_context_attrs(agent_recorded_event)) @validate_custom_event_count(count=2) @@ -75,6 +77,8 @@ 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task() def test_agent(exercise_agent, set_trace_info, single_tool_model): set_trace_info() @@ -97,6 +101,7 @@ def test_agent(exercise_agent, set_trace_info, single_tool_model): assert response.metrics.tool_metrics["add_exclamation"].success_count == 1 +@dt_enabled @reset_core_stats_engine() @disabled_ai_monitoring_record_content_settings @validate_custom_events(agent_recorded_event) @@ -108,6 +113,8 @@ def test_agent(exercise_agent, set_trace_info, single_tool_model): 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task() def test_agent_no_content(exercise_agent, set_trace_info, single_tool_model): set_trace_info() @@ -129,6 +136,7 @@ def test_agent_no_content(exercise_agent, set_trace_info, single_tool_model): assert response.metrics.tool_metrics["add_exclamation"].success_count == 1 +@dt_enabled @reset_core_stats_engine() @validate_custom_event_count(count=0) def test_agent_outside_txn(exercise_agent, single_tool_model): @@ -150,6 +158,7 @@ def test_agent_outside_txn(exercise_agent, single_tool_model): assert response.metrics.tool_metrics["add_exclamation"].success_count == 1 +@dt_enabled @disabled_ai_monitoring_settings @reset_core_stats_engine() @validate_custom_event_count(count=0) @@ -174,6 +183,7 @@ def test_agent_disabled_ai_monitoring_events(exercise_agent, set_trace_info, sin assert response.metrics.tool_metrics["add_exclamation"].success_count == 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": {}}) @@ -186,6 +196,8 @@ def test_agent_disabled_ai_monitoring_events(exercise_agent, set_trace_info, sin background_task=True, ) @validate_attributes("agent", ["llm"]) +# Only an agent span is expected here and not a tool because the error is injected before the tool is called +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "my_agent"}'}) @background_task() def test_agent_execution_error(exercise_agent, set_trace_info, single_tool_model): # Add a wrapper to intentionally force an error in the Agent code diff --git a/tests/mlmodel_strands/test_multiagent_graph.py b/tests/mlmodel_strands/test_multiagent_graph.py index 7bd84fc901..216a6bd3a5 100644 --- a/tests/mlmodel_strands/test_multiagent_graph.py +++ b/tests/mlmodel_strands/test_multiagent_graph.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from testing_support.fixtures import reset_core_stats_engine, validate_attributes +from testing_support.fixtures import dt_enabled, reset_core_stats_engine, validate_attributes from testing_support.ml_testing_utils import 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 @@ -86,6 +87,7 @@ ] +@dt_enabled @reset_core_stats_engine() @validate_custom_events(events_with_context_attrs(tool_recorded_events)) @validate_custom_events(events_with_context_attrs(agent_recorded_events)) @@ -107,6 +109,10 @@ background_task=True, ) @validate_attributes("agent", ["llm"]) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "math_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "calculate_sum"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "analysis_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "analyze_result"}'}) @background_task() def test_multiagent_graph_invoke(set_trace_info, agent_graph): set_trace_info() @@ -123,6 +129,7 @@ def test_multiagent_graph_invoke(set_trace_info, agent_graph): ) +@dt_enabled @reset_core_stats_engine() @validate_custom_events(tool_recorded_events) @validate_custom_events(agent_recorded_events) @@ -144,6 +151,10 @@ def test_multiagent_graph_invoke(set_trace_info, agent_graph): background_task=True, ) @validate_attributes("agent", ["llm"]) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "math_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "calculate_sum"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "analysis_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "analyze_result"}'}) @background_task() def test_multiagent_graph_invoke_async(loop, set_trace_info, agent_graph): set_trace_info() @@ -162,6 +173,7 @@ async def _test(): loop.run_until_complete(_test()) +@dt_enabled @reset_core_stats_engine() @validate_custom_events(tool_recorded_events) @validate_custom_events(agent_recorded_events) @@ -183,6 +195,10 @@ async def _test(): background_task=True, ) @validate_attributes("agent", ["llm"]) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "math_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "calculate_sum"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "analysis_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "analyze_result"}'}) @background_task() def test_multiagent_graph_stream_async(loop, set_trace_info, agent_graph): set_trace_info() @@ -201,6 +217,7 @@ async def _test(): loop.run_until_complete(_test()) +@dt_enabled @disabled_ai_monitoring_settings @reset_core_stats_engine() @validate_custom_event_count(count=0) @@ -219,6 +236,7 @@ def test_multiagent_graph_invoke_disabled_ai_monitoring_events(set_trace_info, a ) +@dt_enabled @reset_core_stats_engine() @validate_custom_event_count(count=0) def test_multiagent_graph_invoke_outside_txn(agent_graph): diff --git a/tests/mlmodel_strands/test_multiagent_swarm.py b/tests/mlmodel_strands/test_multiagent_swarm.py index bbcbb3e27c..dc4ee95628 100644 --- a/tests/mlmodel_strands/test_multiagent_swarm.py +++ b/tests/mlmodel_strands/test_multiagent_swarm.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from testing_support.fixtures import reset_core_stats_engine, validate_attributes +from testing_support.fixtures import dt_enabled, reset_core_stats_engine, validate_attributes from testing_support.ml_testing_utils import 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 @@ -106,6 +107,7 @@ ] +@dt_enabled @reset_core_stats_engine() @validate_custom_events(events_with_context_attrs(tool_recorded_events)) @validate_custom_events(events_with_context_attrs(agent_recorded_events)) @@ -128,6 +130,10 @@ background_task=True, ) @validate_attributes("agent", ["llm"]) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "math_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "calculate_sum"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "analysis_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "analyze_result"}'}) @background_task() def test_multiagent_swarm_invoke(set_trace_info, agent_swarm): set_trace_info() @@ -145,6 +151,7 @@ def test_multiagent_swarm_invoke(set_trace_info, agent_swarm): ) +@dt_enabled @reset_core_stats_engine() @validate_custom_events(tool_recorded_events) @validate_custom_events(agent_recorded_events) @@ -167,6 +174,10 @@ def test_multiagent_swarm_invoke(set_trace_info, agent_swarm): background_task=True, ) @validate_attributes("agent", ["llm"]) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "math_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "calculate_sum"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "analysis_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "analyze_result"}'}) @background_task() def test_multiagent_swarm_invoke_async(loop, set_trace_info, agent_swarm): set_trace_info() @@ -186,6 +197,7 @@ async def _test(): loop.run_until_complete(_test()) +@dt_enabled @reset_core_stats_engine() @validate_custom_events(tool_recorded_events) @validate_custom_events(agent_recorded_events) @@ -208,6 +220,10 @@ async def _test(): background_task=True, ) @validate_attributes("agent", ["llm"]) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "math_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "calculate_sum"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "analysis_agent"}'}) +@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "analyze_result"}'}) @background_task() def test_multiagent_swarm_stream_async(loop, set_trace_info, agent_swarm): set_trace_info() @@ -226,6 +242,7 @@ async def _test(): loop.run_until_complete(_test()) +@dt_enabled @disabled_ai_monitoring_settings @reset_core_stats_engine() @validate_custom_event_count(count=0) @@ -245,6 +262,7 @@ def test_multiagent_swarm_invoke_disabled_ai_monitoring_events(set_trace_info, a ) +@dt_enabled @reset_core_stats_engine() @validate_custom_event_count(count=0) def test_multiagent_swarm_invoke_outside_txn(agent_swarm): diff --git a/tests/mlmodel_strands/test_tools.py b/tests/mlmodel_strands/test_tools.py index a5e62ff3a3..bc9eb233c7 100644 --- a/tests/mlmodel_strands/test_tools.py +++ b/tests/mlmodel_strands/test_tools.py @@ -14,7 +14,7 @@ import pytest from strands import Agent -from testing_support.fixtures import reset_core_stats_engine, validate_attributes +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, events_with_context_attrs, @@ -23,6 +23,7 @@ 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 @@ -94,6 +95,7 @@ EXPECTED_ERROR_MESSAGES = ["Error: RuntimeError - Oops", "Error: Oops"] +@dt_enabled @reset_core_stats_engine() @validate_custom_events(events_with_context_attrs(tool_recorded_event)) @validate_custom_event_count(count=2) @@ -104,6 +106,8 @@ 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task() def test_tool(exercise_agent, set_trace_info, single_tool_model, add_exclamation): set_trace_info() @@ -126,6 +130,7 @@ def test_tool(exercise_agent, set_trace_info, single_tool_model, add_exclamation assert response.metrics.tool_metrics["add_exclamation"].success_count == 1 +@dt_enabled @reset_core_stats_engine() @disabled_ai_monitoring_record_content_settings @validate_custom_events(tool_events_sans_content(tool_recorded_event)) @@ -137,6 +142,8 @@ def test_tool(exercise_agent, set_trace_info, single_tool_model, add_exclamation 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task() def test_tool_no_content(exercise_agent, set_trace_info, single_tool_model, add_exclamation): set_trace_info() @@ -158,6 +165,7 @@ def test_tool_no_content(exercise_agent, set_trace_info, single_tool_model, add_ assert response.metrics.tool_metrics["add_exclamation"].success_count == 1 +@dt_enabled @reset_core_stats_engine() def test_tool_execution_error(exercise_agent, set_trace_info, single_tool_model_error, add_exclamation): from strands.tools import PythonAgentTool @@ -178,6 +186,8 @@ def test_tool_execution_error(exercise_agent, set_trace_info, single_tool_model_ 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task(name="test_tool_execution_error") def _test(): set_trace_info() @@ -201,6 +211,7 @@ def _test(): _test() +@dt_enabled @reset_core_stats_engine() @validate_transaction_error_event_count(1) @validate_error_trace_attributes(callable_name(ValueError), exact_attrs={"agent": {}, "intrinsic": {}, "user": {}}) @@ -213,6 +224,8 @@ def _test(): 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": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'}) @background_task() def test_tool_pre_execution_exception(exercise_agent, set_trace_info, single_tool_model, add_exclamation): # Add a wrapper to intentionally force an error in the ToolExecutor._stream code to hit the exception path in