do cohere specific mapping of roles

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", "chatbot"],
+    GEN_AI_ALLOWED_MESSAGE_ROLES.ASSISTANT: ["assistant", "ai"],
     GEN_AI_ALLOWED_MESSAGE_ROLES.TOOL: ["tool", "tool_call"],
 }
 
@@ -503,7 +503,7 @@ def normalize_message_role(role: str) -> str:
     Normalize a message role to one of the 4 allowed gen_ai role values.
     Maps "ai" -> "assistant" and keeps other standard roles unchanged.
     """
-    return GEN_AI_MESSAGE_ROLE_MAPPING.get(role.lower(), role)
+    return GEN_AI_MESSAGE_ROLE_MAPPING.get(role, role)
 
 
 def normalize_message_roles(messages: "list[dict[str, Any]]") -> "list[dict[str, Any]]":

sentry_sdk/integrations/cohere.py

@@ -4,7 +4,7 @@
 from sentry_sdk import consts
 from sentry_sdk.ai.monitoring import record_token_usage
 from sentry_sdk.consts import OP, SPANDATA
-from sentry_sdk.ai.utils import set_data_normalized, normalize_message_roles
+from sentry_sdk.ai.utils import set_data_normalized
 
 from typing import TYPE_CHECKING
 
@@ -39,6 +39,14 @@
     from cohere import StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse
 
 
+COHERE_ROLE_MAPPING = {
+    "SYSTEM": "system",
+    "USER": "user",
+    "CHATBOT": "assistant",
+    "TOOL": "tool",
+}
+
+
 COLLECTED_CHAT_PARAMS = {
     "model": SPANDATA.GEN_AI_REQUEST_MODEL,
     "k": SPANDATA.GEN_AI_REQUEST_TOP_K,
@@ -157,14 +165,14 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any":
             if should_send_default_pii() and integration.include_prompts:
                 messages = []
                 for x in kwargs.get("chat_history", []):
+                    role = getattr(x, "role", "")
                     messages.append(
                         {
-                            "role": getattr(x, "role", ""),
+                            "role": COHERE_ROLE_MAPPING.get(role, role),
                             "content": getattr(x, "message", ""),
                         }
                     )
                 messages.append({"role": "user", "content": message})
-                messages = normalize_message_roles(messages)
                 set_data_normalized(
                     span,
                     SPANDATA.GEN_AI_REQUEST_MESSAGES,