integrate message truncation in langchain & langgraph

Author: shellmayr

sentry_sdk/integrations/langchain.py

@@ -5,6 +5,7 @@
 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.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
@@ -209,9 +210,11 @@ def on_llm_start(
             _set_tools_on_span(span, all_params.get("tools"))
 
             if should_send_default_pii() and self.include_prompts:
-                set_data_normalized(
-                    span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompts, 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
@@ -262,12 +265,11 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
                         normalized_messages.append(
                             self._normalize_langchain_message(message)
                         )
-                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
@@ -740,9 +742,11 @@ def new_invoke(self, *args, **kwargs):
                 and should_send_default_pii()
                 and integration.include_prompts
             ):
-                set_data_normalized(
-                    span, SPANDATA.GEN_AI_REQUEST_MESSAGES, [input], unpack=False
-                )
+                result = truncate_and_serialize_messages([input])
+                if result["serialized_data"]:
+                    span.set_data(
+                        SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
+                    )
 
             output = result.get("output")
             if (
@@ -791,9 +795,11 @@ def new_stream(self, *args, **kwargs):
             and should_send_default_pii()
             and integration.include_prompts
         ):
-            set_data_normalized(
-                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, [input], unpack=False
-            )
+            result = truncate_and_serialize_messages([input])
+            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
+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
@@ -180,12 +181,11 @@ def new_invoke(self, *args, **kwargs):
             ):
                 input_messages = _parse_langgraph_messages(args[0])
                 if input_messages:
-                    set_data_normalized(
-                        span,
-                        SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                        input_messages,
-                        unpack=False,
-                    )
+                    result = truncate_and_serialize_messages(input_messages)
+                    if result["serialized_data"]:
+                        span.set_data(
+                            SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
+                        )
 
             result = f(self, *args, **kwargs)
 
@@ -230,12 +230,11 @@ async def new_ainvoke(self, *args, **kwargs):
             ):
                 input_messages = _parse_langgraph_messages(args[0])
                 if input_messages:
-                    set_data_normalized(
-                        span,
-                        SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                        input_messages,
-                        unpack=False,
-                    )
+                    result = truncate_and_serialize_messages(input_messages)
+                    if result["serialized_data"]:
+                        span.set_data(
+                            SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
+                        )
 
             result = await f(self, *args, **kwargs)