8
8
)
9
9
10
10
import sentry_sdk
11
+ from sentry_sdk .ai .utils import get_start_span_function
11
12
from sentry_sdk .integrations import DidNotEnable , Integration
12
13
from sentry_sdk .consts import OP , SPANDATA
14
+ from sentry_sdk .tracing import SPANSTATUS
13
15
14
16
15
17
try :
23
25
from .utils import (
24
26
set_span_data_for_request ,
25
27
set_span_data_for_response ,
26
- capture_exception ,
28
+ _capture_exception ,
27
29
prepare_generate_content_args ,
28
30
)
29
31
from .streaming import (
@@ -69,11 +71,12 @@ def new_generate_content_stream(self, *args, **kwargs):
69
71
70
72
_model , contents , model_name = prepare_generate_content_args (args , kwargs )
71
73
72
- span = sentry_sdk . start_span (
74
+ span = get_start_span_function () (
73
75
op = OP .GEN_AI_INVOKE_AGENT ,
74
76
name = "invoke_agent" ,
75
77
origin = ORIGIN ,
76
78
)
79
+ span .__enter__ ()
77
80
span .set_data (SPANDATA .GEN_AI_AGENT_NAME , model_name )
78
81
span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "invoke_agent" )
79
82
set_span_data_for_request (span , integration , model_name , contents , kwargs )
@@ -84,6 +87,7 @@ def new_generate_content_stream(self, *args, **kwargs):
84
87
name = f"chat { model_name } " ,
85
88
origin = ORIGIN ,
86
89
)
90
+ chat_span .__enter__ ()
87
91
chat_span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "chat" )
88
92
chat_span .set_data (SPANDATA .GEN_AI_SYSTEM , GEN_AI_SYSTEM )
89
93
chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
@@ -102,7 +106,8 @@ def new_iterator():
102
106
chunks .append (chunk )
103
107
yield chunk
104
108
except Exception as exc :
105
- capture_exception (exc )
109
+ _capture_exception (exc )
110
+ chat_span .set_status (SPANSTATUS .ERROR )
106
111
raise
107
112
finally :
108
113
# Accumulate all chunks and set final response data on spans
@@ -114,15 +119,15 @@ def new_iterator():
114
119
set_span_data_for_streaming_response (
115
120
span , integration , accumulated_response
116
121
)
117
- chat_span .finish ( )
118
- span .finish ( )
122
+ chat_span .__exit__ ( None , None , None )
123
+ span .__exit__ ( None , None , None )
119
124
120
125
return new_iterator ()
121
126
122
127
except Exception as exc :
123
- capture_exception (exc )
124
- chat_span .finish ( )
125
- span .finish ( )
128
+ _capture_exception (exc )
129
+ chat_span .__exit__ ( None , None , None )
130
+ span .__exit__ ( None , None , None )
126
131
raise
127
132
128
133
return new_generate_content_stream
@@ -139,11 +144,12 @@ async def new_async_generate_content_stream(self, *args, **kwargs):
139
144
140
145
_model , contents , model_name = prepare_generate_content_args (args , kwargs )
141
146
142
- span = sentry_sdk . start_span (
147
+ span = get_start_span_function () (
143
148
op = OP .GEN_AI_INVOKE_AGENT ,
144
149
name = "invoke_agent" ,
145
150
origin = ORIGIN ,
146
151
)
152
+ span .__enter__ ()
147
153
span .set_data (SPANDATA .GEN_AI_AGENT_NAME , model_name )
148
154
span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "invoke_agent" )
149
155
set_span_data_for_request (span , integration , model_name , contents , kwargs )
@@ -154,6 +160,7 @@ async def new_async_generate_content_stream(self, *args, **kwargs):
154
160
name = f"chat { model_name } " ,
155
161
origin = ORIGIN ,
156
162
)
163
+ chat_span .__enter__ ()
157
164
chat_span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "chat" )
158
165
chat_span .set_data (SPANDATA .GEN_AI_SYSTEM , GEN_AI_SYSTEM )
159
166
chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
@@ -172,7 +179,8 @@ async def new_async_iterator():
172
179
chunks .append (chunk )
173
180
yield chunk
174
181
except Exception as exc :
175
- capture_exception (exc )
182
+ _capture_exception (exc )
183
+ chat_span .set_status (SPANSTATUS .ERROR )
176
184
raise
177
185
finally :
178
186
# Accumulate all chunks and set final response data on spans
@@ -184,15 +192,15 @@ async def new_async_iterator():
184
192
set_span_data_for_streaming_response (
185
193
span , integration , accumulated_response
186
194
)
187
- chat_span .finish ( )
188
- span .finish ( )
195
+ chat_span .__exit__ ( None , None , None )
196
+ span .__exit__ ( None , None , None )
189
197
190
198
return new_async_iterator ()
191
199
192
200
except Exception as exc :
193
- capture_exception (exc )
194
- chat_span .finish ( )
195
- span .finish ( )
201
+ _capture_exception (exc )
202
+ chat_span .__exit__ ( None , None , None )
203
+ span .__exit__ ( None , None , None )
196
204
raise
197
205
198
206
return new_async_generate_content_stream
@@ -209,7 +217,7 @@ def new_generate_content(self, *args, **kwargs):
209
217
210
218
model , contents , model_name = prepare_generate_content_args (args , kwargs )
211
219
212
- with sentry_sdk . start_span (
220
+ with get_start_span_function () (
213
221
op = OP .GEN_AI_INVOKE_AGENT ,
214
222
name = "invoke_agent" ,
215
223
origin = ORIGIN ,
@@ -218,28 +226,29 @@ def new_generate_content(self, *args, **kwargs):
218
226
span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "invoke_agent" )
219
227
set_span_data_for_request (span , integration , model_name , contents , kwargs )
220
228
221
- try :
222
- with sentry_sdk .start_span (
223
- op = OP .GEN_AI_CHAT ,
224
- name = f"chat { model_name } " ,
225
- origin = ORIGIN ,
226
- ) as chat_span :
227
- chat_span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "chat" )
228
- chat_span .set_data (SPANDATA .GEN_AI_SYSTEM , GEN_AI_SYSTEM )
229
- chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
230
- set_span_data_for_request (
231
- chat_span , integration , model_name , contents , kwargs
232
- )
229
+ with sentry_sdk .start_span (
230
+ op = OP .GEN_AI_CHAT ,
231
+ name = f"chat { model_name } " ,
232
+ origin = ORIGIN ,
233
+ ) as chat_span :
234
+ chat_span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "chat" )
235
+ chat_span .set_data (SPANDATA .GEN_AI_SYSTEM , GEN_AI_SYSTEM )
236
+ chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
237
+ set_span_data_for_request (
238
+ chat_span , integration , model_name , contents , kwargs
239
+ )
233
240
241
+ try :
234
242
response = f (self , * args , ** kwargs )
243
+ except Exception as exc :
244
+ _capture_exception (exc )
245
+ chat_span .set_status (SPANSTATUS .ERROR )
246
+ raise
235
247
236
- set_span_data_for_response (chat_span , integration , response )
237
- set_span_data_for_response (span , integration , response )
248
+ set_span_data_for_response (chat_span , integration , response )
249
+ set_span_data_for_response (span , integration , response )
238
250
239
- return response
240
- except Exception as exc :
241
- capture_exception (exc )
242
- raise
251
+ return response
243
252
244
253
return new_generate_content
245
254
@@ -255,7 +264,7 @@ async def new_async_generate_content(self, *args, **kwargs):
255
264
256
265
model , contents , model_name = prepare_generate_content_args (args , kwargs )
257
266
258
- with sentry_sdk . start_span (
267
+ with get_start_span_function () (
259
268
op = OP .GEN_AI_INVOKE_AGENT ,
260
269
name = "invoke_agent" ,
261
270
origin = ORIGIN ,
@@ -264,27 +273,27 @@ async def new_async_generate_content(self, *args, **kwargs):
264
273
span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "invoke_agent" )
265
274
set_span_data_for_request (span , integration , model_name , contents , kwargs )
266
275
267
- try :
268
- with sentry_sdk .start_span (
269
- op = OP .GEN_AI_CHAT ,
270
- name = f"chat { model_name } " ,
271
- origin = ORIGIN ,
272
- ) as chat_span :
273
- chat_span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "chat" )
274
- chat_span .set_data (SPANDATA .GEN_AI_SYSTEM , GEN_AI_SYSTEM )
275
- chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
276
- set_span_data_for_request (
277
- chat_span , integration , model_name , contents , kwargs
278
- )
279
-
276
+ with sentry_sdk .start_span (
277
+ op = OP .GEN_AI_CHAT ,
278
+ name = f"chat { model_name } " ,
279
+ origin = ORIGIN ,
280
+ ) as chat_span :
281
+ chat_span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "chat" )
282
+ chat_span .set_data (SPANDATA .GEN_AI_SYSTEM , GEN_AI_SYSTEM )
283
+ chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
284
+ set_span_data_for_request (
285
+ chat_span , integration , model_name , contents , kwargs
286
+ )
287
+ try :
280
288
response = await f (self , * args , ** kwargs )
289
+ except Exception as exc :
290
+ _capture_exception (exc )
291
+ chat_span .set_status (SPANSTATUS .ERROR )
292
+ raise
281
293
282
- set_span_data_for_response (chat_span , integration , response )
283
- set_span_data_for_response (span , integration , response )
294
+ set_span_data_for_response (chat_span , integration , response )
295
+ set_span_data_for_response (span , integration , response )
284
296
285
- return response
286
- except Exception as exc :
287
- capture_exception (exc )
288
- raise
297
+ return response
289
298
290
299
return new_async_generate_content
0 commit comments