ref(openai): Separate output handling (#5543)

Create separate functions for output-handling of `openai` Responses, Completions, and Embeddings API functions. Create streaming variants where applicable.

Author: alexander-alderman-webb

sentry_sdk/integrations/openai.py

@@ -462,7 +462,7 @@ def _set_embeddings_input_data(
     _commmon_set_input_data(span, kwargs)
 
 
-def _set_output_data(
+def _set_common_output_data(
     span: "Span",
     response: "Any",
     kwargs: "dict[str, Any]",
@@ -706,11 +706,109 @@ def _new_chat_completion_common(f: "Any", *args: "Any", **kwargs: "Any") -> "Any
     start_time = time.perf_counter()
     response = yield f, args, kwargs
 
-    _set_output_data(span, response, kwargs, integration, start_time, finish_span=True)
+    is_streaming_response = kwargs.get("stream", False)
+    if is_streaming_response:
+        _set_streaming_completions_api_output_data(
+            span, response, kwargs, integration, start_time, finish_span=True
+        )
+    else:
+        _set_completions_api_output_data(
+            span, response, kwargs, integration, start_time, finish_span=True
+        )
 
     return response
 
 
+def _set_completions_api_output_data(
+    span: "Span",
+    response: "Any",
+    kwargs: "dict[str, Any]",
+    integration: "OpenAIIntegration",
+    start_time: "Optional[float]" = None,
+    finish_span: bool = True,
+) -> None:
+    _set_common_output_data(
+        span,
+        response,
+        kwargs,
+        integration,
+        start_time,
+        finish_span,
+    )
+
+
+def _set_streaming_completions_api_output_data(
+    span: "Span",
+    response: "Any",
+    kwargs: "dict[str, Any]",
+    integration: "OpenAIIntegration",
+    start_time: "Optional[float]" = None,
+    finish_span: bool = True,
+) -> None:
+    _set_common_output_data(
+        span,
+        response,
+        kwargs,
+        integration,
+        start_time,
+        finish_span,
+    )
+
+
+def _set_responses_api_output_data(
+    span: "Span",
+    response: "Any",
+    kwargs: "dict[str, Any]",
+    integration: "OpenAIIntegration",
+    start_time: "Optional[float]" = None,
+    finish_span: bool = True,
+) -> None:
+    _set_common_output_data(
+        span,
+        response,
+        kwargs,
+        integration,
+        start_time,
+        finish_span,
+    )
+
+
+def _set_streaming_responses_api_output_data(
+    span: "Span",
+    response: "Any",
+    kwargs: "dict[str, Any]",
+    integration: "OpenAIIntegration",
+    start_time: "Optional[float]" = None,
+    finish_span: bool = True,
+) -> None:
+    _set_common_output_data(
+        span,
+        response,
+        kwargs,
+        integration,
+        start_time,
+        finish_span,
+    )
+
+
+def _set_embeddings_output_data(
+    span: "Span",
+    response: "Any",
+    kwargs: "dict[str, Any]",
+    integration: "OpenAIIntegration",
+    start_time: "Optional[float]" = None,
+    finish_span: bool = True,
+) -> None:
+    _set_common_output_data(
+        span,
+        response,
+        kwargs,
+        integration,
+        start_time,
+        finish_span,
+    )
+
+
 def _wrap_chat_completion_create(f: "Callable[..., Any]") -> "Callable[..., Any]":
     def _execute_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any":
         gen = _new_chat_completion_common(f, *args, **kwargs)
@@ -795,7 +893,9 @@ def _new_embeddings_create_common(f: "Any", *args: "Any", **kwargs: "Any") -> "A
 
         response = yield f, args, kwargs
 
-        _set_output_data(span, response, kwargs, integration, finish_span=False)
+        _set_embeddings_output_data(
+            span, response, kwargs, integration, finish_span=False
+        )
 
         return response
 
@@ -885,7 +985,15 @@ def _new_responses_create_common(f: "Any", *args: "Any", **kwargs: "Any") -> "An
     start_time = time.perf_counter()
     response = yield f, args, kwargs
 
-    _set_output_data(span, response, kwargs, integration, start_time, finish_span=True)
+    is_streaming_response = kwargs.get("stream", False)
+    if is_streaming_response:
+        _set_streaming_responses_api_output_data(
+            span, response, kwargs, integration, start_time, finish_span=True
+        )
+    else:
+        _set_responses_api_output_data(
+            span, response, kwargs, integration, start_time, finish_span=True
+        )
 
     return response