fix role mapping

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

sentry_sdk/integrations/cohere.py

@@ -4,10 +4,7 @@
 from sentry_sdk import consts
 from sentry_sdk.ai.monitoring import record_token_usage
 from sentry_sdk.consts import OP, SPANDATA
-from sentry_sdk.ai.utils import (
-    set_data_normalized,
-    normalize_message_roles,
-)
+from sentry_sdk.ai.utils import set_data_normalized, normalize_message_roles
 
 from typing import TYPE_CHECKING
 
@@ -154,18 +151,15 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any":
             reraise(*exc_info)
 
         with capture_internal_exceptions():
-            set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "cohere")
-            set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+            span.set_data(SPANDATA.GEN_AI_SYSTEM, "cohere")
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
 
             if should_send_default_pii() and integration.include_prompts:
                 messages = []
                 for x in kwargs.get("chat_history", []):
-                    role = getattr(x, "role", "").lower()
-                    if role == "chatbot":
-                        role = "assistant"
                     messages.append(
                         {
-                            "role": role,
+                            "role": getattr(x, "role", ""),
                             "content": getattr(x, "message", ""),
                         }
                     )