feat(langgraph): Gate prompt/response collection on data_collection option (#7175)

Modify the LangGraph integration to respect the data_collection config
for controlling whether prompts, responses, tool calls, and available
tools are captured in spans.

When data collection is enabled, the gen_ai.inputs flag controls request
messages, tool calls, and available tools, while gen_ai.outputs controls
the response text. Tool calls are gated on inputs because they are fed
back to the model as input. Available tools are only gated once data
collection is configured, since they were never gated on the legacy PII
settings.

When data collection is not configured, falls back to legacy
send_default_pii and include_prompts settings for compatibility.

Refs #6748  <br>Refs #6748

Author: ericapisani

sentry_sdk/integrations/langgraph.py

@@ -19,7 +19,11 @@
     has_span_streaming_enabled,
     should_truncate_gen_ai_input,
 )
-from sentry_sdk.utils import package_version, safe_serialize
+from sentry_sdk.utils import (
+    has_data_collection_enabled,
+    package_version,
+    safe_serialize,
+)
 
 try:
     from langgraph.errors import GraphBubbleUp
@@ -56,6 +60,24 @@ def setup_once() -> None:
             Pregel.ainvoke = _wrap_pregel_ainvoke(Pregel.ainvoke)
 
 
+def _should_record_inputs(integration: "LanggraphIntegration") -> bool:
+    client = sentry_sdk.get_client()
+    if has_data_collection_enabled(client.options):
+        return bool(client.options["data_collection"]["gen_ai"]["inputs"])
+
+    # To remove once data collection has been fully rolled out
+    return should_send_default_pii() and integration.include_prompts
+
+
+def _should_record_outputs(integration: "LanggraphIntegration") -> bool:
+    client = sentry_sdk.get_client()
+    if has_data_collection_enabled(client.options):
+        return bool(client.options["data_collection"]["gen_ai"]["outputs"])
+
+    # To remove once data collection has been fully rolled out
+    return should_send_default_pii() and integration.include_prompts
+
+
 def _get_graph_name(graph_obj: "Any") -> "Optional[str]":
     for attr in ["name", "graph_name", "__name__", "_name"]:
         if hasattr(graph_obj, attr):
@@ -156,7 +178,13 @@ def new_compile(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
                             tools = list(data.tools_by_name.keys())
 
             if tools is not None:
-                span.set_data(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools)
+                # Available tools aren't gated on the legacy PII settings, so they're
+                # only gated when data collection has been configured.
+                if has_data_collection_enabled(client.options):
+                    if client.options["data_collection"]["gen_ai"]["inputs"]:
+                        span.set_data(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools)
+                else:
+                    span.set_data(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools)
 
             return compiled_graph
 
@@ -191,18 +219,13 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
 
                 # Store input messages to later compare with output
                 input_messages = None
-                if (
-                    len(args) > 0
-                    and should_send_default_pii()
-                    and integration.include_prompts
-                ):
+                if len(args) > 0:
                     input_messages = _parse_langgraph_messages(args[0])
-                    if input_messages:
+                    if input_messages and _should_record_inputs(integration):
                         normalized_input_messages = normalize_message_roles(
                             input_messages
                         )
 
-                        client = sentry_sdk.get_client()
                         scope = sentry_sdk.get_current_scope()
                         messages_data = (
                             truncate_and_annotate_messages(
@@ -238,18 +261,13 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
 
                 # Store input messages to later compare with output
                 input_messages = None
-                if (
-                    len(args) > 0
-                    and should_send_default_pii()
-                    and integration.include_prompts
-                ):
+                if len(args) > 0:
                     input_messages = _parse_langgraph_messages(args[0])
-                    if input_messages:
+                    if input_messages and _should_record_inputs(integration):
                         normalized_input_messages = normalize_message_roles(
                             input_messages
                         )
 
-                        client = sentry_sdk.get_client()
                         scope = sentry_sdk.get_current_scope()
                         messages_data = (
                             truncate_and_annotate_messages(
@@ -302,18 +320,13 @@ async def new_ainvoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
                     span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, graph_name)
 
                 input_messages = None
-                if (
-                    len(args) > 0
-                    and should_send_default_pii()
-                    and integration.include_prompts
-                ):
+                if len(args) > 0:
                     input_messages = _parse_langgraph_messages(args[0])
-                    if input_messages:
+                    if input_messages and _should_record_inputs(integration):
                         normalized_input_messages = normalize_message_roles(
                             input_messages
                         )
 
-                        client = sentry_sdk.get_client()
                         scope = sentry_sdk.get_current_scope()
                         messages_data = (
                             truncate_and_annotate_messages(
@@ -348,16 +361,11 @@ async def new_ainvoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
             span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
 
             input_messages = None
-            if (
-                len(args) > 0
-                and should_send_default_pii()
-                and integration.include_prompts
-            ):
+            if len(args) > 0:
                 input_messages = _parse_langgraph_messages(args[0])
-                if input_messages:
+                if input_messages and _should_record_inputs(integration):
                     normalized_input_messages = normalize_message_roles(input_messages)
 
-                    client = sentry_sdk.get_client()
                     scope = sentry_sdk.get_current_scope()
                     messages_data = (
                         truncate_and_annotate_messages(
@@ -497,22 +505,22 @@ def _set_response_attributes(
     _set_usage_data(span, new_messages)
     _set_response_model_name(span, new_messages)
 
-    if not (should_send_default_pii() and integration.include_prompts):
-        return
-
-    llm_response_text = _extract_llm_response_text(new_messages)
-    if llm_response_text:
-        set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, llm_response_text)
-    elif new_messages:
-        set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, new_messages)
-    else:
-        set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, result)
-
-    tool_calls = _extract_tool_calls(new_messages)
-    if tool_calls:
-        set_data_normalized(
-            span,
-            SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
-            safe_serialize(tool_calls),
-            unpack=False,
-        )
+    if _should_record_outputs(integration):
+        llm_response_text = _extract_llm_response_text(new_messages)
+        if llm_response_text:
+            set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, llm_response_text)
+        elif new_messages:
+            set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, new_messages)
+        else:
+            set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, result)
+
+    # Tool calls are an input to the model, so they're gated on inputs
+    if _should_record_inputs(integration):
+        tool_calls = _extract_tool_calls(new_messages)
+        if tool_calls:
+            set_data_normalized(
+                span,
+                SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
+                safe_serialize(tool_calls),
+                unpack=False,
+            )