diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index 66dc4a1c48..a7af385d12 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -23,9 +23,20 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Iterable, List, Optional, Callable, AsyncIterator, Iterator + from typing import ( + Any, + Iterable, + List, + Optional, + Callable, + AsyncIterator, + Iterator, + Union, + ) from sentry_sdk.tracing import Span + from openai.types.responses import ResponseInputParam + try: try: from openai import NotGiven @@ -182,12 +193,9 @@ def _calculate_token_usage( ) -def _set_input_data( - span: "Span", +def _get_input_messages( kwargs: "dict[str, Any]", - operation: str, - integration: "OpenAIIntegration", -) -> None: +) -> "Optional[Union[Iterable[Any], list[str]]]": # Input messages (the prompt or data sent to the model) messages = kwargs.get("messages") if messages is None: @@ -196,29 +204,15 @@ def _set_input_data( if isinstance(messages, str): messages = [messages] - if ( - messages is not None - and len(messages) > 0 - and should_send_default_pii() - and integration.include_prompts - ): - normalized_messages = normalize_message_roles(messages) - scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) - if messages_data is not None: - # Use appropriate field based on operation type - if operation == "embeddings": - set_data_normalized( - span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, messages_data, unpack=False - ) - else: - set_data_normalized( - span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False - ) + return messages + +def _commmon_set_input_data( + span: "Span", + kwargs: "dict[str, Any]", +) -> None: # Input attributes: Common set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "openai") - set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, operation) # Input attributes: Optional kwargs_keys_to_attributes = { @@ -244,6 +238,85 @@ def _set_input_data( ) +def _set_responses_api_input_data( + span: "Span", + kwargs: "dict[str, Any]", + integration: "OpenAIIntegration", +) -> None: + messages: "Optional[Union[ResponseInputParam, list[str]]]" = _get_input_messages( + kwargs + ) + + if ( + messages is not None + and len(messages) > 0 + and should_send_default_pii() + and integration.include_prompts + ): + normalized_messages = normalize_message_roles(messages) # type: ignore + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + if messages_data is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False + ) + + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses") + _commmon_set_input_data(span, kwargs) + + +def _set_completions_api_input_data( + span: "Span", + kwargs: "dict[str, Any]", + integration: "OpenAIIntegration", +) -> None: + messages: "Optional[Union[Iterable[ChatCompletionMessageParam], list[str]]]" = ( + _get_input_messages(kwargs) + ) + + if ( + messages is not None + and len(messages) > 0 # type: ignore + and should_send_default_pii() + and integration.include_prompts + ): + normalized_messages = normalize_message_roles(messages) # type: ignore + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + if messages_data is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False + ) + + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") + _commmon_set_input_data(span, kwargs) + + +def _set_embeddings_input_data( + span: "Span", + kwargs: "dict[str, Any]", + integration: "OpenAIIntegration", +) -> None: + messages = _get_input_messages(kwargs) + + if ( + messages is not None + and len(messages) > 0 # type: ignore + and should_send_default_pii() + and integration.include_prompts + ): + normalized_messages = normalize_message_roles(messages) # type: ignore + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + if messages_data is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, messages_data, unpack=False + ) + + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + _commmon_set_input_data(span, kwargs) + + def _set_output_data( span: "Span", response: "Any", @@ -454,16 +527,15 @@ def _new_chat_completion_common(f: "Any", *args: "Any", **kwargs: "Any") -> "Any return f(*args, **kwargs) model = kwargs.get("model") - operation = "chat" span = sentry_sdk.start_span( op=consts.OP.GEN_AI_CHAT, - name=f"{operation} {model}", + name=f"chat {model}", origin=OpenAIIntegration.origin, ) span.__enter__() - _set_input_data(span, kwargs, operation, integration) + _set_completions_api_input_data(span, kwargs, integration) response = yield f, args, kwargs @@ -546,14 +618,13 @@ def _new_embeddings_create_common(f: "Any", *args: "Any", **kwargs: "Any") -> "A return f(*args, **kwargs) model = kwargs.get("model") - operation = "embeddings" with sentry_sdk.start_span( op=consts.OP.GEN_AI_EMBEDDINGS, - name=f"{operation} {model}", + name=f"embeddings {model}", origin=OpenAIIntegration.origin, ) as span: - _set_input_data(span, kwargs, operation, integration) + _set_embeddings_input_data(span, kwargs, integration) response = yield f, args, kwargs @@ -634,16 +705,15 @@ def _new_responses_create_common(f: "Any", *args: "Any", **kwargs: "Any") -> "An return f(*args, **kwargs) model = kwargs.get("model") - operation = "responses" span = sentry_sdk.start_span( op=consts.OP.GEN_AI_RESPONSES, - name=f"{operation} {model}", + name=f"responses {model}", origin=OpenAIIntegration.origin, ) span.__enter__() - _set_input_data(span, kwargs, operation, integration) + _set_responses_api_input_data(span, kwargs, integration) response = yield f, args, kwargs