telemetry-dev-openrouter 0.2.1__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.
- {telemetry_dev_openrouter-0.2.1 → telemetry_dev_openrouter-0.2.3}/PKG-INFO +7 -4
- {telemetry_dev_openrouter-0.2.1 → telemetry_dev_openrouter-0.2.3}/README.md +6 -3
- {telemetry_dev_openrouter-0.2.1 → telemetry_dev_openrouter-0.2.3}/pyproject.toml +1 -1
- {telemetry_dev_openrouter-0.2.1 → telemetry_dev_openrouter-0.2.3}/pyproject.toml.orig +1 -1
- {telemetry_dev_openrouter-0.2.1 → telemetry_dev_openrouter-0.2.3}/src/telemetry_dev_openrouter/__init__.py +295 -35
- {telemetry_dev_openrouter-0.2.1 → telemetry_dev_openrouter-0.2.3}/LICENSE +0 -0
- {telemetry_dev_openrouter-0.2.1 → telemetry_dev_openrouter-0.2.3}/src/telemetry_dev_openrouter/py.typed +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: telemetry-dev-openrouter
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.3
|
|
4
4
|
Summary: OpenRouter integration for telemetry.dev Python SDK
|
|
5
5
|
Keywords: telemetry,opentelemetry,openrouter,llm,genai,tracing
|
|
6
6
|
Author: telemetry.dev
|
|
@@ -134,9 +134,12 @@ still execute and return their normal SDK values without exporting telemetry.
|
|
|
134
134
|
## Limitations
|
|
135
135
|
|
|
136
136
|
- Unconsumed streams end their spans only when the stream is exhausted, errors, or is closed.
|
|
137
|
-
- Streaming output is
|
|
138
|
-
|
|
139
|
-
|
|
137
|
+
- Streaming output is kept up to 16 MiB, and all streams in flight share the `max_stream_bytes`
|
|
138
|
+
bound of the core (64 MiB by default); terminal metadata and usage can still be recorded after
|
|
139
|
+
the bound is reached. The configured `max_attribute_length` does not shrink this
|
|
140
|
+
retention bound; the core applies it to the exported attribute after the mask runs. These bounds
|
|
141
|
+
need a `telemetry-dev` core with `max_stream_bytes`; with an older core, each stream keeps up to
|
|
142
|
+
64 KiB and 1,024 items.
|
|
140
143
|
- A chat stream that is cut at the bound, closed early, or fails, and a Responses stream that is cut
|
|
141
144
|
at the bound, ends without a terminal response, or ends with an `error` event set
|
|
142
145
|
`telemetry.dev.capture.truncated = true`. With a mask configured, the incomplete reconstructed
|
|
@@ -111,9 +111,12 @@ still execute and return their normal SDK values without exporting telemetry.
|
|
|
111
111
|
## Limitations
|
|
112
112
|
|
|
113
113
|
- Unconsumed streams end their spans only when the stream is exhausted, errors, or is closed.
|
|
114
|
-
- Streaming output is
|
|
115
|
-
|
|
116
|
-
|
|
114
|
+
- Streaming output is kept up to 16 MiB, and all streams in flight share the `max_stream_bytes`
|
|
115
|
+
bound of the core (64 MiB by default); terminal metadata and usage can still be recorded after
|
|
116
|
+
the bound is reached. The configured `max_attribute_length` does not shrink this
|
|
117
|
+
retention bound; the core applies it to the exported attribute after the mask runs. These bounds
|
|
118
|
+
need a `telemetry-dev` core with `max_stream_bytes`; with an older core, each stream keeps up to
|
|
119
|
+
64 KiB and 1,024 items.
|
|
117
120
|
- A chat stream that is cut at the bound, closed early, or fails, and a Responses stream that is cut
|
|
118
121
|
at the bound, ends without a terminal response, or ends with an `error` event set
|
|
119
122
|
`telemetry.dev.capture.truncated = true`. With a mask configured, the incomplete reconstructed
|
|
@@ -1,10 +1,13 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
3
|
import asyncio
|
|
4
|
+
import os
|
|
4
5
|
import threading
|
|
5
6
|
import time
|
|
7
|
+
import weakref
|
|
6
8
|
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
|
|
7
|
-
from functools import wraps
|
|
9
|
+
from functools import partial, wraps
|
|
10
|
+
from io import StringIO
|
|
8
11
|
from typing import Any, TypeVar, cast
|
|
9
12
|
|
|
10
13
|
import telemetry_dev
|
|
@@ -12,7 +15,7 @@ from openrouter.chat import Chat
|
|
|
12
15
|
from openrouter.embeddings import Embeddings
|
|
13
16
|
from openrouter.responses import Responses
|
|
14
17
|
|
|
15
|
-
__version__ = "0.2.
|
|
18
|
+
__version__ = "0.2.3"
|
|
16
19
|
|
|
17
20
|
RequestMapper = Callable[[Mapping[str, Any]], tuple[str, dict[str, Any]]]
|
|
18
21
|
ResponseMapper = Callable[[Any], dict[str, Any]]
|
|
@@ -21,10 +24,8 @@ _PROVIDER = "openrouter"
|
|
|
21
24
|
_WRAPPED_ATTR = "_telemetry_dev_openrouter_wrapped"
|
|
22
25
|
_ORIGINAL_ATTR = "_telemetry_dev_openrouter_original"
|
|
23
26
|
_CAPTURE_TRUNCATED_ATTRIBUTE = "telemetry.dev.capture.truncated"
|
|
24
|
-
#
|
|
25
|
-
|
|
26
|
-
_STREAM_CAPTURE_MAX_BYTES = 64 * 1024
|
|
27
|
-
_STREAM_CAPTURE_MAX_ITEMS = 1024
|
|
27
|
+
# The core applies max_attribute_length after the mask. More choices than this are not kept.
|
|
28
|
+
_MAX_STREAM_CHOICES = 1024
|
|
28
29
|
_FINISH_REASON_MAX_LENGTH = 128
|
|
29
30
|
_TERMINAL_RESPONSE_EVENTS = frozenset(
|
|
30
31
|
{"response.completed", "response.failed", "response.incomplete"}
|
|
@@ -33,6 +34,22 @@ _OMIT = object()
|
|
|
33
34
|
_ORIGINALS: list[tuple[type[Any], str, Any]] = []
|
|
34
35
|
_installed = False
|
|
35
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
|
+
|
|
36
53
|
_T = TypeVar("_T")
|
|
37
54
|
|
|
38
55
|
|
|
@@ -174,6 +191,12 @@ def _cost_usd(raw_usage: Any) -> float | None:
|
|
|
174
191
|
return cost if cost is not None else upstream_cost
|
|
175
192
|
|
|
176
193
|
|
|
194
|
+
def _prompt_attributes(messages: object) -> dict[str, Any]:
|
|
195
|
+
# Older cores have no prompt_attributes_for and record a prompt only inside with_prompt.
|
|
196
|
+
read = getattr(telemetry_dev, "prompt_attributes_for", None)
|
|
197
|
+
return cast("dict[str, Any]", read(messages)) if callable(read) else {}
|
|
198
|
+
|
|
199
|
+
|
|
177
200
|
def _chat_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
|
|
178
201
|
model = _string(params.get("model"))
|
|
179
202
|
max_completion_tokens = _number(params.get("max_completion_tokens"))
|
|
@@ -193,6 +216,7 @@ def _chat_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
|
|
|
193
216
|
"seed": _number(params.get("seed")),
|
|
194
217
|
"frequency_penalty": _number(params.get("frequency_penalty")),
|
|
195
218
|
"presence_penalty": _number(params.get("presence_penalty")),
|
|
219
|
+
"attributes": _prompt_attributes(params.get("messages")) or None,
|
|
196
220
|
},
|
|
197
221
|
)
|
|
198
222
|
|
|
@@ -260,6 +284,7 @@ def _responses_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
|
|
|
260
284
|
"max_tokens": _number(params.get("max_output_tokens")),
|
|
261
285
|
"frequency_penalty": _number(params.get("frequency_penalty")),
|
|
262
286
|
"presence_penalty": _number(params.get("presence_penalty")),
|
|
287
|
+
"attributes": _prompt_attributes(params.get("input")) or None,
|
|
263
288
|
},
|
|
264
289
|
)
|
|
265
290
|
|
|
@@ -382,22 +407,107 @@ def _end_once(handle: telemetry_dev.SpanHandle) -> Callable[..., None]:
|
|
|
382
407
|
return end
|
|
383
408
|
|
|
384
409
|
|
|
410
|
+
# On Python 3.10 and 3.11, StringIO keeps up to 100,000 written strings in a list before it
|
|
411
|
+
# joins them, about 60 bytes each that the stream bound does not charge: 15 bytes a character
|
|
412
|
+
# for 4-character deltas. Deltas are joined 256 at a time, so that those waiting hold under
|
|
413
|
+
# about 22 KB, onto a string that is written once it has 4,096 characters, so that what is
|
|
414
|
+
# written costs about 1.5% more than its text. Python 3.12 dropped the list.
|
|
415
|
+
_DELTA_BATCH = 256
|
|
416
|
+
_WRITE_CHARS = 4096
|
|
417
|
+
|
|
418
|
+
|
|
419
|
+
class _DeltaBuffer:
|
|
420
|
+
# A write takes several steps, unlike one StringIO call, so without the lock a read on
|
|
421
|
+
# another thread, such as one that closes the stream, could copy a batch twice or drop it.
|
|
422
|
+
# A fork copies the lock held if another thread was writing, and that thread does not run
|
|
423
|
+
# in the child to release it. Every fork, also one that runs no at-fork hooks, changes the
|
|
424
|
+
# process ID, so a child takes a new lock. The check is inline because it runs per delta.
|
|
425
|
+
# A signal handler runs on the thread it interrupts, so one that reads the stream, such as
|
|
426
|
+
# one that closes it, can read inside a write. So the lock is reentrant, and each step takes
|
|
427
|
+
# the deltas and the tail out of the buffer before it joins or writes them, so that a read
|
|
428
|
+
# inside the step does not write them as well; that read returns a prefix, which can leave
|
|
429
|
+
# out the deltas being joined. A handler that writes inside a step is not supported. An
|
|
430
|
+
# exception inside a step, such as KeyboardInterrupt from a handler, puts back the text the
|
|
431
|
+
# step took unless the StringIO has it, so that a later read still has every delta.
|
|
432
|
+
__slots__ = ("_buffer", "_lock", "_pending", "_pid", "_tail")
|
|
433
|
+
|
|
434
|
+
def __init__(self) -> None:
|
|
435
|
+
self._buffer = StringIO()
|
|
436
|
+
self._lock = threading.RLock()
|
|
437
|
+
self._pid = os.getpid()
|
|
438
|
+
self._pending: list[str] = []
|
|
439
|
+
self._tail = ""
|
|
440
|
+
|
|
441
|
+
def write(self, text: str) -> None:
|
|
442
|
+
pid = os.getpid()
|
|
443
|
+
if self._pid != pid:
|
|
444
|
+
self._lock = threading.RLock()
|
|
445
|
+
self._pid = pid
|
|
446
|
+
with self._lock:
|
|
447
|
+
self._pending.append(text)
|
|
448
|
+
if len(self._pending) >= _DELTA_BATCH:
|
|
449
|
+
self._join(_WRITE_CHARS)
|
|
450
|
+
|
|
451
|
+
def getvalue(self) -> str:
|
|
452
|
+
pid = os.getpid()
|
|
453
|
+
if self._pid != pid:
|
|
454
|
+
self._lock = threading.RLock()
|
|
455
|
+
self._pid = pid
|
|
456
|
+
with self._lock:
|
|
457
|
+
self._join(1)
|
|
458
|
+
return self._buffer.getvalue()
|
|
459
|
+
|
|
460
|
+
def _join(self, write_at: int) -> None:
|
|
461
|
+
if not self._pending and len(self._tail) < write_at:
|
|
462
|
+
return
|
|
463
|
+
written = self._buffer.tell()
|
|
464
|
+
pending: list[str] = []
|
|
465
|
+
tail = ""
|
|
466
|
+
try:
|
|
467
|
+
pending, self._pending = self._pending, []
|
|
468
|
+
tail, self._tail = self._tail, ""
|
|
469
|
+
joined = "".join([tail, *pending])
|
|
470
|
+
if len(joined) >= write_at:
|
|
471
|
+
self._buffer.write(joined)
|
|
472
|
+
else:
|
|
473
|
+
self._tail = joined
|
|
474
|
+
except BaseException:
|
|
475
|
+
if self._buffer.tell() == written:
|
|
476
|
+
self._pending[:0] = pending
|
|
477
|
+
self._tail = tail + self._tail
|
|
478
|
+
raise
|
|
479
|
+
|
|
480
|
+
|
|
385
481
|
class _ChatChoice:
|
|
386
482
|
def __init__(self) -> None:
|
|
387
483
|
self.role: str | None = None
|
|
388
|
-
|
|
389
|
-
|
|
484
|
+
# Deltas are written to buffers and joined when output is read: repeated concatenation
|
|
485
|
+
# copies the whole text for each delta.
|
|
486
|
+
self.content_fragments = _DeltaBuffer()
|
|
487
|
+
self.reasoning_fragments = _DeltaBuffer()
|
|
390
488
|
self.reasoning_details: list[Any] = []
|
|
391
|
-
self.
|
|
489
|
+
self.refusal_fragments = _DeltaBuffer()
|
|
392
490
|
self.tool_calls: dict[int, dict[str, Any]] = {}
|
|
393
491
|
self.finish_reason: str | None = None
|
|
394
492
|
|
|
493
|
+
@property
|
|
494
|
+
def content(self) -> str:
|
|
495
|
+
return self.content_fragments.getvalue()
|
|
496
|
+
|
|
497
|
+
@property
|
|
498
|
+
def reasoning(self) -> str:
|
|
499
|
+
return self.reasoning_fragments.getvalue()
|
|
500
|
+
|
|
501
|
+
@property
|
|
502
|
+
def refusal(self) -> str:
|
|
503
|
+
return self.refusal_fragments.getvalue()
|
|
504
|
+
|
|
395
505
|
|
|
396
506
|
def _choice_state(
|
|
397
507
|
states: dict[int, _ChatChoice], index: int, budget: telemetry_dev.CaptureBudget
|
|
398
508
|
) -> _ChatChoice | None:
|
|
399
509
|
if index not in states:
|
|
400
|
-
if len(states) >=
|
|
510
|
+
if len(states) >= _MAX_STREAM_CHOICES:
|
|
401
511
|
budget.truncated = True
|
|
402
512
|
return None
|
|
403
513
|
states[index] = _ChatChoice()
|
|
@@ -422,7 +532,10 @@ def _merge_tool_call(state: _ChatChoice, delta: Any) -> None:
|
|
|
422
532
|
if name is not None:
|
|
423
533
|
current_function["name"] = name
|
|
424
534
|
if isinstance(arguments, str):
|
|
425
|
-
|
|
535
|
+
fragments = current_function.get("arguments")
|
|
536
|
+
if not isinstance(fragments, _DeltaBuffer):
|
|
537
|
+
fragments = current_function["arguments"] = _DeltaBuffer()
|
|
538
|
+
fragments.write(arguments)
|
|
426
539
|
current["function"] = current_function
|
|
427
540
|
state.tool_calls[tool_index] = current
|
|
428
541
|
|
|
@@ -443,15 +556,15 @@ def _record_chat_chunk(
|
|
|
443
556
|
if isinstance(role, str):
|
|
444
557
|
state.role = role
|
|
445
558
|
if isinstance(content, str) and budget.accept(content):
|
|
446
|
-
state.content
|
|
559
|
+
state.content_fragments.write(content)
|
|
447
560
|
if isinstance(reasoning, str) and budget.accept(reasoning):
|
|
448
|
-
state.reasoning
|
|
561
|
+
state.reasoning_fragments.write(reasoning)
|
|
449
562
|
for detail in _sequence_items(_field(delta, "reasoning_details")):
|
|
450
563
|
native_detail = _native(detail)
|
|
451
564
|
if native_detail is not None and budget.accept(native_detail):
|
|
452
565
|
state.reasoning_details.append(native_detail)
|
|
453
566
|
if isinstance(refusal, str) and budget.accept(refusal):
|
|
454
|
-
state.refusal
|
|
567
|
+
state.refusal_fragments.write(refusal)
|
|
455
568
|
for tool_call in _sequence_items(_field(delta, "tool_calls")):
|
|
456
569
|
if budget.accept(tool_call):
|
|
457
570
|
_merge_tool_call(state, tool_call)
|
|
@@ -520,6 +633,16 @@ def _response_event_has_output(event: Any) -> bool:
|
|
|
520
633
|
)
|
|
521
634
|
|
|
522
635
|
|
|
636
|
+
def _joined_tool_call(call: dict[str, Any]) -> dict[str, Any]:
|
|
637
|
+
function = call.get("function")
|
|
638
|
+
if not isinstance(function, dict):
|
|
639
|
+
return call
|
|
640
|
+
arguments = cast(dict[str, Any], function).get("arguments")
|
|
641
|
+
if not isinstance(arguments, _DeltaBuffer):
|
|
642
|
+
return call
|
|
643
|
+
return {**call, "function": {**function, "arguments": arguments.getvalue()}}
|
|
644
|
+
|
|
645
|
+
|
|
523
646
|
def _chat_output(states: Mapping[int, _ChatChoice]) -> list[dict[str, Any]]:
|
|
524
647
|
output: list[dict[str, Any]] = []
|
|
525
648
|
for _, state in sorted(states.items()):
|
|
@@ -535,13 +658,57 @@ def _chat_output(states: Mapping[int, _ChatChoice]) -> list[dict[str, Any]]:
|
|
|
535
658
|
if state.refusal:
|
|
536
659
|
message["refusal"] = state.refusal
|
|
537
660
|
if state.tool_calls:
|
|
538
|
-
message["tool_calls"] = [
|
|
661
|
+
message["tool_calls"] = [
|
|
662
|
+
_joined_tool_call(call) for _, call in sorted(state.tool_calls.items())
|
|
663
|
+
]
|
|
539
664
|
output.append(message)
|
|
540
665
|
return output
|
|
541
666
|
|
|
542
667
|
|
|
543
668
|
def _stream_capture_budget() -> telemetry_dev.CaptureBudget:
|
|
544
|
-
|
|
669
|
+
# Older cores have no for_stream and bound each budget at 64 KiB on its own.
|
|
670
|
+
for_stream = getattr(telemetry_dev.CaptureBudget, "for_stream", None)
|
|
671
|
+
budget = for_stream() if callable(for_stream) else telemetry_dev.CaptureBudget()
|
|
672
|
+
return cast(telemetry_dev.CaptureBudget, budget)
|
|
673
|
+
|
|
674
|
+
|
|
675
|
+
def _release_budget(budget: telemetry_dev.CaptureBudget) -> None:
|
|
676
|
+
release = getattr(budget, "release", None)
|
|
677
|
+
if callable(release):
|
|
678
|
+
release()
|
|
679
|
+
|
|
680
|
+
|
|
681
|
+
def _measuring_budget(stream: telemetry_dev.CaptureBudget) -> telemetry_dev.CaptureBudget:
|
|
682
|
+
"""An empty budget that measures a value replacing everything ``stream`` holds, within the
|
|
683
|
+
bound ``stream`` has now, without holding bytes itself."""
|
|
684
|
+
return telemetry_dev.CaptureBudget(stream.max_bytes, stream.max_items)
|
|
685
|
+
|
|
686
|
+
|
|
687
|
+
def _adopt_budget(
|
|
688
|
+
stream: telemetry_dev.CaptureBudget, measured: telemetry_dev.CaptureBudget
|
|
689
|
+
) -> bool:
|
|
690
|
+
"""Takes the counts of a budget that measured a replacement into the stream's own budget.
|
|
691
|
+
Returns False, and takes nothing, when another stream took the room after the measure."""
|
|
692
|
+
# Older cores have no reserve_bytes; there the measure against max_bytes stands.
|
|
693
|
+
reserve = getattr(stream, "reserve_bytes", None)
|
|
694
|
+
if callable(reserve):
|
|
695
|
+
if not reserve(measured.bytes_used):
|
|
696
|
+
return False
|
|
697
|
+
else:
|
|
698
|
+
stream.bytes_used = measured.bytes_used
|
|
699
|
+
stream.items_used = measured.items_used
|
|
700
|
+
stream.truncated = measured.truncated
|
|
701
|
+
return True
|
|
702
|
+
|
|
703
|
+
|
|
704
|
+
def _ending_stream(end: Callable[..., None], release: Callable[[], None]) -> Callable[..., None]:
|
|
705
|
+
def end_stream(**fields: Any) -> None:
|
|
706
|
+
try:
|
|
707
|
+
end(**fields)
|
|
708
|
+
finally:
|
|
709
|
+
release()
|
|
710
|
+
|
|
711
|
+
return end_stream
|
|
545
712
|
|
|
546
713
|
|
|
547
714
|
def _capture_masked(handle: telemetry_dev.SpanHandle) -> bool:
|
|
@@ -611,6 +778,15 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
|
|
|
611
778
|
response = getattr(inner, "response", None)
|
|
612
779
|
if response is None:
|
|
613
780
|
return
|
|
781
|
+
# ``finish`` is a method of the stream wrapper, which holds ``response``: held weakly, the hook
|
|
782
|
+
# adds no reference cycle, so a dropped wrapper and its stream budget are freed at once.
|
|
783
|
+
finish_method = weakref.WeakMethod(finish)
|
|
784
|
+
|
|
785
|
+
def finish_if_alive() -> None:
|
|
786
|
+
method = finish_method()
|
|
787
|
+
if method is not None:
|
|
788
|
+
method()
|
|
789
|
+
|
|
614
790
|
close = getattr(response, "close", None)
|
|
615
791
|
if callable(close):
|
|
616
792
|
|
|
@@ -618,7 +794,7 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
|
|
|
618
794
|
try:
|
|
619
795
|
return close(*args, **kwargs)
|
|
620
796
|
finally:
|
|
621
|
-
|
|
797
|
+
finish_if_alive()
|
|
622
798
|
|
|
623
799
|
response.close = _close_hook
|
|
624
800
|
aclose = getattr(response, "aclose", None)
|
|
@@ -629,7 +805,7 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
|
|
|
629
805
|
try:
|
|
630
806
|
return await aclose_fn(*args, **kwargs)
|
|
631
807
|
finally:
|
|
632
|
-
|
|
808
|
+
finish_if_alive()
|
|
633
809
|
|
|
634
810
|
response.aclose = _aclose_hook
|
|
635
811
|
|
|
@@ -657,7 +833,6 @@ async def _close_async_stream(inner: Any) -> None:
|
|
|
657
833
|
class _InstrumentedStream:
|
|
658
834
|
def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
|
|
659
835
|
self._inner = inner
|
|
660
|
-
self._end = _end_once(handle)
|
|
661
836
|
self._handle = handle
|
|
662
837
|
self._started_at = started_at
|
|
663
838
|
self._states: dict[int, _ChatChoice] = {}
|
|
@@ -669,6 +844,9 @@ class _InstrumentedStream:
|
|
|
669
844
|
self._completed = False
|
|
670
845
|
self._mask_output_when_incomplete = _capture_masked(handle)
|
|
671
846
|
self._budget = _stream_capture_budget()
|
|
847
|
+
# Bound to the budget, not self, so the callback adds no reference cycle.
|
|
848
|
+
self._end_stream = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
|
|
849
|
+
self._ended = False
|
|
672
850
|
_hook_response_close(inner, self._on_response_close)
|
|
673
851
|
|
|
674
852
|
def _iterate(self) -> Iterator[Any]:
|
|
@@ -739,7 +917,15 @@ class _InstrumentedStream:
|
|
|
739
917
|
def __getattr__(self, name: str) -> Any:
|
|
740
918
|
return getattr(self._inner, name)
|
|
741
919
|
|
|
920
|
+
def _end(self, **fields: Any) -> None:
|
|
921
|
+
# The span has ended and the stream budget no longer counts bytes, so later events are
|
|
922
|
+
# not kept.
|
|
923
|
+
self._ended = True
|
|
924
|
+
self._end_stream(**fields)
|
|
925
|
+
|
|
742
926
|
def _record(self, chunk: Any, received_at: float) -> bool:
|
|
927
|
+
if self._ended:
|
|
928
|
+
return True
|
|
743
929
|
update = _record_chat_chunk(chunk, self._states, self._budget)
|
|
744
930
|
if _chat_chunk_has_output(chunk):
|
|
745
931
|
record_output_chunk = getattr(self._handle, "record_output_chunk", None)
|
|
@@ -762,7 +948,6 @@ class _InstrumentedStream:
|
|
|
762
948
|
class _InstrumentedAsyncStream:
|
|
763
949
|
def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
|
|
764
950
|
self._inner = inner
|
|
765
|
-
self._end = _end_once(handle)
|
|
766
951
|
self._handle = handle
|
|
767
952
|
self._started_at = started_at
|
|
768
953
|
self._states: dict[int, _ChatChoice] = {}
|
|
@@ -774,6 +959,9 @@ class _InstrumentedAsyncStream:
|
|
|
774
959
|
self._completed = False
|
|
775
960
|
self._mask_output_when_incomplete = _capture_masked(handle)
|
|
776
961
|
self._budget = _stream_capture_budget()
|
|
962
|
+
# Bound to the budget, not self, so the callback adds no reference cycle.
|
|
963
|
+
self._end_stream = _ending_stream(_end_once(handle), partial(_release_budget, self._budget))
|
|
964
|
+
self._ended = False
|
|
777
965
|
_hook_response_close(inner, self._on_response_close)
|
|
778
966
|
|
|
779
967
|
async def _aiterate(self) -> AsyncIterator[Any]:
|
|
@@ -847,7 +1035,15 @@ class _InstrumentedAsyncStream:
|
|
|
847
1035
|
def __getattr__(self, name: str) -> Any:
|
|
848
1036
|
return getattr(self._inner, name)
|
|
849
1037
|
|
|
1038
|
+
def _end(self, **fields: Any) -> None:
|
|
1039
|
+
# The span has ended and the stream budget no longer counts bytes, so later events are
|
|
1040
|
+
# not kept.
|
|
1041
|
+
self._ended = True
|
|
1042
|
+
self._end_stream(**fields)
|
|
1043
|
+
|
|
850
1044
|
def _record(self, chunk: Any, received_at: float) -> bool:
|
|
1045
|
+
if self._ended:
|
|
1046
|
+
return True
|
|
851
1047
|
update = _record_chat_chunk(chunk, self._states, self._budget)
|
|
852
1048
|
if _chat_chunk_has_output(chunk):
|
|
853
1049
|
record_output_chunk = getattr(self._handle, "record_output_chunk", None)
|
|
@@ -867,6 +1063,37 @@ class _InstrumentedAsyncStream:
|
|
|
867
1063
|
return _chunk_error(chunk) is None
|
|
868
1064
|
|
|
869
1065
|
|
|
1066
|
+
def _write_fragment(target: dict[str, Any], key: str, value: str) -> None:
|
|
1067
|
+
"""Appends a delta to ``target[key]`` in a buffer, since repeated concatenation copies the
|
|
1068
|
+
whole text for each delta. _joined_item reads the text back."""
|
|
1069
|
+
fragments = target.get(key)
|
|
1070
|
+
if not isinstance(fragments, _DeltaBuffer):
|
|
1071
|
+
fragments = _DeltaBuffer()
|
|
1072
|
+
fragments.write(f"{target.get(key, '')}")
|
|
1073
|
+
target[key] = fragments
|
|
1074
|
+
fragments.write(value)
|
|
1075
|
+
|
|
1076
|
+
|
|
1077
|
+
def _joined_value(value: Any) -> Any:
|
|
1078
|
+
return value.getvalue() if isinstance(value, _DeltaBuffer) else value
|
|
1079
|
+
|
|
1080
|
+
|
|
1081
|
+
def _joined_item(item: Any) -> Any:
|
|
1082
|
+
if not isinstance(item, dict):
|
|
1083
|
+
return item
|
|
1084
|
+
joined: dict[str, Any] = {}
|
|
1085
|
+
for key, value in cast(dict[str, Any], item).items():
|
|
1086
|
+
if key in ("content", "summary") and isinstance(value, list):
|
|
1087
|
+
value = [
|
|
1088
|
+
{name: _joined_value(field) for name, field in cast(dict[str, Any], part).items()}
|
|
1089
|
+
if isinstance(part, dict)
|
|
1090
|
+
else part
|
|
1091
|
+
for part in cast(list[Any], value)
|
|
1092
|
+
]
|
|
1093
|
+
joined[key] = _joined_value(value)
|
|
1094
|
+
return joined
|
|
1095
|
+
|
|
1096
|
+
|
|
870
1097
|
class _ResponsesStreamState:
|
|
871
1098
|
def __init__(self) -> None:
|
|
872
1099
|
self.partial: dict[str, Any] = {}
|
|
@@ -895,7 +1122,10 @@ class _ResponsesStreamState:
|
|
|
895
1122
|
|
|
896
1123
|
def _output_with_item(self, output_index: int, item: dict[str, Any]) -> list[dict[str, Any]]:
|
|
897
1124
|
items = {**self._items, output_index: item}
|
|
898
|
-
return [
|
|
1125
|
+
return [
|
|
1126
|
+
*[_joined_item(value) for _, value in sorted(items.items())],
|
|
1127
|
+
*self._synthetic_events,
|
|
1128
|
+
]
|
|
899
1129
|
|
|
900
1130
|
@staticmethod
|
|
901
1131
|
def _can_accept(budget: telemetry_dev.CaptureBudget, value: Any) -> bool:
|
|
@@ -906,14 +1136,19 @@ class _ResponsesStreamState:
|
|
|
906
1136
|
return candidate.accept(value)
|
|
907
1137
|
|
|
908
1138
|
def _replace_output(self, output: Any) -> bool:
|
|
909
|
-
replacement =
|
|
1139
|
+
replacement = _measuring_budget(self._budget)
|
|
910
1140
|
if not replacement.accept(output):
|
|
911
1141
|
self._budget.truncated = True
|
|
912
1142
|
return False
|
|
913
1143
|
replacement.truncated = self._provider_events_truncated
|
|
914
|
-
self._budget
|
|
1144
|
+
if not _adopt_budget(self._budget, replacement):
|
|
1145
|
+
self._budget.truncated = True
|
|
1146
|
+
return False
|
|
915
1147
|
return True
|
|
916
1148
|
|
|
1149
|
+
def release(self) -> None:
|
|
1150
|
+
_release_budget(self._budget)
|
|
1151
|
+
|
|
917
1152
|
def _sync_output(self) -> None:
|
|
918
1153
|
self.retained_output = self._output()
|
|
919
1154
|
self.partial["output"] = self.retained_output
|
|
@@ -999,7 +1234,7 @@ class _ResponsesStreamState:
|
|
|
999
1234
|
if event_type in {"response.output_item.added", "response.output_item.done"}:
|
|
1000
1235
|
raw_item = _field(event, "item")
|
|
1001
1236
|
replacement = event_type.endswith(".done")
|
|
1002
|
-
preflight =
|
|
1237
|
+
preflight = _measuring_budget(self._budget) if replacement else self._budget
|
|
1003
1238
|
if not self._can_accept(preflight, raw_item):
|
|
1004
1239
|
self._budget.truncated = True
|
|
1005
1240
|
return False
|
|
@@ -1017,7 +1252,7 @@ class _ResponsesStreamState:
|
|
|
1017
1252
|
if event_type in {"response.content_part.added", "response.content_part.done"}:
|
|
1018
1253
|
event_part = _field(event, "part")
|
|
1019
1254
|
replacement = event_type.endswith(".done")
|
|
1020
|
-
preflight =
|
|
1255
|
+
preflight = _measuring_budget(self._budget) if replacement else self._budget
|
|
1021
1256
|
if not self._can_accept(preflight, event_part):
|
|
1022
1257
|
self._budget.truncated = True
|
|
1023
1258
|
return False
|
|
@@ -1119,7 +1354,7 @@ class _ResponsesStreamState:
|
|
|
1119
1354
|
else:
|
|
1120
1355
|
if not self._budget.accept(value):
|
|
1121
1356
|
return False
|
|
1122
|
-
item
|
|
1357
|
+
_write_fragment(item, field, value)
|
|
1123
1358
|
name = _field(event, "name")
|
|
1124
1359
|
if name is not None:
|
|
1125
1360
|
item["name"] = name
|
|
@@ -1130,7 +1365,7 @@ class _ResponsesStreamState:
|
|
|
1130
1365
|
}:
|
|
1131
1366
|
event_part = _field(event, "part")
|
|
1132
1367
|
replacement = event_type.endswith(".done")
|
|
1133
|
-
preflight =
|
|
1368
|
+
preflight = _measuring_budget(self._budget) if replacement else self._budget
|
|
1134
1369
|
if not self._can_accept(preflight, event_part):
|
|
1135
1370
|
self._budget.truncated = True
|
|
1136
1371
|
return False
|
|
@@ -1245,7 +1480,7 @@ class _ResponsesStreamState:
|
|
|
1245
1480
|
self._summary[summary_key] = part
|
|
1246
1481
|
self._sync_item(output_index, item)
|
|
1247
1482
|
else:
|
|
1248
|
-
summary
|
|
1483
|
+
_write_fragment(summary, "text", value)
|
|
1249
1484
|
return True
|
|
1250
1485
|
content_index = self._event_index(event, "content_index", "contentIndex")
|
|
1251
1486
|
content_key = (output_index, content_index)
|
|
@@ -1284,9 +1519,9 @@ class _ResponsesStreamState:
|
|
|
1284
1519
|
self._content[content_key] = replacement
|
|
1285
1520
|
self._sync_item(output_index, item)
|
|
1286
1521
|
elif refusal:
|
|
1287
|
-
part
|
|
1522
|
+
_write_fragment(part, "refusal", value)
|
|
1288
1523
|
else:
|
|
1289
|
-
part
|
|
1524
|
+
_write_fragment(part, "text", value)
|
|
1290
1525
|
return True
|
|
1291
1526
|
|
|
1292
1527
|
def record(self, event: Any) -> None:
|
|
@@ -1294,7 +1529,7 @@ class _ResponsesStreamState:
|
|
|
1294
1529
|
response = _field(event, "response")
|
|
1295
1530
|
if response is not None:
|
|
1296
1531
|
fields = _responses_response(response, include_error=False, include_output=False)
|
|
1297
|
-
response_budget =
|
|
1532
|
+
response_budget = _measuring_budget(self._budget)
|
|
1298
1533
|
raw_output = _field(response, "output")
|
|
1299
1534
|
if raw_output is not None:
|
|
1300
1535
|
if not self._can_accept(response_budget, raw_output):
|
|
@@ -1306,7 +1541,11 @@ class _ResponsesStreamState:
|
|
|
1306
1541
|
self._budget.truncated = True
|
|
1307
1542
|
else:
|
|
1308
1543
|
response_artifact = [*response_items, *self._synthetic_events]
|
|
1544
|
+
adopted = False
|
|
1309
1545
|
if response_budget.accept(response_artifact):
|
|
1546
|
+
response_budget.truncated = self._provider_events_truncated
|
|
1547
|
+
adopted = _adopt_budget(self._budget, response_budget)
|
|
1548
|
+
if adopted:
|
|
1310
1549
|
self._items.clear()
|
|
1311
1550
|
self._content.clear()
|
|
1312
1551
|
self._summary.clear()
|
|
@@ -1314,8 +1553,6 @@ class _ResponsesStreamState:
|
|
|
1314
1553
|
if isinstance(item, dict):
|
|
1315
1554
|
self._hydrate_item(index, cast(dict[str, Any], item))
|
|
1316
1555
|
self.retained_output = response_artifact
|
|
1317
|
-
response_budget.truncated = self._provider_events_truncated
|
|
1318
|
-
self._budget = response_budget
|
|
1319
1556
|
else:
|
|
1320
1557
|
self._budget.truncated = True
|
|
1321
1558
|
self.partial = fields
|
|
@@ -1329,6 +1566,9 @@ class _ResponsesStreamState:
|
|
|
1329
1566
|
|
|
1330
1567
|
def fields(self, mask_output: bool, *, failed: bool = False) -> dict[str, Any]:
|
|
1331
1568
|
fields = dict(self.partial)
|
|
1569
|
+
output = fields.get("output")
|
|
1570
|
+
if isinstance(output, list):
|
|
1571
|
+
fields["output"] = [_joined_item(item) for item in cast(list[Any], output)]
|
|
1332
1572
|
if failed or self._budget.truncated or not self.saw_terminal_snapshot:
|
|
1333
1573
|
_mark_capture_incomplete(fields, mask_output)
|
|
1334
1574
|
return fields
|
|
@@ -1338,10 +1578,12 @@ class _InstrumentedResponsesStream:
|
|
|
1338
1578
|
def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
|
|
1339
1579
|
self._inner = inner
|
|
1340
1580
|
self._handle = handle
|
|
1341
|
-
self._end = _end_once(handle)
|
|
1342
1581
|
self._started_at = started_at
|
|
1343
1582
|
self._saw_first = False
|
|
1344
1583
|
self._state = _ResponsesStreamState()
|
|
1584
|
+
# Bound to the state, not self, so the callback adds no reference cycle.
|
|
1585
|
+
self._end_stream = _ending_stream(_end_once(handle), self._state.release)
|
|
1586
|
+
self._ended = False
|
|
1345
1587
|
self._mask_output_when_incomplete = _capture_masked(handle)
|
|
1346
1588
|
self._consume: Iterator[Any] | None = None
|
|
1347
1589
|
self._in_next = False
|
|
@@ -1409,7 +1651,15 @@ class _InstrumentedResponsesStream:
|
|
|
1409
1651
|
def __getattr__(self, name: str) -> Any:
|
|
1410
1652
|
return getattr(self._inner, name)
|
|
1411
1653
|
|
|
1654
|
+
def _end(self, **fields: Any) -> None:
|
|
1655
|
+
# The span has ended and the stream budget no longer counts bytes, so later events are
|
|
1656
|
+
# not kept.
|
|
1657
|
+
self._ended = True
|
|
1658
|
+
self._end_stream(**fields)
|
|
1659
|
+
|
|
1412
1660
|
def _record(self, event: Any, received_at: float) -> None:
|
|
1661
|
+
if self._ended:
|
|
1662
|
+
return
|
|
1413
1663
|
if _response_event_has_output(event):
|
|
1414
1664
|
record_output_chunk = getattr(self._handle, "record_output_chunk", None)
|
|
1415
1665
|
if callable(record_output_chunk):
|
|
@@ -1436,10 +1686,12 @@ class _InstrumentedAsyncResponsesStream:
|
|
|
1436
1686
|
def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
|
|
1437
1687
|
self._inner = inner
|
|
1438
1688
|
self._handle = handle
|
|
1439
|
-
self._end = _end_once(handle)
|
|
1440
1689
|
self._started_at = started_at
|
|
1441
1690
|
self._saw_first = False
|
|
1442
1691
|
self._state = _ResponsesStreamState()
|
|
1692
|
+
# Bound to the state, not self, so the callback adds no reference cycle.
|
|
1693
|
+
self._end_stream = _ending_stream(_end_once(handle), self._state.release)
|
|
1694
|
+
self._ended = False
|
|
1443
1695
|
self._mask_output_when_incomplete = _capture_masked(handle)
|
|
1444
1696
|
self._consume: AsyncIterator[Any] | None = None
|
|
1445
1697
|
self._in_next = False
|
|
@@ -1510,7 +1762,15 @@ class _InstrumentedAsyncResponsesStream:
|
|
|
1510
1762
|
def __getattr__(self, name: str) -> Any:
|
|
1511
1763
|
return getattr(self._inner, name)
|
|
1512
1764
|
|
|
1765
|
+
def _end(self, **fields: Any) -> None:
|
|
1766
|
+
# The span has ended and the stream budget no longer counts bytes, so later events are
|
|
1767
|
+
# not kept.
|
|
1768
|
+
self._ended = True
|
|
1769
|
+
self._end_stream(**fields)
|
|
1770
|
+
|
|
1513
1771
|
def _record(self, event: Any, received_at: float) -> None:
|
|
1772
|
+
if self._ended:
|
|
1773
|
+
return
|
|
1514
1774
|
if _response_event_has_output(event):
|
|
1515
1775
|
record_output_chunk = getattr(self._handle, "record_output_chunk", None)
|
|
1516
1776
|
if callable(record_output_chunk):
|
|
File without changes
|
|
File without changes
|