wip: keeping anthropic wrapper same as openai wrapper.

Author: eternalcuriouslearner

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

@@ -58,7 +58,7 @@ def messages_create(
     Union[
         "AnthropicMessage",
         "AnthropicStream[RawMessageStreamEvent]",
-        MessagesStreamWrapper[RawMessageStreamEvent],
+        MessagesStreamWrapper[None],
     ],
 ]:
     """Wrap the `create` method of the `Messages` class to trace it."""
@@ -78,7 +78,7 @@ def traced_method(
     ) -> Union[
         "AnthropicMessage",
         "AnthropicStream[RawMessageStreamEvent]",
-        MessagesStreamWrapper[RawMessageStreamEvent],
+        MessagesStreamWrapper[None],
     ]:
         params = extract_params(*args, **kwargs)
         attributes = get_llm_request_attributes(params, instance)
@@ -123,6 +123,6 @@ def traced_method(
             raise
 
     return cast(
-        'Callable[..., Union["AnthropicMessage", "AnthropicStream[RawMessageStreamEvent]", MessagesStreamWrapper[RawMessageStreamEvent]]]',
+        'Callable[..., Union["AnthropicMessage", "AnthropicStream[RawMessageStreamEvent]", MessagesStreamWrapper[None]]]',
         traced_method,
     )

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

@@ -19,12 +19,11 @@
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
-    AsyncIterator,
+    Any,
     Callable,
     Generator,
     Generic,
     Iterator,
-    Protocol,
     TypeVar,
     cast,
 )
@@ -45,64 +44,29 @@
     _sdk_accumulate_event = None
 
 if TYPE_CHECKING:
+    from anthropic._streaming import AsyncStream, Stream
     from anthropic.lib.streaming._messages import (  # pylint: disable=no-name-in-module
+        AsyncMessageStream,
         AsyncMessageStreamManager,
+        MessageStream,
         MessageStreamManager,
     )
     from anthropic.lib.streaming._types import (  # pylint: disable=no-name-in-module
-        MessageStreamEvent,
+        ParsedMessageStreamEvent,
     )
     from anthropic.types import (
         Message,
         RawMessageStreamEvent,
     )
+    from anthropic.types.parsed_message import ParsedMessage
 
 
 _logger = logging.getLogger(__name__)
-SyncResponseT = TypeVar("SyncResponseT", bound="_SupportsClose")
-AsyncResponseT = TypeVar("AsyncResponseT", bound="_SupportsAclose")
-StreamEventT = TypeVar(
-    "StreamEventT", "RawMessageStreamEvent", "MessageStreamEvent"
-)
-StreamEventT_co = TypeVar(
-    "StreamEventT_co",
-    "RawMessageStreamEvent",
-    "MessageStreamEvent",
-    covariant=True,
-)
+ResponseT = TypeVar("ResponseT")
+ResponseFormatT = TypeVar("ResponseFormatT")
 accumulate_event = cast("Callable[..., Message] | None", _sdk_accumulate_event)
 
 
-class _SupportsClose(Protocol):
-    def close(self) -> None: ...
-
-
-class _SupportsAclose(_SupportsClose, Protocol):
-    async def aclose(self) -> None: ...
-
-
-class _SyncStream(Protocol[StreamEventT_co]):
-    @property
-    def response(self) -> _SupportsClose: ...
-
-    def __iter__(self) -> Iterator[StreamEventT_co]: ...
-
-    def __next__(self) -> StreamEventT_co: ...
-
-    def close(self) -> None: ...
-
-
-class _AsyncStream(Protocol[StreamEventT_co]):
-    @property
-    def response(self) -> _SupportsAclose: ...
-
-    def __aiter__(self) -> AsyncIterator[StreamEventT_co]: ...
-
-    async def __anext__(self) -> StreamEventT_co: ...
-
-    async def close(self) -> None: ...
-
-
 def _set_response_attributes(
     invocation: LLMInvocation,
     result: "Message | None",
@@ -111,9 +75,9 @@ def _set_response_attributes(
     set_invocation_response_attributes(invocation, result, capture_content)
 
 
-class _ResponseProxy(Generic[SyncResponseT]):
-    def __init__(self, response: SyncResponseT, finalize: Callable[[], None]):
-        self._response = response
+class _ResponseProxy(Generic[ResponseT]):
+    def __init__(self, response: ResponseT, finalize: Callable[[], None]):
+        self._response: Any = response
         self._finalize = finalize
 
     def close(self) -> None:
@@ -126,17 +90,11 @@ def __getattr__(self, name: str):
         return getattr(self._response, name)
 
 
-class _AsyncResponseProxy(Generic[AsyncResponseT]):
-    def __init__(self, response: AsyncResponseT, finalize: Callable[[], None]):
-        self._response = response
+class _AsyncResponseProxy(Generic[ResponseT]):
+    def __init__(self, response: ResponseT, finalize: Callable[[], None]):
+        self._response: Any = response
         self._finalize = finalize
 
-    def close(self) -> None:
-        try:
-            self._response.close()
-        finally:
-            self._finalize()
-
     async def aclose(self) -> None:
         try:
             await self._response.aclose()
@@ -166,24 +124,29 @@ def message(self) -> Message:
         return self._message
 
 
-class MessagesStreamWrapper(Generic[StreamEventT], Iterator[StreamEventT]):
+class MessagesStreamWrapper(
+    Generic[ResponseFormatT],
+    Iterator[
+        "RawMessageStreamEvent | ParsedMessageStreamEvent[ResponseFormatT]"
+    ],
+):
     """Wrapper for Anthropic Stream that handles telemetry."""
 
     def __init__(
         self,
-        stream: _SyncStream[StreamEventT],
+        stream: "Stream[RawMessageStreamEvent] | MessageStream[ResponseFormatT]",
         handler: TelemetryHandler,
         invocation: LLMInvocation,
         capture_content: bool,
     ):
         self.stream = stream
         self.handler = handler
         self.invocation = invocation
-        self._message: "Message | None" = None
+        self._message: "Message | ParsedMessage[ResponseFormatT] | None" = None
         self._capture_content = capture_content
         self._finalized = False
 
-    def __enter__(self) -> "MessagesStreamWrapper[StreamEventT]":
+    def __enter__(self) -> "MessagesStreamWrapper[ResponseFormatT]":
         return self
 
     def __exit__(
@@ -207,10 +170,12 @@ def close(self) -> None:
         finally:
             self._stop()
 
-    def __iter__(self) -> "MessagesStreamWrapper[StreamEventT]":
+    def __iter__(self) -> "MessagesStreamWrapper[ResponseFormatT]":
         return self
 
-    def __next__(self) -> StreamEventT:
+    def __next__(
+        self,
+    ) -> "RawMessageStreamEvent | ParsedMessageStreamEvent[ResponseFormatT]":
         try:
             chunk = next(self.stream)
         except StopIteration:
@@ -227,9 +192,7 @@ def __getattr__(self, name: str) -> object:
         return getattr(self.stream, name)
 
     @property
-    def response(
-        self,
-    ) -> "_ResponseProxy[_SupportsClose] | _AsyncResponseProxy[_SupportsAclose] | None":
+    def response(self):
         return _ResponseProxy(self.stream.response, self._stop)
 
     def _stop(self) -> None:
@@ -266,10 +229,13 @@ def _safe_instrumentation(
                 exc_info=True,
             )
 
-    def _process_chunk(self, chunk: StreamEventT) -> None:
+    def _process_chunk(
+        self,
+        chunk: "RawMessageStreamEvent | ParsedMessageStreamEvent[ResponseFormatT]",
+    ) -> None:
         """Accumulate a final message snapshot from a streaming chunk."""
         snapshot = cast(
-            "Message | None",
+            "ParsedMessage[ResponseFormatT] | None",
             getattr(self.stream, "current_message_snapshot", None),
         )
         if snapshot is not None:
@@ -279,30 +245,32 @@ def _process_chunk(self, chunk: StreamEventT) -> None:
             return
         self._message = accumulate_event(
             event=cast("RawMessageStreamEvent", chunk),
-            current_snapshot=self._message,
+            current_snapshot=cast(
+                "ParsedMessage[ResponseFormatT] | None", self._message
+            ),
         )
 
 
-class AsyncMessagesStreamWrapper(MessagesStreamWrapper[StreamEventT]):
+class AsyncMessagesStreamWrapper(MessagesStreamWrapper[ResponseFormatT]):
     """Wrapper for async Anthropic Stream that handles telemetry."""
 
-    stream: _AsyncStream[StreamEventT]
-
     def __init__(
         self,
-        stream: _AsyncStream[StreamEventT],
+        stream: "AsyncStream[RawMessageStreamEvent] | AsyncMessageStream[ResponseFormatT]",
         handler: TelemetryHandler,
         invocation: LLMInvocation,
         capture_content: bool,
     ):
         self.stream = stream
         self.handler = handler
         self.invocation = invocation
-        self._message: "Message | None" = None
+        self._message: "Message | ParsedMessage[ResponseFormatT] | None" = None
         self._capture_content = capture_content
         self._finalized = False
 
-    async def __aenter__(self) -> "AsyncMessagesStreamWrapper[StreamEventT]":
+    async def __aenter__(
+        self,
+    ) -> "AsyncMessagesStreamWrapper[ResponseFormatT]":
         return self
 
     async def __aexit__(
@@ -326,16 +294,16 @@ async def close(self) -> None:  # type: ignore[override]
         finally:
             self._stop()
 
-    def __aiter__(self) -> "AsyncMessagesStreamWrapper[StreamEventT]":
+    def __aiter__(self) -> "AsyncMessagesStreamWrapper[ResponseFormatT]":
         return self
 
     @property
-    def response(
-        self,
-    ) -> "_ResponseProxy[_SupportsClose] | _AsyncResponseProxy[_SupportsAclose] | None":
+    def response(self) -> Any:
         return _AsyncResponseProxy(self.stream.response, self._stop)
 
-    async def __anext__(self) -> StreamEventT:
+    async def __anext__(
+        self,
+    ) -> "RawMessageStreamEvent | ParsedMessageStreamEvent[ResponseFormatT]":
         try:
             chunk = await self.stream.__anext__()
         except StopAsyncIteration:
@@ -349,12 +317,12 @@ async def __anext__(self) -> StreamEventT:
         return chunk
 
 
-class MessagesStreamManagerWrapper:
+class MessagesStreamManagerWrapper(Generic[ResponseFormatT]):
     """Wrapper for sync Anthropic stream managers."""
 
     def __init__(
         self,
-        manager: "MessageStreamManager",
+        manager: "MessageStreamManager[ResponseFormatT]",
         handler: TelemetryHandler,
         invocation: LLMInvocation,
         capture_content: bool,
@@ -363,15 +331,12 @@ def __init__(
         self._handler = handler
         self._invocation = invocation
         self._capture_content = capture_content
-        self._stream_wrapper: (
-            MessagesStreamWrapper[MessageStreamEvent] | None
-        ) = None
-
-    def __enter__(self) -> MessagesStreamWrapper[MessageStreamEvent]:
-        stream = cast(
-            "_SyncStream[MessageStreamEvent]",
-            self._manager.__enter__(),
+        self._stream_wrapper: MessagesStreamWrapper[ResponseFormatT] | None = (
+            None
         )
+
+    def __enter__(self) -> MessagesStreamWrapper[ResponseFormatT]:
+        stream = self._manager.__enter__()
         self._stream_wrapper = MessagesStreamWrapper(
             stream,
             self._handler,
@@ -406,7 +371,7 @@ def __getattr__(self, name: str) -> object:
         return getattr(self._manager, name)
 
 
-class AsyncMessagesStreamManagerWrapper:
+class AsyncMessagesStreamManagerWrapper(Generic[ResponseFormatT]):
     """Wrapper for AsyncMessageStreamManager that handles telemetry.
 
     Wraps AsyncMessageStreamManager from the Anthropic SDK:
@@ -415,7 +380,7 @@ class AsyncMessagesStreamManagerWrapper:
 
     def __init__(
         self,
-        manager: "AsyncMessageStreamManager",
+        manager: "AsyncMessageStreamManager[ResponseFormatT]",
         handler: TelemetryHandler,
         invocation: LLMInvocation,
         capture_content: bool,
@@ -425,16 +390,13 @@ def __init__(
         self._invocation = invocation
         self._capture_content = capture_content
         self._stream_wrapper: (
-            AsyncMessagesStreamWrapper[MessageStreamEvent] | None
+            AsyncMessagesStreamWrapper[ResponseFormatT] | None
         ) = None
 
     async def __aenter__(
         self,
-    ) -> AsyncMessagesStreamWrapper[MessageStreamEvent]:
-        msg_stream = cast(
-            "_AsyncStream[MessageStreamEvent]",
-            await self._manager.__aenter__(),
-        )
+    ) -> AsyncMessagesStreamWrapper[ResponseFormatT]:
+        msg_stream = await self._manager.__aenter__()
         self._stream_wrapper = AsyncMessagesStreamWrapper(
             msg_stream,
             self._handler,