ref(openai): Separate input handling

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