Add subcomponent attribute to LangChain instrumentation.

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

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)

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,13 +130,15 @@ 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):
     my_agent = create_agent_runnable(tools=[add_exclamation], system_prompt="You are a text manipulation algorithm.")
     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()

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