From 603b9d93eb8a18013f750575b7d3b273bf44e6f6 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 29 Sep 2026 11:33:30 +0200 Subject: [PATCH] ref(cohere): Drop send_default_pii support --- sentry_sdk/integrations/cohere.py | 27 +- tests/integrations/cohere/test_cohere.py | 351 +++++++++++++++-------- 2 files changed, 249 insertions(+), 129 deletions(-) diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 4fbb21e68e..6cbbe9d4fb 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -14,11 +14,9 @@ import sentry_sdk from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version -from sentry_sdk.scope import should_send_default_pii from sentry_sdk.utils import ( capture_internal_exceptions, event_from_exception, - has_data_collection_enabled, parse_version, reraise, ) @@ -90,14 +88,6 @@ def setup_once() -> None: BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True) -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() - - def _capture_exception(exc: "Any") -> None: event, hint = event_from_exception( exc, @@ -177,7 +167,7 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": reraise(*exc_info) with capture_internal_exceptions(): - if _should_record("inputs"): + if sentry_sdk.get_client().options["data_collection"]["gen_ai"]["inputs"]: set_data_normalized( span, SPANDATA.AI_INPUT_MESSAGES, @@ -213,7 +203,9 @@ def new_iterator() -> "Iterator[StreamedChatResponse]": collect_chat_response_fields( span, x.response, - include_pii=_should_record("outputs"), + include_pii=sentry_sdk.get_client().options[ + "data_collection" + ]["gen_ai"]["outputs"], ) yield x span.end() @@ -223,7 +215,9 @@ def new_iterator() -> "Iterator[StreamedChatResponse]": collect_chat_response_fields( span, res, - include_pii=_should_record("outputs"), + include_pii=sentry_sdk.get_client().options["data_collection"][ + "gen_ai" + ]["outputs"], ) span.end() else: @@ -248,7 +242,12 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any": "sentry.origin": CohereIntegration.origin, }, ) as span: - if "texts" in kwargs and _should_record("inputs"): + if ( + "texts" in kwargs + and sentry_sdk.get_client().options["data_collection"]["gen_ai"][ + "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 941dcc3145..e1efdd0652 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -11,19 +11,19 @@ from sentry_sdk.integrations.cohere import CohereIntegration -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) def test_nonstreaming_chat( sentry_init, capture_items, - send_default_pii, ): sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) client = Client(api_key="z") @@ -58,39 +58,89 @@ 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: - assert ( - '{"role": "system", "content": "some context"}' - in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] - ) - assert ( - '{"role": "user", "content": "hello"}' - in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] + assert ( + '{"role": "system", "content": "some context"}' + in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] + ) + assert ( + '{"role": "user", "content": "hello"}' + in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] + ) + assert "the model response" in span["attributes"][SPANDATA.AI_RESPONSES] + + assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 + assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 + + +def test_nonstreaming_chat_no_sensitive_data( + sentry_init, + capture_items, +): + sentry_init( + integrations=[CohereIntegration()], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + ) + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + return_value=httpx.Response( + 200, + json={ + "text": "the model response", + "meta": { + "billed_units": { + "output_tokens": 10, + "input_tokens": 20, + } + }, + }, ) - assert "the model response" in span["attributes"][SPANDATA.AI_RESPONSES] - else: - assert SPANDATA.AI_INPUT_MESSAGES not in span["attributes"] - assert SPANDATA.AI_RESPONSES not in span["attributes"] + ) + items = capture_items("span") + + response = client.chat( + model="some-model", + chat_history=[ChatMessage(role="SYSTEM", message="some context")], + message="hello", + ).text + + assert response == "the model response" + sentry_sdk.flush() + + assert len(items) == 1 + span = items[0].payload + + assert span["attributes"]["sentry.op"] == "ai.chat_completions.create.cohere" + assert span["attributes"][SPANDATA.AI_MODEL_ID] == "some-model" + + assert SPANDATA.AI_INPUT_MESSAGES not in span["attributes"] + assert SPANDATA.AI_RESPONSES not in span["attributes"] assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 -# noinspection PyTypeChecker -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) def test_streaming_chat( sentry_init, capture_items, - send_default_pii, ): sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) client = Client(api_key="z") @@ -140,82 +190,153 @@ 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: - assert ( - '{"role": "system", "content": "some context"}' - in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] - ) - assert ( - '{"role": "user", "content": "hello"}' - in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] - ) - assert "the model response" in span["attributes"][SPANDATA.AI_RESPONSES] - else: - assert SPANDATA.AI_INPUT_MESSAGES not in span["attributes"] - assert SPANDATA.AI_RESPONSES not in span["attributes"] + assert ( + '{"role": "system", "content": "some context"}' + in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] + ) + assert ( + '{"role": "user", "content": "hello"}' + in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] + ) + assert "the model response" in span["attributes"][SPANDATA.AI_RESPONSES] assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 -def test_bad_chat(sentry_init, capture_items): +def test_streaming_chat_no_sensitive_data( + sentry_init, + capture_items, +): sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, ) - items = capture_items("event", "span") - client = Client(api_key="z") - HTTPXClient.request = mock.Mock( - side_effect=httpx.HTTPError("API rate limit reached") + HTTPXClient.send = mock.Mock( + return_value=httpx.Response( + 200, + content="\n".join( + [ + json.dumps({"event_type": "text-generation", "text": "the model "}), + json.dumps({"event_type": "text-generation", "text": "response"}), + json.dumps( + { + "event_type": "stream-end", + "finish_reason": "COMPLETE", + "response": { + "text": "the model response", + "meta": { + "billed_units": { + "output_tokens": 10, + "input_tokens": 20, + } + }, + }, + } + ), + ] + ), + ) ) - with pytest.raises(httpx.HTTPError): - client.chat(model="some-model", message="hello") + items = capture_items("span") - (event,) = (item.payload for item in items if item.type == "event") - assert event["level"] == "error" + responses = list( + client.chat_stream( + model="some-model", + chat_history=[ChatMessage(role="SYSTEM", message="some context")], + message="hello", + ) + ) + response_string = responses[-1].response.text + assert response_string == "the model response" sentry_sdk.flush() - (span,) = (item.payload for item in items if item.type == "span") - assert span["status"] == "error" + assert len(items) == 1 + span = items[0].payload -def test_span_status_error(sentry_init, capture_items): + assert span["attributes"]["sentry.op"] == "ai.chat_completions.create.cohere" + assert span["attributes"][SPANDATA.AI_MODEL_ID] == "some-model" + + assert SPANDATA.AI_INPUT_MESSAGES not in span["attributes"] + assert SPANDATA.AI_RESPONSES not in span["attributes"] + + assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 + assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 + + +def test_embed( + sentry_init, + capture_items, +): sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) - items = capture_items("span") client = Client(api_key="z") HTTPXClient.request = mock.Mock( - side_effect=httpx.HTTPError("API rate limit reached") + return_value=httpx.Response( + 200, + json={ + "response_type": "embeddings_floats", + "id": "1", + "texts": ["hello"], + "embeddings": [[1.0, 2.0, 3.0]], + "meta": { + "billed_units": { + "input_tokens": 10, + } + }, + }, + ) ) - with pytest.raises(httpx.HTTPError): - client.chat(model="some-model", message="hello") + items = capture_items("span") + response = client.embed(texts=["hello"], model="text-embedding-3-large") + + assert len(response.embeddings[0]) == 3 sentry_sdk.flush() assert len(items) == 1 span = items[0].payload - assert span["status"] == "error" + assert span["attributes"]["sentry.op"] == "ai.embeddings.create.cohere" + assert "hello" in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) -def test_embed( + assert span["attributes"]["gen_ai.usage.input_tokens"] == 10 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 10 + + +def test_embed_no_sensitive_data( sentry_init, capture_items, - send_default_pii, ): sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, ) client = Client(api_key="z") @@ -246,19 +367,63 @@ def test_embed( span = items[0].payload assert span["attributes"]["sentry.op"] == "ai.embeddings.create.cohere" - if send_default_pii: - assert "hello" in span["attributes"][SPANDATA.AI_INPUT_MESSAGES] - else: - assert SPANDATA.AI_INPUT_MESSAGES not in span["attributes"] + assert SPANDATA.AI_INPUT_MESSAGES not in span["attributes"] assert span["attributes"]["gen_ai.usage.input_tokens"] == 10 assert span["attributes"]["gen_ai.usage.total_tokens"] == 10 +def test_bad_chat(sentry_init, capture_items): + sentry_init( + integrations=[CohereIntegration()], + traces_sample_rate=1.0, + data_collection={}, + ) + + items = capture_items("event", "span") + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + side_effect=httpx.HTTPError("API rate limit reached") + ) + with pytest.raises(httpx.HTTPError): + client.chat(model="some-model", message="hello") + + (event,) = (item.payload for item in items if item.type == "event") + assert event["level"] == "error" + + sentry_sdk.flush() + (span,) = (item.payload for item in items if item.type == "span") + assert span["status"] == "error" + + +def test_span_status_error(sentry_init, capture_items): + sentry_init( + integrations=[CohereIntegration()], + traces_sample_rate=1.0, + data_collection={}, + ) + items = capture_items("span") + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + side_effect=httpx.HTTPError("API rate limit reached") + ) + with pytest.raises(httpx.HTTPError): + client.chat(model="some-model", message="hello") + + sentry_sdk.flush() + + assert len(items) == 1 + span = items[0].payload + assert span["status"] == "error" + + def test_span_origin_chat(sentry_init, capture_items): sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, + data_collection={}, ) items = capture_items("span") @@ -294,6 +459,7 @@ def test_span_origin_embed(sentry_init, capture_items): sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, + data_collection={}, ) items = capture_items("span") @@ -323,25 +489,10 @@ def test_span_origin_embed(sentry_init, capture_items): assert span["attributes"]["sentry.origin"] == "auto.ai.cohere" -# data_collection config, send_default_pii, expect_inputs, expect_outputs +# data_collection config, expect_inputs, expect_outputs DATA_COLLECTION_CASES = [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - True, - True, - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - False, - False, - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), pytest.param( {"gen_ai": {"inputs": True, "outputs": False}}, - False, True, False, id="gen-ai-inputs-enabled-outputs-disabled", @@ -349,43 +500,19 @@ def test_span_origin_embed(sentry_init, capture_items): pytest.param( {"gen_ai": {"inputs": False, "outputs": True}}, False, - False, True, id="gen-ai-outputs-enabled-inputs-disabled", ), - pytest.param( - {"gen_ai": {}}, - False, - True, - True, - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - True, - True, - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - False, - False, - id="no-gen-ai-config-legacy-pii-disabled", - ), ] def _init_with_data_collection( sentry_init, data_collection, - send_default_pii, ): kwargs = dict( integrations=[CohereIntegration()], traces_sample_rate=1.0, - send_default_pii=send_default_pii, ) if data_collection is not None: kwargs["data_collection"] = data_collection @@ -394,21 +521,19 @@ def _init_with_data_collection( @pytest.mark.parametrize( - "data_collection, send_default_pii, expect_inputs, expect_outputs", + "data_collection, expect_inputs, expect_outputs", DATA_COLLECTION_CASES, ) def test_nonstreaming_chat_data_collection( sentry_init, capture_items, data_collection, - send_default_pii, expect_inputs, expect_outputs, ): _init_with_data_collection( sentry_init, data_collection, - send_default_pii, ) client = Client(api_key="z") @@ -470,21 +595,19 @@ def test_nonstreaming_chat_data_collection( @pytest.mark.parametrize( - "data_collection, send_default_pii, expect_inputs, expect_outputs", + "data_collection, expect_inputs, expect_outputs", DATA_COLLECTION_CASES, ) def test_streaming_chat_data_collection( sentry_init, capture_items, data_collection, - send_default_pii, expect_inputs, expect_outputs, ): _init_with_data_collection( sentry_init, data_collection, - send_default_pii, ) client = Client(api_key="z") @@ -559,21 +682,19 @@ def test_streaming_chat_data_collection( @pytest.mark.parametrize( - "data_collection, send_default_pii, expect_inputs, expect_outputs", + "data_collection, expect_inputs, expect_outputs", DATA_COLLECTION_CASES, ) def test_embed_data_collection( sentry_init, capture_items, data_collection, - send_default_pii, expect_inputs, expect_outputs, ): _init_with_data_collection( sentry_init, data_collection, - send_default_pii, ) client = Client(api_key="z")