From c3c34660353726815f38575f009ffb3872ac6cda Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Wed, 25 Feb 2026 13:57:30 +0100 Subject: [PATCH 01/29] 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 02/29] 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 03/29] 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 04/29] 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 05/29] 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 06/29] 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 07/29] 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 08/29] 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", [])) From a1f4ab7ba7699b30e6c437b7dbcf0694deb577e2 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Mon, 2 Mar 2026 13:46:17 +0100 Subject: [PATCH 09/29] make sure to capture exceptions correctly --- sentry_sdk/integrations/cohere/v1.py | 28 +++++----- sentry_sdk/integrations/cohere/v2.py | 80 +++++++++++++++------------- 2 files changed, 58 insertions(+), 50 deletions(-) diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index c0e33cfbc7..ead4b2c023 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -178,20 +178,24 @@ def new_chat(*args, **kwargs): def new_iterator(): # type: () -> Iterator[StreamedChatResponse] - with capture_internal_exceptions(): + try: for x in old_iterator: - if isinstance(x, ChatStreamEndEvent) or isinstance( - x, StreamEndStreamedChatResponse - ): - collect_chat_response_fields( - span, - x.response, - include_pii=should_send_default_pii() - and integration.include_prompts, - ) + with capture_internal_exceptions(): + if isinstance(x, ChatStreamEndEvent) or isinstance( + x, StreamEndStreamedChatResponse + ): + collect_chat_response_fields( + span, + x.response, + include_pii=should_send_default_pii() + and integration.include_prompts, + ) yield x - - span.__exit__(None, None, None) + except Exception as exc: + _capture_exception(exc) + raise + finally: + span.__exit__(None, None, None) return new_iterator() elif isinstance(res, NonStreamedChatResponse): diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index 0a12828462..0aa2353346 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -218,50 +218,54 @@ def new_chat(*args, **kwargs): def new_iterator(): # type: () -> Iterator[V2ChatStreamResponse] collected_text = [] - with capture_internal_exceptions(): + try: 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)], + with capture_internal_exceptions(): + 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 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"): + if include_pii and collected_text: set_data_normalized( span, - SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, - x.delta.finish_reason, + 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, "usage") - and x.delta.usage is not None - ): - _record_token_usage_v2(span, x.delta.usage) + 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) + except Exception as exc: + _capture_exception(exc) + raise + finally: + span.__exit__(None, None, None) return new_iterator() elif isinstance(res, V2ChatResponse): From 130f509ce15886e800ce2f4706278b2f36d95912 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Wed, 4 Mar 2026 11:03:42 +0100 Subject: [PATCH 10/29] use context managers where appropriate --- sentry_sdk/integrations/cohere/__init__.py | 32 ++-- sentry_sdk/integrations/cohere/v1.py | 168 ++++++++++----------- sentry_sdk/integrations/cohere/v2.py | 131 ++++++++-------- 3 files changed, 159 insertions(+), 172 deletions(-) diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index f8a26f1fc8..8642cf751c 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -34,6 +34,15 @@ } +def _normalize_embedding_input(texts): + # type: (Any) -> Any + if isinstance(texts, list): + return texts + if isinstance(texts, tuple): + return list(texts) + return [texts] + + class CohereIntegration(Integration): identifier = "cohere" origin = f"auto.ai.{identifier}" @@ -85,26 +94,23 @@ def new_embed(*args, **kwargs): set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "cohere") set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + # attach embeddings input + # should this be truncated and annotated? 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"] - ) + set_data_normalized( + span, + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + _normalize_embedding_input(kwargs["texts"]), + ) if "model" in kwargs: set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MODEL, kwargs["model"] ) + + # call the function and capture any exceptions try: res = f(*args, **kwargs) except Exception as e: @@ -112,6 +118,8 @@ def new_embed(*args, **kwargs): with capture_internal_exceptions(): _capture_exception(e) reraise(*exc_info) + + # record token usage if ( hasattr(res, "meta") and hasattr(res.meta, "billed_units") diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index ead4b2c023..b8e87e5e7c 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -13,6 +13,7 @@ if TYPE_CHECKING: from typing import Any, Callable, Iterator + from cohere import StreamedChatResponse from sentry_sdk.tracing import Span import sentry_sdk @@ -25,6 +26,23 @@ _capture_exception, ) +try: + from cohere import ( + ChatStreamEndEvent, + NonStreamedChatResponse, + ) + + try: + from cohere import StreamEndStreamedChatResponse + except ImportError: + from cohere import ( + StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse, + ) + + _has_chat_types = True +except ImportError: + _has_chat_types = False + COLLECTED_PII_CHAT_PARAMS = { "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, "preamble": SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, @@ -56,26 +74,9 @@ def setup_v1(wrap_embed_fn): def _wrap_chat(f, streaming): # type: (Callable[..., Any], bool) -> Callable[..., Any] - - try: - from cohere import ( - ChatStreamEndEvent, - NonStreamedChatResponse, - ) - - if TYPE_CHECKING: - from cohere import StreamedChatResponse - except ImportError: + if not _has_chat_types: return f - try: - # cohere 5.9.3+ - from cohere import StreamEndStreamedChatResponse - except ImportError: - from cohere import ( - StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse, - ) - def collect_chat_response_fields(span, res, include_pii): # type: (Span, NonStreamedChatResponse, bool) -> None if include_pii: @@ -122,65 +123,61 @@ def new_chat(*args, **kwargs): message = kwargs.get("message") model = kwargs.get("model", "") - span = sentry_sdk.start_span( + with sentry_sdk.start_span( op=OP.GEN_AI_CHAT, name=f"chat {model}".strip(), origin=CohereIntegration.origin, - ) - span.__enter__() - try: - res = f(*args, **kwargs) - except Exception as e: - exc_info = sys.exc_info() + ) as span: + try: + res = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e) + reraise(*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_REQUEST_MODEL, model) - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, model) - - 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": "user", "content": message}) - messages = normalize_message_roles(messages) - 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(): + 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_REQUEST_MODEL, model) + + 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": "user", "content": message}) + messages = normalize_message_roles(messages) + 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]) + + 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) - 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 - if streaming: - old_iterator = res - - def new_iterator(): - # type: () -> Iterator[StreamedChatResponse] - try: - for x in old_iterator: - with capture_internal_exceptions(): + def new_iterator(): + # type: () -> Iterator[StreamedChatResponse] + with capture_internal_exceptions(): + for x in old_iterator: if isinstance(x, ChatStreamEndEvent) or isinstance( x, StreamEndStreamedChatResponse ): @@ -190,25 +187,18 @@ def new_iterator(): include_pii=should_send_default_pii() and integration.include_prompts, ) - yield x - except Exception as exc: - _capture_exception(exc) - raise - finally: - span.__exit__(None, None, None) - - return new_iterator() - elif isinstance(res, NonStreamedChatResponse): - collect_chat_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 + yield x + + return new_iterator() + elif isinstance(res, NonStreamedChatResponse): + collect_chat_response_fields( + span, + res, + include_pii=should_send_default_pii() + and integration.include_prompts, + ) + else: + set_data_normalized(span, "unknown_response", True) + return res return new_chat diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index 0aa2353346..fa1313e6a5 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -166,61 +166,57 @@ def new_chat(*args, **kwargs): model = kwargs.get("model", "") - span = sentry_sdk.start_span( + with sentry_sdk.start_span( op=OP.GEN_AI_CHAT, name=f"chat {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_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", [])) - messages = normalize_message_roles(messages) - 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, - 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) + ) as span: + 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 streaming: - old_iterator = res - - def new_iterator(): - # type: () -> Iterator[V2ChatStreamResponse] - collected_text = [] - try: - for x in old_iterator: - with capture_internal_exceptions(): + 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_REQUEST_MODEL, model) + + if should_send_default_pii() and integration.include_prompts: + 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) + 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, + 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" @@ -260,25 +256,18 @@ def new_iterator(): and x.delta.usage is not None ): _record_token_usage_v2(span, x.delta.usage) - yield x - except Exception as exc: - _capture_exception(exc) - raise - finally: - 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 + yield x + + return new_iterator() + elif isinstance(res, V2ChatResponse): + collect_v2_response_fields( + span, + res, + include_pii=should_send_default_pii() + and integration.include_prompts, + ) + else: + set_data_normalized(span, "unknown_response", True) + return res return new_chat From e9840a2318f2b108170ded86d823794172ad8533 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Wed, 4 Mar 2026 11:04:10 +0100 Subject: [PATCH 11/29] format --- sentry_sdk/integrations/cohere/v1.py | 4 +++- sentry_sdk/integrations/cohere/v2.py | 8 ++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index b8e87e5e7c..4afb59b50d 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -154,7 +154,9 @@ def new_chat(*args, **kwargs): 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, diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index fa1313e6a5..8963c9813e 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -189,7 +189,9 @@ 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, @@ -226,7 +228,9 @@ 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 50cfff224aad7905c61e934affc4c52b1c5dd347 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Wed, 4 Mar 2026 11:04:58 +0100 Subject: [PATCH 12/29] remove comments --- sentry_sdk/integrations/cohere/__init__.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index 8642cf751c..75dd0cc13f 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -94,8 +94,6 @@ def new_embed(*args, **kwargs): set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "cohere") set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") - # attach embeddings input - # should this be truncated and annotated? if "texts" in kwargs and ( should_send_default_pii() and integration.include_prompts ): @@ -110,7 +108,6 @@ def new_embed(*args, **kwargs): span, SPANDATA.GEN_AI_REQUEST_MODEL, kwargs["model"] ) - # call the function and capture any exceptions try: res = f(*args, **kwargs) except Exception as e: @@ -119,7 +116,6 @@ def new_embed(*args, **kwargs): _capture_exception(e) reraise(*exc_info) - # record token usage if ( hasattr(res, "meta") and hasattr(res.meta, "billed_units") From 65f92305af0993042c2edc45c3b14a49de450564 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Wed, 4 Mar 2026 13:58:43 +0100 Subject: [PATCH 13/29] wip --- sentry_sdk/ai/span_config.py | 68 +++++++++++ sentry_sdk/integrations/cohere/__init__.py | 30 ++--- sentry_sdk/integrations/cohere/v1.py | 133 ++++++++++----------- sentry_sdk/integrations/cohere/v2.py | 51 +++----- 4 files changed, 159 insertions(+), 123 deletions(-) create mode 100644 sentry_sdk/ai/span_config.py diff --git a/sentry_sdk/ai/span_config.py b/sentry_sdk/ai/span_config.py new file mode 100644 index 0000000000..367e544297 --- /dev/null +++ b/sentry_sdk/ai/span_config.py @@ -0,0 +1,68 @@ +import sentry_sdk +from sentry_sdk.consts import SPANDATA +from sentry_sdk.ai.utils import ( + set_data_normalized, + normalize_message_roles, + truncate_and_annotate_messages, +) +from sentry_sdk.scope import should_send_default_pii + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Dict + from sentry_sdk.tracing import Span + + +def set_input_span_data(span, kwargs, integration, config): + # type: (Span, Dict[str, Any], Any, Dict[str, Any]) -> None + """ + Set input span data from a declarative config. + + Config keys: + system: str - gen_ai.system value + operation: str - gen_ai.operation.name value + params: dict - kwargs key -> span attr (always set if present) + pii_params: dict - kwargs key -> span attr (only when PII allowed) + extract_messages: callable(kwargs) -> list or None + message_target: str - span attr for messages (default: GEN_AI_REQUEST_MESSAGES) + truncation_fn: callable or None - truncation function (default: truncate_and_annotate_messages, None to skip) + is_given: callable(value) -> bool - for NotGiven sentinels + extra_static: dict - additional key/value pairs to set + """ + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, config["system"]) + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, config["operation"]) + + is_given = config.get("is_given") + for kwarg_key, span_attr in config.get("params", {}).items(): + if kwarg_key in kwargs: + value = kwargs[kwarg_key] + if is_given is None or is_given(value): + set_data_normalized(span, span_attr, value) + + if should_send_default_pii() and integration.include_prompts: + extract = config.get("extract_messages") + if extract is not None: + messages = extract(kwargs) + if messages: + messages = normalize_message_roles(messages) + truncation_fn = config.get( + "truncation_fn", truncate_and_annotate_messages + ) + if truncation_fn is not None: + scope = sentry_sdk.get_current_scope() + messages = truncation_fn(messages, span, scope) + if messages is not None: + target = config.get( + "message_target", SPANDATA.GEN_AI_REQUEST_MESSAGES + ) + set_data_normalized(span, target, messages, unpack=False) + + for kwarg_key, span_attr in config.get("pii_params", {}).items(): + if kwarg_key in kwargs: + value = kwargs[kwarg_key] + if is_given is None or is_given(value): + set_data_normalized(span, span_attr, value) + + for key, value in config.get("extra_static", {}).items(): + set_data_normalized(span, key, value) diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index 75dd0cc13f..456b531cf0 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -3,7 +3,7 @@ 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 sentry_sdk.ai.span_config import set_input_span_data from typing import TYPE_CHECKING @@ -13,7 +13,6 @@ 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 @@ -43,6 +42,16 @@ def _normalize_embedding_input(texts): return [texts] +COHERE_EMBED_CONFIG = { + "system": "cohere", + "operation": "embeddings", + "params": {"model": SPANDATA.GEN_AI_REQUEST_MODEL}, + "extract_messages": lambda kw: _normalize_embedding_input(kw["texts"]) if "texts" in kw else None, + "message_target": SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + "truncation_fn": None, +} + + class CohereIntegration(Integration): identifier = "cohere" origin = f"auto.ai.{identifier}" @@ -91,22 +100,7 @@ def new_embed(*args, **kwargs): 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 - ): - set_data_normalized( - span, - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - _normalize_embedding_input(kwargs["texts"]), - ) - - if "model" in kwargs: - set_data_normalized( - span, SPANDATA.GEN_AI_REQUEST_MODEL, kwargs["model"] - ) + set_input_span_data(span, kwargs, integration, COHERE_EMBED_CONFIG) try: res = f(*args, **kwargs) diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index 4afb59b50d..c63990b3cd 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -3,11 +3,8 @@ 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, - truncate_and_annotate_messages, -) +from sentry_sdk.ai.utils import set_data_normalized +from sentry_sdk.ai.span_config import set_input_span_data from typing import TYPE_CHECKING @@ -48,6 +45,32 @@ "preamble": SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, } + +def _extract_messages_v1(kwargs): + # type: (dict[str, Any]) -> list[dict[str, str]] + """Extract role/content dicts from V1-style chat_history + message.""" + messages = [] + for x in kwargs.get("chat_history", []): + messages.append( + { + "role": getattr(x, "role", "").lower(), + "content": getattr(x, "message", ""), + } + ) + message = kwargs.get("message") + if message: + messages.append({"role": "user", "content": message}) + return messages + + +COHERE_V1_CHAT_CONFIG = { + "system": "cohere", + "operation": "chat", + "params": COLLECTED_CHAT_PARAMS, + "pii_params": COLLECTED_PII_CHAT_PARAMS, + "extract_messages": _extract_messages_v1, +} + COLLECTED_CHAT_RESP_ATTRS = { "generation_id": SPANDATA.GEN_AI_RESPONSE_ID, "finish_reason": SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, @@ -77,36 +100,6 @@ def _wrap_chat(f, streaming): if not _has_chat_types: return f - def collect_chat_response_fields(span, res, include_pii): - # type: (Span, NonStreamedChatResponse, bool) -> None - if include_pii: - if hasattr(res, "text"): - set_data_normalized( - span, - SPANDATA.GEN_AI_RESPONSE_TEXT, - [res.text], - ) - 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, spandata_key in COLLECTED_CHAT_RESP_ATTRS.items(): - if hasattr(res, attr): - set_data_normalized(span, spandata_key, getattr(res, attr)) - - if hasattr(res, "meta"): - if hasattr(res.meta, "billed_units"): - record_token_usage( - span, - input_tokens=res.meta.billed_units.input_tokens, - output_tokens=res.meta.billed_units.output_tokens, - ) - elif hasattr(res.meta, "tokens"): - record_token_usage( - span, - input_tokens=res.meta.tokens.input_tokens, - output_tokens=res.meta.tokens.output_tokens, - ) @wraps(f) def new_chat(*args, **kwargs): @@ -120,7 +113,6 @@ def new_chat(*args, **kwargs): ): return f(*args, **kwargs) - message = kwargs.get("message") model = kwargs.get("model", "") with sentry_sdk.start_span( @@ -137,41 +129,10 @@ def new_chat(*args, **kwargs): 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_REQUEST_MODEL, model) - - 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": "user", "content": message}) - messages = normalize_message_roles(messages) - 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]) - - 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) + set_input_span_data(span, kwargs, integration, { + **COHERE_V1_CHAT_CONFIG, + "extra_static": {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming}, + }) if streaming: old_iterator = res @@ -203,4 +164,34 @@ def new_iterator(): set_data_normalized(span, "unknown_response", True) return res + def collect_chat_response_fields(span, res, include_pii): + # type: (Span, NonStreamedChatResponse, bool) -> None + if include_pii: + if hasattr(res, "text"): + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TEXT, + [res.text], + ) + 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, spandata_key in COLLECTED_CHAT_RESP_ATTRS.items(): + if hasattr(res, attr): + set_data_normalized(span, spandata_key, getattr(res, attr)) + + if hasattr(res, "meta"): + if hasattr(res.meta, "billed_units"): + record_token_usage( + span, + input_tokens=res.meta.billed_units.input_tokens, + output_tokens=res.meta.billed_units.output_tokens, + ) + elif hasattr(res.meta, "tokens"): + record_token_usage( + span, + input_tokens=res.meta.tokens.input_tokens, + output_tokens=res.meta.tokens.output_tokens, + ) return new_chat diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index 8963c9813e..c8dd590103 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -3,11 +3,8 @@ 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, - truncate_and_annotate_messages, -) +from sentry_sdk.ai.utils import set_data_normalized +from sentry_sdk.ai.span_config import set_input_span_data from typing import TYPE_CHECKING @@ -101,6 +98,15 @@ def _extract_messages_v2(messages): return result +COHERE_V2_CHAT_CONFIG = { + "system": "cohere", + "operation": "chat", + "params": COLLECTED_CHAT_PARAMS, + "pii_params": {"tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS}, + "extract_messages": lambda kw: _extract_messages_v2(kw.get("messages", [])), +} + + def _record_token_usage_v2(span, usage): # type: (Span, Any) -> None """Extract and record token usage from a V2 Usage object.""" @@ -180,36 +186,13 @@ def new_chat(*args, **kwargs): 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") + extra = {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming} if 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", [])) - messages = normalize_message_roles(messages) - 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, - 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) + extra[SPANDATA.GEN_AI_RESPONSE_MODEL] = model + set_input_span_data(span, kwargs, integration, { + **COHERE_V2_CHAT_CONFIG, + "extra_static": extra, + }) if streaming: old_iterator = res From 953217523b2dd9dcb8975a58156c9d65e1388925 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Wed, 4 Mar 2026 14:10:02 +0100 Subject: [PATCH 14/29] format --- sentry_sdk/integrations/cohere/__init__.py | 4 +++- sentry_sdk/integrations/cohere/v1.py | 15 ++++++++++----- sentry_sdk/integrations/cohere/v2.py | 13 +++++++++---- 3 files changed, 22 insertions(+), 10 deletions(-) diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index 456b531cf0..9917bf23db 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -46,7 +46,9 @@ def _normalize_embedding_input(texts): "system": "cohere", "operation": "embeddings", "params": {"model": SPANDATA.GEN_AI_REQUEST_MODEL}, - "extract_messages": lambda kw: _normalize_embedding_input(kw["texts"]) if "texts" in kw else None, + "extract_messages": lambda kw: ( + _normalize_embedding_input(kw["texts"]) if "texts" in kw else None + ), "message_target": SPANDATA.GEN_AI_EMBEDDINGS_INPUT, "truncation_fn": None, } diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index c63990b3cd..22dca47a47 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -100,7 +100,6 @@ def _wrap_chat(f, streaming): if not _has_chat_types: return f - @wraps(f) def new_chat(*args, **kwargs): # type: (*Any, **Any) -> Any @@ -129,10 +128,15 @@ def new_chat(*args, **kwargs): reraise(*exc_info) with capture_internal_exceptions(): - set_input_span_data(span, kwargs, integration, { - **COHERE_V1_CHAT_CONFIG, - "extra_static": {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming}, - }) + set_input_span_data( + span, + kwargs, + integration, + { + **COHERE_V1_CHAT_CONFIG, + "extra_static": {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming}, + }, + ) if streaming: old_iterator = res @@ -194,4 +198,5 @@ def collect_chat_response_fields(span, res, include_pii): input_tokens=res.meta.tokens.input_tokens, output_tokens=res.meta.tokens.output_tokens, ) + return new_chat diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index c8dd590103..98b04e40a1 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -189,10 +189,15 @@ def new_chat(*args, **kwargs): extra = {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming} if model: extra[SPANDATA.GEN_AI_RESPONSE_MODEL] = model - set_input_span_data(span, kwargs, integration, { - **COHERE_V2_CHAT_CONFIG, - "extra_static": extra, - }) + set_input_span_data( + span, + kwargs, + integration, + { + **COHERE_V2_CHAT_CONFIG, + "extra_static": extra, + }, + ) if streaming: old_iterator = res From b1d4fbd09a5ba026ab3712de62cbfbd131d86486 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Wed, 4 Mar 2026 17:12:26 +0100 Subject: [PATCH 15/29] refacotr --- sentry_sdk/integrations/cohere/__init__.py | 12 - sentry_sdk/integrations/cohere/utils.py | 35 +++ sentry_sdk/integrations/cohere/v1.py | 196 +++++++------- sentry_sdk/integrations/cohere/v2.py | 292 +++++++++------------ 4 files changed, 244 insertions(+), 291 deletions(-) create mode 100644 sentry_sdk/integrations/cohere/utils.py diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index 9917bf23db..deb159a23a 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -21,17 +21,6 @@ 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, -} - def _normalize_embedding_input(texts): # type: (Any) -> Any @@ -66,7 +55,6 @@ def __init__(self, include_prompts=True): 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 diff --git a/sentry_sdk/integrations/cohere/utils.py b/sentry_sdk/integrations/cohere/utils.py new file mode 100644 index 0000000000..d9aba318fc --- /dev/null +++ b/sentry_sdk/integrations/cohere/utils.py @@ -0,0 +1,35 @@ +from sentry_sdk.ai.utils import set_data_normalized + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + +_MISSING = object() + + +def transitive_getattr(obj, *attrs): + # type: (Any, str) -> Any + current = obj + for attr in attrs: + current = getattr(current, attr, _MISSING) + if current is _MISSING: + return None + return current + + +def get_first_from_sources(obj, source_paths, require_truthy=False): + # type: (Any, list[tuple[str, ...]], bool) -> Any + for source_path in source_paths: + value = transitive_getattr(obj, *source_path) + if value if require_truthy else value is not None: + return value + return None + + +def set_span_data_from_sources(span, obj, target_sources, require_truthy): + # type: (Any, Any, dict[str, list[tuple[str, ...]]], bool) -> None + for spandata_key, source_paths in target_sources.items(): + value = get_first_from_sources(obj, source_paths, require_truthy=require_truthy) + if value is not None: + set_data_normalized(span, spandata_key, value) diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index 22dca47a47..3f787151c0 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -2,32 +2,30 @@ 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 sentry_sdk.ai.span_config import set_input_span_data +from sentry_sdk.ai.utils import set_data_normalized, transform_message_content +from sentry_sdk.consts import OP, SPANDATA from typing import TYPE_CHECKING if TYPE_CHECKING: from typing import Any, Callable, Iterator from cohere import StreamedChatResponse - 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, reraise - from sentry_sdk.integrations.cohere import ( CohereIntegration, - COLLECTED_CHAT_PARAMS, _capture_exception, ) +from sentry_sdk.integrations.cohere.utils import ( + get_first_from_sources, + set_span_data_from_sources, +) +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import capture_internal_exceptions, reraise try: - from cohere import ( - ChatStreamEndEvent, - NonStreamedChatResponse, - ) + from cohere import ChatStreamEndEvent, NonStreamedChatResponse try: from cohere import StreamEndStreamedChatResponse @@ -40,53 +38,32 @@ except ImportError: _has_chat_types = False -COLLECTED_PII_CHAT_PARAMS = { - "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - "preamble": SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, -} - - -def _extract_messages_v1(kwargs): - # type: (dict[str, Any]) -> list[dict[str, str]] - """Extract role/content dicts from V1-style chat_history + message.""" - messages = [] - for x in kwargs.get("chat_history", []): - messages.append( - { - "role": getattr(x, "role", "").lower(), - "content": getattr(x, "message", ""), - } - ) - message = kwargs.get("message") - if message: - messages.append({"role": "user", "content": message}) - return messages - - -COHERE_V1_CHAT_CONFIG = { - "system": "cohere", - "operation": "chat", - "params": COLLECTED_CHAT_PARAMS, - "pii_params": COLLECTED_PII_CHAT_PARAMS, - "extract_messages": _extract_messages_v1, +CHAT_RESPONSE_SOURCES = { + SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], } - -COLLECTED_CHAT_RESP_ATTRS = { - "generation_id": SPANDATA.GEN_AI_RESPONSE_ID, - "finish_reason": SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, +PII_CHAT_RESPONSE_SOURCES = { + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("tool_calls",)], } - -COLLECTED_PII_CHAT_RESP_ATTRS = { - "tool_calls": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, +CHAT_RESPONSE_TEXT_SOURCES = [("text",)] +CHAT_USAGE_TOKEN_SOURCES = { + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: [ + ("meta", "billed_units", "input_tokens"), + ("meta", "tokens", "input_tokens"), + ], + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: [ + ("meta", "billed_units", "output_tokens"), + ("meta", "tokens", "output_tokens"), + ], } +STREAM_RESPONSE_SOURCES = [("response",)] 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 + from cohere.client import Client except ImportError: return @@ -104,7 +81,6 @@ def _wrap_chat(f, streaming): def new_chat(*args, **kwargs): # type: (*Any, **Any) -> Any integration = sentry_sdk.get_client().get_integration(CohereIntegration) - if ( integration is None or "message" not in kwargs @@ -113,6 +89,7 @@ def new_chat(*args, **kwargs): return f(*args, **kwargs) model = kwargs.get("model", "") + include_pii = should_send_default_pii() and integration.include_prompts with sentry_sdk.start_span( op=OP.GEN_AI_CHAT, @@ -128,75 +105,84 @@ def new_chat(*args, **kwargs): reraise(*exc_info) with capture_internal_exceptions(): + if model: + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) set_input_span_data( span, kwargs, integration, { - **COHERE_V1_CHAT_CONFIG, + "system": "cohere", + "operation": "chat", + "extract_messages": _extract_messages_v1, "extra_static": {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming}, }, ) if streaming: - old_iterator = res - - def new_iterator(): - # type: () -> Iterator[StreamedChatResponse] - with capture_internal_exceptions(): - for x in old_iterator: - if isinstance(x, ChatStreamEndEvent) or isinstance( - x, StreamEndStreamedChatResponse - ): - collect_chat_response_fields( - span, - x.response, - include_pii=should_send_default_pii() - and integration.include_prompts, - ) - yield x - - return new_iterator() - elif isinstance(res, NonStreamedChatResponse): - collect_chat_response_fields( - span, - res, - include_pii=should_send_default_pii() - and integration.include_prompts, - ) + return _iter_v1_stream_events(res, span, include_pii) + if isinstance(res, NonStreamedChatResponse): + _collect_v1_response_fields(span, res, include_pii=include_pii) else: set_data_normalized(span, "unknown_response", True) return res - def collect_chat_response_fields(span, res, include_pii): - # type: (Span, NonStreamedChatResponse, bool) -> None - if include_pii: - if hasattr(res, "text"): - set_data_normalized( - span, - SPANDATA.GEN_AI_RESPONSE_TEXT, - [res.text], - ) - for attr, spandata_key in COLLECTED_PII_CHAT_RESP_ATTRS.items(): - if hasattr(res, attr): - set_data_normalized(span, spandata_key, getattr(res, attr)) + return new_chat - for attr, spandata_key in COLLECTED_CHAT_RESP_ATTRS.items(): - if hasattr(res, attr): - set_data_normalized(span, spandata_key, getattr(res, attr)) - if hasattr(res, "meta"): - if hasattr(res.meta, "billed_units"): - record_token_usage( - span, - input_tokens=res.meta.billed_units.input_tokens, - output_tokens=res.meta.billed_units.output_tokens, - ) - elif hasattr(res.meta, "tokens"): - record_token_usage( - span, - input_tokens=res.meta.tokens.input_tokens, - output_tokens=res.meta.tokens.output_tokens, - ) +def _extract_messages_v1(kwargs): + # type: (dict[str, Any]) -> list[dict[str, str]] + messages = [] + for x in kwargs.get("chat_history", []): + messages.append( + { + "role": getattr(x, "role", "").lower(), + "content": transform_message_content(getattr(x, "message", "")), + } + ) + message = kwargs.get("message") + if message: + messages.append({"role": "user", "content": transform_message_content(message)}) + return messages - return new_chat + +def _iter_v1_stream_events(old_iterator, span, include_pii): + # type: (Any, Any, bool) -> Iterator[StreamedChatResponse] + with capture_internal_exceptions(): + for x in old_iterator: + if isinstance(x, ChatStreamEndEvent) or isinstance( + x, StreamEndStreamedChatResponse + ): + _collect_v1_stream_end_fields(span, x, include_pii) + yield x + + +def _collect_v1_stream_end_fields(span, event, include_pii): + # type: (Any, Any, bool) -> None + response = get_first_from_sources(event, STREAM_RESPONSE_SOURCES) + if response is not None: + _collect_v1_response_fields(span, response, include_pii) + + +def _collect_v1_response_fields(span, response, include_pii): + # type: (Any, Any, bool) -> None + if include_pii: + text = get_first_from_sources(response, CHAT_RESPONSE_TEXT_SOURCES) + if text is not None: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, [text]) + set_span_data_from_sources( + span, response, PII_CHAT_RESPONSE_SOURCES, require_truthy=False + ) + + set_span_data_from_sources( + span, response, CHAT_RESPONSE_SOURCES, require_truthy=False + ) + record_token_usage( + span, + input_tokens=get_first_from_sources( + response, CHAT_USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] + ), + output_tokens=get_first_from_sources( + response, CHAT_USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] + ), + ) diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index 98b04e40a1..c7f9b3f305 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -2,9 +2,9 @@ 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 sentry_sdk.ai.span_config import set_input_span_data +from sentry_sdk.ai.utils import set_data_normalized, transform_message_content +from sentry_sdk.consts import OP, SPANDATA from typing import TYPE_CHECKING @@ -13,24 +13,23 @@ 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, reraise - from sentry_sdk.integrations.cohere import ( CohereIntegration, - COLLECTED_CHAT_PARAMS, _capture_exception, ) +from sentry_sdk.integrations.cohere.utils import ( + get_first_from_sources, + set_span_data_from_sources, + transitive_getattr, +) +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import capture_internal_exceptions, reraise try: from cohere.v2.client import V2Client as CohereV2Client - # 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 + from cohere.v2.types import MessageEndV2ChatStreamResponse, V2ChatResponse if TYPE_CHECKING: from cohere.v2.types import V2ChatStreamResponse @@ -47,18 +46,31 @@ except ImportError: _has_v2 = False +CHAT_RESPONSE_SOURCES = { + SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], +} +PII_CHAT_RESPONSE_SOURCES = { + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("message", "tool_calls")], +} +CHAT_USAGE_SOURCES = [("usage",)] +STREAM_DELTA_TEXT_SOURCES = [("delta", "message", "content", "text")] +STREAM_CHAT_RESPONSE_SOURCES = { + SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("delta", "finish_reason")], +} +STREAM_CHAT_USAGE_SOURCES = [("delta", "usage")] +USAGE_UNIT_SOURCES = [("billed_units",), ("tokens",)] +USAGE_TOKEN_SOURCES = { + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: [("input_tokens",)], + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: [("output_tokens",)], +} + 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 __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. - """ 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 @@ -66,111 +78,17 @@ def setup_v2(wrap_embed_fn): 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. - - Handles both plain dicts and Pydantic model instances. - """ - result = [] - for msg in messages: - 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( - ( - 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 "" - result.append({"role": role, "content": text}) - return result - - -COHERE_V2_CHAT_CONFIG = { - "system": "cohere", - "operation": "chat", - "params": COLLECTED_CHAT_PARAMS, - "pii_params": {"tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS}, - "extract_messages": lambda kw: _extract_messages_v2(kw.get("messages", [])), -} - - -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", "") + include_pii = should_send_default_pii() and integration.include_prompts with sentry_sdk.start_span( op=OP.GEN_AI_CHAT, @@ -188,78 +106,104 @@ def new_chat(*args, **kwargs): with capture_internal_exceptions(): extra = {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming} if model: + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) extra[SPANDATA.GEN_AI_RESPONSE_MODEL] = model set_input_span_data( span, kwargs, integration, { - **COHERE_V2_CHAT_CONFIG, + "system": "cohere", + "operation": "chat", + "pii_params": { + "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + }, + "extract_messages": lambda kw: _extract_messages_v2( + kw.get("messages", []) + ), "extra_static": extra, }, ) - if streaming: - old_iterator = res + return _iter_v2_stream_events(res, span, include_pii) + _collect_v2_response_fields(span, res, include_pii=include_pii) + return 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) + return new_chat - 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 - return new_iterator() - elif isinstance(res, V2ChatResponse): - collect_v2_response_fields( - span, - res, - include_pii=should_send_default_pii() - and integration.include_prompts, - ) - else: - set_data_normalized(span, "unknown_response", True) - return res +def _extract_messages_v2(messages): + # type: (Any) -> list[dict[str, Any]] + result = [] + for msg in messages: + role = msg["role"] if isinstance(msg, dict) else getattr(msg, "role", "unknown") + content = msg["content"] if isinstance(msg, dict) else getattr(msg, "content", "") + result.append({"role": role, "content": transform_message_content(content)}) + return result - return new_chat + +def _iter_v2_stream_events(old_iterator, span, include_pii): + # type: (Any, Span, bool) -> Iterator[V2ChatStreamResponse] + collected_text = [] + with capture_internal_exceptions(): + for x in old_iterator: + _append_stream_delta_text(collected_text, x) + if isinstance(x, MessageEndV2ChatStreamResponse): + _collect_v2_stream_end_fields(span, x, include_pii, collected_text) + yield x + + +def _append_stream_delta_text(collected_text, event): + # type: (list[str], Any) -> None + if transitive_getattr(event, "type") != "content-delta": + return + content_text = get_first_from_sources(event, STREAM_DELTA_TEXT_SOURCES) + if content_text is not None: + collected_text.append(content_text) + + +def _collect_v2_stream_end_fields(span, event, include_pii, collected_text): + # type: (Span, Any, bool, list[str]) -> None + if include_pii and collected_text: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, ["".join(collected_text)]) + set_span_data_from_sources( + span, event, STREAM_CHAT_RESPONSE_SOURCES, require_truthy=False + ) + stream_usage = get_first_from_sources(event, STREAM_CHAT_USAGE_SOURCES) + if stream_usage is not None: + _record_token_usage_v2(span, stream_usage) + + +def _collect_v2_response_fields(span, response, include_pii): + # type: (Span, V2ChatResponse, bool) -> None + if include_pii: + content = get_first_from_sources(response, [("message", "content")], True) + if content: + texts = [item.text for item in content if hasattr(item, "text")] + if texts: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts) + set_span_data_from_sources( + span, response, PII_CHAT_RESPONSE_SOURCES, require_truthy=True + ) + set_span_data_from_sources( + span, response, CHAT_RESPONSE_SOURCES, require_truthy=False + ) + usage = get_first_from_sources(response, CHAT_USAGE_SOURCES) + if usage is not None: + _record_token_usage_v2(span, usage) + + +def _record_token_usage_v2(span, usage): + # type: (Span, Any) -> None + usage_obj = get_first_from_sources(usage, USAGE_UNIT_SOURCES) + if usage_obj is None: + return + record_token_usage( + span, + input_tokens=get_first_from_sources( + usage_obj, USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] + ), + output_tokens=get_first_from_sources( + usage_obj, USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] + ), + ) From f3563dd62d11facd4565b1cb6e35f93e6bb437b1 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 09:27:06 +0100 Subject: [PATCH 16/29] format --- sentry_sdk/integrations/cohere/v2.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index c7f9b3f305..a9e307d574 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -137,7 +137,9 @@ def _extract_messages_v2(messages): result = [] for msg in messages: role = msg["role"] if isinstance(msg, dict) else getattr(msg, "role", "unknown") - content = msg["content"] if isinstance(msg, dict) else getattr(msg, "content", "") + content = ( + msg["content"] if isinstance(msg, dict) else getattr(msg, "content", "") + ) result.append({"role": role, "content": transform_message_content(content)}) return result @@ -165,7 +167,9 @@ def _append_stream_delta_text(collected_text, event): def _collect_v2_stream_end_fields(span, event, include_pii, collected_text): # type: (Span, Any, bool, list[str]) -> None if include_pii and collected_text: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, ["".join(collected_text)]) + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, ["".join(collected_text)] + ) set_span_data_from_sources( span, event, STREAM_CHAT_RESPONSE_SOURCES, require_truthy=False ) From 3227de82c739ff8fc58bb3a2395ca9d536106769 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 09:29:50 +0100 Subject: [PATCH 17/29] style --- sentry_sdk/integrations/cohere/utils.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/sentry_sdk/integrations/cohere/utils.py b/sentry_sdk/integrations/cohere/utils.py index d9aba318fc..2e41d227a0 100644 --- a/sentry_sdk/integrations/cohere/utils.py +++ b/sentry_sdk/integrations/cohere/utils.py @@ -5,15 +5,13 @@ if TYPE_CHECKING: from typing import Any -_MISSING = object() - def transitive_getattr(obj, *attrs): # type: (Any, str) -> Any current = obj for attr in attrs: - current = getattr(current, attr, _MISSING) - if current is _MISSING: + current = getattr(current, attr, None) + if not current: return None return current From 068ce5d4c9d05719c6137aa4bd5c6f2fac0d089c Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 09:35:29 +0100 Subject: [PATCH 18/29] fix getattr --- sentry_sdk/integrations/cohere/utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/sentry_sdk/integrations/cohere/utils.py b/sentry_sdk/integrations/cohere/utils.py index 2e41d227a0..78d662d614 100644 --- a/sentry_sdk/integrations/cohere/utils.py +++ b/sentry_sdk/integrations/cohere/utils.py @@ -11,7 +11,7 @@ def transitive_getattr(obj, *attrs): current = obj for attr in attrs: current = getattr(current, attr, None) - if not current: + if current is None: return None return current @@ -20,7 +20,9 @@ def get_first_from_sources(obj, source_paths, require_truthy=False): # type: (Any, list[tuple[str, ...]], bool) -> Any for source_path in source_paths: value = transitive_getattr(obj, *source_path) - if value if require_truthy else value is not None: + if not value: + continue + if not require_truthy or value: return value return None From c7a1b581a36c1bbdedf336f7c650f17ea4cb87a6 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 10:25:11 +0100 Subject: [PATCH 19/29] simplify --- sentry_sdk/ai/span_config.py | 37 ++++++++----------- sentry_sdk/ai/utils.py | 3 +- sentry_sdk/integrations/cohere/__init__.py | 6 ++-- sentry_sdk/integrations/cohere/utils.py | 6 ++-- sentry_sdk/integrations/cohere/v1.py | 35 +++++++++--------- sentry_sdk/integrations/cohere/v2.py | 41 +++++++++++----------- 6 files changed, 61 insertions(+), 67 deletions(-) diff --git a/sentry_sdk/ai/span_config.py b/sentry_sdk/ai/span_config.py index 367e544297..ef38160ef1 100644 --- a/sentry_sdk/ai/span_config.py +++ b/sentry_sdk/ai/span_config.py @@ -14,31 +14,30 @@ from sentry_sdk.tracing import Span -def set_input_span_data(span, kwargs, integration, config): - # type: (Span, Dict[str, Any], Any, Dict[str, Any]) -> None +def set_input_span_data(span, kwargs, integration, config, span_data=None): + # type: (Span, Dict[str, Any], Any, Dict[str, Any], Dict[str, Any] | None) -> None """ Set input span data from a declarative config. Config keys: - system: str - gen_ai.system value - operation: str - gen_ai.operation.name value + static: dict - key/value pairs to set unconditionally params: dict - kwargs key -> span attr (always set if present) pii_params: dict - kwargs key -> span attr (only when PII allowed) extract_messages: callable(kwargs) -> list or None message_target: str - span attr for messages (default: GEN_AI_REQUEST_MESSAGES) - truncation_fn: callable or None - truncation function (default: truncate_and_annotate_messages, None to skip) - is_given: callable(value) -> bool - for NotGiven sentinels - extra_static: dict - additional key/value pairs to set + + span_data: additional key/value pairs for dynamic per-call values """ - set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, config["system"]) - set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, config["operation"]) + for key, value in config.get("static", {}).items(): + set_data_normalized(span, key, value) + if span_data: + for key, value in span_data.items(): + set_data_normalized(span, key, value) - is_given = config.get("is_given") for kwarg_key, span_attr in config.get("params", {}).items(): if kwarg_key in kwargs: value = kwargs[kwarg_key] - if is_given is None or is_given(value): - set_data_normalized(span, span_attr, value) + set_data_normalized(span, span_attr, value) if should_send_default_pii() and integration.include_prompts: extract = config.get("extract_messages") @@ -46,12 +45,8 @@ def set_input_span_data(span, kwargs, integration, config): messages = extract(kwargs) if messages: messages = normalize_message_roles(messages) - truncation_fn = config.get( - "truncation_fn", truncate_and_annotate_messages - ) - if truncation_fn is not None: - scope = sentry_sdk.get_current_scope() - messages = truncation_fn(messages, span, scope) + scope = sentry_sdk.get_current_scope() + messages = truncate_and_annotate_messages(messages, span, scope) if messages is not None: target = config.get( "message_target", SPANDATA.GEN_AI_REQUEST_MESSAGES @@ -61,8 +56,4 @@ def set_input_span_data(span, kwargs, integration, config): for kwarg_key, span_attr in config.get("pii_params", {}).items(): if kwarg_key in kwargs: value = kwargs[kwarg_key] - if is_given is None or is_given(value): - set_data_normalized(span, span_attr, value) - - for key, value in config.get("extra_static", {}).items(): - set_data_normalized(span, key, value) + set_data_normalized(span, span_attr, value) diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py index 0f104cc8f5..7e14c15e24 100644 --- a/sentry_sdk/ai/utils.py +++ b/sentry_sdk/ai/utils.py @@ -503,7 +503,8 @@ 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) + role_lower = role.lower() + return GEN_AI_MESSAGE_ROLE_MAPPING.get(role_lower, role_lower) def normalize_message_roles(messages: "list[dict[str, Any]]") -> "list[dict[str, Any]]": diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index deb159a23a..ac6b7c474d 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -32,8 +32,10 @@ def _normalize_embedding_input(texts): COHERE_EMBED_CONFIG = { - "system": "cohere", - "operation": "embeddings", + "static": { + SPANDATA.GEN_AI_SYSTEM: "cohere", + SPANDATA.GEN_AI_OPERATION_NAME: "embeddings", + }, "params": {"model": SPANDATA.GEN_AI_REQUEST_MODEL}, "extract_messages": lambda kw: ( _normalize_embedding_input(kw["texts"]) if "texts" in kw else None diff --git a/sentry_sdk/integrations/cohere/utils.py b/sentry_sdk/integrations/cohere/utils.py index 78d662d614..346ac72a6a 100644 --- a/sentry_sdk/integrations/cohere/utils.py +++ b/sentry_sdk/integrations/cohere/utils.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any + from typing import Any, Mapping, Sequence def transitive_getattr(obj, *attrs): @@ -17,7 +17,7 @@ def transitive_getattr(obj, *attrs): def get_first_from_sources(obj, source_paths, require_truthy=False): - # type: (Any, list[tuple[str, ...]], bool) -> Any + # type: (Any, Sequence[tuple[str, ...]], bool) -> Any for source_path in source_paths: value = transitive_getattr(obj, *source_path) if not value: @@ -28,7 +28,7 @@ def get_first_from_sources(obj, source_paths, require_truthy=False): def set_span_data_from_sources(span, obj, target_sources, require_truthy): - # type: (Any, Any, dict[str, list[tuple[str, ...]]], bool) -> None + # type: (Any, Any, Mapping[str, Sequence[tuple[str, ...]]], bool) -> None for spandata_key, source_paths in target_sources.items(): value = get_first_from_sources(obj, source_paths, require_truthy=require_truthy) if value is not None: diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index 3f787151c0..6b12d5d67d 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -38,6 +38,14 @@ except ImportError: _has_chat_types = False +COHERE_V1_CHAT_CONFIG = { + "static": { + SPANDATA.GEN_AI_SYSTEM: "cohere", + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + }, + "extract_messages": lambda kw: _extract_messages(kw), +} + CHAT_RESPONSE_SOURCES = { SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)], SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], @@ -105,24 +113,17 @@ def new_chat(*args, **kwargs): reraise(*exc_info) with capture_internal_exceptions(): + span_data = {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming} if model: - set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) + span_data[SPANDATA.GEN_AI_REQUEST_MODEL] = model set_input_span_data( - span, - kwargs, - integration, - { - "system": "cohere", - "operation": "chat", - "extract_messages": _extract_messages_v1, - "extra_static": {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming}, - }, + span, kwargs, integration, COHERE_V1_CHAT_CONFIG, span_data ) if streaming: - return _iter_v1_stream_events(res, span, include_pii) + return _iter_stream_events(res, span, include_pii) if isinstance(res, NonStreamedChatResponse): - _collect_v1_response_fields(span, res, include_pii=include_pii) + _collect_response_fields(span, res, include_pii=include_pii) else: set_data_normalized(span, "unknown_response", True) return res @@ -130,13 +131,13 @@ def new_chat(*args, **kwargs): return new_chat -def _extract_messages_v1(kwargs): +def _extract_messages(kwargs): # type: (dict[str, Any]) -> list[dict[str, str]] messages = [] for x in kwargs.get("chat_history", []): messages.append( { - "role": getattr(x, "role", "").lower(), + "role": getattr(x, "role", ""), "content": transform_message_content(getattr(x, "message", "")), } ) @@ -146,7 +147,7 @@ def _extract_messages_v1(kwargs): return messages -def _iter_v1_stream_events(old_iterator, span, include_pii): +def _iter_stream_events(old_iterator, span, include_pii): # type: (Any, Any, bool) -> Iterator[StreamedChatResponse] with capture_internal_exceptions(): for x in old_iterator: @@ -161,10 +162,10 @@ def _collect_v1_stream_end_fields(span, event, include_pii): # type: (Any, Any, bool) -> None response = get_first_from_sources(event, STREAM_RESPONSE_SOURCES) if response is not None: - _collect_v1_response_fields(span, response, include_pii) + _collect_response_fields(span, response, include_pii) -def _collect_v1_response_fields(span, response, include_pii): +def _collect_response_fields(span, response, include_pii): # type: (Any, Any, bool) -> None if include_pii: text = get_first_from_sources(response, CHAT_RESPONSE_TEXT_SOURCES) diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index a9e307d574..3d9a4b9d26 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -46,6 +46,17 @@ except ImportError: _has_v2 = False +COHERE_V2_CHAT_CONFIG = { + "static": { + SPANDATA.GEN_AI_SYSTEM: "cohere", + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + }, + "pii_params": { + "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + }, + "extract_messages": lambda kw: _extract_messages_v2(kw.get("messages", [])), +} + CHAT_RESPONSE_SOURCES = { SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], @@ -55,7 +66,7 @@ } CHAT_USAGE_SOURCES = [("usage",)] STREAM_DELTA_TEXT_SOURCES = [("delta", "message", "content", "text")] -STREAM_CHAT_RESPONSE_SOURCES = { +STREAM_CHAT_RESPONSE_SOURCES: "dict[str, list[tuple[str, ...]]]" = { SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("delta", "finish_reason")], } @@ -104,29 +115,17 @@ def new_chat(*args, **kwargs): reraise(*exc_info) with capture_internal_exceptions(): - extra = {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming} - if model: - set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) - extra[SPANDATA.GEN_AI_RESPONSE_MODEL] = model + span_data = { + SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming, + SPANDATA.GEN_AI_REQUEST_MODEL: model if model else None, + SPANDATA.GEN_AI_RESPONSE_MODEL: model if model else None, + } set_input_span_data( - span, - kwargs, - integration, - { - "system": "cohere", - "operation": "chat", - "pii_params": { - "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - }, - "extract_messages": lambda kw: _extract_messages_v2( - kw.get("messages", []) - ), - "extra_static": extra, - }, + span, kwargs, integration, COHERE_V2_CHAT_CONFIG, span_data ) if streaming: return _iter_v2_stream_events(res, span, include_pii) - _collect_v2_response_fields(span, res, include_pii=include_pii) + _collect_v2_response_fields(span, res, include_pii) return res return new_chat @@ -146,7 +145,7 @@ def _extract_messages_v2(messages): def _iter_v2_stream_events(old_iterator, span, include_pii): # type: (Any, Span, bool) -> Iterator[V2ChatStreamResponse] - collected_text = [] + collected_text = [] # type: list[str] with capture_internal_exceptions(): for x in old_iterator: _append_stream_delta_text(collected_text, x) From 33c63f6e6d9788441c0ca604a3d9765d6a8ac6c6 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 10:41:36 +0100 Subject: [PATCH 20/29] fix --- sentry_sdk/ai/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py index 7e14c15e24..0f104cc8f5 100644 --- a/sentry_sdk/ai/utils.py +++ b/sentry_sdk/ai/utils.py @@ -503,8 +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. """ - role_lower = role.lower() - return GEN_AI_MESSAGE_ROLE_MAPPING.get(role_lower, role_lower) + return GEN_AI_MESSAGE_ROLE_MAPPING.get(role, role) def normalize_message_roles(messages: "list[dict[str, Any]]") -> "list[dict[str, Any]]": From dadbc7957c1dfacd0c532e07c68a0fc8227e8a52 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 11:08:20 +0100 Subject: [PATCH 21/29] clarify --- sentry_sdk/ai/span_config.py | 53 +++++++- sentry_sdk/ai/utils.py | 31 ++++- sentry_sdk/integrations/cohere/__init__.py | 4 +- sentry_sdk/integrations/cohere/utils.py | 35 ----- sentry_sdk/integrations/cohere/v1.py | 113 +++++++--------- sentry_sdk/integrations/cohere/v2.py | 148 ++++++++++----------- 6 files changed, 198 insertions(+), 186 deletions(-) delete mode 100644 sentry_sdk/integrations/cohere/utils.py diff --git a/sentry_sdk/ai/span_config.py b/sentry_sdk/ai/span_config.py index ef38160ef1..60ad9883de 100644 --- a/sentry_sdk/ai/span_config.py +++ b/sentry_sdk/ai/span_config.py @@ -1,7 +1,10 @@ import sentry_sdk from sentry_sdk.consts import SPANDATA +from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.ai.utils import ( + get_first_from_sources, set_data_normalized, + set_span_data_from_sources, normalize_message_roles, truncate_and_annotate_messages, ) @@ -10,11 +13,11 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Dict + from typing import Any, Dict, List, Optional from sentry_sdk.tracing import Span -def set_input_span_data(span, kwargs, integration, config, span_data=None): +def set_request_span_data(span, kwargs, integration, config, span_data=None): # type: (Span, Dict[str, Any], Any, Dict[str, Any], Dict[str, Any] | None) -> None """ Set input span data from a declarative config. @@ -57,3 +60,49 @@ def set_input_span_data(span, kwargs, integration, config, span_data=None): if kwarg_key in kwargs: value = kwargs[kwarg_key] set_data_normalized(span, span_attr, value) + + +def set_response_span_data(span, response, include_pii, response_config, collected_text=None): + # type: (Span, Any, bool, Dict[str, Any], Optional[List[str]]) -> None + """ + Set response span data from a declarative config. + + response_config keys: + sources: dict - always set from response object + pii_sources: dict - only when PII allowed + extract_text: (response) -> list[str] | None (PII only) + usage: dict with input_tokens/output_tokens source paths + collected_text: pre-collected streaming text (overrides extract_text) + """ + set_span_data_from_sources( + span, response, response_config.get("sources", {}), require_truthy=False + ) + + if include_pii: + pii_sources = response_config.get("pii_sources") + if pii_sources: + set_span_data_from_sources( + span, response, pii_sources, require_truthy=True + ) + if collected_text: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, ["".join(collected_text)] + ) + else: + extract_text = response_config.get("extract_text") + if extract_text: + texts = extract_text(response) + if texts: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts) + + usage_config = response_config.get("usage") + if usage_config: + record_token_usage( + span, + input_tokens=get_first_from_sources( + response, usage_config.get("input_tokens", []) + ), + output_tokens=get_first_from_sources( + response, usage_config.get("output_tokens", []) + ), + ) diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py index 0f104cc8f5..a05a068097 100644 --- a/sentry_sdk/ai/utils.py +++ b/sentry_sdk/ai/utils.py @@ -8,7 +8,7 @@ from sentry_sdk._types import BLOB_DATA_SUBSTITUTE if TYPE_CHECKING: - from typing import Any, Callable, Dict, List, Optional, Tuple + from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple from sentry_sdk.tracing import Span @@ -725,3 +725,32 @@ def set_conversation_id(conversation_id: str) -> None: """ scope = sentry_sdk.get_current_scope() scope.set_conversation_id(conversation_id) + + +def transitive_getattr(obj, *attrs): + # type: (Any, str) -> Any + current = obj + for attr in attrs: + current = getattr(current, attr, None) + if current is None: + return None + return current + + +def get_first_from_sources(obj, source_paths, require_truthy=False): + # type: (Any, Sequence[tuple[str, ...]], bool) -> Any + for source_path in source_paths: + value = transitive_getattr(obj, *source_path) + if not value: + continue + if not require_truthy or value: + return value + return None + + +def set_span_data_from_sources(span, obj, target_sources, require_truthy): + # type: (Any, Any, Mapping[str, Sequence[tuple[str, ...]]], bool) -> None + for spandata_key, source_paths in target_sources.items(): + value = get_first_from_sources(obj, source_paths, require_truthy=require_truthy) + if value is not None: + set_data_normalized(span, spandata_key, value) diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index ac6b7c474d..06bd11ac5d 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -3,7 +3,7 @@ from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.ai.span_config import set_input_span_data +from sentry_sdk.ai.span_config import set_request_span_data from typing import TYPE_CHECKING @@ -92,7 +92,7 @@ def new_embed(*args, **kwargs): name=f"embeddings {model}".strip(), origin=CohereIntegration.origin, ) as span: - set_input_span_data(span, kwargs, integration, COHERE_EMBED_CONFIG) + set_request_span_data(span, kwargs, integration, COHERE_EMBED_CONFIG) try: res = f(*args, **kwargs) diff --git a/sentry_sdk/integrations/cohere/utils.py b/sentry_sdk/integrations/cohere/utils.py deleted file mode 100644 index 346ac72a6a..0000000000 --- a/sentry_sdk/integrations/cohere/utils.py +++ /dev/null @@ -1,35 +0,0 @@ -from sentry_sdk.ai.utils import set_data_normalized - -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from typing import Any, Mapping, Sequence - - -def transitive_getattr(obj, *attrs): - # type: (Any, str) -> Any - current = obj - for attr in attrs: - current = getattr(current, attr, None) - if current is None: - return None - return current - - -def get_first_from_sources(obj, source_paths, require_truthy=False): - # type: (Any, Sequence[tuple[str, ...]], bool) -> Any - for source_path in source_paths: - value = transitive_getattr(obj, *source_path) - if not value: - continue - if not require_truthy or value: - return value - return None - - -def set_span_data_from_sources(span, obj, target_sources, require_truthy): - # type: (Any, Any, Mapping[str, Sequence[tuple[str, ...]]], bool) -> None - for spandata_key, source_paths in target_sources.items(): - value = get_first_from_sources(obj, source_paths, require_truthy=require_truthy) - if value is not None: - set_data_normalized(span, spandata_key, value) diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index 6b12d5d67d..b85fbfad12 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -1,9 +1,11 @@ import sys from functools import wraps -from sentry_sdk.ai.monitoring import record_token_usage -from sentry_sdk.ai.span_config import set_input_span_data -from sentry_sdk.ai.utils import set_data_normalized, transform_message_content +from sentry_sdk.ai.span_config import set_request_span_data, set_response_span_data +from sentry_sdk.ai.utils import ( + get_first_from_sources, + transform_message_content, +) from sentry_sdk.consts import OP, SPANDATA from typing import TYPE_CHECKING @@ -17,10 +19,6 @@ CohereIntegration, _capture_exception, ) -from sentry_sdk.integrations.cohere.utils import ( - get_first_from_sources, - set_span_data_from_sources, -) from sentry_sdk.scope import should_send_default_pii from sentry_sdk.utils import capture_internal_exceptions, reraise @@ -38,34 +36,41 @@ except ImportError: _has_chat_types = False +def _extract_response_text(response): + # type: (Any) -> list[str] | None + text = getattr(response, "text", None) + return [text] if text is not None else None + + COHERE_V1_CHAT_CONFIG = { "static": { SPANDATA.GEN_AI_SYSTEM: "cohere", SPANDATA.GEN_AI_OPERATION_NAME: "chat", }, "extract_messages": lambda kw: _extract_messages(kw), + "response": { + "sources": { + SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], + }, + "pii_sources": { + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("tool_calls",)], + }, + "extract_text": _extract_response_text, + "usage": { + "input_tokens": [ + ("meta", "billed_units", "input_tokens"), + ("meta", "tokens", "input_tokens"), + ], + "output_tokens": [ + ("meta", "billed_units", "output_tokens"), + ("meta", "tokens", "output_tokens"), + ], + }, + }, + "stream_response_object": [("response",)], } -CHAT_RESPONSE_SOURCES = { - SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)], - SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], -} -PII_CHAT_RESPONSE_SOURCES = { - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("tool_calls",)], -} -CHAT_RESPONSE_TEXT_SOURCES = [("text",)] -CHAT_USAGE_TOKEN_SOURCES = { - SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: [ - ("meta", "billed_units", "input_tokens"), - ("meta", "tokens", "input_tokens"), - ], - SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: [ - ("meta", "billed_units", "output_tokens"), - ("meta", "tokens", "output_tokens"), - ], -} -STREAM_RESPONSE_SOURCES = [("response",)] - def setup_v1(wrap_embed_fn): # type: (Callable[..., Any]) -> None @@ -113,19 +118,20 @@ def new_chat(*args, **kwargs): reraise(*exc_info) with capture_internal_exceptions(): - span_data = {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming} - if model: - span_data[SPANDATA.GEN_AI_REQUEST_MODEL] = model - set_input_span_data( + span_data = { + SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming, + SPANDATA.GEN_AI_REQUEST_MODEL: model if model else None, + } + set_request_span_data( span, kwargs, integration, COHERE_V1_CHAT_CONFIG, span_data ) if streaming: return _iter_stream_events(res, span, include_pii) - if isinstance(res, NonStreamedChatResponse): - _collect_response_fields(span, res, include_pii=include_pii) else: - set_data_normalized(span, "unknown_response", True) + set_response_span_data( + span, res, include_pii, COHERE_V1_CHAT_CONFIG["response"] + ) return res return new_chat @@ -154,36 +160,11 @@ def _iter_stream_events(old_iterator, span, include_pii): if isinstance(x, ChatStreamEndEvent) or isinstance( x, StreamEndStreamedChatResponse ): - _collect_v1_stream_end_fields(span, x, include_pii) + response = get_first_from_sources( + x, COHERE_V1_CHAT_CONFIG["stream_response_object"] + ) + if response is not None: + set_response_span_data( + span, response, include_pii, COHERE_V1_CHAT_CONFIG["response"] + ) yield x - - -def _collect_v1_stream_end_fields(span, event, include_pii): - # type: (Any, Any, bool) -> None - response = get_first_from_sources(event, STREAM_RESPONSE_SOURCES) - if response is not None: - _collect_response_fields(span, response, include_pii) - - -def _collect_response_fields(span, response, include_pii): - # type: (Any, Any, bool) -> None - if include_pii: - text = get_first_from_sources(response, CHAT_RESPONSE_TEXT_SOURCES) - if text is not None: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, [text]) - set_span_data_from_sources( - span, response, PII_CHAT_RESPONSE_SOURCES, require_truthy=False - ) - - set_span_data_from_sources( - span, response, CHAT_RESPONSE_SOURCES, require_truthy=False - ) - record_token_usage( - span, - input_tokens=get_first_from_sources( - response, CHAT_USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] - ), - output_tokens=get_first_from_sources( - response, CHAT_USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] - ), - ) diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index 3d9a4b9d26..fb556956b5 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -1,9 +1,12 @@ import sys from functools import wraps -from sentry_sdk.ai.monitoring import record_token_usage -from sentry_sdk.ai.span_config import set_input_span_data -from sentry_sdk.ai.utils import set_data_normalized, transform_message_content +from sentry_sdk.ai.span_config import set_request_span_data, set_response_span_data +from sentry_sdk.ai.utils import ( + get_first_from_sources, + transitive_getattr, + transform_message_content, +) from sentry_sdk.consts import OP, SPANDATA from typing import TYPE_CHECKING @@ -17,11 +20,6 @@ CohereIntegration, _capture_exception, ) -from sentry_sdk.integrations.cohere.utils import ( - get_first_from_sources, - set_span_data_from_sources, - transitive_getattr, -) from sentry_sdk.scope import should_send_default_pii from sentry_sdk.utils import capture_internal_exceptions, reraise @@ -46,6 +44,30 @@ except ImportError: _has_v2 = False +STREAM_DELTA_TEXT_SOURCES = [("delta", "message", "content", "text")] + +_V2_USAGE = { + "input_tokens": [ + ("usage", "billed_units", "input_tokens"), + ("usage", "tokens", "input_tokens"), + ], + "output_tokens": [ + ("usage", "billed_units", "output_tokens"), + ("usage", "tokens", "output_tokens"), + ], +} + + +def _extract_v2_response_text(response): + # type: (Any) -> list[str] | None + content = get_first_from_sources(response, [("message", "content")], True) + if content: + texts = [item.text for item in content if hasattr(item, "text")] + if texts: + return texts + return None + + COHERE_V2_CHAT_CONFIG = { "static": { SPANDATA.GEN_AI_SYSTEM: "cohere", @@ -55,26 +77,33 @@ "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, }, "extract_messages": lambda kw: _extract_messages_v2(kw.get("messages", [])), -} - -CHAT_RESPONSE_SOURCES = { - SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], - SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], -} -PII_CHAT_RESPONSE_SOURCES = { - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("message", "tool_calls")], -} -CHAT_USAGE_SOURCES = [("usage",)] -STREAM_DELTA_TEXT_SOURCES = [("delta", "message", "content", "text")] -STREAM_CHAT_RESPONSE_SOURCES: "dict[str, list[tuple[str, ...]]]" = { - SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], - SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("delta", "finish_reason")], -} -STREAM_CHAT_USAGE_SOURCES = [("delta", "usage")] -USAGE_UNIT_SOURCES = [("billed_units",), ("tokens",)] -USAGE_TOKEN_SOURCES = { - SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: [("input_tokens",)], - SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: [("output_tokens",)], + "response": { + "sources": { + SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], + }, + "pii_sources": { + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("message", "tool_calls")], + }, + "extract_text": _extract_v2_response_text, + "usage": _V2_USAGE, + }, + "stream_response": { + "sources": { + SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("delta", "finish_reason")], + }, + "usage": { + "input_tokens": [ + ("delta", "usage", "billed_units", "input_tokens"), + ("delta", "usage", "tokens", "input_tokens"), + ], + "output_tokens": [ + ("delta", "usage", "billed_units", "output_tokens"), + ("delta", "usage", "tokens", "output_tokens"), + ], + }, + }, } @@ -120,12 +149,14 @@ def new_chat(*args, **kwargs): SPANDATA.GEN_AI_REQUEST_MODEL: model if model else None, SPANDATA.GEN_AI_RESPONSE_MODEL: model if model else None, } - set_input_span_data( + set_request_span_data( span, kwargs, integration, COHERE_V2_CHAT_CONFIG, span_data ) if streaming: return _iter_v2_stream_events(res, span, include_pii) - _collect_v2_response_fields(span, res, include_pii) + set_response_span_data( + span, res, include_pii, COHERE_V2_CHAT_CONFIG["response"] + ) return res return new_chat @@ -150,7 +181,13 @@ def _iter_v2_stream_events(old_iterator, span, include_pii): for x in old_iterator: _append_stream_delta_text(collected_text, x) if isinstance(x, MessageEndV2ChatStreamResponse): - _collect_v2_stream_end_fields(span, x, include_pii, collected_text) + set_response_span_data( + span, + x, + include_pii, + COHERE_V2_CHAT_CONFIG["stream_response"], + collected_text, + ) yield x @@ -161,52 +198,3 @@ def _append_stream_delta_text(collected_text, event): content_text = get_first_from_sources(event, STREAM_DELTA_TEXT_SOURCES) if content_text is not None: collected_text.append(content_text) - - -def _collect_v2_stream_end_fields(span, event, include_pii, collected_text): - # type: (Span, Any, bool, list[str]) -> None - if include_pii and collected_text: - set_data_normalized( - span, SPANDATA.GEN_AI_RESPONSE_TEXT, ["".join(collected_text)] - ) - set_span_data_from_sources( - span, event, STREAM_CHAT_RESPONSE_SOURCES, require_truthy=False - ) - stream_usage = get_first_from_sources(event, STREAM_CHAT_USAGE_SOURCES) - if stream_usage is not None: - _record_token_usage_v2(span, stream_usage) - - -def _collect_v2_response_fields(span, response, include_pii): - # type: (Span, V2ChatResponse, bool) -> None - if include_pii: - content = get_first_from_sources(response, [("message", "content")], True) - if content: - texts = [item.text for item in content if hasattr(item, "text")] - if texts: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts) - set_span_data_from_sources( - span, response, PII_CHAT_RESPONSE_SOURCES, require_truthy=True - ) - set_span_data_from_sources( - span, response, CHAT_RESPONSE_SOURCES, require_truthy=False - ) - usage = get_first_from_sources(response, CHAT_USAGE_SOURCES) - if usage is not None: - _record_token_usage_v2(span, usage) - - -def _record_token_usage_v2(span, usage): - # type: (Span, Any) -> None - usage_obj = get_first_from_sources(usage, USAGE_UNIT_SOURCES) - if usage_obj is None: - return - record_token_usage( - span, - input_tokens=get_first_from_sources( - usage_obj, USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] - ), - output_tokens=get_first_from_sources( - usage_obj, USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] - ), - ) From 1cba43308c5610b6d3a924717d71d0eb464bdbcb Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 11:30:17 +0100 Subject: [PATCH 22/29] response model from response --- sentry_sdk/ai/span_config.py | 8 ++++---- sentry_sdk/integrations/cohere/v1.py | 2 ++ sentry_sdk/integrations/cohere/v2.py | 2 +- tests/integrations/cohere/test_cohere.py | 1 + 4 files changed, 8 insertions(+), 5 deletions(-) diff --git a/sentry_sdk/ai/span_config.py b/sentry_sdk/ai/span_config.py index 60ad9883de..60aaf51fcb 100644 --- a/sentry_sdk/ai/span_config.py +++ b/sentry_sdk/ai/span_config.py @@ -62,7 +62,9 @@ def set_request_span_data(span, kwargs, integration, config, span_data=None): set_data_normalized(span, span_attr, value) -def set_response_span_data(span, response, include_pii, response_config, collected_text=None): +def set_response_span_data( + span, response, include_pii, response_config, collected_text=None +): # type: (Span, Any, bool, Dict[str, Any], Optional[List[str]]) -> None """ Set response span data from a declarative config. @@ -81,9 +83,7 @@ def set_response_span_data(span, response, include_pii, response_config, collect if include_pii: pii_sources = response_config.get("pii_sources") if pii_sources: - set_span_data_from_sources( - span, response, pii_sources, require_truthy=True - ) + set_span_data_from_sources(span, response, pii_sources, require_truthy=True) if collected_text: set_data_normalized( span, SPANDATA.GEN_AI_RESPONSE_TEXT, ["".join(collected_text)] diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index b85fbfad12..ffcff603db 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -36,6 +36,7 @@ except ImportError: _has_chat_types = False + def _extract_response_text(response): # type: (Any) -> list[str] | None text = getattr(response, "text", None) @@ -50,6 +51,7 @@ def _extract_response_text(response): "extract_messages": lambda kw: _extract_messages(kw), "response": { "sources": { + SPANDATA.GEN_AI_RESPONSE_MODEL: [("model",)], SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)], SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], }, diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index fb556956b5..3014df1047 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -79,6 +79,7 @@ def _extract_v2_response_text(response): "extract_messages": lambda kw: _extract_messages_v2(kw.get("messages", [])), "response": { "sources": { + SPANDATA.GEN_AI_RESPONSE_MODEL: [("model",)], SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], }, @@ -147,7 +148,6 @@ def new_chat(*args, **kwargs): span_data = { SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming, SPANDATA.GEN_AI_REQUEST_MODEL: model if model else None, - SPANDATA.GEN_AI_RESPONSE_MODEL: model if model else None, } set_request_span_data( span, kwargs, integration, COHERE_V2_CHAT_CONFIG, span_data diff --git a/tests/integrations/cohere/test_cohere.py b/tests/integrations/cohere/test_cohere.py index 3da2f616ed..78c8e97a27 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -280,6 +280,7 @@ def test_v2_nonstreaming_chat( 200, json={ "id": "resp-123", + "model": "some-model", "finish_reason": "COMPLETE", "message": { "role": "assistant", From 43bad0ca1696955118d0459160fa439935cef1a2 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 11:34:35 +0100 Subject: [PATCH 23/29] make cost declarative in embeddings --- sentry_sdk/ai/span_config.py | 3 +++ sentry_sdk/integrations/cohere/__init__.py | 22 ++++++++++------------ 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/sentry_sdk/ai/span_config.py b/sentry_sdk/ai/span_config.py index 60aaf51fcb..b3f77baab8 100644 --- a/sentry_sdk/ai/span_config.py +++ b/sentry_sdk/ai/span_config.py @@ -105,4 +105,7 @@ def set_response_span_data( output_tokens=get_first_from_sources( response, usage_config.get("output_tokens", []) ), + total_tokens=get_first_from_sources( + response, usage_config.get("total_tokens", []) + ), ) diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index 06bd11ac5d..2d98adea74 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -1,9 +1,8 @@ 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.span_config import set_request_span_data +from sentry_sdk.ai.span_config import set_request_span_data, set_response_span_data from typing import TYPE_CHECKING @@ -42,6 +41,12 @@ def _normalize_embedding_input(texts): ), "message_target": SPANDATA.GEN_AI_EMBEDDINGS_INPUT, "truncation_fn": None, + "response": { + "usage": { + "input_tokens": [("meta", "billed_units", "input_tokens")], + "total_tokens": [("meta", "billed_units", "input_tokens")], + }, + }, } @@ -102,16 +107,9 @@ def new_embed(*args, **kwargs): _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, - ) + set_response_span_data( + span, res, False, COHERE_EMBED_CONFIG["response"] + ) return res return new_embed From 3e077fe168c9937cc8649722fd63bbe668ce465f Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 11:49:13 +0100 Subject: [PATCH 24/29] typing --- sentry_sdk/ai/span_config.py | 77 +++++++++++++++------- sentry_sdk/integrations/cohere/__init__.py | 3 +- sentry_sdk/integrations/cohere/v1.py | 37 ++++++----- sentry_sdk/integrations/cohere/v2.py | 56 ++++++++-------- 4 files changed, 103 insertions(+), 70 deletions(-) diff --git a/sentry_sdk/ai/span_config.py b/sentry_sdk/ai/span_config.py index b3f77baab8..ba4ee84610 100644 --- a/sentry_sdk/ai/span_config.py +++ b/sentry_sdk/ai/span_config.py @@ -13,24 +13,64 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Dict, List, Optional + from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + from typing_extensions import TypedDict + from sentry_sdk.tracing import Span + # Source paths: list of attribute chains to try in order. + # e.g. [("meta", "billed_units", "input_tokens"), ("meta", "tokens", "input_tokens")] + SourcePaths = Sequence[Tuple[str, ...]] + + # Maps a SPANDATA key to source paths on the response object. + # e.g. {SPANDATA.GEN_AI_RESPONSE_ID: [("id",)]} + SourceMapping = Dict[str, SourcePaths] + + class UsageConfig(TypedDict, total=False): + """Declarative token usage extraction paths (from response object).""" + + input_tokens: SourcePaths + output_tokens: SourcePaths + total_tokens: SourcePaths + + class ResponseConfig(TypedDict, total=False): + """Declarative response span data config.""" + + # Attributes always extracted from the response object. + sources: SourceMapping + # Attributes extracted only when PII sending is enabled. + pii_sources: SourceMapping + # Custom extractor for response text (PII only). + # Returns list of text strings, or None. + extract_text: Callable[[Any], Optional[List[str]]] + # Declarative token usage paths. + usage: UsageConfig + + class OperationConfig(TypedDict, total=False): + """Full declarative config for an AI operation (chat, embeddings, etc.).""" + + # Key/value pairs set on every span unconditionally. + static: Dict[str, Any] + # Maps kwarg names to SPANDATA keys (always set if present in kwargs). + params: Dict[str, str] + # Maps kwarg names to SPANDATA keys (only set when PII is enabled). + pii_params: Dict[str, str] + # Extracts messages from kwargs for the span. + extract_messages: Callable[[Dict[str, Any]], Optional[List[Dict[str, Any]]]] + # SPANDATA key for messages (default: GEN_AI_REQUEST_MESSAGES). + message_target: str + # Non-streaming response config. + response: ResponseConfig + # Streaming response config (different attribute paths). + stream_response: ResponseConfig + # Source paths to extract a full response object from a stream-end event + # (V1 pattern: reuse "response" config after extracting). + stream_response_object: SourcePaths + def set_request_span_data(span, kwargs, integration, config, span_data=None): # type: (Span, Dict[str, Any], Any, Dict[str, Any], Dict[str, Any] | None) -> None - """ - Set input span data from a declarative config. - - Config keys: - static: dict - key/value pairs to set unconditionally - params: dict - kwargs key -> span attr (always set if present) - pii_params: dict - kwargs key -> span attr (only when PII allowed) - extract_messages: callable(kwargs) -> list or None - message_target: str - span attr for messages (default: GEN_AI_REQUEST_MESSAGES) - - span_data: additional key/value pairs for dynamic per-call values - """ + """Set request/static span data from a declarative config.""" for key, value in config.get("static", {}).items(): set_data_normalized(span, key, value) if span_data: @@ -66,16 +106,7 @@ def set_response_span_data( span, response, include_pii, response_config, collected_text=None ): # type: (Span, Any, bool, Dict[str, Any], Optional[List[str]]) -> None - """ - Set response span data from a declarative config. - - response_config keys: - sources: dict - always set from response object - pii_sources: dict - only when PII allowed - extract_text: (response) -> list[str] | None (PII only) - usage: dict with input_tokens/output_tokens source paths - collected_text: pre-collected streaming text (overrides extract_text) - """ + """Set response span data from a declarative config.""" set_span_data_from_sources( span, response, response_config.get("sources", {}), require_truthy=False ) diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index 2d98adea74..b346664c60 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -10,6 +10,7 @@ if TYPE_CHECKING: from typing import Any, Callable + from sentry_sdk.ai.span_config import OperationConfig import sentry_sdk from sentry_sdk.integrations import DidNotEnable, Integration @@ -30,7 +31,7 @@ def _normalize_embedding_input(texts): return [texts] -COHERE_EMBED_CONFIG = { +COHERE_EMBED_CONFIG: "OperationConfig" = { "static": { SPANDATA.GEN_AI_SYSTEM: "cohere", SPANDATA.GEN_AI_OPERATION_NAME: "embeddings", diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index ffcff603db..cad7c4b2ec 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -13,6 +13,7 @@ if TYPE_CHECKING: from typing import Any, Callable, Iterator from cohere import StreamedChatResponse + from sentry_sdk.ai.span_config import OperationConfig import sentry_sdk from sentry_sdk.integrations.cohere import ( @@ -37,13 +38,26 @@ _has_chat_types = False +def setup_v1(wrap_embed_fn): + # type: (Callable[..., Any]) -> None + try: + from cohere.base_client import BaseCohere + from cohere.client import Client + except ImportError: + return + + 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 _extract_response_text(response): # type: (Any) -> list[str] | None text = getattr(response, "text", None) return [text] if text is not None else None -COHERE_V1_CHAT_CONFIG = { +COHERE_V1_CHAT_CONFIG: "OperationConfig" = { "static": { SPANDATA.GEN_AI_SYSTEM: "cohere", SPANDATA.GEN_AI_OPERATION_NAME: "chat", @@ -74,19 +88,6 @@ def _extract_response_text(response): } -def setup_v1(wrap_embed_fn): - # type: (Callable[..., Any]) -> None - try: - from cohere.base_client import BaseCohere - from cohere.client import Client - except ImportError: - return - - 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 _wrap_chat(f, streaming): # type: (Callable[..., Any], bool) -> Callable[..., Any] if not _has_chat_types: @@ -112,7 +113,7 @@ def new_chat(*args, **kwargs): origin=CohereIntegration.origin, ) as span: try: - res = f(*args, **kwargs) + response = f(*args, **kwargs) except Exception as e: exc_info = sys.exc_info() with capture_internal_exceptions(): @@ -129,12 +130,12 @@ def new_chat(*args, **kwargs): ) if streaming: - return _iter_stream_events(res, span, include_pii) + return _iter_stream_events(response, span, include_pii) else: set_response_span_data( - span, res, include_pii, COHERE_V1_CHAT_CONFIG["response"] + span, response, include_pii, COHERE_V1_CHAT_CONFIG["response"] ) - return res + return response return new_chat diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index 3014df1047..94863d7045 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: from typing import Any, Callable, Iterator from sentry_sdk.tracing import Span + from sentry_sdk.ai.span_config import OperationConfig import sentry_sdk from sentry_sdk.integrations.cohere import ( @@ -44,18 +45,19 @@ except ImportError: _has_v2 = False -STREAM_DELTA_TEXT_SOURCES = [("delta", "message", "content", "text")] -_V2_USAGE = { - "input_tokens": [ - ("usage", "billed_units", "input_tokens"), - ("usage", "tokens", "input_tokens"), - ], - "output_tokens": [ - ("usage", "billed_units", "output_tokens"), - ("usage", "tokens", "output_tokens"), - ], -} +def setup_v2(wrap_embed_fn): + # type: (Callable[..., Any]) -> None + 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) + + +STREAM_DELTA_TEXT_SOURCES = [("delta", "message", "content", "text")] def _extract_v2_response_text(response): @@ -68,7 +70,7 @@ def _extract_v2_response_text(response): return None -COHERE_V2_CHAT_CONFIG = { +COHERE_V2_CHAT_CONFIG: "OperationConfig" = { "static": { SPANDATA.GEN_AI_SYSTEM: "cohere", SPANDATA.GEN_AI_OPERATION_NAME: "chat", @@ -87,7 +89,16 @@ def _extract_v2_response_text(response): SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("message", "tool_calls")], }, "extract_text": _extract_v2_response_text, - "usage": _V2_USAGE, + "usage": { + "input_tokens": [ + ("usage", "billed_units", "input_tokens"), + ("usage", "tokens", "input_tokens"), + ], + "output_tokens": [ + ("usage", "billed_units", "output_tokens"), + ("usage", "tokens", "output_tokens"), + ], + }, }, "stream_response": { "sources": { @@ -108,17 +119,6 @@ def _extract_v2_response_text(response): } -def setup_v2(wrap_embed_fn): - # type: (Callable[..., Any]) -> None - 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 _wrap_chat_v2(f, streaming): # type: (Callable[..., Any], bool) -> Callable[..., Any] @wraps(f) @@ -137,7 +137,7 @@ def new_chat(*args, **kwargs): origin=CohereIntegration.origin, ) as span: try: - res = f(*args, **kwargs) + response = f(*args, **kwargs) except Exception as e: exc_info = sys.exc_info() with capture_internal_exceptions(): @@ -153,11 +153,11 @@ def new_chat(*args, **kwargs): span, kwargs, integration, COHERE_V2_CHAT_CONFIG, span_data ) if streaming: - return _iter_v2_stream_events(res, span, include_pii) + return _iter_v2_stream_events(response, span, include_pii) set_response_span_data( - span, res, include_pii, COHERE_V2_CHAT_CONFIG["response"] + span, response, include_pii, COHERE_V2_CHAT_CONFIG["response"] ) - return res + return response return new_chat From 1e975ceb085c017aaa2af02ec75607782bece043 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 11:53:41 +0100 Subject: [PATCH 25/29] move configs to a separate file --- sentry_sdk/integrations/cohere/__init__.py | 33 +--- sentry_sdk/integrations/cohere/configs.py | 172 +++++++++++++++++++++ sentry_sdk/integrations/cohere/v1.py | 60 +------ sentry_sdk/integrations/cohere/v2.py | 85 +--------- 4 files changed, 181 insertions(+), 169 deletions(-) create mode 100644 sentry_sdk/integrations/cohere/configs.py diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index b346664c60..0f2989fefe 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -1,8 +1,9 @@ import sys from functools import wraps -from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.consts import OP from sentry_sdk.ai.span_config import set_request_span_data, set_response_span_data +from sentry_sdk.integrations.cohere.configs import COHERE_EMBED_CONFIG from typing import TYPE_CHECKING @@ -10,7 +11,6 @@ if TYPE_CHECKING: from typing import Any, Callable - from sentry_sdk.ai.span_config import OperationConfig import sentry_sdk from sentry_sdk.integrations import DidNotEnable, Integration @@ -22,35 +22,6 @@ raise DidNotEnable("Cohere not installed") -def _normalize_embedding_input(texts): - # type: (Any) -> Any - if isinstance(texts, list): - return texts - if isinstance(texts, tuple): - return list(texts) - return [texts] - - -COHERE_EMBED_CONFIG: "OperationConfig" = { - "static": { - SPANDATA.GEN_AI_SYSTEM: "cohere", - SPANDATA.GEN_AI_OPERATION_NAME: "embeddings", - }, - "params": {"model": SPANDATA.GEN_AI_REQUEST_MODEL}, - "extract_messages": lambda kw: ( - _normalize_embedding_input(kw["texts"]) if "texts" in kw else None - ), - "message_target": SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - "truncation_fn": None, - "response": { - "usage": { - "input_tokens": [("meta", "billed_units", "input_tokens")], - "total_tokens": [("meta", "billed_units", "input_tokens")], - }, - }, -} - - class CohereIntegration(Integration): identifier = "cohere" origin = f"auto.ai.{identifier}" diff --git a/sentry_sdk/integrations/cohere/configs.py b/sentry_sdk/integrations/cohere/configs.py new file mode 100644 index 0000000000..9b17b503fb --- /dev/null +++ b/sentry_sdk/integrations/cohere/configs.py @@ -0,0 +1,172 @@ +from sentry_sdk.ai.utils import ( + get_first_from_sources, + transform_message_content, +) +from sentry_sdk.consts import SPANDATA + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from sentry_sdk.ai.span_config import OperationConfig + + +# ── Helpers ────────────────────────────────────────────────────────────────── + + +def _normalize_embedding_input(texts): + # type: (Any) -> Any + if isinstance(texts, list): + return texts + if isinstance(texts, tuple): + return list(texts) + return [texts] + + +def _extract_v1_messages(kwargs): + # type: (dict[str, Any]) -> list[dict[str, str]] + messages = [] + for x in kwargs.get("chat_history", []): + messages.append( + { + "role": getattr(x, "role", ""), + "content": transform_message_content(getattr(x, "message", "")), + } + ) + message = kwargs.get("message") + if message: + messages.append({"role": "user", "content": transform_message_content(message)}) + return messages + + +def _extract_v1_response_text(response): + # type: (Any) -> list[str] | None + text = getattr(response, "text", None) + return [text] if text is not None else None + + +def _extract_v2_messages(messages): + # type: (Any) -> list[dict[str, Any]] + result = [] + for msg in messages: + role = msg["role"] if isinstance(msg, dict) else getattr(msg, "role", "unknown") + content = ( + msg["content"] if isinstance(msg, dict) else getattr(msg, "content", "") + ) + result.append({"role": role, "content": transform_message_content(content)}) + return result + + +def _extract_v2_response_text(response): + # type: (Any) -> list[str] | None + content = get_first_from_sources(response, [("message", "content")], True) + if content: + texts = [item.text for item in content if hasattr(item, "text")] + if texts: + return texts + return None + + +# ── Configs ────────────────────────────────────────────────────────────────── + + +COHERE_EMBED_CONFIG: "OperationConfig" = { + "static": { + SPANDATA.GEN_AI_SYSTEM: "cohere", + SPANDATA.GEN_AI_OPERATION_NAME: "embeddings", + }, + "params": {"model": SPANDATA.GEN_AI_REQUEST_MODEL}, + "extract_messages": lambda kw: ( + _normalize_embedding_input(kw["texts"]) if "texts" in kw else None + ), + "message_target": SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + "response": { + "usage": { + "input_tokens": [("meta", "billed_units", "input_tokens")], + "total_tokens": [("meta", "billed_units", "input_tokens")], + }, + }, +} + + +COHERE_V1_CHAT_CONFIG: "OperationConfig" = { + "static": { + SPANDATA.GEN_AI_SYSTEM: "cohere", + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + }, + "extract_messages": lambda kw: _extract_v1_messages(kw), + "response": { + "sources": { + SPANDATA.GEN_AI_RESPONSE_MODEL: [("model",)], + SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], + }, + "pii_sources": { + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("tool_calls",)], + }, + "extract_text": _extract_v1_response_text, + "usage": { + "input_tokens": [ + ("meta", "billed_units", "input_tokens"), + ("meta", "tokens", "input_tokens"), + ], + "output_tokens": [ + ("meta", "billed_units", "output_tokens"), + ("meta", "tokens", "output_tokens"), + ], + }, + }, + "stream_response_object": [("response",)], +} + + +STREAM_DELTA_TEXT_SOURCES = [("delta", "message", "content", "text")] + + +COHERE_V2_CHAT_CONFIG: "OperationConfig" = { + "static": { + SPANDATA.GEN_AI_SYSTEM: "cohere", + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + }, + "pii_params": { + "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + }, + "extract_messages": lambda kw: _extract_v2_messages(kw.get("messages", [])), + "response": { + "sources": { + SPANDATA.GEN_AI_RESPONSE_MODEL: [("model",)], + SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], + }, + "pii_sources": { + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("message", "tool_calls")], + }, + "extract_text": _extract_v2_response_text, + "usage": { + "input_tokens": [ + ("usage", "billed_units", "input_tokens"), + ("usage", "tokens", "input_tokens"), + ], + "output_tokens": [ + ("usage", "billed_units", "output_tokens"), + ("usage", "tokens", "output_tokens"), + ], + }, + }, + "stream_response": { + "sources": { + SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("delta", "finish_reason")], + }, + "usage": { + "input_tokens": [ + ("delta", "usage", "billed_units", "input_tokens"), + ("delta", "usage", "tokens", "input_tokens"), + ], + "output_tokens": [ + ("delta", "usage", "billed_units", "output_tokens"), + ("delta", "usage", "tokens", "output_tokens"), + ], + }, + }, +} diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index cad7c4b2ec..2bac3bd415 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -2,18 +2,15 @@ from functools import wraps from sentry_sdk.ai.span_config import set_request_span_data, set_response_span_data -from sentry_sdk.ai.utils import ( - get_first_from_sources, - transform_message_content, -) +from sentry_sdk.ai.utils import get_first_from_sources from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations.cohere.configs import COHERE_V1_CHAT_CONFIG from typing import TYPE_CHECKING if TYPE_CHECKING: from typing import Any, Callable, Iterator from cohere import StreamedChatResponse - from sentry_sdk.ai.span_config import OperationConfig import sentry_sdk from sentry_sdk.integrations.cohere import ( @@ -51,43 +48,6 @@ def setup_v1(wrap_embed_fn): Client.embed = wrap_embed_fn(Client.embed) -def _extract_response_text(response): - # type: (Any) -> list[str] | None - text = getattr(response, "text", None) - return [text] if text is not None else None - - -COHERE_V1_CHAT_CONFIG: "OperationConfig" = { - "static": { - SPANDATA.GEN_AI_SYSTEM: "cohere", - SPANDATA.GEN_AI_OPERATION_NAME: "chat", - }, - "extract_messages": lambda kw: _extract_messages(kw), - "response": { - "sources": { - SPANDATA.GEN_AI_RESPONSE_MODEL: [("model",)], - SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)], - SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], - }, - "pii_sources": { - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("tool_calls",)], - }, - "extract_text": _extract_response_text, - "usage": { - "input_tokens": [ - ("meta", "billed_units", "input_tokens"), - ("meta", "tokens", "input_tokens"), - ], - "output_tokens": [ - ("meta", "billed_units", "output_tokens"), - ("meta", "tokens", "output_tokens"), - ], - }, - }, - "stream_response_object": [("response",)], -} - - def _wrap_chat(f, streaming): # type: (Callable[..., Any], bool) -> Callable[..., Any] if not _has_chat_types: @@ -140,22 +100,6 @@ def new_chat(*args, **kwargs): return new_chat -def _extract_messages(kwargs): - # type: (dict[str, Any]) -> list[dict[str, str]] - messages = [] - for x in kwargs.get("chat_history", []): - messages.append( - { - "role": getattr(x, "role", ""), - "content": transform_message_content(getattr(x, "message", "")), - } - ) - message = kwargs.get("message") - if message: - messages.append({"role": "user", "content": transform_message_content(message)}) - return messages - - def _iter_stream_events(old_iterator, span, include_pii): # type: (Any, Any, bool) -> Iterator[StreamedChatResponse] with capture_internal_exceptions(): diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index 94863d7045..bbb82e3c64 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -2,19 +2,18 @@ from functools import wraps from sentry_sdk.ai.span_config import set_request_span_data, set_response_span_data -from sentry_sdk.ai.utils import ( - get_first_from_sources, - transitive_getattr, - transform_message_content, -) +from sentry_sdk.ai.utils import get_first_from_sources, transitive_getattr from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations.cohere.configs import ( + COHERE_V2_CHAT_CONFIG, + STREAM_DELTA_TEXT_SOURCES, +) from typing import TYPE_CHECKING if TYPE_CHECKING: from typing import Any, Callable, Iterator from sentry_sdk.tracing import Span - from sentry_sdk.ai.span_config import OperationConfig import sentry_sdk from sentry_sdk.integrations.cohere import ( @@ -57,68 +56,6 @@ def setup_v2(wrap_embed_fn): CohereV2Client.embed = wrap_embed_fn(CohereV2Client.embed) -STREAM_DELTA_TEXT_SOURCES = [("delta", "message", "content", "text")] - - -def _extract_v2_response_text(response): - # type: (Any) -> list[str] | None - content = get_first_from_sources(response, [("message", "content")], True) - if content: - texts = [item.text for item in content if hasattr(item, "text")] - if texts: - return texts - return None - - -COHERE_V2_CHAT_CONFIG: "OperationConfig" = { - "static": { - SPANDATA.GEN_AI_SYSTEM: "cohere", - SPANDATA.GEN_AI_OPERATION_NAME: "chat", - }, - "pii_params": { - "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - }, - "extract_messages": lambda kw: _extract_messages_v2(kw.get("messages", [])), - "response": { - "sources": { - SPANDATA.GEN_AI_RESPONSE_MODEL: [("model",)], - SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], - SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)], - }, - "pii_sources": { - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("message", "tool_calls")], - }, - "extract_text": _extract_v2_response_text, - "usage": { - "input_tokens": [ - ("usage", "billed_units", "input_tokens"), - ("usage", "tokens", "input_tokens"), - ], - "output_tokens": [ - ("usage", "billed_units", "output_tokens"), - ("usage", "tokens", "output_tokens"), - ], - }, - }, - "stream_response": { - "sources": { - SPANDATA.GEN_AI_RESPONSE_ID: [("id",)], - SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("delta", "finish_reason")], - }, - "usage": { - "input_tokens": [ - ("delta", "usage", "billed_units", "input_tokens"), - ("delta", "usage", "tokens", "input_tokens"), - ], - "output_tokens": [ - ("delta", "usage", "billed_units", "output_tokens"), - ("delta", "usage", "tokens", "output_tokens"), - ], - }, - }, -} - - def _wrap_chat_v2(f, streaming): # type: (Callable[..., Any], bool) -> Callable[..., Any] @wraps(f) @@ -162,18 +99,6 @@ def new_chat(*args, **kwargs): return new_chat -def _extract_messages_v2(messages): - # type: (Any) -> list[dict[str, Any]] - result = [] - for msg in messages: - role = msg["role"] if isinstance(msg, dict) else getattr(msg, "role", "unknown") - content = ( - msg["content"] if isinstance(msg, dict) else getattr(msg, "content", "") - ) - result.append({"role": role, "content": transform_message_content(content)}) - return result - - def _iter_v2_stream_events(old_iterator, span, include_pii): # type: (Any, Span, bool) -> Iterator[V2ChatStreamResponse] collected_text = [] # type: list[str] From aa1dbfab7077e6226c8eac84f73b3485e0415f09 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 11:55:21 +0100 Subject: [PATCH 26/29] move type definition --- sentry_sdk/ai/span_config.py | 52 +--------------------- sentry_sdk/integrations/cohere/configs.py | 53 ++++++++++++++++++++++- 2 files changed, 52 insertions(+), 53 deletions(-) diff --git a/sentry_sdk/ai/span_config.py b/sentry_sdk/ai/span_config.py index ba4ee84610..becd6710d4 100644 --- a/sentry_sdk/ai/span_config.py +++ b/sentry_sdk/ai/span_config.py @@ -13,60 +13,10 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple - from typing_extensions import TypedDict + from typing import Any, Dict, List, Optional from sentry_sdk.tracing import Span - # Source paths: list of attribute chains to try in order. - # e.g. [("meta", "billed_units", "input_tokens"), ("meta", "tokens", "input_tokens")] - SourcePaths = Sequence[Tuple[str, ...]] - - # Maps a SPANDATA key to source paths on the response object. - # e.g. {SPANDATA.GEN_AI_RESPONSE_ID: [("id",)]} - SourceMapping = Dict[str, SourcePaths] - - class UsageConfig(TypedDict, total=False): - """Declarative token usage extraction paths (from response object).""" - - input_tokens: SourcePaths - output_tokens: SourcePaths - total_tokens: SourcePaths - - class ResponseConfig(TypedDict, total=False): - """Declarative response span data config.""" - - # Attributes always extracted from the response object. - sources: SourceMapping - # Attributes extracted only when PII sending is enabled. - pii_sources: SourceMapping - # Custom extractor for response text (PII only). - # Returns list of text strings, or None. - extract_text: Callable[[Any], Optional[List[str]]] - # Declarative token usage paths. - usage: UsageConfig - - class OperationConfig(TypedDict, total=False): - """Full declarative config for an AI operation (chat, embeddings, etc.).""" - - # Key/value pairs set on every span unconditionally. - static: Dict[str, Any] - # Maps kwarg names to SPANDATA keys (always set if present in kwargs). - params: Dict[str, str] - # Maps kwarg names to SPANDATA keys (only set when PII is enabled). - pii_params: Dict[str, str] - # Extracts messages from kwargs for the span. - extract_messages: Callable[[Dict[str, Any]], Optional[List[Dict[str, Any]]]] - # SPANDATA key for messages (default: GEN_AI_REQUEST_MESSAGES). - message_target: str - # Non-streaming response config. - response: ResponseConfig - # Streaming response config (different attribute paths). - stream_response: ResponseConfig - # Source paths to extract a full response object from a stream-end event - # (V1 pattern: reuse "response" config after extracting). - stream_response_object: SourcePaths - def set_request_span_data(span, kwargs, integration, config, span_data=None): # type: (Span, Dict[str, Any], Any, Dict[str, Any], Dict[str, Any] | None) -> None diff --git a/sentry_sdk/integrations/cohere/configs.py b/sentry_sdk/integrations/cohere/configs.py index 9b17b503fb..33d52952c3 100644 --- a/sentry_sdk/integrations/cohere/configs.py +++ b/sentry_sdk/integrations/cohere/configs.py @@ -7,8 +7,57 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any - from sentry_sdk.ai.span_config import OperationConfig + from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + from typing_extensions import TypedDict + + # Source paths: list of attribute chains to try in order. + # e.g. [("meta", "billed_units", "input_tokens"), ("meta", "tokens", "input_tokens")] + SourcePaths = Sequence[Tuple[str, ...]] + + # Maps a SPANDATA key to source paths on the response object. + # e.g. {SPANDATA.GEN_AI_RESPONSE_ID: [("id",)]} + SourceMapping = Dict[str, SourcePaths] + + class UsageConfig(TypedDict, total=False): + """Declarative token usage extraction paths (from response object).""" + + input_tokens: SourcePaths + output_tokens: SourcePaths + total_tokens: SourcePaths + + class ResponseConfig(TypedDict, total=False): + """Declarative response span data config.""" + + # Attributes always extracted from the response object. + sources: SourceMapping + # Attributes extracted only when PII sending is enabled. + pii_sources: SourceMapping + # Custom extractor for response text (PII only). + # Returns list of text strings, or None. + extract_text: Callable[[Any], Optional[List[str]]] + # Declarative token usage paths. + usage: UsageConfig + + class OperationConfig(TypedDict, total=False): + """Full declarative config for an AI operation (chat, embeddings, etc.).""" + + # Key/value pairs set on every span unconditionally. + static: Dict[str, Any] + # Maps kwarg names to SPANDATA keys (always set if present in kwargs). + params: Dict[str, str] + # Maps kwarg names to SPANDATA keys (only set when PII is enabled). + pii_params: Dict[str, str] + # Extracts messages from kwargs for the span. + extract_messages: Callable[[Dict[str, Any]], Optional[List[Dict[str, Any]]]] + # SPANDATA key for messages (default: GEN_AI_REQUEST_MESSAGES). + message_target: str + # Non-streaming response config. + response: ResponseConfig + # Streaming response config (different attribute paths). + stream_response: ResponseConfig + # Source paths to extract a full response object from a stream-end event + # (V1 pattern: reuse "response" config after extracting). + stream_response_object: SourcePaths # ── Helpers ────────────────────────────────────────────────────────────────── From 31769eecc1da11e264bf73739b13cfbf7eec53cd Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 11:56:11 +0100 Subject: [PATCH 27/29] fix --- sentry_sdk/integrations/cohere/v1.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index 2bac3bd415..e8a6da2695 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -91,10 +91,9 @@ def new_chat(*args, **kwargs): if streaming: return _iter_stream_events(response, span, include_pii) - else: - set_response_span_data( - span, response, include_pii, COHERE_V1_CHAT_CONFIG["response"] - ) + set_response_span_data( + span, response, include_pii, COHERE_V1_CHAT_CONFIG["response"] + ) return response return new_chat From c0a0e101d5b89f7dab75ce3371f81b9973ebd96e Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 12:44:41 +0100 Subject: [PATCH 28/29] fix swalling exception bug --- sentry_sdk/integrations/cohere/v1.py | 6 ++-- sentry_sdk/integrations/cohere/v2.py | 6 ++-- tests/integrations/cohere/test_cohere.py | 37 ++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 6 deletions(-) diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index e8a6da2695..4e3c4741e9 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -101,8 +101,8 @@ def new_chat(*args, **kwargs): def _iter_stream_events(old_iterator, span, include_pii): # type: (Any, Any, bool) -> Iterator[StreamedChatResponse] - with capture_internal_exceptions(): - for x in old_iterator: + for x in old_iterator: + with capture_internal_exceptions(): if isinstance(x, ChatStreamEndEvent) or isinstance( x, StreamEndStreamedChatResponse ): @@ -113,4 +113,4 @@ def _iter_stream_events(old_iterator, span, include_pii): set_response_span_data( span, response, include_pii, COHERE_V1_CHAT_CONFIG["response"] ) - yield x + yield x diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index bbb82e3c64..cf6677ff75 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -102,8 +102,8 @@ def new_chat(*args, **kwargs): def _iter_v2_stream_events(old_iterator, span, include_pii): # type: (Any, Span, bool) -> Iterator[V2ChatStreamResponse] collected_text = [] # type: list[str] - with capture_internal_exceptions(): - for x in old_iterator: + for x in old_iterator: + with capture_internal_exceptions(): _append_stream_delta_text(collected_text, x) if isinstance(x, MessageEndV2ChatStreamResponse): set_response_span_data( @@ -113,7 +113,7 @@ def _iter_v2_stream_events(old_iterator, span, include_pii): COHERE_V2_CHAT_CONFIG["stream_response"], collected_text, ) - yield x + yield x def _append_stream_delta_text(collected_text, event): diff --git a/tests/integrations/cohere/test_cohere.py b/tests/integrations/cohere/test_cohere.py index 78c8e97a27..d1a0657e54 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -182,6 +182,24 @@ def test_v1_bad_chat(sentry_init, capture_events): assert event["level"] == "error" +def test_v1_streaming_error_propagates(sentry_init, capture_events): + """Stream errors must not be silently swallowed by capture_internal_exceptions.""" + sentry_init(integrations=[CohereIntegration()], traces_sample_rate=1.0) + events = capture_events() + + from sentry_sdk.integrations.cohere.v1 import _iter_stream_events + + def failing_iterator(): + yield "event1" + raise ConnectionError("stream interrupted") + + with start_transaction(name="cohere tx") as tx: + span = tx.start_child(op="gen_ai.chat") + with pytest.raises(ConnectionError, match="stream interrupted"): + list(_iter_stream_events(failing_iterator(), span, False)) + span.finish() + + def test_v1_span_status_error(sentry_init, capture_events): sentry_init(integrations=[CohereIntegration()], traces_sample_rate=1.0) events = capture_events() @@ -410,6 +428,25 @@ def test_v2_streaming_chat( # --- V2 Error --- +@pytest.mark.skipif(not has_v2, reason="Cohere V2 client not available") +def test_v2_streaming_error_propagates(sentry_init, capture_events): + """Stream errors must not be silently swallowed by capture_internal_exceptions.""" + sentry_init(integrations=[CohereIntegration()], traces_sample_rate=1.0) + events = capture_events() + + from sentry_sdk.integrations.cohere.v2 import _iter_v2_stream_events + + def failing_iterator(): + yield "event1" + raise ConnectionError("stream interrupted") + + with start_transaction(name="cohere tx") as tx: + span = tx.start_child(op="gen_ai.chat") + with pytest.raises(ConnectionError, match="stream interrupted"): + list(_iter_v2_stream_events(failing_iterator(), span, False)) + span.finish() + + @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) From 3cbaad63d715c18e247015d8a500623f6f408b05 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 5 Mar 2026 12:56:17 +0100 Subject: [PATCH 29/29] move text massaging to core file --- sentry_sdk/ai/span_config.py | 44 ++++++------- sentry_sdk/integrations/cohere/__init__.py | 32 +++++++-- sentry_sdk/integrations/cohere/configs.py | 77 +--------------------- sentry_sdk/integrations/cohere/v1.py | 50 ++++++++++++-- sentry_sdk/integrations/cohere/v2.py | 55 ++++++++++++++-- 5 files changed, 145 insertions(+), 113 deletions(-) diff --git a/sentry_sdk/ai/span_config.py b/sentry_sdk/ai/span_config.py index becd6710d4..0520c87de4 100644 --- a/sentry_sdk/ai/span_config.py +++ b/sentry_sdk/ai/span_config.py @@ -33,27 +33,31 @@ def set_request_span_data(span, kwargs, integration, config, span_data=None): set_data_normalized(span, span_attr, value) if should_send_default_pii() and integration.include_prompts: - extract = config.get("extract_messages") - if extract is not None: - messages = extract(kwargs) - if messages: - messages = normalize_message_roles(messages) - scope = sentry_sdk.get_current_scope() - messages = truncate_and_annotate_messages(messages, span, scope) - if messages is not None: - target = config.get( - "message_target", SPANDATA.GEN_AI_REQUEST_MESSAGES - ) - set_data_normalized(span, target, messages, unpack=False) - for kwarg_key, span_attr in config.get("pii_params", {}).items(): if kwarg_key in kwargs: value = kwargs[kwarg_key] set_data_normalized(span, span_attr, value) +def set_request_messages(span, messages, target=None): + # type: (Span, Any, Optional[str]) -> None + """Normalize, truncate, and set request messages on the span. + + Caller is responsible for PII gating. + """ + if not messages: + return + messages = normalize_message_roles(messages) + scope = sentry_sdk.get_current_scope() + messages = truncate_and_annotate_messages(messages, span, scope) + if messages is not None: + set_data_normalized( + span, target or SPANDATA.GEN_AI_REQUEST_MESSAGES, messages, unpack=False + ) + + def set_response_span_data( - span, response, include_pii, response_config, collected_text=None + span, response, include_pii, response_config, response_text=None ): # type: (Span, Any, bool, Dict[str, Any], Optional[List[str]]) -> None """Set response span data from a declarative config.""" @@ -65,16 +69,8 @@ def set_response_span_data( pii_sources = response_config.get("pii_sources") if pii_sources: set_span_data_from_sources(span, response, pii_sources, require_truthy=True) - if collected_text: - set_data_normalized( - span, SPANDATA.GEN_AI_RESPONSE_TEXT, ["".join(collected_text)] - ) - else: - extract_text = response_config.get("extract_text") - if extract_text: - texts = extract_text(response) - if texts: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts) + if response_text: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_text) usage_config = response_config.get("usage") if usage_config: diff --git a/sentry_sdk/integrations/cohere/__init__.py b/sentry_sdk/integrations/cohere/__init__.py index 0f2989fefe..01d1f2e53c 100644 --- a/sentry_sdk/integrations/cohere/__init__.py +++ b/sentry_sdk/integrations/cohere/__init__.py @@ -1,12 +1,17 @@ import sys from functools import wraps -from sentry_sdk.consts import OP -from sentry_sdk.ai.span_config import set_request_span_data, set_response_span_data +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.ai.span_config import ( + set_request_span_data, + set_request_messages, + set_response_span_data, +) from sentry_sdk.integrations.cohere.configs import COHERE_EMBED_CONFIG from typing import TYPE_CHECKING +from sentry_sdk.scope import should_send_default_pii from sentry_sdk.tracing_utils import set_span_errored if TYPE_CHECKING: @@ -64,15 +69,23 @@ def new_embed(*args, **kwargs): model = kwargs.get("model", "") + include_pii = should_send_default_pii() and integration.include_prompts + with sentry_sdk.start_span( op=OP.GEN_AI_EMBEDDINGS, name=f"embeddings {model}".strip(), origin=CohereIntegration.origin, ) as span: set_request_span_data(span, kwargs, integration, COHERE_EMBED_CONFIG) + if include_pii and "texts" in kwargs: + set_request_messages( + span, + _normalize_embedding_input(kwargs["texts"]), + target=SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ) try: - res = f(*args, **kwargs) + response = f(*args, **kwargs) except Exception as e: exc_info = sys.exc_info() with capture_internal_exceptions(): @@ -80,8 +93,17 @@ def new_embed(*args, **kwargs): reraise(*exc_info) set_response_span_data( - span, res, False, COHERE_EMBED_CONFIG["response"] + span, response, False, COHERE_EMBED_CONFIG["response"] ) - return res + return response return new_embed + + +def _normalize_embedding_input(texts): + # type: (Any) -> Any + if isinstance(texts, list): + return texts + if isinstance(texts, tuple): + return list(texts) + return [texts] diff --git a/sentry_sdk/integrations/cohere/configs.py b/sentry_sdk/integrations/cohere/configs.py index 33d52952c3..67db83b7d3 100644 --- a/sentry_sdk/integrations/cohere/configs.py +++ b/sentry_sdk/integrations/cohere/configs.py @@ -1,13 +1,9 @@ -from sentry_sdk.ai.utils import ( - get_first_from_sources, - transform_message_content, -) from sentry_sdk.consts import SPANDATA from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + from typing import Any, Dict, Sequence, Tuple from typing_extensions import TypedDict # Source paths: list of attribute chains to try in order. @@ -32,9 +28,6 @@ class ResponseConfig(TypedDict, total=False): sources: SourceMapping # Attributes extracted only when PII sending is enabled. pii_sources: SourceMapping - # Custom extractor for response text (PII only). - # Returns list of text strings, or None. - extract_text: Callable[[Any], Optional[List[str]]] # Declarative token usage paths. usage: UsageConfig @@ -47,10 +40,6 @@ class OperationConfig(TypedDict, total=False): params: Dict[str, str] # Maps kwarg names to SPANDATA keys (only set when PII is enabled). pii_params: Dict[str, str] - # Extracts messages from kwargs for the span. - extract_messages: Callable[[Dict[str, Any]], Optional[List[Dict[str, Any]]]] - # SPANDATA key for messages (default: GEN_AI_REQUEST_MESSAGES). - message_target: str # Non-streaming response config. response: ResponseConfig # Streaming response config (different attribute paths). @@ -60,62 +49,6 @@ class OperationConfig(TypedDict, total=False): stream_response_object: SourcePaths -# ── Helpers ────────────────────────────────────────────────────────────────── - - -def _normalize_embedding_input(texts): - # type: (Any) -> Any - if isinstance(texts, list): - return texts - if isinstance(texts, tuple): - return list(texts) - return [texts] - - -def _extract_v1_messages(kwargs): - # type: (dict[str, Any]) -> list[dict[str, str]] - messages = [] - for x in kwargs.get("chat_history", []): - messages.append( - { - "role": getattr(x, "role", ""), - "content": transform_message_content(getattr(x, "message", "")), - } - ) - message = kwargs.get("message") - if message: - messages.append({"role": "user", "content": transform_message_content(message)}) - return messages - - -def _extract_v1_response_text(response): - # type: (Any) -> list[str] | None - text = getattr(response, "text", None) - return [text] if text is not None else None - - -def _extract_v2_messages(messages): - # type: (Any) -> list[dict[str, Any]] - result = [] - for msg in messages: - role = msg["role"] if isinstance(msg, dict) else getattr(msg, "role", "unknown") - content = ( - msg["content"] if isinstance(msg, dict) else getattr(msg, "content", "") - ) - result.append({"role": role, "content": transform_message_content(content)}) - return result - - -def _extract_v2_response_text(response): - # type: (Any) -> list[str] | None - content = get_first_from_sources(response, [("message", "content")], True) - if content: - texts = [item.text for item in content if hasattr(item, "text")] - if texts: - return texts - return None - - # ── Configs ────────────────────────────────────────────────────────────────── @@ -125,10 +58,6 @@ def _extract_v2_response_text(response): SPANDATA.GEN_AI_OPERATION_NAME: "embeddings", }, "params": {"model": SPANDATA.GEN_AI_REQUEST_MODEL}, - "extract_messages": lambda kw: ( - _normalize_embedding_input(kw["texts"]) if "texts" in kw else None - ), - "message_target": SPANDATA.GEN_AI_EMBEDDINGS_INPUT, "response": { "usage": { "input_tokens": [("meta", "billed_units", "input_tokens")], @@ -143,7 +72,6 @@ def _extract_v2_response_text(response): SPANDATA.GEN_AI_SYSTEM: "cohere", SPANDATA.GEN_AI_OPERATION_NAME: "chat", }, - "extract_messages": lambda kw: _extract_v1_messages(kw), "response": { "sources": { SPANDATA.GEN_AI_RESPONSE_MODEL: [("model",)], @@ -153,7 +81,6 @@ def _extract_v2_response_text(response): "pii_sources": { SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("tool_calls",)], }, - "extract_text": _extract_v1_response_text, "usage": { "input_tokens": [ ("meta", "billed_units", "input_tokens"), @@ -180,7 +107,6 @@ def _extract_v2_response_text(response): "pii_params": { "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, }, - "extract_messages": lambda kw: _extract_v2_messages(kw.get("messages", [])), "response": { "sources": { SPANDATA.GEN_AI_RESPONSE_MODEL: [("model",)], @@ -190,7 +116,6 @@ def _extract_v2_response_text(response): "pii_sources": { SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("message", "tool_calls")], }, - "extract_text": _extract_v2_response_text, "usage": { "input_tokens": [ ("usage", "billed_units", "input_tokens"), diff --git a/sentry_sdk/integrations/cohere/v1.py b/sentry_sdk/integrations/cohere/v1.py index 4e3c4741e9..193f863a33 100644 --- a/sentry_sdk/integrations/cohere/v1.py +++ b/sentry_sdk/integrations/cohere/v1.py @@ -1,8 +1,12 @@ import sys from functools import wraps -from sentry_sdk.ai.span_config import set_request_span_data, set_response_span_data -from sentry_sdk.ai.utils import get_first_from_sources +from sentry_sdk.ai.span_config import ( + set_request_span_data, + set_request_messages, + set_response_span_data, +) +from sentry_sdk.ai.utils import get_first_from_sources, transform_message_content from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations.cohere.configs import COHERE_V1_CHAT_CONFIG @@ -88,11 +92,20 @@ def new_chat(*args, **kwargs): set_request_span_data( span, kwargs, integration, COHERE_V1_CHAT_CONFIG, span_data ) + if include_pii: + set_request_messages(span, _extract_v1_messages(kwargs)) if streaming: return _iter_stream_events(response, span, include_pii) + response_text = ( + _extract_v1_response_text(response) if include_pii else None + ) set_response_span_data( - span, response, include_pii, COHERE_V1_CHAT_CONFIG["response"] + span, + response, + include_pii, + COHERE_V1_CHAT_CONFIG["response"], + response_text, ) return response @@ -110,7 +123,36 @@ def _iter_stream_events(old_iterator, span, include_pii): x, COHERE_V1_CHAT_CONFIG["stream_response_object"] ) if response is not None: + response_text = ( + _extract_v1_response_text(response) if include_pii else None + ) set_response_span_data( - span, response, include_pii, COHERE_V1_CHAT_CONFIG["response"] + span, + response, + include_pii, + COHERE_V1_CHAT_CONFIG["response"], + response_text, ) yield x + + +def _extract_v1_messages(kwargs): + # type: (Any) -> list[dict[str, str]] + messages = [] + for x in kwargs.get("chat_history", []): + messages.append( + { + "role": getattr(x, "role", ""), + "content": transform_message_content(getattr(x, "message", "")), + } + ) + message = kwargs.get("message") + if message: + messages.append({"role": "user", "content": transform_message_content(message)}) + return messages + + +def _extract_v1_response_text(response): + # type: (Any) -> list[str] | None + text = getattr(response, "text", None) + return [text] if text is not None else None diff --git a/sentry_sdk/integrations/cohere/v2.py b/sentry_sdk/integrations/cohere/v2.py index cf6677ff75..752551ac73 100644 --- a/sentry_sdk/integrations/cohere/v2.py +++ b/sentry_sdk/integrations/cohere/v2.py @@ -1,8 +1,16 @@ import sys from functools import wraps -from sentry_sdk.ai.span_config import set_request_span_data, set_response_span_data -from sentry_sdk.ai.utils import get_first_from_sources, transitive_getattr +from sentry_sdk.ai.span_config import ( + set_request_span_data, + set_request_messages, + set_response_span_data, +) +from sentry_sdk.ai.utils import ( + get_first_from_sources, + transform_message_content, + transitive_getattr, +) from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations.cohere.configs import ( COHERE_V2_CHAT_CONFIG, @@ -89,10 +97,22 @@ def new_chat(*args, **kwargs): set_request_span_data( span, kwargs, integration, COHERE_V2_CHAT_CONFIG, span_data ) + if include_pii: + set_request_messages( + span, _extract_v2_messages(kwargs.get("messages", [])) + ) + if streaming: return _iter_v2_stream_events(response, span, include_pii) + response_text = ( + _extract_v2_response_text(response) if include_pii else None + ) set_response_span_data( - span, response, include_pii, COHERE_V2_CHAT_CONFIG["response"] + span, + response, + include_pii, + COHERE_V2_CHAT_CONFIG["response"], + response_text, ) return response @@ -106,12 +126,17 @@ def _iter_v2_stream_events(old_iterator, span, include_pii): with capture_internal_exceptions(): _append_stream_delta_text(collected_text, x) if isinstance(x, MessageEndV2ChatStreamResponse): + response_text = ( + ["".join(collected_text)] + if include_pii and collected_text + else None + ) set_response_span_data( span, x, include_pii, COHERE_V2_CHAT_CONFIG["stream_response"], - collected_text, + response_text, ) yield x @@ -123,3 +148,25 @@ def _append_stream_delta_text(collected_text, event): content_text = get_first_from_sources(event, STREAM_DELTA_TEXT_SOURCES) if content_text is not None: collected_text.append(content_text) + + +def _extract_v2_messages(messages): + # type: (Any) -> list[dict[str, Any]] + result = [] + for msg in messages: + role = msg["role"] if isinstance(msg, dict) else getattr(msg, "role", "unknown") + content = ( + msg["content"] if isinstance(msg, dict) else getattr(msg, "content", "") + ) + result.append({"role": role, "content": transform_message_content(content)}) + return result + + +def _extract_v2_response_text(response): + # type: (Any) -> list[str] | None + content = get_first_from_sources(response, [("message", "content")], True) + if content: + texts = [item.text for item in content if hasattr(item, "text")] + if texts: + return texts + return None