Skip to content

Commit 02c4f37

Browse files
umaannamalaimergify[bot]
umaannamalai
and
mergify[bot]
authored
Sync autogen feature branch (#1472)
* Add instrumentation for autogen MCP tool adapter. (#1409) * Add instrumentation for MCP tool adapter. * [MegaLinter] Apply linters fixes * Add autogen agent instrumentation. * Revert "Add autogen agent instrumentation." This reverts commit a26766a. * Add autogen agent instrumentation. * Revert "Add autogen agent instrumentation." This reverts commit a26766a. * Autogen Agent Instrumentation (#1412) * Add autogen agent instrumentation. * Address review feedback. * [MegaLinter] Apply linters fixes --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> * Add instrumentation for autogen MCP tool adapter. (#1409) * Add instrumentation for MCP tool adapter. * [MegaLinter] Apply linters fixes * Add autogen agent instrumentation. * Revert "Add autogen agent instrumentation." This reverts commit a26766a. * Add autogen agent instrumentation. * Revert "Add autogen agent instrumentation." This reverts commit a26766a. * Autogen Agent Instrumentation (#1412) * Add autogen agent instrumentation. * Address review feedback. * [MegaLinter] Apply linters fixes --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> * Save linking metadata on generator proxy. * Add comments to changes in instrumentation. * Add LLMAgent events to autogen instrumentation. (#1455) --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
1 parent 5bc8b67 commit 02c4f37

File tree

7 files changed

+254
-81
lines changed

7 files changed

+254
-81
lines changed

‎newrelic/hooks/mlmodel_autogen.py‎

Lines changed: 51 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,11 +59,40 @@ def wrap_on_messages_stream(wrapped, instance, args, kwargs):
5959
if not transaction:
6060
return wrapped(*args, **kwargs)
6161

62+
settings = transaction.settings or global_settings()
63+
if not settings.ai_monitoring.enabled:
64+
return wrapped(*args, **kwargs)
65+
66+
# Framework metric also used for entity tagging in the UI
67+
transaction.add_ml_model_info("Autogen", AUTOGEN_VERSION)
68+
transaction._add_agent_attribute("llm", True)
69+
6270
agent_name = getattr(instance, "name", "agent")
71+
agent_id = str(uuid.uuid4())
72+
agent_event_dict = _construct_base_agent_event_dict(agent_name, agent_id, transaction)
6373
func_name = callable_name(wrapped)
6474
function_trace_name = f"{func_name}/{agent_name}"
65-
with FunctionTrace(name=function_trace_name, group="Llm", source=wrapped):
66-
return wrapped(*args, **kwargs)
75+
76+
ft = FunctionTrace(name=function_trace_name, group="Llm/agent/Autogen")
77+
ft.__enter__()
78+
79+
try:
80+
return_val = wrapped(*args, **kwargs)
81+
except Exception:
82+
ft.notice_error(attributes={"agent_id": agent_id})
83+
ft.__exit__(*sys.exc_info())
84+
# If we hit an exception, append the error attribute and duration from the exited function trace
85+
agent_event_dict.update({"duration": ft.duration * 1000, "error": True})
86+
transaction.record_custom_event("LlmAgent", agent_event_dict)
87+
raise
88+
89+
ft.__exit__(None, None, None)
90+
agent_event_dict.update({"duration": ft.duration * 1000})
91+
92+
transaction.record_custom_event("LlmAgent", agent_event_dict)
93+
94+
return return_val
95+
6796

6897

6998
def _get_llm_metadata(transaction):
@@ -115,6 +144,26 @@ def _construct_base_tool_event_dict(bound_args, tool_call_data, tool_id, transac
115144
return tool_event_dict
116145

117146

147+
def _construct_base_agent_event_dict(agent_name, agent_id, transaction):
148+
try:
149+
linking_metadata = get_trace_linking_metadata()
150+
151+
agent_event_dict = {
152+
"id": agent_id,
153+
"name": agent_name,
154+
"span_id": linking_metadata.get("span.id"),
155+
"trace_id": linking_metadata.get("trace.id"),
156+
"vendor": "autogen",
157+
"ingest_source": "Python",
158+
}
159+
agent_event_dict.update(_get_llm_metadata(transaction))
160+
except Exception:
161+
agent_event_dict = {}
162+
_logger.warning(RECORD_EVENTS_FAILURE_LOG_MESSAGE, exc_info=True)
163+
164+
return agent_event_dict
165+
166+
118167
async def wrap__execute_tool_call(wrapped, instance, args, kwargs):
119168
transaction = current_transaction()
120169
if not transaction:
@@ -150,7 +199,6 @@ async def wrap__execute_tool_call(wrapped, instance, args, kwargs):
150199
raise
151200

152201
ft.__exit__(None, None, None)
153-
154202
tool_event_dict.update({"duration": ft.duration * 1000})
155203

156204
# If the tool was executed successfully, we can grab the tool output from the result

‎newrelic/hooks/mlmodel_openai.py‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -442,6 +442,7 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa
442442
# The function trace will be exited when in the final iteration of the response
443443
# generator.
444444
return_val._nr_ft = ft
445+
return_val._nr_metadata = linking_metadata
445446
return_val._nr_openai_attrs = getattr(return_val, "_nr_openai_attrs", {})
446447
return_val._nr_openai_attrs["messages"] = kwargs.get("messages", [])
447448
return_val._nr_openai_attrs["temperature"] = kwargs.get("temperature")
@@ -492,10 +493,14 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa
492493
response_model = kwargs.get("response.model")
493494
response_id = kwargs.get("id")
494495
output_message_list = []
495-
finish_reason = None
496+
finish_reason = kwargs.get("finish_reason")
496497
if "content" in kwargs:
497498
output_message_list = [{"content": kwargs.get("content"), "role": kwargs.get("role")}]
498-
finish_reason = kwargs.get("finish_reason")
499+
# When tools are involved, the content key may hold an empty string which we do not want to report
500+
# In this case, the content we are interested in capturing will already be covered in the input_message_list
501+
# We empty out the output_message_list so that we do not report an empty message
502+
if "tool_call" in finish_reason and not kwargs.get("content"):
503+
output_message_list = []
499504
request_model = kwargs.get("model") or kwargs.get("engine")
500505

501506
request_id = response_headers.get("x-request-id")
@@ -765,7 +770,10 @@ def _record_stream_chunk(self, return_val):
765770

766771
def _record_events_on_stop_iteration(self, transaction):
767772
if hasattr(self, "_nr_ft"):
768-
linking_metadata = get_trace_linking_metadata()
773+
# We first check for our saved linking metadata before making a new call to get_trace_linking_metadata
774+
# Directly calling get_trace_linking_metadata() causes the incorrect span ID to be captured and associated with the LLM call
775+
# This leads to incorrect linking of the LLM call in the UI
776+
linking_metadata = self._nr_metadata or get_trace_linking_metadata()
769777
self._nr_ft.__exit__(None, None, None)
770778
try:
771779
openai_attrs = getattr(self, "_nr_openai_attrs", {})
@@ -872,6 +880,8 @@ def set_attrs_on_generator_proxy(proxy, instance):
872880
proxy._nr_response_headers = instance._nr_response_headers
873881
if hasattr(instance, "_nr_openai_attrs"):
874882
proxy._nr_openai_attrs = instance._nr_openai_attrs
883+
if hasattr(instance, "_nr_metadata"):
884+
proxy._nr_metadata = instance._nr_metadata
875885

876886

877887
def wrap_engine_api_resource_create_sync(wrapped, instance, args, kwargs):

‎tests/mlmodel_autogen/conftest.py‎

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,21 +13,14 @@
1313
# limitations under the License.
1414

1515
import json
16-
1716
import pytest
18-
from autogen_agentchat.agents import AssistantAgent
19-
from autogen_agentchat.base import TaskResult
20-
from autogen_agentchat.teams import RoundRobinGroupChat
21-
from autogen_core import ComponentModel, FunctionCall, Image
17+
18+
from autogen_core import FunctionCall
2219
from autogen_core.models import CreateResult, RequestUsage
23-
from autogen_core.models._model_client import ModelFamily
2420
from autogen_ext.models.replay import ReplayChatCompletionClient
25-
from pydantic import BaseModel, ValidationError
2621
from testing_support.fixture.event_loop import event_loop as loop
2722
from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture
2823

29-
from newrelic.common.object_names import callable_name
30-
3124
_default_settings = {
3225
"package_reporting.enabled": False, # Turn off package reporting for testing as it causes slowdowns.
3326
"transaction_tracer.explain_threshold": 0.0,
@@ -104,13 +97,13 @@ def multi_tool_model_client():
10497
),
10598
CreateResult(
10699
finish_reason="function_calls",
107-
content=[FunctionCall(id="2", name="add_exclamation", arguments=json.dumps({"message": "Goodbye"}))],
100+
content=[FunctionCall(id="3", name="compute_sum", arguments=json.dumps({"a": 5, "b": 3}))],
108101
usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
109102
cached=False,
110103
),
111104
CreateResult(
112105
finish_reason="function_calls",
113-
content=[FunctionCall(id="3", name="compute_sum", arguments=json.dumps({"a": 5, "b": 3}))],
106+
content=[FunctionCall(id="2", name="add_exclamation", arguments=json.dumps({"message": "Goodbye"}))],
114107
usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
115108
cached=False,
116109
),

‎tests/mlmodel_autogen/test_assistant_agent.py‎

Lines changed: 70 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
import copy
1615

1716
import pytest
17+
1818
from autogen_agentchat.agents import AssistantAgent
1919
from autogen_agentchat.base import TaskResult
2020
from testing_support.fixtures import reset_core_stats_engine, validate_attributes
@@ -34,6 +34,10 @@
3434
from newrelic.api.background_task import background_task
3535
from newrelic.api.llm_custom_attributes import WithLlmCustomAttributes
3636
from newrelic.common.object_names import callable_name
37+
from newrelic.common.package_version_utils import get_package_version_tuple
38+
39+
40+
AUTOGEN_VERSION = get_package_version_tuple("autogen-agentchat")
3741

3842
tool_recorded_event = [
3943
(
@@ -74,25 +78,49 @@
7478
]
7579

7680

81+
agent_recorded_event = [
82+
(
83+
{"type": "LlmAgent"},
84+
{
85+
"id": None,
86+
"name": "pirate_agent",
87+
"span_id": None,
88+
"trace_id": "trace-id",
89+
"vendor": "autogen",
90+
"ingest_source": "Python",
91+
"duration": None,
92+
},
93+
)
94+
]
95+
96+
7797
# Example tool for testing purposes
7898
def add_exclamation(message: str) -> str:
7999
return f"{message}!"
80100

81101

82102
@reset_core_stats_engine()
83-
@validate_custom_events(events_with_context_attrs(tool_recorded_event))
84-
@validate_custom_event_count(count=1)
103+
@validate_custom_events(
104+
events_with_context_attrs(tool_recorded_event) + events_with_context_attrs(agent_recorded_event)
105+
)
106+
@validate_custom_event_count(count=2)
85107
@validate_transaction_metrics(
86108
"test_assistant_agent:test_run_assistant_agent",
87109
scoped_metrics=[
88-
("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
110+
(
111+
"Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
112+
1,
113+
),
89114
(
90115
"Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
91116
1,
92117
),
93118
],
94119
rollup_metrics=[
95-
("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
120+
(
121+
"Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
122+
1,
123+
),
96124
(
97125
"Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
98126
1,
@@ -117,19 +145,25 @@ async def _test():
117145

118146

119147
@reset_core_stats_engine()
120-
@validate_custom_events(tool_recorded_event)
121-
@validate_custom_event_count(count=1)
148+
@validate_custom_events(tool_recorded_event + agent_recorded_event)
149+
@validate_custom_event_count(count=2)
122150
@validate_transaction_metrics(
123151
"test_assistant_agent:test_run_stream_assistant_agent",
124152
scoped_metrics=[
125-
("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
153+
(
154+
"Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
155+
1,
156+
),
126157
(
127158
"Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
128159
1,
129160
),
130161
],
131162
rollup_metrics=[
132-
("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
163+
(
164+
"Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
165+
1,
166+
),
133167
(
134168
"Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
135169
1,
@@ -162,19 +196,25 @@ async def _test():
162196

163197
@reset_core_stats_engine()
164198
@disabled_ai_monitoring_record_content_settings
165-
@validate_custom_events(tool_events_sans_content(tool_recorded_event))
166-
@validate_custom_event_count(count=1)
199+
@validate_custom_events(tool_events_sans_content(tool_recorded_event) + agent_recorded_event)
200+
@validate_custom_event_count(count=2)
167201
@validate_transaction_metrics(
168202
"test_assistant_agent:test_run_assistant_agent_no_content",
169203
scoped_metrics=[
170-
("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
204+
(
205+
"Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
206+
1,
207+
),
171208
(
172209
"Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
173210
1,
174211
),
175212
],
176213
rollup_metrics=[
177-
("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
214+
(
215+
"Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
216+
1,
217+
),
178218
(
179219
"Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
180220
1,
@@ -214,22 +254,35 @@ async def _test():
214254
loop.run_until_complete(_test())
215255

216256

257+
SKIP_IF_AUTOGEN_062 = pytest.mark.skipif(
258+
AUTOGEN_VERSION > (0, 6, 1),
259+
reason="Forcing invalid tool call arguments causes a hang on autogen versions above 0.6.1",
260+
)
261+
262+
263+
@SKIP_IF_AUTOGEN_062
217264
@reset_core_stats_engine()
218265
@validate_transaction_error_event_count(1)
219266
@validate_error_trace_attributes(callable_name(TypeError), exact_attrs={"agent": {}, "intrinsic": {}, "user": {}})
220267
@validate_custom_events(tool_recorded_event_error)
221-
@validate_custom_event_count(count=1)
268+
@validate_custom_event_count(count=2)
222269
@validate_transaction_metrics(
223270
"test_assistant_agent:test_run_assistant_agent_error",
224271
scoped_metrics=[
225-
("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
272+
(
273+
"Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
274+
1,
275+
),
226276
(
227277
"Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
228278
1,
229279
),
230280
],
231281
rollup_metrics=[
232-
("Llm/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent", 1),
282+
(
283+
"Llm/agent/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent.on_messages_stream/pirate_agent",
284+
1,
285+
),
233286
(
234287
"Llm/tool/Autogen/autogen_agentchat.agents._assistant_agent:AssistantAgent._execute_tool_call/add_exclamation",
235288
1,
@@ -250,7 +303,7 @@ def test_run_assistant_agent_error(loop, set_trace_info, single_tool_model_clien
250303

251304
async def _test():
252305
with pytest.raises(TypeError):
253-
response = await pirate_agent.run()
306+
await pirate_agent.run()
254307

255308
loop.run_until_complete(_test())
256309

0 commit comments

Comments
 (0)