From 2e8a9655b21faa2203383eececa5509951bec18b Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Fri, 28 Aug 2026 13:40:12 -0400 Subject: [PATCH] feat(cohere): Support data_collection filtering for inputs and outputs Respect the data_collection.gen_ai.inputs/outputs config option when capturing chat/embed request and response data, falling back to the legacy send_default_pii + include_prompts behavior when not set. Fixes PY-2747 Fixes #7282 --- sentry_sdk/integrations/cohere.py | 27 +- tests/integrations/cohere/test_cohere.py | 325 +++++++++++++++++++++++ 2 files changed, 343 insertions(+), 9 deletions(-) diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 7abf3f6808..9cac912a2b 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -17,7 +17,12 @@ import sentry_sdk from sentry_sdk.integrations import DidNotEnable, Integration from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.utils import capture_internal_exceptions, event_from_exception, reraise +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, + has_data_collection_enabled, + reraise, +) try: from cohere import ( @@ -83,6 +88,14 @@ def setup_once() -> None: BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True) +def _should_record(integration: "CohereIntegration", 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 + + def _capture_exception(exc: "Any") -> None: event, hint = event_from_exception( exc, @@ -179,7 +192,7 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": reraise(*exc_info) with capture_internal_exceptions(): - if should_send_default_pii() and integration.include_prompts: + if _should_record(integration, "inputs"): set_data_normalized( span, SPANDATA.AI_INPUT_MESSAGES, @@ -215,8 +228,7 @@ def new_iterator() -> "Iterator[StreamedChatResponse]": collect_chat_response_fields( span, x.response, - include_pii=should_send_default_pii() - and integration.include_prompts, + include_pii=_should_record(integration, "outputs"), ) yield x _end_span(span) @@ -226,8 +238,7 @@ def new_iterator() -> "Iterator[StreamedChatResponse]": collect_chat_response_fields( span, res, - include_pii=should_send_default_pii() - and integration.include_prompts, + include_pii=_should_record(integration, "outputs"), ) _end_span(span) else: @@ -265,9 +276,7 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any": ) with span_ctx as span: - if "texts" in kwargs and ( - should_send_default_pii() and integration.include_prompts - ): + if "texts" in kwargs and _should_record(integration, "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 73d4204727..b27b8e5d4e 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -454,3 +454,328 @@ def test_span_origin_embed(sentry_init, capture_events): assert event["contexts"]["trace"]["origin"] == "manual" assert event["spans"][0]["origin"] == "auto.ai.cohere" + + +# data_collection config, send_default_pii, include_prompts, 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", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": False}}, + True, + True, + False, + False, + id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", + ), + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + False, + False, + True, + False, + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"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", + ), + pytest.param( + None, + True, + True, + True, + True, + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + False, + False, + id="no-gen-ai-config-legacy-pii-disabled", + ), +] + + +def _init_with_data_collection( + sentry_init, data_collection, send_default_pii, include_prompts, span_streaming +): + kwargs = dict( + integrations=[CohereIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + trace_lifecycle="stream" if span_streaming else "static", + ) + if data_collection is not None: + kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**kwargs) + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize( + "data_collection, send_default_pii, include_prompts, expect_inputs, expect_outputs", + DATA_COLLECTION_CASES, +) +def test_nonstreaming_chat_data_collection( + sentry_init, + capture_events, + capture_items, + data_collection, + send_default_pii, + include_prompts, + expect_inputs, + expect_outputs, + span_streaming, +): + _init_with_data_collection( + sentry_init, data_collection, send_default_pii, include_prompts, span_streaming + ) + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + return_value=httpx.Response( + 200, + json={ + "text": "the model response", + "generation_id": "gen-1", + "citations": [ + { + "start": 0, + "end": 3, + "text": "the", + "document_ids": ["doc-1"], + } + ], + "meta": { + "billed_units": { + "output_tokens": 10, + "input_tokens": 20, + } + }, + }, + ) + ) + + if span_streaming: + items = capture_items("span") + else: + events = capture_events() + + with start_transaction(name="cohere tx"): + client.chat( + model="some-model", + chat_history=[ChatMessage(role="SYSTEM", message="some context")], + message="hello", + preamble="be concise", + ) + + if span_streaming: + sentry_sdk.flush() + assert len(items) == 1 + attributes = items[0].payload["attributes"] + else: + attributes = events[0]["spans"][0]["data"] + + assert attributes[SPANDATA.AI_MODEL_ID] == "some-model" + assert attributes["gen_ai.usage.input_tokens"] == 20 + assert attributes["gen_ai.usage.output_tokens"] == 10 + assert attributes["ai.generation_id"] == "gen-1" + + if expect_inputs: + assert '{"role": "user", "content": "hello"}' in str( + attributes[SPANDATA.AI_INPUT_MESSAGES] + ) + assert attributes[SPANDATA.AI_PREAMBLE] == "be concise" + else: + assert SPANDATA.AI_INPUT_MESSAGES not in attributes + assert SPANDATA.AI_PREAMBLE not in attributes + + if expect_outputs: + assert "the model response" in str(attributes[SPANDATA.AI_RESPONSES]) + assert "doc-1" in str(attributes["ai.citations"]) + else: + assert SPANDATA.AI_RESPONSES not in attributes + assert "ai.citations" not in attributes + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize( + "data_collection, send_default_pii, include_prompts, expect_inputs, expect_outputs", + DATA_COLLECTION_CASES, +) +def test_streaming_chat_data_collection( + sentry_init, + capture_events, + capture_items, + data_collection, + send_default_pii, + include_prompts, + expect_inputs, + expect_outputs, + span_streaming, +): + _init_with_data_collection( + sentry_init, data_collection, send_default_pii, include_prompts, span_streaming + ) + + client = Client(api_key="z") + 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", + "generation_id": "gen-1", + "citations": [ + { + "start": 0, + "end": 3, + "text": "the", + "document_ids": ["doc-1"], + } + ], + "meta": { + "billed_units": { + "output_tokens": 10, + "input_tokens": 20, + } + }, + }, + } + ), + ] + ), + ) + ) + + if span_streaming: + items = capture_items("span") + else: + events = capture_events() + + with start_transaction(name="cohere tx"): + list( + client.chat_stream( + model="some-model", + chat_history=[ChatMessage(role="SYSTEM", message="some context")], + message="hello", + preamble="be concise", + ) + ) + + if span_streaming: + sentry_sdk.flush() + assert len(items) == 1 + attributes = items[0].payload["attributes"] + else: + attributes = events[0]["spans"][0]["data"] + + assert attributes[SPANDATA.AI_MODEL_ID] == "some-model" + assert attributes["gen_ai.usage.input_tokens"] == 20 + assert attributes["gen_ai.usage.output_tokens"] == 10 + + if expect_inputs: + assert '{"role": "user", "content": "hello"}' in str( + attributes[SPANDATA.AI_INPUT_MESSAGES] + ) + assert attributes[SPANDATA.AI_PREAMBLE] == "be concise" + else: + assert SPANDATA.AI_INPUT_MESSAGES not in attributes + assert SPANDATA.AI_PREAMBLE not in attributes + + if expect_outputs: + assert "the model response" in str(attributes[SPANDATA.AI_RESPONSES]) + assert "doc-1" in str(attributes["ai.citations"]) + else: + assert SPANDATA.AI_RESPONSES not in attributes + assert "ai.citations" not in attributes + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize( + "data_collection, send_default_pii, include_prompts, expect_inputs, expect_outputs", + DATA_COLLECTION_CASES, +) +def test_embed_data_collection( + sentry_init, + capture_events, + capture_items, + data_collection, + send_default_pii, + include_prompts, + expect_inputs, + expect_outputs, + span_streaming, +): + _init_with_data_collection( + sentry_init, data_collection, send_default_pii, include_prompts, span_streaming + ) + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + 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, + } + }, + }, + ) + ) + + if span_streaming: + items = capture_items("span") + else: + events = capture_events() + + with start_transaction(name="cohere tx"): + client.embed(texts=["hello"], model="text-embedding-3-large") + + if span_streaming: + sentry_sdk.flush() + assert len(items) == 1 + attributes = items[0].payload["attributes"] + else: + attributes = events[0]["spans"][0]["data"] + + assert attributes[SPANDATA.AI_MODEL_ID] == "text-embedding-3-large" + assert attributes["gen_ai.usage.input_tokens"] == 10 + + if expect_inputs: + assert "hello" in str(attributes[SPANDATA.AI_INPUT_MESSAGES]) + else: + assert SPANDATA.AI_INPUT_MESSAGES not in attributes