wip

Author: shellmayr

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"],
 }
 

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])

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,