telemetry-dev-openai 0.2.2__tar.gz → 0.2.4__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.
- {telemetry_dev_openai-0.2.2 → telemetry_dev_openai-0.2.4}/PKG-INFO +4 -4
- {telemetry_dev_openai-0.2.2 → telemetry_dev_openai-0.2.4}/README.md +3 -3
- {telemetry_dev_openai-0.2.2 → telemetry_dev_openai-0.2.4}/pyproject.toml +1 -1
- {telemetry_dev_openai-0.2.2 → telemetry_dev_openai-0.2.4}/pyproject.toml.orig +1 -1
- {telemetry_dev_openai-0.2.2 → telemetry_dev_openai-0.2.4}/src/telemetry_dev_openai/__init__.py +323 -74
- {telemetry_dev_openai-0.2.2 → telemetry_dev_openai-0.2.4}/LICENSE +0 -0
- {telemetry_dev_openai-0.2.2 → telemetry_dev_openai-0.2.4}/src/telemetry_dev_openai/py.typed +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: telemetry-dev-openai
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.4
|
|
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
|
|
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`.
|
|
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
|
|
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
|
|
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`.
|
|
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
|
|
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
|
|
{telemetry_dev_openai-0.2.2 → telemetry_dev_openai-0.2.4}/src/telemetry_dev_openai/__init__.py
RENAMED
|
@@ -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.
|
|
26
|
+
__version__ = "0.2.4"
|
|
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
|
-
|
|
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
|
-
|
|
161
|
-
|
|
162
|
-
|
|
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
|
|
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
|
|
317
|
-
|
|
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
|
-
|
|
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 =
|
|
328
|
-
|
|
329
|
-
|
|
330
|
-
|
|
331
|
-
|
|
332
|
-
|
|
333
|
-
|
|
334
|
-
|
|
335
|
-
|
|
336
|
-
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
character_bytes
|
|
340
|
-
|
|
341
|
-
|
|
342
|
-
|
|
343
|
-
|
|
344
|
-
|
|
345
|
-
|
|
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 =
|
|
700
|
-
self.refusal_fragments =
|
|
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 =
|
|
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
|
|
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
|
|
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
|
|
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,
|
|
915
|
-
value_fragments = fragments if had_value else
|
|
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,
|
|
1069
|
-
argument_fragments = fragments if had_arguments else
|
|
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,
|
|
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,
|
|
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
|
-
|
|
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
|
-
|
|
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
|
|
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
|
|
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.
|
|
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
|
-
|
|
2153
|
+
budget=measured,
|
|
1964
2154
|
report_error=self._report_error,
|
|
1965
2155
|
)
|
|
1966
|
-
|
|
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.
|
|
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
|
-
|
|
2328
|
+
budget=measured,
|
|
2115
2329
|
report_error=self._report_error,
|
|
2116
2330
|
)
|
|
2117
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
2184
|
-
|
|
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 =
|
|
2251
|
-
self
|
|
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.
|
|
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.
|
|
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 =
|
|
2368
|
-
self
|
|
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.
|
|
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.
|
|
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,
|
|
File without changes
|
|
File without changes
|