diff --git a/.github/workflows/test-integrations-ai.yml b/.github/workflows/test-integrations-ai.yml index 72a4253744..26a8bdb8bb 100644 --- a/.github/workflows/test-integrations-ai.yml +++ b/.github/workflows/test-integrations-ai.yml @@ -74,6 +74,10 @@ jobs: run: | set -x # print commands that are executed ./scripts/runtox.sh --exclude-latest "py${{ matrix.python-version }}-openai-notiktoken" + - name: Test langgraph pinned + run: | + set -x # print commands that are executed + ./scripts/runtox.sh --exclude-latest "py${{ matrix.python-version }}-langgraph" - name: Test openai_agents pinned run: | set -x # print commands that are executed diff --git a/pyproject.toml b/pyproject.toml index deba247e39..44eded7641 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -130,6 +130,10 @@ ignore_missing_imports = true module = "langchain.*" ignore_missing_imports = true +[[tool.mypy.overrides]] +module = "langgraph.*" +ignore_missing_imports = true + [[tool.mypy.overrides]] module = "executing.*" ignore_missing_imports = true diff --git a/scripts/populate_tox/config.py b/scripts/populate_tox/config.py index 689253e889..6795e36303 100644 --- a/scripts/populate_tox/config.py +++ b/scripts/populate_tox/config.py @@ -157,6 +157,9 @@ }, "include": "<1.0", }, + "langgraph": { + "package": "langgraph", + }, "launchdarkly": { "package": "launchdarkly-server-sdk", }, diff --git a/scripts/populate_tox/tox.jinja b/scripts/populate_tox/tox.jinja index 65a5ba3f36..241e0ca288 100755 --- a/scripts/populate_tox/tox.jinja +++ b/scripts/populate_tox/tox.jinja @@ -338,6 +338,7 @@ setenv = huggingface_hub: TESTPATH=tests/integrations/huggingface_hub langchain-base: TESTPATH=tests/integrations/langchain langchain-notiktoken: TESTPATH=tests/integrations/langchain + langgraph: TESTPATH=tests/integrations/langgraph launchdarkly: TESTPATH=tests/integrations/launchdarkly litestar: TESTPATH=tests/integrations/litestar loguru: TESTPATH=tests/integrations/loguru diff --git a/scripts/split_tox_gh_actions/split_tox_gh_actions.py b/scripts/split_tox_gh_actions/split_tox_gh_actions.py index cf83e0a3fe..51ee614d04 100755 --- a/scripts/split_tox_gh_actions/split_tox_gh_actions.py +++ b/scripts/split_tox_gh_actions/split_tox_gh_actions.py @@ -78,6 +78,7 @@ "langchain-notiktoken", "openai-base", "openai-notiktoken", + "langgraph", "openai_agents", "huggingface_hub", ], diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py index 3ed8efd506..5480ef5dce 100644 --- a/sentry_sdk/consts.py +++ b/sentry_sdk/consts.py @@ -792,6 +792,7 @@ class OP: FUNCTION_AWS = "function.aws" FUNCTION_GCP = "function.gcp" GEN_AI_CHAT = "gen_ai.chat" + GEN_AI_CREATE_AGENT = "gen_ai.create_agent" GEN_AI_EMBEDDINGS = "gen_ai.embeddings" GEN_AI_EXECUTE_TOOL = "gen_ai.execute_tool" GEN_AI_HANDOFF = "gen_ai.handoff" diff --git a/sentry_sdk/integrations/__init__.py b/sentry_sdk/integrations/__init__.py index 6f0109aced..7f202221a7 100644 --- a/sentry_sdk/integrations/__init__.py +++ b/sentry_sdk/integrations/__init__.py @@ -95,6 +95,7 @@ def iter_default_integrations(with_auto_enabling_integrations): "sentry_sdk.integrations.huey.HueyIntegration", "sentry_sdk.integrations.huggingface_hub.HuggingfaceHubIntegration", "sentry_sdk.integrations.langchain.LangchainIntegration", + "sentry_sdk.integrations.langgraph.LanggraphIntegration", "sentry_sdk.integrations.litestar.LitestarIntegration", "sentry_sdk.integrations.loguru.LoguruIntegration", "sentry_sdk.integrations.openai.OpenAIIntegration", @@ -142,6 +143,7 @@ def iter_default_integrations(with_auto_enabling_integrations): "grpc": (1, 32, 0), # grpcio "huggingface_hub": (0, 22), "langchain": (0, 1, 0), + "langgraph": (0, 6, 6), "launchdarkly": (9, 8, 0), "loguru": (0, 7, 0), "openai": (1, 0, 0), diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py new file mode 100644 index 0000000000..4b241fe895 --- /dev/null +++ b/sentry_sdk/integrations/langgraph.py @@ -0,0 +1,321 @@ +from functools import wraps +from typing import Any, Callable, List, Optional + +import sentry_sdk +from sentry_sdk.ai.utils import set_data_normalized +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import safe_serialize + + +try: + from langgraph.graph import StateGraph + from langgraph.pregel import Pregel +except ImportError: + raise DidNotEnable("langgraph not installed") + + +class LanggraphIntegration(Integration): + identifier = "langgraph" + origin = f"auto.ai.{identifier}" + + def __init__(self, include_prompts=True): + # type: (LanggraphIntegration, bool) -> None + self.include_prompts = include_prompts + + @staticmethod + def setup_once(): + # type: () -> None + # LangGraph lets users create agents using a StateGraph or the Functional API. + # StateGraphs are then compiled to a CompiledStateGraph. Both CompiledStateGraph and + # the functional API execute on a Pregel instance. Pregel is the runtime for the graph + # and the invocation happens on Pregel, so patching the invoke methods takes care of both. + # The streaming methods are not patched, because due to some internal reasons, LangGraph + # will automatically patch the streaming methods to run through invoke, and by doing this + # we prevent duplicate spans for invocations. + StateGraph.compile = _wrap_state_graph_compile(StateGraph.compile) + if hasattr(Pregel, "invoke"): + Pregel.invoke = _wrap_pregel_invoke(Pregel.invoke) + if hasattr(Pregel, "ainvoke"): + Pregel.ainvoke = _wrap_pregel_ainvoke(Pregel.ainvoke) + + +def _get_graph_name(graph_obj): + # type: (Any) -> Optional[str] + for attr in ["name", "graph_name", "__name__", "_name"]: + if hasattr(graph_obj, attr): + name = getattr(graph_obj, attr) + if name and isinstance(name, str): + return name + return None + + +def _normalize_langgraph_message(message): + # type: (Any) -> Any + if not hasattr(message, "content"): + return None + + parsed = {"role": getattr(message, "type", None), "content": message.content} + + for attr in ["name", "tool_calls", "function_call", "tool_call_id"]: + if hasattr(message, attr): + value = getattr(message, attr) + if value is not None: + parsed[attr] = value + + return parsed + + +def _parse_langgraph_messages(state): + # type: (Any) -> Optional[List[Any]] + if not state: + return None + + messages = None + + if isinstance(state, dict): + messages = state.get("messages") + elif hasattr(state, "messages"): + messages = state.messages + elif hasattr(state, "get") and callable(state.get): + try: + messages = state.get("messages") + except Exception: + pass + + if not messages or not isinstance(messages, (list, tuple)): + return None + + normalized_messages = [] + for message in messages: + try: + normalized = _normalize_langgraph_message(message) + if normalized: + normalized_messages.append(normalized) + except Exception: + continue + + return normalized_messages if normalized_messages else None + + +def _wrap_state_graph_compile(f): + # type: (Callable[..., Any]) -> Callable[..., Any] + @wraps(f) + def new_compile(self, *args, **kwargs): + # type: (Any, Any, Any) -> Any + integration = sentry_sdk.get_client().get_integration(LanggraphIntegration) + if integration is None: + return f(self, *args, **kwargs) + with sentry_sdk.start_span( + op=OP.GEN_AI_CREATE_AGENT, + origin=LanggraphIntegration.origin, + ) as span: + compiled_graph = f(self, *args, **kwargs) + + compiled_graph_name = getattr(compiled_graph, "name", None) + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "create_agent") + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, compiled_graph_name) + + if compiled_graph_name: + span.description = f"create_agent {compiled_graph_name}" + else: + span.description = "create_agent" + + if kwargs.get("model", None) is not None: + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, kwargs.get("model")) + + tools = None + get_graph = getattr(compiled_graph, "get_graph", None) + if get_graph and callable(get_graph): + graph_obj = compiled_graph.get_graph() + nodes = getattr(graph_obj, "nodes", None) + if nodes and isinstance(nodes, dict): + tools_node = nodes.get("tools") + if tools_node: + data = getattr(tools_node, "data", None) + if data and hasattr(data, "tools_by_name"): + tools = list(data.tools_by_name.keys()) + + if tools is not None: + span.set_data(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools) + + return compiled_graph + + return new_compile + + +def _wrap_pregel_invoke(f): + # type: (Callable[..., Any]) -> Callable[..., Any] + + @wraps(f) + def new_invoke(self, *args, **kwargs): + # type: (Any, Any, Any) -> Any + integration = sentry_sdk.get_client().get_integration(LanggraphIntegration) + if integration is None: + return f(self, *args, **kwargs) + + graph_name = _get_graph_name(self) + span_name = ( + f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent" + ) + + with sentry_sdk.start_span( + op=OP.GEN_AI_INVOKE_AGENT, + name=span_name, + origin=LanggraphIntegration.origin, + ) as span: + if graph_name: + span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, graph_name) + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + + # Store input messages to later compare with output + input_messages = None + if ( + len(args) > 0 + and should_send_default_pii() + and integration.include_prompts + ): + input_messages = _parse_langgraph_messages(args[0]) + if input_messages: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + safe_serialize(input_messages), + ) + + result = f(self, *args, **kwargs) + + _set_response_attributes(span, input_messages, result, integration) + + return result + + return new_invoke + + +def _wrap_pregel_ainvoke(f): + # type: (Callable[..., Any]) -> Callable[..., Any] + + @wraps(f) + async def new_ainvoke(self, *args, **kwargs): + # type: (Any, Any, Any) -> Any + integration = sentry_sdk.get_client().get_integration(LanggraphIntegration) + if integration is None: + return await f(self, *args, **kwargs) + + graph_name = _get_graph_name(self) + span_name = ( + f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent" + ) + + with sentry_sdk.start_span( + op=OP.GEN_AI_INVOKE_AGENT, + name=span_name, + origin=LanggraphIntegration.origin, + ) as span: + if graph_name: + span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, graph_name) + + 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 + ): + input_messages = _parse_langgraph_messages(args[0]) + if input_messages: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + safe_serialize(input_messages), + ) + + result = await f(self, *args, **kwargs) + + _set_response_attributes(span, input_messages, result, integration) + + return result + + return new_ainvoke + + +def _get_new_messages(input_messages, output_messages): + # type: (Optional[List[Any]], Optional[List[Any]]) -> Optional[List[Any]] + """Extract only the new messages added during this invocation.""" + if not output_messages: + return None + + if not input_messages: + return output_messages + + # only return the new messages, aka the output messages that are not in the input messages + input_count = len(input_messages) + new_messages = ( + output_messages[input_count:] if len(output_messages) > input_count else [] + ) + + return new_messages if new_messages else None + + +def _extract_llm_response_text(messages): + # type: (Optional[List[Any]]) -> Optional[str] + if not messages: + return None + + for message in reversed(messages): + if isinstance(message, dict): + role = message.get("role") + if role in ["assistant", "ai"]: + content = message.get("content") + if content and isinstance(content, str): + return content + + return None + + +def _extract_tool_calls(messages): + # type: (Optional[List[Any]]) -> Optional[List[Any]] + if not messages: + return None + + tool_calls = [] + for message in messages: + if isinstance(message, dict): + msg_tool_calls = message.get("tool_calls") + if msg_tool_calls and isinstance(msg_tool_calls, list): + tool_calls.extend(msg_tool_calls) + + return tool_calls if tool_calls else None + + +def _set_response_attributes(span, input_messages, result, integration): + # type: (Any, Optional[List[Any]], Any, LanggraphIntegration) -> None + if not (should_send_default_pii() and integration.include_prompts): + return + + parsed_response_messages = _parse_langgraph_messages(result) + new_messages = _get_new_messages(input_messages, parsed_response_messages) + + 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, safe_serialize(new_messages) + ) + else: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(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, + ) diff --git a/setup.py b/setup.py index 828dd43461..ca6e7ec534 100644 --- a/setup.py +++ b/setup.py @@ -63,6 +63,7 @@ def get_file_text(file_name): "huey": ["huey>=2"], "huggingface_hub": ["huggingface_hub>=0.22"], "langchain": ["langchain>=0.0.210"], + "langgraph": ["langgraph>=0.6.6"], "launchdarkly": ["launchdarkly-server-sdk>=9.8.0"], "litestar": ["litestar>=2.0.0"], "loguru": ["loguru>=0.5"], diff --git a/tests/integrations/langgraph/__init__.py b/tests/integrations/langgraph/__init__.py new file mode 100644 index 0000000000..b7dd1cb562 --- /dev/null +++ b/tests/integrations/langgraph/__init__.py @@ -0,0 +1,3 @@ +import pytest + +pytest.importorskip("langgraph") diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py new file mode 100644 index 0000000000..5e35f772f5 --- /dev/null +++ b/tests/integrations/langgraph/test_langgraph.py @@ -0,0 +1,632 @@ +import asyncio +import sys +from unittest.mock import MagicMock, patch + +import pytest + +from sentry_sdk import start_transaction +from sentry_sdk.consts import SPANDATA, OP + + +def mock_langgraph_imports(): + """Mock langgraph modules to prevent import errors.""" + mock_state_graph = MagicMock() + mock_pregel = MagicMock() + + langgraph_graph_mock = MagicMock() + langgraph_graph_mock.StateGraph = mock_state_graph + + langgraph_pregel_mock = MagicMock() + langgraph_pregel_mock.Pregel = mock_pregel + + sys.modules["langgraph"] = MagicMock() + sys.modules["langgraph.graph"] = langgraph_graph_mock + sys.modules["langgraph.pregel"] = langgraph_pregel_mock + + return mock_state_graph, mock_pregel + + +mock_state_graph, mock_pregel = mock_langgraph_imports() + +from sentry_sdk.integrations.langgraph import ( # noqa: E402 + LanggraphIntegration, + _parse_langgraph_messages, + _wrap_state_graph_compile, + _wrap_pregel_invoke, + _wrap_pregel_ainvoke, +) + + +class MockStateGraph: + def __init__(self, schema=None): + self.name = "test_graph" + self.schema = schema + self._compiled_graph = None + + def compile(self, *args, **kwargs): + compiled = MockCompiledGraph(self.name) + compiled.graph = self + return compiled + + +class MockCompiledGraph: + def __init__(self, name="test_graph"): + self.name = name + self._graph = None + + def get_graph(self): + return MockGraphRepresentation() + + def invoke(self, state, config=None): + return {"messages": [MockMessage("Response from graph")]} + + async def ainvoke(self, state, config=None): + return {"messages": [MockMessage("Async response from graph")]} + + +class MockGraphRepresentation: + def __init__(self): + self.nodes = {"tools": MockToolsNode()} + + +class MockToolsNode: + def __init__(self): + self.data = MockToolsData() + + +class MockToolsData: + def __init__(self): + self.tools_by_name = { + "search_tool": MockTool("search_tool"), + "calculator": MockTool("calculator"), + } + + +class MockTool: + def __init__(self, name): + self.name = name + + +class MockMessage: + def __init__( + self, + content, + name=None, + tool_calls=None, + function_call=None, + role=None, + type=None, + ): + self.content = content + self.name = name + self.tool_calls = tool_calls + self.function_call = function_call + self.role = role + # The integration uses getattr(message, "type", None) for the role in _normalize_langgraph_message + # Set default type based on name if type not explicitly provided + if type is None and name in ["assistant", "ai", "user", "system", "function"]: + self.type = name + else: + self.type = type + + +class MockPregelInstance: + def __init__(self, name="test_pregel"): + self.name = name + self.graph_name = name + + def invoke(self, state, config=None): + return {"messages": [MockMessage("Pregel response")]} + + async def ainvoke(self, state, config=None): + return {"messages": [MockMessage("Async Pregel response")]} + + +def test_langgraph_integration_init(): + """Test LanggraphIntegration initialization with different parameters.""" + integration = LanggraphIntegration() + assert integration.include_prompts is True + assert integration.identifier == "langgraph" + assert integration.origin == "auto.ai.langgraph" + + integration = LanggraphIntegration(include_prompts=False) + assert integration.include_prompts is False + assert integration.identifier == "langgraph" + assert integration.origin == "auto.ai.langgraph" + + +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [ + (True, True), + (True, False), + (False, True), + (False, False), + ], +) +def test_state_graph_compile( + sentry_init, capture_events, send_default_pii, include_prompts +): + """Test StateGraph.compile() wrapper creates proper create_agent span.""" + sentry_init( + integrations=[LanggraphIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + graph = MockStateGraph() + + def original_compile(self, *args, **kwargs): + return MockCompiledGraph(self.name) + + with patch("sentry_sdk.integrations.langgraph.StateGraph"): + with start_transaction(): + wrapped_compile = _wrap_state_graph_compile(original_compile) + compiled_graph = wrapped_compile( + graph, model="test-model", checkpointer=None + ) + + assert compiled_graph is not None + assert compiled_graph.name == "test_graph" + + tx = events[0] + assert tx["type"] == "transaction" + + agent_spans = [span for span in tx["spans"] if span["op"] == OP.GEN_AI_CREATE_AGENT] + assert len(agent_spans) == 1 + + agent_span = agent_spans[0] + assert agent_span["description"] == "create_agent test_graph" + assert agent_span["origin"] == "auto.ai.langgraph" + assert agent_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "create_agent" + assert agent_span["data"][SPANDATA.GEN_AI_AGENT_NAME] == "test_graph" + assert agent_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "test-model" + assert SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in agent_span["data"] + + tools_data = agent_span["data"][SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS] + assert tools_data == ["search_tool", "calculator"] + assert len(tools_data) == 2 + assert "search_tool" in tools_data + assert "calculator" in tools_data + + +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [ + (True, True), + (True, False), + (False, True), + (False, False), + ], +) +def test_pregel_invoke(sentry_init, capture_events, send_default_pii, include_prompts): + """Test Pregel.invoke() wrapper creates proper invoke_agent span.""" + sentry_init( + integrations=[LanggraphIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + + test_state = { + "messages": [ + MockMessage("Hello, can you help me?", name="user"), + MockMessage("Of course! How can I assist you?", name="assistant"), + ] + } + + pregel = MockPregelInstance("test_graph") + + expected_assistant_response = "I'll help you with that task!" + expected_tool_calls = [ + { + "id": "call_test_123", + "type": "function", + "function": {"name": "search_tool", "arguments": '{"query": "help"}'}, + } + ] + + def original_invoke(self, *args, **kwargs): + input_messages = args[0].get("messages", []) + new_messages = input_messages + [ + MockMessage( + content=expected_assistant_response, + name="assistant", + tool_calls=expected_tool_calls, + ) + ] + return {"messages": new_messages} + + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) + + assert result is not None + + tx = events[0] + assert tx["type"] == "transaction" + + invoke_spans = [ + span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + + invoke_span = invoke_spans[0] + assert invoke_span["description"] == "invoke_agent test_graph" + assert invoke_span["origin"] == "auto.ai.langgraph" + assert invoke_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" + assert invoke_span["data"][SPANDATA.GEN_AI_PIPELINE_NAME] == "test_graph" + assert invoke_span["data"][SPANDATA.GEN_AI_AGENT_NAME] == "test_graph" + + if send_default_pii and include_prompts: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["data"] + + request_messages = invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + + if isinstance(request_messages, str): + import json + + request_messages = json.loads(request_messages) + 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["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] + assert response_text == expected_assistant_response + + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["data"] + tool_calls_data = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + if isinstance(tool_calls_data, str): + import json + + tool_calls_data = json.loads(tool_calls_data) + + assert len(tool_calls_data) == 1 + assert tool_calls_data[0]["id"] == "call_test_123" + assert tool_calls_data[0]["function"]["name"] == "search_tool" + else: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get("data", {}) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get("data", {}) + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get("data", {}) + + +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [ + (True, True), + (True, False), + (False, True), + (False, False), + ], +) +def test_pregel_ainvoke(sentry_init, capture_events, send_default_pii, include_prompts): + """Test Pregel.ainvoke() async wrapper creates proper invoke_agent span.""" + sentry_init( + integrations=[LanggraphIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + test_state = {"messages": [MockMessage("What's the weather like?", name="user")]} + pregel = MockPregelInstance("async_graph") + + expected_assistant_response = "It's sunny and 72°F today!" + expected_tool_calls = [ + { + "id": "call_weather_456", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"location": "current"}'}, + } + ] + + async def original_ainvoke(self, *args, **kwargs): + input_messages = args[0].get("messages", []) + new_messages = input_messages + [ + MockMessage( + content=expected_assistant_response, + name="assistant", + tool_calls=expected_tool_calls, + ) + ] + return {"messages": new_messages} + + async def run_test(): + with start_transaction(): + + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + result = await wrapped_ainvoke(pregel, test_state) + return result + + result = asyncio.run(run_test()) + assert result is not None + + tx = events[0] + assert tx["type"] == "transaction" + + invoke_spans = [ + span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + + invoke_span = invoke_spans[0] + assert invoke_span["description"] == "invoke_agent async_graph" + assert invoke_span["origin"] == "auto.ai.langgraph" + assert invoke_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" + assert invoke_span["data"][SPANDATA.GEN_AI_PIPELINE_NAME] == "async_graph" + assert invoke_span["data"][SPANDATA.GEN_AI_AGENT_NAME] == "async_graph" + + if send_default_pii and include_prompts: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["data"] + + response_text = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] + assert response_text == expected_assistant_response + + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["data"] + tool_calls_data = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + if isinstance(tool_calls_data, str): + import json + + tool_calls_data = json.loads(tool_calls_data) + + assert len(tool_calls_data) == 1 + assert tool_calls_data[0]["id"] == "call_weather_456" + assert tool_calls_data[0]["function"]["name"] == "get_weather" + else: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get("data", {}) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get("data", {}) + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get("data", {}) + + +def test_pregel_invoke_error(sentry_init, capture_events): + """Test error handling during graph execution.""" + sentry_init( + integrations=[LanggraphIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + test_state = {"messages": [MockMessage("This will fail")]} + pregel = MockPregelInstance("error_graph") + + def original_invoke(self, *args, **kwargs): + raise Exception("Graph execution failed") + + with start_transaction(), pytest.raises(Exception, match="Graph execution failed"): + + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, test_state) + + tx = events[0] + invoke_spans = [ + span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + + invoke_span = invoke_spans[0] + assert invoke_span.get("tags", {}).get("status") == "internal_error" + + +def test_pregel_ainvoke_error(sentry_init, capture_events): + """Test error handling during async graph execution.""" + sentry_init( + integrations=[LanggraphIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + test_state = {"messages": [MockMessage("This will fail async")]} + pregel = MockPregelInstance("async_error_graph") + + async def original_ainvoke(self, *args, **kwargs): + raise Exception("Async graph execution failed") + + async def run_error_test(): + with start_transaction(), pytest.raises( + Exception, match="Async graph execution failed" + ): + + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + await wrapped_ainvoke(pregel, test_state) + + asyncio.run(run_error_test()) + + tx = events[0] + invoke_spans = [ + span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + + invoke_span = invoke_spans[0] + assert invoke_span.get("tags", {}).get("status") == "internal_error" + + +def test_span_origin(sentry_init, capture_events): + """Test that span origins are correctly set.""" + sentry_init( + integrations=[LanggraphIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + graph = MockStateGraph() + + def original_compile(self, *args, **kwargs): + return MockCompiledGraph(self.name) + + with start_transaction(): + from sentry_sdk.integrations.langgraph import _wrap_state_graph_compile + + wrapped_compile = _wrap_state_graph_compile(original_compile) + wrapped_compile(graph) + + tx = events[0] + assert tx["contexts"]["trace"]["origin"] == "manual" + + for span in tx["spans"]: + assert span["origin"] == "auto.ai.langgraph" + + +@pytest.mark.parametrize("graph_name", ["my_graph", None, ""]) +def test_pregel_invoke_with_different_graph_names( + sentry_init, capture_events, graph_name +): + """Test Pregel.invoke() with different graph name scenarios.""" + sentry_init( + integrations=[LanggraphIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + pregel = MockPregelInstance(graph_name) if graph_name else MockPregelInstance() + if not graph_name: + + delattr(pregel, "name") + delattr(pregel, "graph_name") + + def original_invoke(self, *args, **kwargs): + return {"result": "test"} + + with start_transaction(): + + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, {"messages": []}) + + tx = events[0] + invoke_spans = [ + span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + + invoke_span = invoke_spans[0] + + if graph_name and graph_name.strip(): + assert invoke_span["description"] == "invoke_agent my_graph" + assert invoke_span["data"][SPANDATA.GEN_AI_PIPELINE_NAME] == graph_name + assert invoke_span["data"][SPANDATA.GEN_AI_AGENT_NAME] == graph_name + else: + assert invoke_span["description"] == "invoke_agent" + assert SPANDATA.GEN_AI_PIPELINE_NAME not in invoke_span.get("data", {}) + assert SPANDATA.GEN_AI_AGENT_NAME not in invoke_span.get("data", {}) + + +def test_complex_message_parsing(): + """Test message parsing with complex message structures.""" + messages = [ + MockMessage(content="User query", name="user"), + MockMessage( + content="Assistant response with tools", + name="assistant", + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": {"name": "search", "arguments": "{}"}, + }, + { + "id": "call_2", + "type": "function", + "function": {"name": "calculate", "arguments": '{"x": 5}'}, + }, + ], + ), + MockMessage( + content="Function call response", + name="function", + function_call={"name": "search", "arguments": '{"query": "test"}'}, + ), + ] + + state = {"messages": messages} + result = _parse_langgraph_messages(state) + + assert result is not None + assert len(result) == 3 + + assert result[0]["content"] == "User query" + assert result[0]["name"] == "user" + assert "tool_calls" not in result[0] + assert "function_call" not in result[0] + + assert result[1]["content"] == "Assistant response with tools" + assert result[1]["name"] == "assistant" + assert len(result[1]["tool_calls"]) == 2 + + assert result[2]["content"] == "Function call response" + assert result[2]["name"] == "function" + assert result[2]["function_call"]["name"] == "search" + + +def test_extraction_functions_complex_scenario(sentry_init, capture_events): + """Test extraction functions with complex scenarios including multiple messages and edge cases.""" + sentry_init( + integrations=[LanggraphIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + pregel = MockPregelInstance("complex_graph") + test_state = {"messages": [MockMessage("Complex request", name="user")]} + + def original_invoke(self, *args, **kwargs): + input_messages = args[0].get("messages", []) + new_messages = input_messages + [ + MockMessage( + content="I'll help with multiple tasks", + name="assistant", + tool_calls=[ + { + "id": "call_multi_1", + "type": "function", + "function": { + "name": "search", + "arguments": '{"query": "complex"}', + }, + }, + { + "id": "call_multi_2", + "type": "function", + "function": { + "name": "calculate", + "arguments": '{"expr": "2+2"}', + }, + }, + ], + ), + MockMessage("", name="assistant"), + MockMessage("Final response", name="ai", type="ai"), + ] + return {"messages": new_messages} + + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) + + assert result is not None + + tx = events[0] + invoke_spans = [ + span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + + invoke_span = invoke_spans[0] + assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["data"] + response_text = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] + assert response_text == "Final response" + + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["data"] + import json + + tool_calls_data = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + if isinstance(tool_calls_data, str): + tool_calls_data = json.loads(tool_calls_data) + + assert len(tool_calls_data) == 2 + assert tool_calls_data[0]["id"] == "call_multi_1" + assert tool_calls_data[0]["function"]["name"] == "search" + assert tool_calls_data[1]["id"] == "call_multi_2" + assert tool_calls_data[1]["function"]["name"] == "calculate" diff --git a/tox.ini b/tox.ini index fd633654be..40afc2a6a7 100644 --- a/tox.ini +++ b/tox.ini @@ -10,7 +10,7 @@ # The file (and all resulting CI YAMLs) then need to be regenerated via # "scripts/generate-test-files.sh". # -# Last generated: 2025-09-04T10:35:13.756355+00:00 +# Last generated: 2025-09-04T12:59:44.328902+00:00 [tox] requires = @@ -142,6 +142,9 @@ envlist = {py3.8,py3.11,py3.12}-openai-notiktoken-v1.71.0 {py3.8,py3.12,py3.13}-openai-notiktoken-v1.105.0 + {py3.9,py3.12,py3.13}-langgraph-v0.6.6 + {py3.10,py3.12,py3.13}-langgraph-v1.0.0a2 + {py3.10,py3.11,py3.12}-openai_agents-v0.0.19 {py3.10,py3.12,py3.13}-openai_agents-v0.1.0 {py3.10,py3.12,py3.13}-openai_agents-v0.2.11 @@ -525,6 +528,9 @@ deps = openai-notiktoken-v1.0.1: httpx<0.28 openai-notiktoken-v1.36.1: httpx<0.28 + langgraph-v0.6.6: langgraph==0.6.6 + langgraph-v1.0.0a2: langgraph==1.0.0a2 + openai_agents-v0.0.19: openai-agents==0.0.19 openai_agents-v0.1.0: openai-agents==0.1.0 openai_agents-v0.2.11: openai-agents==0.2.11 @@ -841,6 +847,7 @@ setenv = huggingface_hub: TESTPATH=tests/integrations/huggingface_hub langchain-base: TESTPATH=tests/integrations/langchain langchain-notiktoken: TESTPATH=tests/integrations/langchain + langgraph: TESTPATH=tests/integrations/langgraph launchdarkly: TESTPATH=tests/integrations/launchdarkly litestar: TESTPATH=tests/integrations/litestar loguru: TESTPATH=tests/integrations/loguru