diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py index 3d3856a913..09152959dd 100644 --- a/sentry_sdk/integrations/langgraph.py +++ b/sentry_sdk/integrations/langgraph.py @@ -19,7 +19,7 @@ has_span_streaming_enabled, should_truncate_gen_ai_input, ) -from sentry_sdk.utils import safe_serialize +from sentry_sdk.utils import has_data_collection_enabled, safe_serialize try: from langgraph.errors import GraphBubbleUp @@ -53,6 +53,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): @@ -153,7 +171,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 @@ -188,18 +212,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( @@ -235,18 +254,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( @@ -299,18 +313,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( @@ -345,16 +354,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( @@ -494,22 +498,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, + ) diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py index 3785e2be9e..f16bd620a2 100644 --- a/tests/integrations/langgraph/test_langgraph.py +++ b/tests/integrations/langgraph/test_langgraph.py @@ -2132,3 +2132,539 @@ def test_graph_bubble_up_ignored(sentry_init, capture_items): model.invoke([HumanMessage(content="hi")]) assert len(events) == 0 + + +def _invoke_span_data(items_or_events, span_streaming): + if span_streaming: + sentry_sdk.flush() + spans = [item.payload for item in items_or_events] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + return invoke_spans[0]["attributes"] + + tx = items_or_events[0] + invoke_spans = [ + span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + return invoke_spans[0]["data"] + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize( + "data_collection, send_default_pii, expect_inputs", + [ + pytest.param( + {"gen_ai": {"inputs": True}}, + False, + True, + id="gen-ai-inputs-enabled-overrides-pii-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False}}, + True, + False, + id="gen-ai-inputs-disabled-overrides-pii-enabled", + ), + pytest.param( + {}, + False, + True, + id="gen-ai-omitted-defaults-to-enabled", + ), + pytest.param( + {"gen_ai": {"outputs": False}}, + False, + True, + id="gen-ai-outputs-disabled-does-not-affect-inputs", + ), + pytest.param( + None, + False, + False, + id="no-data-collection-falls-back-to-send-default-pii", + ), + pytest.param( + None, + True, + True, + id="no-data-collection-pii-enabled-collects", + ), + ], +) +def test_pregel_invoke_gates_request_messages_and_tool_calls_on_inputs_setting( + sentry_init, + capture_events, + capture_items, + data_collection, + send_default_pii, + expect_inputs, + span_streaming, +): + init_kwargs = { + "integrations": [LanggraphIntegration()], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": span_streaming, + "trace_lifecycle": "stream" if span_streaming else "static", + } + if data_collection is not None: + init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**init_kwargs) + + test_state = {"messages": [MockMessage("Hello, can you help me?", name="user")]} + pregel = MockPregelInstance("test_graph") + expected_tool_calls = [ + { + "id": "call_test_123", + "type": "function", + "function": {"name": "search_tool", "arguments": '{"query": "help"}'}, + } + ] + + def original_invoke(self, *args, **kwargs): + return { + "messages": args[0].get("messages", []) + + [ + MockMessage( + content="I'll help you with that task!", + name="assistant", + tool_calls=expected_tool_calls, + ) + ] + } + + captured = capture_items("span") if span_streaming else capture_events() + + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, test_state) + + data = _invoke_span_data(captured, span_streaming) + + if expect_inputs: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in data + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in data + else: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in data + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize( + "data_collection, send_default_pii, expect_outputs", + [ + pytest.param( + {"gen_ai": {"outputs": True}}, + False, + True, + id="gen-ai-outputs-enabled-overrides-pii-disabled", + ), + pytest.param( + {"gen_ai": {"outputs": False}}, + True, + False, + id="gen-ai-outputs-disabled-overrides-pii-enabled", + ), + pytest.param( + {}, + False, + True, + id="gen-ai-omitted-defaults-to-enabled", + ), + pytest.param( + {"gen_ai": {"inputs": False}}, + False, + True, + id="gen-ai-inputs-disabled-does-not-affect-outputs", + ), + pytest.param( + None, + False, + False, + id="no-data-collection-falls-back-to-send-default-pii", + ), + pytest.param( + None, + True, + True, + id="no-data-collection-pii-enabled-collects", + ), + ], +) +def test_pregel_invoke_gates_response_text_on_outputs_setting( + sentry_init, + capture_events, + capture_items, + data_collection, + send_default_pii, + expect_outputs, + span_streaming, +): + init_kwargs = { + "integrations": [LanggraphIntegration()], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": span_streaming, + "trace_lifecycle": "stream" if span_streaming else "static", + } + if data_collection is not None: + init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**init_kwargs) + + test_state = {"messages": [MockMessage("Hello, can you help me?", name="user")]} + pregel = MockPregelInstance("test_graph") + expected_assistant_response = "I'll help you with that task!" + + def original_invoke(self, *args, **kwargs): + return { + "messages": args[0].get("messages", []) + + [MockMessage(content=expected_assistant_response, name="assistant")] + } + + captured = capture_items("span") if span_streaming else capture_events() + + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, test_state) + + data = _invoke_span_data(captured, span_streaming) + + if expect_outputs: + assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response + else: + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in data + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize( + "data_collection, send_default_pii, expect_inputs, expect_outputs", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + False, + True, + False, + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + True, + False, + True, + id="gen-ai-inputs-disabled-outputs-enabled", + ), + pytest.param( + None, + False, + False, + False, + id="no-data-collection-falls-back-to-send-default-pii", + ), + pytest.param( + None, + True, + True, + True, + id="no-data-collection-pii-enabled-collects", + ), + ], +) +def test_pregel_ainvoke_gates_inputs_and_outputs_independently( + sentry_init, + capture_events, + capture_items, + data_collection, + send_default_pii, + expect_inputs, + expect_outputs, + span_streaming, +): + init_kwargs = { + "integrations": [LanggraphIntegration()], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": span_streaming, + "trace_lifecycle": "stream" if span_streaming else "static", + } + if data_collection is not None: + init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**init_kwargs) + + test_state = {"messages": [MockMessage("What is the weather?", name="user")]} + pregel = MockPregelInstance("async_graph") + expected_assistant_response = "Let me check the weather for you!" + expected_tool_calls = [ + { + "id": "call_weather_456", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"location": "Berlin"}'}, + } + ] + + async def original_ainvoke(self, *args, **kwargs): + return { + "messages": args[0].get("messages", []) + + [ + MockMessage( + content=expected_assistant_response, + name="assistant", + tool_calls=expected_tool_calls, + ) + ] + } + + async def run_test(): + with start_transaction(): + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + return await wrapped_ainvoke(pregel, test_state) + + captured = capture_items("span") if span_streaming else capture_events() + + asyncio.run(run_test()) + + data = _invoke_span_data(captured, span_streaming) + + if expect_inputs: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in data + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in data + else: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in data + + if expect_outputs: + assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response + else: + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in data + + +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_pregel_invoke_message_delta_ignores_gen_ai_inputs_setting( + sentry_init, + capture_events, + capture_items, + span_streaming, +): + sentry_init( + integrations=[LanggraphIntegration()], + traces_sample_rate=1.0, + stream_gen_ai_spans=span_streaming, + trace_lifecycle="stream" if span_streaming else "static", + _experiments={ + "data_collection": {"gen_ai": {"inputs": False, "outputs": True}} + }, + ) + + prior_response = "Of course! How can I assist you?" + test_state = { + "messages": [ + MockMessage("Hello, can you help me?", name="user"), + MockMessage( + prior_response, + name="assistant", + response_metadata={ + "token_usage": { + "total_tokens": 300, + "prompt_tokens": 100, + "completion_tokens": 200, + }, + "model_name": "gpt-3.5-turbo", + }, + ), + ] + } + pregel = MockPregelInstance("test_graph") + expected_assistant_response = "I'll help you with that task!" + + def original_invoke(self, *args, **kwargs): + return { + "messages": args[0].get("messages", []) + + [ + MockMessage( + content=expected_assistant_response, + name="assistant", + response_metadata={ + "token_usage": { + "total_tokens": 30, + "prompt_tokens": 10, + "completion_tokens": 20, + }, + "model_name": "gpt-4.1-2025-04-14", + }, + ) + ] + } + + captured = capture_items("span") if span_streaming else capture_events() + + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, test_state) + + data = _invoke_span_data(captured, span_streaming) + + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data + assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response + assert prior_response not in data[SPANDATA.GEN_AI_RESPONSE_TEXT] + assert data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + assert data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4.1-2025-04-14" + + +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_pregel_ainvoke_message_delta_ignores_gen_ai_inputs_setting( + sentry_init, + capture_events, + capture_items, + span_streaming, +): + sentry_init( + integrations=[LanggraphIntegration()], + traces_sample_rate=1.0, + stream_gen_ai_spans=span_streaming, + trace_lifecycle="stream" if span_streaming else "static", + _experiments={ + "data_collection": {"gen_ai": {"inputs": False, "outputs": True}} + }, + ) + + prior_response = "It is sunny in Berlin." + test_state = { + "messages": [ + MockMessage("What is the weather?", name="user"), + MockMessage( + prior_response, + name="assistant", + response_metadata={ + "token_usage": { + "total_tokens": 300, + "prompt_tokens": 100, + "completion_tokens": 200, + }, + "model_name": "gpt-3.5-turbo", + }, + ), + ] + } + pregel = MockPregelInstance("async_graph") + expected_assistant_response = "Let me check the weather for you!" + + async def original_ainvoke(self, *args, **kwargs): + return { + "messages": args[0].get("messages", []) + + [ + MockMessage( + content=expected_assistant_response, + name="assistant", + response_metadata={ + "token_usage": { + "total_tokens": 30, + "prompt_tokens": 10, + "completion_tokens": 20, + }, + "model_name": "gpt-4.1-2025-04-14", + }, + ) + ] + } + + async def run_test(): + with start_transaction(): + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + return await wrapped_ainvoke(pregel, test_state) + + captured = capture_items("span") if span_streaming else capture_events() + + asyncio.run(run_test()) + + data = _invoke_span_data(captured, span_streaming) + + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data + assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response + assert prior_response not in data[SPANDATA.GEN_AI_RESPONSE_TEXT] + assert data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + assert data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4.1-2025-04-14" + + +@pytest.mark.parametrize( + "data_collection, send_default_pii, expect_available_tools", + [ + pytest.param( + {"gen_ai": {"inputs": True}}, + False, + True, + id="gen-ai-inputs-enabled-overrides-pii-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False}}, + True, + False, + id="gen-ai-inputs-disabled-overrides-pii-enabled", + ), + pytest.param( + {}, + False, + True, + id="gen-ai-omitted-defaults-to-enabled", + ), + pytest.param( + None, + False, + True, + id="no-data-collection-collects-regardless-of-pii", + ), + ], +) +def test_state_graph_compile_gates_available_tools_only_when_data_collection_configured( + sentry_init, + capture_events, + data_collection, + send_default_pii, + expect_available_tools, +): + init_kwargs = { + "integrations": [LanggraphIntegration()], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": False, + } + if data_collection is not None: + init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**init_kwargs) + + graph = MockStateGraph() + + def original_compile(self, *args, **kwargs): + return MockCompiledGraph(self.name) + + events = capture_events() + + with patch("sentry_sdk.integrations.langgraph.StateGraph"), start_transaction(): + wrapped_compile = _wrap_state_graph_compile(original_compile) + wrapped_compile(graph, model="test-model", checkpointer=None) + + tx = events[0] + agent_spans = [span for span in tx["spans"] if span["op"] == OP.GEN_AI_CREATE_AGENT] + assert len(agent_spans) == 1 + data = agent_spans[0]["data"] + + if expect_available_tools: + assert data[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS] == [ + "search_tool", + "calculator", + ] + else: + assert SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS not in data + + assert data[SPANDATA.GEN_AI_AGENT_NAME] == "test_graph"