12
12
# See the License for the specific language governing permissions and
13
13
# limitations under the License.
14
14
15
- import logging
16
- import json
17
- import sys
18
- import traceback
19
- import uuid
20
-
21
15
from newrelic .api .function_trace import FunctionTrace
22
- from newrelic .api .time_trace import get_trace_linking_metadata
23
- from newrelic .common .package_version_utils import get_package_version
24
16
from newrelic .api .transaction import current_transaction
25
17
from newrelic .common .object_names import callable_name
26
18
from newrelic .common .object_wrapper import wrap_function_wrapper
27
19
from newrelic .common .signature import bind_args
28
- from newrelic .core .config import global_settings
29
-
30
- # Check for the presence of the autogen-core, autogen-agentchat, or autogen-ext package as they should all have the
31
- # same version and one or multiple could be installed
32
- AUTOGEN_VERSION = (
33
- get_package_version ("autogen-core" )
34
- or get_package_version ("autogen-agentchat" )
35
- or get_package_version ("autogen-ext" )
36
- )
37
-
38
- RECORD_EVENTS_FAILURE_LOG_MESSAGE = "Exception occurred in Autogen instrumentation: Failed to record LLM events. Please report this issue to New Relic Support.\n %s"
39
-
40
- _logger = logging .getLogger (__name__ )
41
20
42
21
43
22
async def wrap_from_server_params (wrapped , instance , args , kwargs ):
@@ -53,126 +32,6 @@ async def wrap_from_server_params(wrapped, instance, args, kwargs):
53
32
return await wrapped (* args , ** kwargs )
54
33
55
34
56
- def wrap_on_messages_stream (wrapped , instance , args , kwargs ):
57
- transaction = current_transaction ()
58
- if not transaction :
59
- return wrapped (* args , ** kwargs )
60
-
61
- agent_name = getattr (instance , "name" , "agent" )
62
- func_name = callable_name (wrapped )
63
- function_trace_name = f"{ func_name } /{ agent_name } "
64
- with FunctionTrace (name = function_trace_name , group = "Llm" , source = wrapped ):
65
- return wrapped (* args , ** kwargs )
66
-
67
-
68
- def _get_llm_metadata (transaction ):
69
- # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events
70
- custom_attrs_dict = transaction ._custom_params
71
- llm_metadata_dict = {key : value for key , value in custom_attrs_dict .items () if key .startswith ("llm." )}
72
- llm_context_attrs = getattr (transaction , "_llm_context_attrs" , None )
73
- if llm_context_attrs :
74
- llm_metadata_dict .update (llm_context_attrs )
75
-
76
- return llm_metadata_dict
77
-
78
-
79
- def _extract_tool_output (return_val , tool_name ):
80
- try :
81
- output = getattr (return_val [1 ], "content" , None )
82
- return output
83
- except Exception :
84
- _logger .warning (f"Unable to parse tool output value from { tool_name } . Omitting output from LlmTool event." )
85
- return None
86
-
87
-
88
- def _construct_base_tool_event_dict (bound_args , tool_call_data , tool_id , transaction , settings ):
89
- try :
90
- input = getattr (tool_call_data , "arguments" , None )
91
- tool_input = str (input ) if input else None
92
- run_id = getattr (tool_call_data , "id" , None )
93
- tool_name = getattr (tool_call_data , "name" , "tool" )
94
- agent_name = bound_args .get ("agent_name" )
95
- linking_metadata = get_trace_linking_metadata ()
96
-
97
- tool_event_dict = {
98
- "id" : tool_id ,
99
- "run_id" : run_id ,
100
- "name" : tool_name ,
101
- "span_id" : linking_metadata .get ("span.id" ),
102
- "trace_id" : linking_metadata .get ("trace.id" ),
103
- "agent_name" : agent_name ,
104
- "vendor" : "autogen" ,
105
- "ingest_source" : "Python" ,
106
- }
107
- if settings .ai_monitoring .record_content .enabled :
108
- tool_event_dict .update ({"input" : tool_input })
109
- tool_event_dict .update (_get_llm_metadata (transaction ))
110
- except Exception :
111
- tool_event_dict = {}
112
- _logger .warning (RECORD_EVENTS_FAILURE_LOG_MESSAGE , exc_info = True )
113
-
114
- return tool_event_dict
115
-
116
-
117
- async def wrap__execute_tool_call (wrapped , instance , args , kwargs ):
118
- transaction = current_transaction ()
119
- if not transaction :
120
- return await wrapped (* args , ** kwargs )
121
-
122
- settings = transaction .settings if transaction .settings is not None else global_settings ()
123
- if not settings .ai_monitoring .enabled :
124
- return await wrapped (* args , ** kwargs )
125
-
126
- # Framework metric also used for entity tagging in the UI
127
- transaction .add_ml_model_info ("Autogen" , AUTOGEN_VERSION )
128
- transaction ._add_agent_attribute ("llm" , True )
129
-
130
- tool_id = str (uuid .uuid4 ())
131
- bound_args = bind_args (wrapped , args , kwargs )
132
- tool_call_data = bound_args .get ("tool_call" )
133
- tool_event_dict = _construct_base_tool_event_dict (bound_args , tool_call_data , tool_id , transaction , settings )
134
-
135
- tool_name = getattr (tool_call_data , "name" , "tool" )
136
-
137
- func_name = callable_name (wrapped )
138
- ft = FunctionTrace (name = f"{ func_name } /{ tool_name } " , group = "Llm/tool/Autogen" )
139
- ft .__enter__ ()
140
-
141
- try :
142
- return_val = await wrapped (* args , ** kwargs )
143
- except Exception :
144
- ft .notice_error (attributes = {"tool_id" : tool_id })
145
- ft .__exit__ (* sys .exc_info ())
146
- # If we hit an exception, append the error attribute and duration from the exited function trace
147
- tool_event_dict .update ({"duration" : ft .duration * 1000 , "error" : True })
148
- transaction .record_custom_event ("LlmTool" , tool_event_dict )
149
- raise
150
-
151
- ft .__exit__ (None , None , None )
152
-
153
- if not return_val :
154
- return return_val
155
-
156
- tool_event_dict .update ({"duration" : ft .duration * 1000 })
157
-
158
- # If the tool was executed successfully, we can grab the tool output from the result
159
- tool_output = _extract_tool_output (return_val , tool_name )
160
- if settings .ai_monitoring .record_content .enabled :
161
- tool_event_dict .update ({"output" : tool_output })
162
-
163
- transaction .record_custom_event ("LlmTool" , tool_event_dict )
164
-
165
- return return_val
166
-
167
-
168
- def instrument_autogen_agentchat_agents__assistant_agent (module ):
169
- if hasattr (module , "AssistantAgent" ):
170
- if hasattr (module .AssistantAgent , "on_messages_stream" ):
171
- wrap_function_wrapper (module , "AssistantAgent.on_messages_stream" , wrap_on_messages_stream )
172
- if hasattr (module .AssistantAgent , "_execute_tool_call" ):
173
- wrap_function_wrapper (module , "AssistantAgent._execute_tool_call" , wrap__execute_tool_call )
174
-
175
-
176
35
def instrument_autogen_ext_tools_mcp__base (module ):
177
36
if hasattr (module , "McpToolAdapter" ):
178
37
if hasattr (module .McpToolAdapter , "from_server_params" ):
0 commit comments