Add subcomponent attribute support for Strands and Autogen agent/tool spans

newrelic/hooks/mlmodel_autogen.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+import json
 import logging
 import sys
 import uuid
@@ -73,8 +73,11 @@ def wrap_on_messages_stream(wrapped, instance, args, kwargs):
     func_name = callable_name(wrapped)
     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/Autogen")
     ft.__enter__()
+    ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data))
 
     try:
         return_val = wrapped(*args, **kwargs)
@@ -180,12 +183,14 @@ async def wrap__execute_tool_call(wrapped, instance, args, kwargs):
     bound_args = bind_args(wrapped, args, kwargs)
     tool_call_data = bound_args.get("tool_call")
     tool_event_dict = _construct_base_tool_event_dict(bound_args, tool_call_data, tool_id, transaction, settings)
-
     tool_name = getattr(tool_call_data, "name", "tool")
-
     func_name = callable_name(wrapped)
+
+    agentic_subcomponent_data = {"type": "APM-AI_TOOL", "name": tool_name}
+
     ft = FunctionTrace(name=f"{func_name}/{tool_name}", group="Llm/tool/Autogen")
     ft.__enter__()
+    ft._add_agent_attribute("subcomponent", json.dumps(agentic_subcomponent_data))
 
     try:
         return_val = await wrapped(*args, **kwargs)

newrelic/hooks/mlmodel_strands.py

@@ -97,6 +97,8 @@ def wrap_stream_async(wrapped, instance, args, kwargs):
     function_trace_name = f"{func_name}/{agent_name}"
     agentic_subcomponent_data = {"type": "APM-AI_AGENT", "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))
@@ -356,6 +358,7 @@ def wrap_tool_executor__stream(wrapped, instance, args, kwargs):
     function_trace_name = f"{func_name}/{tool_name}"
     agentic_subcomponent_data = {"type": "APM-AI_TOOL", "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))

tests/mlmodel_autogen/test_assistant_agent.py

@@ -27,6 +27,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
 
@@ -127,6 +128,8 @@ def add_exclamation(message: str) -> str:
     background_task=True,
 )
 @validate_attributes("agent", ["llm"])
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "pirate_agent"}'})
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'})
 @background_task()
 def test_run_assistant_agent(loop, set_trace_info, single_tool_model_client):
     set_trace_info()
@@ -170,6 +173,8 @@ async def _test():
     background_task=True,
 )
 @validate_attributes("agent", ["llm"])
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "pirate_agent"}'})
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'})
 @background_task()
 def test_run_stream_assistant_agent(loop, set_trace_info, single_tool_model_client):
     set_trace_info()
@@ -221,6 +226,8 @@ async def _test():
     background_task=True,
 )
 @validate_attributes("agent", ["llm"])
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "pirate_agent"}'})
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'})
 @background_task()
 def test_run_assistant_agent_no_content(loop, set_trace_info, single_tool_model_client):
     set_trace_info()
@@ -289,6 +296,8 @@ async def _test():
     background_task=True,
 )
 @validate_attributes("agent", ["llm"])
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "pirate_agent"}'})
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'})
 @background_task()
 def test_run_assistant_agent_error(loop, set_trace_info, single_tool_model_client_error):
     set_trace_info()

tests/mlmodel_autogen/test_teams.py

@@ -28,6 +28,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
 
@@ -213,6 +214,10 @@ def compute_sum(a: int, b: int) -> int:
     background_task=True,
 )
 @validate_attributes("agent", ["llm"])
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "pirate_agent"}'})
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'})
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "robot_agent"}'})
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "compute_sum"}'})
 @background_task()
 def test_run_stream_round_robin_group(loop, set_trace_info, multi_tool_model_client):
     set_trace_info()
@@ -287,6 +292,10 @@ async def _test():
     background_task=True,
 )
 @validate_attributes("agent", ["llm"])
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "pirate_agent"}'})
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'})
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "robot_agent"}'})
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "compute_sum"}'})
 @background_task()
 def test_run_round_robin_group(loop, set_trace_info, multi_tool_model_client):
     set_trace_info()
@@ -358,6 +367,10 @@ async def _test():
     background_task=True,
 )
 @validate_attributes("agent", ["llm"])
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "pirate_agent"}'})
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'})
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "robot_agent"}'})
+@validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "compute_sum"}'})
 @background_task()
 def test_run_round_robin_group_no_content(loop, set_trace_info, multi_tool_model_client):
     set_trace_info()
@@ -456,6 +469,10 @@ async def _test():
     background_task=True,
 )
 @validate_attributes("agent", ["llm"])
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "pirate_agent"}'})
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "add_exclamation"}'})
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "robot_agent"}'})
+@validate_span_events(count=2, exact_agents={"subcomponent": '{"type": "APM-AI_TOOL", "name": "compute_sum"}'})
 @background_task()
 def test_run_round_robin_group_error(loop, set_trace_info, multi_tool_model_client_error):
     set_trace_info()