Skip to content

fix(anthropic): async stream manager #3220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nirga merged 3 commits into from
Aug 3, 2025
Merged

fix(anthropic): async stream manager #3220

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0b4a7f5
fix(anthropic): async stream manager
dinmukhamedm Aug 3, 2025
1f080a5
Merge branch 'main' into fix/anthropic-stream-manager
dinmukhamedm Aug 3, 2025
3acff3c
Merge branch 'main' into fix/anthropic-stream-manager
nirga Aug 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 82 additions & 6 deletions ...entelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
from opentelemetry.instrumentation.anthropic.streaming import (
abuild_from_streaming_response,
build_from_streaming_response,
WrappedAsyncMessageStreamManager,
WrappedMessageStreamManager,
)
from opentelemetry.instrumentation.anthropic.utils import (
acount_prompt_tokens_from_request,
Expand Down Expand Up @@ -70,6 +72,15 @@
"method": "stream",
"span_name": "anthropic.chat",
},
# This method is on an async resource, but is meant to be called as
# an async context manager (async with), which we don't need to await;
# thus, we wrap it with a sync wrapper
{
"package": "anthropic.resources.messages",
"object": "AsyncMessages",
"method": "stream",
"span_name": "anthropic.chat",
},
]

WRAPPED_AMETHODS = [
Expand All @@ -85,19 +96,30 @@
"method": "create",
"span_name": "anthropic.chat",
},
{
"package": "anthropic.resources.messages",
"object": "AsyncMessages",
"method": "stream",
"span_name": "anthropic.chat",
},
]


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


def is_stream_manager(response):
"""Check if response is a MessageStreamManager or AsyncMessageStreamManager"""
try:
from anthropic.lib.streaming._messages import (
MessageStreamManager,
AsyncMessageStreamManager,
)

return isinstance(response, (MessageStreamManager, AsyncMessageStreamManager))
except ImportError:
# Check by class name as fallback
return (
response.__class__.__name__ == "MessageStreamManager"
or response.__class__.__name__ == "AsyncMessageStreamManager"
)


@dont_throw
async def _aset_token_usage(
span,
Expand Down Expand Up @@ -435,6 +457,33 @@ def _wrap(
event_logger,
kwargs,
)
elif is_stream_manager(response):
if response.__class__.__name__ == "AsyncMessageStreamManager":
return WrappedAsyncMessageStreamManager(
response,
span,
instance._client,
start_time,
token_histogram,
choice_counter,
duration_histogram,
exception_counter,
event_logger,
kwargs,
)
else:
return WrappedMessageStreamManager(
response,
span,
instance._client,
start_time,
token_histogram,
choice_counter,
duration_histogram,
exception_counter,
event_logger,
kwargs,
)
elif response:
try:
metric_attributes = shared_metrics_attributes(response)
Expand Down Expand Up @@ -529,6 +578,33 @@ async def _awrap(
event_logger,
kwargs,
)
elif is_stream_manager(response):
if response.__class__.__name__ == "AsyncMessageStreamManager":
return WrappedAsyncMessageStreamManager(
response,
span,
instance._client,
start_time,
token_histogram,
choice_counter,
duration_histogram,
exception_counter,
event_logger,
kwargs,
)
else:
return WrappedMessageStreamManager(
response,
span,
instance._client,
start_time,
token_histogram,
choice_counter,
duration_histogram,
exception_counter,
event_logger,
kwargs,
)
elif response:
metric_attributes = shared_metrics_attributes(response)

Expand Down
96 changes: 96 additions & 0 deletions ...ntelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/streaming.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -296,3 +296,99 @@ async def abuild_from_streaming_response(
if span.is_recording():
span.set_status(Status(StatusCode.OK))
span.end()


class WrappedMessageStreamManager:
"""Wrapper for MessageStreamManager that handles instrumentation"""

def __init__(
self,
stream_manager,
span,
instance,
start_time,
token_histogram,
choice_counter,
duration_histogram,
exception_counter,
event_logger,
kwargs,
):
self._stream_manager = stream_manager
self._span = span
self._instance = instance
self._start_time = start_time
self._token_histogram = token_histogram
self._choice_counter = choice_counter
self._duration_histogram = duration_histogram
self._exception_counter = exception_counter
self._event_logger = event_logger
self._kwargs = kwargs

def __enter__(self):
# Call the original stream manager's __enter__ to get the actual stream
stream = self._stream_manager.__enter__()
# Return the wrapped stream
return build_from_streaming_response(
self._span,
stream,
self._instance,
self._start_time,
self._token_histogram,
self._choice_counter,
self._duration_histogram,
self._exception_counter,
self._event_logger,
self._kwargs,
)

def __exit__(self, exc_type, exc_val, exc_tb):
return self._stream_manager.__exit__(exc_type, exc_val, exc_tb)


class WrappedAsyncMessageStreamManager:
"""Wrapper for AsyncMessageStreamManager that handles instrumentation"""

def __init__(
self,
stream_manager,
span,
instance,
start_time,
token_histogram,
choice_counter,
duration_histogram,
exception_counter,
event_logger,
kwargs,
):
self._stream_manager = stream_manager
self._span = span
self._instance = instance
self._start_time = start_time
self._token_histogram = token_histogram
self._choice_counter = choice_counter
self._duration_histogram = duration_histogram
self._exception_counter = exception_counter
self._event_logger = event_logger
self._kwargs = kwargs

async def __aenter__(self):
# Call the original stream manager's __aenter__ to get the actual stream
stream = await self._stream_manager.__aenter__()
# Return the wrapped stream
return abuild_from_streaming_response(
self._span,
stream,
self._instance,
self._start_time,
self._token_histogram,
self._choice_counter,
self._duration_histogram,
self._exception_counter,
self._event_logger,
self._kwargs,
)

async def __aexit__(self, exc_type, exc_val, exc_tb):
return await self._stream_manager.__aexit__(exc_type, exc_val, exc_tb)
166 changes: 166 additions & 0 deletions ...anthropic/tests/cassettes/test_messages/test_anthropic_message_stream_manager_legacy.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,166 @@
interactions:
- request:
body: '{"max_tokens": 1024, "messages": [{"role": "user", "content": "Tell me
a joke about OpenTelemetry"}], "model": "claude-3-5-haiku-20241022", "stream":
true}'
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate
anthropic-version:
- '2023-06-01'
connection:
- keep-alive
content-length:
- '155'
content-type:
- application/json
host:
- api.anthropic.com
user-agent:
- Anthropic/Python 0.49.0
x-stainless-arch:
- arm64
x-stainless-async:
- 'false'
x-stainless-lang:
- python
x-stainless-os:
- MacOS
x-stainless-package-version:
- 0.49.0
x-stainless-read-timeout:
- '600'
x-stainless-retry-count:
- '0'
x-stainless-runtime:
- CPython
x-stainless-runtime-version:
- 3.9.6
x-stainless-stream-helper:
- messages
x-stainless-timeout:
- NOT_GIVEN
method: POST
uri: https://api.anthropic.com/v1/messages
response:
body:
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\"}}
\ }\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}
\ }\n\nevent: ping\ndata: {\"type\": \"ping\"}\n\nevent: content_block_delta\ndata:
{\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Here\"}}\n\nevent:
content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"'s
a joke about Open\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Telemetry:\"}
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\\n\\nWhy
did the developer\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
love\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
OpenTelemetry\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"?\"}
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\\n\\nBecause
it\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
helpe\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"d
them\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
trace\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
their problems\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
instea\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"d
of just\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
tr\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"acing
their\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
coffee\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
m\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"ug!\"}
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\\n\\n(\"}
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Ba\"}
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
dum tss!\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\U0001F941
\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"It\"}
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"'s
a\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
play\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
on the wor\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"d
\\\"trace\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\\\"
-\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
which\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
in\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
OpenTelemetry\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
means tracking\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
system\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
performance\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
an\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"d
interactions\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\",\"}
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
but\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
also can\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
mean physically\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
following\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
a path\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\".)\"}
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\\n\\nWoul\"}
\ }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"d
you like me\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
to explain\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
the joke\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
or\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
tell\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
another\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
tech\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
humor\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"
one\"} }\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"?\"}
\ }\n\nevent: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0
\ }\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":108}
\ }\n\nevent: message_stop\ndata: {\"type\":\"message_stop\" }\n\n"
headers:
CF-RAY:
- 968d6aa8fff6ed07-LHR
Cache-Control:
- no-cache
Connection:
- keep-alive
Content-Type:
- text/event-stream; charset=utf-8
Date:
- Sat, 02 Aug 2025 12:00:45 GMT
Server:
- cloudflare
Transfer-Encoding:
- chunked
X-Robots-Tag:
- none
anthropic-organization-id:
- 04aa8588-6567-40cb-9042-a54b20ebaf4f
anthropic-ratelimit-input-tokens-limit:
- '400000'
anthropic-ratelimit-input-tokens-remaining:
- '400000'
anthropic-ratelimit-input-tokens-reset:
- '2025-08-02T12:00:45Z'
anthropic-ratelimit-output-tokens-limit:
- '80000'
anthropic-ratelimit-output-tokens-remaining:
- '80000'
anthropic-ratelimit-output-tokens-reset:
- '2025-08-02T12:00:45Z'
anthropic-ratelimit-requests-limit:
- '4000'
anthropic-ratelimit-requests-remaining:
- '3999'
anthropic-ratelimit-requests-reset:
- '2025-08-02T12:00:45Z'
anthropic-ratelimit-tokens-limit:
- '480000'
anthropic-ratelimit-tokens-remaining:
- '480000'
anthropic-ratelimit-tokens-reset:
- '2025-08-02T12:00:45Z'
cf-cache-status:
- DYNAMIC
request-id:
- req_011CRir4jvenjRy5HDFm6Z4m
strict-transport-security:
- max-age=31536000; includeSubDomains; preload
via:
- 1.1 google
status:
code: 200
message: OK
version: 1
Loading