Skip to content
31 changes: 10 additions & 21 deletions sentry_sdk/integrations/langgraph.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
from sentry_sdk.ai.utils import (
set_data_normalized,
normalize_message_roles,
truncate_and_annotate_messages,
)
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations import DidNotEnable, Integration
Expand Down Expand Up @@ -181,17 +180,12 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
input_messages = _parse_langgraph_messages(args[0])
if input_messages:
normalized_input_messages = normalize_message_roles(input_messages)
scope = sentry_sdk.get_current_scope()
messages_data = truncate_and_annotate_messages(
normalized_input_messages, span, scope
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
normalized_input_messages,
unpack=False,
)
if messages_data is not None:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
messages_data,
unpack=False,
)

result = f(self, *args, **kwargs)

Expand Down Expand Up @@ -234,17 +228,12 @@ async def new_ainvoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
input_messages = _parse_langgraph_messages(args[0])
if input_messages:
normalized_input_messages = normalize_message_roles(input_messages)
scope = sentry_sdk.get_current_scope()
messages_data = truncate_and_annotate_messages(
normalized_input_messages, span, scope
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
normalized_input_messages,
unpack=False,
)
if messages_data is not None:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
messages_data,
unpack=False,
)

result = await f(self, *args, **kwargs)

Expand Down
63 changes: 3 additions & 60 deletions tests/integrations/langgraph/test_langgraph.py
Original file line number Diff line number Diff line change
Expand Up @@ -272,8 +272,9 @@ def original_invoke(self, *args, **kwargs):
import json

request_messages = json.loads(request_messages)
assert len(request_messages) == 1
assert request_messages[0]["content"] == "Of course! How can I assist you?"
assert len(request_messages) == 2
assert request_messages[0]["content"] == "Hello, can you help me?"
assert request_messages[1]["content"] == "Of course! How can I assist you?"

response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT]
assert response_text == expected_assistant_response
Expand Down Expand Up @@ -1351,61 +1352,3 @@ def __init__(self, content, message_type="human"):
# Verify no "ai" roles remain
roles = [msg["role"] for msg in stored_messages if "role" in msg]
assert "ai" not in roles


def test_langgraph_message_truncation(sentry_init, capture_items):
"""Test that large messages are truncated properly in Langgraph integration."""
import json

sentry_init(
integrations=[LanggraphIntegration(include_prompts=True)],
traces_sample_rate=1.0,
send_default_pii=True,
)
items = capture_items("transaction", "span")

large_content = (
"This is a very long message that will exceed our size limits. " * 1000
)
test_state = {
"messages": [
MockMessage("small message 1", name="user"),
MockMessage(large_content, name="assistant"),
MockMessage(large_content, name="user"),
MockMessage("small message 4", name="assistant"),
MockMessage("small message 5", name="user"),
]
}

pregel = MockPregelInstance("test_graph")

def original_invoke(self, *args, **kwargs):
return {"messages": args[0].get("messages", [])}

with start_transaction():
wrapped_invoke = _wrap_pregel_invoke(original_invoke)
result = wrapped_invoke(pregel, test_state)

assert result is not None

spans = [item.payload for item in items if item.type == "span"]
invoke_spans = [
span
for span in spans
if span["attributes"].get("sentry.op") == OP.GEN_AI_INVOKE_AGENT
]
assert len(invoke_spans) > 0

invoke_span = invoke_spans[0]
assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"]

messages_data = invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]
assert isinstance(messages_data, str)

parsed_messages = json.loads(messages_data)
assert isinstance(parsed_messages, list)
assert len(parsed_messages) == 1
assert "small message 5" in str(parsed_messages[0])

(tx,) = (item.payload for item in items if item.type == "transaction")
assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5
Loading