telemetry-dev-openai 0.2.2__tar.gz → 0.2.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: telemetry-dev-openai
3
- Version: 0.2.2
3
+ Version: 0.2.3
4
4
  Summary: OpenAI integration for telemetry.dev Python SDK
5
5
  Keywords: telemetry,opentelemetry,openai,llm,genai,tracing
6
6
  Author: telemetry.dev
@@ -101,9 +101,9 @@ The integration maps native OpenAI request/response shapes directly into telemet
101
101
 
102
102
  Chat completion streams are traced. Requests are sent unchanged by default, so token usage is only captured when the caller sets `stream_options={"include_usage": True}` themselves. Pass `inject_stream_usage=True` to `wrap_openai` or `instrument_openai` to inject it automatically; the synthetic usage-only chunk is then hidden from the caller. Injection is opt-in because some providers reject `stream_options` — for example Azure OpenAI "on your data" (`data_sources`) returns 400 for it while plain `stream=True` works.
103
103
 
104
- Streamed chat output and finish reasons use separate budgets of up to 48 KiB and 1,000 items; the configured `max_attribute_length` does not shrink these retention bounds; the core applies it to the exported attribute after the mask runs, with its `...[truncated]` marker, without setting `telemetry.dev.capture.truncated`. When present, `gen_ai.output.messages` contains bounded partial output. If output and/or finish-reason capture is incomplete, including when a choice ends without `finish_reason`, `telemetry.dev.capture.truncated` is `true`. With a mask configured, incomplete output is omitted because the mask cannot inspect the complete value. When output capture is disabled or omitted, response IDs and models, finish reasons, usage, timing, and errors may still be recorded.
104
+ Streamed chat output is kept up to 16 MiB, and all streams in flight share the `max_stream_bytes` bound of the core (64 MiB by default). These bounds need a `telemetry-dev` core with `max_stream_bytes`; with an older core, each stream keeps up to 64 KiB and 1,024 items. Each streamed finish reason is cut to its first 128 characters, without setting `telemetry.dev.capture.truncated`, and the finish reasons of a stream have their own budget, which fits at least 1,000 of them; the configured `max_attribute_length` does not shrink these retention bounds; the core applies it to the exported attribute after the mask runs, with its `...[truncated]` marker, without setting `telemetry.dev.capture.truncated`. A chat stream inspects at most 1,000 choices and 1,000 tool calls in each chunk; past these caps the rest is not kept and `telemetry.dev.capture.truncated` is set. When present, `gen_ai.output.messages` contains bounded partial output. If output and/or finish-reason capture is incomplete, including when a choice ends without `finish_reason`, `telemetry.dev.capture.truncated` is `true`. With a mask configured, incomplete output is omitted because the mask cannot inspect the complete value. When output capture is disabled or omitted, response IDs and models, finish reasons, usage, timing, and errors may still be recorded.
105
105
 
106
- Capture flags do not gate stop sequences or caller-supplied metadata or raw attributes; redact those separately when needed. Exception messages and stack traces are recorded only when both `capture_input` and `capture_output` are on, and pass through `mask`. Older `telemetry-dev` versions record them regardless of capture flags.
106
+ Capture flags do not gate stop sequences or caller-supplied metadata or raw attributes; redact those separately when needed. Exception messages and stack traces are recorded only when both `capture_input` and `capture_output` are on, and pass through `mask`.
107
107
 
108
108
  ```py
109
109
  client = wrap_openai(OpenAI(), inject_stream_usage=True)
@@ -111,7 +111,7 @@ client = wrap_openai(OpenAI(), inject_stream_usage=True)
111
111
 
112
112
  Responses API streams are traced through `responses.create(stream=True)` and `responses.stream(response_id=...)`; terminal `response.completed`, `response.failed`, and `response.incomplete` events close the span.
113
113
 
114
- Responses and transcription streams follow the same rule. Responses snapshots are retained within 48 KiB and the core default of 1,024 items. A Responses stream that ends without a terminal snapshot, ends with an `error` event, fails to map an event, or has a terminal snapshot over the limit sets `telemetry.dev.capture.truncated`; it keeps the last snapshot that fit, unless a mask is configured, in which case output is omitted. Transcription text has its own retention limit of up to 64 KiB with no item limit. A complete non-streamed transcript is passed to the core whole. A transcription stream that ends before `transcript.text.done` (an error, early close, or end of stream) or whose text exceeds the limit is flagged the same way and keeps its bounded partial text only when no mask is configured.
114
+ Responses and transcription streams follow the same rule. Responses snapshots and streamed transcription text are kept within the same stream bounds. A Responses stream that ends without a terminal snapshot, ends with an `error` event, fails to map an event, or has a terminal snapshot over the limit sets `telemetry.dev.capture.truncated`; it keeps the last snapshot that fit, unless a mask is configured, in which case output is omitted. A complete non-streamed transcript is passed to the core whole. A transcription stream that ends before `transcript.text.done` (an error, early close, or end of stream) or whose text exceeds the limit is flagged the same way and keeps its bounded partial text only when no mask is configured.
115
115
 
116
116
  ## Embeddings
117
117
 
@@ -78,9 +78,9 @@ The integration maps native OpenAI request/response shapes directly into telemet
78
78
 
79
79
  Chat completion streams are traced. Requests are sent unchanged by default, so token usage is only captured when the caller sets `stream_options={"include_usage": True}` themselves. Pass `inject_stream_usage=True` to `wrap_openai` or `instrument_openai` to inject it automatically; the synthetic usage-only chunk is then hidden from the caller. Injection is opt-in because some providers reject `stream_options` — for example Azure OpenAI "on your data" (`data_sources`) returns 400 for it while plain `stream=True` works.
80
80
 
81
- Streamed chat output and finish reasons use separate budgets of up to 48 KiB and 1,000 items; the configured `max_attribute_length` does not shrink these retention bounds; the core applies it to the exported attribute after the mask runs, with its `...[truncated]` marker, without setting `telemetry.dev.capture.truncated`. When present, `gen_ai.output.messages` contains bounded partial output. If output and/or finish-reason capture is incomplete, including when a choice ends without `finish_reason`, `telemetry.dev.capture.truncated` is `true`. With a mask configured, incomplete output is omitted because the mask cannot inspect the complete value. When output capture is disabled or omitted, response IDs and models, finish reasons, usage, timing, and errors may still be recorded.
81
+ Streamed chat output is kept up to 16 MiB, and all streams in flight share the `max_stream_bytes` bound of the core (64 MiB by default). These bounds need a `telemetry-dev` core with `max_stream_bytes`; with an older core, each stream keeps up to 64 KiB and 1,024 items. Each streamed finish reason is cut to its first 128 characters, without setting `telemetry.dev.capture.truncated`, and the finish reasons of a stream have their own budget, which fits at least 1,000 of them; the configured `max_attribute_length` does not shrink these retention bounds; the core applies it to the exported attribute after the mask runs, with its `...[truncated]` marker, without setting `telemetry.dev.capture.truncated`. A chat stream inspects at most 1,000 choices and 1,000 tool calls in each chunk; past these caps the rest is not kept and `telemetry.dev.capture.truncated` is set. When present, `gen_ai.output.messages` contains bounded partial output. If output and/or finish-reason capture is incomplete, including when a choice ends without `finish_reason`, `telemetry.dev.capture.truncated` is `true`. With a mask configured, incomplete output is omitted because the mask cannot inspect the complete value. When output capture is disabled or omitted, response IDs and models, finish reasons, usage, timing, and errors may still be recorded.
82
82
 
83
- Capture flags do not gate stop sequences or caller-supplied metadata or raw attributes; redact those separately when needed. Exception messages and stack traces are recorded only when both `capture_input` and `capture_output` are on, and pass through `mask`. Older `telemetry-dev` versions record them regardless of capture flags.
83
+ Capture flags do not gate stop sequences or caller-supplied metadata or raw attributes; redact those separately when needed. Exception messages and stack traces are recorded only when both `capture_input` and `capture_output` are on, and pass through `mask`.
84
84
 
85
85
  ```py
86
86
  client = wrap_openai(OpenAI(), inject_stream_usage=True)
@@ -88,7 +88,7 @@ client = wrap_openai(OpenAI(), inject_stream_usage=True)
88
88
 
89
89
  Responses API streams are traced through `responses.create(stream=True)` and `responses.stream(response_id=...)`; terminal `response.completed`, `response.failed`, and `response.incomplete` events close the span.
90
90
 
91
- Responses and transcription streams follow the same rule. Responses snapshots are retained within 48 KiB and the core default of 1,024 items. A Responses stream that ends without a terminal snapshot, ends with an `error` event, fails to map an event, or has a terminal snapshot over the limit sets `telemetry.dev.capture.truncated`; it keeps the last snapshot that fit, unless a mask is configured, in which case output is omitted. Transcription text has its own retention limit of up to 64 KiB with no item limit. A complete non-streamed transcript is passed to the core whole. A transcription stream that ends before `transcript.text.done` (an error, early close, or end of stream) or whose text exceeds the limit is flagged the same way and keeps its bounded partial text only when no mask is configured.
91
+ Responses and transcription streams follow the same rule. Responses snapshots and streamed transcription text are kept within the same stream bounds. A Responses stream that ends without a terminal snapshot, ends with an `error` event, fails to map an event, or has a terminal snapshot over the limit sets `telemetry.dev.capture.truncated`; it keeps the last snapshot that fit, unless a mask is configured, in which case output is omitted. A complete non-streamed transcript is passed to the core whole. A transcription stream that ends before `transcript.text.done` (an error, early close, or end of stream) or whose text exceeds the limit is flagged the same way and keeps its bounded partial text only when no mask is configured.
92
92
 
93
93
  ## Embeddings
94
94
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "telemetry-dev-openai"
3
- version = "0.2.2"
3
+ version = "0.2.3"
4
4
  description = "OpenAI integration for telemetry.dev Python SDK"
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "telemetry-dev-openai"
3
- version = "0.2.2"
3
+ version = "0.2.3"
4
4
  description = "OpenAI integration for telemetry.dev Python SDK"
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -1,10 +1,13 @@
1
1
  from __future__ import annotations
2
2
 
3
+ import os
4
+ import re
3
5
  import threading
4
6
  import time
5
7
  import types
8
+ import weakref
6
9
  from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
7
- from functools import wraps
10
+ from functools import partial, wraps
8
11
  from io import StringIO
9
12
  from typing import Any, Literal, TypeVar, cast
10
13
  from urllib.parse import urlsplit
@@ -20,7 +23,7 @@ from openai.resources.embeddings import AsyncEmbeddings, Embeddings
20
23
  from openai.resources.images import AsyncImages, Images
21
24
  from openai.resources.responses.responses import AsyncResponses, Responses
22
25
 
23
- __version__ = "0.2.2"
26
+ __version__ = "0.2.3"
24
27
 
25
28
  ProviderResolver = Callable[[object | None], str]
26
29
  RequestMapper = Callable[[Mapping[str, Any]], tuple[str, dict[str, Any]]]
@@ -31,11 +34,30 @@ _ORIGINAL_ATTR = "_telemetry_dev_openai_original"
31
34
  _ORIGINALS: list[tuple[type[Any], str, Any, Any]] = []
32
35
  _installed = False
33
36
  _install_lock = threading.Lock()
37
+
38
+
39
+ def _reset_install_lock_in_child() -> None:
40
+ """A thread of the parent can hold the lock at a fork, and only this thread exists here."""
41
+ global _install_lock
42
+ _install_lock = threading.Lock()
43
+
44
+
45
+ # A hook is enough: a fork that runs no hooks, as uWSGI's fork of its workers by default,
46
+ # starts from its master, which serves no requests. os.register_at_fork cannot remove a hook,
47
+ # and importlib.reload runs this module again with the same globals, so the hook of the first
48
+ # run stays. Windows has no fork.
49
+ if hasattr(os, "register_at_fork") and "_install_fork_hook" not in globals():
50
+ _install_fork_hook = True
51
+ os.register_at_fork(after_in_child=_reset_install_lock_in_child)
52
+
34
53
  _T = TypeVar("_T")
35
54
  _OMIT = object()
36
- _CHAT_STREAM_CAPTURE_MAX_BYTES = 48 * 1024
37
- _TRANSCRIPT_CAPTURE_MAX_BYTES = 64 * 1024
55
+ # More choices, tool calls, or finish reasons than this are not kept.
38
56
  _CHAT_STREAM_CAPTURE_MAX_ITEMS = 1000
57
+ # Each streamed finish reason is cut to this many characters, as in the other integrations.
58
+ _FINISH_REASON_MAX_LENGTH = 128
59
+ _TEXT_CHUNK_CHARS = 4096
60
+ _SURROGATES = re.compile("[\ud800-\udfff]")
39
61
 
40
62
 
41
63
  class _CaptureLimit(Exception):
@@ -65,6 +87,12 @@ def _structure_cost(shape: object) -> tuple[int, int]:
65
87
 
66
88
 
67
89
  _FINISH_REASON_BYTES, _FINISH_REASON_ITEMS = _structure_cost({"finish_reason": "", "index": 0})
90
+ # Finish reasons stay outside the stream budget so a terminal one is still exported. Their own
91
+ # budget fits _CHAT_STREAM_CAPTURE_MAX_ITEMS of the largest records a cut reason can take.
92
+ _FINISH_REASON_BUDGET_BYTES = _CHAT_STREAM_CAPTURE_MAX_ITEMS * (
93
+ _FINISH_REASON_BYTES
94
+ + telemetry_dev.CaptureBudget.json_string_bytes("\x00" * _FINISH_REASON_MAX_LENGTH)
95
+ )
68
96
  _TOOL_CALLS_FIELD_BYTES = _field_bytes("tool_calls")
69
97
  _CONTENT_FIELD_BYTES = _field_bytes("content")
70
98
  _FUNCTION_CALL_FIELD_BYTES = _field_bytes("function_call")
@@ -147,6 +175,7 @@ def _bounded_responses_capture(
147
175
  *,
148
176
  capture_enabled: bool | None = None,
149
177
  max_bytes: int | None = None,
178
+ budget: telemetry_dev.CaptureBudget | None = None,
150
179
  report_error: Callable[[BaseException], None] | None = None,
151
180
  ) -> tuple[Any | None, bool]:
152
181
  client = telemetry_dev.get_client()
@@ -156,11 +185,12 @@ def _bounded_responses_capture(
156
185
  capture_enabled = client.capture_input if capture == "input" else client.capture_output
157
186
  if not capture_enabled:
158
187
  return None, False
159
- budget = (
160
- telemetry_dev.CaptureBudget(max_bytes=max_bytes)
161
- if max_bytes is not None
162
- else telemetry_dev.CaptureBudget()
163
- )
188
+ if budget is None:
189
+ budget = (
190
+ telemetry_dev.CaptureBudget(max_bytes=max_bytes)
191
+ if max_bytes is not None
192
+ else telemetry_dev.CaptureBudget()
193
+ )
164
194
  ancestors: set[int] = set()
165
195
 
166
196
  def reserve(byte_count: int) -> None:
@@ -169,7 +199,8 @@ def _bounded_responses_capture(
169
199
  or budget.items_used >= budget.max_items
170
200
  ):
171
201
  raise _CaptureLimit
172
- budget.bytes_used += byte_count
202
+ if not _reserve_bytes(budget, budget.bytes_used + byte_count):
203
+ raise _CaptureLimit
173
204
  budget.items_used += 1
174
205
 
175
206
  def convert(item: Any, parent_type: str | None, depth: int) -> Any:
@@ -313,36 +344,85 @@ def _sent_params(params: Mapping[str, Any]) -> dict[str, Any]:
313
344
  }
314
345
 
315
346
 
316
- def _transcript_capture_budget() -> telemetry_dev.CaptureBudget:
317
- return telemetry_dev.CaptureBudget(max_bytes=_TRANSCRIPT_CAPTURE_MAX_BYTES)
347
+ def _stream_budget() -> telemetry_dev.CaptureBudget:
348
+ # Older cores have no for_stream and bound each budget at 64 KiB on its own.
349
+ for_stream = getattr(telemetry_dev.CaptureBudget, "for_stream", None)
350
+ budget = for_stream() if callable(for_stream) else telemetry_dev.CaptureBudget()
351
+ return cast(telemetry_dev.CaptureBudget, budget)
352
+
353
+
354
+ def _release_budget(budget: telemetry_dev.CaptureBudget) -> None:
355
+ release = getattr(budget, "release", None)
356
+ if callable(release):
357
+ release()
358
+
359
+
360
+ def _reserve_bytes(budget: telemetry_dev.CaptureBudget, total: int) -> bool:
361
+ # Older cores have no reserve_bytes; there the caller's check against max_bytes stands.
362
+ reserve = getattr(budget, "reserve_bytes", None)
363
+ if callable(reserve):
364
+ return bool(reserve(total))
365
+ budget.bytes_used = total
366
+ return True
367
+
368
+
369
+ def _ending_stream(end: Callable[..., None], release: Callable[[], None]) -> Callable[..., None]:
370
+ def end_stream(**fields: Any) -> None:
371
+ try:
372
+ end(**fields)
373
+ finally:
374
+ release()
375
+
376
+ return end_stream
377
+
378
+
379
+ def _character_bytes(character: str) -> int:
380
+ code = ord(character)
381
+ if code <= 0x7F:
382
+ return 1
383
+ if code <= 0x7FF:
384
+ return 2
385
+ if 0xD800 <= code <= 0xDFFF:
386
+ return 6
387
+ if code <= 0xFFFF:
388
+ return 3
389
+ return 4
390
+
391
+
392
+ def _text_bytes(value: str) -> int:
393
+ """The bytes _character_bytes charges for ``value``: its UTF-8 width, and 6 per surrogate."""
394
+ return len(value.encode("utf-8", "surrogatepass")) + 3 * len(_SURROGATES.findall(value))
318
395
 
319
396
 
320
397
  def _capture_text(value: str, budget: telemetry_dev.CaptureBudget) -> str:
321
398
  if not value or budget.truncated:
322
399
  return ""
323
- if budget.remaining_bytes <= 0:
400
+ # Read once: the remaining bytes of a stream budget take a lock.
401
+ remaining = budget.remaining_bytes
402
+ if remaining <= 0:
324
403
  budget.truncated = True
325
404
  return ""
326
405
  byte_count = 0
327
- retained_characters = 0
328
- for index, character in enumerate(value):
329
- code = ord(character)
330
- if code <= 0x7F:
331
- character_bytes = 1
332
- elif code <= 0x7FF:
333
- character_bytes = 2
334
- elif 0xD800 <= code <= 0xDFFF:
335
- character_bytes = 6
336
- elif code <= 0xFFFF:
337
- character_bytes = 3
338
- else:
339
- character_bytes = 4
340
- if byte_count + character_bytes > budget.remaining_bytes:
341
- budget.truncated = True
342
- break
343
- byte_count += character_bytes
344
- retained_characters = index + 1
345
- budget.bytes_used += byte_count
406
+ retained_characters = len(value)
407
+ # Whole chunks are measured at once; only the chunk that does not fit is walked by character.
408
+ # Every character costs at least one byte, so the scan stops within remaining + one chunk.
409
+ for start in range(0, len(value), _TEXT_CHUNK_CHARS):
410
+ chunk = value[start : start + _TEXT_CHUNK_CHARS]
411
+ chunk_bytes = _text_bytes(chunk)
412
+ if byte_count + chunk_bytes <= remaining:
413
+ byte_count += chunk_bytes
414
+ continue
415
+ retained_characters = start
416
+ for character in chunk:
417
+ character_bytes = _character_bytes(character)
418
+ if byte_count + character_bytes > remaining:
419
+ break
420
+ byte_count += character_bytes
421
+ retained_characters += 1
422
+ break
423
+ if not _reserve_bytes(budget, budget.bytes_used + byte_count):
424
+ budget.truncated = True
425
+ return ""
346
426
  if retained_characters < len(value):
347
427
  budget.truncated = True
348
428
  return value[:retained_characters]
@@ -692,12 +772,92 @@ def _end_once(handle: telemetry_dev.SpanHandle) -> Callable[..., None]:
692
772
  return end
693
773
 
694
774
 
775
+ # On Python 3.10 and 3.11, StringIO keeps up to 100,000 written strings in a list before it
776
+ # joins them, about 60 bytes each that the stream bound does not charge: 15 bytes a character
777
+ # for 4-character deltas. Deltas are joined 256 at a time, so that those waiting hold under
778
+ # about 22 KB, onto a string that is written once it has 4,096 characters, so that what is
779
+ # written costs about 1.5% more than its text. Python 3.12 dropped the list.
780
+ _DELTA_BATCH = 256
781
+ _WRITE_CHARS = 4096
782
+
783
+
784
+ class _DeltaBuffer:
785
+ # A write takes several steps, unlike one StringIO call, so without the lock a read on
786
+ # another thread, such as one that closes the stream, could copy a batch twice or drop it.
787
+ # A fork copies the lock held if another thread was writing, and that thread does not run
788
+ # in the child to release it. Every fork, also one that runs no at-fork hooks, changes the
789
+ # process ID, so a child takes a new lock. The check is inline because it runs per delta.
790
+ # A signal handler runs on the thread it interrupts, so one that reads the stream, such as
791
+ # one that closes it, can read inside a write. So the lock is reentrant, and each step takes
792
+ # the deltas and the tail out of the buffer before it joins or writes them, so that a read
793
+ # inside the step does not write them as well; that read returns a prefix, which can leave
794
+ # out the deltas being joined. A handler that writes inside a step is not supported. An
795
+ # exception inside a step, such as KeyboardInterrupt from a handler, puts back the text the
796
+ # step took unless the StringIO has it, so that a later read still has every delta.
797
+ __slots__ = ("_buffer", "_length", "_lock", "_pending", "_pid", "_tail")
798
+
799
+ def __init__(self) -> None:
800
+ self._buffer = StringIO()
801
+ self._length = 0
802
+ self._lock = threading.RLock()
803
+ self._pid = os.getpid()
804
+ self._pending: list[str] = []
805
+ self._tail = ""
806
+
807
+ def write(self, text: str) -> None:
808
+ pid = os.getpid()
809
+ if self._pid != pid:
810
+ self._lock = threading.RLock()
811
+ self._pid = pid
812
+ with self._lock:
813
+ # Counted first: a signal handler runs only once a call returns, so none runs
814
+ # between the count and the append, and one that raises after the append, or in
815
+ # a join, finds the delta counted.
816
+ self._length += len(text)
817
+ self._pending.append(text)
818
+ if len(self._pending) >= _DELTA_BATCH:
819
+ self._join(_WRITE_CHARS)
820
+
821
+ def tell(self) -> int:
822
+ # Read on every content delta, so it counts the pending deltas instead of writing them.
823
+ return self._length
824
+
825
+ def getvalue(self) -> str:
826
+ pid = os.getpid()
827
+ if self._pid != pid:
828
+ self._lock = threading.RLock()
829
+ self._pid = pid
830
+ with self._lock:
831
+ self._join(1)
832
+ return self._buffer.getvalue()
833
+
834
+ def _join(self, write_at: int) -> None:
835
+ if not self._pending and len(self._tail) < write_at:
836
+ return
837
+ written = self._buffer.tell()
838
+ pending: list[str] = []
839
+ tail = ""
840
+ try:
841
+ pending, self._pending = self._pending, []
842
+ tail, self._tail = self._tail, ""
843
+ joined = "".join([tail, *pending])
844
+ if len(joined) >= write_at:
845
+ self._buffer.write(joined)
846
+ else:
847
+ self._tail = joined
848
+ except BaseException:
849
+ if self._buffer.tell() == written:
850
+ self._pending[:0] = pending
851
+ self._tail = tail + self._tail
852
+ raise
853
+
854
+
695
855
  class _ChatChoice:
696
856
  def __init__(self) -> None:
697
857
  self.role: str | None = None
698
858
  self.role_resolved = True
699
- self.content_fragments = StringIO()
700
- self.refusal_fragments = StringIO()
859
+ self.content_fragments = _DeltaBuffer()
860
+ self.refusal_fragments = _DeltaBuffer()
701
861
  self.function_call: dict[str, Any] | None = None
702
862
  self.unresolved_function_scalars: set[str] = set()
703
863
  self.tool_calls: dict[int, dict[str, Any]] = {}
@@ -737,10 +897,7 @@ class _ChatToolCallDelta:
737
897
 
738
898
 
739
899
  def _chat_capture_budget(*, reserve_output_list: bool = False) -> telemetry_dev.CaptureBudget:
740
- budget = telemetry_dev.CaptureBudget(
741
- _CHAT_STREAM_CAPTURE_MAX_BYTES,
742
- _CHAT_STREAM_CAPTURE_MAX_ITEMS,
743
- )
900
+ budget = _stream_budget()
744
901
  if reserve_output_list:
745
902
  budget.accept([])
746
903
  return budget
@@ -768,7 +925,11 @@ def _replace_finish_reason(
768
925
  reservations.pop(index, None)
769
926
  return False
770
927
  reservation = (_FINISH_REASON_BYTES + value_bytes, _FINISH_REASON_ITEMS)
771
- budget.bytes_used = base_bytes + reservation[0]
928
+ if not _reserve_bytes(budget, base_bytes + reservation[0]):
929
+ budget.bytes_used = base_bytes
930
+ budget.items_used = base_items
931
+ reservations.pop(index, None)
932
+ return False
772
933
  budget.items_used = base_items + reservation[1]
773
934
  reservations[index] = reservation
774
935
  return True
@@ -789,7 +950,10 @@ def _reserve_chat_budget(
789
950
  if not recoverable:
790
951
  budget.truncated = True
791
952
  return False
792
- budget.bytes_used += byte_count
953
+ if not _reserve_bytes(budget, budget.bytes_used + byte_count):
954
+ if not recoverable:
955
+ budget.truncated = True
956
+ return False
793
957
  budget.items_used += item_count
794
958
  return True
795
959
 
@@ -837,8 +1001,7 @@ def _replace_chat_scalar(
837
1001
  value_bytes = telemetry_dev.CaptureBudget.json_string_bytes(value, available_bytes)
838
1002
  if value_bytes > available_bytes:
839
1003
  return False
840
- budget.bytes_used = base_bytes + _BASE_ITEM_BYTES + value_bytes
841
- return True
1004
+ return _reserve_bytes(budget, base_bytes + _BASE_ITEM_BYTES + value_bytes)
842
1005
 
843
1006
 
844
1007
  def _release_chat_field(field_name: str, value: str, budget: telemetry_dev.CaptureBudget) -> None:
@@ -911,8 +1074,8 @@ def _capture_tool_call_payload(
911
1074
  unresolved.add(unresolved_name)
912
1075
  if not budget.truncated and value is not None:
913
1076
  fragments = payload.get(value_key)
914
- had_value = isinstance(fragments, StringIO)
915
- value_fragments = fragments if had_value else StringIO()
1077
+ had_value = isinstance(fragments, _DeltaBuffer)
1078
+ value_fragments = fragments if had_value else _DeltaBuffer()
916
1079
  if value or not had_value:
917
1080
  if _capture_chat_string(
918
1081
  value,
@@ -1065,8 +1228,8 @@ def _capture_function_call_delta(
1065
1228
  state.unresolved_function_scalars.add("name")
1066
1229
  if not budget.truncated and arguments is not None:
1067
1230
  fragments = current.get("arguments")
1068
- had_arguments = isinstance(fragments, StringIO)
1069
- argument_fragments = fragments if had_arguments else StringIO()
1231
+ had_arguments = isinstance(fragments, _DeltaBuffer)
1232
+ argument_fragments = fragments if had_arguments else _DeltaBuffer()
1070
1233
  if arguments or not had_arguments:
1071
1234
  if _capture_chat_string(
1072
1235
  arguments,
@@ -1139,6 +1302,7 @@ def _record_chat_chunk(
1139
1302
  )
1140
1303
  )
1141
1304
  if finish_reason is not None:
1305
+ finish_reason = finish_reason[:_FINISH_REASON_MAX_LENGTH]
1142
1306
  if unterminated_choices is not None:
1143
1307
  unterminated_choices.discard(choice_index)
1144
1308
  if _replace_finish_reason(
@@ -1303,7 +1467,7 @@ def _chat_message(state: _ChatChoice) -> dict[str, Any]:
1303
1467
  if state.function_call is not None:
1304
1468
  captured_function_call = dict(state.function_call)
1305
1469
  arguments = captured_function_call.get("arguments")
1306
- if isinstance(arguments, StringIO):
1470
+ if isinstance(arguments, _DeltaBuffer):
1307
1471
  captured_function_call["arguments"] = arguments.getvalue()
1308
1472
  if "name" in state.unresolved_function_scalars:
1309
1473
  captured_function_call.pop("name", None)
@@ -1322,7 +1486,7 @@ def _chat_message(state: _ChatChoice) -> dict[str, Any]:
1322
1486
  if isinstance(raw_payload, Mapping):
1323
1487
  captured_payload = dict(cast(Mapping[str, Any], raw_payload))
1324
1488
  fragments = captured_payload.get(value_key)
1325
- if isinstance(fragments, StringIO):
1489
+ if isinstance(fragments, _DeltaBuffer):
1326
1490
  captured_payload[value_key] = fragments.getvalue()
1327
1491
  else:
1328
1492
  captured_payload = {}
@@ -1443,6 +1607,15 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
1443
1607
  response = getattr(inner, "response", None)
1444
1608
  if response is None:
1445
1609
  return
1610
+ # ``finish`` is a method of the stream wrapper, which holds ``response``: held weakly, the hook
1611
+ # adds no reference cycle, so a dropped wrapper and its stream budget are freed at once.
1612
+ finish_method = weakref.WeakMethod(finish)
1613
+
1614
+ def finish_if_alive() -> None:
1615
+ method = finish_method()
1616
+ if method is not None:
1617
+ method()
1618
+
1446
1619
  close = getattr(response, "close", None)
1447
1620
  if callable(close):
1448
1621
 
@@ -1450,7 +1623,7 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
1450
1623
  try:
1451
1624
  return close(*args, **kwargs)
1452
1625
  finally:
1453
- finish()
1626
+ finish_if_alive()
1454
1627
 
1455
1628
  response.close = _close_hook
1456
1629
  aclose = getattr(response, "aclose", None)
@@ -1461,7 +1634,7 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
1461
1634
  try:
1462
1635
  return await aclose_fn(*args, **kwargs)
1463
1636
  finally:
1464
- finish()
1637
+ finish_if_alive()
1465
1638
 
1466
1639
  response.aclose = _aclose_hook
1467
1640
 
@@ -1516,7 +1689,6 @@ class _InstrumentedStream:
1516
1689
  started_at: float,
1517
1690
  ) -> None:
1518
1691
  self._inner = inner
1519
- self._end = _end_once(handle)
1520
1692
  self._handle = handle
1521
1693
  self._injected_usage = injected_usage
1522
1694
  self._started_at = started_at
@@ -1536,7 +1708,9 @@ class _InstrumentedStream:
1536
1708
  self._consume: Iterator[Any] | None = None
1537
1709
  self._in_next = False
1538
1710
  self._budget = _chat_capture_budget(reserve_output_list=self._capture_output)
1539
- self._finish_reason_budget = _chat_capture_budget()
1711
+ # Bound to the budget, not self, so the callback adds no reference cycle.
1712
+ self._end = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
1713
+ self._finish_reason_budget = telemetry_dev.CaptureBudget(_FINISH_REASON_BUDGET_BYTES)
1540
1714
  _hook_response_close(inner, self._on_response_close)
1541
1715
 
1542
1716
  def _iterate(self) -> Iterator[Any]:
@@ -1679,7 +1853,6 @@ class _InstrumentedAsyncStream:
1679
1853
  started_at: float,
1680
1854
  ) -> None:
1681
1855
  self._inner = inner
1682
- self._end = _end_once(handle)
1683
1856
  self._handle = handle
1684
1857
  self._injected_usage = injected_usage
1685
1858
  self._started_at = started_at
@@ -1699,7 +1872,9 @@ class _InstrumentedAsyncStream:
1699
1872
  self._consume: AsyncIterator[Any] | None = None
1700
1873
  self._in_next = False
1701
1874
  self._budget = _chat_capture_budget(reserve_output_list=self._capture_output)
1702
- self._finish_reason_budget = _chat_capture_budget()
1875
+ # Bound to the budget, not self, so the callback adds no reference cycle.
1876
+ self._end = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
1877
+ self._finish_reason_budget = telemetry_dev.CaptureBudget(_FINISH_REASON_BUDGET_BYTES)
1703
1878
  _hook_response_close(inner, self._on_response_close)
1704
1879
 
1705
1880
  async def _aiterate(self) -> AsyncIterator[Any]:
@@ -1851,7 +2026,12 @@ class _InstrumentedResponsesStream:
1851
2026
  def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
1852
2027
  self._inner = inner
1853
2028
  self._handle = handle
1854
- self._end = _end_once(handle)
2029
+ self._output_budget = _stream_budget()
2030
+ # Bound to the budget, not self, so the callback adds no reference cycle.
2031
+ self._end_stream = _ending_stream(
2032
+ _end_once(handle), partial(_release_budget, self._output_budget)
2033
+ )
2034
+ self._ended = False
1855
2035
  (
1856
2036
  self._capture_output,
1857
2037
  self._mask_output_when_incomplete,
@@ -1939,7 +2119,15 @@ class _InstrumentedResponsesStream:
1939
2119
  def __getattr__(self, name: str) -> Any:
1940
2120
  return getattr(self._inner, name)
1941
2121
 
2122
+ def _end(self, **fields: Any) -> None:
2123
+ # The span has ended and the stream budget no longer counts bytes, so later events are
2124
+ # not kept.
2125
+ self._ended = True
2126
+ self._end_stream(**fields)
2127
+
1942
2128
  def _record(self, event: Any, received_at: float) -> None:
2129
+ if self._ended:
2130
+ return
1943
2131
  if _response_event_has_output(event):
1944
2132
  self._handle.record_output_chunk(received_at * 1000)
1945
2133
  if not self._saw_first:
@@ -1957,13 +2145,24 @@ class _InstrumentedResponsesStream:
1957
2145
  output = None
1958
2146
  truncated = False
1959
2147
  if raw_output is not None:
2148
+ # The snapshot replaces the output kept so far.
2149
+ measured = telemetry_dev.CaptureBudget(self._output_budget.max_bytes)
1960
2150
  output, truncated = _bounded_responses_capture(
1961
2151
  raw_output,
1962
2152
  capture_enabled=self._capture_output,
1963
- max_bytes=_CHAT_STREAM_CAPTURE_MAX_BYTES,
2153
+ budget=measured,
1964
2154
  report_error=self._report_error,
1965
2155
  )
1966
- if not truncated and output:
2156
+ # A snapshot that fits replaces the output kept so far. Until one fits, the
2157
+ # stream keeps the cut one, so it holds that one's bytes.
2158
+ keeps = output and (
2159
+ not truncated
2160
+ or (self._retained_output is None and not self._mask_output_when_incomplete)
2161
+ )
2162
+ if keeps and not _reserve_bytes(self._output_budget, measured.bytes_used):
2163
+ # Another stream took the room after this one measured.
2164
+ output, truncated = None, True
2165
+ elif keeps and not truncated:
1967
2166
  self._retained_output = output
1968
2167
  self._partial = fields
1969
2168
  if self._retained_output is not None and not (
@@ -1997,7 +2196,12 @@ class _InstrumentedAsyncResponsesStream:
1997
2196
  def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
1998
2197
  self._inner = inner
1999
2198
  self._handle = handle
2000
- self._end = _end_once(handle)
2199
+ self._output_budget = _stream_budget()
2200
+ # Bound to the budget, not self, so the callback adds no reference cycle.
2201
+ self._end_stream = _ending_stream(
2202
+ _end_once(handle), partial(_release_budget, self._output_budget)
2203
+ )
2204
+ self._ended = False
2001
2205
  (
2002
2206
  self._capture_output,
2003
2207
  self._mask_output_when_incomplete,
@@ -2090,7 +2294,15 @@ class _InstrumentedAsyncResponsesStream:
2090
2294
  def __getattr__(self, name: str) -> Any:
2091
2295
  return getattr(self._inner, name)
2092
2296
 
2297
+ def _end(self, **fields: Any) -> None:
2298
+ # The span has ended and the stream budget no longer counts bytes, so later events are
2299
+ # not kept.
2300
+ self._ended = True
2301
+ self._end_stream(**fields)
2302
+
2093
2303
  def _record(self, event: Any, received_at: float) -> None:
2304
+ if self._ended:
2305
+ return
2094
2306
  if _response_event_has_output(event):
2095
2307
  self._handle.record_output_chunk(received_at * 1000)
2096
2308
  if not self._saw_first:
@@ -2108,13 +2320,24 @@ class _InstrumentedAsyncResponsesStream:
2108
2320
  output = None
2109
2321
  truncated = False
2110
2322
  if raw_output is not None:
2323
+ # The snapshot replaces the output kept so far.
2324
+ measured = telemetry_dev.CaptureBudget(self._output_budget.max_bytes)
2111
2325
  output, truncated = _bounded_responses_capture(
2112
2326
  raw_output,
2113
2327
  capture_enabled=self._capture_output,
2114
- max_bytes=_CHAT_STREAM_CAPTURE_MAX_BYTES,
2328
+ budget=measured,
2115
2329
  report_error=self._report_error,
2116
2330
  )
2117
- if not truncated and output:
2331
+ # A snapshot that fits replaces the output kept so far. Until one fits, the
2332
+ # stream keeps the cut one, so it holds that one's bytes.
2333
+ keeps = output and (
2334
+ not truncated
2335
+ or (self._retained_output is None and not self._mask_output_when_incomplete)
2336
+ )
2337
+ if keeps and not _reserve_bytes(self._output_budget, measured.bytes_used):
2338
+ # Another stream took the room after this one measured.
2339
+ output, truncated = None, True
2340
+ elif keeps and not truncated:
2118
2341
  self._retained_output = output
2119
2342
  self._partial = fields
2120
2343
  if self._retained_output is not None and not (
@@ -2148,6 +2371,7 @@ def _media_stream_event_fields(
2148
2371
  event: Any,
2149
2372
  response_mapper: ResponseMapper,
2150
2373
  capture_output: bool,
2374
+ budget: telemetry_dev.CaptureBudget,
2151
2375
  ) -> dict[str, Any]:
2152
2376
  if response_mapper is not _text_media_response:
2153
2377
  return response_mapper(event)
@@ -2155,7 +2379,9 @@ def _media_stream_event_fields(
2155
2379
  if capture_output and _string(_field(event, "type")) == "transcript.text.done":
2156
2380
  text = _string(_field(event, "text"))
2157
2381
  if text is not None:
2158
- budget = _transcript_capture_budget()
2382
+ # The terminal text replaces the deltas, within the room the stream has.
2383
+ budget.bytes_used = 0
2384
+ budget.truncated = False
2159
2385
  fields["output"] = _capture_text(text, budget)
2160
2386
  if budget.truncated:
2161
2387
  fields["attributes"] = {"telemetry.dev.capture.truncated": True}
@@ -2165,7 +2391,7 @@ def _media_stream_event_fields(
2165
2391
  def _media_stream_end_fields(
2166
2392
  fields: Mapping[str, Any],
2167
2393
  response_mapper: ResponseMapper,
2168
- text_parts: Sequence[str],
2394
+ text_fragments: _DeltaBuffer,
2169
2395
  budget: telemetry_dev.CaptureBudget,
2170
2396
  mask_output_when_incomplete: bool,
2171
2397
  stream_incomplete: bool,
@@ -2180,8 +2406,9 @@ def _media_stream_end_fields(
2180
2406
  )
2181
2407
  if "output" not in result:
2182
2408
  capture_incomplete = capture_incomplete or budget.truncated
2183
- if text_parts and not (mask_output_when_incomplete and capture_incomplete):
2184
- result["output"] = "".join(text_parts)
2409
+ text = text_fragments.getvalue()
2410
+ if text and not (mask_output_when_incomplete and capture_incomplete):
2411
+ result["output"] = text
2185
2412
  else:
2186
2413
  capture_incomplete = capture_incomplete or terminal_truncated
2187
2414
  if mask_output_when_incomplete and capture_incomplete:
@@ -2240,15 +2467,17 @@ class _InstrumentedMediaStream:
2240
2467
  self, inner: Any, handle: telemetry_dev.SpanHandle, response_mapper: ResponseMapper
2241
2468
  ) -> None:
2242
2469
  self._inner = inner
2243
- self._end = _end_once(handle)
2244
2470
  self._fields: dict[str, Any] = {}
2245
2471
  self._response_mapper = response_mapper
2246
2472
  (
2247
2473
  self._capture_output,
2248
2474
  self._mask_output_when_incomplete,
2249
2475
  ) = _chat_span_capture_policy(handle)
2250
- self._budget = _transcript_capture_budget()
2251
- self._text_parts: list[str] = []
2476
+ self._budget = _stream_budget()
2477
+ # Bound to the budget, not self, so the callback adds no reference cycle.
2478
+ self._end_stream = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
2479
+ self._ended = False
2480
+ self._text_fragments = _DeltaBuffer()
2252
2481
  self._saw_delta = False
2253
2482
  self._saw_terminal = False
2254
2483
  self._report_error = _chat_span_reporter(handle)
@@ -2304,7 +2533,15 @@ class _InstrumentedMediaStream:
2304
2533
  if failure is not None:
2305
2534
  self._end(**self._finish_fields(), error=failure)
2306
2535
 
2536
+ def _end(self, **fields: Any) -> None:
2537
+ # The span has ended and the stream budget no longer counts bytes, so later events are
2538
+ # not kept.
2539
+ self._ended = True
2540
+ self._end_stream(**fields)
2541
+
2307
2542
  def _record(self, event: Any) -> None:
2543
+ if self._ended:
2544
+ return
2308
2545
  event_type = _string(_field(event, "type")) or ""
2309
2546
  if event_type in _MEDIA_TERMINAL_EVENTS:
2310
2547
  self._saw_terminal = True
@@ -2312,6 +2549,7 @@ class _InstrumentedMediaStream:
2312
2549
  event,
2313
2550
  self._response_mapper,
2314
2551
  self._capture_output,
2552
+ self._budget,
2315
2553
  )
2316
2554
  if "output" in mapped:
2317
2555
  self._fields = _clean_fields(mapped)
@@ -2327,7 +2565,7 @@ class _InstrumentedMediaStream:
2327
2565
  if text:
2328
2566
  captured = _capture_text(text, self._budget)
2329
2567
  if captured:
2330
- self._text_parts.append(captured)
2568
+ self._text_fragments.write(captured)
2331
2569
  if event_type in _MEDIA_TERMINAL_EVENTS:
2332
2570
  error = _media_stream_error(event, event_type)
2333
2571
  self._end(**self._finish_fields(), **({"error": error} if error is not None else {}))
@@ -2336,7 +2574,7 @@ class _InstrumentedMediaStream:
2336
2574
  return _media_stream_end_fields(
2337
2575
  self._fields,
2338
2576
  self._response_mapper,
2339
- self._text_parts,
2577
+ self._text_fragments,
2340
2578
  self._budget,
2341
2579
  self._mask_output_when_incomplete,
2342
2580
  not self._saw_terminal or self._mapping_failed,
@@ -2357,15 +2595,17 @@ class _InstrumentedAsyncMediaStream:
2357
2595
  self, inner: Any, handle: telemetry_dev.SpanHandle, response_mapper: ResponseMapper
2358
2596
  ) -> None:
2359
2597
  self._inner = inner
2360
- self._end = _end_once(handle)
2361
2598
  self._fields: dict[str, Any] = {}
2362
2599
  self._response_mapper = response_mapper
2363
2600
  (
2364
2601
  self._capture_output,
2365
2602
  self._mask_output_when_incomplete,
2366
2603
  ) = _chat_span_capture_policy(handle)
2367
- self._budget = _transcript_capture_budget()
2368
- self._text_parts: list[str] = []
2604
+ self._budget = _stream_budget()
2605
+ # Bound to the budget, not self, so the callback adds no reference cycle.
2606
+ self._end_stream = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
2607
+ self._ended = False
2608
+ self._text_fragments = _DeltaBuffer()
2369
2609
  self._saw_delta = False
2370
2610
  self._saw_terminal = False
2371
2611
  self._report_error = _chat_span_reporter(handle)
@@ -2428,7 +2668,15 @@ class _InstrumentedAsyncMediaStream:
2428
2668
  if failure is not None:
2429
2669
  self._end(**self._finish_fields(), error=failure)
2430
2670
 
2671
+ def _end(self, **fields: Any) -> None:
2672
+ # The span has ended and the stream budget no longer counts bytes, so later events are
2673
+ # not kept.
2674
+ self._ended = True
2675
+ self._end_stream(**fields)
2676
+
2431
2677
  def _record(self, event: Any) -> None:
2678
+ if self._ended:
2679
+ return
2432
2680
  event_type = _string(_field(event, "type")) or ""
2433
2681
  if event_type in _MEDIA_TERMINAL_EVENTS:
2434
2682
  self._saw_terminal = True
@@ -2436,6 +2684,7 @@ class _InstrumentedAsyncMediaStream:
2436
2684
  event,
2437
2685
  self._response_mapper,
2438
2686
  self._capture_output,
2687
+ self._budget,
2439
2688
  )
2440
2689
  if "output" in mapped:
2441
2690
  self._fields = _clean_fields(mapped)
@@ -2451,7 +2700,7 @@ class _InstrumentedAsyncMediaStream:
2451
2700
  if text:
2452
2701
  captured = _capture_text(text, self._budget)
2453
2702
  if captured:
2454
- self._text_parts.append(captured)
2703
+ self._text_fragments.write(captured)
2455
2704
  if event_type in _MEDIA_TERMINAL_EVENTS:
2456
2705
  error = _media_stream_error(event, event_type)
2457
2706
  self._end(**self._finish_fields(), **({"error": error} if error is not None else {}))
@@ -2460,7 +2709,7 @@ class _InstrumentedAsyncMediaStream:
2460
2709
  return _media_stream_end_fields(
2461
2710
  self._fields,
2462
2711
  self._response_mapper,
2463
- self._text_parts,
2712
+ self._text_fragments,
2464
2713
  self._budget,
2465
2714
  self._mask_output_when_incomplete,
2466
2715
  not self._saw_terminal or self._mapping_failed,