Skip to content

Commit d82510e

Browse files
dinmukhamedmdinmukhamedm
authored
fix(anthropic): async stream manager (#3220)
1 parent 1406790 commit d82510e

9 files changed

+1472
-6
lines changed

‎packages/opentelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/__init__.py

Lines changed: 82 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
from opentelemetry.instrumentation.anthropic.streaming import (
2020
abuild_from_streaming_response,
2121
build_from_streaming_response,
22+
WrappedAsyncMessageStreamManager,
23+
WrappedMessageStreamManager,
2224
)
2325
from opentelemetry.instrumentation.anthropic.utils import (
2426
acount_prompt_tokens_from_request,
@@ -70,6 +72,15 @@
7072
"method": "stream",
7173
"span_name": "anthropic.chat",
7274
},
75+
# This method is on an async resource, but is meant to be called as
76+
# an async context manager (async with), which we don't need to await;
77+
# thus, we wrap it with a sync wrapper
78+
{
79+
"package": "anthropic.resources.messages",
80+
"object": "AsyncMessages",
81+
"method": "stream",
82+
"span_name": "anthropic.chat",
83+
},
7384
]
7485

7586
WRAPPED_AMETHODS = [
@@ -85,19 +96,30 @@
8596
"method": "create",
8697
"span_name": "anthropic.chat",
8798
},
88-
{
89-
"package": "anthropic.resources.messages",
90-
"object": "AsyncMessages",
91-
"method": "stream",
92-
"span_name": "anthropic.chat",
93-
},
9499
]
95100

96101

97102
def is_streaming_response(response):
98103
return isinstance(response, Stream) or isinstance(response, AsyncStream)
99104

100105

106+
def is_stream_manager(response):
107+
"""Check if response is a MessageStreamManager or AsyncMessageStreamManager"""
108+
try:
109+
from anthropic.lib.streaming._messages import (
110+
MessageStreamManager,
111+
AsyncMessageStreamManager,
112+
)
113+
114+
return isinstance(response, (MessageStreamManager, AsyncMessageStreamManager))
115+
except ImportError:
116+
# Check by class name as fallback
117+
return (
118+
response.__class__.__name__ == "MessageStreamManager"
119+
or response.__class__.__name__ == "AsyncMessageStreamManager"
120+
)
121+
122+
101123
@dont_throw
102124
async def _aset_token_usage(
103125
span,
@@ -435,6 +457,33 @@ def _wrap(
435457
event_logger,
436458
kwargs,
437459
)
460+
elif is_stream_manager(response):
461+
if response.__class__.__name__ == "AsyncMessageStreamManager":
462+
return WrappedAsyncMessageStreamManager(
463+
response,
464+
span,
465+
instance._client,
466+
start_time,
467+
token_histogram,
468+
choice_counter,
469+
duration_histogram,
470+
exception_counter,
471+
event_logger,
472+
kwargs,
473+
)
474+
else:
475+
return WrappedMessageStreamManager(
476+
response,
477+
span,
478+
instance._client,
479+
start_time,
480+
token_histogram,
481+
choice_counter,
482+
duration_histogram,
483+
exception_counter,
484+
event_logger,
485+
kwargs,
486+
)
438487
elif response:
439488
try:
440489
metric_attributes = shared_metrics_attributes(response)
@@ -529,6 +578,33 @@ async def _awrap(
529578
event_logger,
530579
kwargs,
531580
)
581+
elif is_stream_manager(response):
582+
if response.__class__.__name__ == "AsyncMessageStreamManager":
583+
return WrappedAsyncMessageStreamManager(
584+
response,
585+
span,
586+
instance._client,
587+
start_time,
588+
token_histogram,
589+
choice_counter,
590+
duration_histogram,
591+
exception_counter,
592+
event_logger,
593+
kwargs,
594+
)
595+
else:
596+
return WrappedMessageStreamManager(
597+
response,
598+
span,
599+
instance._client,
600+
start_time,
601+
token_histogram,
602+
choice_counter,
603+
duration_histogram,
604+
exception_counter,
605+
event_logger,
606+
kwargs,
607+
)
532608
elif response:
533609
metric_attributes = shared_metrics_attributes(response)
534610

‎packages/opentelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/streaming.py

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,3 +296,99 @@ async def abuild_from_streaming_response(
296296
if span.is_recording():
297297
span.set_status(Status(StatusCode.OK))
298298
span.end()
299+
300+
301+
class WrappedMessageStreamManager:
302+
"""Wrapper for MessageStreamManager that handles instrumentation"""
303+
304+
def __init__(
305+
self,
306+
stream_manager,
307+
span,
308+
instance,
309+
start_time,
310+
token_histogram,
311+
choice_counter,
312+
duration_histogram,
313+
exception_counter,
314+
event_logger,
315+
kwargs,
316+
):
317+
self._stream_manager = stream_manager
318+
self._span = span
319+
self._instance = instance
320+
self._start_time = start_time
321+
self._token_histogram = token_histogram
322+
self._choice_counter = choice_counter
323+
self._duration_histogram = duration_histogram
324+
self._exception_counter = exception_counter
325+
self._event_logger = event_logger
326+
self._kwargs = kwargs
327+
328+
def __enter__(self):
329+
# Call the original stream manager's __enter__ to get the actual stream
330+
stream = self._stream_manager.__enter__()
331+
# Return the wrapped stream
332+
return build_from_streaming_response(
333+
self._span,
334+
stream,
335+
self._instance,
336+
self._start_time,
337+
self._token_histogram,
338+
self._choice_counter,
339+
self._duration_histogram,
340+
self._exception_counter,
341+
self._event_logger,
342+
self._kwargs,
343+
)
344+
345+
def __exit__(self, exc_type, exc_val, exc_tb):
346+
return self._stream_manager.__exit__(exc_type, exc_val, exc_tb)
347+
348+
349+
class WrappedAsyncMessageStreamManager:
350+
"""Wrapper for AsyncMessageStreamManager that handles instrumentation"""
351+
352+
def __init__(
353+
self,
354+
stream_manager,
355+
span,
356+
instance,
357+
start_time,
358+
token_histogram,
359+
choice_counter,
360+
duration_histogram,
361+
exception_counter,
362+
event_logger,
363+
kwargs,
364+
):
365+
self._stream_manager = stream_manager
366+
self._span = span
367+
self._instance = instance
368+
self._start_time = start_time
369+
self._token_histogram = token_histogram
370+
self._choice_counter = choice_counter
371+
self._duration_histogram = duration_histogram
372+
self._exception_counter = exception_counter
373+
self._event_logger = event_logger
374+
self._kwargs = kwargs
375+
376+
async def __aenter__(self):
377+
# Call the original stream manager's __aenter__ to get the actual stream
378+
stream = await self._stream_manager.__aenter__()
379+
# Return the wrapped stream
380+
return abuild_from_streaming_response(
381+
self._span,
382+
stream,
383+
self._instance,
384+
self._start_time,
385+
self._token_histogram,
386+
self._choice_counter,
387+
self._duration_histogram,
388+
self._exception_counter,
389+
self._event_logger,
390+
self._kwargs,
391+
)
392+
393+
async def __aexit__(self, exc_type, exc_val, exc_tb):
394+
return await self._stream_manager.__aexit__(exc_type, exc_val, exc_tb)

‎packages/opentelemetry-instrumentation-anthropic/tests/cassettes/test_messages/test_anthropic_message_stream_manager_legacy.yaml

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
interactions:
2+
- request:
3+
body: '{"max_tokens": 1024, "messages": [{"role": "user", "content": "Tell me
4+
a joke about OpenTelemetry"}], "model": "claude-3-5-haiku-20241022", "stream":
5+
true}'
6+
headers:
7+
accept:
8+
- application/json
9+
accept-encoding:
10+
- gzip, deflate
11+
anthropic-version:
12+
- '2023-06-01'
13+
connection:
14+
- keep-alive
15+
content-length:
16+
- '155'
17+
content-type:
18+
- application/json
19+
host:
20+
- api.anthropic.com
21+
user-agent:
22+
- Anthropic/Python 0.49.0
23+
x-stainless-arch:
24+
- arm64
25+
x-stainless-async:
26+
- 'false'
27+
x-stainless-lang:
28+
- python
29+
x-stainless-os:
30+
- MacOS
31+
x-stainless-package-version:
32+
- 0.49.0
33+
x-stainless-read-timeout:
34+
- '600'
35+
x-stainless-retry-count:
36+
- '0'
37+
x-stainless-runtime:
38+
- CPython
39+
x-stainless-runtime-version:
40+
- 3.9.6
41+
x-stainless-stream-helper:
42+
- messages
43+
x-stainless-timeout:
44+
- NOT_GIVEN
45+
method: POST
46+
uri: https://api.anthropic.com/v1/messages
47+
response:
48+
body:
49+
string: "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_01MCkQZZtEKF3nVbFaExwATe\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-3-5-haiku-20241022\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":17,\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0,\"output_tokens\":1,\"service_tier\":\"standard\"}}
50+
\ }\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}
51+
\ }\n\nevent: ping\ndata: {\"type\": \"ping\"}\n\nevent: content_block_delta\ndata:
52+
{\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Here\"}}\n\nevent:
53+
content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"'s
54+
a joke about Open\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Telemetry:\"}
55+
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\\n\\nWhy
56+
did the developer\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
57+
love\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
58+
OpenTelemetry\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"?\"}
59+
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\\n\\nBecause
60+
it\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
61+
helpe\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"d
62+
them\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
63+
trace\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
64+
their problems\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
65+
instea\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"d
66+
of just\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
67+
tr\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"acing
68+
their\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
69+
coffee\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
70+
m\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"ug!\"}
71+
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
72+
\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\\n\\n(\"}
73+
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Ba\"}
74+
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
75+
dum tss!\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
76+
\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\U0001F941
77+
\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"It\"}
78+
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"'s
79+
a\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
80+
play\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
81+
on the wor\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"d
82+
\\\"trace\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\\\"
83+
-\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
84+
which\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
85+
in\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
86+
OpenTelemetry\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
87+
means tracking\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
88+
system\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
89+
performance\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
90+
an\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"d
91+
interactions\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\",\"}
92+
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
93+
but\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
94+
also can\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
95+
mean physically\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
96+
following\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
97+
a path\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\".)\"}
98+
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\\n\\nWoul\"}
99+
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"d
100+
you like me\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
101+
to explain\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
102+
the joke\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
103+
or\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
104+
tell\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
105+
another\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
106+
tech\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
107+
humor\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
108+
one\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"?\"}
109+
\ }\n\nevent: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0
110+
\ }\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":108}
111+
\ }\n\nevent: message_stop\ndata: {\"type\":\"message_stop\" }\n\n"
112+
headers:
113+
CF-RAY:
114+
- 968d6aa8fff6ed07-LHR
115+
Cache-Control:
116+
- no-cache
117+
Connection:
118+
- keep-alive
119+
Content-Type:
120+
- text/event-stream; charset=utf-8
121+
Date:
122+
- Sat, 02 Aug 2025 12:00:45 GMT
123+
Server:
124+
- cloudflare
125+
Transfer-Encoding:
126+
- chunked
127+
X-Robots-Tag:
128+
- none
129+
anthropic-organization-id:
130+
- 04aa8588-6567-40cb-9042-a54b20ebaf4f
131+
anthropic-ratelimit-input-tokens-limit:
132+
- '400000'
133+
anthropic-ratelimit-input-tokens-remaining:
134+
- '400000'
135+
anthropic-ratelimit-input-tokens-reset:
136+
- '2025-08-02T12:00:45Z'
137+
anthropic-ratelimit-output-tokens-limit:
138+
- '80000'
139+
anthropic-ratelimit-output-tokens-remaining:
140+
- '80000'
141+
anthropic-ratelimit-output-tokens-reset:
142+
- '2025-08-02T12:00:45Z'
143+
anthropic-ratelimit-requests-limit:
144+
- '4000'
145+
anthropic-ratelimit-requests-remaining:
146+
- '3999'
147+
anthropic-ratelimit-requests-reset:
148+
- '2025-08-02T12:00:45Z'
149+
anthropic-ratelimit-tokens-limit:
150+
- '480000'
151+
anthropic-ratelimit-tokens-remaining:
152+
- '480000'
153+
anthropic-ratelimit-tokens-reset:
154+
- '2025-08-02T12:00:45Z'
155+
cf-cache-status:
156+
- DYNAMIC
157+
request-id:
158+
- req_011CRir4jvenjRy5HDFm6Z4m
159+
strict-transport-security:
160+
- max-age=31536000; includeSubDomains; preload
161+
via:
162+
- 1.1 google
163+
status:
164+
code: 200
165+
message: OK
166+
version: 1

0 commit comments

Comments
 (0)