diff --git a/sentry_sdk/integrations/anthropic.py b/sentry_sdk/integrations/anthropic.py index efc2f70ffd..ca9e60e59d 100644 --- a/sentry_sdk/integrations/anthropic.py +++ b/sentry_sdk/integrations/anthropic.py @@ -438,9 +438,13 @@ def _set_common_input_data( normalized_messages.append(transformed_message) role_normalized_messages = normalize_message_roles(normalized_messages) + + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - role_normalized_messages, span, scope + messages_data = ( + role_normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(role_normalized_messages, span, scope) ) if messages_data is not None: set_data_normalized( diff --git a/sentry_sdk/integrations/google_genai/utils.py b/sentry_sdk/integrations/google_genai/utils.py index 25763ebe07..55a5b80233 100644 --- a/sentry_sdk/integrations/google_genai/utils.py +++ b/sentry_sdk/integrations/google_genai/utils.py @@ -892,9 +892,12 @@ def set_span_data_for_request( if messages: normalized_messages = normalize_message_roles(messages) + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - normalized_messages, span, scope + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(normalized_messages, span, scope) ) if messages_data is not None: set_data_normalized( diff --git a/sentry_sdk/integrations/langchain.py b/sentry_sdk/integrations/langchain.py index 8acf215bfe..4f5a1b4939 100644 --- a/sentry_sdk/integrations/langchain.py +++ b/sentry_sdk/integrations/langchain.py @@ -374,9 +374,15 @@ def on_llm_start( } for prompt in prompts ] + + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - normalized_messages, span, scope + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages( + normalized_messages, span, scope + ) ) if messages_data is not None: set_data_normalized( @@ -463,9 +469,15 @@ def on_chat_model_start( self._normalize_langchain_message(message) ) normalized_messages = normalize_message_roles(normalized_messages) + + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - normalized_messages, span, scope + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages( + normalized_messages, span, scope + ) ) if messages_data is not None: set_data_normalized( @@ -992,9 +1004,15 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": and integration.include_prompts ): normalized_messages = normalize_message_roles([input]) + + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - normalized_messages, span, scope + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages( + normalized_messages, span, scope + ) ) if messages_data is not None: set_data_normalized( @@ -1049,9 +1067,13 @@ def new_stream(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": and integration.include_prompts ): normalized_messages = normalize_message_roles([input]) + + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - normalized_messages, span, scope + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(normalized_messages, span, scope) ) if messages_data is not None: set_data_normalized( diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py index e5ea12b90a..1454d151f4 100644 --- a/sentry_sdk/integrations/langgraph.py +++ b/sentry_sdk/integrations/langgraph.py @@ -181,9 +181,15 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": input_messages = _parse_langgraph_messages(args[0]) if input_messages: normalized_input_messages = normalize_message_roles(input_messages) + + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - normalized_input_messages, span, scope + messages_data = ( + normalized_input_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages( + normalized_input_messages, span, scope + ) ) if messages_data is not None: set_data_normalized( @@ -234,9 +240,15 @@ async def new_ainvoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": input_messages = _parse_langgraph_messages(args[0]) if input_messages: normalized_input_messages = normalize_message_roles(input_messages) + + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - normalized_input_messages, span, scope + messages_data = ( + normalized_input_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages( + normalized_input_messages, span, scope + ) ) if messages_data is not None: set_data_normalized( diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py index 3cff0fbc23..9561bd61f3 100644 --- a/sentry_sdk/integrations/litellm.py +++ b/sentry_sdk/integrations/litellm.py @@ -119,8 +119,11 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None: if isinstance(embedding_input, list) else [embedding_input] ) - messages_data = truncate_and_annotate_embedding_inputs( - input_list, span, scope + client = sentry_sdk.get_client() + messages_data = ( + input_list + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_embedding_inputs(input_list, span, scope) ) if messages_data is not None: set_data_normalized( @@ -133,9 +136,14 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None: # For chat, look for the 'messages' parameter messages = kwargs.get("messages", []) if messages: + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() messages = _convert_message_parts(messages) - messages_data = truncate_and_annotate_messages(messages, span, scope) + messages_data = ( + messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(messages, span, scope) + ) if messages_data is not None: set_data_normalized( span, diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index b3919d1a9d..7bb328741e 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -398,8 +398,13 @@ def _set_responses_api_input_data( if isinstance(messages, str): normalized_messages = normalize_message_roles([messages]) # type: ignore + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(normalized_messages, span, scope) + ) if messages_data is not None: set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False @@ -413,8 +418,13 @@ def _set_responses_api_input_data( ] if len(non_system_messages) > 0: normalized_messages = normalize_message_roles(non_system_messages) + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(normalized_messages, span, scope) + ) if messages_data is not None: set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False @@ -472,8 +482,13 @@ def _set_completions_api_input_data( if isinstance(messages, str): normalized_messages = normalize_message_roles([messages]) # type: ignore + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(normalized_messages, span, scope) + ) if messages_data is not None: set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False @@ -503,8 +518,13 @@ def _set_completions_api_input_data( ] if len(non_system_messages) > 0: normalized_messages = normalize_message_roles(non_system_messages) + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(normalized_messages, span, scope) + ) if messages_data is not None: set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False @@ -539,9 +559,14 @@ def _set_embeddings_input_data( set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") normalized_messages = normalize_message_roles([messages]) # type: ignore + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_embedding_inputs( - normalized_messages, span, scope + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_embedding_inputs( + normalized_messages, span, scope + ) ) if messages_data is not None: set_data_normalized( @@ -560,9 +585,14 @@ def _set_embeddings_input_data( if len(messages) > 0: normalized_messages = normalize_message_roles(messages) + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_embedding_inputs( - normalized_messages, span, scope + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_embedding_inputs( + normalized_messages, span, scope + ) ) if messages_data is not None: set_data_normalized( diff --git a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py index 27f9fdab25..2346189a96 100644 --- a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py +++ b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py @@ -63,9 +63,12 @@ def invoke_agent_span( if len(messages) > 0: normalized_messages = normalize_message_roles(messages) + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - normalized_messages, span, scope + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(normalized_messages, span, scope) ) if messages_data is not None: set_data_normalized( diff --git a/sentry_sdk/integrations/openai_agents/utils.py b/sentry_sdk/integrations/openai_agents/utils.py index ee504b3496..ea1faefde7 100644 --- a/sentry_sdk/integrations/openai_agents/utils.py +++ b/sentry_sdk/integrations/openai_agents/utils.py @@ -173,8 +173,13 @@ def _set_input_data( ) normalized_messages = normalize_message_roles(request_messages) + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(normalized_messages, span, scope) + ) if messages_data is not None: set_data_normalized( span, diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index dc95acad45..e549083fed 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -182,9 +182,12 @@ def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> Non if formatted_messages: normalized_messages = normalize_message_roles(formatted_messages) + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - normalized_messages, span, scope + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(normalized_messages, span, scope) ) set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index ee08ca7036..c507315dcd 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -122,9 +122,12 @@ def invoke_agent_span( if messages: normalized_messages = normalize_message_roles(messages) + client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - normalized_messages, span, scope + messages_data = ( + normalized_messages + if client.options.get("stream_gen_ai_spans", False) + else truncate_and_annotate_messages(normalized_messages, span, scope) ) set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False diff --git a/tests/integrations/anthropic/test_anthropic.py b/tests/integrations/anthropic/test_anthropic.py index 31f487aef2..4255a0e6fc 100644 --- a/tests/integrations/anthropic/test_anthropic.py +++ b/tests/integrations/anthropic/test_anthropic.py @@ -3625,20 +3625,14 @@ def mock_messages_create(*args, **kwargs): assert stored_messages[0]["role"] == expected_role -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_anthropic_message_truncation( - sentry_init, - capture_events, - capture_items, - stream_gen_ai_spans, -): +def test_anthropic_message_truncation(sentry_init, capture_events): """Test that large messages are truncated properly in Anthropic integration.""" sentry_init( integrations=[AnthropicIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) + events = capture_events() client = Anthropic(api_key="z") client.messages._post = mock.Mock(return_value=EXAMPLE_MESSAGE) @@ -3654,82 +3648,43 @@ def test_anthropic_message_truncation( {"role": "user", "content": "small message 5"}, ] - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(): - client.messages.create(max_tokens=1024, messages=messages, model="model") - - spans = [item.payload for item in items if item.type == "span"] - chat_spans = [ - span - for span in spans - if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT - ] - - assert len(chat_spans) > 0 - - chat_span = chat_spans[0] - assert chat_span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "anthropic" - assert chat_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["attributes"] - - messages_data = chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) - - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert "small message 5" in str(parsed_messages[0]) - - tx = next(item.payload for item in items if item.type == "transaction") - else: - events = capture_events() - - with start_transaction(): - client.messages.create(max_tokens=1024, messages=messages, model="model") + with start_transaction(): + client.messages.create(max_tokens=1024, messages=messages, model="model") - assert len(events) > 0 - tx = events[0] - assert tx["type"] == "transaction" + assert len(events) > 0 + tx = events[0] + assert tx["type"] == "transaction" - chat_spans = [ - span for span in tx.get("spans", []) if span.get("op") == OP.GEN_AI_CHAT - ] - - assert len(chat_spans) > 0 + chat_spans = [ + span for span in tx.get("spans", []) if span.get("op") == OP.GEN_AI_CHAT + ] + assert len(chat_spans) > 0 - chat_span = chat_spans[0] - assert chat_span["data"][SPANDATA.GEN_AI_SYSTEM] == "anthropic" - assert chat_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["data"] + chat_span = chat_spans[0] + assert chat_span["data"][SPANDATA.GEN_AI_SYSTEM] == "anthropic" + assert chat_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["data"] - messages_data = chat_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) + messages_data = chat_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert isinstance(messages_data, str) - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert "small message 5" in str(parsed_messages[0]) + parsed_messages = json.loads(messages_data) + assert isinstance(parsed_messages, list) + assert len(parsed_messages) == 1 + assert "small message 5" in str(parsed_messages[0]) assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5 -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio -async def test_anthropic_message_truncation_async( - sentry_init, - capture_events, - capture_items, - stream_gen_ai_spans, -): +async def test_anthropic_message_truncation_async(sentry_init, capture_events): """Test that large messages are truncated properly in Anthropic integration.""" sentry_init( integrations=[AnthropicIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) + events = capture_events() client = AsyncAnthropic(api_key="z") client.messages._post = mock.AsyncMock(return_value=EXAMPLE_MESSAGE) @@ -3745,68 +3700,30 @@ async def test_anthropic_message_truncation_async( {"role": "user", "content": "small message 5"}, ] - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(): - await client.messages.create( - max_tokens=1024, messages=messages, model="model" - ) - - spans = [item.payload for item in items if item.type == "span"] - chat_spans = [ - span - for span in spans - if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT - ] - assert len(chat_spans) > 0 - - chat_span = chat_spans[0] - - assert chat_span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "anthropic" - assert chat_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["attributes"] - - messages_data = chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + with start_transaction(): + await client.messages.create(max_tokens=1024, messages=messages, model="model") - assert isinstance(messages_data, str) + assert len(events) > 0 + tx = events[0] + assert tx["type"] == "transaction" - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert "small message 5" in str(parsed_messages[0]) - - tx = next(item.payload for item in items if item.type == "transaction") - else: - events = capture_events() - - with start_transaction(): - await client.messages.create( - max_tokens=1024, messages=messages, model="model" - ) - - assert len(events) > 0 - tx = events[0] - assert tx["type"] == "transaction" - - chat_spans = [ - span for span in tx.get("spans", []) if span.get("op") == OP.GEN_AI_CHAT - ] - assert len(chat_spans) > 0 - - chat_span = chat_spans[0] + chat_spans = [ + span for span in tx.get("spans", []) if span.get("op") == OP.GEN_AI_CHAT + ] + assert len(chat_spans) > 0 - assert chat_span["data"][SPANDATA.GEN_AI_SYSTEM] == "anthropic" - assert chat_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["data"] + chat_span = chat_spans[0] + assert chat_span["data"][SPANDATA.GEN_AI_SYSTEM] == "anthropic" + assert chat_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["data"] - messages_data = chat_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) + messages_data = chat_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert isinstance(messages_data, str) - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert "small message 5" in str(parsed_messages[0]) + parsed_messages = json.loads(messages_data) + assert isinstance(parsed_messages, list) + assert len(parsed_messages) == 1 + assert "small message 5" in str(parsed_messages[0]) assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5 @@ -5200,21 +5117,14 @@ def test_transform_message_content_list_anthropic(): # Integration tests for binary data in messages -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_message_with_base64_image( - sentry_init, - capture_events, - capture_items, - stream_gen_ai_spans, -): +def test_message_with_base64_image(sentry_init, capture_events): """Test that messages with base64 images are properly captured.""" sentry_init( integrations=[AnthropicIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) - + events = capture_events() client = Anthropic(api_key="z") client.messages._post = mock.Mock(return_value=EXAMPLE_MESSAGE) @@ -5235,31 +5145,15 @@ def test_message_with_base64_image( } ] - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(name="anthropic"): - client.messages.create(max_tokens=1024, messages=messages, model="model") - - spans = [item.payload for item in items if item.type == "span"] - (span,) = spans - - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["attributes"] - stored_messages = json.loads( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - else: - events = capture_events() - - with start_transaction(name="anthropic"): - client.messages.create(max_tokens=1024, messages=messages, model="model") + with start_transaction(name="anthropic"): + client.messages.create(max_tokens=1024, messages=messages, model="model") - assert len(events) == 1 - (event,) = events - (span,) = event["spans"] + assert len(events) == 1 + (event,) = events + (span,) = event["spans"] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["data"] - stored_messages = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["data"] + stored_messages = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(stored_messages) == 1 assert stored_messages[0]["role"] == "user" @@ -5409,21 +5303,14 @@ def test_message_with_file_image( } -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_message_with_base64_pdf( - sentry_init, - capture_events, - capture_items, - stream_gen_ai_spans, -): +def test_message_with_base64_pdf(sentry_init, capture_events): """Test that messages with base64-encoded PDF documents are properly captured.""" sentry_init( integrations=[AnthropicIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) - + events = capture_events() client = Anthropic(api_key="z") client.messages._post = mock.Mock(return_value=EXAMPLE_MESSAGE) @@ -5444,30 +5331,14 @@ def test_message_with_base64_pdf( } ] - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(name="anthropic"): - client.messages.create(max_tokens=1024, messages=messages, model="model") + with start_transaction(name="anthropic"): + client.messages.create(max_tokens=1024, messages=messages, model="model") - spans = [item.payload for item in items if item.type == "span"] - (span,) = spans - - stored_messages = json.loads( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - else: - events = capture_events() - - with start_transaction(name="anthropic"): - client.messages.create(max_tokens=1024, messages=messages, model="model") - - assert len(events) == 1 - (event,) = events - (span,) = event["spans"] - - stored_messages = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert len(events) == 1 + (event,) = events + (span,) = event["spans"] + stored_messages = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) content = stored_messages[0]["content"] assert content[1] == { "type": "blob", @@ -5612,21 +5483,14 @@ def test_message_with_file_document( } -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_message_with_mixed_content( - sentry_init, - capture_events, - capture_items, - stream_gen_ai_spans, -): +def test_message_with_mixed_content(sentry_init, capture_events): """Test that messages with mixed content (text, images, documents) are properly captured.""" sentry_init( integrations=[AnthropicIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) - + events = capture_events() client = Anthropic(api_key="z") client.messages._post = mock.Mock(return_value=EXAMPLE_MESSAGE) @@ -5663,30 +5527,14 @@ def test_message_with_mixed_content( } ] - if stream_gen_ai_spans: - items = capture_items("transaction", "span") + with start_transaction(name="anthropic"): + client.messages.create(max_tokens=1024, messages=messages, model="model") - with start_transaction(name="anthropic"): - client.messages.create(max_tokens=1024, messages=messages, model="model") - - spans = [item.payload for item in items if item.type == "span"] - (span,) = spans - - stored_messages = json.loads( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - else: - events = capture_events() - - with start_transaction(name="anthropic"): - client.messages.create(max_tokens=1024, messages=messages, model="model") - - assert len(events) == 1 - (event,) = events - (span,) = event["spans"] - - stored_messages = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert len(events) == 1 + (event,) = events + (span,) = event["spans"] + stored_messages = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) content = stored_messages[0]["content"] assert len(content) == 5 @@ -5718,21 +5566,14 @@ def test_message_with_mixed_content( } -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_message_with_multiple_images_different_formats( - sentry_init, - capture_events, - capture_items, - stream_gen_ai_spans, -): +def test_message_with_multiple_images_different_formats(sentry_init, capture_events): """Test that messages with multiple images of different source types are handled.""" sentry_init( integrations=[AnthropicIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) - + events = capture_events() client = Anthropic(api_key="z") client.messages._post = mock.Mock(return_value=EXAMPLE_MESSAGE) @@ -5768,30 +5609,14 @@ def test_message_with_multiple_images_different_formats( } ] - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(name="anthropic"): - client.messages.create(max_tokens=1024, messages=messages, model="model") - - spans = [item.payload for item in items if item.type == "span"] - (span,) = spans - - stored_messages = json.loads( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - else: - events = capture_events() - - with start_transaction(name="anthropic"): - client.messages.create(max_tokens=1024, messages=messages, model="model") + with start_transaction(name="anthropic"): + client.messages.create(max_tokens=1024, messages=messages, model="model") - assert len(events) == 1 - (event,) = events - (span,) = event["spans"] - - stored_messages = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert len(events) == 1 + (event,) = events + (span,) = event["spans"] + stored_messages = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) content = stored_messages[0]["content"] assert len(content) == 4 diff --git a/tests/integrations/google_genai/test_google_genai.py b/tests/integrations/google_genai/test_google_genai.py index 79318eaea5..ff0b59178b 100644 --- a/tests/integrations/google_genai/test_google_genai.py +++ b/tests/integrations/google_genai/test_google_genai.py @@ -1390,21 +1390,16 @@ def test_tool_calls_extraction( assert json.loads(tool_calls[1]["arguments"]) == {"timezone": "PST"} -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_google_genai_message_truncation( - sentry_init, - capture_events, - capture_items, - mock_genai_client, - stream_gen_ai_spans, + sentry_init, capture_events, mock_genai_client ): """Test that large messages are truncated properly in Google GenAI integration.""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) + events = capture_events() large_content = ( "This is a very long message that will exceed our size limits. " * 1000 @@ -1413,39 +1408,21 @@ def test_google_genai_message_truncation( mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - if stream_gen_ai_spans: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents=[large_content, small_content], - config=create_test_config(), - ) - - invoke_span = next(item.payload for item in items if item.type == "span") - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"] - - messages_data = invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + with start_transaction(name="google_genai"): mock_genai_client.models.generate_content( model="gemini-1.5-flash", contents=[large_content, small_content], config=create_test_config(), ) - (event,) = events - invoke_span = event["spans"][0] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["data"] + (event,) = events + invoke_span = event["spans"][0] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["data"] - messages_data = invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + messages_data = invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] assert isinstance(messages_data, str) parsed_messages = json.loads(messages_data) @@ -2528,21 +2505,16 @@ def test_generate_content_with_inline_data( assert messages[0]["content"][1]["content"] == BLOB_DATA_SUBSTITUTE -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_with_function_response( - sentry_init, - capture_events, - capture_items, - mock_genai_client, - stream_gen_ai_spans, + sentry_init, capture_events, mock_genai_client ): """Test generate_content with function_response (tool result).""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) + events = capture_events() mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2568,36 +2540,18 @@ def test_generate_content_with_function_response( ), ] - if stream_gen_ai_spans: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + with start_transaction(name="google_genai"): mock_genai_client.models.generate_content( model="gemini-1.5-flash", contents=contents, config=create_test_config() ) - invoke_span = next(item.payload for item in items if item.type == "span") - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - - (event,) = events - invoke_span = event["spans"][0] - - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + (event,) = events + invoke_span = event["spans"][0] + messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 # First message is user message assert messages[0]["role"] == "tool" @@ -2606,21 +2560,16 @@ def test_generate_content_with_function_response( assert messages[0]["content"]["output"] == "Sunny, 72F" -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_with_mixed_string_and_content( - sentry_init, - capture_events, - capture_items, - mock_genai_client, - stream_gen_ai_spans, + sentry_init, capture_events, mock_genai_client ): """Test generate_content with mixed string and Content objects in list.""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) + events = capture_events() mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2637,36 +2586,18 @@ def test_generate_content_with_mixed_string_and_content( ), ] - if stream_gen_ai_spans: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + with start_transaction(name="google_genai"): mock_genai_client.models.generate_content( model="gemini-1.5-flash", contents=contents, config=create_test_config() ) - invoke_span = next(item.payload for item in items if item.type == "span") - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - - (event,) = events - invoke_span = event["spans"][0] - - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + (event,) = events + invoke_span = event["spans"][0] + messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 # User message assert messages[0]["role"] == "user" @@ -2729,13 +2660,8 @@ def test_generate_content_with_part_object_directly( assert messages[0]["content"] == [{"text": "Direct Part object", "type": "text"}] -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_with_list_of_dicts( - sentry_init, - capture_events, - capture_items, - mock_genai_client, - stream_gen_ai_spans, + sentry_init, capture_events, mock_genai_client ): """ Test generate_content with list of dict format inputs. @@ -2748,8 +2674,8 @@ def test_generate_content_with_list_of_dicts( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) + events = capture_events() mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2760,36 +2686,18 @@ def test_generate_content_with_list_of_dicts( {"role": "user", "parts": [{"text": "Second user message"}]}, ] - if stream_gen_ai_spans: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + with start_transaction(name="google_genai"): mock_genai_client.models.generate_content( model="gemini-1.5-flash", contents=contents, config=create_test_config() ) - invoke_span = next(item.payload for item in items if item.type == "span") - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - - (event,) = events - invoke_span = event["spans"][0] - - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + (event,) = events + invoke_span = event["spans"][0] + messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 assert messages[0]["role"] == "user" assert messages[0]["content"] == [{"text": "Second user message", "type": "text"}] diff --git a/tests/integrations/langchain/test_langchain.py b/tests/integrations/langchain/test_langchain.py index 79ecc7e96b..2c0f8af977 100644 --- a/tests/integrations/langchain/test_langchain.py +++ b/tests/integrations/langchain/test_langchain.py @@ -2943,13 +2943,7 @@ def test_langchain_message_role_normalization_units(): assert normalized[5] == "string message" # String message unchanged -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_langchain_message_truncation( - sentry_init, - capture_events, - capture_items, - stream_gen_ai_spans, -): +def test_langchain_message_truncation(sentry_init, capture_events): """Test that large messages are truncated properly in Langchain integration.""" from langchain_core.outputs import LLMResult, Generation @@ -2957,8 +2951,8 @@ def test_langchain_message_truncation( integrations=[LangchainIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) + events = capture_events() callback = SentryLangchainCallback(max_span_map_size=100, include_prompts=True) @@ -2976,101 +2970,48 @@ def test_langchain_message_truncation( "small message 5", ] - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(): - callback.on_llm_start( - serialized=serialized, - prompts=prompts, - run_id=run_id, - name="my_pipeline", - invocation_params={ - "temperature": 0.7, - "max_tokens": 100, - "model": "gpt-3.5-turbo", - }, - ) - - response = LLMResult( - generations=[[Generation(text="The response")]], - llm_output={ - "token_usage": { - "total_tokens": 25, - "prompt_tokens": 10, - "completion_tokens": 15, - } - }, - ) - callback.on_llm_end(response=response, run_id=run_id) - - tx = next(item.payload for item in items if item.type == "transaction") - assert tx["type"] == "transaction" - - spans = [item.payload for item in items if item.type == "span"] - llm_spans = [ - span - for span in spans - if span["attributes"].get("sentry.op") == "gen_ai.text_completion" - ] - - assert len(llm_spans) > 0 - - llm_span = llm_spans[0] - - assert llm_span["attributes"]["gen_ai.operation.name"] == "text_completion" - assert llm_span["attributes"][SPANDATA.GEN_AI_FUNCTION_ID] == "my_pipeline" - - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in llm_span["attributes"] - messages_data = llm_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - else: - events = capture_events() - - with start_transaction(): - callback.on_llm_start( - serialized=serialized, - prompts=prompts, - run_id=run_id, - name="my_pipeline", - invocation_params={ - "temperature": 0.7, - "max_tokens": 100, - "model": "gpt-3.5-turbo", - }, - ) - - response = LLMResult( - generations=[[Generation(text="The response")]], - llm_output={ - "token_usage": { - "total_tokens": 25, - "prompt_tokens": 10, - "completion_tokens": 15, - } - }, - ) - callback.on_llm_end(response=response, run_id=run_id) - - assert len(events) > 0 - tx = events[0] - assert tx["type"] == "transaction" - - llm_spans = [ - span - for span in tx.get("spans", []) - if span.get("op") == "gen_ai.text_completion" - ] + with start_transaction(): + callback.on_llm_start( + serialized=serialized, + prompts=prompts, + run_id=run_id, + name="my_pipeline", + invocation_params={ + "temperature": 0.7, + "max_tokens": 100, + "model": "gpt-3.5-turbo", + }, + ) - assert len(llm_spans) > 0 + response = LLMResult( + generations=[[Generation(text="The response")]], + llm_output={ + "token_usage": { + "total_tokens": 25, + "prompt_tokens": 10, + "completion_tokens": 15, + } + }, + ) + callback.on_llm_end(response=response, run_id=run_id) - llm_span = llm_spans[0] + assert len(events) > 0 + tx = events[0] + assert tx["type"] == "transaction" - assert llm_span["data"]["gen_ai.operation.name"] == "text_completion" - assert llm_span["data"][SPANDATA.GEN_AI_FUNCTION_ID] == "my_pipeline" + llm_spans = [ + span + for span in tx.get("spans", []) + if span.get("op") == "gen_ai.text_completion" + ] + assert len(llm_spans) > 0 - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in llm_span["data"] - messages_data = llm_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + llm_span = llm_spans[0] + assert llm_span["data"]["gen_ai.operation.name"] == "text_completion" + assert llm_span["data"][SPANDATA.GEN_AI_FUNCTION_ID] == "my_pipeline" + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in llm_span["data"] + messages_data = llm_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] assert isinstance(messages_data, str) parsed_messages = json.loads(messages_data) diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py index f8df60739f..b8554f2f60 100644 --- a/tests/integrations/langgraph/test_langgraph.py +++ b/tests/integrations/langgraph/test_langgraph.py @@ -332,8 +332,9 @@ def original_invoke(self, *args, **kwargs): import json request_messages = json.loads(request_messages) - assert len(request_messages) == 1 - assert request_messages[0]["content"] == "Of course! How can I assist you?" + assert len(request_messages) == 2 + assert request_messages[0]["content"] == "Hello, can you help me?" + assert request_messages[1]["content"] == "Of course! How can I assist you?" response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] assert response_text == expected_assistant_response @@ -1987,13 +1988,7 @@ def __init__(self, content, message_type="human"): assert "ai" not in roles -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_langgraph_message_truncation( - sentry_init, - capture_events, - capture_items, - stream_gen_ai_spans, -): +def test_langgraph_message_truncation(sentry_init, capture_events): """Test that large messages are truncated properly in Langgraph integration.""" import json @@ -2001,8 +1996,8 @@ def test_langgraph_message_truncation( integrations=[LanggraphIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) + events = capture_events() large_content = ( "This is a very long message that will exceed our size limits. " * 1000 @@ -2022,66 +2017,28 @@ def test_langgraph_message_truncation( def original_invoke(self, *args, **kwargs): return {"messages": args[0].get("messages", [])} - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None - - spans = [item.payload for item in items if item.type == "span"] - invoke_spans = [ - span - for span in spans - if span["attributes"].get("sentry.op") == OP.GEN_AI_INVOKE_AGENT - ] - - assert len(invoke_spans) > 0 + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) - invoke_span = invoke_spans[0] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"] - - messages_data = invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - - assert isinstance(messages_data, str) - - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert "small message 5" in str(parsed_messages[0]) - (tx,) = (item.payload for item in items if item.type == "transaction") - else: - events = capture_events() - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None - - assert len(events) > 0 - tx = events[0] - assert tx["type"] == "transaction" - - invoke_spans = [ - span - for span in tx.get("spans", []) - if span.get("op") == OP.GEN_AI_INVOKE_AGENT - ] - - assert len(invoke_spans) > 0 + assert result is not None + assert len(events) > 0 + tx = events[0] + assert tx["type"] == "transaction" - invoke_span = invoke_spans[0] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["data"] + invoke_spans = [ + span for span in tx.get("spans", []) if span.get("op") == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) > 0 - messages_data = invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) + invoke_span = invoke_spans[0] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["data"] - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert "small message 5" in str(parsed_messages[0]) + messages_data = invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert isinstance(messages_data, str) + parsed_messages = json.loads(messages_data) + assert isinstance(parsed_messages, list) + assert len(parsed_messages) == 1 + assert "small message 5" in str(parsed_messages[0]) assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5 diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index 22663f9472..aab289b28f 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -2325,20 +2325,14 @@ def test_integration_setup(sentry_init): assert _failure_callback in (litellm.failure_callback or []) -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_litellm_message_truncation( - sentry_init, - capture_events, - capture_items, - stream_gen_ai_spans, -): +def test_litellm_message_truncation(sentry_init, capture_events): """Test that large messages are truncated properly in LiteLLM integration.""" sentry_init( integrations=[LiteLLMIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) + events = capture_events() large_content = ( "This is a very long message that will exceed our size limits. " * 1000 @@ -2352,78 +2346,39 @@ def test_litellm_message_truncation( ] mock_response = MockCompletionResponse() - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(name="litellm test"): - kwargs = { - "model": "gpt-3.5-turbo", - "messages": messages, - } - - _input_callback(kwargs) - _success_callback( - kwargs, - mock_response, - datetime.now(), - datetime.now(), - ) - - spans = [item.payload for item in items if item.type == "span"] - chat_spans = [ - span - for span in spans - if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT - ] - assert len(chat_spans) > 0 - - chat_span = chat_spans[0] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["attributes"] - - messages_data = chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) - - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert "small message 5" in str(parsed_messages[0]) - tx = next(item.payload for item in items if item.type == "transaction") - else: - events = capture_events() - - with start_transaction(name="litellm test"): - kwargs = { - "model": "gpt-3.5-turbo", - "messages": messages, - } + with start_transaction(name="litellm test"): + kwargs = { + "model": "gpt-3.5-turbo", + "messages": messages, + } - _input_callback(kwargs) - _success_callback( - kwargs, - mock_response, - datetime.now(), - datetime.now(), - ) + _input_callback(kwargs) + _success_callback( + kwargs, + mock_response, + datetime.now(), + datetime.now(), + ) - assert len(events) > 0 - tx = events[0] - assert tx["type"] == "transaction" + assert len(events) > 0 + tx = events[0] + assert tx["type"] == "transaction" - chat_spans = [ - span for span in tx.get("spans", []) if span.get("op") == OP.GEN_AI_CHAT - ] - assert len(chat_spans) > 0 + chat_spans = [ + span for span in tx.get("spans", []) if span.get("op") == OP.GEN_AI_CHAT + ] + assert len(chat_spans) > 0 - chat_span = chat_spans[0] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["data"] + chat_span = chat_spans[0] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["data"] - messages_data = chat_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) + messages_data = chat_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert isinstance(messages_data, str) - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert "small message 5" in str(parsed_messages[0]) + parsed_messages = json.loads(messages_data) + assert isinstance(parsed_messages, list) + assert len(parsed_messages) == 1 + assert "small message 5" in str(parsed_messages[0]) assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5 diff --git a/tests/integrations/openai/test_openai.py b/tests/integrations/openai/test_openai.py index 6c113078a3..c9e734da69 100644 --- a/tests/integrations/openai/test_openai.py +++ b/tests/integrations/openai/test_openai.py @@ -5853,21 +5853,16 @@ def test_openai_message_role_mapping( assert stored_messages[0]["role"] == expected_role -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_openai_message_truncation( - sentry_init, - capture_events, - capture_items, - nonstreaming_chat_completions_model_response, - stream_gen_ai_spans, + sentry_init, capture_events, nonstreaming_chat_completions_model_response ): """Test that large messages are truncated properly in OpenAI integration.""" sentry_init( integrations=[OpenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, ) + events = capture_events() client = OpenAI(api_key="z") client.chat.completions._post = mock.Mock( @@ -5894,45 +5889,22 @@ def test_openai_message_truncation( {"role": "user", "content": large_content}, ] - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(name="openai tx"): - client.chat.completions.create( - model="some-model", - messages=large_messages, - ) - - span = next(item.payload for item in items if item.type == "span") - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["attributes"] - - messages_data = span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) - - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) <= len(large_messages) - - (event,) = (item.payload for item in items if item.type == "transaction") - else: - events = capture_events() - - with start_transaction(name="openai tx"): - client.chat.completions.create( - model="some-model", - messages=large_messages, - ) + with start_transaction(name="openai tx"): + client.chat.completions.create( + model="some-model", + messages=large_messages, + ) - (event,) = events - span = event["spans"][0] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["data"] + (event,) = events + span = event["spans"][0] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["data"] - messages_data = span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) + messages_data = span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert isinstance(messages_data, str) - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) <= len(large_messages) + parsed_messages = json.loads(messages_data) + assert isinstance(parsed_messages, list) + assert len(parsed_messages) <= len(large_messages) meta_path = event["_meta"] span_meta = meta_path["spans"]["0"]["data"] diff --git a/tests/integrations/openai_agents/test_openai_agents.py b/tests/integrations/openai_agents/test_openai_agents.py index bf44562b14..beb44471de 100644 --- a/tests/integrations/openai_agents/test_openai_agents.py +++ b/tests/integrations/openai_agents/test_openai_agents.py @@ -1932,6 +1932,24 @@ def simple_test_tool(message: str) -> str: "gen_ai.request.messages" ] == safe_serialize( [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Please use the simple test tool"} + ], + }, + { + "role": "assistant", + "content": [ + { + "arguments": '{"message": "hello"}', + "call_id": "call_123", + "name": "simple_test_tool", + "type": "function_call", + "id": "call_123", + } + ], + }, { "role": "tool", "content": [