add to langchain

Author: shellmayr

sentry_sdk/ai/utils.py

@@ -19,7 +19,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"],
+    GEN_AI_ALLOWED_MESSAGE_ROLES.USER: ["user", "human"],
     GEN_AI_ALLOWED_MESSAGE_ROLES.ASSISTANT: ["assistant", "ai"],
     GEN_AI_ALLOWED_MESSAGE_ROLES.TOOL_CALL: ["tool", "tool_call"],
 }

sentry_sdk/integrations/langchain.py

@@ -4,7 +4,11 @@
 
 import sentry_sdk
 from sentry_sdk.ai.monitoring import set_ai_pipeline_name
-from sentry_sdk.ai.utils import set_data_normalized, get_start_span_function
+from sentry_sdk.ai.utils import (
+    normalize_message_roles,
+    set_data_normalized,
+    get_start_span_function,
+)
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
@@ -209,8 +213,12 @@ def on_llm_start(
             _set_tools_on_span(span, all_params.get("tools"))
 
             if should_send_default_pii() and self.include_prompts:
+                normalized_messages = normalize_message_roles(prompts)
                 set_data_normalized(
-                    span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompts, unpack=False
+                    span,
+                    SPANDATA.GEN_AI_REQUEST_MESSAGES,
+                    normalized_messages,
+                    unpack=False,
                 )
 
     def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
@@ -262,6 +270,8 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
                         normalized_messages.append(
                             self._normalize_langchain_message(message)
                         )
+                normalized_messages = normalize_message_roles(normalized_messages)
+
                 set_data_normalized(
                     span,
                     SPANDATA.GEN_AI_REQUEST_MESSAGES,
@@ -740,8 +750,12 @@ def new_invoke(self, *args, **kwargs):
                 and should_send_default_pii()
                 and integration.include_prompts
             ):
+                normalized_messages = normalize_message_roles([input])
                 set_data_normalized(
-                    span, SPANDATA.GEN_AI_REQUEST_MESSAGES, [input], unpack=False
+                    span,
+                    SPANDATA.GEN_AI_REQUEST_MESSAGES,
+                    normalized_messages,
+                    unpack=False,
                 )
 
             output = result.get("output")
@@ -791,8 +805,12 @@ def new_stream(self, *args, **kwargs):
             and should_send_default_pii()
             and integration.include_prompts
         ):
+            normalized_messages = normalize_message_roles([input])
             set_data_normalized(
-                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, [input], unpack=False
+                span,
+                SPANDATA.GEN_AI_REQUEST_MESSAGES,
+                normalized_messages,
+                unpack=False,
             )
 
         # Run the agent