diff --git a/docs/supported-integrations/langgraph.mdx b/docs/supported-integrations/langgraph.mdx index 595a7e596..02debf25f 100644 --- a/docs/supported-integrations/langgraph.mdx +++ b/docs/supported-integrations/langgraph.mdx @@ -64,6 +64,53 @@ with nemo_relay.scope.scope("langgraph-request", nemo_relay.ScopeType.Agent): print(result) ``` +## Managed ToolNode Calls + +For a standalone LangGraph `ToolNode`, construct the node with +`create_tool_node`. It routes each model-requested tool call through Relay's +managed execution pipeline while preserving LangGraph state, store, and runtime +argument injection. + +```python +from langchain_core.messages import AIMessage +from langchain_core.tools import tool +from langgraph.graph import END, START, MessagesState, StateGraph +from nemo_relay.integrations.langgraph import ( + NemoRelayCallbackHandler, + create_tool_node, +) + +@tool +def get_weather(location: str) -> str: + """Return the weather for a location.""" + return f"Sunny in {location}" + +builder = StateGraph(MessagesState) +builder.add_node("tools", create_tool_node([get_weather])) +builder.add_edge(START, "tools") +builder.add_edge("tools", END) +graph = builder.compile() + +with nemo_relay.scope.scope("langgraph-request", nemo_relay.ScopeType.Agent): + result = graph.invoke( + { + "messages": [ + AIMessage( + content="", + tool_calls=[{"name": "get_weather", "args": {"location": "Boston"}, "id": "call-1"}], + ) + ] + }, + config={"callbacks": [NemoRelayCallbackHandler()]}, + ) +``` + +`NemoRelayCallbackHandler` records the graph lifecycle and provides its scopes. +`create_tool_node` provides managed tool execution; use both for complete +standalone LangGraph instrumentation. For custom ToolNode wrapper composition, +construct `ToolNode` directly and pass `wrap_tool_call` and `awrap_tool_call` +from this integration. + For LangChain agents inside a LangGraph workflow, use `NemoRelayMiddleware` from this package the same way as the LangChain integration and pass the LangGraph `config` into the nested agent call: diff --git a/python/nemo_relay/integrations/langgraph/__init__.py b/python/nemo_relay/integrations/langgraph/__init__.py index 79c021c4c..a9da9218b 100644 --- a/python/nemo_relay/integrations/langgraph/__init__.py +++ b/python/nemo_relay/integrations/langgraph/__init__.py @@ -5,8 +5,12 @@ from nemo_relay.integrations.langchain import NemoRelayMiddleware from nemo_relay.integrations.langgraph.callbacks import NemoRelayCallbackHandler +from nemo_relay.integrations.langgraph.tool_node import awrap_tool_call, create_tool_node, wrap_tool_call __all__ = [ "NemoRelayCallbackHandler", "NemoRelayMiddleware", + "awrap_tool_call", + "create_tool_node", + "wrap_tool_call", ] diff --git a/python/nemo_relay/integrations/langgraph/tool_node.py b/python/nemo_relay/integrations/langgraph/tool_node.py new file mode 100644 index 000000000..78372e55f --- /dev/null +++ b/python/nemo_relay/integrations/langgraph/tool_node.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Managed NeMo Relay wrappers for standalone LangGraph ``ToolNode`` objects.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import replace +from typing import Any, cast + +from langchain_core.messages import ToolMessage +from langchain_core.tools import BaseTool +from langgraph.errors import GraphBubbleUp +from langgraph.prebuilt import ToolNode +from langgraph.prebuilt.tool_node import ToolCallRequest, ToolInvocationError +from langgraph.types import Command + +import nemo_relay +from nemo_relay.typed import Codec +from nemo_relay.utils import run_sync + + +class _ToolNodeResultCodec(nemo_relay.typed.BestEffortAnyCodec): + """Restore ToolMessages nested in LangGraph Command updates.""" + + _LIST_TAG = "__nemo_relay_langgraph_list__" + + def to_json(self, value: object) -> nemo_relay.Json: + if isinstance(value, list): + return {self._LIST_TAG: [self.to_json(item) for item in value]} + return super().to_json(value) + + def from_json(self, data: nemo_relay.Json) -> object: + if isinstance(data, dict) and isinstance(data.get(self._LIST_TAG), list): + return [self.from_json(item) for item in data[self._LIST_TAG]] + + result = super().from_json(data) + if not isinstance(result, Command): + return result + + if isinstance(result.update, dict): + messages = result.update.get("messages") + if isinstance(messages, list): + result.update["messages"] = _restore_tool_messages(messages) + elif isinstance(result.update, list): + return replace(result, update=_restore_tool_messages(result.update)) + return result + + +def _restore_tool_messages(messages: list[object]) -> list[object]: + """Reconstruct serialized ToolMessages in a LangGraph command update.""" + return [ToolMessage.model_validate(message) if isinstance(message, dict) else message for message in messages] + + +_DEFAULT_HANDLE_TOOL_ERRORS = object() +_TOOL_CALL_ERROR_TEMPLATE = "Error: {error}\n Please fix your mistakes." +_GRAPH_BUBBLE_RESULT = {"__nemo_relay_langgraph_graph_bubble_up__": True} + + +def _handle_tool_error(error: Exception, policy: object) -> str: + """Preserve ToolNode error handling while allowing graph bubbles to escape.""" + if isinstance(error, GraphBubbleUp): + raise error + if policy is _DEFAULT_HANDLE_TOOL_ERRORS: + if isinstance(error, ToolInvocationError): + return error.message + raise error + if isinstance(policy, tuple): + if not isinstance(error, policy): + raise error + return _TOOL_CALL_ERROR_TEMPLATE.format(error=repr(error)) + if isinstance(policy, type) and issubclass(policy, Exception): + if not isinstance(error, policy): + raise error + return _TOOL_CALL_ERROR_TEMPLATE.format(error=repr(error)) + if policy is True: + return _TOOL_CALL_ERROR_TEMPLATE.format(error=repr(error)) + if isinstance(policy, str): + return policy + if callable(policy): + return cast(Callable[[Exception], str], policy)(error) + raise ValueError(f"unexpected handle_tool_errors value: {policy}") + + +def _tool_details(request: ToolCallRequest) -> tuple[nemo_relay.ScopeHandle, str, dict[str, Any], str | None]: + """Extract the model-controlled tool-call fields managed by Relay.""" + return ( + nemo_relay.scope.get_handle(), + request.tool_call["name"], + request.tool_call.get("args") or {}, + request.tool_call.get("id"), + ) + + +def wrap_tool_call( + request: ToolCallRequest, + execute: Callable[[ToolCallRequest], ToolMessage | Command[Any]], +) -> ToolMessage | Command[Any]: + """Run one synchronous LangGraph tool call through NeMo Relay. + + Args: + request: LangGraph's tool-call request. + execute: LangGraph callback that invokes the requested tool. + + Returns: + The LangGraph tool result after managed Relay execution. + """ + parent, tool_name, tool_args, tool_call_id = _tool_details(request) + args_codec = cast(Codec[dict[str, Any]], nemo_relay.typed.BestEffortAnyCodec()) + result_codec = cast(Codec[ToolMessage | Command[Any] | dict[str, bool]], _ToolNodeResultCodec()) + graph_bubble: GraphBubbleUp | None = None + + def _call(args: dict[str, Any]) -> nemo_relay.ToolExecutionResult[ToolMessage | Command[Any] | dict[str, bool]]: + nonlocal graph_bubble + try: + result = execute(request.override(tool_call={**request.tool_call, "args": args})) + except GraphBubbleUp as error: + # Relay's native callback boundary cannot propagate arbitrary Python + # exceptions. Preserve the original graph bubble locally and re-raise + # it after managed execution returns. + graph_bubble = error + result = _GRAPH_BUBBLE_RESULT + return nemo_relay.ToolExecutionResult(result) + + outcome = run_sync( + nemo_relay.typed.tool_execute( + name=tool_name, + args=tool_args, + func=_call, + args_codec=args_codec, + result_codec=result_codec, + handle=parent, + tool_call_id=tool_call_id, + ) + ) + if graph_bubble is not None: + raise graph_bubble + return cast(ToolMessage | Command[Any], outcome.result) + + +async def awrap_tool_call( + request: ToolCallRequest, + execute: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]], +) -> ToolMessage | Command[Any]: + """Run one asynchronous LangGraph tool call through NeMo Relay. + + Args: + request: LangGraph's tool-call request. + execute: Async LangGraph callback that invokes the requested tool. + + Returns: + The LangGraph tool result after managed Relay execution. + """ + parent, tool_name, tool_args, tool_call_id = _tool_details(request) + args_codec = cast(Codec[dict[str, Any]], nemo_relay.typed.BestEffortAnyCodec()) + result_codec = cast(Codec[ToolMessage | Command[Any] | dict[str, bool]], _ToolNodeResultCodec()) + graph_bubble: GraphBubbleUp | None = None + + async def _call( + args: dict[str, Any], + ) -> nemo_relay.ToolExecutionResult[ToolMessage | Command[Any] | dict[str, bool]]: + nonlocal graph_bubble + try: + result = await execute(request.override(tool_call={**request.tool_call, "args": args})) + except GraphBubbleUp as error: + # See the synchronous wrapper: retain the original exception instead + # of letting the native callback boundary convert it to RuntimeError. + graph_bubble = error + result = _GRAPH_BUBBLE_RESULT + return nemo_relay.ToolExecutionResult(result) + + outcome = await nemo_relay.typed.tool_execute( + name=tool_name, + args=tool_args, + func=_call, + args_codec=args_codec, + result_codec=result_codec, + handle=parent, + tool_call_id=tool_call_id, + ) + if graph_bubble is not None: + raise graph_bubble + return cast(ToolMessage | Command[Any], outcome.result) + + +def create_tool_node( + tools: Sequence[BaseTool | Callable[..., Any]], + **tool_node_kwargs: Any, +) -> ToolNode: + """Create a LangGraph ``ToolNode`` whose tool calls use managed Relay execution. + + For custom LangGraph wrapper composition, construct ``ToolNode`` directly + and pass :func:`wrap_tool_call` and :func:`awrap_tool_call` explicitly. + + Args: + tools: LangGraph tools available to the returned node. + **tool_node_kwargs: Remaining native ``ToolNode`` constructor options, + excluding the Relay-managed wrapper options. + + Returns: + A native ``ToolNode`` configured with Relay's sync and async wrappers. + """ + configured_wrappers = {"wrap_tool_call", "awrap_tool_call"}.intersection(tool_node_kwargs) + if configured_wrappers: + names = ", ".join(sorted(configured_wrappers)) + raise ValueError(f"create_tool_node configures {names}; construct ToolNode directly to compose custom wrappers") + error_policy = tool_node_kwargs.pop("handle_tool_errors", _DEFAULT_HANDLE_TOOL_ERRORS) + if error_policy is False: + error_handler: bool | Callable[[Exception], str] = False + else: + + def error_handler(error: Exception) -> str: + return _handle_tool_error(error, error_policy) + + return ToolNode( + tools, + wrap_tool_call=wrap_tool_call, + awrap_tool_call=awrap_tool_call, + handle_tool_errors=error_handler, + **tool_node_kwargs, + ) + + +__all__ = ["awrap_tool_call", "create_tool_node", "wrap_tool_call"] diff --git a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py index 7696c241e..57b845979 100644 --- a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py +++ b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py @@ -25,6 +25,11 @@ class State(TypedDict): value: int +class ToolNodeState(TypedDict): + messages: Annotated[list[Any], operator.add] + offset: int + + def increment(state: State) -> State: return {"value": state["value"] + 1} @@ -86,6 +91,331 @@ def test_handler_type(callback_handler: NemoRelayCallbackHandler): assert isinstance(callback_handler, GraphCallbackHandler) +@pytest.mark.parametrize("use_async", [False, True]) +def test_create_tool_node_routes_standalone_tool_calls_through_relay( + use_async: bool, + subscribed_events: list[nemo_relay.Event], +): + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, StateGraph + from langgraph.prebuilt import InjectedState + + from nemo_relay.integrations.langgraph import create_tool_node + + def add_offset(value: int, state: Any) -> str: + """Add the graph state's offset to a model-provided value.""" + return str(value + state["offset"]) + + add_offset.__annotations__["state"] = Annotated[dict[str, Any], InjectedState] + add_offset_tool = tool(add_offset) + + async def rewrite_tool_args(_name: str, args: nemo_relay.Json, next_call: Any) -> Any: + downstream = await next_call({**cast(dict[str, Any], args), "value": 4}) + return nemo_relay.ToolExecutionInterceptOutcome( + downstream.result, + annotation=downstream.annotation, + ) + + node = create_tool_node([add_offset_tool], name="managed-tools") + builder = StateGraph(ToolNodeState) + builder.add_node("tools", node) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile() + nemo_relay.intercepts.register_tool_execution("langgraph-rewrite-tool-args", 1, rewrite_tool_args) + try: + with nemo_relay.scope.scope("request", nemo_relay.ScopeType.Agent) as request: + input_state = { + "offset": 3, + "messages": [ + AIMessage( + content="", + tool_calls=[ + {"name": "add_offset", "args": {"value": 1}, "id": "call-1"}, + {"name": "add_offset", "args": {"value": 2}, "id": "call-3"}, + ], + ) + ], + } + if use_async: + result = asyncio.run(graph.ainvoke(input_state)) + else: + result = graph.invoke(input_state) + finally: + nemo_relay.intercepts.deregister_tool_execution("langgraph-rewrite-tool-args") + + nemo_relay.subscribers.flush() + assert [message.content for message in result["messages"][-2:]] == ["7", "7"] + tool_events = [ + event for event in subscribed_events if isinstance(event, nemo_relay.ScopeEvent) and event.name == "add_offset" + ] + assert len(tool_events) == 4 + assert {event.scope_category for event in tool_events} == {"start", "end"} + assert all(event.parent_uuid == request.uuid for event in tool_events) + assert {event.category_profile["tool_call_id"] for event in tool_events if event.category_profile} == { + "call-1", + "call-3", + } + + +def test_exported_tool_node_wrappers_support_direct_tool_node_construction( + subscribed_events: list[nemo_relay.Event], +): + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, MessagesState, StateGraph + from langgraph.prebuilt import ToolNode + + from nemo_relay.integrations.langgraph import awrap_tool_call, wrap_tool_call + + @tool + def echo(value: str) -> str: + """Echo a value.""" + return value + + builder = StateGraph(MessagesState) + builder.add_node("tools", ToolNode([echo], wrap_tool_call=wrap_tool_call, awrap_tool_call=awrap_tool_call)) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile() + with nemo_relay.scope.scope("request", nemo_relay.ScopeType.Agent): + result = graph.invoke( + { + "messages": [ + AIMessage( + content="", + tool_calls=[{"name": "echo", "args": {"value": "managed"}, "id": "call-2"}], + ) + ] + } + ) + + nemo_relay.subscribers.flush() + assert result["messages"][-1].content == "managed" + assert any( + isinstance(event, nemo_relay.ScopeEvent) + and event.name == "echo" + and event.category_profile == {"tool_call_id": "call-2"} + for event in subscribed_events + ) + + +def test_create_tool_node_preserves_command_and_error_handling( + subscribed_events: list[nemo_relay.Event], +): + from langchain_core.messages import AIMessage, ToolMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, StateGraph + from langgraph.types import Command + + from nemo_relay.integrations.langgraph import create_tool_node + + @tool + def update_offset(): + """Update graph state through a LangGraph command.""" + return Command( + update={ + "offset": 9, + "messages": [ToolMessage(content="updated", tool_call_id="call-command")], + } + ) + + @tool + def fail() -> str: + """Raise a tool failure handled by ToolNode.""" + raise ValueError("expected failure") + + command_builder = StateGraph(ToolNodeState) + command_builder.add_node("command", create_tool_node([update_offset])) + command_builder.add_edge(START, "command") + command_builder.add_edge("command", END) + command_graph = command_builder.compile() + + failure_builder = StateGraph(ToolNodeState) + failure_builder.add_node("failure", create_tool_node([fail], handle_tool_errors=True)) + failure_builder.add_edge(START, "failure") + failure_builder.add_edge("failure", END) + failure_graph = failure_builder.compile() + + with nemo_relay.scope.scope("request", nemo_relay.ScopeType.Agent): + command_result = command_graph.invoke( + { + "offset": 0, + "messages": [ + AIMessage( + content="", + tool_calls=[{"name": "update_offset", "args": {}, "id": "call-command"}], + ) + ], + } + ) + failure_result = failure_graph.invoke( + { + "offset": 0, + "messages": [AIMessage(content="", tool_calls=[{"name": "fail", "args": {}, "id": "call-failure"}])], + } + ) + + nemo_relay.subscribers.flush() + assert command_result["offset"] == 9 + assert command_result["messages"][-1].content == "updated" + assert failure_result["messages"][-1].status == "error" + assert {event.name for event in subscribed_events if isinstance(event, nemo_relay.ScopeEvent)} == { + "request", + "update_offset", + "fail", + } + + +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("policy", [ValueError, (ValueError,)]) +def test_create_tool_node_preserves_selective_error_handling( + use_async: bool, + policy: type[ValueError] | tuple[type[ValueError], ...], +): + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, StateGraph + + from nemo_relay.integrations.langgraph import create_tool_node + + @tool + def value_error() -> str: + """Raise an error selected by the ToolNode policy.""" + raise ValueError("handled") + + @tool + def runtime_error() -> str: + """Raise an error outside the ToolNode policy.""" + raise RuntimeError("unhandled") + + def build_graph(tool: Any): + builder = StateGraph(ToolNodeState) + builder.add_node("tools", create_tool_node([tool], handle_tool_errors=policy)) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + return builder.compile() + + handled_graph = build_graph(value_error) + unhandled_graph = build_graph(runtime_error) + handled_input = { + "offset": 0, + "messages": [AIMessage(content="", tool_calls=[{"name": "value_error", "args": {}, "id": "call-value"}])], + } + unhandled_input = { + "offset": 0, + "messages": [AIMessage(content="", tool_calls=[{"name": "runtime_error", "args": {}, "id": "call-runtime"}])], + } + + if use_async: + + async def invoke_handled() -> ToolNodeState: + return await handled_graph.ainvoke(handled_input) + + async def invoke_unhandled() -> ToolNodeState: + return await unhandled_graph.ainvoke(unhandled_input) + + handled_result = asyncio.run(invoke_handled()) + with pytest.raises(RuntimeError, match="internal error"): + asyncio.run(invoke_unhandled()) + else: + handled_result = handled_graph.invoke(handled_input) + with pytest.raises(RuntimeError, match="internal error"): + unhandled_graph.invoke(unhandled_input) + + assert handled_result["messages"][-1].status == "error" + + +@pytest.mark.parametrize("use_async", [False, True]) +def test_create_tool_node_preserves_list_command_results(use_async: bool): + from langchain_core.messages import AIMessage, ToolMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, StateGraph + from langgraph.types import Command + + from nemo_relay.integrations.langgraph import create_tool_node + + @tool + def update_offset(): + """Return a list containing a graph state update command.""" + return [ + Command( + update={ + "offset": 9, + "messages": [ToolMessage(content="updated", tool_call_id="call-list")], + } + ) + ] + + builder = StateGraph(ToolNodeState) + builder.add_node("tools", create_tool_node([update_offset])) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile() + input_state = { + "offset": 0, + "messages": [AIMessage(content="", tool_calls=[{"name": "update_offset", "args": {}, "id": "call-list"}])], + } + + if use_async: + result = asyncio.run(graph.ainvoke(input_state)) + else: + result = graph.invoke(input_state) + + assert result["offset"] == 9 + assert result["messages"][-1].content == "updated" + + +@pytest.mark.parametrize("use_async", [False, True]) +def test_create_tool_node_propagates_graph_interrupts(use_async: bool): + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + from langgraph.checkpoint.memory import MemorySaver + from langgraph.graph import END, START, StateGraph + from langgraph.types import interrupt + + from nemo_relay.integrations.langgraph import create_tool_node + + @tool + def await_approval() -> str: + """Pause the graph until approval is supplied.""" + interrupt("approval required") + return "approved" + + builder = StateGraph(ToolNodeState) + builder.add_node("tools", create_tool_node([await_approval])) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile(checkpointer=MemorySaver()) + config = {"configurable": {"thread_id": str(uuid4())}} + input_state = { + "offset": 0, + "messages": [ + AIMessage(content="", tool_calls=[{"name": "await_approval", "args": {}, "id": "call-interrupt"}]) + ], + } + + if use_async: + + async def collect_updates() -> list[dict[str, Any]]: + return [update async for update in graph.astream(input_state, config, stream_mode="updates")] + + updates = asyncio.run(collect_updates()) + else: + updates = list(graph.stream(input_state, config, stream_mode="updates")) + + assert any("__interrupt__" in update for update in updates) + + +@pytest.mark.parametrize("wrapper_name", ["wrap_tool_call", "awrap_tool_call"]) +def test_create_tool_node_rejects_custom_tool_wrappers(wrapper_name: str): + from nemo_relay.integrations.langgraph import create_tool_node + + with pytest.raises(ValueError, match="construct ToolNode directly"): + create_tool_node([], **{wrapper_name: lambda *_args: None}) + + class TestGraphCallbacks: _expected_events = [ "scope.start.request",