@@ -73,18 +73,7 @@ def new_generate_content_stream(
7373
7474 _model , contents , model_name = prepare_generate_content_args (args , kwargs )
7575
76- span = get_start_span_function ()(
77- op = OP .GEN_AI_INVOKE_AGENT ,
78- name = "invoke_agent" ,
79- origin = ORIGIN ,
80- )
81- span .__enter__ ()
82- span .set_data (SPANDATA .GEN_AI_AGENT_NAME , model_name )
83- span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "invoke_agent" )
84- set_span_data_for_request (span , integration , model_name , contents , kwargs )
85- span .set_data (SPANDATA .GEN_AI_RESPONSE_STREAMING , True )
86-
87- chat_span = sentry_sdk .start_span (
76+ chat_span = get_start_span_function ()(
8877 op = OP .GEN_AI_CHAT ,
8978 name = f"chat { model_name } " ,
9079 origin = ORIGIN ,
@@ -95,7 +84,6 @@ def new_generate_content_stream(
9584 chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
9685 set_span_data_for_request (chat_span , integration , model_name , contents , kwargs )
9786 chat_span .set_data (SPANDATA .GEN_AI_RESPONSE_STREAMING , True )
98- chat_span .set_data (SPANDATA .GEN_AI_AGENT_NAME , model_name )
9987
10088 try :
10189 stream = f (self , * args , ** kwargs )
@@ -118,18 +106,13 @@ def new_iterator() -> "Iterator[Any]":
118106 set_span_data_for_streaming_response (
119107 chat_span , integration , accumulated_response
120108 )
121- set_span_data_for_streaming_response (
122- span , integration , accumulated_response
123- )
124109 chat_span .__exit__ (None , None , None )
125- span .__exit__ (None , None , None )
126110
127111 return new_iterator ()
128112
129113 except Exception as exc :
130114 _capture_exception (exc )
131115 chat_span .__exit__ (None , None , None )
132- span .__exit__ (None , None , None )
133116 raise
134117
135118 return new_generate_content_stream
@@ -148,18 +131,7 @@ async def new_async_generate_content_stream(
148131
149132 _model , contents , model_name = prepare_generate_content_args (args , kwargs )
150133
151- span = get_start_span_function ()(
152- op = OP .GEN_AI_INVOKE_AGENT ,
153- name = "invoke_agent" ,
154- origin = ORIGIN ,
155- )
156- span .__enter__ ()
157- span .set_data (SPANDATA .GEN_AI_AGENT_NAME , model_name )
158- span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "invoke_agent" )
159- set_span_data_for_request (span , integration , model_name , contents , kwargs )
160- span .set_data (SPANDATA .GEN_AI_RESPONSE_STREAMING , True )
161-
162- chat_span = sentry_sdk .start_span (
134+ chat_span = get_start_span_function ()(
163135 op = OP .GEN_AI_CHAT ,
164136 name = f"chat { model_name } " ,
165137 origin = ORIGIN ,
@@ -170,7 +142,6 @@ async def new_async_generate_content_stream(
170142 chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
171143 set_span_data_for_request (chat_span , integration , model_name , contents , kwargs )
172144 chat_span .set_data (SPANDATA .GEN_AI_RESPONSE_STREAMING , True )
173- chat_span .set_data (SPANDATA .GEN_AI_AGENT_NAME , model_name )
174145
175146 try :
176147 stream = await f (self , * args , ** kwargs )
@@ -193,18 +164,13 @@ async def new_async_iterator() -> "AsyncIterator[Any]":
193164 set_span_data_for_streaming_response (
194165 chat_span , integration , accumulated_response
195166 )
196- set_span_data_for_streaming_response (
197- span , integration , accumulated_response
198- )
199167 chat_span .__exit__ (None , None , None )
200- span .__exit__ (None , None , None )
201168
202169 return new_async_iterator ()
203170
204171 except Exception as exc :
205172 _capture_exception (exc )
206173 chat_span .__exit__ (None , None , None )
207- span .__exit__ (None , None , None )
208174 raise
209175
210176 return new_async_generate_content_stream
@@ -220,38 +186,27 @@ def new_generate_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
220186 model , contents , model_name = prepare_generate_content_args (args , kwargs )
221187
222188 with get_start_span_function ()(
223- op = OP .GEN_AI_INVOKE_AGENT ,
224- name = "invoke_agent " ,
189+ op = OP .GEN_AI_CHAT ,
190+ name = f"chat { model_name } " ,
225191 origin = ORIGIN ,
226- ) as span :
227- span .set_data (SPANDATA .GEN_AI_AGENT_NAME , model_name )
228- span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "invoke_agent" )
229- set_span_data_for_request (span , integration , model_name , contents , kwargs )
230-
231- with sentry_sdk .start_span (
232- op = OP .GEN_AI_CHAT ,
233- name = f"chat { model_name } " ,
234- origin = ORIGIN ,
235- ) as chat_span :
236- chat_span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "chat" )
237- chat_span .set_data (SPANDATA .GEN_AI_SYSTEM , GEN_AI_SYSTEM )
238- chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
239- chat_span .set_data (SPANDATA .GEN_AI_AGENT_NAME , model_name )
240- set_span_data_for_request (
241- chat_span , integration , model_name , contents , kwargs
242- )
192+ ) as chat_span :
193+ chat_span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "chat" )
194+ chat_span .set_data (SPANDATA .GEN_AI_SYSTEM , GEN_AI_SYSTEM )
195+ chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
196+ set_span_data_for_request (
197+ chat_span , integration , model_name , contents , kwargs
198+ )
243199
244- try :
245- response = f (self , * args , ** kwargs )
246- except Exception as exc :
247- _capture_exception (exc )
248- chat_span .set_status (SPANSTATUS .INTERNAL_ERROR )
249- raise
200+ try :
201+ response = f (self , * args , ** kwargs )
202+ except Exception as exc :
203+ _capture_exception (exc )
204+ chat_span .set_status (SPANSTATUS .INTERNAL_ERROR )
205+ raise
250206
251- set_span_data_for_response (chat_span , integration , response )
252- set_span_data_for_response (span , integration , response )
207+ set_span_data_for_response (chat_span , integration , response )
253208
254- return response
209+ return response
255210
256211 return new_generate_content
257212
@@ -268,36 +223,26 @@ async def new_async_generate_content(
268223 model , contents , model_name = prepare_generate_content_args (args , kwargs )
269224
270225 with get_start_span_function ()(
271- op = OP .GEN_AI_INVOKE_AGENT ,
272- name = "invoke_agent " ,
226+ op = OP .GEN_AI_CHAT ,
227+ name = f"chat { model_name } " ,
273228 origin = ORIGIN ,
274- ) as span :
275- span .set_data (SPANDATA .GEN_AI_AGENT_NAME , model_name )
276- span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "invoke_agent" )
277- set_span_data_for_request (span , integration , model_name , contents , kwargs )
278-
279- with sentry_sdk .start_span (
280- op = OP .GEN_AI_CHAT ,
281- name = f"chat { model_name } " ,
282- origin = ORIGIN ,
283- ) as chat_span :
284- chat_span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "chat" )
285- chat_span .set_data (SPANDATA .GEN_AI_SYSTEM , GEN_AI_SYSTEM )
286- chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
287- set_span_data_for_request (
288- chat_span , integration , model_name , contents , kwargs
289- )
290- try :
291- response = await f (self , * args , ** kwargs )
292- except Exception as exc :
293- _capture_exception (exc )
294- chat_span .set_status (SPANSTATUS .INTERNAL_ERROR )
295- raise
229+ ) as chat_span :
230+ chat_span .set_data (SPANDATA .GEN_AI_OPERATION_NAME , "chat" )
231+ chat_span .set_data (SPANDATA .GEN_AI_SYSTEM , GEN_AI_SYSTEM )
232+ chat_span .set_data (SPANDATA .GEN_AI_REQUEST_MODEL , model_name )
233+ set_span_data_for_request (
234+ chat_span , integration , model_name , contents , kwargs
235+ )
236+ try :
237+ response = await f (self , * args , ** kwargs )
238+ except Exception as exc :
239+ _capture_exception (exc )
240+ chat_span .set_status (SPANSTATUS .INTERNAL_ERROR )
241+ raise
296242
297- set_span_data_for_response (chat_span , integration , response )
298- set_span_data_for_response (span , integration , response )
243+ set_span_data_for_response (chat_span , integration , response )
299244
300- return response
245+ return response
301246
302247 return new_async_generate_content
303248
@@ -311,7 +256,7 @@ def new_embed_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
311256
312257 model_name , contents = prepare_embed_content_args (args , kwargs )
313258
314- with sentry_sdk . start_span (
259+ with get_start_span_function () (
315260 op = OP .GEN_AI_EMBEDDINGS ,
316261 name = f"embeddings { model_name } " ,
317262 origin = ORIGIN ,
@@ -346,7 +291,7 @@ async def new_async_embed_content(
346291
347292 model_name , contents = prepare_embed_content_args (args , kwargs )
348293
349- with sentry_sdk . start_span (
294+ with get_start_span_function () (
350295 op = OP .GEN_AI_EMBEDDINGS ,
351296 name = f"embeddings { model_name } " ,
352297 origin = ORIGIN ,
0 commit comments