telemetry-dev-openrouter 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_openrouter-0.2.2 → telemetry_dev_openrouter-0.2.4}/PKG-INFO +7 -4
- {telemetry_dev_openrouter-0.2.2 → telemetry_dev_openrouter-0.2.4}/README.md +6 -3
- {telemetry_dev_openrouter-0.2.2 → telemetry_dev_openrouter-0.2.4}/pyproject.toml +1 -1
- {telemetry_dev_openrouter-0.2.2 → telemetry_dev_openrouter-0.2.4}/pyproject.toml.orig +1 -1
- {telemetry_dev_openrouter-0.2.2 → telemetry_dev_openrouter-0.2.4}/src/telemetry_dev_openrouter/__init__.py +287 -35
- {telemetry_dev_openrouter-0.2.2 → telemetry_dev_openrouter-0.2.4}/LICENSE +0 -0
- {telemetry_dev_openrouter-0.2.2 → telemetry_dev_openrouter-0.2.4}/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.4
|
|
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.4"
|
|
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
|
|
|
@@ -390,22 +407,107 @@ def _end_once(handle: telemetry_dev.SpanHandle) -> Callable[..., None]:
|
|
|
390
407
|
return end
|
|
391
408
|
|
|
392
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
|
+
|
|
393
481
|
class _ChatChoice:
|
|
394
482
|
def __init__(self) -> None:
|
|
395
483
|
self.role: str | None = None
|
|
396
|
-
|
|
397
|
-
|
|
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()
|
|
398
488
|
self.reasoning_details: list[Any] = []
|
|
399
|
-
self.
|
|
489
|
+
self.refusal_fragments = _DeltaBuffer()
|
|
400
490
|
self.tool_calls: dict[int, dict[str, Any]] = {}
|
|
401
491
|
self.finish_reason: str | None = None
|
|
402
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
|
+
|
|
403
505
|
|
|
404
506
|
def _choice_state(
|
|
405
507
|
states: dict[int, _ChatChoice], index: int, budget: telemetry_dev.CaptureBudget
|
|
406
508
|
) -> _ChatChoice | None:
|
|
407
509
|
if index not in states:
|
|
408
|
-
if len(states) >=
|
|
510
|
+
if len(states) >= _MAX_STREAM_CHOICES:
|
|
409
511
|
budget.truncated = True
|
|
410
512
|
return None
|
|
411
513
|
states[index] = _ChatChoice()
|
|
@@ -430,7 +532,10 @@ def _merge_tool_call(state: _ChatChoice, delta: Any) -> None:
|
|
|
430
532
|
if name is not None:
|
|
431
533
|
current_function["name"] = name
|
|
432
534
|
if isinstance(arguments, str):
|
|
433
|
-
|
|
535
|
+
fragments = current_function.get("arguments")
|
|
536
|
+
if not isinstance(fragments, _DeltaBuffer):
|
|
537
|
+
fragments = current_function["arguments"] = _DeltaBuffer()
|
|
538
|
+
fragments.write(arguments)
|
|
434
539
|
current["function"] = current_function
|
|
435
540
|
state.tool_calls[tool_index] = current
|
|
436
541
|
|
|
@@ -451,15 +556,15 @@ def _record_chat_chunk(
|
|
|
451
556
|
if isinstance(role, str):
|
|
452
557
|
state.role = role
|
|
453
558
|
if isinstance(content, str) and budget.accept(content):
|
|
454
|
-
state.content
|
|
559
|
+
state.content_fragments.write(content)
|
|
455
560
|
if isinstance(reasoning, str) and budget.accept(reasoning):
|
|
456
|
-
state.reasoning
|
|
561
|
+
state.reasoning_fragments.write(reasoning)
|
|
457
562
|
for detail in _sequence_items(_field(delta, "reasoning_details")):
|
|
458
563
|
native_detail = _native(detail)
|
|
459
564
|
if native_detail is not None and budget.accept(native_detail):
|
|
460
565
|
state.reasoning_details.append(native_detail)
|
|
461
566
|
if isinstance(refusal, str) and budget.accept(refusal):
|
|
462
|
-
state.refusal
|
|
567
|
+
state.refusal_fragments.write(refusal)
|
|
463
568
|
for tool_call in _sequence_items(_field(delta, "tool_calls")):
|
|
464
569
|
if budget.accept(tool_call):
|
|
465
570
|
_merge_tool_call(state, tool_call)
|
|
@@ -528,6 +633,16 @@ def _response_event_has_output(event: Any) -> bool:
|
|
|
528
633
|
)
|
|
529
634
|
|
|
530
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
|
+
|
|
531
646
|
def _chat_output(states: Mapping[int, _ChatChoice]) -> list[dict[str, Any]]:
|
|
532
647
|
output: list[dict[str, Any]] = []
|
|
533
648
|
for _, state in sorted(states.items()):
|
|
@@ -543,13 +658,57 @@ def _chat_output(states: Mapping[int, _ChatChoice]) -> list[dict[str, Any]]:
|
|
|
543
658
|
if state.refusal:
|
|
544
659
|
message["refusal"] = state.refusal
|
|
545
660
|
if state.tool_calls:
|
|
546
|
-
message["tool_calls"] = [
|
|
661
|
+
message["tool_calls"] = [
|
|
662
|
+
_joined_tool_call(call) for _, call in sorted(state.tool_calls.items())
|
|
663
|
+
]
|
|
547
664
|
output.append(message)
|
|
548
665
|
return output
|
|
549
666
|
|
|
550
667
|
|
|
551
668
|
def _stream_capture_budget() -> telemetry_dev.CaptureBudget:
|
|
552
|
-
|
|
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
|
|
553
712
|
|
|
554
713
|
|
|
555
714
|
def _capture_masked(handle: telemetry_dev.SpanHandle) -> bool:
|
|
@@ -619,6 +778,15 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
|
|
|
619
778
|
response = getattr(inner, "response", None)
|
|
620
779
|
if response is None:
|
|
621
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
|
+
|
|
622
790
|
close = getattr(response, "close", None)
|
|
623
791
|
if callable(close):
|
|
624
792
|
|
|
@@ -626,7 +794,7 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
|
|
|
626
794
|
try:
|
|
627
795
|
return close(*args, **kwargs)
|
|
628
796
|
finally:
|
|
629
|
-
|
|
797
|
+
finish_if_alive()
|
|
630
798
|
|
|
631
799
|
response.close = _close_hook
|
|
632
800
|
aclose = getattr(response, "aclose", None)
|
|
@@ -637,7 +805,7 @@ def _hook_response_close(inner: Any, finish: Callable[[], None]) -> None:
|
|
|
637
805
|
try:
|
|
638
806
|
return await aclose_fn(*args, **kwargs)
|
|
639
807
|
finally:
|
|
640
|
-
|
|
808
|
+
finish_if_alive()
|
|
641
809
|
|
|
642
810
|
response.aclose = _aclose_hook
|
|
643
811
|
|
|
@@ -665,7 +833,6 @@ async def _close_async_stream(inner: Any) -> None:
|
|
|
665
833
|
class _InstrumentedStream:
|
|
666
834
|
def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
|
|
667
835
|
self._inner = inner
|
|
668
|
-
self._end = _end_once(handle)
|
|
669
836
|
self._handle = handle
|
|
670
837
|
self._started_at = started_at
|
|
671
838
|
self._states: dict[int, _ChatChoice] = {}
|
|
@@ -677,6 +844,9 @@ class _InstrumentedStream:
|
|
|
677
844
|
self._completed = False
|
|
678
845
|
self._mask_output_when_incomplete = _capture_masked(handle)
|
|
679
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
|
|
680
850
|
_hook_response_close(inner, self._on_response_close)
|
|
681
851
|
|
|
682
852
|
def _iterate(self) -> Iterator[Any]:
|
|
@@ -747,7 +917,15 @@ class _InstrumentedStream:
|
|
|
747
917
|
def __getattr__(self, name: str) -> Any:
|
|
748
918
|
return getattr(self._inner, name)
|
|
749
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
|
+
|
|
750
926
|
def _record(self, chunk: Any, received_at: float) -> bool:
|
|
927
|
+
if self._ended:
|
|
928
|
+
return True
|
|
751
929
|
update = _record_chat_chunk(chunk, self._states, self._budget)
|
|
752
930
|
if _chat_chunk_has_output(chunk):
|
|
753
931
|
record_output_chunk = getattr(self._handle, "record_output_chunk", None)
|
|
@@ -770,7 +948,6 @@ class _InstrumentedStream:
|
|
|
770
948
|
class _InstrumentedAsyncStream:
|
|
771
949
|
def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
|
|
772
950
|
self._inner = inner
|
|
773
|
-
self._end = _end_once(handle)
|
|
774
951
|
self._handle = handle
|
|
775
952
|
self._started_at = started_at
|
|
776
953
|
self._states: dict[int, _ChatChoice] = {}
|
|
@@ -782,6 +959,9 @@ class _InstrumentedAsyncStream:
|
|
|
782
959
|
self._completed = False
|
|
783
960
|
self._mask_output_when_incomplete = _capture_masked(handle)
|
|
784
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
|
|
785
965
|
_hook_response_close(inner, self._on_response_close)
|
|
786
966
|
|
|
787
967
|
async def _aiterate(self) -> AsyncIterator[Any]:
|
|
@@ -855,7 +1035,15 @@ class _InstrumentedAsyncStream:
|
|
|
855
1035
|
def __getattr__(self, name: str) -> Any:
|
|
856
1036
|
return getattr(self._inner, name)
|
|
857
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
|
+
|
|
858
1044
|
def _record(self, chunk: Any, received_at: float) -> bool:
|
|
1045
|
+
if self._ended:
|
|
1046
|
+
return True
|
|
859
1047
|
update = _record_chat_chunk(chunk, self._states, self._budget)
|
|
860
1048
|
if _chat_chunk_has_output(chunk):
|
|
861
1049
|
record_output_chunk = getattr(self._handle, "record_output_chunk", None)
|
|
@@ -875,6 +1063,37 @@ class _InstrumentedAsyncStream:
|
|
|
875
1063
|
return _chunk_error(chunk) is None
|
|
876
1064
|
|
|
877
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
|
+
|
|
878
1097
|
class _ResponsesStreamState:
|
|
879
1098
|
def __init__(self) -> None:
|
|
880
1099
|
self.partial: dict[str, Any] = {}
|
|
@@ -903,7 +1122,10 @@ class _ResponsesStreamState:
|
|
|
903
1122
|
|
|
904
1123
|
def _output_with_item(self, output_index: int, item: dict[str, Any]) -> list[dict[str, Any]]:
|
|
905
1124
|
items = {**self._items, output_index: item}
|
|
906
|
-
return [
|
|
1125
|
+
return [
|
|
1126
|
+
*[_joined_item(value) for _, value in sorted(items.items())],
|
|
1127
|
+
*self._synthetic_events,
|
|
1128
|
+
]
|
|
907
1129
|
|
|
908
1130
|
@staticmethod
|
|
909
1131
|
def _can_accept(budget: telemetry_dev.CaptureBudget, value: Any) -> bool:
|
|
@@ -914,14 +1136,19 @@ class _ResponsesStreamState:
|
|
|
914
1136
|
return candidate.accept(value)
|
|
915
1137
|
|
|
916
1138
|
def _replace_output(self, output: Any) -> bool:
|
|
917
|
-
replacement =
|
|
1139
|
+
replacement = _measuring_budget(self._budget)
|
|
918
1140
|
if not replacement.accept(output):
|
|
919
1141
|
self._budget.truncated = True
|
|
920
1142
|
return False
|
|
921
1143
|
replacement.truncated = self._provider_events_truncated
|
|
922
|
-
self._budget
|
|
1144
|
+
if not _adopt_budget(self._budget, replacement):
|
|
1145
|
+
self._budget.truncated = True
|
|
1146
|
+
return False
|
|
923
1147
|
return True
|
|
924
1148
|
|
|
1149
|
+
def release(self) -> None:
|
|
1150
|
+
_release_budget(self._budget)
|
|
1151
|
+
|
|
925
1152
|
def _sync_output(self) -> None:
|
|
926
1153
|
self.retained_output = self._output()
|
|
927
1154
|
self.partial["output"] = self.retained_output
|
|
@@ -1007,7 +1234,7 @@ class _ResponsesStreamState:
|
|
|
1007
1234
|
if event_type in {"response.output_item.added", "response.output_item.done"}:
|
|
1008
1235
|
raw_item = _field(event, "item")
|
|
1009
1236
|
replacement = event_type.endswith(".done")
|
|
1010
|
-
preflight =
|
|
1237
|
+
preflight = _measuring_budget(self._budget) if replacement else self._budget
|
|
1011
1238
|
if not self._can_accept(preflight, raw_item):
|
|
1012
1239
|
self._budget.truncated = True
|
|
1013
1240
|
return False
|
|
@@ -1025,7 +1252,7 @@ class _ResponsesStreamState:
|
|
|
1025
1252
|
if event_type in {"response.content_part.added", "response.content_part.done"}:
|
|
1026
1253
|
event_part = _field(event, "part")
|
|
1027
1254
|
replacement = event_type.endswith(".done")
|
|
1028
|
-
preflight =
|
|
1255
|
+
preflight = _measuring_budget(self._budget) if replacement else self._budget
|
|
1029
1256
|
if not self._can_accept(preflight, event_part):
|
|
1030
1257
|
self._budget.truncated = True
|
|
1031
1258
|
return False
|
|
@@ -1127,7 +1354,7 @@ class _ResponsesStreamState:
|
|
|
1127
1354
|
else:
|
|
1128
1355
|
if not self._budget.accept(value):
|
|
1129
1356
|
return False
|
|
1130
|
-
item
|
|
1357
|
+
_write_fragment(item, field, value)
|
|
1131
1358
|
name = _field(event, "name")
|
|
1132
1359
|
if name is not None:
|
|
1133
1360
|
item["name"] = name
|
|
@@ -1138,7 +1365,7 @@ class _ResponsesStreamState:
|
|
|
1138
1365
|
}:
|
|
1139
1366
|
event_part = _field(event, "part")
|
|
1140
1367
|
replacement = event_type.endswith(".done")
|
|
1141
|
-
preflight =
|
|
1368
|
+
preflight = _measuring_budget(self._budget) if replacement else self._budget
|
|
1142
1369
|
if not self._can_accept(preflight, event_part):
|
|
1143
1370
|
self._budget.truncated = True
|
|
1144
1371
|
return False
|
|
@@ -1253,7 +1480,7 @@ class _ResponsesStreamState:
|
|
|
1253
1480
|
self._summary[summary_key] = part
|
|
1254
1481
|
self._sync_item(output_index, item)
|
|
1255
1482
|
else:
|
|
1256
|
-
summary
|
|
1483
|
+
_write_fragment(summary, "text", value)
|
|
1257
1484
|
return True
|
|
1258
1485
|
content_index = self._event_index(event, "content_index", "contentIndex")
|
|
1259
1486
|
content_key = (output_index, content_index)
|
|
@@ -1292,9 +1519,9 @@ class _ResponsesStreamState:
|
|
|
1292
1519
|
self._content[content_key] = replacement
|
|
1293
1520
|
self._sync_item(output_index, item)
|
|
1294
1521
|
elif refusal:
|
|
1295
|
-
part
|
|
1522
|
+
_write_fragment(part, "refusal", value)
|
|
1296
1523
|
else:
|
|
1297
|
-
part
|
|
1524
|
+
_write_fragment(part, "text", value)
|
|
1298
1525
|
return True
|
|
1299
1526
|
|
|
1300
1527
|
def record(self, event: Any) -> None:
|
|
@@ -1302,7 +1529,7 @@ class _ResponsesStreamState:
|
|
|
1302
1529
|
response = _field(event, "response")
|
|
1303
1530
|
if response is not None:
|
|
1304
1531
|
fields = _responses_response(response, include_error=False, include_output=False)
|
|
1305
|
-
response_budget =
|
|
1532
|
+
response_budget = _measuring_budget(self._budget)
|
|
1306
1533
|
raw_output = _field(response, "output")
|
|
1307
1534
|
if raw_output is not None:
|
|
1308
1535
|
if not self._can_accept(response_budget, raw_output):
|
|
@@ -1314,7 +1541,11 @@ class _ResponsesStreamState:
|
|
|
1314
1541
|
self._budget.truncated = True
|
|
1315
1542
|
else:
|
|
1316
1543
|
response_artifact = [*response_items, *self._synthetic_events]
|
|
1544
|
+
adopted = False
|
|
1317
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:
|
|
1318
1549
|
self._items.clear()
|
|
1319
1550
|
self._content.clear()
|
|
1320
1551
|
self._summary.clear()
|
|
@@ -1322,8 +1553,6 @@ class _ResponsesStreamState:
|
|
|
1322
1553
|
if isinstance(item, dict):
|
|
1323
1554
|
self._hydrate_item(index, cast(dict[str, Any], item))
|
|
1324
1555
|
self.retained_output = response_artifact
|
|
1325
|
-
response_budget.truncated = self._provider_events_truncated
|
|
1326
|
-
self._budget = response_budget
|
|
1327
1556
|
else:
|
|
1328
1557
|
self._budget.truncated = True
|
|
1329
1558
|
self.partial = fields
|
|
@@ -1337,6 +1566,9 @@ class _ResponsesStreamState:
|
|
|
1337
1566
|
|
|
1338
1567
|
def fields(self, mask_output: bool, *, failed: bool = False) -> dict[str, Any]:
|
|
1339
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)]
|
|
1340
1572
|
if failed or self._budget.truncated or not self.saw_terminal_snapshot:
|
|
1341
1573
|
_mark_capture_incomplete(fields, mask_output)
|
|
1342
1574
|
return fields
|
|
@@ -1346,10 +1578,12 @@ class _InstrumentedResponsesStream:
|
|
|
1346
1578
|
def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
|
|
1347
1579
|
self._inner = inner
|
|
1348
1580
|
self._handle = handle
|
|
1349
|
-
self._end = _end_once(handle)
|
|
1350
1581
|
self._started_at = started_at
|
|
1351
1582
|
self._saw_first = False
|
|
1352
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
|
|
1353
1587
|
self._mask_output_when_incomplete = _capture_masked(handle)
|
|
1354
1588
|
self._consume: Iterator[Any] | None = None
|
|
1355
1589
|
self._in_next = False
|
|
@@ -1417,7 +1651,15 @@ class _InstrumentedResponsesStream:
|
|
|
1417
1651
|
def __getattr__(self, name: str) -> Any:
|
|
1418
1652
|
return getattr(self._inner, name)
|
|
1419
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
|
+
|
|
1420
1660
|
def _record(self, event: Any, received_at: float) -> None:
|
|
1661
|
+
if self._ended:
|
|
1662
|
+
return
|
|
1421
1663
|
if _response_event_has_output(event):
|
|
1422
1664
|
record_output_chunk = getattr(self._handle, "record_output_chunk", None)
|
|
1423
1665
|
if callable(record_output_chunk):
|
|
@@ -1444,10 +1686,12 @@ class _InstrumentedAsyncResponsesStream:
|
|
|
1444
1686
|
def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
|
|
1445
1687
|
self._inner = inner
|
|
1446
1688
|
self._handle = handle
|
|
1447
|
-
self._end = _end_once(handle)
|
|
1448
1689
|
self._started_at = started_at
|
|
1449
1690
|
self._saw_first = False
|
|
1450
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
|
|
1451
1695
|
self._mask_output_when_incomplete = _capture_masked(handle)
|
|
1452
1696
|
self._consume: AsyncIterator[Any] | None = None
|
|
1453
1697
|
self._in_next = False
|
|
@@ -1518,7 +1762,15 @@ class _InstrumentedAsyncResponsesStream:
|
|
|
1518
1762
|
def __getattr__(self, name: str) -> Any:
|
|
1519
1763
|
return getattr(self._inner, name)
|
|
1520
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
|
+
|
|
1521
1771
|
def _record(self, event: Any, received_at: float) -> None:
|
|
1772
|
+
if self._ended:
|
|
1773
|
+
return
|
|
1522
1774
|
if _response_event_has_output(event):
|
|
1523
1775
|
record_output_chunk = getattr(self._handle, "record_output_chunk", None)
|
|
1524
1776
|
if callable(record_output_chunk):
|
|
File without changes
|
|
File without changes
|