Skip to content

Commit ac6eb5f

Browse files
umaannamalaiumaannamalai
authored
Merge branch 'main' into sgoel/fixed_tool_type_bug
2 parents 661efd8 + c31c334 commit ac6eb5f

10 files changed

Lines changed: 928 additions & 166 deletions

‎newrelic/config.py‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2948,12 +2948,26 @@ def _process_module_builtin_defaults():
29482948
"newrelic.hooks.mlmodel_autogen",
29492949
"instrument_autogen_agentchat_agents__assistant_agent",
29502950
)
2951-
_process_module_definition("strands.agent.agent", "newrelic.hooks.mlmodel_strands", "instrument_agent_agent")
29522951
_process_module_definition(
2953-
"strands.tools.executors._executor", "newrelic.hooks.mlmodel_strands", "instrument_tools_executors__executor"
2952+
"strands.agent.agent", "newrelic.hooks.mlmodel_strands", "instrument_strands_agent_agent"
2953+
)
2954+
_process_module_definition(
2955+
"strands.multiagent.graph", "newrelic.hooks.mlmodel_strands", "instrument_strands_multiagent_graph"
2956+
)
2957+
_process_module_definition(
2958+
"strands.multiagent.swarm", "newrelic.hooks.mlmodel_strands", "instrument_strands_multiagent_swarm"
2959+
)
2960+
_process_module_definition(
2961+
"strands.tools.executors._executor",
2962+
"newrelic.hooks.mlmodel_strands",
2963+
"instrument_strands_tools_executors__executor",
2964+
)
2965+
_process_module_definition(
2966+
"strands.tools.registry", "newrelic.hooks.mlmodel_strands", "instrument_strands_tools_registry"
2967+
)
2968+
_process_module_definition(
2969+
"strands.models.bedrock", "newrelic.hooks.mlmodel_strands", "instrument_strands_models_bedrock"
29542970
)
2955-
_process_module_definition("strands.tools.registry", "newrelic.hooks.mlmodel_strands", "instrument_tools_registry")
2956-
_process_module_definition("strands.models.bedrock", "newrelic.hooks.mlmodel_strands", "instrument_models_bedrock")
29572971

29582972
_process_module_definition("mcp.client.session", "newrelic.hooks.adapter_mcp", "instrument_mcp_client_session")
29592973
_process_module_definition(

‎newrelic/hooks/mlmodel_strands.py‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -464,7 +464,7 @@ def wrap_bedrock_model__stream(wrapped, instance, args, kwargs):
464464
return wrapped(*args, **kwargs)
465465

466466

467-
def instrument_agent_agent(module):
467+
def instrument_strands_agent_agent(module):
468468
if hasattr(module, "Agent"):
469469
if hasattr(module.Agent, "__call__"): # noqa: B004
470470
wrap_function_wrapper(module, "Agent.__call__", wrap_agent__call__)
@@ -474,19 +474,35 @@ def instrument_agent_agent(module):
474474
wrap_function_wrapper(module, "Agent.stream_async", wrap_stream_async)
475475

476476

477-
def instrument_tools_executors__executor(module):
477+
def instrument_strands_multiagent_graph(module):
478+
if hasattr(module, "Graph"):
479+
if hasattr(module.Graph, "__call__"): # noqa: B004
480+
wrap_function_wrapper(module, "Graph.__call__", wrap_agent__call__)
481+
if hasattr(module.Graph, "invoke_async"):
482+
wrap_function_wrapper(module, "Graph.invoke_async", wrap_agent_invoke_async)
483+
484+
485+
def instrument_strands_multiagent_swarm(module):
486+
if hasattr(module, "Swarm"):
487+
if hasattr(module.Swarm, "__call__"): # noqa: B004
488+
wrap_function_wrapper(module, "Swarm.__call__", wrap_agent__call__)
489+
if hasattr(module.Swarm, "invoke_async"):
490+
wrap_function_wrapper(module, "Swarm.invoke_async", wrap_agent_invoke_async)
491+
492+
493+
def instrument_strands_tools_executors__executor(module):
478494
if hasattr(module, "ToolExecutor"):
479495
if hasattr(module.ToolExecutor, "_stream"):
480496
wrap_function_wrapper(module, "ToolExecutor._stream", wrap_tool_executor__stream)
481497

482498

483-
def instrument_tools_registry(module):
499+
def instrument_strands_tools_registry(module):
484500
if hasattr(module, "ToolRegistry"):
485501
if hasattr(module.ToolRegistry, "register_tool"):
486502
wrap_function_wrapper(module, "ToolRegistry.register_tool", wrap_ToolRegister_register_tool)
487503

488504

489-
def instrument_models_bedrock(module):
505+
def instrument_strands_models_bedrock(module):
490506
# This instrumentation only exists to pass trace context due to bedrock models using a separate thread.
491507
if hasattr(module, "BedrockModel"):
492508
if hasattr(module.BedrockModel, "stream"):

‎tests/mlmodel_strands/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright 2010 New Relic, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

‎tests/mlmodel_strands/_test_agent.py‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
# Copyright 2010 New Relic, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import pytest
16+
from strands import tool
17+
18+
from ._mock_model_provider import MockedModelProvider
19+
20+
21+
# Example tool for testing purposes
22+
@tool
23+
async def add_exclamation(message: str) -> str:
24+
return f"{message}!"
25+
26+
27+
@tool
28+
async def throw_exception_coro(message: str) -> str:
29+
raise RuntimeError("Oops")
30+
31+
32+
@tool
33+
async def throw_exception_agen(message: str) -> str:
34+
raise RuntimeError("Oops")
35+
yield
36+
37+
38+
@pytest.fixture
39+
def single_tool_model():
40+
model = MockedModelProvider(
41+
[
42+
{
43+
"role": "assistant",
44+
"content": [
45+
{"text": "Calling add_exclamation tool"},
46+
{"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Hello"}}},
47+
],
48+
},
49+
{"role": "assistant", "content": [{"text": "Success!"}]},
50+
]
51+
)
52+
return model
53+
54+
55+
@pytest.fixture
56+
def single_tool_model_runtime_error_coro():
57+
model = MockedModelProvider(
58+
[
59+
{
60+
"role": "assistant",
61+
"content": [
62+
{"text": "Calling throw_exception_coro tool"},
63+
# Set arguments to an invalid type to trigger error in tool
64+
{"toolUse": {"name": "throw_exception_coro", "toolUseId": "123", "input": {"message": "Hello"}}},
65+
],
66+
},
67+
{"role": "assistant", "content": [{"text": "Success!"}]},
68+
]
69+
)
70+
return model
71+
72+
73+
@pytest.fixture
74+
def single_tool_model_runtime_error_agen():
75+
model = MockedModelProvider(
76+
[
77+
{
78+
"role": "assistant",
79+
"content": [
80+
{"text": "Calling throw_exception_agen tool"},
81+
# Set arguments to an invalid type to trigger error in tool
82+
{"toolUse": {"name": "throw_exception_agen", "toolUseId": "123", "input": {"message": "Hello"}}},
83+
],
84+
},
85+
{"role": "assistant", "content": [{"text": "Success!"}]},
86+
]
87+
)
88+
return model
89+
90+
91+
@pytest.fixture
92+
def multi_tool_model():
93+
model = MockedModelProvider(
94+
[
95+
{
96+
"role": "assistant",
97+
"content": [
98+
{"text": "Calling add_exclamation tool"},
99+
{"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Hello"}}},
100+
],
101+
},
102+
{
103+
"role": "assistant",
104+
"content": [
105+
{"text": "Calling compute_sum tool"},
106+
{"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 5, "b": 3}}},
107+
],
108+
},
109+
{
110+
"role": "assistant",
111+
"content": [
112+
{"text": "Calling add_exclamation tool"},
113+
{"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Goodbye"}}},
114+
],
115+
},
116+
{
117+
"role": "assistant",
118+
"content": [
119+
{"text": "Calling compute_sum tool"},
120+
{"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 123, "b": 2}}},
121+
],
122+
},
123+
{"role": "assistant", "content": [{"text": "Success!"}]},
124+
]
125+
)
126+
return model
127+
128+
129+
@pytest.fixture
130+
def multi_tool_model_error():
131+
model = MockedModelProvider(
132+
[
133+
{
134+
"role": "assistant",
135+
"content": [
136+
{"text": "Calling add_exclamation tool"},
137+
{"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Hello"}}},
138+
],
139+
},
140+
{
141+
"role": "assistant",
142+
"content": [
143+
{"text": "Calling compute_sum tool"},
144+
{"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 5, "b": 3}}},
145+
],
146+
},
147+
{
148+
"role": "assistant",
149+
"content": [
150+
{"text": "Calling add_exclamation tool"},
151+
{"toolUse": {"name": "add_exclamation", "toolUseId": "123", "input": {"message": "Goodbye"}}},
152+
],
153+
},
154+
{
155+
"role": "assistant",
156+
"content": [
157+
{"text": "Calling compute_sum tool"},
158+
# Set insufficient arguments to trigger error in tool
159+
{"toolUse": {"name": "compute_sum", "toolUseId": "123", "input": {"a": 123}}},
160+
],
161+
},
162+
{"role": "assistant", "content": [{"text": "Success!"}]},
163+
]
164+
)
165+
return model

‎tests/mlmodel_strands/_test_multiagent_graph.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
# Copyright 2010 New Relic, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import pytest
16+
from strands import Agent, tool
17+
from strands.multiagent.graph import GraphBuilder
18+
19+
from ._mock_model_provider import MockedModelProvider
20+
21+
22+
@pytest.fixture
23+
def math_model():
24+
model = MockedModelProvider(
25+
[
26+
{
27+
"role": "assistant",
28+
"content": [
29+
{"text": "I'll calculate the sum of 15 and 27 for you."},
30+
{"toolUse": {"name": "calculate_sum", "toolUseId": "123", "input": {"a": 15, "b": 27}}},
31+
],
32+
},
33+
{"role": "assistant", "content": [{"text": "The sum of 15 and 27 is 42."}]},
34+
]
35+
)
36+
return model
37+
38+
39+
@pytest.fixture
40+
def analysis_model():
41+
model = MockedModelProvider(
42+
[
43+
{
44+
"role": "assistant",
45+
"content": [
46+
{"text": "I'll validate the calculation result of 42 from the calculator."},
47+
{"toolUse": {"name": "analyze_result", "toolUseId": "456", "input": {"value": 42}}},
48+
],
49+
},
50+
{
51+
"role": "assistant",
52+
"content": [{"text": "The calculation is correct, and 42 is a positive integer result."}],
53+
},
54+
]
55+
)
56+
return model
57+
58+
59+
# Example tool for testing purposes
60+
@tool
61+
async def calculate_sum(a: int, b: int) -> int:
62+
"""Calculate the sum of two numbers."""
63+
return a + b
64+
65+
66+
@tool
67+
async def analyze_result(value: int) -> str:
68+
"""Analyze a numeric result."""
69+
return f"The result {value} is {'positive' if value > 0 else 'zero or negative'}"
70+
71+
72+
@pytest.fixture
73+
def math_agent(math_model):
74+
return Agent(name="math_agent", model=math_model, tools=[calculate_sum])
75+
76+
77+
@pytest.fixture
78+
def analysis_agent(analysis_model):
79+
return Agent(name="analysis_agent", model=analysis_model, tools=[analyze_result])
80+
81+
82+
@pytest.fixture
83+
def agent_graph(math_agent, analysis_agent):
84+
# Build graph
85+
builder = GraphBuilder()
86+
builder.add_node(math_agent, "math")
87+
builder.add_node(analysis_agent, "analysis")
88+
builder.add_edge("math", "analysis")
89+
builder.set_entry_point("math")
90+
91+
return builder.build()

0 commit comments

Comments
 (0)