Skip to content

Commit 9e03a64

Browse files
committed
apply all the feedback
1 parent 3d165f3 commit 9e03a64

File tree

5 files changed

+92
-104
lines changed

5 files changed

+92
-104
lines changed

sentry_sdk/integrations/google_genai/__init__.py

Lines changed: 63 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,10 @@
88
)
99

1010
import sentry_sdk
11+
from sentry_sdk.ai.utils import get_start_span_function
1112
from sentry_sdk.integrations import DidNotEnable, Integration
1213
from sentry_sdk.consts import OP, SPANDATA
14+
from sentry_sdk.tracing import SPANSTATUS
1315

1416

1517
try:
@@ -23,7 +25,7 @@
2325
from .utils import (
2426
set_span_data_for_request,
2527
set_span_data_for_response,
26-
capture_exception,
28+
_capture_exception,
2729
prepare_generate_content_args,
2830
)
2931
from .streaming import (
@@ -69,11 +71,12 @@ def new_generate_content_stream(self, *args, **kwargs):
6971

7072
_model, contents, model_name = prepare_generate_content_args(args, kwargs)
7173

72-
span = sentry_sdk.start_span(
74+
span = get_start_span_function()(
7375
op=OP.GEN_AI_INVOKE_AGENT,
7476
name="invoke_agent",
7577
origin=ORIGIN,
7678
)
79+
span.__enter__()
7780
span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
7881
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
7982
set_span_data_for_request(span, integration, model_name, contents, kwargs)
@@ -84,6 +87,7 @@ def new_generate_content_stream(self, *args, **kwargs):
8487
name=f"chat {model_name}",
8588
origin=ORIGIN,
8689
)
90+
chat_span.__enter__()
8791
chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
8892
chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM)
8993
chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
@@ -102,7 +106,8 @@ def new_iterator():
102106
chunks.append(chunk)
103107
yield chunk
104108
except Exception as exc:
105-
capture_exception(exc)
109+
_capture_exception(exc)
110+
chat_span.set_status(SPANSTATUS.ERROR)
106111
raise
107112
finally:
108113
# Accumulate all chunks and set final response data on spans
@@ -114,15 +119,15 @@ def new_iterator():
114119
set_span_data_for_streaming_response(
115120
span, integration, accumulated_response
116121
)
117-
chat_span.finish()
118-
span.finish()
122+
chat_span.__exit__(None, None, None)
123+
span.__exit__(None, None, None)
119124

120125
return new_iterator()
121126

122127
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)
126131
raise
127132

128133
return new_generate_content_stream
@@ -139,11 +144,12 @@ async def new_async_generate_content_stream(self, *args, **kwargs):
139144

140145
_model, contents, model_name = prepare_generate_content_args(args, kwargs)
141146

142-
span = sentry_sdk.start_span(
147+
span = get_start_span_function()(
143148
op=OP.GEN_AI_INVOKE_AGENT,
144149
name="invoke_agent",
145150
origin=ORIGIN,
146151
)
152+
span.__enter__()
147153
span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
148154
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
149155
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):
154160
name=f"chat {model_name}",
155161
origin=ORIGIN,
156162
)
163+
chat_span.__enter__()
157164
chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
158165
chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM)
159166
chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
@@ -172,7 +179,8 @@ async def new_async_iterator():
172179
chunks.append(chunk)
173180
yield chunk
174181
except Exception as exc:
175-
capture_exception(exc)
182+
_capture_exception(exc)
183+
chat_span.set_status(SPANSTATUS.ERROR)
176184
raise
177185
finally:
178186
# Accumulate all chunks and set final response data on spans
@@ -184,15 +192,15 @@ async def new_async_iterator():
184192
set_span_data_for_streaming_response(
185193
span, integration, accumulated_response
186194
)
187-
chat_span.finish()
188-
span.finish()
195+
chat_span.__exit__(None, None, None)
196+
span.__exit__(None, None, None)
189197

190198
return new_async_iterator()
191199

192200
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)
196204
raise
197205

198206
return new_async_generate_content_stream
@@ -209,7 +217,7 @@ def new_generate_content(self, *args, **kwargs):
209217

210218
model, contents, model_name = prepare_generate_content_args(args, kwargs)
211219

212-
with sentry_sdk.start_span(
220+
with get_start_span_function()(
213221
op=OP.GEN_AI_INVOKE_AGENT,
214222
name="invoke_agent",
215223
origin=ORIGIN,
@@ -218,28 +226,29 @@ def new_generate_content(self, *args, **kwargs):
218226
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
219227
set_span_data_for_request(span, integration, model_name, contents, kwargs)
220228

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+
)
233240

241+
try:
234242
response = f(self, *args, **kwargs)
243+
except Exception as exc:
244+
_capture_exception(exc)
245+
chat_span.set_status(SPANSTATUS.ERROR)
246+
raise
235247

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)
238250

239-
return response
240-
except Exception as exc:
241-
capture_exception(exc)
242-
raise
251+
return response
243252

244253
return new_generate_content
245254

@@ -255,7 +264,7 @@ async def new_async_generate_content(self, *args, **kwargs):
255264

256265
model, contents, model_name = prepare_generate_content_args(args, kwargs)
257266

258-
with sentry_sdk.start_span(
267+
with get_start_span_function()(
259268
op=OP.GEN_AI_INVOKE_AGENT,
260269
name="invoke_agent",
261270
origin=ORIGIN,
@@ -264,27 +273,27 @@ async def new_async_generate_content(self, *args, **kwargs):
264273
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
265274
set_span_data_for_request(span, integration, model_name, contents, kwargs)
266275

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:
280288
response = await f(self, *args, **kwargs)
289+
except Exception as exc:
290+
_capture_exception(exc)
291+
chat_span.set_status(SPANSTATUS.ERROR)
292+
raise
281293

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)
284296

285-
return response
286-
except Exception as exc:
287-
capture_exception(exc)
288-
raise
297+
return response
289298

290299
return new_async_generate_content

sentry_sdk/integrations/google_genai/streaming.py

Lines changed: 10 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,10 @@ def accumulate_streaming_response(chunks):
4141

4242
for chunk in chunks:
4343
# Extract text and tool calls
44-
if hasattr(chunk, "candidates") and chunk.candidates:
44+
if getattr(chunk, "candidates", None):
4545
for candidate in chunk.candidates:
46-
if hasattr(candidate, "content") and hasattr(
47-
candidate.content, "parts"
46+
if hasattr(candidate, "content") and getattr(
47+
candidate.content, "parts", []
4848
):
4949
extracted_text = extract_contents_text(candidate.content)
5050
if extracted_text:
@@ -59,41 +59,23 @@ def accumulate_streaming_response(chunks):
5959
tool_calls.extend(extracted_tool_calls)
6060

6161
# Accumulate token usage
62-
if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
62+
if getattr(chunk, "usage_metadata", None):
6363
usage = chunk.usage_metadata
64-
if (
65-
hasattr(usage, "prompt_token_count")
66-
and usage.prompt_token_count is not None
67-
):
64+
if getattr(usage, "prompt_token_count", None):
6865
total_prompt_tokens = max(total_prompt_tokens, usage.prompt_token_count)
69-
if (
70-
hasattr(usage, "tool_use_prompt_token_count")
71-
and usage.tool_use_prompt_token_count is not None
72-
):
66+
if getattr(usage, "tool_use_prompt_token_count", None):
7367
total_tool_use_prompt_tokens = max(
7468
total_tool_use_prompt_tokens, usage.tool_use_prompt_token_count
7569
)
76-
if (
77-
hasattr(usage, "candidates_token_count")
78-
and usage.candidates_token_count is not None
79-
):
70+
if getattr(usage, "candidates_token_count", None):
8071
total_output_tokens += usage.candidates_token_count
81-
if (
82-
hasattr(usage, "cached_content_token_count")
83-
and usage.cached_content_token_count is not None
84-
):
72+
if getattr(usage, "cached_content_token_count", None):
8573
total_cached_tokens = max(
8674
total_cached_tokens, usage.cached_content_token_count
8775
)
88-
if (
89-
hasattr(usage, "thoughts_token_count")
90-
and usage.thoughts_token_count is not None
91-
):
76+
if getattr(usage, "thoughts_token_count", None):
9277
total_reasoning_tokens += usage.thoughts_token_count
93-
if (
94-
hasattr(usage, "total_token_count")
95-
and usage.total_token_count is not None
96-
):
78+
if getattr(usage, "total_token_count", None):
9779
# Only use the final total_token_count from the last chunk
9880
total_tokens = usage.total_token_count
9981

0 commit comments

Comments
 (0)