refactor: responses and messages genai streams integration (#92)

* wip: responses and messages genai streams integration

* polish: added changelog, fixed precommit and typechecks.

* Address stream wrapper review feedback

Assisted-by: ChatGPT 5.5

* wip: Code design refactor. Assisted by: GPT 5.5.

---------

Co-authored-by: Leighton Chen <lechen@microsoft.com>

Author: eternalcuriouslearner

instrumentation/opentelemetry-instrumentation-genai-anthropic/.changelog/92.changed

@@ -0,0 +1 @@
+Use shared GenAI stream wrappers for Messages API streams.

instrumentation/opentelemetry-instrumentation-genai-anthropic/src/opentelemetry/instrumentation/genai/anthropic/wrappers.py

@@ -9,11 +9,16 @@
     Any,
     Callable,
     Generic,
-    Iterator,
+    Protocol,
     TypeVar,
     cast,
 )
 
+from opentelemetry.util.genai.stream import (
+    AsyncStreamWrapper,
+    SyncStreamWrapper,
+)
+
 from .messages_extractors import set_invocation_response_attributes
 
 try:
@@ -48,6 +53,11 @@
 accumulate_event = cast("Callable[..., Message] | None", _sdk_accumulate_event)
 
 
+class _StreamWrapperWithStream(Protocol):
+    @property
+    def stream(self) -> object: ...
+
+
 def _set_response_attributes(
     invocation: InferenceInvocation,
     result: Message | None,
@@ -105,174 +115,144 @@ def message(self) -> Message:
         return self._message
 
 
-class MessagesStreamWrapper(
-    Generic[ResponseFormatT],
-    Iterator[
-        "RawMessageStreamEvent | ParsedMessageStreamEvent[ResponseFormatT]"
-    ],
-):
-    """Wrapper for Anthropic Stream that handles telemetry."""
-
-    def __init__(
-        self,
-        stream: Stream[RawMessageStreamEvent] | MessageStream[ResponseFormatT],
-        invocation: InferenceInvocation,
-        capture_content: bool,
-    ):
-        self.stream = stream
-        self.invocation = invocation
-        self._message: Message | ParsedMessage[ResponseFormatT] | None = None
-        self._capture_content = capture_content
-        self._finalized = False
-
-    def __enter__(self) -> MessagesStreamWrapper[ResponseFormatT]:
-        return self
-
-    def __exit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: TracebackType | None,
-    ) -> bool:
-        try:
-            if exc_val is not None:
-                self._fail(exc_val)
-        finally:
-            self.close()
-        return False
-
-    def close(self) -> None:
-        try:
-            self.stream.close()
-        except Exception as exc:
-            self._fail(exc)
-            raise
-        self._stop()
-
-    def __iter__(self) -> MessagesStreamWrapper[ResponseFormatT]:
-        return self
-
-    def __next__(
-        self,
-    ) -> RawMessageStreamEvent | ParsedMessageStreamEvent[ResponseFormatT]:
-        try:
-            chunk = next(self.stream)
-        except StopIteration:
-            self._stop()
-            raise
-        except Exception as exc:
-            self._fail(exc)
-            raise
-        self._process_chunk(chunk)
-        return chunk
-
-    def __getattr__(self, name: str) -> object:
-        return getattr(self.stream, name)
-
-    @property
-    def response(self):
-        return _ResponseProxy(self.stream.response, self._stop)
+class _MessagesStreamMixin(Generic[ResponseFormatT]):
+    _self_invocation: InferenceInvocation
+    _self_message: Message | ParsedMessage[ResponseFormatT] | None
+    _self_capture_content: bool
+    _self_message_telemetry_finalized: bool
 
     def _stop(self) -> None:
-        if self._finalized:
+        if self._self_message_telemetry_finalized:
             return
         _set_response_attributes(
-            self.invocation, self._message, self._capture_content
+            self._self_invocation,
+            self._self_message,
+            self._self_capture_content,
         )
-        self.invocation.stop()
-        self._finalized = True
+        self._self_invocation.stop()
+        self._self_message_telemetry_finalized = True
 
     def _fail(self, exc: BaseException) -> None:
-        if self._finalized:
+        if self._self_message_telemetry_finalized:
             return
-        self.invocation.fail(exc)
-        self._finalized = True
+        self._self_invocation.fail(exc)
+        self._self_message_telemetry_finalized = True
+
+    def _on_stream_end(self) -> None:
+        self._stop()
+
+    def _on_stream_error(self, error: BaseException) -> None:
+        self._fail(error)
 
     def _process_chunk(
         self,
         chunk: RawMessageStreamEvent
         | ParsedMessageStreamEvent[ResponseFormatT],
     ) -> None:
         """Accumulate a final message snapshot from a streaming chunk."""
+        stream = cast(_StreamWrapperWithStream, self).stream
         snapshot = cast(
             "ParsedMessage[ResponseFormatT] | None",
-            getattr(self.stream, "current_message_snapshot", None),
+            getattr(stream, "current_message_snapshot", None),
         )
         if snapshot is not None:
-            self._message = snapshot
+            self._self_message = snapshot
             return
         if accumulate_event is None:
             return
-        self._message = accumulate_event(
+        self._self_message = accumulate_event(
             event=cast("RawMessageStreamEvent", chunk),
             current_snapshot=cast(
-                "ParsedMessage[ResponseFormatT] | None", self._message
+                "ParsedMessage[ResponseFormatT] | None", self._self_message
             ),
         )
 
 
-class AsyncMessagesStreamWrapper(MessagesStreamWrapper[ResponseFormatT]):
-    """Wrapper for async Anthropic Stream that handles telemetry."""
+class MessagesStreamWrapper(
+    _MessagesStreamMixin[ResponseFormatT],
+    SyncStreamWrapper[
+        "RawMessageStreamEvent | ParsedMessageStreamEvent[ResponseFormatT]"
+    ],
+    Generic[ResponseFormatT],
+):
+    """Wrapper for Anthropic Stream that handles telemetry."""
 
     def __init__(
         self,
-        stream: AsyncStream[RawMessageStreamEvent]
-        | AsyncMessageStream[ResponseFormatT],
+        stream: Stream[RawMessageStreamEvent] | MessageStream[ResponseFormatT],
         invocation: InferenceInvocation,
         capture_content: bool,
     ):
-        self.stream = stream
-        self.invocation = invocation
-        self._message: Message | ParsedMessage[ResponseFormatT] | None = None
-        self._capture_content = capture_content
-        self._finalized = False
+        super().__init__(stream)
+        self._self_invocation = invocation
+        self._self_message = None
+        self._self_capture_content = capture_content
+        self._self_message_telemetry_finalized = False
 
-    async def __aenter__(
+    @property
+    def response(self) -> _ResponseProxy[object]:
+        return _ResponseProxy(self.stream.response, self._stop)
+
+    @property
+    def stream(
         self,
-    ) -> AsyncMessagesStreamWrapper[ResponseFormatT]:
-        return self
+    ) -> Stream[RawMessageStreamEvent] | MessageStream[ResponseFormatT]:
+        return self._self_stream
 
-    async def __aexit__(
+    @stream.setter
+    def stream(
         self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: TracebackType | None,
-    ) -> bool:
-        try:
-            if exc_val is not None:
-                self._fail(exc_val)
-        finally:
-            await self.close()
-        return False
+        stream: Stream[RawMessageStreamEvent] | MessageStream[ResponseFormatT],
+    ) -> None:
+        self.__wrapped__ = stream
+        self._self_stream = stream
+        self._self_iterator = iter(stream)
 
-    async def close(self) -> None:  # type: ignore[override]
-        try:
-            await self.stream.close()
-        except Exception as exc:
-            self._fail(exc)
-            raise
-        self._stop()
 
-    def __aiter__(self) -> AsyncMessagesStreamWrapper[ResponseFormatT]:
-        return self
+class AsyncMessagesStreamWrapper(
+    _MessagesStreamMixin[ResponseFormatT],
+    AsyncStreamWrapper[
+        "RawMessageStreamEvent | ParsedMessageStreamEvent[ResponseFormatT]"
+    ],
+    Generic[ResponseFormatT],
+):
+    """Wrapper for async Anthropic Stream that handles telemetry."""
+
+    def __init__(
+        self,
+        stream: AsyncStream[RawMessageStreamEvent]
+        | AsyncMessageStream[ResponseFormatT],
+        invocation: InferenceInvocation,
+        capture_content: bool,
+    ):
+        super().__init__(stream)
+        self._self_invocation = invocation
+        self._self_message = None
+        self._self_capture_content = capture_content
+        self._self_message_telemetry_finalized = False
 
     @property
     def response(self) -> Any:
         return _AsyncResponseProxy(self.stream.response, self._stop)
 
-    async def __anext__(
+    @property
+    def stream(
         self,
-    ) -> RawMessageStreamEvent | ParsedMessageStreamEvent[ResponseFormatT]:
-        try:
-            chunk = await self.stream.__anext__()
-        except StopAsyncIteration:
-            self._stop()
-            raise
-        except Exception as exc:
-            self._fail(exc)
-            raise
-        self._process_chunk(chunk)
-        return chunk
+    ) -> (
+        AsyncStream[RawMessageStreamEvent]
+        | AsyncMessageStream[ResponseFormatT]
+    ):
+        return self._self_stream
+
+    @stream.setter
+    def stream(
+        self,
+        stream: AsyncStream[RawMessageStreamEvent]
+        | AsyncMessageStream[ResponseFormatT],
+    ) -> None:
+        self.__wrapped__ = stream
+        self._self_stream = stream
+        self._self_aiter = aiter(stream)
 
 
 class MessagesStreamManagerWrapper(Generic[ResponseFormatT]):