From f8ca30d8e0fa3e409261a2cb97bd396bcab5e1d7 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Wed, 19 Nov 2025 15:54:25 -0500 Subject: [PATCH] better typing --- libs/prebuilt/langgraph/prebuilt/tool_node.py | 24 ++++++----- libs/prebuilt/tests/test_on_tool_call.py | 40 ++++++++++++++----- libs/prebuilt/tests/test_tool_node.py | 19 ++++++--- 3 files changed, 58 insertions(+), 25 deletions(-) diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index 86084c0bb..805f5ef26 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -124,7 +124,7 @@ class _ToolCallRequestOverrides(TypedDict, total=False): @dataclass -class ToolCallRequest: +class ToolCallRequest(Generic[ContextT, StateT]): """Tool execution request passed to tool call interceptors. Attributes: @@ -133,14 +133,22 @@ class ToolCallRequest: registered with the `ToolNode`. When tool is `None`, interceptors can handle the request without validation. If the interceptor calls `execute()`, validation will occur and raise an error for unregistered tools. - state: Agent state (`dict`, `list`, or `BaseModel`). runtime: LangGraph runtime context (optional, `None` if outside graph). + state: Agent state (`dict`, `list`, or `BaseModel`). Pulled from `runtime.state`. """ tool_call: ToolCall tool: BaseTool | None - state: Any - runtime: ToolRuntime + runtime: ToolRuntime[ContextT, StateT] + + @property + def state(self) -> StateT: + """Get the state from the runtime. + + Returns: + The current graph state from the runtime context. + """ + return self.runtime.state def __setattr__(self, name: str, value: Any) -> None: """Raise deprecation warning when setting attributes directly. @@ -163,7 +171,7 @@ class ToolCallRequest: def override( self, **overrides: Unpack[_ToolCallRequestOverrides] - ) -> ToolCallRequest: + ) -> ToolCallRequest[ContextT, StateT]: """Replace the request with a new request with the given overrides. Returns a new `ToolCallRequest` instance with the specified attributes replaced. @@ -947,11 +955,10 @@ class ToolNode(RunnableCallable): # to short-circuit requests for unregistered tools tool = self.tools_by_name.get(call["name"]) - # Create the tool request with state and runtime + # Create the tool request with runtime tool_request = ToolCallRequest( tool_call=call, tool=tool, - state=tool_runtime.state, runtime=tool_runtime, ) @@ -1104,11 +1111,10 @@ class ToolNode(RunnableCallable): # to short-circuit requests for unregistered tools tool = self.tools_by_name.get(call["name"]) - # Create the tool request with state and runtime + # Create the tool request with runtime tool_request = ToolCallRequest( tool_call=call, tool=tool, - state=tool_runtime.state, runtime=tool_runtime, ) diff --git a/libs/prebuilt/tests/test_on_tool_call.py b/libs/prebuilt/tests/test_on_tool_call.py index bdff99222..33c7a226d 100644 --- a/libs/prebuilt/tests/test_on_tool_call.py +++ b/libs/prebuilt/tests/test_on_tool_call.py @@ -13,6 +13,7 @@ from langgraph.types import Command from langgraph.prebuilt.tool_node import ( ToolCallRequest, ToolNode, + ToolRuntime, ) pytestmark = pytest.mark.anyio @@ -342,16 +343,23 @@ def test_tool_call_request_dataclass() -> None: """Test ToolCallRequest dataclass.""" tool_call: ToolCall = {"name": "add", "args": {"a": 1, "b": 2}, "id": "call_1"} state: dict = {"messages": []} - runtime = None + tool_runtime = ToolRuntime( + state=state, + config={}, + context=None, + store=None, + stream_writer=Mock(), + tool_call_id="call_1", + ) request = ToolCallRequest( - tool_call=tool_call, tool=add, state=state, runtime=runtime - ) # type: ignore[arg-type] + tool_call=tool_call, tool=add, runtime=tool_runtime + ) assert request.tool_call == tool_call assert request.tool == add assert request.state == state - assert request.runtime is None + assert request.runtime is tool_runtime assert request.tool_call["name"] == "add" @@ -1324,11 +1332,18 @@ def test_tool_call_request_is_frozen() -> None: """Test that ToolCallRequest raises deprecation warnings on direct attribute reassignment.""" tool_call: ToolCall = {"name": "add", "args": {"a": 1, "b": 2}, "id": "call_1"} state: dict = {"messages": []} - runtime = None + tool_runtime = ToolRuntime( + state=state, + config={}, + context=None, + store=None, + stream_writer=Mock(), + tool_call_id="call_1", + ) request = ToolCallRequest( - tool_call=tool_call, tool=add, state=state, runtime=runtime - ) # type: ignore[arg-type] + tool_call=tool_call, tool=add, runtime=tool_runtime + ) # Test that direct attribute reassignment raises DeprecationWarning with pytest.warns( @@ -1343,11 +1358,14 @@ def test_tool_call_request_is_frozen() -> None: ): request.tool = None # type: ignore[misc] + # state is now a property, so setting it will raise a deprecation warning + # (and then fail with AttributeError after the warning) with pytest.warns( DeprecationWarning, match="Setting attribute 'state' on ToolCallRequest is deprecated", ): - request.state = {} # type: ignore[misc] + with pytest.raises(AttributeError): + request.state = {} # type: ignore[misc] with pytest.warns( DeprecationWarning, @@ -1365,8 +1383,8 @@ def test_tool_call_request_is_frozen() -> None: # Original request should be unchanged (note: it was modified by the warnings tests above) # So we create a fresh request to test override properly fresh_request = ToolCallRequest( - tool_call=tool_call, tool=add, state=state, runtime=runtime - ) # type: ignore[arg-type] + tool_call=tool_call, tool=add, runtime=tool_runtime + ) fresh_new_request = fresh_request.override(tool_call=new_tool_call) # Original request should be unchanged @@ -1378,4 +1396,4 @@ def test_tool_call_request_is_frozen() -> None: assert fresh_new_request.tool_call["name"] == "multiply" assert fresh_new_request.tool == add # Other fields should remain the same assert fresh_new_request.state == state - assert fresh_new_request.runtime is None + assert fresh_new_request.runtime is tool_runtime diff --git a/libs/prebuilt/tests/test_tool_node.py b/libs/prebuilt/tests/test_tool_node.py index c6e143f4f..eae3fa06e 100644 --- a/libs/prebuilt/tests/test_tool_node.py +++ b/libs/prebuilt/tests/test_tool_node.py @@ -1615,18 +1615,28 @@ def test_tool_node_stream_writer() -> None: def test_tool_call_request_setattr_deprecation_warning(): """Test that ToolCallRequest raises a deprecation warning on direct attribute modification.""" import warnings + from unittest.mock import Mock - from langgraph.prebuilt.tool_node import ToolCallRequest + from langgraph.prebuilt.tool_node import ToolCallRequest, ToolRuntime # Create a mock ToolCall tool_call = {"name": "test", "args": {"a": 1}, "id": "call_1", "type": "tool_call"} + # Create a ToolRuntime + tool_runtime = ToolRuntime( + state={"messages": []}, + config={}, + context=None, + store=None, + stream_writer=Mock(), + tool_call_id="call_1", + ) + # Create a ToolCallRequest request = ToolCallRequest( tool_call=tool_call, tool=None, - state={"messages": []}, - runtime=None, + runtime=tool_runtime, ) # Test 1: Direct attribute assignment should raise deprecation warning but still work @@ -1667,8 +1677,7 @@ def test_tool_call_request_setattr_deprecation_warning(): ToolCallRequest( tool_call=tool_call, tool=None, - state={"messages": []}, - runtime=None, + runtime=tool_runtime, ) # Verify no warning was raised during initialization assert len(w) == 0