From c3c34660353726815f38575f009ffb3872ac6cda Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Wed, 25 Feb 2026 13:57:30 +0100 Subject: [PATCH 1/8] 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 b17b79d894ad78d426075093bf1c4b00e49f3c9f Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Wed, 25 Feb 2026 15:21:46 +0100 Subject: [PATCH 2/8] add instrumentation for cohere v2 --- sentry_sdk/integrations/cohere.py | 12 +- sentry_sdk/integrations/cohere_v2.py | 248 +++++++++++++++++++++++++++ 2 files changed, 257 insertions(+), 3 deletions(-) create mode 100644 sentry_sdk/integrations/cohere_v2.py diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 3f0d53d099..78d4107503 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -1,7 +1,6 @@ import sys from functools import wraps -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 ( @@ -41,9 +40,10 @@ except ImportError: from cohere import StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse - COLLECTED_CHAT_PARAMS = { "model": SPANDATA.GEN_AI_REQUEST_MODEL, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, "k": SPANDATA.GEN_AI_REQUEST_TOP_K, "p": SPANDATA.GEN_AI_REQUEST_TOP_P, "seed": SPANDATA.GEN_AI_REQUEST_SEED, @@ -79,6 +79,10 @@ def setup_once() -> None: Client.embed = _wrap_embed(Client.embed) BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True) + from sentry_sdk.integrations.cohere_v2 import setup_v2 + + setup_v2(_wrap_embed) + def _capture_exception(exc: "Any") -> None: set_span_errored() @@ -156,6 +160,8 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": with capture_internal_exceptions(): set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "cohere") set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") + if model: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, model) if should_send_default_pii() and integration.include_prompts: messages = [] @@ -182,7 +188,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/sentry_sdk/integrations/cohere_v2.py b/sentry_sdk/integrations/cohere_v2.py new file mode 100644 index 0000000000..6860ce21fd --- /dev/null +++ b/sentry_sdk/integrations/cohere_v2.py @@ -0,0 +1,248 @@ +import sys +from functools import wraps + +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 typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable, Iterator + from sentry_sdk.tracing import Span + +import sentry_sdk +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import capture_internal_exceptions, event_from_exception, reraise + +from sentry_sdk.integrations.cohere import ( + CohereIntegration, + COLLECTED_CHAT_PARAMS, + _capture_exception, +) + +try: + from cohere.v2.client import V2Client as CohereV2Client + from cohere.v2.types import V2ChatResponse + from cohere.v2.types.v2chat_stream_response import MessageEndV2ChatStreamResponse + + if TYPE_CHECKING: + from cohere.v2.types import V2ChatStreamResponse + + _has_v2 = True +except ImportError: + _has_v2 = False + + +def setup_v2(wrap_embed_fn): + # type: (Callable[..., Any]) -> None + """Called from CohereIntegration.setup_once() to patch V2Client methods.""" + if not _has_v2: + return + + CohereV2Client.chat = _wrap_chat_v2(CohereV2Client.chat, streaming=False) + CohereV2Client.chat_stream = _wrap_chat_v2( + CohereV2Client.chat_stream, streaming=True + ) + CohereV2Client.embed = wrap_embed_fn(CohereV2Client.embed) + + +def _extract_messages_v2(messages): + # type: (Any) -> list[dict[str, str]] + """Extract role/content dicts from V2-style message objects.""" + result = [] + for msg in messages: + role = getattr(msg, "role", "unknown") + content = getattr(msg, "content", "") + if isinstance(content, str): + text = content + elif isinstance(content, list): + text = " ".join( + getattr(item, "text", "") for item in content if hasattr(item, "text") + ) + else: + text = str(content) if content else "" + result.append({"role": role, "content": text}) + return result + + +def _record_token_usage_v2(span, usage): + # type: (Span, Any) -> None + """Extract and record token usage from a V2 Usage object.""" + if hasattr(usage, "billed_units") and usage.billed_units is not None: + record_token_usage( + span, + input_tokens=getattr(usage.billed_units, "input_tokens", None), + output_tokens=getattr(usage.billed_units, "output_tokens", None), + ) + elif hasattr(usage, "tokens") and usage.tokens is not None: + record_token_usage( + span, + input_tokens=getattr(usage.tokens, "input_tokens", None), + output_tokens=getattr(usage.tokens, "output_tokens", None), + ) + + +def _wrap_chat_v2(f, streaming): + # type: (Callable[..., Any], bool) -> Callable[..., Any] + def collect_v2_response_fields(span, res, include_pii): + # type: (Span, V2ChatResponse, bool) -> None + if include_pii: + if ( + hasattr(res, "message") + and hasattr(res.message, "content") + and res.message.content + ): + texts = [ + item.text + for item in res.message.content + if hasattr(item, "text") + ] + if texts: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts) + + if ( + hasattr(res, "message") + and hasattr(res.message, "tool_calls") + and res.message.tool_calls + ): + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + res.message.tool_calls, + ) + + if hasattr(res, "id"): + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_ID, res.id) + + if hasattr(res, "finish_reason"): + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, res.finish_reason + ) + + if hasattr(res, "usage") and res.usage is not None: + _record_token_usage_v2(span, res.usage) + + @wraps(f) + def new_chat(*args, **kwargs): + # type: (*Any, **Any) -> Any + integration = sentry_sdk.get_client().get_integration(CohereIntegration) + + if integration is None or "messages" not in kwargs: + return f(*args, **kwargs) + + model = kwargs.get("model", "") + + span = sentry_sdk.start_span( + op=OP.GEN_AI_CHAT, + name="chat {}".format(model).strip(), + origin=CohereIntegration.origin, + ) + span.__enter__() + try: + res = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e) + span.__exit__(None, None, None) + 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 model: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, model) + + if should_send_default_pii() and integration.include_prompts: + messages = _extract_messages_v2(kwargs.get("messages", [])) + messages = normalize_message_roles(messages) + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages, + unpack=False, + ) + if "tools" in kwargs: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + kwargs["tools"], + ) + + 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, streaming) + + if streaming: + old_iterator = res + + def new_iterator(): + # type: () -> Iterator[V2ChatStreamResponse] + collected_text = [] + with capture_internal_exceptions(): + for x in old_iterator: + if ( + hasattr(x, "type") + and x.type == "content-delta" + and hasattr(x, "delta") + and x.delta is not None + ): + msg = getattr(x.delta, "message", None) + if msg is not None: + content = getattr(msg, "content", None) + if content is not None and hasattr( + content, "text" + ): + collected_text.append(content.text) + + if isinstance(x, MessageEndV2ChatStreamResponse): + include_pii = ( + should_send_default_pii() + and integration.include_prompts + ) + if include_pii and collected_text: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TEXT, + ["".join(collected_text)], + ) + if hasattr(x, "id"): + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_ID, x.id + ) + if hasattr(x, "delta") and x.delta is not None: + if hasattr(x.delta, "finish_reason"): + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, + x.delta.finish_reason, + ) + if ( + hasattr(x.delta, "usage") + and x.delta.usage is not None + ): + _record_token_usage_v2(span, x.delta.usage) + yield x + + span.__exit__(None, None, None) + + return new_iterator() + elif isinstance(res, V2ChatResponse): + collect_v2_response_fields( + span, + res, + include_pii=should_send_default_pii() + and integration.include_prompts, + ) + span.__exit__(None, None, None) + else: + set_data_normalized(span, "unknown_response", True) + span.__exit__(None, None, None) + return res + + return new_chat From 280fbd882d85195cca9258dff3b94cbe4fcb6dae Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Fri, 27 Feb 2026 12:19:38 +0100 Subject: [PATCH 3/8] wip --- sentry_sdk/ai/utils.py | 2 +- sentry_sdk/integrations/cohere.py | 21 +- sentry_sdk/integrations/cohere_v2.py | 65 ++++-- tests/integrations/cohere/test_cohere.py | 275 +++++++++++++++++++---- 4 files changed, 288 insertions(+), 75 deletions(-) diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py index 5acc501172..0f104cc8f5 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"], } diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 78d4107503..84e870b0d5 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -6,6 +6,7 @@ from sentry_sdk.ai.utils import ( set_data_normalized, normalize_message_roles, + truncate_and_annotate_messages, ) from typing import TYPE_CHECKING @@ -166,21 +167,23 @@ 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", "").lower() - if role == "chatbot": - role = "assistant" messages.append({ - "role": role, + "role": getattr(x, "role", "").lower(), "content": getattr(x, "message", ""), }) messages.append({"role": "user", "content": message}) messages = normalize_message_roles(messages) - set_data_normalized( - span, - SPANDATA.GEN_AI_REQUEST_MESSAGES, - messages, - unpack=False, + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages( + messages, span, scope ) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages_data, + unpack=False, + ) for k, v in COLLECTED_PII_CHAT_PARAMS.items(): if k in kwargs: set_data_normalized(span, v, kwargs[k]) diff --git a/sentry_sdk/integrations/cohere_v2.py b/sentry_sdk/integrations/cohere_v2.py index 6860ce21fd..dbaa54287e 100644 --- a/sentry_sdk/integrations/cohere_v2.py +++ b/sentry_sdk/integrations/cohere_v2.py @@ -6,6 +6,7 @@ from sentry_sdk.ai.utils import ( set_data_normalized, normalize_message_roles, + truncate_and_annotate_messages, ) from typing import TYPE_CHECKING @@ -16,7 +17,7 @@ import sentry_sdk 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, reraise from sentry_sdk.integrations.cohere import ( CohereIntegration, @@ -26,11 +27,24 @@ try: from cohere.v2.client import V2Client as CohereV2Client - from cohere.v2.types import V2ChatResponse - from cohere.v2.types.v2chat_stream_response import MessageEndV2ChatStreamResponse - if TYPE_CHECKING: - from cohere.v2.types import V2ChatStreamResponse + # Type locations changed between cohere versions: + # 5.13.x: cohere.types (ChatResponse, MessageEndStreamedChatResponseV2) + # 5.20+: cohere.v2.types (V2ChatResponse, MessageEndV2ChatStreamResponse) + try: + from cohere.v2.types import V2ChatResponse + from cohere.v2.types import MessageEndV2ChatStreamResponse + + if TYPE_CHECKING: + from cohere.v2.types import V2ChatStreamResponse + except ImportError: + from cohere.types import ChatResponse as V2ChatResponse + from cohere.types import ( + MessageEndStreamedChatResponseV2 as MessageEndV2ChatStreamResponse, + ) + + if TYPE_CHECKING: + from cohere.types import StreamedChatResponseV2 as V2ChatStreamResponse _has_v2 = True except ImportError: @@ -39,7 +53,12 @@ def setup_v2(wrap_embed_fn): # type: (Callable[..., Any]) -> None - """Called from CohereIntegration.setup_once() to patch V2Client methods.""" + """Called from CohereIntegration.setup_once() to patch V2Client methods. + + The embed wrapper is passed in from cohere.py to reuse the same _wrap_embed + for both V1 and V2, since the embed response format (.meta.billed_units) + is identical across both API versions. + """ if not _has_v2: return @@ -52,16 +71,25 @@ def setup_v2(wrap_embed_fn): def _extract_messages_v2(messages): # type: (Any) -> list[dict[str, str]] - """Extract role/content dicts from V2-style message objects.""" + """Extract role/content dicts from V2-style message objects. + + Handles both plain dicts and Pydantic model instances. + """ result = [] for msg in messages: - role = getattr(msg, "role", "unknown") - content = getattr(msg, "content", "") + if isinstance(msg, dict): + role = msg.get("role", "unknown") + content = msg.get("content", "") + else: + role = getattr(msg, "role", "unknown") + content = getattr(msg, "content", "") if isinstance(content, str): text = content elif isinstance(content, list): text = " ".join( - getattr(item, "text", "") for item in content if hasattr(item, "text") + (item.get("text", "") if isinstance(item, dict) else getattr(item, "text", "")) + for item in content + if (isinstance(item, dict) and "text" in item) or hasattr(item, "text") ) else: text = str(content) if content else "" @@ -138,7 +166,7 @@ def new_chat(*args, **kwargs): span = sentry_sdk.start_span( op=OP.GEN_AI_CHAT, - name="chat {}".format(model).strip(), + name=f"chat {model}".strip(), origin=CohereIntegration.origin, ) span.__enter__() @@ -160,12 +188,17 @@ def new_chat(*args, **kwargs): if should_send_default_pii() and integration.include_prompts: messages = _extract_messages_v2(kwargs.get("messages", [])) messages = normalize_message_roles(messages) - set_data_normalized( - span, - SPANDATA.GEN_AI_REQUEST_MESSAGES, - messages, - unpack=False, + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages( + messages, span, scope ) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages_data, + unpack=False, + ) if "tools" in kwargs: set_data_normalized( span, diff --git a/tests/integrations/cohere/test_cohere.py b/tests/integrations/cohere/test_cohere.py index 5018c0cca7..3da2f616ed 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -2,21 +2,32 @@ import httpx import pytest +from unittest import mock + +from httpx import Client as HTTPXClient + from cohere import Client, ChatMessage from sentry_sdk import start_transaction from sentry_sdk.consts import SPANDATA from sentry_sdk.integrations.cohere import CohereIntegration -from unittest import mock # python 3.3 and above -from httpx import Client as HTTPXClient +try: + from cohere import ClientV2 + + has_v2 = True +except ImportError: + has_v2 = False + + +# --- V1 Chat (non-streaming) --- @pytest.mark.parametrize( "send_default_pii, include_prompts", [(True, True), (True, False), (False, True), (False, False)], ) -def test_nonstreaming_chat( +def test_v1_nonstreaming_chat( sentry_init, capture_events, send_default_pii, include_prompts ): sentry_init( @@ -32,6 +43,8 @@ def test_nonstreaming_chat( 200, json={ "text": "the model response", + "generation_id": "gen-123", + "finish_reason": "COMPLETE", "meta": { "billed_units": { "output_tokens": 10, @@ -47,26 +60,21 @@ def test_nonstreaming_chat( model="some-model", chat_history=[ChatMessage(role="SYSTEM", message="some context")], message="hello", - ).text + ) - assert response == "the model response" + assert response.text == "the model response" tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] assert span["op"] == "gen_ai.chat" - assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["origin"] == "auto.ai.cohere" assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "cohere" assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False if send_default_pii and include_prompts: - assert ( - '{"role": "system", "content": "some context"}' - in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - assert ( - '{"role": "user", "content": "hello"}' - in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) + assert "hello" in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] assert "the model response" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] @@ -77,12 +85,16 @@ def test_nonstreaming_chat( assert span["data"]["gen_ai.usage.total_tokens"] == 30 -# noinspection PyTypeChecker +# --- V1 Chat (streaming) --- + + @pytest.mark.parametrize( "send_default_pii, include_prompts", [(True, True), (True, False), (False, True), (False, False)], ) -def test_streaming_chat(sentry_init, capture_events, send_default_pii, include_prompts): +def test_v1_streaming_chat( + sentry_init, capture_events, send_default_pii, include_prompts +): sentry_init( integrations=[CohereIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, @@ -104,6 +116,7 @@ def test_streaming_chat(sentry_init, capture_events, send_default_pii, include_p "finish_reason": "COMPLETE", "response": { "text": "the model response", + "generation_id": "gen-123", "meta": { "billed_units": { "output_tokens": 10, @@ -133,19 +146,14 @@ def test_streaming_chat(sentry_init, capture_events, send_default_pii, include_p assert tx["type"] == "transaction" span = tx["spans"][0] assert span["op"] == "gen_ai.chat" - assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["origin"] == "auto.ai.cohere" assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "cohere" assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True if send_default_pii and include_prompts: - assert ( - '{"role": "system", "content": "some context"}' - in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - assert ( - '{"role": "user", "content": "hello"}' - in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) + assert "hello" in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] assert "the model response" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] @@ -156,7 +164,10 @@ def test_streaming_chat(sentry_init, capture_events, send_default_pii, include_p assert span["data"]["gen_ai.usage.total_tokens"] == 30 -def test_bad_chat(sentry_init, capture_events): +# --- V1 Error --- + + +def test_v1_bad_chat(sentry_init, capture_events): sentry_init(integrations=[CohereIntegration()], traces_sample_rate=1.0) events = capture_events() @@ -171,7 +182,7 @@ def test_bad_chat(sentry_init, capture_events): assert event["level"] == "error" -def test_span_status_error(sentry_init, capture_events): +def test_v1_span_status_error(sentry_init, capture_events): sentry_init(integrations=[CohereIntegration()], traces_sample_rate=1.0) events = capture_events() @@ -190,11 +201,14 @@ def test_span_status_error(sentry_init, capture_events): assert transaction["contexts"]["trace"]["status"] == "internal_error" +# --- V1 Embed --- + + @pytest.mark.parametrize( "send_default_pii, include_prompts", [(True, True), (True, False), (False, True), (False, False)], ) -def test_embed(sentry_init, capture_events, send_default_pii, include_prompts): +def test_v1_embed(sentry_init, capture_events, send_default_pii, include_prompts): sentry_init( integrations=[CohereIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, @@ -221,7 +235,7 @@ def test_embed(sentry_init, capture_events, send_default_pii, include_prompts): ) with start_transaction(name="cohere tx"): - response = client.embed(texts=["hello"], model="text-embedding-3-large") + response = client.embed(texts=["hello"], model="embed-english-v3.0") assert len(response.embeddings[0]) == 3 @@ -229,8 +243,10 @@ def test_embed(sentry_init, capture_events, send_default_pii, include_prompts): assert tx["type"] == "transaction" span = tx["spans"][0] assert span["op"] == "gen_ai.embeddings" + assert span["origin"] == "auto.ai.cohere" 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.GEN_AI_EMBEDDINGS_INPUT] else: @@ -240,50 +256,195 @@ def test_embed(sentry_init, capture_events, send_default_pii, include_prompts): assert span["data"]["gen_ai.usage.total_tokens"] == 10 -def test_span_origin_chat(sentry_init, capture_events): +# --- V2 Chat (non-streaming) --- + + +@pytest.mark.skipif(not has_v2, reason="Cohere V2 client not available") +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (True, False), (False, True), (False, False)], +) +def test_v2_nonstreaming_chat( + sentry_init, capture_events, send_default_pii, include_prompts +): sentry_init( - integrations=[CohereIntegration()], + integrations=[CohereIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, + send_default_pii=send_default_pii, ) events = capture_events() - client = Client(api_key="z") + client = ClientV2(api_key="z") HTTPXClient.request = mock.Mock( return_value=httpx.Response( 200, json={ - "text": "the model response", - "meta": { + "id": "resp-123", + "finish_reason": "COMPLETE", + "message": { + "role": "assistant", + "content": [{"type": "text", "text": "the model response"}], + }, + "usage": { "billed_units": { - "output_tokens": 10, "input_tokens": 20, - } + "output_tokens": 10, + }, + "tokens": { + "input_tokens": 25, + "output_tokens": 15, + }, }, }, ) ) with start_transaction(name="cohere tx"): + response = client.chat( + model="some-model", + messages=[ + {"role": "system", "content": "some context"}, + {"role": "user", "content": "hello"}, + ], + ) + + assert response.message.content[0].text == "the model response" + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.chat" + assert span["origin"] == "auto.ai.cohere" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "cohere" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "some-model" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + assert span["data"][SPANDATA.GEN_AI_RESPONSE_ID] == "resp-123" + + if send_default_pii and include_prompts: + assert "hello" in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert "the model response" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] + else: + 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 + assert span["data"]["gen_ai.usage.total_tokens"] == 30 + + +# --- V2 Chat (streaming) --- + + +@pytest.mark.skipif(not has_v2, reason="Cohere V2 client not available") +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (True, False), (False, True), (False, False)], +) +def test_v2_streaming_chat( + sentry_init, capture_events, send_default_pii, include_prompts +): + sentry_init( + integrations=[CohereIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + + client = ClientV2(api_key="z") + + # SSE format: each event is "data: ...\n\n" + sse_content = "".join( + [ + 'data: {"type":"message-start","id":"resp-123"}\n', + "\n", + 'data: {"type":"content-delta","index":0,"delta":{"type":"content-delta","message":{"role":"assistant","content":{"type":"text","text":"the model "}}}}\n', + "\n", + 'data: {"type":"content-delta","index":0,"delta":{"type":"content-delta","message":{"role":"assistant","content":{"type":"text","text":"response"}}}}\n', + "\n", + 'data: {"type":"message-end","id":"resp-123","delta":{"finish_reason":"COMPLETE","usage":{"billed_units":{"input_tokens":20,"output_tokens":10},"tokens":{"input_tokens":25,"output_tokens":15}}}}\n', + "\n", + ] + ) + + HTTPXClient.send = mock.Mock( + return_value=httpx.Response( + 200, + content=sse_content, + headers={"content-type": "text/event-stream"}, + ) + ) + + with start_transaction(name="cohere tx"): + responses = list( + client.chat_stream( + model="some-model", + messages=[ + {"role": "user", "content": "hello"}, + ], + ) + ) + + assert len(responses) > 0 + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.chat" + assert span["origin"] == "auto.ai.cohere" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "cohere" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + + if send_default_pii and include_prompts: + assert "hello" in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert "the model response" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] + else: + 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 + assert span["data"]["gen_ai.usage.total_tokens"] == 30 + + +# --- V2 Error --- + + +@pytest.mark.skipif(not has_v2, reason="Cohere V2 client not available") +def test_v2_bad_chat(sentry_init, capture_events): + sentry_init(integrations=[CohereIntegration()], traces_sample_rate=1.0) + events = capture_events() + + client = ClientV2(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", - chat_history=[ChatMessage(role="SYSTEM", message="some context")], - message="hello", - ).text + messages=[{"role": "user", "content": "hello"}], + ) (event,) = events + assert event["level"] == "error" - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.cohere" +# --- V2 Embed --- -def test_span_origin_embed(sentry_init, capture_events): + +@pytest.mark.skipif(not has_v2, reason="Cohere V2 client not available") +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (True, False), (False, True), (False, False)], +) +def test_v2_embed(sentry_init, capture_events, send_default_pii, include_prompts): sentry_init( - integrations=[CohereIntegration()], + integrations=[CohereIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, + send_default_pii=send_default_pii, ) events = capture_events() - client = Client(api_key="z") + client = ClientV2(api_key="z") HTTPXClient.request = mock.Mock( return_value=httpx.Response( 200, @@ -291,7 +452,7 @@ def test_span_origin_embed(sentry_init, capture_events): "response_type": "embeddings_floats", "id": "1", "texts": ["hello"], - "embeddings": [[1.0, 2.0, 3.0]], + "embeddings": {"float": [[1.0, 2.0, 3.0]]}, "meta": { "billed_units": { "input_tokens": 10, @@ -302,9 +463,25 @@ def test_span_origin_embed(sentry_init, capture_events): ) with start_transaction(name="cohere tx"): - client.embed(texts=["hello"], model="text-embedding-3-large") + client.embed( + texts=["hello"], + model="embed-english-v3.0", + input_type="search_document", + embedding_types=["float"], + ) - (event,) = events + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.embeddings" + assert span["origin"] == "auto.ai.cohere" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "cohere" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.cohere" + if send_default_pii and include_prompts: + assert "hello" in span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + else: + 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 041844e7dd08de2da2cf14dcd4f940684d023ebe Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Fri, 27 Feb 2026 12:36:37 +0100 Subject: [PATCH 4/8] format --- sentry_sdk/integrations/cohere.py | 18 ++++++++++-------- sentry_sdk/integrations/cohere_v2.py | 18 ++++++++---------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 84e870b0d5..fc7eba4433 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -167,16 +167,16 @@ 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", []): - messages.append({ - "role": getattr(x, "role", "").lower(), - "content": getattr(x, "message", ""), - }) + messages.append( + { + "role": getattr(x, "role", "").lower(), + "content": getattr(x, "message", ""), + } + ) messages.append({"role": "user", "content": message}) messages = normalize_message_roles(messages) scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - messages, span, scope - ) + messages_data = truncate_and_annotate_messages(messages, span, scope) if messages_data is not None: set_data_normalized( span, @@ -263,7 +263,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: diff --git a/sentry_sdk/integrations/cohere_v2.py b/sentry_sdk/integrations/cohere_v2.py index dbaa54287e..0c54d6517a 100644 --- a/sentry_sdk/integrations/cohere_v2.py +++ b/sentry_sdk/integrations/cohere_v2.py @@ -87,7 +87,11 @@ def _extract_messages_v2(messages): text = content elif isinstance(content, list): text = " ".join( - (item.get("text", "") if isinstance(item, dict) else getattr(item, "text", "")) + ( + item.get("text", "") + if isinstance(item, dict) + else getattr(item, "text", "") + ) for item in content if (isinstance(item, dict) and "text" in item) or hasattr(item, "text") ) @@ -125,9 +129,7 @@ def collect_v2_response_fields(span, res, include_pii): and res.message.content ): texts = [ - item.text - for item in res.message.content - if hasattr(item, "text") + item.text for item in res.message.content if hasattr(item, "text") ] if texts: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts) @@ -189,9 +191,7 @@ def new_chat(*args, **kwargs): messages = _extract_messages_v2(kwargs.get("messages", [])) messages = normalize_message_roles(messages) scope = sentry_sdk.get_current_scope() - messages_data = truncate_and_annotate_messages( - messages, span, scope - ) + messages_data = truncate_and_annotate_messages(messages, span, scope) if messages_data is not None: set_data_normalized( span, @@ -228,9 +228,7 @@ def new_iterator(): msg = getattr(x.delta, "message", None) if msg is not None: content = getattr(msg, "content", None) - if content is not None and hasattr( - content, "text" - ): + if content is not None and hasattr(content, "text"): collected_text.append(content.text) if isinstance(x, MessageEndV2ChatStreamResponse): From 08dc0c29334521f287b98aefa1be23f17753e73d Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Fri, 27 Feb 2026 14:29:30 +0100 Subject: [PATCH 5/8] correct model --- sentry_sdk/integrations/cohere.py | 11 ++++++++++- sentry_sdk/integrations/cohere_v2.py | 2 +- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index fc7eba4433..0619a45445 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -162,14 +162,23 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "cohere") set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") if model: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, model) + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) if should_send_default_pii() and integration.include_prompts: messages = [] for x in kwargs.get("chat_history", []): +<<<<<<< HEAD messages.append( { "role": getattr(x, "role", "").lower(), +======= + role = getattr(x, "role", "").lower() + if role == "chatbot": + role = "assistant" + messages.append( + { + "role": role, +>>>>>>> c51eeb90 (correct model) "content": getattr(x, "message", ""), } ) diff --git a/sentry_sdk/integrations/cohere_v2.py b/sentry_sdk/integrations/cohere_v2.py index 0c54d6517a..4610ec06ea 100644 --- a/sentry_sdk/integrations/cohere_v2.py +++ b/sentry_sdk/integrations/cohere_v2.py @@ -185,7 +185,7 @@ def new_chat(*args, **kwargs): set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "cohere") set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") if model: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, model) + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) if should_send_default_pii() and integration.include_prompts: messages = _extract_messages_v2(kwargs.get("messages", [])) From 54e99fe0197bbe2d1b589c6ccf437b96dde4f241 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Fri, 27 Feb 2026 14:31:43 +0100 Subject: [PATCH 6/8] wip --- sentry_sdk/integrations/cohere.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 0619a45445..8ac20b6ed8 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -21,6 +21,7 @@ from sentry_sdk.scope import should_send_default_pii from sentry_sdk.integrations import DidNotEnable, Integration from sentry_sdk.utils import capture_internal_exceptions, event_from_exception, reraise +from sentry_sdk.integrations.cohere_v2 import setup_v2 try: from cohere.client import Client @@ -79,9 +80,6 @@ def setup_once() -> None: BaseCohere.chat = _wrap_chat(BaseCohere.chat, streaming=False) Client.embed = _wrap_embed(Client.embed) BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True) - - from sentry_sdk.integrations.cohere_v2 import setup_v2 - setup_v2(_wrap_embed) @@ -167,18 +165,9 @@ 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", []): -<<<<<<< HEAD messages.append( { "role": getattr(x, "role", "").lower(), -======= - role = getattr(x, "role", "").lower() - if role == "chatbot": - role = "assistant" - messages.append( - { - "role": role, ->>>>>>> c51eeb90 (correct model) "content": getattr(x, "message", ""), } ) From ec237b491d066cc52bd88f95075ceec7501ecf0a Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Fri, 27 Feb 2026 14:46:47 +0100 Subject: [PATCH 7/8] move to separate folder --- sentry_sdk/integrations/cohere/__init__.py | 127 ++++++++++++++ .../integrations/{cohere.py => cohere/v1.py} | 161 +++++------------- .../{cohere_v2.py => cohere/v2.py} | 2 +- 3 files changed, 170 insertions(+), 120 deletions(-) create mode 100644 sentry_sdk/integrations/cohere/__init__.py rename sentry_sdk/integrations/{cohere.py => cohere/v1.py} (56%) rename sentry_sdk/integrations/{cohere_v2.py => cohere/v2.py} (99%) diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py new file mode 100644 index 0000000000..f8a26f1fc8 --- /dev/null +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -0,0 +1,127 @@ +import sys +from functools import wraps + +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 + +from typing import TYPE_CHECKING + +from sentry_sdk.tracing_utils import set_span_errored + +if TYPE_CHECKING: + from typing import Any, Callable + +import sentry_sdk +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.utils import capture_internal_exceptions, event_from_exception, reraise + +try: + from cohere import __version__ as cohere_version # noqa: F401 +except ImportError: + raise DidNotEnable("Cohere not installed") + +COLLECTED_CHAT_PARAMS = { + "model": SPANDATA.GEN_AI_REQUEST_MODEL, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "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, +} + + +class CohereIntegration(Integration): + identifier = "cohere" + origin = f"auto.ai.{identifier}" + + def __init__(self, include_prompts=True): + # type: (bool) -> None + self.include_prompts = include_prompts + + @staticmethod + def setup_once(): + # type: () -> None + # Lazy imports to avoid circular dependencies: + # v1/v2 import COLLECTED_CHAT_PARAMS and _capture_exception from this module. + from sentry_sdk.integrations.cohere.v1 import setup_v1 + from sentry_sdk.integrations.cohere.v2 import setup_v2 + + setup_v1(_wrap_embed) + setup_v2(_wrap_embed) + + +def _capture_exception(exc): + # type: (Any) -> None + set_span_errored() + + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "cohere", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _wrap_embed(f): + # type: (Callable[..., Any]) -> Callable[..., Any] + @wraps(f) + def new_embed(*args, **kwargs): + # type: (*Any, **Any) -> Any + integration = sentry_sdk.get_client().get_integration(CohereIntegration) + if integration is None: + return f(*args, **kwargs) + + model = kwargs.get("model", "") + + with sentry_sdk.start_span( + 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.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.GEN_AI_EMBEDDINGS_INPUT, kwargs["texts"] + ) + + if "model" in kwargs: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MODEL, kwargs["model"] + ) + try: + res = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e) + reraise(*exc_info) + if ( + hasattr(res, "meta") + and hasattr(res.meta, "billed_units") + and hasattr(res.meta.billed_units, "input_tokens") + ): + record_token_usage( + span, + input_tokens=res.meta.billed_units.input_tokens, + total_tokens=res.meta.billed_units.input_tokens, + ) + return res + + return new_embed diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere/v1.py similarity index 56% rename from sentry_sdk/integrations/cohere.py rename to sentry_sdk/integrations/cohere/v1.py index 8ac20b6ed8..6e9f5b42cc 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -11,47 +11,19 @@ from typing import TYPE_CHECKING -from sentry_sdk.tracing_utils import set_span_errored - if TYPE_CHECKING: from typing import Any, Callable, Iterator from sentry_sdk.tracing import Span import sentry_sdk from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.integrations import DidNotEnable, Integration -from sentry_sdk.utils import capture_internal_exceptions, event_from_exception, reraise -from sentry_sdk.integrations.cohere_v2 import setup_v2 - -try: - from cohere.client import Client - from cohere.base_client import BaseCohere - from cohere import ( - ChatStreamEndEvent, - NonStreamedChatResponse, - ) - - if TYPE_CHECKING: - from cohere import StreamedChatResponse -except ImportError: - raise DidNotEnable("Cohere not installed") +from sentry_sdk.utils import capture_internal_exceptions, reraise -try: - # cohere 5.9.3+ - from cohere import StreamEndStreamedChatResponse -except ImportError: - from cohere import StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse - -COLLECTED_CHAT_PARAMS = { - "model": SPANDATA.GEN_AI_REQUEST_MODEL, - "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, - "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, - "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, -} +from sentry_sdk.integrations.cohere import ( + CohereIntegration, + COLLECTED_CHAT_PARAMS, + _capture_exception, +) COLLECTED_PII_CHAT_PARAMS = { "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, @@ -68,36 +40,44 @@ } -class CohereIntegration(Integration): - identifier = "cohere" - origin = f"auto.ai.{identifier}" - - def __init__(self: "CohereIntegration", include_prompts: bool = True) -> None: - self.include_prompts = include_prompts +def setup_v1(wrap_embed_fn): + # type: (Callable[..., Any]) -> None + """Called from CohereIntegration.setup_once() to patch V1 Client methods.""" + try: + from cohere.client import Client + from cohere.base_client import BaseCohere + except ImportError: + return - @staticmethod - def setup_once() -> None: - BaseCohere.chat = _wrap_chat(BaseCohere.chat, streaming=False) - Client.embed = _wrap_embed(Client.embed) - BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True) - setup_v2(_wrap_embed) + BaseCohere.chat = _wrap_chat(BaseCohere.chat, streaming=False) + BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True) + Client.embed = wrap_embed_fn(Client.embed) -def _capture_exception(exc: "Any") -> None: - set_span_errored() +def _wrap_chat(f, streaming): + # type: (Callable[..., Any], bool) -> Callable[..., Any] - event, hint = event_from_exception( - exc, - client_options=sentry_sdk.get_client().options, - mechanism={"type": "cohere", "handled": False}, - ) - sentry_sdk.capture_event(event, hint=hint) + try: + from cohere import ( + ChatStreamEndEvent, + NonStreamedChatResponse, + ) + if TYPE_CHECKING: + from cohere import StreamedChatResponse + except ImportError: + return f + + try: + # cohere 5.9.3+ + from cohere import StreamEndStreamedChatResponse + except ImportError: + from cohere import ( + StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse, + ) -def _wrap_chat(f: "Callable[..., Any]", streaming: bool) -> "Callable[..., Any]": - def collect_chat_response_fields( - span: "Span", res: "NonStreamedChatResponse", include_pii: bool - ) -> None: + def collect_chat_response_fields(span, res, include_pii): + # type: (Span, NonStreamedChatResponse, bool) -> None if include_pii: if hasattr(res, "text"): set_data_normalized( @@ -128,7 +108,8 @@ def collect_chat_response_fields( ) @wraps(f) - def new_chat(*args: "Any", **kwargs: "Any") -> "Any": + def new_chat(*args, **kwargs): + # type: (*Any, **Any) -> Any integration = sentry_sdk.get_client().get_integration(CohereIntegration) if ( @@ -194,7 +175,8 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": if streaming: old_iterator = res - def new_iterator() -> "Iterator[StreamedChatResponse]": + def new_iterator(): + # type: () -> Iterator[StreamedChatResponse] with capture_internal_exceptions(): for x in old_iterator: if isinstance(x, ChatStreamEndEvent) or isinstance( @@ -225,62 +207,3 @@ def new_iterator() -> "Iterator[StreamedChatResponse]": return res return new_chat - - -def _wrap_embed(f: "Callable[..., Any]") -> "Callable[..., Any]": - @wraps(f) - def new_embed(*args: "Any", **kwargs: "Any") -> "Any": - integration = sentry_sdk.get_client().get_integration(CohereIntegration) - if integration is None: - return f(*args, **kwargs) - - model = kwargs.get("model", "") - - with sentry_sdk.start_span( - 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.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.GEN_AI_EMBEDDINGS_INPUT, kwargs["texts"] - ) - - if "model" in kwargs: - set_data_normalized( - span, SPANDATA.GEN_AI_REQUEST_MODEL, kwargs["model"] - ) - try: - res = f(*args, **kwargs) - except Exception as e: - exc_info = sys.exc_info() - with capture_internal_exceptions(): - _capture_exception(e) - reraise(*exc_info) - if ( - hasattr(res, "meta") - and hasattr(res.meta, "billed_units") - and hasattr(res.meta.billed_units, "input_tokens") - ): - record_token_usage( - span, - input_tokens=res.meta.billed_units.input_tokens, - total_tokens=res.meta.billed_units.input_tokens, - ) - return res - - return new_embed diff --git a/sentry_sdk/integrations/cohere_v2.py b/sentry_sdk/integrations/cohere/v2.py similarity index 99% rename from sentry_sdk/integrations/cohere_v2.py rename to sentry_sdk/integrations/cohere/v2.py index 4610ec06ea..d26b9ee3c3 100644 --- a/sentry_sdk/integrations/cohere_v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -55,7 +55,7 @@ def setup_v2(wrap_embed_fn): # type: (Callable[..., Any]) -> None """Called from CohereIntegration.setup_once() to patch V2Client methods. - The embed wrapper is passed in from cohere.py to reuse the same _wrap_embed + The embed wrapper is passed in from __init__.py to reuse the same _wrap_embed for both V1 and V2, since the embed response format (.meta.billed_units) is identical across both API versions. """ From d7814b00ce4918245ce741a2c3d4081236dbdeeb Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Fri, 27 Feb 2026 15:04:01 +0100 Subject: [PATCH 8/8] wip --- sentry_sdk/integrations/cohere/v1.py | 1 + sentry_sdk/integrations/cohere/v2.py | 1 + 2 files changed, 2 insertions(+) diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index 6e9f5b42cc..c0e33cfbc7 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -142,6 +142,7 @@ def new_chat(*args, **kwargs): set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") if model: set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, model) if should_send_default_pii() and integration.include_prompts: messages = [] diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index d26b9ee3c3..0a12828462 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -186,6 +186,7 @@ def new_chat(*args, **kwargs): set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") if model: set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, model) if should_send_default_pii() and integration.include_prompts: messages = _extract_messages_v2(kwargs.get("messages", []))