diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 06bba671cb..4fbb21e68e 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -80,9 +80,6 @@ class CohereIntegration(Integration): identifier = "cohere" origin = f"auto.ai.{identifier}" - def __init__(self: "CohereIntegration", include_prompts: bool = True) -> None: - self.include_prompts = include_prompts - @staticmethod def setup_once() -> None: version = parse_version(COHERE_VERSION) @@ -93,12 +90,12 @@ def setup_once() -> None: BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True) -def _should_record(integration: "CohereIntegration", category: str) -> bool: +def _should_record(category: str) -> bool: client = sentry_sdk.get_client() if has_data_collection_enabled(client.options): return bool(client.options["data_collection"]["gen_ai"][category]) - return should_send_default_pii() and integration.include_prompts + return should_send_default_pii() def _capture_exception(exc: "Any") -> None: @@ -180,7 +177,7 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": reraise(*exc_info) with capture_internal_exceptions(): - if _should_record(integration, "inputs"): + if _should_record("inputs"): set_data_normalized( span, SPANDATA.AI_INPUT_MESSAGES, @@ -216,7 +213,7 @@ def new_iterator() -> "Iterator[StreamedChatResponse]": collect_chat_response_fields( span, x.response, - include_pii=_should_record(integration, "outputs"), + include_pii=_should_record("outputs"), ) yield x span.end() @@ -226,7 +223,7 @@ def new_iterator() -> "Iterator[StreamedChatResponse]": collect_chat_response_fields( span, res, - include_pii=_should_record(integration, "outputs"), + include_pii=_should_record("outputs"), ) span.end() else: @@ -251,7 +248,7 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any": "sentry.origin": CohereIntegration.origin, }, ) as span: - if "texts" in kwargs and _should_record(integration, "inputs"): + if "texts" in kwargs and _should_record("inputs"): if isinstance(kwargs["texts"], str): set_data_normalized(span, SPANDATA.AI_TEXTS, [kwargs["texts"]]) elif ( diff --git a/tests/integrations/cohere/test_cohere.py b/tests/integrations/cohere/test_cohere.py index 2ce4e09396..941dcc3145 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -12,17 +12,16 @@ @pytest.mark.parametrize( - "send_default_pii, include_prompts", - [(True, True), (True, False), (False, True), (False, False)], + "send_default_pii", + [True, False], ) def test_nonstreaming_chat( sentry_init, capture_items, send_default_pii, - include_prompts, ): sentry_init( - integrations=[CohereIntegration(include_prompts=include_prompts)], + integrations=[CohereIntegration()], traces_sample_rate=1.0, send_default_pii=send_default_pii, ) @@ -59,7 +58,7 @@ def test_nonstreaming_chat( assert span["attributes"]["sentry.op"] == "ai.chat_completions.create.cohere" assert span["attributes"][SPANDATA.AI_MODEL_ID] == "some-model" - if send_default_pii and include_prompts: + if send_default_pii: assert ( '{"role": "system", "content": "some context"}' in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] @@ -80,17 +79,16 @@ def test_nonstreaming_chat( # noinspection PyTypeChecker @pytest.mark.parametrize( - "send_default_pii, include_prompts", - [(True, True), (True, False), (False, True), (False, False)], + "send_default_pii", + [True, False], ) def test_streaming_chat( sentry_init, capture_items, send_default_pii, - include_prompts, ): sentry_init( - integrations=[CohereIntegration(include_prompts=include_prompts)], + integrations=[CohereIntegration()], traces_sample_rate=1.0, send_default_pii=send_default_pii, ) @@ -142,7 +140,7 @@ def test_streaming_chat( assert span["attributes"]["sentry.op"] == "ai.chat_completions.create.cohere" assert span["attributes"][SPANDATA.AI_MODEL_ID] == "some-model" - if send_default_pii and include_prompts: + if send_default_pii: assert ( '{"role": "system", "content": "some context"}' in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] @@ -206,17 +204,16 @@ def test_span_status_error(sentry_init, capture_items): @pytest.mark.parametrize( - "send_default_pii, include_prompts", - [(True, True), (True, False), (False, True), (False, False)], + "send_default_pii", + [True, False], ) def test_embed( sentry_init, capture_items, send_default_pii, - include_prompts, ): sentry_init( - integrations=[CohereIntegration(include_prompts=include_prompts)], + integrations=[CohereIntegration()], traces_sample_rate=1.0, send_default_pii=send_default_pii, ) @@ -249,7 +246,7 @@ def test_embed( span = items[0].payload assert span["attributes"]["sentry.op"] == "ai.embeddings.create.cohere" - if send_default_pii and include_prompts: + if send_default_pii: assert "hello" in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] else: assert SPANDATA.AI_INPUT_MESSAGES not in span["attributes"] @@ -326,12 +323,11 @@ def test_span_origin_embed(sentry_init, capture_items): assert span["attributes"]["sentry.origin"] == "auto.ai.cohere" -# data_collection config, send_default_pii, include_prompts, expect_inputs, expect_outputs +# data_collection config, send_default_pii, expect_inputs, expect_outputs DATA_COLLECTION_CASES = [ pytest.param( {"gen_ai": {"inputs": True, "outputs": True}}, False, - False, True, True, id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", @@ -339,7 +335,6 @@ def test_span_origin_embed(sentry_init, capture_items): pytest.param( {"gen_ai": {"inputs": False, "outputs": False}}, True, - True, False, False, id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", @@ -347,7 +342,6 @@ def test_span_origin_embed(sentry_init, capture_items): pytest.param( {"gen_ai": {"inputs": True, "outputs": False}}, False, - False, True, False, id="gen-ai-inputs-enabled-outputs-disabled", @@ -356,14 +350,12 @@ def test_span_origin_embed(sentry_init, capture_items): {"gen_ai": {"inputs": False, "outputs": True}}, False, False, - False, True, id="gen-ai-outputs-enabled-inputs-disabled", ), pytest.param( {"gen_ai": {}}, False, - False, True, True, id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", @@ -373,13 +365,11 @@ def test_span_origin_embed(sentry_init, capture_items): True, True, True, - True, - id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + id="no-gen-ai-config-legacy-pii", ), pytest.param( None, False, - True, False, False, id="no-gen-ai-config-legacy-pii-disabled", @@ -391,10 +381,9 @@ def _init_with_data_collection( sentry_init, data_collection, send_default_pii, - include_prompts, ): kwargs = dict( - integrations=[CohereIntegration(include_prompts=include_prompts)], + integrations=[CohereIntegration()], traces_sample_rate=1.0, send_default_pii=send_default_pii, ) @@ -405,7 +394,7 @@ def _init_with_data_collection( @pytest.mark.parametrize( - "data_collection, send_default_pii, include_prompts, expect_inputs, expect_outputs", + "data_collection, send_default_pii, expect_inputs, expect_outputs", DATA_COLLECTION_CASES, ) def test_nonstreaming_chat_data_collection( @@ -413,7 +402,6 @@ def test_nonstreaming_chat_data_collection( capture_items, data_collection, send_default_pii, - include_prompts, expect_inputs, expect_outputs, ): @@ -421,7 +409,6 @@ def test_nonstreaming_chat_data_collection( sentry_init, data_collection, send_default_pii, - include_prompts, ) client = Client(api_key="z") @@ -483,7 +470,7 @@ def test_nonstreaming_chat_data_collection( @pytest.mark.parametrize( - "data_collection, send_default_pii, include_prompts, expect_inputs, expect_outputs", + "data_collection, send_default_pii, expect_inputs, expect_outputs", DATA_COLLECTION_CASES, ) def test_streaming_chat_data_collection( @@ -491,7 +478,6 @@ def test_streaming_chat_data_collection( capture_items, data_collection, send_default_pii, - include_prompts, expect_inputs, expect_outputs, ): @@ -499,7 +485,6 @@ def test_streaming_chat_data_collection( sentry_init, data_collection, send_default_pii, - include_prompts, ) client = Client(api_key="z") @@ -574,7 +559,7 @@ def test_streaming_chat_data_collection( @pytest.mark.parametrize( - "data_collection, send_default_pii, include_prompts, expect_inputs, expect_outputs", + "data_collection, send_default_pii, expect_inputs, expect_outputs", DATA_COLLECTION_CASES, ) def test_embed_data_collection( @@ -582,7 +567,6 @@ def test_embed_data_collection( capture_items, data_collection, send_default_pii, - include_prompts, expect_inputs, expect_outputs, ): @@ -590,7 +574,6 @@ def test_embed_data_collection( sentry_init, data_collection, send_default_pii, - include_prompts, ) client = Client(api_key="z")