From 83602dadc682396d8d2f500936344552f5bfe86f Mon Sep 17 00:00:00 2001 From: Aikins Laryea Date: Wed, 4 Nov 2026 00:57:54 +0000 Subject: [PATCH] refactor(sdk): call span handle members directly in the python integrations --- .../src/telemetry_dev_bedrock/_streams.py | 19 +++-------- .../telemetry_dev_google_genai/__init__.py | 28 +++------------ .../src/telemetry_dev_litellm/__init__.py | 18 ++-------- .../src/telemetry_dev_openrouter/__init__.py | 34 +++++-------------- 4 files changed, 20 insertions(+), 79 deletions(-) diff --git a/sdks/python-bedrock/src/telemetry_dev_bedrock/_streams.py b/sdks/python-bedrock/src/telemetry_dev_bedrock/_streams.py index e113d7e..3a4e584 100644 --- a/sdks/python-bedrock/src/telemetry_dev_bedrock/_streams.py +++ b/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 diff --git a/sdks/python-google-genai/src/telemetry_dev_google_genai/__init__.py b/sdks/python-google-genai/src/telemetry_dev_google_genai/__init__.py index dd90683..790ad33 100644 --- a/sdks/python-google-genai/src/telemetry_dev_google_genai/__init__.py +++ b/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) diff --git a/sdks/python-litellm/src/telemetry_dev_litellm/__init__.py b/sdks/python-litellm/src/telemetry_dev_litellm/__init__.py index affff2a..163ca1f 100644 --- a/sdks/python-litellm/src/telemetry_dev_litellm/__init__.py +++ b/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]] = {} diff --git a/sdks/python-openrouter/src/telemetry_dev_openrouter/__init__.py b/sdks/python-openrouter/src/telemetry_dev_openrouter/__init__.py index b7244c1..a510356 100644 --- a/sdks/python-openrouter/src/telemetry_dev_openrouter/__init__.py +++ b/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(