integrate message truncation in langchain & langgraph

Author: shellmayr

sentry_sdk/integrations/langchain.py

@@ -10,6 +10,8 @@
     set_data_normalized,
     get_start_span_function,
 )
+from sentry_sdk.ai.utils import set_data_normalized, get_start_span_function
+from sentry_sdk.ai.message_utils import truncate_and_serialize_messages
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
@@ -221,12 +223,11 @@ def on_llm_start(
                     }
                     for prompt in prompts
                 ]
-                set_data_normalized(
-                    span,
-                    SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                    normalized_messages,
-                    unpack=False,
-                )
+                result = truncate_and_serialize_messages(prompts)
+                if result["serialized_data"]:
+                    span.set_data(
+                        SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
+                    )
 
     def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
         # type: (SentryLangchainCallback, Dict[str, Any], List[List[BaseMessage]], UUID, Any) -> Any
@@ -278,13 +279,11 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
                             self._normalize_langchain_message(message)
                         )
                 normalized_messages = normalize_message_roles(normalized_messages)
-
-                set_data_normalized(
-                    span,
-                    SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                    normalized_messages,
-                    unpack=False,
-                )
+                result = truncate_and_serialize_messages(normalized_messages)
+                if result["serialized_data"]:
+                    span.set_data(
+                        SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
+                    )
 
     def on_chat_model_end(self, response, *, run_id, **kwargs):
         # type: (SentryLangchainCallback, LLMResult, UUID, Any) -> Any
@@ -758,12 +757,11 @@ def new_invoke(self, *args, **kwargs):
                 and integration.include_prompts
             ):
                 normalized_messages = normalize_message_roles([input])
-                set_data_normalized(
-                    span,
-                    SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                    normalized_messages,
-                    unpack=False,
-                )
+                result = truncate_and_serialize_messages(normalized_messages)
+                if result["serialized_data"]:
+                    span.set_data(
+                        SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
+                    )
 
             output = result.get("output")
             if (
@@ -813,12 +811,11 @@ def new_stream(self, *args, **kwargs):
             and integration.include_prompts
         ):
             normalized_messages = normalize_message_roles([input])
-            set_data_normalized(
-                span,
-                SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                normalized_messages,
-                unpack=False,
-            )
+            result = truncate_and_serialize_messages(normalized_messages)
+            if result["serialized_data"]:
+                span.set_data(
+                    SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
+                )
 
         # Run the agent
         result = f(self, *args, **kwargs)

sentry_sdk/integrations/langgraph.py

@@ -3,6 +3,7 @@
 
 import sentry_sdk
 from sentry_sdk.ai.utils import set_data_normalized, normalize_message_roles
+from sentry_sdk.ai.message_utils import truncate_and_serialize_messages
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
@@ -181,12 +182,11 @@ def new_invoke(self, *args, **kwargs):
                 input_messages = _parse_langgraph_messages(args[0])
                 if input_messages:
                     normalized_input_messages = normalize_message_roles(input_messages)
-                    set_data_normalized(
-                        span,
-                        SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                        normalized_input_messages,
-                        unpack=False,
-                    )
+                    result = truncate_and_serialize_messages(normalized_input_messages)
+                    if result["serialized_data"]:
+                        span.set_data(
+                            SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
+                        )
 
             result = f(self, *args, **kwargs)
 
@@ -232,12 +232,11 @@ async def new_ainvoke(self, *args, **kwargs):
                 input_messages = _parse_langgraph_messages(args[0])
                 if input_messages:
                     normalized_input_messages = normalize_message_roles(input_messages)
-                    set_data_normalized(
-                        span,
-                        SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                        normalized_input_messages,
-                        unpack=False,
-                    )
+                    result = truncate_and_serialize_messages(normalized_input_messages)
+                    if result["serialized_data"]:
+                        span.set_data(
+                            SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
+                        )
 
             result = await f(self, *args, **kwargs)