refactor(sdk): call span handle members directly in the python integrations

sdks/python-bedrock/src/telemetry_dev_bedrock/_streams.py

@@ -19,7 +19,6 @@
 
 # The core applies max_attribute_length after the mask.
 _CAPTURE_TRUNCATED_ATTRIBUTE = "telemetry.dev.capture.truncated"
-_OMIT = object()
 
 
 def stream_capture_budget() -> telemetry_dev.CaptureBudget:
@@ -106,14 +105,6 @@ def _join(self, write_at: int) -> None:
             raise
 
 
-def capture_masked(handle: Any) -> bool:
-    masked = getattr(handle, "capture_masked", None)
-    if isinstance(masked, bool):
-        return masked
-    client = getattr(handle, "_client", None)
-    return getattr(client, "mask", _OMIT) is not None if client is not None else True
-
-
 def mark_capture_incomplete(fields: dict[str, Any], mask_output: bool) -> dict[str, Any]:
     marked = merge_fields(fields, {"attributes": {_CAPTURE_TRUNCATED_ATTRIBUTE: True}})
     if mask_output:
@@ -580,7 +571,7 @@ def __init__(self, *, capture_trace: bool = False) -> None:
         # Metadata bypasses the capture flags and the mask, so trace content is filtered here,
         # and a mask that cannot see it withholds it.
         client = telemetry_dev.get_client()
-        masked = client is None or getattr(client, "mask", _OMIT) is not None
+        masked = client is None or client.mask is not None
         self.trace_capture = {
             "input": client is not None and client.capture_input and not masked,
             "output": client is not None and client.capture_output and not masked,
@@ -792,7 +783,7 @@ def __init__(
         state: StreamState,
         finish: Callable[[dict[str, Any]], None],
         started_at: float,
-        handle: Any,
+        handle: telemetry_dev.SpanHandle,
         report_error: Callable[[BaseException], None],
     ) -> None:
         self._stream = stream
@@ -803,7 +794,7 @@ def __init__(
         self._ended = False
         self._first = False
         self._track_output_chunks = isinstance(state, ConverseStreamState | InvokeModelStreamState)
-        self._mask_output_when_incomplete = capture_masked(handle)
+        self._mask_output_when_incomplete = handle.capture_masked
         self._report_error = report_error
         state.report_error = self._report_error
 
@@ -841,9 +832,7 @@ def __iter__(self) -> Iterator[Any]:
                             if state_has_output is not None
                             else _event_has_output(event)
                         ):
-                            record_output_chunk = getattr(self._handle, "record_output_chunk", None)
-                            if callable(record_output_chunk):
-                                record_output_chunk(received_at * 1000)
+                            self._handle.record_output_chunk(received_at * 1000)
                     except Exception as exc:
                         self._report_error(exc)
                 yield event

sdks/python-google-genai/src/telemetry_dev_google_genai/__init__.py

@@ -65,7 +65,6 @@ def _reset_install_lock_in_child() -> None:
 # so each is capped instead. Gemini's longest is 25 characters.
 _FINISH_REASON_MAX_LENGTH = 128
 _CAPTURE_TRUNCATED_ATTRIBUTE = "telemetry.dev.capture.truncated"
-_OMIT = object()
 
 _BUILTIN_TOOL_FIELDS = (
     ("google_search", "googleSearch"),
@@ -992,21 +991,8 @@ def end_stream(**fields: Any) -> None:
     return end_stream
 
 
-def _capture_masked(handle: telemetry_dev.SpanHandle) -> bool:
-    compatible = cast(Any, handle)
-    masked = getattr(compatible, "capture_masked", None)
-    if isinstance(masked, bool):
-        return masked
-    client = getattr(compatible, "_client", None)
-    return getattr(client, "mask", _OMIT) is not None if client is not None else True
-
-
 def _capture_metadata_content(handle: telemetry_dev.SpanHandle) -> bool:
-    compatible = cast(Any, handle)
-    capture_output = getattr(compatible, "capture_output", None)
-    if not isinstance(capture_output, bool):
-        capture_output = getattr(getattr(compatible, "_state", None), "capture_output", False)
-    return capture_output is True and not _capture_masked(handle)
+    return handle.capture_output and not handle.capture_masked
 
 
 def _mark_stream_capture(
@@ -1368,7 +1354,7 @@ def __init__(
             # An AFC turn that fails to map reports through the stream, once per stream.
             afc_usage_state["report_error"] = self._state.report_error
         self._completed = False
-        self._mask_output_when_incomplete = _capture_masked(handle)
+        self._mask_output_when_incomplete = handle.capture_masked
 
     def _fields(self) -> dict[str, Any]:
         return _mark_stream_capture(
@@ -1429,9 +1415,7 @@ def _advance(self, pull: Callable[[], Any]) -> Any:
                 _AFC_USAGE_STATE.reset(token)
         try:
             if _chunk_has_output(chunk):
-                record_output_chunk = getattr(self._handle, "record_output_chunk", None)
-                if callable(record_output_chunk):
-                    record_output_chunk(received_at * 1000)
+                self._handle.record_output_chunk(received_at * 1000)
         except Exception as exc:
             # Only the chunk's timing is lost; the chunk is still recorded.
             self._state.report_error(exc)
@@ -1468,7 +1452,7 @@ def __init__(
             # An AFC turn that fails to map reports through the stream, once per stream.
             afc_usage_state["report_error"] = self._state.report_error
         self._completed = False
-        self._mask_output_when_incomplete = _capture_masked(handle)
+        self._mask_output_when_incomplete = handle.capture_masked
 
     def _fields(self) -> dict[str, Any]:
         return _mark_stream_capture(
@@ -1532,9 +1516,7 @@ async def _advance(self, pull: Callable[[], Awaitable[Any]]) -> Any:
                 _AFC_USAGE_STATE.reset(token)
         try:
             if _chunk_has_output(chunk):
-                record_output_chunk = getattr(self._handle, "record_output_chunk", None)
-                if callable(record_output_chunk):
-                    record_output_chunk(received_at * 1000)
+                self._handle.record_output_chunk(received_at * 1000)
         except Exception as exc:
             # Only the chunk's timing is lost; the chunk is still recorded.
             self._state.report_error(exc)

sdks/python-litellm/src/telemetry_dev_litellm/__init__.py

@@ -47,7 +47,6 @@ def _reset_install_lock_in_child() -> None:
     _install_fork_hook = True
     os.register_at_fork(after_in_child=_reset_install_lock_in_child)
 
-_OMIT = object()
 # The core applies max_attribute_length after the mask. More choices, or more responses items
 # and rejected parts tracked for truncation, than this are not kept.
 _STREAM_CAPTURE_MAX_ITEMS = 1000
@@ -160,15 +159,6 @@ def _adopt_budget(
     return True
 
 
-def _capture_masked(handle: telemetry_dev.SpanHandle) -> bool:
-    compatible = cast(Any, handle)
-    masked = getattr(compatible, "capture_masked", None)
-    if isinstance(masked, bool):
-        return masked
-    client = getattr(compatible, "_client", None)
-    return getattr(client, "mask", _OMIT) is not None if client is not None else True
-
-
 def _mark_capture_incomplete(fields: dict[str, Any], mask_output: bool) -> None:
     attributes = fields.get("attributes")
     fields["attributes"] = {
@@ -908,7 +898,7 @@ def __init__(
         self._saw_first = False
         self._finished = False
         self._completed = False
-        self._mask_output_when_incomplete = _capture_masked(handle)
+        self._mask_output_when_incomplete = handle.capture_masked
         self._budget = _stream_capture_budget()
         self._response_id: str | None = None
         self._response_model: str | None = None
@@ -1063,9 +1053,7 @@ def _record(self, chunk: Any, received_at: float) -> None:
                 )
             if has_output:
                 self._saw_output = True
-                record_output_chunk = getattr(self._handle, "record_output_chunk", None)
-                if callable(record_output_chunk):
-                    record_output_chunk(received_at * 1000)
+                self._handle.record_output_chunk(received_at * 1000)
         except Exception as exc:
             self._report(exc)
         # The chunk is kept before its metadata is read, so a failure in the metadata (ID,
@@ -1288,7 +1276,7 @@ def __init__(
         self._retained_item_count = 0
         self._truncated_key_count = 0
         self._truncation_overflow = False
-        self._mask_output_when_incomplete = _capture_masked(handle)
+        self._mask_output_when_incomplete = handle.capture_masked
         self._budget = _stream_capture_budget()
         self._charges: dict[object, tuple[int, int]] = {}
         self._annotation_charge_keys: dict[object, dict[int, object]] = {}

sdks/python-openrouter/src/telemetry_dev_openrouter/__init__.py

@@ -30,7 +30,6 @@
 _TERMINAL_RESPONSE_EVENTS = frozenset(
     {"response.completed", "response.failed", "response.incomplete"}
 )
-_OMIT = object()
 _ORIGINALS: list[tuple[type[Any], str, Any]] = []
 _installed = False
 _install_lock = threading.Lock()
@@ -711,15 +710,6 @@ def end_stream(**fields: Any) -> None:
     return end_stream
 
 
-def _capture_masked(handle: telemetry_dev.SpanHandle) -> bool:
-    compatible = cast(Any, handle)
-    masked = getattr(compatible, "capture_masked", None)
-    if isinstance(masked, bool):
-        return masked
-    client = getattr(compatible, "_client", None)
-    return getattr(client, "mask", _OMIT) is not None if client is not None else True
-
-
 def _mark_capture_incomplete(fields: dict[str, Any], mask_output: bool) -> dict[str, Any]:
     fields["attributes"] = {**(fields.get("attributes") or {}), _CAPTURE_TRUNCATED_ATTRIBUTE: True}
     if mask_output:
@@ -842,7 +832,7 @@ def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: flo
         self._consume: Iterator[Any] | None = None
         self._in_next = False
         self._completed = False
-        self._mask_output_when_incomplete = _capture_masked(handle)
+        self._mask_output_when_incomplete = handle.capture_masked
         self._budget = _stream_capture_budget()
         # Bound to the budget, not self, so the callback adds no reference cycle.
         self._end_stream = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
@@ -928,9 +918,7 @@ def _record(self, chunk: Any, received_at: float) -> bool:
             return True
         update = _record_chat_chunk(chunk, self._states, self._budget)
         if _chat_chunk_has_output(chunk):
-            record_output_chunk = getattr(self._handle, "record_output_chunk", None)
-            if callable(record_output_chunk):
-                record_output_chunk(received_at * 1000)
+            self._handle.record_output_chunk(received_at * 1000)
         if not self._saw_first:
             self._saw_first = True
             self._handle.update(
@@ -957,7 +945,7 @@ def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: flo
         self._consume: AsyncIterator[Any] | None = None
         self._in_next = False
         self._completed = False
-        self._mask_output_when_incomplete = _capture_masked(handle)
+        self._mask_output_when_incomplete = handle.capture_masked
         self._budget = _stream_capture_budget()
         # Bound to the budget, not self, so the callback adds no reference cycle.
         self._end_stream = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
@@ -1046,9 +1034,7 @@ def _record(self, chunk: Any, received_at: float) -> bool:
             return True
         update = _record_chat_chunk(chunk, self._states, self._budget)
         if _chat_chunk_has_output(chunk):
-            record_output_chunk = getattr(self._handle, "record_output_chunk", None)
-            if callable(record_output_chunk):
-                record_output_chunk(received_at * 1000)
+            self._handle.record_output_chunk(received_at * 1000)
         if not self._saw_first:
             self._saw_first = True
             self._handle.update(
@@ -1584,7 +1570,7 @@ def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: flo
         # Bound to the state, not self, so the callback adds no reference cycle.
         self._end_stream = _ending_stream(_end_once(handle), self._state.release)
         self._ended = False
-        self._mask_output_when_incomplete = _capture_masked(handle)
+        self._mask_output_when_incomplete = handle.capture_masked
         self._consume: Iterator[Any] | None = None
         self._in_next = False
         _hook_response_close(inner, self._on_response_close)
@@ -1661,9 +1647,7 @@ def _record(self, event: Any, received_at: float) -> None:
         if self._ended:
             return
         if _response_event_has_output(event):
-            record_output_chunk = getattr(self._handle, "record_output_chunk", None)
-            if callable(record_output_chunk):
-                record_output_chunk(received_at * 1000)
+            self._handle.record_output_chunk(received_at * 1000)
         if not self._saw_first:
             self._saw_first = True
             self._handle.update(
@@ -1692,7 +1676,7 @@ def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: flo
         # Bound to the state, not self, so the callback adds no reference cycle.
         self._end_stream = _ending_stream(_end_once(handle), self._state.release)
         self._ended = False
-        self._mask_output_when_incomplete = _capture_masked(handle)
+        self._mask_output_when_incomplete = handle.capture_masked
         self._consume: AsyncIterator[Any] | None = None
         self._in_next = False
         _hook_response_close(inner, self._on_response_close)
@@ -1772,9 +1756,7 @@ def _record(self, event: Any, received_at: float) -> None:
         if self._ended:
             return
         if _response_event_has_output(event):
-            record_output_chunk = getattr(self._handle, "record_output_chunk", None)
-            if callable(record_output_chunk):
-                record_output_chunk(received_at * 1000)
+            self._handle.record_output_chunk(received_at * 1000)
         if not self._saw_first:
             self._saw_first = True
             self._handle.update(