From 9672d8d7a1e21cea3414f3009722f9a99c5e8a63 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Wed, 25 Feb 2026 13:57:30 +0100 Subject: [PATCH 1/5] feat(cohere): upgrade integration from ai to gen_ai --- sentry_sdk/integrations/cohere.py | 102 ++++++++++++----------- tests/integrations/cohere/test_cohere.py | 40 +++++---- 2 files changed, 78 insertions(+), 64 deletions(-) diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index f45a02f2b5..3f0d53d099 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -3,8 +3,11 @@ from sentry_sdk import consts from sentry_sdk.ai.monitoring import record_token_usage -from sentry_sdk.consts import SPANDATA -from sentry_sdk.ai.utils import set_data_normalized +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.ai.utils import ( + set_data_normalized, + normalize_message_roles, +) from typing import TYPE_CHECKING @@ -40,32 +43,26 @@ COLLECTED_CHAT_PARAMS = { - "model": SPANDATA.AI_MODEL_ID, - "k": SPANDATA.AI_TOP_K, - "p": SPANDATA.AI_TOP_P, - "seed": SPANDATA.AI_SEED, - "frequency_penalty": SPANDATA.AI_FREQUENCY_PENALTY, - "presence_penalty": SPANDATA.AI_PRESENCE_PENALTY, - "raw_prompting": SPANDATA.AI_RAW_PROMPTING, + "model": SPANDATA.GEN_AI_REQUEST_MODEL, + "k": SPANDATA.GEN_AI_REQUEST_TOP_K, + "p": SPANDATA.GEN_AI_REQUEST_TOP_P, + "seed": SPANDATA.GEN_AI_REQUEST_SEED, + "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, + "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, } COLLECTED_PII_CHAT_PARAMS = { - "tools": SPANDATA.AI_TOOLS, - "preamble": SPANDATA.AI_PREAMBLE, + "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + "preamble": SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, } COLLECTED_CHAT_RESP_ATTRS = { - "generation_id": SPANDATA.AI_GENERATION_ID, - "is_search_required": SPANDATA.AI_SEARCH_REQUIRED, - "finish_reason": SPANDATA.AI_FINISH_REASON, + "generation_id": SPANDATA.GEN_AI_RESPONSE_ID, + "finish_reason": SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, } COLLECTED_PII_CHAT_RESP_ATTRS = { - "citations": SPANDATA.AI_CITATIONS, - "documents": SPANDATA.AI_DOCUMENTS, - "search_queries": SPANDATA.AI_SEARCH_QUERIES, - "search_results": SPANDATA.AI_SEARCH_RESULTS, - "tool_calls": SPANDATA.AI_TOOL_CALLS, + "tool_calls": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, } @@ -102,16 +99,16 @@ def collect_chat_response_fields( if hasattr(res, "text"): set_data_normalized( span, - SPANDATA.AI_RESPONSES, + SPANDATA.GEN_AI_RESPONSE_TEXT, [res.text], ) - for pii_attr in COLLECTED_PII_CHAT_RESP_ATTRS: - if hasattr(res, pii_attr): - set_data_normalized(span, "ai." + pii_attr, getattr(res, pii_attr)) + for attr, spandata_key in COLLECTED_PII_CHAT_RESP_ATTRS.items(): + if hasattr(res, attr): + set_data_normalized(span, spandata_key, getattr(res, attr)) - for attr in COLLECTED_CHAT_RESP_ATTRS: + for attr, spandata_key in COLLECTED_CHAT_RESP_ATTRS.items(): if hasattr(res, attr): - set_data_normalized(span, "ai." + attr, getattr(res, attr)) + set_data_normalized(span, spandata_key, getattr(res, attr)) if hasattr(res, "meta"): if hasattr(res.meta, "billed_units"): @@ -127,9 +124,6 @@ def collect_chat_response_fields( output_tokens=res.meta.tokens.output_tokens, ) - if hasattr(res.meta, "warnings"): - set_data_normalized(span, SPANDATA.AI_WARNINGS, res.meta.warnings) - @wraps(f) def new_chat(*args: "Any", **kwargs: "Any") -> "Any": integration = sentry_sdk.get_client().get_integration(CohereIntegration) @@ -142,10 +136,11 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": return f(*args, **kwargs) message = kwargs.get("message") + model = kwargs.get("model", "") span = sentry_sdk.start_span( - op=consts.OP.COHERE_CHAT_COMPLETIONS_CREATE, - name="cohere.client.Chat", + op=OP.GEN_AI_CHAT, + name=f"chat {model}".strip(), origin=CohereIntegration.origin, ) span.__enter__() @@ -159,20 +154,26 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": reraise(*exc_info) with capture_internal_exceptions(): + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "cohere") + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") + if should_send_default_pii() and integration.include_prompts: + messages = [] + for x in kwargs.get("chat_history", []): + role = getattr(x, "role", "").lower() + if role == "chatbot": + role = "assistant" + messages.append({ + "role": role, + "content": getattr(x, "message", ""), + }) + messages.append({"role": "user", "content": message}) + messages = normalize_message_roles(messages) set_data_normalized( span, - SPANDATA.AI_INPUT_MESSAGES, - list( - map( - lambda x: { - "role": getattr(x, "role", "").lower(), - "content": getattr(x, "message", ""), - }, - kwargs.get("chat_history", []), - ) - ) - + [{"role": "user", "content": message}], + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages, + unpack=False, ) for k, v in COLLECTED_PII_CHAT_PARAMS.items(): if k in kwargs: @@ -181,7 +182,7 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": for k, v in COLLECTED_CHAT_PARAMS.items(): if k in kwargs: set_data_normalized(span, v, kwargs[k]) - set_data_normalized(span, SPANDATA.AI_STREAMING, False) + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_STREAMING, False) if streaming: old_iterator = res @@ -226,27 +227,34 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any": if integration is None: return f(*args, **kwargs) + model = kwargs.get("model", "") + with sentry_sdk.start_span( - op=consts.OP.COHERE_EMBEDDINGS_CREATE, - name="Cohere Embedding Creation", + op=OP.GEN_AI_EMBEDDINGS, + name=f"embeddings {model}".strip(), origin=CohereIntegration.origin, ) as span: + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "cohere") + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + if "texts" in kwargs and ( should_send_default_pii() and integration.include_prompts ): if isinstance(kwargs["texts"], str): - set_data_normalized(span, SPANDATA.AI_TEXTS, [kwargs["texts"]]) + set_data_normalized( + span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, [kwargs["texts"]] + ) elif ( isinstance(kwargs["texts"], list) and len(kwargs["texts"]) > 0 and isinstance(kwargs["texts"][0], str) ): set_data_normalized( - span, SPANDATA.AI_INPUT_MESSAGES, kwargs["texts"] + span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, kwargs["texts"] ) if "model" in kwargs: - set_data_normalized(span, SPANDATA.AI_MODEL_ID, kwargs["model"]) + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, kwargs["model"]) try: res = f(*args, **kwargs) except Exception as e: diff --git a/tests/integrations/cohere/test_cohere.py b/tests/integrations/cohere/test_cohere.py index 9ff56ed697..5018c0cca7 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -53,22 +53,24 @@ def test_nonstreaming_chat( tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] - assert span["op"] == "ai.chat_completions.create.cohere" - assert span["data"][SPANDATA.AI_MODEL_ID] == "some-model" + assert span["op"] == "gen_ai.chat" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "cohere" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" if send_default_pii and include_prompts: assert ( '{"role": "system", "content": "some context"}' - in span["data"][SPANDATA.AI_INPUT_MESSAGES] + in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] ) assert ( '{"role": "user", "content": "hello"}' - in span["data"][SPANDATA.AI_INPUT_MESSAGES] + in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] ) - assert "the model response" in span["data"][SPANDATA.AI_RESPONSES] + assert "the model response" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] else: - assert SPANDATA.AI_INPUT_MESSAGES not in span["data"] - assert SPANDATA.AI_RESPONSES not in span["data"] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] assert span["data"]["gen_ai.usage.output_tokens"] == 10 assert span["data"]["gen_ai.usage.input_tokens"] == 20 @@ -130,22 +132,24 @@ def test_streaming_chat(sentry_init, capture_events, send_default_pii, include_p tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] - assert span["op"] == "ai.chat_completions.create.cohere" - assert span["data"][SPANDATA.AI_MODEL_ID] == "some-model" + assert span["op"] == "gen_ai.chat" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "cohere" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" if send_default_pii and include_prompts: assert ( '{"role": "system", "content": "some context"}' - in span["data"][SPANDATA.AI_INPUT_MESSAGES] + in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] ) assert ( '{"role": "user", "content": "hello"}' - in span["data"][SPANDATA.AI_INPUT_MESSAGES] + in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] ) - assert "the model response" in span["data"][SPANDATA.AI_RESPONSES] + assert "the model response" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] else: - assert SPANDATA.AI_INPUT_MESSAGES not in span["data"] - assert SPANDATA.AI_RESPONSES not in span["data"] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] assert span["data"]["gen_ai.usage.output_tokens"] == 10 assert span["data"]["gen_ai.usage.input_tokens"] == 20 @@ -224,11 +228,13 @@ def test_embed(sentry_init, capture_events, send_default_pii, include_prompts): tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] - assert span["op"] == "ai.embeddings.create.cohere" + assert span["op"] == "gen_ai.embeddings" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "cohere" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" if send_default_pii and include_prompts: - assert "hello" in span["data"][SPANDATA.AI_INPUT_MESSAGES] + assert "hello" in span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] else: - assert SPANDATA.AI_INPUT_MESSAGES not in span["data"] + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["data"] assert span["data"]["gen_ai.usage.input_tokens"] == 10 assert span["data"]["gen_ai.usage.total_tokens"] == 10 From 2b2d3a6a3b896630b49d384a99bf04defd2bbd42 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Fri, 6 Mar 2026 13:55:45 +0100 Subject: [PATCH 2/5] format --- sentry_sdk/integrations/cohere.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 3f0d53d099..2a8ef6e720 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -163,10 +163,12 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": role = getattr(x, "role", "").lower() if role == "chatbot": role = "assistant" - messages.append({ - "role": role, - "content": getattr(x, "message", ""), - }) + messages.append( + { + "role": role, + "content": getattr(x, "message", ""), + } + ) messages.append({"role": "user", "content": message}) messages = normalize_message_roles(messages) set_data_normalized( @@ -254,7 +256,9 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any": ) if "model" in kwargs: - set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, kwargs["model"]) + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MODEL, kwargs["model"] + ) try: res = f(*args, **kwargs) except Exception as e: From b6346a3bf2f0849cb1f3c14ff465ccaf8c533ccd Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Mon, 9 Mar 2026 14:09:14 +0100 Subject: [PATCH 3/5] fix role mapping --- sentry_sdk/ai/utils.py | 4 ++-- sentry_sdk/integrations/cohere.py | 14 ++++---------- 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py index 5acc501172..426fa42e3f 100644 --- a/sentry_sdk/ai/utils.py +++ b/sentry_sdk/ai/utils.py @@ -30,7 +30,7 @@ class GEN_AI_ALLOWED_MESSAGE_ROLES: GEN_AI_MESSAGE_ROLE_REVERSE_MAPPING = { GEN_AI_ALLOWED_MESSAGE_ROLES.SYSTEM: ["system"], GEN_AI_ALLOWED_MESSAGE_ROLES.USER: ["user", "human"], - GEN_AI_ALLOWED_MESSAGE_ROLES.ASSISTANT: ["assistant", "ai"], + GEN_AI_ALLOWED_MESSAGE_ROLES.ASSISTANT: ["assistant", "ai", "chatbot"], GEN_AI_ALLOWED_MESSAGE_ROLES.TOOL: ["tool", "tool_call"], } @@ -503,7 +503,7 @@ def normalize_message_role(role: str) -> str: Normalize a message role to one of the 4 allowed gen_ai role values. Maps "ai" -> "assistant" and keeps other standard roles unchanged. """ - return GEN_AI_MESSAGE_ROLE_MAPPING.get(role, role) + return GEN_AI_MESSAGE_ROLE_MAPPING.get(role.lower(), role) def normalize_message_roles(messages: "list[dict[str, Any]]") -> "list[dict[str, Any]]": diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 2a8ef6e720..c1d1458de3 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -4,10 +4,7 @@ from sentry_sdk import consts from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.ai.utils import ( - set_data_normalized, - normalize_message_roles, -) +from sentry_sdk.ai.utils import set_data_normalized, normalize_message_roles from typing import TYPE_CHECKING @@ -154,18 +151,15 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": reraise(*exc_info) with capture_internal_exceptions(): - set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "cohere") - set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") + span.set_data(SPANDATA.GEN_AI_SYSTEM, "cohere") + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") if should_send_default_pii() and integration.include_prompts: messages = [] for x in kwargs.get("chat_history", []): - role = getattr(x, "role", "").lower() - if role == "chatbot": - role = "assistant" messages.append( { - "role": role, + "role": getattr(x, "role", ""), "content": getattr(x, "message", ""), } ) From e34424ef962bc674b941a419164f56401aa322fd Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Mon, 9 Mar 2026 14:32:44 +0100 Subject: [PATCH 4/5] dont hardcode streaming parameter --- sentry_sdk/integrations/cohere.py | 2 +- tests/integrations/cohere/test_cohere.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index c1d1458de3..20fc57b266 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -178,7 +178,7 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": for k, v in COLLECTED_CHAT_PARAMS.items(): if k in kwargs: set_data_normalized(span, v, kwargs[k]) - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_STREAMING, False) + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_STREAMING, streaming) if streaming: old_iterator = res diff --git a/tests/integrations/cohere/test_cohere.py b/tests/integrations/cohere/test_cohere.py index 5018c0cca7..a87f42e294 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -72,6 +72,7 @@ def test_nonstreaming_chat( assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False assert span["data"]["gen_ai.usage.output_tokens"] == 10 assert span["data"]["gen_ai.usage.input_tokens"] == 20 assert span["data"]["gen_ai.usage.total_tokens"] == 30 @@ -151,6 +152,7 @@ def test_streaming_chat(sentry_init, capture_events, send_default_pii, include_p assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True assert span["data"]["gen_ai.usage.output_tokens"] == 10 assert span["data"]["gen_ai.usage.input_tokens"] == 20 assert span["data"]["gen_ai.usage.total_tokens"] == 30 From 4a737ca52b456ea9df917374633fad7b92ba4e28 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Mon, 9 Mar 2026 15:34:49 +0100 Subject: [PATCH 5/5] do cohere specific mapping of roles --- sentry_sdk/ai/utils.py | 4 ++-- sentry_sdk/integrations/cohere.py | 14 +++++++++++--- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py index 426fa42e3f..5acc501172 100644 --- a/sentry_sdk/ai/utils.py +++ b/sentry_sdk/ai/utils.py @@ -30,7 +30,7 @@ class GEN_AI_ALLOWED_MESSAGE_ROLES: GEN_AI_MESSAGE_ROLE_REVERSE_MAPPING = { GEN_AI_ALLOWED_MESSAGE_ROLES.SYSTEM: ["system"], GEN_AI_ALLOWED_MESSAGE_ROLES.USER: ["user", "human"], - GEN_AI_ALLOWED_MESSAGE_ROLES.ASSISTANT: ["assistant", "ai", "chatbot"], + GEN_AI_ALLOWED_MESSAGE_ROLES.ASSISTANT: ["assistant", "ai"], GEN_AI_ALLOWED_MESSAGE_ROLES.TOOL: ["tool", "tool_call"], } @@ -503,7 +503,7 @@ def normalize_message_role(role: str) -> str: Normalize a message role to one of the 4 allowed gen_ai role values. Maps "ai" -> "assistant" and keeps other standard roles unchanged. """ - return GEN_AI_MESSAGE_ROLE_MAPPING.get(role.lower(), role) + return GEN_AI_MESSAGE_ROLE_MAPPING.get(role, role) def normalize_message_roles(messages: "list[dict[str, Any]]") -> "list[dict[str, Any]]": diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 20fc57b266..458a89ce84 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -4,7 +4,7 @@ from sentry_sdk import consts from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.ai.utils import set_data_normalized, normalize_message_roles +from sentry_sdk.ai.utils import set_data_normalized from typing import TYPE_CHECKING @@ -39,6 +39,14 @@ from cohere import StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse +COHERE_ROLE_MAPPING = { + "SYSTEM": "system", + "USER": "user", + "CHATBOT": "assistant", + "TOOL": "tool", +} + + COLLECTED_CHAT_PARAMS = { "model": SPANDATA.GEN_AI_REQUEST_MODEL, "k": SPANDATA.GEN_AI_REQUEST_TOP_K, @@ -157,14 +165,14 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": if should_send_default_pii() and integration.include_prompts: messages = [] for x in kwargs.get("chat_history", []): + role = getattr(x, "role", "") messages.append( { - "role": getattr(x, "role", ""), + "role": COHERE_ROLE_MAPPING.get(role, role), "content": getattr(x, "message", ""), } ) messages.append({"role": "user", "content": message}) - messages = normalize_message_roles(messages) set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES,