diff --git a/sentry_sdk/integrations/mistral.py b/sentry_sdk/integrations/mistral.py index a7e6c5a2a2..021108784c 100644 --- a/sentry_sdk/integrations/mistral.py +++ b/sentry_sdk/integrations/mistral.py @@ -6,8 +6,6 @@ import sentry_sdk from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration -from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.utils import has_data_collection_enabled if TYPE_CHECKING: from typing import Any, Callable, Iterable, Optional, TypeGuard, Union @@ -245,15 +243,9 @@ def wrap_complete(self: "Chat", *args: "Any", **kwargs: "Any") -> "Any": span.set_attribute(SPANDATA.GEN_AI_RESPONSE_STREAMING, False) messages = kwargs.get("messages") - if isinstance(messages, Sequence) and ( - ( - has_data_collection_enabled(client.options) - and client.options["data_collection"]["gen_ai"]["inputs"] - ) - or ( - not has_data_collection_enabled(client.options) - and should_send_default_pii() - ) + if ( + isinstance(messages, Sequence) + and client.options["data_collection"]["gen_ai"]["inputs"] ): system_instructions = [ message for message in messages if _is_system_instruction(message) @@ -320,28 +312,24 @@ def wrap_complete(self: "Chat", *args: "Any", **kwargs: "Any") -> "Any": SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, response.usage.total_tokens ) - if ( - has_data_collection_enabled(client.options) - and client.options["data_collection"]["gen_ai"]["outputs"] - ) or ( - not has_data_collection_enabled(client.options) - and should_send_default_pii() - ): - output_messages: "list[OutputMessage]" = [] - for choice in response.choices: - if choice.message is None: - continue + if not client.options["data_collection"]["gen_ai"]["outputs"]: + return response - transformed_message = _transform_output_message(choice.message) - if transformed_message is None: - continue + output_messages: "list[OutputMessage]" = [] + for choice in response.choices: + if choice.message is None: + continue - output_messages.append(transformed_message) + transformed_message = _transform_output_message(choice.message) + if transformed_message is None: + continue - span.set_attribute( - SPANDATA.GEN_AI_OUTPUT_MESSAGES, - json.dumps(output_messages), - ) + output_messages.append(transformed_message) + + span.set_attribute( + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + json.dumps(output_messages), + ) return response @@ -375,15 +363,9 @@ async def wrap_complete_async(self: "Chat", *args: "Any", **kwargs: "Any") -> "A messages: "Optional[Union[Iterable[ChatCompletionRequestMessage], Iterable[ChatCompletionRequestMessageTypedDict]]]" = kwargs.get( "messages" ) - if isinstance(messages, Sequence) and ( - ( - has_data_collection_enabled(client.options) - and client.options["data_collection"]["gen_ai"]["inputs"] - ) - or ( - not has_data_collection_enabled(client.options) - and should_send_default_pii() - ) + if ( + isinstance(messages, Sequence) + and client.options["data_collection"]["gen_ai"]["inputs"] ): system_instructions = [ message for message in messages if _is_system_instruction(message) @@ -450,28 +432,24 @@ async def wrap_complete_async(self: "Chat", *args: "Any", **kwargs: "Any") -> "A SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, response.usage.total_tokens ) - if ( - has_data_collection_enabled(client.options) - and client.options["data_collection"]["gen_ai"]["outputs"] - ) or ( - not has_data_collection_enabled(client.options) - and should_send_default_pii() - ): - output_messages: "list[OutputMessage]" = [] - for choice in response.choices: - if choice.message is None: - continue + if not client.options["data_collection"]["gen_ai"]["outputs"]: + return response - transformed_message = _transform_output_message(choice.message) - if transformed_message is None: - continue + output_messages: "list[OutputMessage]" = [] + for choice in response.choices: + if choice.message is None: + continue - output_messages.append(transformed_message) + transformed_message = _transform_output_message(choice.message) + if transformed_message is None: + continue - span.set_attribute( - SPANDATA.GEN_AI_OUTPUT_MESSAGES, - json.dumps(output_messages), - ) + output_messages.append(transformed_message) + + span.set_attribute( + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + json.dumps(output_messages), + ) return response diff --git a/tests/integrations/mistral/test_mistral.py b/tests/integrations/mistral/test_mistral.py index c336125a4a..8a35a57979 100644 --- a/tests/integrations/mistral/test_mistral.py +++ b/tests/integrations/mistral/test_mistral.py @@ -60,6 +60,12 @@ def test_nonstreaming_chat( sentry_init( integrations=[MistralIntegration()], traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + }, + }, ) client = Mistral(api_key="z") @@ -121,6 +127,12 @@ async def test_nonstreaming_chat_async( sentry_init( integrations=[MistralIntegration()], traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, ) client = Mistral(api_key="z") @@ -365,7 +377,6 @@ async def test_nonstreaming_chat_async( ), ), ) -@pytest.mark.parametrize("data_collection", [True, False]) def test_input_attributes_nonstreaming_chat( sentry_init, capture_items, @@ -374,20 +385,17 @@ def test_input_attributes_nonstreaming_chat( messages, expected_system_instructions, expected_input_messages, - data_collection, ): - if data_collection: - sentry_init( - integrations=[MistralIntegration()], - traces_sample_rate=1.0, - data_collection={}, - ) - else: - sentry_init( - integrations=[MistralIntegration()], - traces_sample_rate=1.0, - send_default_pii=True, - ) + sentry_init( + integrations=[MistralIntegration()], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": False, + } + }, + ) client = Mistral(api_key="z") @@ -618,7 +626,6 @@ def test_input_attributes_nonstreaming_chat( ), ) @pytest.mark.asyncio -@pytest.mark.parametrize("data_collection", [True, False]) async def test_input_attributes_nonstreaming_chat_async( sentry_init, capture_items, @@ -627,20 +634,18 @@ async def test_input_attributes_nonstreaming_chat_async( messages, expected_system_instructions, expected_input_messages, - data_collection, ): - if data_collection: - sentry_init( - integrations=[MistralIntegration()], - traces_sample_rate=1.0, - data_collection={}, - ) - else: - sentry_init( - integrations=[MistralIntegration()], - traces_sample_rate=1.0, - send_default_pii=True, - ) + sentry_init( + integrations=[MistralIntegration()], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": False, + } + }, + ) + client = Mistral(api_key="z") model_response = get_model_response( @@ -677,26 +682,22 @@ async def test_input_attributes_nonstreaming_chat_async( ) -@pytest.mark.parametrize("data_collection", [True, False]) def test_output_attributes_nonstreaming_chat( sentry_init, capture_items, get_model_response, mistral_response, - data_collection, ): - if data_collection: - sentry_init( - integrations=[MistralIntegration()], - traces_sample_rate=1.0, - data_collection={}, - ) - else: - sentry_init( - integrations=[MistralIntegration()], - traces_sample_rate=1.0, - send_default_pii=True, - ) + sentry_init( + integrations=[MistralIntegration()], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": True, + } + }, + ) client = Mistral(api_key="z") @@ -742,26 +743,22 @@ def test_output_attributes_nonstreaming_chat( @pytest.mark.asyncio -@pytest.mark.parametrize("data_collection", [True, False]) async def test_output_attributes_nonstreaming_chat_async( sentry_init, capture_items, get_model_response, mistral_response, - data_collection, ): - if data_collection: - sentry_init( - integrations=[MistralIntegration()], - traces_sample_rate=1.0, - data_collection={}, - ) - else: - sentry_init( - integrations=[MistralIntegration()], - traces_sample_rate=1.0, - send_default_pii=True, - ) + sentry_init( + integrations=[MistralIntegration()], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": True, + } + }, + ) client = Mistral(api_key="z")