diff --git a/sentry_sdk/integrations/anthropic.py b/sentry_sdk/integrations/anthropic.py index 4b04bc1cfc..27f1492109 100644 --- a/sentry_sdk/integrations/anthropic.py +++ b/sentry_sdk/integrations/anthropic.py @@ -1,3 +1,4 @@ +import inspect import json import sys from collections.abc import Iterable @@ -732,6 +733,175 @@ def _set_output_data( set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, input_tokens + output_tokens) +def _set_message_output_data( + span: "Any", integration: "AnthropicIntegration", message: "Any" +) -> None: + if hasattr(message, "content"): + ( + input_tokens, + output_tokens, + cache_read_input_tokens, + cache_write_input_tokens, + ) = _get_token_usage(message) + + content_blocks = [] + for content_block in message.content: + if hasattr(content_block, "to_dict"): + content_blocks.append(content_block.to_dict()) + elif hasattr(content_block, "model_dump"): + content_blocks.append(content_block.model_dump()) + elif hasattr(content_block, "text"): + content_blocks.append({"type": "text", "text": content_block.text}) + + _set_output_data( + span=span, + integration=integration, + model=getattr(message, "model", None), + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_read_input_tokens=cache_read_input_tokens, + cache_write_input_tokens=cache_write_input_tokens, + content_blocks=content_blocks, + response_id=getattr(message, "id", None), + finish_reason=getattr(message, "stop_reason", None), + ) + elif isinstance(span, Span): + span.set_data("unknown_response", True) + + +def _finish_span_once(span: "Any") -> "Callable[[], None]": + state = {"finished": False} + + def finish() -> None: + if state["finished"]: + return + state["finished"] = True + with capture_internal_exceptions(): + span.__exit__(None, None, None) + + return finish + + +def _finish_raw_response( + span: "Any", integration: "AnthropicIntegration", response: "Any" +) -> "Any": + """ + `.with_raw_response` and `.with_streaming_response` hand back the response + wrapper instead of a `Message`, so the span would end without any output. + + A body that is already read is parsed here, the way `create` parses it; + `parse()` caches its result, so the caller pays nothing for it. A body the + caller asked to stream must not be read here, so the span is filled when the + caller parses and closed when the caller closes the response. + """ + finish = _finish_span_once(span) + + if not hasattr(response, "parse"): + with capture_internal_exceptions(): + if isinstance(span, Span): + span.set_data("unknown_response", True) + finish() + return response + + if getattr(response.http_response, "is_stream_consumed", False): + with capture_internal_exceptions(): + _set_message_output_data(span, integration, response.parse()) + finish() + return response + + original_parse = response.parse + + def parse(*args: "Any", **kwargs: "Any") -> "Any": + parsed = original_parse(*args, **kwargs) + with capture_internal_exceptions(): + _set_message_output_data(span, integration, parsed) + finish() + return parsed + + response.parse = parse + + original_close = getattr(response, "close", None) + + if original_close is not None: + + def close() -> None: + finish() + original_close() + + response.close = close + return response + + +async def _finish_raw_response_async( + span: "Any", integration: "AnthropicIntegration", response: "Any" +) -> "Any": + """ + The async twin of `_finish_raw_response`. + + `parse()` and `close()` are coroutines on some SDK versions and plain + functions on others, so each wrapper keeps the shape of the method it + replaces and only awaits when the method it calls is awaitable. + """ + finish = _finish_span_once(span) + + if not hasattr(response, "parse"): + with capture_internal_exceptions(): + if isinstance(span, Span): + span.set_data("unknown_response", True) + finish() + return response + + def record(parsed: "Any") -> None: + with capture_internal_exceptions(): + _set_message_output_data(span, integration, parsed) + finish() + + if getattr(response.http_response, "is_stream_consumed", False): + with capture_internal_exceptions(): + parsed = response.parse() + if inspect.isawaitable(parsed): + parsed = await parsed + _set_message_output_data(span, integration, parsed) + finish() + return response + + original_parse = response.parse + + if inspect.iscoroutinefunction(original_parse): + + async def parse(*args: "Any", **kwargs: "Any") -> "Any": + parsed = await original_parse(*args, **kwargs) + record(parsed) + return parsed + + else: + + def parse(*args: "Any", **kwargs: "Any") -> "Any": # type: ignore[misc] + parsed = original_parse(*args, **kwargs) + record(parsed) + return parsed + + response.parse = parse + + original_close = getattr(response, "close", None) + + if original_close is not None: + if inspect.iscoroutinefunction(original_close): + + async def close() -> None: + finish() + await original_close() + + else: + + def close() -> None: # type: ignore[misc] + finish() + original_close() + + response.close = close + return response + + def _sentry_patched_create_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": """ Creates and manages an AI Client Span for both non-streaming and streaming calls. @@ -790,39 +960,11 @@ def _sentry_patched_create_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any return result - with capture_internal_exceptions(): - if hasattr(result, "content"): - ( - input_tokens, - output_tokens, - cache_read_input_tokens, - cache_write_input_tokens, - ) = _get_token_usage(result) - - content_blocks = [] - for content_block in result.content: - if hasattr(content_block, "to_dict"): - content_blocks.append(content_block.to_dict()) - elif hasattr(content_block, "model_dump"): - content_blocks.append(content_block.model_dump()) - elif hasattr(content_block, "text"): - content_blocks.append({"type": "text", "text": content_block.text}) - - _set_output_data( - span=span, - integration=integration, - model=getattr(result, "model", None), - input_tokens=input_tokens, - output_tokens=output_tokens, - cache_read_input_tokens=cache_read_input_tokens, - cache_write_input_tokens=cache_write_input_tokens, - content_blocks=content_blocks, - response_id=getattr(result, "id", None), - finish_reason=getattr(result, "stop_reason", None), - ) - elif isinstance(span, Span): - span.set_data("unknown_response", True) + if hasattr(result, "http_response"): + return _finish_raw_response(span, integration, result) + with capture_internal_exceptions(): + _set_message_output_data(span, integration, result) span.__exit__(None, None, None) return result @@ -888,39 +1030,11 @@ async def _sentry_patched_create_async( return result - with capture_internal_exceptions(): - if hasattr(result, "content"): - ( - input_tokens, - output_tokens, - cache_read_input_tokens, - cache_write_input_tokens, - ) = _get_token_usage(result) - - content_blocks = [] - for content_block in result.content: - if hasattr(content_block, "to_dict"): - content_blocks.append(content_block.to_dict()) - elif hasattr(content_block, "model_dump"): - content_blocks.append(content_block.model_dump()) - elif hasattr(content_block, "text"): - content_blocks.append({"type": "text", "text": content_block.text}) - - _set_output_data( - span=span, - integration=integration, - model=getattr(result, "model", None), - input_tokens=input_tokens, - output_tokens=output_tokens, - cache_read_input_tokens=cache_read_input_tokens, - cache_write_input_tokens=cache_write_input_tokens, - content_blocks=content_blocks, - response_id=getattr(result, "id", None), - finish_reason=getattr(result, "stop_reason", None), - ) - elif isinstance(span, Span): - span.set_data("unknown_response", True) + if hasattr(result, "http_response"): + return await _finish_raw_response_async(span, integration, result) + with capture_internal_exceptions(): + _set_message_output_data(span, integration, result) span.__exit__(None, None, None) return result diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index 04ae33daff..da31a39d08 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -1,3 +1,4 @@ +import inspect import json import sys import time @@ -836,6 +837,140 @@ def _set_common_output_data( span.__exit__(None, None, None) +def _finish_span_once(span: "Any") -> "Callable[[], None]": + state = {"finished": False} + + def finish() -> None: + if state["finished"]: + return + state["finished"] = True + with capture_internal_exceptions(): + span.__exit__(None, None, None) + + return finish + + +def _finish_raw_response( + span: "Any", response: "Any", kwargs: "Any", integration: "Any" +) -> "Any": + """ + `.with_raw_response` and `.with_streaming_response` hand back the response + wrapper instead of a `ChatCompletion`, so the span would end without any + output. + + A body that is already read is parsed here; `parse()` caches its result, so + the caller pays nothing for it. A body the caller asked to stream must not be + read here, so the span is filled when the caller parses and closed when the + caller closes the response. + """ + finish = _finish_span_once(span) + + def record(parsed: "Any") -> None: + with capture_internal_exceptions(): + _set_completions_api_output_data( + span, parsed, kwargs, integration, finish_span=False + ) + + if not hasattr(response, "parse"): + finish() + return response + + if getattr(response.http_response, "is_stream_consumed", False): + record(response.parse()) + finish() + return response + + original_parse = response.parse + + def parse(*args: "Any", **kwargs_: "Any") -> "Any": + parsed = original_parse(*args, **kwargs_) + record(parsed) + finish() + return parsed + + response.parse = parse + + original_close = getattr(response, "close", None) + + if original_close is not None: + + def close() -> None: + finish() + original_close() + + response.close = close + return response + + +async def _finish_raw_response_async( + span: "Any", response: "Any", kwargs: "Any", integration: "Any" +) -> "Any": + """ + The async twin of `_finish_raw_response`. + + `parse()` and `close()` are coroutines on some SDK versions and plain + functions on others, so each wrapper keeps the shape of the method it + replaces and only awaits when the method it calls is awaitable. + """ + finish = _finish_span_once(span) + + def record(parsed: "Any") -> None: + with capture_internal_exceptions(): + _set_completions_api_output_data( + span, parsed, kwargs, integration, finish_span=False + ) + + if not hasattr(response, "parse"): + finish() + return response + + if getattr(response.http_response, "is_stream_consumed", False): + parsed = response.parse() + if inspect.isawaitable(parsed): + parsed = await parsed + record(parsed) + finish() + return response + + original_parse = response.parse + + if inspect.iscoroutinefunction(original_parse): + + async def parse(*args: "Any", **kwargs_: "Any") -> "Any": + parsed = await original_parse(*args, **kwargs_) + record(parsed) + finish() + return parsed + + else: + + def parse(*args: "Any", **kwargs_: "Any") -> "Any": # type: ignore[misc] + parsed = original_parse(*args, **kwargs_) + record(parsed) + finish() + return parsed + + response.parse = parse + + original_close = getattr(response, "close", None) + + if original_close is not None: + if inspect.iscoroutinefunction(original_close): + + async def close() -> None: + finish() + await original_close() + + else: + + def close() -> None: # type: ignore[misc] + finish() + original_close() + + response.close = close + return response + + def _new_sync_chat_completion( f: "Callable[..., Union[ChatCompletion, Stream[ChatCompletionChunk]]]", *args: "Any", @@ -915,6 +1050,9 @@ def _new_sync_chat_completion( finish_span=True, ) + elif hasattr(response, "http_response"): + _finish_raw_response(span, response, kwargs, integration) + else: _set_completions_api_output_data( span, response, kwargs, integration, finish_span=True @@ -998,6 +1136,8 @@ async def _new_async_chat_completion( old_iterator=response._iterator, finish_span=True, ) + elif hasattr(response, "http_response"): + await _finish_raw_response_async(span, response, kwargs, integration) else: _set_completions_api_output_data( span, response, kwargs, integration, finish_span=True diff --git a/tests/integrations/anthropic/test_raw_response.py b/tests/integrations/anthropic/test_raw_response.py new file mode 100644 index 0000000000..c3996a3103 --- /dev/null +++ b/tests/integrations/anthropic/test_raw_response.py @@ -0,0 +1,162 @@ +"""`.with_raw_response` and `.with_streaming_response` must fill the span.""" + +import inspect +from types import SimpleNamespace +from unittest import mock + +import pytest + +import sentry_sdk +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations.anthropic import AnthropicIntegration +from sentry_sdk.integrations.stdlib import StdlibIntegration + +anthropic = pytest.importorskip("anthropic") +httpx2 = pytest.importorskip("httpx2") + +from anthropic import Anthropic, AsyncAnthropic # noqa: E402 +from anthropic.types import Message, TextBlock, Usage # noqa: E402 + +MESSAGE = Message( + id="msg_01", + model="claude-sonnet-4-5", + role="assistant", + content=[TextBlock(type="text", text="hi there")], + type="message", + stop_reason="end_turn", + usage=Usage(input_tokens=7, output_tokens=3), +) + +MESSAGE_JSON = { + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "hi there"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 7, "output_tokens": 3}, +} + +KWARGS = dict( + model="claude-sonnet-4-5", + max_tokens=10, + messages=[{"role": "user", "content": "hello"}], +) + + +class UnreadResponse: + """A response wrapper whose body has not been read yet.""" + + def __init__(self, parsed, consumed): + self.http_response = SimpleNamespace(is_stream_consumed=consumed) + self._parsed = parsed + self.parse_calls = 0 + self.closed = False + + def parse(self): + self.parse_calls += 1 + return self._parsed + + def close(self): + self.closed = True + + +def _init(sentry_init): + sentry_init( + integrations=[AnthropicIntegration(include_prompts=True)], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + send_default_pii=True, + stream_gen_ai_spans=False, + ) + + +def _gen_ai_spans(event): + return [span for span in event["spans"] if span["op"] == OP.GEN_AI_CHAT] + + +def _assert_has_output(span): + data = span["data"] + assert data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 7 + assert data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 3 + assert data[SPANDATA.GEN_AI_RESPONSE_ID] == "msg_01" + assert data[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == ["end_turn"] + assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == "hi there" + assert "unknown_response" not in data + + +def test_raw_and_streaming_responses_record_output(sentry_init, capture_events): + _init(sentry_init) + events = capture_events() + + client = Anthropic( + api_key="z", + http_client=httpx2.Client( + transport=httpx2.MockTransport( + lambda request: httpx2.Response(200, json=MESSAGE_JSON) + ) + ), + ) + + with sentry_sdk.start_transaction(name="t"): + client.messages.with_raw_response.create(**KWARGS).parse() + with client.messages.with_streaming_response.create(**KWARGS) as response: + response.parse() + + sentry_sdk.flush() + (event,) = events + spans = _gen_ai_spans(event) + assert len(spans) == 2 + for span in spans: + _assert_has_output(span) + + +def test_unread_body_is_parsed_by_the_caller(sentry_init, capture_events): + _init(sentry_init) + events = capture_events() + + client = Anthropic(api_key="z") + unread = UnreadResponse(MESSAGE, consumed=False) + client.messages._post = mock.Mock(return_value=unread) + + with sentry_sdk.start_transaction(name="t"): + with client.messages.with_streaming_response.create(**KWARGS) as response: + assert response.parse_calls == 0 + response.parse() + assert response.parse_calls == 1 + + assert response is unread + assert unread.closed + + sentry_sdk.flush() + (event,) = events + (span,) = _gen_ai_spans(event) + _assert_has_output(span) + + +@pytest.mark.asyncio +async def test_async_raw_response_records_output(sentry_init, capture_events): + _init(sentry_init) + events = capture_events() + + client = AsyncAnthropic( + api_key="z", + http_client=httpx2.AsyncClient( + transport=httpx2.MockTransport( + lambda request: httpx2.Response(200, json=MESSAGE_JSON) + ) + ), + ) + + with sentry_sdk.start_transaction(name="t"): + response = await client.messages.with_raw_response.create(**KWARGS) + parsed = response.parse() + if inspect.isawaitable(parsed): + parsed = await parsed + assert parsed is not None + + sentry_sdk.flush() + (event,) = events + (span,) = _gen_ai_spans(event) + _assert_has_output(span) diff --git a/tests/integrations/openai/test_raw_response.py b/tests/integrations/openai/test_raw_response.py new file mode 100644 index 0000000000..72e4b680c2 --- /dev/null +++ b/tests/integrations/openai/test_raw_response.py @@ -0,0 +1,169 @@ +"""`.with_raw_response` and `.with_streaming_response` must fill the span.""" + +import inspect +from types import SimpleNamespace +from unittest import mock + +import pytest + +import sentry_sdk +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations.openai import OpenAIIntegration +from sentry_sdk.integrations.stdlib import StdlibIntegration + +pytest.importorskip("openai") +httpx2 = pytest.importorskip("httpx2") + +from openai import AsyncOpenAI, OpenAI # noqa: E402 +from openai.types import CompletionUsage # noqa: E402 +from openai.types.chat import ChatCompletion, ChatCompletionMessage # noqa: E402 +from openai.types.chat.chat_completion import Choice # noqa: E402 + +COMPLETION = ChatCompletion( + id="chatcmpl-1", + choices=[ + Choice( + index=0, + finish_reason="stop", + message=ChatCompletionMessage(role="assistant", content="hi there"), + ) + ], + created=1, + model="gpt-4o-mini", + object="chat.completion", + usage=CompletionUsage(prompt_tokens=7, completion_tokens=3, total_tokens=10), +) + +COMPLETION_JSON = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi there"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}, +} + +KWARGS = dict(model="gpt-4o-mini", messages=[{"role": "user", "content": "hello"}]) + + +class UnreadResponse: + """A response wrapper whose body has not been read yet.""" + + def __init__(self, parsed, consumed): + self.http_response = SimpleNamespace(is_stream_consumed=consumed) + self._parsed = parsed + self.parse_calls = 0 + self.closed = False + + def parse(self): + self.parse_calls += 1 + return self._parsed + + def close(self): + self.closed = True + + +def _init(sentry_init): + sentry_init( + integrations=[OpenAIIntegration(include_prompts=True)], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + send_default_pii=True, + stream_gen_ai_spans=False, + ) + + +def _gen_ai_spans(event): + return [span for span in event["spans"] if span["op"] == OP.GEN_AI_CHAT] + + +def _assert_has_output(span): + data = span["data"] + assert data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 7 + assert data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 3 + assert "hi there" in str(data[SPANDATA.GEN_AI_RESPONSE_TEXT]) + + +def test_raw_and_streaming_responses_record_output(sentry_init, capture_events): + _init(sentry_init) + events = capture_events() + + client = OpenAI( + api_key="z", + http_client=httpx2.Client( + transport=httpx2.MockTransport( + lambda request: httpx2.Response(200, json=COMPLETION_JSON) + ) + ), + ) + + with sentry_sdk.start_transaction(name="t"): + client.chat.completions.with_raw_response.create(**KWARGS).parse() + with client.chat.completions.with_streaming_response.create( + **KWARGS + ) as response: + response.parse() + + sentry_sdk.flush() + (event,) = events + spans = _gen_ai_spans(event) + assert len(spans) == 2 + for span in spans: + _assert_has_output(span) + + +def test_unread_body_is_parsed_by_the_caller(sentry_init, capture_events): + _init(sentry_init) + events = capture_events() + + client = OpenAI(api_key="z") + unread = UnreadResponse(COMPLETION, consumed=False) + client.chat.completions._post = mock.Mock(return_value=unread) + + with sentry_sdk.start_transaction(name="t"): + response = client.chat.completions.with_raw_response.create(**KWARGS) + assert response.parse_calls == 0 + response.parse() + assert response.parse_calls == 1 + response.close() + + assert response is unread + assert unread.closed + + sentry_sdk.flush() + (event,) = events + (span,) = _gen_ai_spans(event) + _assert_has_output(span) + + +@pytest.mark.asyncio +async def test_async_raw_response_records_output(sentry_init, capture_events): + _init(sentry_init) + events = capture_events() + + client = AsyncOpenAI( + api_key="z", + http_client=httpx2.AsyncClient( + transport=httpx2.MockTransport( + lambda request: httpx2.Response(200, json=COMPLETION_JSON) + ) + ), + ) + + with sentry_sdk.start_transaction(name="t"): + response = await client.chat.completions.with_raw_response.create(**KWARGS) + parsed = response.parse() + if inspect.isawaitable(parsed): + parsed = await parsed + assert parsed is not None + + sentry_sdk.flush() + (event,) = events + (span,) = _gen_ai_spans(event) + _assert_has_output(span)