feat: add instrumentation around openai responses stream method.

instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/131.added

@@ -0,0 +1 @@
+Add OpenAI Responses API stream instrumentation for sync and async clients.

instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/__init__.py

@@ -80,7 +80,9 @@
 )
 from .patch_responses import (
     async_responses_create,
+    async_responses_stream,
     responses_create,
+    responses_stream,
 )
 
 
@@ -171,11 +173,21 @@ def _instrument(self, **kwargs):
                 "Responses.create",
                 responses_create(handler),
             )
+            wrap_function_wrapper(
+                "openai.resources.responses.responses",
+                "Responses.stream",
+                responses_stream(handler),
+            )
             wrap_function_wrapper(
                 "openai.resources.responses.responses",
                 "AsyncResponses.create",
                 async_responses_create(handler),
             )
+            wrap_function_wrapper(
+                "openai.resources.responses.responses",
+                "AsyncResponses.stream",
+                async_responses_stream(handler),
+            )
 
     def _uninstrument(self, **kwargs):
         import openai  # pylint: disable=import-outside-toplevel  # noqa: PLC0415
@@ -190,8 +202,9 @@ def _uninstrument(self, **kwargs):
         responses_module = _get_responses_module()
         if responses_module is not None:
             unwrap(responses_module.Responses, "create")
-            if hasattr(responses_module, "AsyncResponses"):
-                unwrap(responses_module.AsyncResponses, "create")
+            unwrap(responses_module.Responses, "stream")
+            unwrap(responses_module.AsyncResponses, "create")
+            unwrap(responses_module.AsyncResponses, "stream")
 
 
 def _get_responses_module():

instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/patch_responses.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Union, cast
+
 from opentelemetry.util.genai.handler import TelemetryHandler
 
 from .response_extractors import (
@@ -12,18 +14,84 @@
     set_invocation_response_attributes,
 )
 from .response_wrappers import (
+    AsyncResponseStreamManagerWrapper,
     AsyncResponseStreamWrapper,
+    ResponseStreamManagerWrapper,
     ResponseStreamWrapper,
+    responses_stream_context,
 )
-from .utils import is_streaming
+
+if TYPE_CHECKING:
+    from openai import AsyncStream as OpenAIAsyncStream
+    from openai import Stream as OpenAIStream
+    from openai.lib.streaming.responses._responses import (  # pylint: disable=no-name-in-module
+        AsyncResponseStream,
+        AsyncResponseStreamManager,
+        ResponseStream,
+        ResponseStreamManager,
+    )
+    from openai.resources.responses.responses import AsyncResponses, Responses
+    from openai.types.responses import (  # pylint: disable=no-name-in-module
+        ParsedResponse,
+        Response,
+    )
+
+try:
+    from openai import AsyncStream as _OpenAIAsyncStream
+    from openai import Stream as _OpenAIStream
+    from openai.lib.streaming.responses._responses import (  # pylint: disable=no-name-in-module
+        AsyncResponseStream as _AsyncResponseStream,
+    )
+    from openai.lib.streaming.responses._responses import (  # pylint: disable=no-name-in-module
+        ResponseStream as _ResponseStream,
+    )
+except ImportError:
+    _AsyncResponseStream = None
+    _OpenAIAsyncStream = None
+    _OpenAIStream = None
+    _ResponseStream = None
+
+ResponseResult = Union["ParsedResponse[Any]", "Response"]
+ResponseStreamResult = Union["OpenAIStream[Any]", "ResponseStream[Any]"]
+AsyncResponseStreamResult = Union[
+    "OpenAIAsyncStream[Any]", "AsyncResponseStream[Any]"
+]
 
 
-def responses_create(handler: TelemetryHandler):
-    """Wrap the `create` method of the `Responses` class to trace it."""
+def responses_create(
+    handler: TelemetryHandler,
+) -> Callable[
+    ...,
+    Union[
+        ResponseResult,
+        ResponseStreamResult,
+        ResponseStreamWrapper[Any],
+    ],
+]:
+    """Wrap ``Responses.create`` to trace Responses API calls.
+
+    Traces :meth:`openai.resources.responses.responses.Responses.create`.
+    OpenAI SDK source:
+    https://github.com/openai/openai-python/blob/main/src/openai/resources/responses/responses.py#L914
+    """
 
     capture_content = handler.should_capture_content()
 
-    def traced_method(wrapped, instance, args, kwargs):
+    def traced_method(
+        wrapped: Callable[..., Union[ResponseResult, ResponseStreamResult]],
+        instance: "Responses",
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Union[
+        ResponseResult,
+        ResponseStreamResult,
+        ResponseStreamWrapper[Any],
+    ]:
+        stream_context = responses_stream_context.get()
+        if stream_context is not None:
+            result = wrapped(*args, **kwargs)
+            return _get_response_stream_result(result)
+
         params = extract_params(**kwargs)
         invocation = handler.inference(
             **get_inference_creation_kwargs(params, instance)
@@ -34,33 +102,77 @@ def traced_method(wrapped, instance, args, kwargs):
             result = wrapped(*args, **kwargs)
             parsed_result = _get_response_stream_result(result)
 
-            if is_streaming(kwargs):
+            if (
+                _ResponseStream is not None
+                and isinstance(parsed_result, _ResponseStream)
+            ) or (
+                _OpenAIStream is not None
+                and isinstance(parsed_result, _OpenAIStream)
+            ):
                 return ResponseStreamWrapper(
-                    parsed_result,
+                    cast("ResponseStreamResult", parsed_result),
                     invocation,
                     capture_content,
                 )
 
             set_invocation_response_attributes(
-                invocation, parsed_result, capture_content
+                invocation,
+                cast("ResponseResult", parsed_result),
+                capture_content,
             )
             invocation.stop()
             return result
         except Exception as error:
             invocation.fail(error)
             raise
 
-    return traced_method
+    return cast(
+        'Callable[..., Union["ResponseResult", "ResponseStreamResult", ResponseStreamWrapper[Any]]]',
+        traced_method,
+    )
+
 
+def async_responses_create(
+    handler: TelemetryHandler,
+) -> Callable[
+    ...,
+    Awaitable[
+        Union[
+            ResponseResult,
+            AsyncResponseStreamResult,
+            AsyncResponseStreamWrapper[Any],
+        ]
+    ],
+]:
+    """Wrap ``AsyncResponses.create`` to trace async Responses API calls.
 
-def async_responses_create(handler: TelemetryHandler):
-    """Wrap the `create` method of the `AsyncResponses` class to trace it."""
+    Traces :meth:`openai.resources.responses.responses.AsyncResponses.create`.
+    OpenAI SDK source:
+    https://github.com/openai/openai-python/blob/main/src/openai/resources/responses/responses.py#L2661
+    """
 
     capture_content = handler.should_capture_content()
 
-    async def traced_method(wrapped, instance, args, kwargs):
+    async def traced_method(
+        wrapped: Callable[
+            ...,
+            Awaitable[Union[ResponseResult, AsyncResponseStreamResult]],
+        ],
+        instance: "AsyncResponses",
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Union[
+        ResponseResult,
+        AsyncResponseStreamResult,
+        AsyncResponseStreamWrapper[Any],
+    ]:
+        stream_context = responses_stream_context.get()
+        if stream_context is not None:
+            result = await wrapped(*args, **kwargs)
+            return _get_response_stream_result(result)
+
         params = extract_params(**kwargs)
-        invocation = handler.start_inference(
+        invocation = handler.inference(
             **get_inference_creation_kwargs(params, instance)
         )
         apply_request_attributes(invocation, params, capture_content)
@@ -69,23 +181,108 @@ async def traced_method(wrapped, instance, args, kwargs):
             result = await wrapped(*args, **kwargs)
             parsed_result = _get_response_stream_result(result)
 
-            if is_streaming(kwargs):
+            if (
+                _AsyncResponseStream is not None
+                and isinstance(parsed_result, _AsyncResponseStream)
+            ) or (
+                _OpenAIAsyncStream is not None
+                and isinstance(parsed_result, _OpenAIAsyncStream)
+            ):
                 return AsyncResponseStreamWrapper(
-                    parsed_result,
+                    cast("AsyncResponseStreamResult", parsed_result),
                     invocation,
                     capture_content,
                 )
 
             set_invocation_response_attributes(
-                invocation, parsed_result, capture_content
+                invocation,
+                cast("ResponseResult", parsed_result),
+                capture_content,
             )
             invocation.stop()
             return result
         except Exception as error:
             invocation.fail(error)
             raise
 
-    return traced_method
+    return cast(
+        'Callable[..., Awaitable[Union["ResponseResult", "AsyncResponseStreamResult", AsyncResponseStreamWrapper[Any]]]]',
+        traced_method,
+    )
+
+
+def responses_stream(
+    handler: TelemetryHandler,
+) -> Callable[..., ResponseStreamManagerWrapper[Any]]:
+    """Wrap ``Responses.stream`` to trace sync stream manager calls.
+
+    Traces :meth:`openai.resources.responses.responses.Responses.stream`.
+    OpenAI SDK source:
+    https://github.com/openai/openai-python/blob/main/src/openai/resources/responses/responses.py#L1062
+    """
+
+    capture_content = handler.should_capture_content()
+
+    def traced_method(
+        wrapped: Callable[..., "ResponseStreamManager[Any]"],
+        instance: "Responses",
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> ResponseStreamManagerWrapper[Any]:
+        def invocation_factory() -> Any:
+            params = extract_params(**kwargs)
+            invocation = handler.inference(
+                **get_inference_creation_kwargs(params, instance)
+            )
+            apply_request_attributes(invocation, params, capture_content)
+            return invocation
+
+        return ResponseStreamManagerWrapper(
+            wrapped(*args, **kwargs),
+            invocation_factory,
+            capture_content,
+        )
+
+    return cast(
+        "Callable[..., ResponseStreamManagerWrapper[Any]]", traced_method
+    )
+
+
+def async_responses_stream(
+    handler: TelemetryHandler,
+) -> Callable[..., AsyncResponseStreamManagerWrapper[Any]]:
+    """Wrap ``AsyncResponses.stream`` to trace async stream manager calls.
+
+    Traces :meth:`openai.resources.responses.responses.AsyncResponses.stream`.
+    OpenAI SDK source:
+    https://github.com/openai/openai-python/blob/main/src/openai/resources/responses/responses.py#L2809
+    """
+
+    capture_content = handler.should_capture_content()
+
+    def traced_method(
+        wrapped: Callable[..., "AsyncResponseStreamManager[Any]"],
+        instance: "AsyncResponses",
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> AsyncResponseStreamManagerWrapper[Any]:
+        def invocation_factory() -> Any:
+            params = extract_params(**kwargs)
+            invocation = handler.inference(
+                **get_inference_creation_kwargs(params, instance)
+            )
+            apply_request_attributes(invocation, params, capture_content)
+            return invocation
+
+        return AsyncResponseStreamManagerWrapper(
+            wrapped(*args, **kwargs),
+            invocation_factory,
+            capture_content,
+        )
+
+    return cast(
+        "Callable[..., AsyncResponseStreamManagerWrapper[Any]]", traced_method
+    )
 
 
 def _get_response_stream_result(result):