diff --git a/newrelic/core/attribute.py b/newrelic/core/attribute.py index 79b9a56cb..ed3e5bffa 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/mlmodel_langchain.py b/newrelic/hooks/mlmodel_langchain.py index e682f1bff..a1c22e331 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/tests/mlmodel_langchain/test_agents.py b/tests/mlmodel_langchain/test_agents.py index 9ec7b20df..6a1c471ec 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 19778997d..3ad250fb4 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()