Anthropic streaming response does not get logged

[qria]

Description

Logfire does not log anthropic streaming response.

This is the resulting dashboard when I stream both openai and anthropic:

Anthropic streaming request gets logged but response does not.

Simple reproduction here: https://github.com/qria/logfire-bug-reproduction

Python, Logfire & OS Versions, related packages (not required)

logfire="2.7.1"
platform="macOS-14.6.1-arm64-arm-64bit"
python="3.11.9 (main, Jul 11 2024, 18:34:45) [Clang 15.0.0 (clang-1500.3.9.4)]"
[related_packages]
requests="2.32.3"
pydantic="2.10.3"
openai="1.57.4"
protobuf="5.29.1"
rich="13.9.4"
executing="2.1.0"
opentelemetry-api="1.29.0"
opentelemetry-exporter-otlp-proto-common="1.29.0"
opentelemetry-exporter-otlp-proto-http="1.29.0"
opentelemetry-instrumentation="0.50b0"
opentelemetry-proto="1.29.0"
opentelemetry-sdk="1.29.0"
opentelemetry-semantic-conventions="0.50b0"

[qria]

I think what's happening is that anthropic streaming response is wrapped again to anthropic MessageStream class, overwriting the stream_cls that logfire sets up.

    self.__stream = MessageStream(
        cast_to=raw_stream._cast_to,
        response=raw_stream.response,
        client=raw_stream._client,
    )

https://github.com/anthropics/anthropic-sdk-python/blob/main/src/anthropic/lib/streaming/_messages.py#L152-L161

Here raw_stream is correctly the type of LogfireInstrumentedStream but the response becomes MessageStream type because of this line.

Not sure how to fix or circumvent this issue.

[qria]

I'm just going to monkey-patch it for now

from anthropic.lib.streaming._messages import MessageStream, MessageStreamManager


def monkey_patch_enter(self: MessageStreamManager) -> MessageStream:
    """Monkey patch the MessageStreamManager for logfire instrumentation."""
    raw_stream = self._MessageStreamManager__api_request()  # type: ignore
    message_stream = MessageStream(  # type: ignore
        cast_to=raw_stream._cast_to,
        response=raw_stream.response,
        client=raw_stream._client,
    )
    message_stream._raw_stream._iterator = raw_stream._iterator

    self._MessageStreamManager__stream = message_stream  # type: ignore

    return message_stream


MessageStreamManager.__enter__ = monkey_patch_enter

[alexmojaki]

Hi @qria, thanks for the report. Your diagnosis and overall solution were spot on, I couldn't see another way from within logfire. So I did start a branch with my own version of this patching: https://github.com/pydantic/logfire/compare/alex/anthropic-streaming

But then I decided it would be better to fix the underlying issue in the anthropic SDK, so I opened the PR linked above. After that gets released, upgrading should fix the problem automatically. However, the maintainer who merged it said:

note that this is technically a breaking change as MessageStream.__init__ is public API but I'd be surprised if anyone is relying on it

which is funny because actually your patch is very much relying on it - you will need to remove it to not get an error.

Closing because I don't think there's more to do here.