Skip to content

Commit f12c534

Browse files
ericapisaniericapisani
authored
Merge branch 'master' into py-2089/use-heroku-build-commit-env-var
2 parents c5a2c07 + 594dce5 commit f12c534

File tree

4 files changed

+382
-180
lines changed

4 files changed

+382
-180
lines changed

‎sentry_sdk/integrations/google_genai/__init__.py‎

Lines changed: 38 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -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,

‎sentry_sdk/integrations/openai.py‎

Lines changed: 56 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@
5050
from sentry_sdk.tracing import Span
5151
from sentry_sdk._types import TextPart
5252

53-
from openai.types.responses import ResponseInputParam
53+
from openai.types.responses import ResponseInputParam, SequenceNotStr
5454
from openai import Omit
5555

5656
try:
@@ -220,20 +220,6 @@ def _calculate_token_usage(
220220
)
221221

222222

223-
def _get_input_messages(
224-
kwargs: "dict[str, Any]",
225-
) -> "Optional[Union[Iterable[Any], list[str]]]":
226-
# Input messages (the prompt or data sent to the model)
227-
messages = kwargs.get("messages")
228-
if messages is None:
229-
messages = kwargs.get("input")
230-
231-
if isinstance(messages, str):
232-
messages = [messages]
233-
234-
return messages
235-
236-
237223
def _commmon_set_input_data(
238224
span: "Span",
239225
kwargs: "dict[str, Any]",
@@ -371,13 +357,6 @@ def _set_completions_api_input_data(
371357
_commmon_set_input_data(span, kwargs)
372358
return
373359

374-
system_instructions = _get_system_instructions_completions(messages)
375-
if len(system_instructions) > 0:
376-
span.set_data(
377-
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
378-
json.dumps(_transform_system_instructions(system_instructions)),
379-
)
380-
381360
if isinstance(messages, str):
382361
normalized_messages = normalize_message_roles([messages]) # type: ignore
383362
scope = sentry_sdk.get_current_scope()
@@ -390,13 +369,29 @@ def _set_completions_api_input_data(
390369
_commmon_set_input_data(span, kwargs)
391370
return
392371

372+
# dict special case following https://github.com/openai/openai-python/blob/3e0c05b84a2056870abf3bd6a5e7849020209cc3/src/openai/_utils/_transform.py#L194-L197
373+
if not isinstance(messages, Iterable) or isinstance(messages, dict):
374+
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
375+
_commmon_set_input_data(span, kwargs)
376+
return
377+
378+
messages = list(messages)
379+
kwargs["messages"] = messages
380+
381+
system_instructions = _get_system_instructions_completions(messages)
382+
if len(system_instructions) > 0:
383+
span.set_data(
384+
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
385+
json.dumps(_transform_system_instructions(system_instructions)),
386+
)
387+
393388
non_system_messages = [
394389
message
395390
for message in messages
396391
if not _is_system_instruction_completions(message)
397392
]
398393
if len(non_system_messages) > 0:
399-
normalized_messages = normalize_message_roles(non_system_messages) # type: ignore
394+
normalized_messages = normalize_message_roles(non_system_messages)
400395
scope = sentry_sdk.get_current_scope()
401396
messages_data = truncate_and_annotate_messages(normalized_messages, span, scope)
402397
if messages_data is not None:
@@ -413,15 +408,47 @@ def _set_embeddings_input_data(
413408
kwargs: "dict[str, Any]",
414409
integration: "OpenAIIntegration",
415410
) -> None:
416-
messages = _get_input_messages(kwargs)
411+
messages: "Union[str, SequenceNotStr[str], Iterable[int], Iterable[Iterable[int]]]" = kwargs.get(
412+
"input"
413+
)
417414

418415
if (
419-
messages is not None
420-
and len(messages) > 0 # type: ignore
421-
and should_send_default_pii()
422-
and integration.include_prompts
416+
not should_send_default_pii()
417+
or not integration.include_prompts
418+
or messages is None
423419
):
424-
normalized_messages = normalize_message_roles(messages) # type: ignore
420+
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
421+
_commmon_set_input_data(span, kwargs)
422+
423+
return
424+
425+
if isinstance(messages, str):
426+
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
427+
_commmon_set_input_data(span, kwargs)
428+
429+
normalized_messages = normalize_message_roles([messages]) # type: ignore
430+
scope = sentry_sdk.get_current_scope()
431+
messages_data = truncate_and_annotate_embedding_inputs(
432+
normalized_messages, span, scope
433+
)
434+
if messages_data is not None:
435+
set_data_normalized(
436+
span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, messages_data, unpack=False
437+
)
438+
439+
return
440+
441+
# dict special case following https://github.com/openai/openai-python/blob/3e0c05b84a2056870abf3bd6a5e7849020209cc3/src/openai/_utils/_transform.py#L194-L197
442+
if not isinstance(messages, Iterable) or isinstance(messages, dict):
443+
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
444+
_commmon_set_input_data(span, kwargs)
445+
return
446+
447+
messages = list(messages)
448+
kwargs["input"] = messages
449+
450+
if len(messages) > 0:
451+
normalized_messages = normalize_message_roles(messages)
425452
scope = sentry_sdk.get_current_scope()
426453
messages_data = truncate_and_annotate_embedding_inputs(
427454
normalized_messages, span, scope

0 commit comments

Comments
 (0)