Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ericapisani wants to merge 1 commit into from
+343 −9
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 18 additions & 9 deletions sentry_sdk/integrations/cohere.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand All @@ -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:
Expand Down Expand Up @@ -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 (
Expand Down
325 changes: 325 additions & 0 deletions tests/integrations/cohere/test_cohere.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading