polish: pr feedback.

Author: eternalcuriouslearner

AGENTS.md

@@ -76,6 +76,11 @@ Apply to packages under `instrumentation/` and `instrumentation-genai/`.
 - When catching exceptions from the underlying library to record telemetry, always re-raise the
   original exception unmodified.
 - Do not raise new exceptions in instrumentation/telemetry code.
+- For GenAI streaming wrappers, prefer the shared `SyncStreamWrapper` and `AsyncStreamWrapper`
+  helpers from `opentelemetry.util.genai._stream` instead of reimplementing iteration,
+  close/context-manager, and finalization behavior in provider packages.
+- Put provider-specific chunk parsing and telemetry finalization in private hook methods or a
+  narrow mixin. Do not make async stream wrappers inherit from sync stream wrappers.
 
 ### Semantic conventions
 

CLAUDE.md

@@ -1 +1,8 @@
 @AGENTS.md
+
+For GenAI streaming wrappers, use `SyncStreamWrapper` and
+`AsyncStreamWrapper` from `opentelemetry.util.genai._stream` instead of
+reimplementing iteration, close/context-manager, and finalization behavior in
+provider packages. Keep provider-specific chunk parsing and telemetry
+finalization in private hooks or a narrow mixin, and do not make async stream
+wrappers inherit from sync stream wrappers.

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

@@ -17,7 +17,7 @@
 import json
 from typing import Any, Optional
 
-from openai import Stream
+from openai import AsyncStream, Stream
 
 from opentelemetry.semconv._incubating.attributes import (
     openai_attributes as OpenAIAttributes,
@@ -36,28 +36,19 @@
 from .chat_buffers import ChoiceBuffer
 
 
-class ChatStreamWrapper(SyncStreamWrapper[Any]):
+class _ChatStreamMixin:
+    """Chat-specific hooks shared by sync and async stream wrappers."""
+
     invocation: InferenceInvocation
+    capture_content: bool
+    choice_buffers: list
     response_id: Optional[str] = None
     response_model: Optional[str] = None
     service_tier: Optional[str] = None
     finish_reasons: list = []
     prompt_tokens: Optional[int] = None
     completion_tokens: Optional[int] = None
 
-    def __init__(
-        self,
-        stream: Stream,
-        invocation: InferenceInvocation,
-        capture_content: bool,
-    ):
-        SyncStreamWrapper.__init__(self, stream)
-        self.stream = stream
-        self.invocation = invocation
-        self.choice_buffers = []
-        self._started = True
-        self.capture_content = capture_content
-
     def _set_response_model(self, chunk):
         if self.response_model:
             return
@@ -155,19 +146,16 @@ def _set_output_messages(self):
         self.invocation.output_messages = output_messages
 
     def _stop_stream(self) -> None:
-        self.cleanup()
+        self._cleanup()
 
     def _fail_stream(self, error: BaseException) -> None:
-        self.cleanup(error)
+        self._cleanup(error)
 
     def parse(self):
         """Called when using with_raw_response with stream=True."""
         return self
 
-    def cleanup(self, error: Optional[BaseException] = None):
-        if not self._started:
-            return
-
+    def _cleanup(self, error: Optional[BaseException] = None) -> None:
         self.invocation.response_model_name = self.response_model
         self.invocation.response_id = self.response_id
         self.invocation.input_tokens = self.prompt_tokens
@@ -188,31 +176,39 @@ def cleanup(self, error: Optional[BaseException] = None):
             self.invocation.fail(error)
         else:
             self.invocation.stop()
-        self._started = False
 
 
-class AsyncChatStreamWrapper(
-    AsyncStreamWrapper[Any],
-    ChatStreamWrapper,
+class ChatStreamWrapper(
+    _ChatStreamMixin,
+    SyncStreamWrapper[Any],
 ):
-    invocation: InferenceInvocation
-
     def __init__(
         self,
         stream: Stream,
         invocation: InferenceInvocation,
         capture_content: bool,
     ):
-        ChatStreamWrapper.__init__(self, stream, invocation, capture_content)
+        super().__init__(stream)
+        self.invocation = invocation
+        self.choice_buffers = []
+        self.capture_content = capture_content
 
-    def _process_chunk(self, chunk):
-        ChatStreamWrapper._process_chunk(self, chunk)
 
-    def _stop_stream(self) -> None:
-        ChatStreamWrapper._stop_stream(self)
+class AsyncChatStreamWrapper(
+    _ChatStreamMixin,
+    AsyncStreamWrapper[Any],
+):
 
-    def _fail_stream(self, error: BaseException) -> None:
-        ChatStreamWrapper._fail_stream(self, error)
+    def __init__(
+        self,
+        stream: AsyncStream,
+        invocation: InferenceInvocation,
+        capture_content: bool,
+    ):
+        super().__init__(stream)
+        self.invocation = invocation
+        self.choice_buffers = []
+        self.capture_content = capture_content
 
 
 __all__ = [

util/opentelemetry-util-genai/src/opentelemetry/util/genai/_stream.py

@@ -24,7 +24,18 @@
 
 
 class SyncStreamWrapper(ABC, Generic[ChunkT]):
-    """Base class for synchronous instrumented stream wrappers."""
+    """Base class for synchronous instrumented stream wrappers.
+
+    Subclass this when wrapping a provider SDK stream that is consumed with
+    normal iteration. The subclass should pass the SDK stream to
+    ``super().__init__(stream)`` and implement the three telemetry hooks:
+    ``_process_chunk`` for per-chunk state, ``_stop_stream`` for successful
+    finalization, and ``_fail_stream`` for failure finalization.
+
+    Users should consume subclasses as normal streams, for example with
+    ``for chunk in wrapper`` or ``with wrapper``. The hook methods are called
+    internally by the wrapper lifecycle and are not part of the public API.
+    """
 
     def __init__(self, stream: Any):
         self.stream = stream
@@ -106,7 +117,19 @@ def _handle_process_chunk_error(_error: Exception) -> None:
 
 
 class AsyncStreamWrapper(ABC, Generic[ChunkT]):
-    """Base class for asynchronous instrumented stream wrappers."""
+    """Base class for asynchronous instrumented stream wrappers.
+
+    Subclass this when wrapping a provider SDK stream that is consumed with
+    async iteration. The subclass should pass the SDK stream to
+    ``super().__init__(stream)`` and implement the three telemetry hooks:
+    ``_process_chunk`` for per-chunk state, ``_stop_stream`` for successful
+    finalization, and ``_fail_stream`` for failure finalization.
+
+    Users should consume subclasses as normal async streams, for example with
+    ``async for chunk in wrapper`` or ``async with wrapper``. The hook methods
+    remain synchronous telemetry hooks; async stream reads and close handling
+    are owned by this base class.
+    """
 
     def __init__(self, stream: Any):
         self.stream = stream