diff --git a/newrelic/config.py b/newrelic/config.py index 94955293d..4b8627772 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2948,12 +2948,26 @@ def _process_module_builtin_defaults(): "newrelic.hooks.mlmodel_autogen", "instrument_autogen_agentchat_agents__assistant_agent", ) - _process_module_definition("strands.agent.agent", "newrelic.hooks.mlmodel_strands", "instrument_agent_agent") _process_module_definition( - "strands.tools.executors._executor", "newrelic.hooks.mlmodel_strands", "instrument_tools_executors__executor" + "strands.agent.agent", "newrelic.hooks.mlmodel_strands", "instrument_strands_agent_agent" + ) + _process_module_definition( + "strands.multiagent.graph", "newrelic.hooks.mlmodel_strands", "instrument_strands_multiagent_graph" + ) + _process_module_definition( + "strands.multiagent.swarm", "newrelic.hooks.mlmodel_strands", "instrument_strands_multiagent_swarm" + ) + _process_module_definition( + "strands.tools.executors._executor", + "newrelic.hooks.mlmodel_strands", + "instrument_strands_tools_executors__executor", + ) + _process_module_definition( + "strands.tools.registry", "newrelic.hooks.mlmodel_strands", "instrument_strands_tools_registry" + ) + _process_module_definition( + "strands.models.bedrock", "newrelic.hooks.mlmodel_strands", "instrument_strands_models_bedrock" ) - _process_module_definition("strands.tools.registry", "newrelic.hooks.mlmodel_strands", "instrument_tools_registry") - _process_module_definition("strands.models.bedrock", "newrelic.hooks.mlmodel_strands", "instrument_models_bedrock") _process_module_definition("mcp.client.session", "newrelic.hooks.adapter_mcp", "instrument_mcp_client_session") _process_module_definition( diff --git a/newrelic/hooks/mlmodel_strands.py b/newrelic/hooks/mlmodel_strands.py index bf849fd71..20317626d 100644 --- a/newrelic/hooks/mlmodel_strands.py +++ b/newrelic/hooks/mlmodel_strands.py @@ -461,7 +461,7 @@ def wrap_bedrock_model__stream(wrapped, instance, args, kwargs): return wrapped(*args, **kwargs) -def instrument_agent_agent(module): +def instrument_strands_agent_agent(module): if hasattr(module, "Agent"): if hasattr(module.Agent, "__call__"): # noqa: B004 wrap_function_wrapper(module, "Agent.__call__", wrap_agent__call__) @@ -471,19 +471,35 @@ def instrument_agent_agent(module): wrap_function_wrapper(module, "Agent.stream_async", wrap_stream_async) -def instrument_tools_executors__executor(module): +def instrument_strands_multiagent_graph(module): + if hasattr(module, "Graph"): + if hasattr(module.Graph, "__call__"): # noqa: B004 + wrap_function_wrapper(module, "Graph.__call__", wrap_agent__call__) + if hasattr(module.Graph, "invoke_async"): + wrap_function_wrapper(module, "Graph.invoke_async", wrap_agent_invoke_async) + + +def instrument_strands_multiagent_swarm(module): + if hasattr(module, "Swarm"): + if hasattr(module.Swarm, "__call__"): # noqa: B004 + wrap_function_wrapper(module, "Swarm.__call__", wrap_agent__call__) + if hasattr(module.Swarm, "invoke_async"): + wrap_function_wrapper(module, "Swarm.invoke_async", wrap_agent_invoke_async) + + +def instrument_strands_tools_executors__executor(module): if hasattr(module, "ToolExecutor"): if hasattr(module.ToolExecutor, "_stream"): wrap_function_wrapper(module, "ToolExecutor._stream", wrap_tool_executor__stream) -def instrument_tools_registry(module): +def instrument_strands_tools_registry(module): if hasattr(module, "ToolRegistry"): if hasattr(module.ToolRegistry, "register_tool"): wrap_function_wrapper(module, "ToolRegistry.register_tool", wrap_ToolRegister_register_tool) -def instrument_models_bedrock(module): +def instrument_strands_models_bedrock(module): # This instrumentation only exists to pass trace context due to bedrock models using a separate thread. if hasattr(module, "BedrockModel"): if hasattr(module.BedrockModel, "stream"): diff --git a/tests/mlmodel_strands/__init__.py b/tests/mlmodel_strands/__init__.py new file mode 100644 index 000000000..8030baccf --- /dev/null +++ b/tests/mlmodel_strands/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/mlmodel_strands/_test_agent.py b/tests/mlmodel_strands/_test_agent.py new file mode 100644 index 000000000..15aa79a5a --- /dev/null +++ b/tests/mlmodel_strands/_test_agent.py @@ -0,0 +1,165 @@ +# 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 pytest +from strands import tool + +from ._mock_model_provider import MockedModelProvider + + +# Example tool for testing purposes +@tool +async def add_exclamation(message: str) -> str: + return f"{message}!" + + +@tool +async def throw_exception_coro(message: str) -> str: + raise RuntimeError("Oops") + + +@tool +async def throw_exception_agen(message: str) -> str: + raise RuntimeError("Oops") + yield + + +@pytest.fixture +def single_tool_model(): + model = MockedModelProvider( + [ + { + "role": "assistant", + "content": [ + {"text": "Calling add_exclamation tool"}, + {"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Hello"}}}, + ], + }, + {"role": "assistant", "content": [{"text": "Success!"}]}, + ] + ) + return model + + +@pytest.fixture +def single_tool_model_runtime_error_coro(): + model = MockedModelProvider( + [ + { + "role": "assistant", + "content": [ + {"text": "Calling throw_exception_coro tool"}, + # Set arguments to an invalid type to trigger error in tool + {"toolUse": {"name": "throw_exception_coro", "toolUseId": "123", "input": {"message": "Hello"}}}, + ], + }, + {"role": "assistant", "content": [{"text": "Success!"}]}, + ] + ) + return model + + +@pytest.fixture +def single_tool_model_runtime_error_agen(): + model = MockedModelProvider( + [ + { + "role": "assistant", + "content": [ + {"text": "Calling throw_exception_agen tool"}, + # Set arguments to an invalid type to trigger error in tool + {"toolUse": {"name": "throw_exception_agen", "toolUseId": "123", "input": {"message": "Hello"}}}, + ], + }, + {"role": "assistant", "content": [{"text": "Success!"}]}, + ] + ) + return model + + +@pytest.fixture +def multi_tool_model(): + model = MockedModelProvider( + [ + { + "role": "assistant", + "content": [ + {"text": "Calling add_exclamation tool"}, + {"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Hello"}}}, + ], + }, + { + "role": "assistant", + "content": [ + {"text": "Calling compute_sum tool"}, + {"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 5, "b": 3}}}, + ], + }, + { + "role": "assistant", + "content": [ + {"text": "Calling add_exclamation tool"}, + {"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Goodbye"}}}, + ], + }, + { + "role": "assistant", + "content": [ + {"text": "Calling compute_sum tool"}, + {"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 123, "b": 2}}}, + ], + }, + {"role": "assistant", "content": [{"text": "Success!"}]}, + ] + ) + return model + + +@pytest.fixture +def multi_tool_model_error(): + model = MockedModelProvider( + [ + { + "role": "assistant", + "content": [ + {"text": "Calling add_exclamation tool"}, + {"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Hello"}}}, + ], + }, + { + "role": "assistant", + "content": [ + {"text": "Calling compute_sum tool"}, + {"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 5, "b": 3}}}, + ], + }, + { + "role": "assistant", + "content": [ + {"text": "Calling add_exclamation tool"}, + {"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Goodbye"}}}, + ], + }, + { + "role": "assistant", + "content": [ + {"text": "Calling compute_sum tool"}, + # Set insufficient arguments to trigger error in tool + {"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 123}}}, + ], + }, + {"role": "assistant", "content": [{"text": "Success!"}]}, + ] + ) + return model diff --git a/tests/mlmodel_strands/_test_multiagent_graph.py b/tests/mlmodel_strands/_test_multiagent_graph.py new file mode 100644 index 000000000..73c167970 --- /dev/null +++ b/tests/mlmodel_strands/_test_multiagent_graph.py @@ -0,0 +1,91 @@ +# 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 pytest +from strands import Agent, tool +from strands.multiagent.graph import GraphBuilder + +from ._mock_model_provider import MockedModelProvider + + +@pytest.fixture +def math_model(): + model = MockedModelProvider( + [ + { + "role": "assistant", + "content": [ + {"text": "I'll calculate the sum of 15 and 27 for you."}, + {"toolUse": {"name": "calculate_sum", "toolUseId": "123", "input": {"a": 15, "b": 27}}}, + ], + }, + {"role": "assistant", "content": [{"text": "The sum of 15 and 27 is 42."}]}, + ] + ) + return model + + +@pytest.fixture +def analysis_model(): + model = MockedModelProvider( + [ + { + "role": "assistant", + "content": [ + {"text": "I'll validate the calculation result of 42 from the calculator."}, + {"toolUse": {"name": "analyze_result", "toolUseId": "456", "input": {"value": 42}}}, + ], + }, + { + "role": "assistant", + "content": [{"text": "The calculation is correct, and 42 is a positive integer result."}], + }, + ] + ) + return model + + +# Example tool for testing purposes +@tool +async def calculate_sum(a: int, b: int) -> int: + """Calculate the sum of two numbers.""" + return a + b + + +@tool +async def analyze_result(value: int) -> str: + """Analyze a numeric result.""" + return f"The result {value} is {'positive' if value > 0 else 'zero or negative'}" + + +@pytest.fixture +def math_agent(math_model): + return Agent(name="math_agent", model=math_model, tools=[calculate_sum]) + + +@pytest.fixture +def analysis_agent(analysis_model): + return Agent(name="analysis_agent", model=analysis_model, tools=[analyze_result]) + + +@pytest.fixture +def agent_graph(math_agent, analysis_agent): + # Build graph + builder = GraphBuilder() + builder.add_node(math_agent, "math") + builder.add_node(analysis_agent, "analysis") + builder.add_edge("math", "analysis") + builder.set_entry_point("math") + + return builder.build() diff --git a/tests/mlmodel_strands/_test_multiagent_swarm.py b/tests/mlmodel_strands/_test_multiagent_swarm.py new file mode 100644 index 000000000..4b7916c27 --- /dev/null +++ b/tests/mlmodel_strands/_test_multiagent_swarm.py @@ -0,0 +1,108 @@ +# 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 pytest +from strands import Agent, tool +from strands.multiagent.swarm import Swarm + +from ._mock_model_provider import MockedModelProvider + + +@pytest.fixture +def math_model(): + model = MockedModelProvider( + [ + { + "role": "assistant", + "content": [ + {"text": "I'll calculate the sum of 15 and 27 for you."}, + {"toolUse": {"name": "calculate_sum", "toolUseId": "123", "input": {"a": 15, "b": 27}}}, + ], + }, + { + "role": "assistant", + "content": [ + { + "toolUse": { + "name": "handoff_to_agent", + "toolUseId": "789", + "input": { + "agent_name": "analysis_agent", + "message": "Analyze the result of the calculation done by the math_agent.", + "context": {"result": 42}, + }, + } + } + ], + }, + {"role": "assistant", "content": [{"text": "The sum of 15 and 27 is 42."}]}, + ] + ) + return model + + +@pytest.fixture +def analysis_model(): + model = MockedModelProvider( + [ + { + "role": "assistant", + "content": [ + {"text": "I'll validate the calculation result of 42 from the calculator."}, + {"toolUse": {"name": "analyze_result", "toolUseId": "456", "input": {"value": 42}}}, + ], + }, + { + "role": "assistant", + "content": [{"text": "The calculation is correct, and 42 is a positive integer result."}], + }, + ] + ) + return model + + +# Example tool for testing purposes +@tool +async def calculate_sum(a: int, b: int) -> int: + """Calculate the sum of two numbers.""" + return a + b + + +@tool +async def analyze_result(value: int) -> str: + """Analyze a numeric result.""" + return f"The result {value} is {'positive' if value > 0 else 'zero or negative'}" + + +@pytest.fixture +def math_agent(math_model): + return Agent(name="math_agent", model=math_model, tools=[calculate_sum]) + + +@pytest.fixture +def analysis_agent(analysis_model): + return Agent(name="analysis_agent", model=analysis_model, tools=[analyze_result]) + + +@pytest.fixture +def agent_swarm(math_agent, analysis_agent): + # Build graph with conditional edge + return Swarm( + [math_agent, analysis_agent], + entry_point=math_agent, + execution_timeout=60, + node_timeout=30, + max_handoffs=5, + max_iterations=5, + ) diff --git a/tests/mlmodel_strands/conftest.py b/tests/mlmodel_strands/conftest.py index a2ad9b8dd..abbc29b96 100644 --- a/tests/mlmodel_strands/conftest.py +++ b/tests/mlmodel_strands/conftest.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest -from _mock_model_provider import MockedModelProvider from testing_support.fixture.event_loop import event_loop as loop from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture from testing_support.ml_testing_utils import set_trace_info @@ -31,133 +29,3 @@ collector_agent_registration = collector_agent_registration_fixture( app_name="Python Agent Test (mlmodel_strands)", default_settings=_default_settings ) - - -@pytest.fixture -def single_tool_model(): - model = MockedModelProvider( - [ - { - "role": "assistant", - "content": [ - {"text": "Calling add_exclamation tool"}, - {"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Hello"}}}, - ], - }, - {"role": "assistant", "content": [{"text": "Success!"}]}, - ] - ) - return model - - -@pytest.fixture -def single_tool_model_runtime_error_coro(): - model = MockedModelProvider( - [ - { - "role": "assistant", - "content": [ - {"text": "Calling throw_exception_coro tool"}, - # Set arguments to an invalid type to trigger error in tool - {"toolUse": {"name": "throw_exception_coro", "toolUseId": "123", "input": {"message": "Hello"}}}, - ], - }, - {"role": "assistant", "content": [{"text": "Success!"}]}, - ] - ) - return model - - -@pytest.fixture -def single_tool_model_runtime_error_agen(): - model = MockedModelProvider( - [ - { - "role": "assistant", - "content": [ - {"text": "Calling throw_exception_agen tool"}, - # Set arguments to an invalid type to trigger error in tool - {"toolUse": {"name": "throw_exception_agen", "toolUseId": "123", "input": {"message": "Hello"}}}, - ], - }, - {"role": "assistant", "content": [{"text": "Success!"}]}, - ] - ) - return model - - -@pytest.fixture -def multi_tool_model(): - model = MockedModelProvider( - [ - { - "role": "assistant", - "content": [ - {"text": "Calling add_exclamation tool"}, - {"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Hello"}}}, - ], - }, - { - "role": "assistant", - "content": [ - {"text": "Calling compute_sum tool"}, - {"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 5, "b": 3}}}, - ], - }, - { - "role": "assistant", - "content": [ - {"text": "Calling add_exclamation tool"}, - {"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Goodbye"}}}, - ], - }, - { - "role": "assistant", - "content": [ - {"text": "Calling compute_sum tool"}, - {"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 123, "b": 2}}}, - ], - }, - {"role": "assistant", "content": [{"text": "Success!"}]}, - ] - ) - return model - - -@pytest.fixture -def multi_tool_model_error(): - model = MockedModelProvider( - [ - { - "role": "assistant", - "content": [ - {"text": "Calling add_exclamation tool"}, - {"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Hello"}}}, - ], - }, - { - "role": "assistant", - "content": [ - {"text": "Calling compute_sum tool"}, - {"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 5, "b": 3}}}, - ], - }, - { - "role": "assistant", - "content": [ - {"text": "Calling add_exclamation tool"}, - {"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Goodbye"}}}, - ], - }, - { - "role": "assistant", - "content": [ - {"text": "Calling compute_sum tool"}, - # Set insufficient arguments to trigger error in tool - {"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 123}}}, - ], - }, - {"role": "assistant", "content": [{"text": "Success!"}]}, - ] - ) - return model diff --git a/tests/mlmodel_strands/test_agent.py b/tests/mlmodel_strands/test_agent.py index af685668a..6fa5e56a6 100644 --- a/tests/mlmodel_strands/test_agent.py +++ b/tests/mlmodel_strands/test_agent.py @@ -13,7 +13,7 @@ # limitations under the License. import pytest -from strands import Agent, tool +from strands import Agent from testing_support.fixtures import reset_core_stats_engine, validate_attributes from testing_support.ml_testing_utils import ( disabled_ai_monitoring_record_content_settings, @@ -32,6 +32,17 @@ from newrelic.common.object_names import callable_name from newrelic.common.object_wrapper import transient_function_wrapper +from ._test_agent import ( + add_exclamation, + multi_tool_model, + multi_tool_model_error, + single_tool_model, + single_tool_model_runtime_error_agen, + single_tool_model_runtime_error_coro, + throw_exception_agen, + throw_exception_coro, +) + tool_recorded_event = [ ( {"type": "LlmTool"}, @@ -144,29 +155,12 @@ ] -# Example tool for testing purposes -@tool -async def add_exclamation(message: str) -> str: - return f"{message}!" - - -@tool -async def throw_exception_coro(message: str) -> str: - raise RuntimeError("Oops") - - -@tool -async def throw_exception_agen(message: str) -> str: - raise RuntimeError("Oops") - yield - - @reset_core_stats_engine() @validate_custom_events(events_with_context_attrs(tool_recorded_event)) @validate_custom_events(events_with_context_attrs(agent_recorded_event)) @validate_custom_event_count(count=2) @validate_transaction_metrics( - "test_agent:test_agent_invoke", + "mlmodel_strands.test_agent:test_agent_invoke", scoped_metrics=[ ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/my_agent", 1), ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/add_exclamation", 1), @@ -194,7 +188,7 @@ def test_agent_invoke(set_trace_info, single_tool_model): @validate_custom_events(agent_recorded_event) @validate_custom_event_count(count=2) @validate_transaction_metrics( - "test_agent:test_agent_invoke_async", + "mlmodel_strands.test_agent:test_agent_invoke_async", scoped_metrics=[ ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/my_agent", 1), ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/add_exclamation", 1), @@ -224,7 +218,7 @@ async def _test(): @validate_custom_events(agent_recorded_event) @validate_custom_event_count(count=2) @validate_transaction_metrics( - "test_agent:test_agent_stream_async", + "mlmodel_strands.test_agent:test_agent_stream_async", scoped_metrics=[ ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/my_agent", 1), ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/add_exclamation", 1), @@ -260,7 +254,7 @@ async def _test(): @validate_custom_events(tool_events_sans_content(tool_recorded_event)) @validate_custom_event_count(count=2) @validate_transaction_metrics( - "test_agent:test_agent_invoke_no_content", + "mlmodel_strands.test_agent:test_agent_invoke_no_content", scoped_metrics=[ ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/my_agent", 1), ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/add_exclamation", 1), @@ -301,7 +295,7 @@ def test_agent_invoke_disabled_ai_monitoring_events(set_trace_info, single_tool_ @validate_custom_events(agent_recorded_event_error) @validate_custom_event_count(count=1) @validate_transaction_metrics( - "test_agent:test_agent_invoke_error", + "mlmodel_strands.test_agent:test_agent_invoke_error", scoped_metrics=[("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/my_agent", 1)], rollup_metrics=[("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/my_agent", 1)], background_task=True, @@ -330,7 +324,7 @@ def _test(): @validate_custom_events(tool_recorded_event_error_coro) @validate_custom_event_count(count=2) @validate_transaction_metrics( - "test_agent:test_agent_invoke_tool_coro_runtime_error", + "mlmodel_strands.test_agent:test_agent_invoke_tool_coro_runtime_error", scoped_metrics=[ ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/my_agent", 1), ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/throw_exception_coro", 1), @@ -358,7 +352,7 @@ def test_agent_invoke_tool_coro_runtime_error(set_trace_info, single_tool_model_ @validate_custom_events(tool_recorded_event_error_agen) @validate_custom_event_count(count=2) @validate_transaction_metrics( - "test_agent:test_agent_invoke_tool_agen_runtime_error", + "mlmodel_strands.test_agent:test_agent_invoke_tool_agen_runtime_error", scoped_metrics=[ ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/my_agent", 1), ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/throw_exception_agen", 1), @@ -387,7 +381,7 @@ def test_agent_invoke_tool_agen_runtime_error(set_trace_info, single_tool_model_ @validate_custom_events(tool_recorded_event_forced_internal_error) @validate_custom_event_count(count=2) @validate_transaction_metrics( - "test_agent:test_agent_tool_forced_exception", + "mlmodel_strands.test_agent:test_agent_tool_forced_exception", scoped_metrics=[ ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/my_agent", 1), ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/add_exclamation", 1), diff --git a/tests/mlmodel_strands/test_multiagent_graph.py b/tests/mlmodel_strands/test_multiagent_graph.py new file mode 100644 index 000000000..7bd84fc90 --- /dev/null +++ b/tests/mlmodel_strands/test_multiagent_graph.py @@ -0,0 +1,233 @@ +# 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. + +from testing_support.fixtures import 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_transaction_metrics import validate_transaction_metrics + +from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes + +from ._test_multiagent_graph import agent_graph, analysis_agent, analysis_model, math_agent, math_model + +agent_recorded_events = [ + [ + {"type": "LlmAgent"}, + { + "duration": None, + "id": None, + "ingest_source": "Python", + "name": "math_agent", + "span_id": None, + "trace_id": "trace-id", + "vendor": "strands", + }, + ], + [ + {"type": "LlmAgent"}, + { + "duration": None, + "id": None, + "ingest_source": "Python", + "name": "analysis_agent", + "span_id": None, + "trace_id": "trace-id", + "vendor": "strands", + }, + ], +] + +tool_recorded_events = [ + [ + {"type": "LlmTool"}, + { + "agent_name": "math_agent", + "duration": None, + "id": None, + "ingest_source": "Python", + "input": "{'a': 15, 'b': 27}", + "name": "calculate_sum", + "output": "{'text': '42'}", + "run_id": "123", + "span_id": None, + "trace_id": "trace-id", + "vendor": "strands", + }, + ], + [ + {"type": "LlmTool"}, + { + "agent_name": "analysis_agent", + "duration": None, + "id": None, + "ingest_source": "Python", + "input": "{'value': 42}", + "name": "analyze_result", + "output": "{'text': 'The result 42 is positive'}", + "run_id": "456", + "span_id": None, + "trace_id": "trace-id", + "vendor": "strands", + }, + ], +] + + +@reset_core_stats_engine() +@validate_custom_events(events_with_context_attrs(tool_recorded_events)) +@validate_custom_events(events_with_context_attrs(agent_recorded_events)) +@validate_custom_event_count(count=4) # 2 LlmTool events, 2 LlmAgent events +@validate_transaction_metrics( + "mlmodel_strands.test_multiagent_graph:test_multiagent_graph_invoke", + scoped_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + rollup_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + background_task=True, +) +@validate_attributes("agent", ["llm"]) +@background_task() +def test_multiagent_graph_invoke(set_trace_info, agent_graph): + set_trace_info() + + with WithLlmCustomAttributes({"context": "attr"}): + response = agent_graph("Calculate the sum of 15 and 27.") + + assert response.execution_count == 2 + assert not response.failed_nodes + assert response.results["math"].result.message["content"][0]["text"] == "The sum of 15 and 27 is 42." + assert ( + response.results["analysis"].result.message["content"][0]["text"] + == "The calculation is correct, and 42 is a positive integer result." + ) + + +@reset_core_stats_engine() +@validate_custom_events(tool_recorded_events) +@validate_custom_events(agent_recorded_events) +@validate_custom_event_count(count=4) # 2 LlmTool events, 2 LlmAgent events +@validate_transaction_metrics( + "mlmodel_strands.test_multiagent_graph:test_multiagent_graph_invoke_async", + scoped_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + rollup_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + background_task=True, +) +@validate_attributes("agent", ["llm"]) +@background_task() +def test_multiagent_graph_invoke_async(loop, set_trace_info, agent_graph): + set_trace_info() + + async def _test(): + response = await agent_graph.invoke_async("Calculate the sum of 15 and 27.") + + assert response.execution_count == 2 + assert not response.failed_nodes + assert response.results["math"].result.message["content"][0]["text"] == "The sum of 15 and 27 is 42." + assert ( + response.results["analysis"].result.message["content"][0]["text"] + == "The calculation is correct, and 42 is a positive integer result." + ) + + loop.run_until_complete(_test()) + + +@reset_core_stats_engine() +@validate_custom_events(tool_recorded_events) +@validate_custom_events(agent_recorded_events) +@validate_custom_event_count(count=4) # 2 LlmTool events, 2 LlmAgent events +@validate_transaction_metrics( + "mlmodel_strands.test_multiagent_graph:test_multiagent_graph_stream_async", + scoped_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + rollup_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + background_task=True, +) +@validate_attributes("agent", ["llm"]) +@background_task() +def test_multiagent_graph_stream_async(loop, set_trace_info, agent_graph): + set_trace_info() + + async def _test(): + response = agent_graph.stream_async("Calculate the sum of 15 and 27.") + messages = [ + event["node_result"].result.message async for event in response if event["type"] == "multiagent_node_stop" + ] + + assert len(messages) == 2 + + assert messages[0]["content"][0]["text"] == "The sum of 15 and 27 is 42." + assert messages[1]["content"][0]["text"] == "The calculation is correct, and 42 is a positive integer result." + + loop.run_until_complete(_test()) + + +@disabled_ai_monitoring_settings +@reset_core_stats_engine() +@validate_custom_event_count(count=0) +@background_task() +def test_multiagent_graph_invoke_disabled_ai_monitoring_events(set_trace_info, agent_graph): + set_trace_info() + + response = agent_graph("Calculate the sum of 15 and 27.") + + assert response.execution_count == 2 + assert not response.failed_nodes + assert response.results["math"].result.message["content"][0]["text"] == "The sum of 15 and 27 is 42." + assert ( + response.results["analysis"].result.message["content"][0]["text"] + == "The calculation is correct, and 42 is a positive integer result." + ) + + +@reset_core_stats_engine() +@validate_custom_event_count(count=0) +def test_multiagent_graph_invoke_outside_txn(agent_graph): + response = agent_graph("Calculate the sum of 15 and 27.") + + assert response.execution_count == 2 + assert not response.failed_nodes + assert response.results["math"].result.message["content"][0]["text"] == "The sum of 15 and 27 is 42." + assert ( + response.results["analysis"].result.message["content"][0]["text"] + == "The calculation is correct, and 42 is a positive integer result." + ) diff --git a/tests/mlmodel_strands/test_multiagent_swarm.py b/tests/mlmodel_strands/test_multiagent_swarm.py new file mode 100644 index 000000000..bbcbb3e27 --- /dev/null +++ b/tests/mlmodel_strands/test_multiagent_swarm.py @@ -0,0 +1,260 @@ +# 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. + +from testing_support.fixtures import 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_transaction_metrics import validate_transaction_metrics + +from newrelic.api.background_task import background_task +from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes + +from ._test_multiagent_swarm import agent_swarm, analysis_agent, analysis_model, math_agent, math_model + +agent_recorded_events = [ + [ + {"type": "LlmAgent"}, + { + "duration": None, + "id": None, + "ingest_source": "Python", + "name": "math_agent", + "span_id": None, + "trace_id": "trace-id", + "vendor": "strands", + }, + ], + [ + {"type": "LlmAgent"}, + { + "duration": None, + "id": None, + "ingest_source": "Python", + "name": "analysis_agent", + "span_id": None, + "trace_id": "trace-id", + "vendor": "strands", + }, + ], +] + +tool_recorded_events = [ + [ + {"type": "LlmTool"}, + { + "agent_name": "math_agent", + "duration": None, + "id": None, + "ingest_source": "Python", + "input": "{'a': 15, 'b': 27}", + "name": "calculate_sum", + "output": "{'text': '42'}", + "run_id": "123", + "span_id": None, + "trace_id": "trace-id", + "vendor": "strands", + }, + ], + [ + {"type": "LlmTool"}, + { + "agent_name": "analysis_agent", + "duration": None, + "id": None, + "ingest_source": "Python", + "input": "{'value': 42}", + "name": "analyze_result", + "output": "{'text': 'The result 42 is positive'}", + "run_id": "456", + "span_id": None, + "trace_id": "trace-id", + "vendor": "strands", + }, + ], +] + +handoff_recorded_event = [ + [ + {"type": "LlmTool"}, + { + "agent_name": "math_agent", + "duration": None, + "id": None, + "ingest_source": "Python", + # This is the output from math_agent being sent to the handoff_to_agent tool, which will then be input to the analysis_agent + "input": "{'agent_name': 'analysis_agent', 'message': 'Analyze the result of the calculation done by the math_agent.', 'context': {'result': 42}}", + "name": "handoff_to_agent", + "output": "{'text': 'Handing off to analysis_agent: Analyze the result of the calculation done by the math_agent.'}", + "run_id": "789", + "span_id": None, + "trace_id": "trace-id", + "vendor": "strands", + }, + ] +] + + +@reset_core_stats_engine() +@validate_custom_events(events_with_context_attrs(tool_recorded_events)) +@validate_custom_events(events_with_context_attrs(agent_recorded_events)) +@validate_custom_events(events_with_context_attrs(handoff_recorded_event)) +@validate_custom_event_count(count=5) # 2 LlmTool events, 2 LlmAgent events, 1 LlmTool Handoff event +@validate_transaction_metrics( + "mlmodel_strands.test_multiagent_swarm:test_multiagent_swarm_invoke", + scoped_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + rollup_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + background_task=True, +) +@validate_attributes("agent", ["llm"]) +@background_task() +def test_multiagent_swarm_invoke(set_trace_info, agent_swarm): + set_trace_info() + + with WithLlmCustomAttributes({"context": "attr"}): + response = agent_swarm("Calculate the sum of 15 and 27.") + + assert response.execution_count == 2 + node_history = [node.node_id for node in response.node_history] + assert node_history == ["math_agent", "analysis_agent"] + assert response.results["math_agent"].result.message["content"][0]["text"] == "The sum of 15 and 27 is 42." + assert ( + response.results["analysis_agent"].result.message["content"][0]["text"] + == "The calculation is correct, and 42 is a positive integer result." + ) + + +@reset_core_stats_engine() +@validate_custom_events(tool_recorded_events) +@validate_custom_events(agent_recorded_events) +@validate_custom_events(handoff_recorded_event) +@validate_custom_event_count(count=5) # 2 LlmTool events, 2 LlmAgent events, 1 LlmTool Handoff event +@validate_transaction_metrics( + "mlmodel_strands.test_multiagent_swarm:test_multiagent_swarm_invoke_async", + scoped_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + rollup_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + background_task=True, +) +@validate_attributes("agent", ["llm"]) +@background_task() +def test_multiagent_swarm_invoke_async(loop, set_trace_info, agent_swarm): + set_trace_info() + + async def _test(): + response = await agent_swarm.invoke_async("Calculate the sum of 15 and 27.") + + assert response.execution_count == 2 + node_history = [node.node_id for node in response.node_history] + assert node_history == ["math_agent", "analysis_agent"] + assert response.results["math_agent"].result.message["content"][0]["text"] == "The sum of 15 and 27 is 42." + assert ( + response.results["analysis_agent"].result.message["content"][0]["text"] + == "The calculation is correct, and 42 is a positive integer result." + ) + + loop.run_until_complete(_test()) + + +@reset_core_stats_engine() +@validate_custom_events(tool_recorded_events) +@validate_custom_events(agent_recorded_events) +@validate_custom_events(handoff_recorded_event) +@validate_custom_event_count(count=5) # 2 LlmTool events, 2 LlmAgent events, 1 LlmTool Handoff event +@validate_transaction_metrics( + "mlmodel_strands.test_multiagent_swarm:test_multiagent_swarm_stream_async", + scoped_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + rollup_metrics=[ + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/math_agent", 1), + ("Llm/agent/Strands/strands.agent.agent:Agent.stream_async/analysis_agent", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/calculate_sum", 1), + ("Llm/tool/Strands/strands.tools.executors._executor:ToolExecutor._stream/analyze_result", 1), + ], + background_task=True, +) +@validate_attributes("agent", ["llm"]) +@background_task() +def test_multiagent_swarm_stream_async(loop, set_trace_info, agent_swarm): + set_trace_info() + + async def _test(): + response = agent_swarm.stream_async("Calculate the sum of 15 and 27.") + messages = [ + event["node_result"].result.message async for event in response if event["type"] == "multiagent_node_stop" + ] + + assert len(messages) == 2 + + assert messages[0]["content"][0]["text"] == "The sum of 15 and 27 is 42." + assert messages[1]["content"][0]["text"] == "The calculation is correct, and 42 is a positive integer result." + + loop.run_until_complete(_test()) + + +@disabled_ai_monitoring_settings +@reset_core_stats_engine() +@validate_custom_event_count(count=0) +@background_task() +def test_multiagent_swarm_invoke_disabled_ai_monitoring_events(set_trace_info, agent_swarm): + set_trace_info() + + response = agent_swarm("Calculate the sum of 15 and 27.") + + assert response.execution_count == 2 + node_history = [node.node_id for node in response.node_history] + assert node_history == ["math_agent", "analysis_agent"] + assert response.results["math_agent"].result.message["content"][0]["text"] == "The sum of 15 and 27 is 42." + assert ( + response.results["analysis_agent"].result.message["content"][0]["text"] + == "The calculation is correct, and 42 is a positive integer result." + ) + + +@reset_core_stats_engine() +@validate_custom_event_count(count=0) +def test_multiagent_swarm_invoke_outside_txn(agent_swarm): + response = agent_swarm("Calculate the sum of 15 and 27.") + + assert response.execution_count == 2 + node_history = [node.node_id for node in response.node_history] + assert node_history == ["math_agent", "analysis_agent"] + assert response.results["math_agent"].result.message["content"][0]["text"] == "The sum of 15 and 27 is 42." + assert ( + response.results["analysis_agent"].result.message["content"][0]["text"] + == "The calculation is correct, and 42 is a positive integer result." + )