feat(cohere): Support data_collection filtering for inputs and outputs (#7289)

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

Author: ericapisani

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 (

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