@@ -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