Add response wrappers for OpenAI Responses API streams.

instrumentation-genai/opentelemetry-instrumentation-openai-v2/CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `OTEL_SEMCONV_STABILITY_OPT_IN` to `gen_ai_latest_experimental` to enable.
   Add dependency on `opentelemetry-util-genai` pypi package.
   ([#3715](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3715))
+- Add wrappers for OpenAI Responses API streams and response stream managers
+  ([#4280](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4280))
 
 ## Version 2.3b0 (2025-12-24)
 

instrumentation-genai/opentelemetry-instrumentation-openai-v2/src/opentelemetry/instrumentation/openai_v2/response_wrappers.py

@@ -0,0 +1,314 @@
+"""Wrappers for OpenAI Responses API streams and stream managers."""
+
+from __future__ import annotations
+
+import logging
+from types import TracebackType
+from typing import TYPE_CHECKING, Callable, Generic, TypeVar
+
+try:
+    from openai.lib.streaming.responses._events import (  # pylint: disable=no-name-in-module
+        ResponseCompletedEvent,
+    )
+    from openai.types.responses import (  # pylint: disable=no-name-in-module
+        ResponseCreatedEvent,
+        ResponseErrorEvent,
+        ResponseFailedEvent,
+        ResponseIncompleteEvent,
+        ResponseInProgressEvent,
+    )
+except ImportError:  # pragma: no cover
+    ResponseCompletedEvent = None
+    ResponseCreatedEvent = None
+    ResponseErrorEvent = None
+    ResponseFailedEvent = None
+    ResponseIncompleteEvent = None
+    ResponseInProgressEvent = None
+
+if (
+    ResponseCreatedEvent is not None
+    and ResponseInProgressEvent is not None
+    and ResponseFailedEvent is not None
+    and ResponseIncompleteEvent is not None
+    and ResponseCompletedEvent is not None
+):
+    _RESPONSE_EVENTS_WITH_RESPONSE = (
+        ResponseCreatedEvent,
+        ResponseInProgressEvent,
+        ResponseFailedEvent,
+        ResponseIncompleteEvent,
+        ResponseCompletedEvent,
+    )
+else:
+    _RESPONSE_EVENTS_WITH_RESPONSE = ()
+
+try:
+    from opentelemetry.instrumentation.openai_v2.response_extractors import (  # pylint: disable=no-name-in-module
+        _set_invocation_response_attributes,
+    )
+except ImportError:  # pragma: no cover
+    _set_invocation_response_attributes = None
+
+try:
+    from opentelemetry.util.genai.types import (  # pylint: disable=no-name-in-module
+        Error,
+    )
+except ImportError:  # pragma: no cover
+    Error = None
+
+if TYPE_CHECKING:
+    from openai.lib.streaming.responses._events import (  # pylint: disable=no-name-in-module
+        ResponseStreamEvent,
+    )
+    from openai.lib.streaming.responses._responses import (
+        ResponseStream,
+        ResponseStreamManager,
+    )  # pylint: disable=no-name-in-module
+    from openai.types.responses import (  # pylint: disable=no-name-in-module
+        ParsedResponse,
+        Response,
+    )
+
+    from opentelemetry.util.genai.handler import TelemetryHandler
+    from opentelemetry.util.genai.types import (  # pylint: disable=no-name-in-module
+        LLMInvocation,
+    )
+
+
+_logger = logging.getLogger(__name__)
+TextFormatT = TypeVar("TextFormatT")
+ResponseT = TypeVar("ResponseT")
+
+
+def _set_response_attributes(
+    invocation: "LLMInvocation",
+    result: "ParsedResponse[TextFormatT] | Response | None",
+    capture_content: bool,
+) -> None:
+    if _set_invocation_response_attributes is None:
+        return
+    _set_invocation_response_attributes(invocation, result, capture_content)
+
+
+class _ResponseProxy(Generic[ResponseT]):
+    def __init__(self, response: ResponseT, finalize: Callable[[], None]):
+        self._response = response
+        self._finalize = finalize
+
+    def close(self) -> None:
+        try:
+            self._response.close()
+        finally:
+            self._finalize()
+
+    def __getattr__(self, name: str):
+        return getattr(self._response, name)
+
+
+class ResponseStreamWrapper(Generic[TextFormatT]):
+    """Wrapper for OpenAI Responses API stream objects.
+
+    Wraps ResponseStream from the OpenAI SDK:
+    https://github.com/openai/openai-python/blob/656e3cab4a18262a49b961d41293367e45ee71b9/src/openai/_streaming.py#L55
+    """
+
+    def __init__(
+        self,
+        stream: "ResponseStream[TextFormatT]",
+        handler: "TelemetryHandler",
+        invocation: "LLMInvocation",
+        capture_content: bool,
+    ):
+        self.stream = stream
+        self.handler = handler
+        self.invocation = invocation
+        self._capture_content = capture_content
+        self._finalized = False
+
+    def __enter__(self) -> "ResponseStreamWrapper":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        try:
+            if exc_type is not None:
+                self._fail(
+                    str(exc_val), type(exc_val) if exc_val else Exception
+                )
+        finally:
+            self.close()
+        return False
+
+    def close(self) -> None:
+        try:
+            self.stream.close()
+        finally:
+            self._stop(None)
+
+    def __iter__(self) -> "ResponseStreamWrapper":
+        return self
+
+    def __next__(self) -> "ResponseStreamEvent[TextFormatT]":
+        try:
+            event = next(self.stream)
+        except StopIteration:
+            self._stop(None)
+            raise
+        except Exception as error:
+            self._fail(str(error), type(error))
+            raise
+        self._safe_instrumentation(
+            lambda: self.process_event(event),
+            "event processing",
+        )
+        return event
+
+    def get_final_response(self) -> "ParsedResponse[TextFormatT]":
+        self.until_done()
+        return self.stream.get_final_response()
+
+    def until_done(self) -> "ResponseStreamWrapper":
+        for _ in self:
+            pass
+        return self
+
+    def parse(self) -> "ResponseStreamWrapper":
+        return self
+
+    def __getattr__(self, name: str):
+        return getattr(self.stream, name)
+
+    @property
+    def response(self):
+        response = self.stream.response
+        if response is None:
+            return None
+        return _ResponseProxy(response, lambda: self._stop(None))
+
+    def _stop(
+        self, result: "ParsedResponse[TextFormatT] | Response | None"
+    ) -> None:
+        if self._finalized:
+            return
+        self._safe_instrumentation(
+            lambda: _set_response_attributes(
+                self.invocation, result, self._capture_content
+            ),
+            "response attribute extraction",
+        )
+        self._safe_instrumentation(
+            lambda: self.handler.stop_llm(self.invocation),
+            "stop_llm",
+        )
+        self._finalized = True
+
+    def _fail(self, message: str, error_type: type[BaseException]) -> None:
+        if self._finalized:
+            return
+        if Error is None:
+            return
+        self._safe_instrumentation(
+            lambda: self.handler.fail_llm(
+                self.invocation, Error(message=message, type=error_type)
+            ),
+            "fail_llm",
+        )
+        self._finalized = True
+
+    @staticmethod
+    def _safe_instrumentation(
+        callback: Callable[[], None], context: str
+    ) -> None:
+        try:
+            callback()
+        except Exception:  # pylint: disable=broad-exception-caught
+            _logger.debug(
+                "OpenAI responses instrumentation error during %s",
+                context,
+                exc_info=True,
+            )
+
+    def process_event(self, event: "ResponseStreamEvent[TextFormatT]") -> None:
+        event_type = event.type
+        response: "ParsedResponse[TextFormatT] | Response | None" = None
+
+        if isinstance(event, _RESPONSE_EVENTS_WITH_RESPONSE):
+            response = event.response
+
+        if response and not self.invocation.request_model:
+            model = response.model
+            if model:
+                self.invocation.request_model = model
+
+        if ResponseCompletedEvent is not None and isinstance(
+            event, ResponseCompletedEvent
+        ):
+            self._stop(response)
+            return
+
+        if (
+            ResponseFailedEvent is not None
+            and ResponseIncompleteEvent is not None
+            and isinstance(
+                event, (ResponseFailedEvent, ResponseIncompleteEvent)
+            )
+        ):
+            self._safe_instrumentation(
+                lambda: _set_response_attributes(
+                    self.invocation, response, self._capture_content
+                ),
+                "response attribute extraction",
+            )
+            self._fail(event_type, RuntimeError)
+            return
+
+        if ResponseErrorEvent is not None and isinstance(
+            event, ResponseErrorEvent
+        ):
+            error_type = event.code or "response.error"
+            message = event.message or error_type
+            self._fail(message, RuntimeError)
+
+
+class ResponseStreamManagerWrapper(Generic[TextFormatT]):
+    """Wrapper for OpenAI Responses API stream managers.
+
+    Wraps ResponseStreamManager from the OpenAI SDK:
+    https://github.com/openai/openai-python/blob/656e3cab4a18262a49b961d41293367e45ee71b9/src/openai/lib/streaming/responses/_responses.py#L95
+    """
+
+    def __init__(
+        self,
+        manager: "ResponseStreamManager[TextFormatT]",
+        handler: "TelemetryHandler",
+        invocation: "LLMInvocation",
+        capture_content: bool,
+    ):
+        self._manager = manager
+        self._handler = handler
+        self._invocation = invocation
+        self._capture_content = capture_content
+
+    def __enter__(self) -> ResponseStreamWrapper[TextFormatT]:
+        stream = self._manager.__enter__()
+        return ResponseStreamWrapper(
+            stream,
+            self._handler,
+            self._invocation,
+            self._capture_content,
+        )
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        return self._manager.__exit__(exc_type, exc_val, exc_tb)
+
+    def __getattr__(self, name: str):
+        return getattr(self._manager, name)