mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-14 12:15:25 +02:00
Baihao WangandGitHub
5965d72ff7
**Problem.** `interrupt(response_schema=...)` checks a human's answer when the graph resumes. If the answer doesn't match the schema, the resume fails, nothing is saved, and the human can answer again (#8886). That works when the interrupt runs in a graph node, but not when it runs inside a tool: either the tool calls `interrupt()` itself, or it starts a graph that does, such as a subagent called as a tool. There, the resume returns normally and the interrupt disappears. **Why.** `interrupt()` raises a `pydantic.ValidationError` for an invalid answer. When one comes out of a tool, `ToolNode` assumes the model passed bad arguments, so it hands the error to the model ("Error invoking tool … Please fix the error and try again") and the run carries on. The human never sees the error, the model is blamed for it, and the interrupt is no longer pending. **Fix.** `interrupt()` still raises the same `ValidationError`, but marks it so a new helper, `langgraph.errors.is_invalid_resume(error)`, can recognize it. `ToolNode` re-raises these errors instead of handing them to the model, whatever `handle_tool_errors` is set to, so a bad answer fails the resume the same way inside a tool as outside one: nothing is saved and the interrupt can be answered again. Every other `ValidationError` from a tool is handled as before. **Also fixed, in the same error handling.** With `wrap_tool_call` set and a `handle_tool_errors` other than the default (for example `True`), `ToolNode` turned any exception from the wrapper into an error `ToolMessage`, including `GraphInterrupt`. So an interrupt in that setup never paused: the model got `Error: GraphInterrupt(...)` instead. Interrupts now propagate there too, as they already do without a wrapper. `create_agent` uses the default handler, so it was not affected. **Compatibility.** Not a breaking change: the error type is unchanged. Worth a careful look: a wrapped `ToolNode` with a non-default `handle_tool_errors` now pauses on interrupts instead of returning them to the model as errors. With an older `langgraph` that doesn't have `is_invalid_resume`, `ToolNode` keeps today's behavior.
2569 lines
81 KiB
Python
2569 lines
81 KiB
Python
import contextlib
|
|
import dataclasses
|
|
import json
|
|
import sys
|
|
import warnings
|
|
from functools import partial
|
|
from typing import (
|
|
Annotated,
|
|
Any,
|
|
NoReturn,
|
|
TypeVar,
|
|
)
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
from langchain_core.messages import (
|
|
AIMessage,
|
|
AnyMessage,
|
|
HumanMessage,
|
|
RemoveMessage,
|
|
ToolCall,
|
|
ToolMessage,
|
|
)
|
|
from langchain_core.runnables.config import RunnableConfig
|
|
from langchain_core.tools import BaseTool, InjectedToolArg, ToolException
|
|
from langchain_core.tools import tool as dec_tool
|
|
from langchain_core.tools.base import InjectedToolCallId
|
|
from langgraph.checkpoint.base import BaseCheckpointSaver
|
|
from langgraph.config import get_stream_writer
|
|
from langgraph.errors import GraphBubbleUp, GraphInterrupt
|
|
from langgraph.graph import START, MessagesState, StateGraph
|
|
from langgraph.graph.message import REMOVE_ALL_MESSAGES, add_messages
|
|
from langgraph.runtime import ExecutionInfo, ServerInfo
|
|
from langgraph.store.base import BaseStore
|
|
from langgraph.store.memory import InMemoryStore
|
|
from langgraph.types import Command, Send, interrupt
|
|
from pydantic import BaseModel, ValidationError
|
|
from pydantic.v1 import BaseModel as BaseModelV1
|
|
from typing_extensions import TypedDict
|
|
|
|
from langgraph.prebuilt import (
|
|
InjectedState,
|
|
InjectedStore,
|
|
ToolNode,
|
|
)
|
|
from langgraph.prebuilt.tool_node import (
|
|
TOOL_CALL_ERROR_TEMPLATE,
|
|
ToolCallRequest,
|
|
ToolInvocationError,
|
|
ToolRuntime,
|
|
tools_condition,
|
|
)
|
|
|
|
from .messages import _AnyIdHumanMessage, _AnyIdToolMessage
|
|
from .model import FakeToolCallingModel
|
|
|
|
pytestmark = pytest.mark.anyio
|
|
|
|
|
|
def _create_mock_runtime(store: BaseStore | None = None) -> Mock:
|
|
"""Create a mock Runtime object for testing ToolNode outside of graph context.
|
|
|
|
This helper is needed because ToolNode._func expects a Runtime parameter
|
|
which is injected by RunnableCallable from config["configurable"]["__pregel_runtime"].
|
|
When testing ToolNode directly (outside a graph), we need to provide this manually.
|
|
"""
|
|
|
|
mock_runtime = Mock()
|
|
mock_runtime.store = store
|
|
mock_runtime.context = None
|
|
mock_runtime.stream_writer = lambda *args, **kwargs: None
|
|
mock_runtime.execution_info = ExecutionInfo(
|
|
checkpoint_id="test-cp", checkpoint_ns="", task_id="test-task"
|
|
)
|
|
mock_runtime.server_info = None
|
|
return mock_runtime
|
|
|
|
|
|
def _create_config_with_runtime(store: BaseStore | None = None) -> RunnableConfig:
|
|
"""Create a RunnableConfig with mock Runtime for testing ToolNode.
|
|
|
|
Returns:
|
|
RunnableConfig with __pregel_runtime in configurable dict.
|
|
"""
|
|
return {"configurable": {"__pregel_runtime": _create_mock_runtime(store)}}
|
|
|
|
|
|
def tool1(some_val: int, some_other_val: str) -> str:
|
|
"""Tool 1 docstring."""
|
|
if some_val == 0:
|
|
msg = "Test error"
|
|
raise ValueError(msg)
|
|
return f"{some_val} - {some_other_val}"
|
|
|
|
|
|
async def tool2(some_val: int, some_other_val: str) -> str:
|
|
"""Tool 2 docstring."""
|
|
if some_val == 0:
|
|
msg = "Test error"
|
|
raise ToolException(msg)
|
|
return f"tool2: {some_val} - {some_other_val}"
|
|
|
|
|
|
async def tool3(some_val: int, some_other_val: str) -> str:
|
|
"""Tool 3 docstring."""
|
|
return [
|
|
{"key_1": some_val, "key_2": "foo"},
|
|
{"key_1": some_other_val, "key_2": "baz"},
|
|
]
|
|
|
|
|
|
async def tool4(some_val: int, some_other_val: str) -> str:
|
|
"""Tool 4 docstring."""
|
|
return [
|
|
{"type": "image_url", "image_url": {"url": "abdc"}},
|
|
]
|
|
|
|
|
|
@dec_tool
|
|
def tool5(some_val: int) -> NoReturn:
|
|
"""Tool 5 docstring."""
|
|
msg = "Test error"
|
|
raise ToolException(msg)
|
|
|
|
|
|
tool5.handle_tool_error = "foo"
|
|
|
|
|
|
async def test_tool_node() -> None:
|
|
"""Test tool node."""
|
|
result = ToolNode([tool1]).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool1",
|
|
"args": {"some_val": 1, "some_other_val": "foo"},
|
|
"id": "some 0",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
tool_message: ToolMessage = result["messages"][-1]
|
|
assert tool_message.type == "tool"
|
|
assert tool_message.content == "1 - foo"
|
|
assert tool_message.tool_call_id == "some 0"
|
|
|
|
result2 = await ToolNode([tool2]).ainvoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool2",
|
|
"args": {"some_val": 2, "some_other_val": "bar"},
|
|
"id": "some 1",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
tool_message: ToolMessage = result2["messages"][-1]
|
|
assert tool_message.type == "tool"
|
|
assert tool_message.content == "tool2: 2 - bar"
|
|
|
|
# list of dicts tool content
|
|
result3 = await ToolNode([tool3]).ainvoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool3",
|
|
"args": {"some_val": 2, "some_other_val": "bar"},
|
|
"id": "some 2",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
tool_message: ToolMessage = result3["messages"][-1]
|
|
assert tool_message.type == "tool"
|
|
assert (
|
|
tool_message.content
|
|
== '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]'
|
|
)
|
|
assert tool_message.tool_call_id == "some 2"
|
|
|
|
# list of content blocks tool content
|
|
result4 = await ToolNode([tool4]).ainvoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool4",
|
|
"args": {"some_val": 2, "some_other_val": "bar"},
|
|
"id": "some 3",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
tool_message: ToolMessage = result4["messages"][-1]
|
|
assert tool_message.type == "tool"
|
|
assert tool_message.content == [{"type": "image_url", "image_url": {"url": "abdc"}}]
|
|
assert tool_message.tool_call_id == "some 3"
|
|
|
|
|
|
async def test_tool_node_tool_call_input() -> None:
|
|
# Single tool call
|
|
tool_call_1 = {
|
|
"name": "tool1",
|
|
"args": {"some_val": 1, "some_other_val": "foo"},
|
|
"id": "some 0",
|
|
"type": "tool_call",
|
|
}
|
|
result = ToolNode([tool1]).invoke(
|
|
[tool_call_1], config=_create_config_with_runtime()
|
|
)
|
|
assert result["messages"] == [
|
|
ToolMessage(content="1 - foo", tool_call_id="some 0", name="tool1"),
|
|
]
|
|
|
|
# Multiple tool calls
|
|
tool_call_2 = {
|
|
"name": "tool1",
|
|
"args": {"some_val": 2, "some_other_val": "bar"},
|
|
"id": "some 1",
|
|
"type": "tool_call",
|
|
}
|
|
result = ToolNode([tool1]).invoke(
|
|
[tool_call_1, tool_call_2], config=_create_config_with_runtime()
|
|
)
|
|
assert result["messages"] == [
|
|
ToolMessage(content="1 - foo", tool_call_id="some 0", name="tool1"),
|
|
ToolMessage(content="2 - bar", tool_call_id="some 1", name="tool1"),
|
|
]
|
|
|
|
# Test with unknown tool
|
|
tool_call_3 = tool_call_1.copy()
|
|
tool_call_3["name"] = "tool2"
|
|
result = ToolNode([tool1]).invoke(
|
|
[tool_call_1, tool_call_3], config=_create_config_with_runtime()
|
|
)
|
|
assert result["messages"] == [
|
|
ToolMessage(content="1 - foo", tool_call_id="some 0", name="tool1"),
|
|
ToolMessage(
|
|
content="Error: tool2 is not a valid tool, try one of [tool1].",
|
|
name="tool2",
|
|
tool_call_id="some 0",
|
|
status="error",
|
|
),
|
|
]
|
|
|
|
|
|
def test_tool_node_error_handling_default_invocation() -> None:
|
|
tn = ToolNode([tool1])
|
|
result = tn.invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool1",
|
|
"args": {"invalid": 0, "args": "foo"},
|
|
"id": "some id",
|
|
},
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
assert all(m.type == "tool" for m in result["messages"])
|
|
assert all(m.status == "error" for m in result["messages"])
|
|
assert (
|
|
"Error invoking tool 'tool1' with kwargs {'invalid': 0, 'args': 'foo'} with error:\n"
|
|
in result["messages"][0].content
|
|
)
|
|
|
|
|
|
def test_tool_node_error_handling_default_exception() -> None:
|
|
tn = ToolNode([tool1])
|
|
with pytest.raises(ValueError):
|
|
tn.invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool1",
|
|
"args": {"some_val": 0, "some_other_val": "foo"},
|
|
"id": "some id",
|
|
},
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
|
|
async def test_tool_node_error_handling() -> None:
|
|
def handle_all(e: ValueError | ToolException | ToolInvocationError):
|
|
return TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
|
|
|
|
# test catching all exceptions, via:
|
|
# - handle_tool_errors = True
|
|
# - passing a tuple of all exceptions
|
|
# - passing a callable with all exceptions in the signature
|
|
for handle_tool_errors in (
|
|
True,
|
|
(ValueError, ToolException, ToolInvocationError),
|
|
handle_all,
|
|
):
|
|
result_error = await ToolNode(
|
|
[tool1, tool2, tool3], handle_tool_errors=handle_tool_errors
|
|
).ainvoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool1",
|
|
"args": {"some_val": 0, "some_other_val": "foo"},
|
|
"id": "some id",
|
|
},
|
|
{
|
|
"name": "tool2",
|
|
"args": {"some_val": 0, "some_other_val": "bar"},
|
|
"id": "some other id",
|
|
},
|
|
{
|
|
"name": "tool3",
|
|
"args": {"some_val": 0},
|
|
"id": "another id",
|
|
},
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
assert all(m.type == "tool" for m in result_error["messages"])
|
|
assert all(m.status == "error" for m in result_error["messages"])
|
|
assert (
|
|
result_error["messages"][0].content
|
|
== f"Error: {ValueError('Test error')!r}\n Please fix your mistakes."
|
|
)
|
|
assert (
|
|
result_error["messages"][1].content
|
|
== f"Error: {ToolException('Test error')!r}\n Please fix your mistakes."
|
|
)
|
|
# Check that the validation error contains the field name
|
|
assert "some_other_val" in result_error["messages"][2].content
|
|
|
|
assert result_error["messages"][0].tool_call_id == "some id"
|
|
assert result_error["messages"][1].tool_call_id == "some other id"
|
|
assert result_error["messages"][2].tool_call_id == "another id"
|
|
|
|
|
|
async def test_tool_node_error_handling_callable() -> None:
|
|
def handle_value_error(e: ValueError) -> str:
|
|
return "Value error"
|
|
|
|
def handle_tool_exception(e: ToolException) -> str:
|
|
return "Tool exception"
|
|
|
|
for handle_tool_errors in ("Value error", handle_value_error):
|
|
result_error = await ToolNode(
|
|
[tool1], handle_tool_errors=handle_tool_errors
|
|
).ainvoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool1",
|
|
"args": {"some_val": 0, "some_other_val": "foo"},
|
|
"id": "some id",
|
|
},
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
tool_message: ToolMessage = result_error["messages"][-1]
|
|
assert tool_message.type == "tool"
|
|
assert tool_message.status == "error"
|
|
assert tool_message.content == "Value error"
|
|
|
|
# test raising for an unhandled exception, via:
|
|
# - passing a tuple of all exceptions
|
|
# - passing a callable with all exceptions in the signature
|
|
for handle_tool_errors in ((ValueError,), handle_value_error):
|
|
with pytest.raises(ToolException) as exc_info:
|
|
await ToolNode(
|
|
[tool1, tool2], handle_tool_errors=handle_tool_errors
|
|
).ainvoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool1",
|
|
"args": {"some_val": 0, "some_other_val": "foo"},
|
|
"id": "some id",
|
|
},
|
|
{
|
|
"name": "tool2",
|
|
"args": {"some_val": 0, "some_other_val": "bar"},
|
|
"id": "some other id",
|
|
},
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert str(exc_info.value) == "Test error"
|
|
|
|
for handle_tool_errors in ((ToolException,), handle_tool_exception):
|
|
with pytest.raises(ValueError) as exc_info:
|
|
await ToolNode(
|
|
[tool1, tool2], handle_tool_errors=handle_tool_errors
|
|
).ainvoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool1",
|
|
"args": {"some_val": 0, "some_other_val": "foo"},
|
|
"id": "some id",
|
|
},
|
|
{
|
|
"name": "tool2",
|
|
"args": {"some_val": 0, "some_other_val": "bar"},
|
|
"id": "some other id",
|
|
},
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert str(exc_info.value) == "Test error"
|
|
|
|
|
|
async def test_tool_node_handle_tool_errors_false() -> None:
|
|
with pytest.raises(ValueError) as exc_info:
|
|
ToolNode([tool1], handle_tool_errors=False).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool1",
|
|
"args": {"some_val": 0, "some_other_val": "foo"},
|
|
"id": "some id",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
assert str(exc_info.value) == "Test error"
|
|
|
|
with pytest.raises(ToolException):
|
|
await ToolNode([tool2], handle_tool_errors=False).ainvoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool2",
|
|
"args": {"some_val": 0, "some_other_val": "bar"},
|
|
"id": "some id",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
assert str(exc_info.value) == "Test error"
|
|
|
|
# test validation errors get raised if handle_tool_errors is False
|
|
with pytest.raises(ToolInvocationError):
|
|
ToolNode([tool1], handle_tool_errors=False).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool1",
|
|
"args": {"some_val": 0},
|
|
"id": "some id",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
|
|
def test_tool_node_individual_tool_error_handling() -> None:
|
|
# test error handling on individual tools (and that it overrides overall error handling!)
|
|
result_individual_tool_error_handler = ToolNode(
|
|
[tool5], handle_tool_errors="bar"
|
|
).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool5",
|
|
"args": {"some_val": 0},
|
|
"id": "some 0",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
tool_message: ToolMessage = result_individual_tool_error_handler["messages"][-1]
|
|
assert tool_message.type == "tool"
|
|
assert tool_message.status == "error"
|
|
assert tool_message.content == "foo"
|
|
assert tool_message.tool_call_id == "some 0"
|
|
|
|
|
|
def test_tool_node_incorrect_tool_name() -> None:
|
|
result_incorrect_name = ToolNode([tool1, tool2]).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool3",
|
|
"args": {"some_val": 1, "some_other_val": "foo"},
|
|
"id": "some 0",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
tool_message: ToolMessage = result_incorrect_name["messages"][-1]
|
|
assert tool_message.type == "tool"
|
|
assert tool_message.status == "error"
|
|
assert (
|
|
tool_message.content
|
|
== "Error: tool3 is not a valid tool, try one of [tool1, tool2]."
|
|
)
|
|
assert tool_message.tool_call_id == "some 0"
|
|
|
|
|
|
def test_tool_node_node_interrupt() -> None:
|
|
def tool_interrupt(some_val: int) -> None:
|
|
"""Tool docstring."""
|
|
msg = "foo"
|
|
raise GraphBubbleUp(msg)
|
|
|
|
def handle(e: GraphInterrupt) -> str:
|
|
return "handled"
|
|
|
|
for handle_tool_errors in (True, (GraphBubbleUp,), "handled", handle, False):
|
|
node = ToolNode([tool_interrupt], handle_tool_errors=handle_tool_errors)
|
|
with pytest.raises(GraphBubbleUp) as exc_info:
|
|
node.invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"hi?",
|
|
tool_calls=[
|
|
{
|
|
"name": "tool_interrupt",
|
|
"args": {"some_val": 0},
|
|
"id": "some 0",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert exc_info.value == "foo"
|
|
|
|
|
|
class _Approval(BaseModel):
|
|
approved: bool
|
|
|
|
|
|
class _AskState(TypedDict, total=False):
|
|
answer: str
|
|
|
|
|
|
def _approval_graph():
|
|
"""A graph that asks a human for approval with a typed interrupt."""
|
|
|
|
def ask(state: _AskState) -> _AskState:
|
|
approval = interrupt("Approve?", response_schema=_Approval)
|
|
return {"answer": f"approved={approval.approved}"}
|
|
|
|
return StateGraph(_AskState).add_node("ask", ask).add_edge(START, "ask").compile()
|
|
|
|
|
|
def _ask_human_call() -> dict[str, list[AnyMessage]]:
|
|
call = ToolCall(name="ask_human", args={}, id="call_1")
|
|
return {"messages": [AIMessage("", tool_calls=[call])]}
|
|
|
|
|
|
def _handle_any(e): # no annotation: handles every error
|
|
return "handled"
|
|
|
|
|
|
# A bad answer to an interrupt must fail the run whatever `handle_tool_errors` is,
|
|
# including settings that cover `ValidationError` (a `ValueError`), with or without
|
|
# a wrapper. `create_agent` always runs tools through a wrapper (its middleware).
|
|
_TOOL_NODES = pytest.mark.parametrize(
|
|
("wrapped", "handle_tool_errors"),
|
|
[
|
|
(wrapped, handler)
|
|
for wrapped in (False, True)
|
|
for handler in (None, True, (ValueError,), _handle_any)
|
|
],
|
|
ids=[
|
|
f"{wrapped}-{handler}"
|
|
for wrapped in ("plain", "wrapped")
|
|
for handler in ("default", "handle_true", "handle_value_error", "untyped")
|
|
],
|
|
)
|
|
# The interrupt either runs in a graph the tool starts (a subagent) or in the tool.
|
|
_SHAPE = pytest.mark.parametrize("nested", [True, False], ids=["nested", "direct"])
|
|
|
|
|
|
@_TOOL_NODES
|
|
@_SHAPE
|
|
def test_tool_node_reraises_invalid_resume(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
wrapped: bool,
|
|
nested: bool,
|
|
handle_tool_errors: Any,
|
|
) -> None:
|
|
asker = _approval_graph()
|
|
|
|
@dec_tool
|
|
def ask_human() -> str:
|
|
"""Ask a human for approval."""
|
|
if nested:
|
|
return asker.invoke({})["answer"]
|
|
approval = interrupt("Approve?", response_schema=_Approval)
|
|
return f"approved={approval.approved}"
|
|
|
|
def pass_through(request, handler):
|
|
return handler(request)
|
|
|
|
errors = (
|
|
{} if handle_tool_errors is None else {"handle_tool_errors": handle_tool_errors}
|
|
)
|
|
graph = (
|
|
StateGraph(MessagesState)
|
|
.add_node(
|
|
"tools",
|
|
ToolNode(
|
|
[ask_human], wrap_tool_call=pass_through if wrapped else None, **errors
|
|
),
|
|
)
|
|
.add_edge(START, "tools")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
|
[pending] = graph.invoke(_ask_human_call(), config)["__interrupt__"]
|
|
|
|
# A bad answer isn't a bad tool argument: the run fails without saving, so
|
|
# the same interrupt can be answered again.
|
|
with pytest.raises(ValidationError, match="approved"):
|
|
graph.invoke(Command(resume={pending.id: {"approved": "maybe"}}), config)
|
|
assert [i.id for i in graph.get_state(config).interrupts] == [pending.id]
|
|
|
|
result = graph.invoke(Command(resume={pending.id: {"approved": True}}), config)
|
|
assert result["messages"][-1].content == "approved=True"
|
|
|
|
|
|
@_TOOL_NODES
|
|
@_SHAPE
|
|
async def test_tool_node_reraises_invalid_resume_async(
|
|
async_checkpointer: BaseCheckpointSaver,
|
|
wrapped: bool,
|
|
nested: bool,
|
|
handle_tool_errors: Any,
|
|
) -> None:
|
|
asker = _approval_graph()
|
|
|
|
@dec_tool
|
|
async def ask_human() -> str:
|
|
"""Ask a human for approval."""
|
|
if nested:
|
|
return (await asker.ainvoke({}))["answer"]
|
|
approval = interrupt("Approve?", response_schema=_Approval)
|
|
return f"approved={approval.approved}"
|
|
|
|
async def pass_through(request, handler):
|
|
return await handler(request)
|
|
|
|
errors = (
|
|
{} if handle_tool_errors is None else {"handle_tool_errors": handle_tool_errors}
|
|
)
|
|
graph = (
|
|
StateGraph(MessagesState)
|
|
.add_node(
|
|
"tools",
|
|
ToolNode(
|
|
[ask_human], awrap_tool_call=pass_through if wrapped else None, **errors
|
|
),
|
|
)
|
|
.add_edge(START, "tools")
|
|
.compile(checkpointer=async_checkpointer)
|
|
)
|
|
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
|
[pending] = (await graph.ainvoke(_ask_human_call(), config))["__interrupt__"]
|
|
|
|
with pytest.raises(ValidationError, match="approved"):
|
|
await graph.ainvoke(Command(resume={pending.id: {"approved": "maybe"}}), config)
|
|
state = await graph.aget_state(config)
|
|
assert [i.id for i in state.interrupts] == [pending.id]
|
|
|
|
resume = Command(resume={pending.id: {"approved": True}})
|
|
result = await graph.ainvoke(resume, config)
|
|
assert result["messages"][-1].content == "approved=True"
|
|
|
|
|
|
@pytest.mark.parametrize("input_type", ["dict", "tool_calls"])
|
|
async def test_tool_node_command(input_type: str) -> None:
|
|
|
|
@dec_tool
|
|
def transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
|
|
"""Transfer to Bob"""
|
|
return Command(
|
|
update={
|
|
"messages": [
|
|
ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
|
|
]
|
|
},
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
@dec_tool
|
|
async def async_transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
|
|
"""Transfer to Bob"""
|
|
return Command(
|
|
update={
|
|
"messages": [
|
|
ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
|
|
]
|
|
},
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
class CustomToolSchema(BaseModel):
|
|
tool_call_id: Annotated[str, InjectedToolCallId]
|
|
|
|
class MyCustomTool(BaseTool):
|
|
def _run(*args: Any, **kwargs: Any):
|
|
return Command(
|
|
update={
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id=kwargs["tool_call_id"],
|
|
)
|
|
]
|
|
},
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
async def _arun(*args: Any, **kwargs: Any):
|
|
return Command(
|
|
update={
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id=kwargs["tool_call_id"],
|
|
)
|
|
]
|
|
},
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
custom_tool = MyCustomTool(
|
|
name="custom_transfer_to_bob",
|
|
description="Transfer to bob",
|
|
args_schema=CustomToolSchema,
|
|
)
|
|
async_custom_tool = MyCustomTool(
|
|
name="async_custom_transfer_to_bob",
|
|
description="Transfer to bob",
|
|
args_schema=CustomToolSchema,
|
|
)
|
|
|
|
# test mixing regular tools and tools returning commands
|
|
def add(a: int, b: int) -> int:
|
|
"""Add two numbers"""
|
|
return a + b
|
|
|
|
tool_calls = [
|
|
{"args": {"a": 1, "b": 2}, "id": "1", "name": "add", "type": "tool_call"},
|
|
{"args": {}, "id": "2", "name": "transfer_to_bob", "type": "tool_call"},
|
|
]
|
|
if input_type == "dict":
|
|
input_ = {"messages": [AIMessage("", tool_calls=tool_calls)]}
|
|
elif input_type == "tool_calls":
|
|
input_ = tool_calls
|
|
result = ToolNode([add, transfer_to_bob]).invoke(
|
|
input_, config=_create_config_with_runtime()
|
|
)
|
|
|
|
assert result == [
|
|
{
|
|
"messages": [
|
|
ToolMessage(
|
|
content="3",
|
|
tool_call_id="1",
|
|
name="add",
|
|
)
|
|
]
|
|
},
|
|
Command(
|
|
update={
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id="2",
|
|
name="transfer_to_bob",
|
|
)
|
|
]
|
|
},
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
),
|
|
]
|
|
|
|
# test tools returning commands
|
|
|
|
# test sync tools
|
|
for tool in [transfer_to_bob, custom_tool]:
|
|
result = ToolNode([tool]).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"", tool_calls=[{"args": {}, "id": "1", "name": tool.name}]
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert result == [
|
|
Command(
|
|
update={
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id="1",
|
|
name=tool.name,
|
|
)
|
|
]
|
|
},
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
]
|
|
|
|
# test async tools
|
|
for tool in [async_transfer_to_bob, async_custom_tool]:
|
|
result = await ToolNode([tool]).ainvoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"", tool_calls=[{"args": {}, "id": "1", "name": tool.name}]
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert result == [
|
|
Command(
|
|
update={
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id="1",
|
|
name=tool.name,
|
|
)
|
|
]
|
|
},
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
]
|
|
|
|
# test multiple commands
|
|
result = ToolNode([transfer_to_bob, custom_tool]).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[
|
|
{"args": {}, "id": "1", "name": "transfer_to_bob"},
|
|
{"args": {}, "id": "2", "name": "custom_transfer_to_bob"},
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert result == [
|
|
Command(
|
|
update={
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id="1",
|
|
name="transfer_to_bob",
|
|
)
|
|
]
|
|
},
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
),
|
|
Command(
|
|
update={
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id="2",
|
|
name="custom_transfer_to_bob",
|
|
)
|
|
]
|
|
},
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
),
|
|
]
|
|
|
|
# test validation (mismatch between input type and command.update type)
|
|
with pytest.raises(ValueError):
|
|
|
|
@dec_tool
|
|
def list_update_tool(tool_call_id: Annotated[str, InjectedToolCallId]):
|
|
"""My tool"""
|
|
return Command(
|
|
update=[ToolMessage(content="foo", tool_call_id=tool_call_id)]
|
|
)
|
|
|
|
ToolNode([list_update_tool]).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[
|
|
{"args": {}, "id": "1", "name": "list_update_tool"}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
# test validation (missing tool message in the update for current graph)
|
|
with pytest.raises(ValueError):
|
|
|
|
@dec_tool
|
|
def no_update_tool():
|
|
"""My tool"""
|
|
return Command(update={"messages": []})
|
|
|
|
ToolNode([no_update_tool]).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[{"args": {}, "id": "1", "name": "no_update_tool"}],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
# test validation (tool message with a wrong tool call ID)
|
|
with pytest.raises(ValueError):
|
|
|
|
@dec_tool
|
|
def mismatching_tool_call_id_tool():
|
|
"""My tool"""
|
|
return Command(
|
|
update={"messages": [ToolMessage(content="foo", tool_call_id="2")]}
|
|
)
|
|
|
|
ToolNode([mismatching_tool_call_id_tool]).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[
|
|
{
|
|
"args": {},
|
|
"id": "1",
|
|
"name": "mismatching_tool_call_id_tool",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
# test validation (missing tool message in the update for parent graph is OK)
|
|
@dec_tool
|
|
def node_update_parent_tool():
|
|
"""No update"""
|
|
return Command(update={"messages": []}, graph=Command.PARENT)
|
|
|
|
assert ToolNode([node_update_parent_tool]).invoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[
|
|
{"args": {}, "id": "1", "name": "node_update_parent_tool"}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=_create_config_with_runtime(),
|
|
) == [Command(update={"messages": []}, graph=Command.PARENT)]
|
|
|
|
|
|
async def test_tool_node_command_list_input() -> None:
|
|
|
|
@dec_tool
|
|
def transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
|
|
"""Transfer to Bob"""
|
|
return Command(
|
|
update=[
|
|
ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
|
|
],
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
@dec_tool
|
|
async def async_transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
|
|
"""Transfer to Bob"""
|
|
return Command(
|
|
update=[
|
|
ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
|
|
],
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
class CustomToolSchema(BaseModel):
|
|
tool_call_id: Annotated[str, InjectedToolCallId]
|
|
|
|
class MyCustomTool(BaseTool):
|
|
def _run(*args: Any, **kwargs: Any):
|
|
return Command(
|
|
update=[
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id=kwargs["tool_call_id"],
|
|
)
|
|
],
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
async def _arun(*args: Any, **kwargs: Any):
|
|
return Command(
|
|
update=[
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id=kwargs["tool_call_id"],
|
|
)
|
|
],
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
custom_tool = MyCustomTool(
|
|
name="custom_transfer_to_bob",
|
|
description="Transfer to bob",
|
|
args_schema=CustomToolSchema,
|
|
)
|
|
async_custom_tool = MyCustomTool(
|
|
name="async_custom_transfer_to_bob",
|
|
description="Transfer to bob",
|
|
args_schema=CustomToolSchema,
|
|
)
|
|
|
|
# test mixing regular tools and tools returning commands
|
|
def add(a: int, b: int) -> int:
|
|
"""Add two numbers"""
|
|
return a + b
|
|
|
|
result = ToolNode([add, transfer_to_bob]).invoke(
|
|
[
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[
|
|
{"args": {"a": 1, "b": 2}, "id": "1", "name": "add"},
|
|
{"args": {}, "id": "2", "name": "transfer_to_bob"},
|
|
],
|
|
)
|
|
],
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
assert result == [
|
|
[
|
|
ToolMessage(
|
|
content="3",
|
|
tool_call_id="1",
|
|
name="add",
|
|
)
|
|
],
|
|
Command(
|
|
update=[
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id="2",
|
|
name="transfer_to_bob",
|
|
)
|
|
],
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
),
|
|
]
|
|
|
|
# test tools returning commands
|
|
|
|
# test sync tools
|
|
for tool in [transfer_to_bob, custom_tool]:
|
|
result = ToolNode([tool]).invoke(
|
|
[AIMessage("", tool_calls=[{"args": {}, "id": "1", "name": tool.name}])],
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert result == [
|
|
Command(
|
|
update=[
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id="1",
|
|
name=tool.name,
|
|
)
|
|
],
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
]
|
|
|
|
# test async tools
|
|
for tool in [async_transfer_to_bob, async_custom_tool]:
|
|
result = await ToolNode([tool]).ainvoke(
|
|
[AIMessage("", tool_calls=[{"args": {}, "id": "1", "name": tool.name}])],
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert result == [
|
|
Command(
|
|
update=[
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id="1",
|
|
name=tool.name,
|
|
)
|
|
],
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
)
|
|
]
|
|
|
|
# test multiple commands
|
|
result = ToolNode([transfer_to_bob, custom_tool]).invoke(
|
|
[
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[
|
|
{"args": {}, "id": "1", "name": "transfer_to_bob"},
|
|
{"args": {}, "id": "2", "name": "custom_transfer_to_bob"},
|
|
],
|
|
)
|
|
],
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert result == [
|
|
Command(
|
|
update=[
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id="1",
|
|
name="transfer_to_bob",
|
|
)
|
|
],
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
),
|
|
Command(
|
|
update=[
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
tool_call_id="2",
|
|
name="custom_transfer_to_bob",
|
|
)
|
|
],
|
|
goto="bob",
|
|
graph=Command.PARENT,
|
|
),
|
|
]
|
|
|
|
# test validation (mismatch between input type and command.update type)
|
|
with pytest.raises(ValueError):
|
|
|
|
@dec_tool
|
|
def list_update_tool(tool_call_id: Annotated[str, InjectedToolCallId]):
|
|
"""My tool"""
|
|
return Command(
|
|
update={
|
|
"messages": [ToolMessage(content="foo", tool_call_id=tool_call_id)]
|
|
}
|
|
)
|
|
|
|
ToolNode([list_update_tool]).invoke(
|
|
[
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[{"args": {}, "id": "1", "name": "list_update_tool"}],
|
|
)
|
|
],
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
# test validation (missing tool message in the update for current graph)
|
|
with pytest.raises(ValueError):
|
|
|
|
@dec_tool
|
|
def no_update_tool():
|
|
"""My tool"""
|
|
return Command(update=[])
|
|
|
|
ToolNode([no_update_tool]).invoke(
|
|
[
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[{"args": {}, "id": "1", "name": "no_update_tool"}],
|
|
)
|
|
],
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
# test validation (tool message with a wrong tool call ID)
|
|
with pytest.raises(ValueError):
|
|
|
|
@dec_tool
|
|
def mismatching_tool_call_id_tool():
|
|
"""My tool"""
|
|
return Command(update=[ToolMessage(content="foo", tool_call_id="2")])
|
|
|
|
ToolNode([mismatching_tool_call_id_tool]).invoke(
|
|
[
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[
|
|
{"args": {}, "id": "1", "name": "mismatching_tool_call_id_tool"}
|
|
],
|
|
)
|
|
],
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
# test validation (missing tool message in the update for parent graph is OK)
|
|
@dec_tool
|
|
def node_update_parent_tool():
|
|
"""No update"""
|
|
return Command(update=[], graph=Command.PARENT)
|
|
|
|
assert ToolNode([node_update_parent_tool]).invoke(
|
|
[
|
|
AIMessage(
|
|
"",
|
|
tool_calls=[{"args": {}, "id": "1", "name": "node_update_parent_tool"}],
|
|
)
|
|
],
|
|
config=_create_config_with_runtime(),
|
|
) == [Command(update=[], graph=Command.PARENT)]
|
|
|
|
|
|
def test_tool_node_parent_command_with_send() -> None:
|
|
|
|
@dec_tool
|
|
def transfer_to_alice(tool_call_id: Annotated[str, InjectedToolCallId]):
|
|
"""Transfer to Alice"""
|
|
return Command(
|
|
goto=[
|
|
Send(
|
|
"alice",
|
|
{
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Alice",
|
|
name="transfer_to_alice",
|
|
tool_call_id=tool_call_id,
|
|
)
|
|
]
|
|
},
|
|
)
|
|
],
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
@dec_tool
|
|
def transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
|
|
"""Transfer to Bob"""
|
|
return Command(
|
|
goto=[
|
|
Send(
|
|
"bob",
|
|
{
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
name="transfer_to_bob",
|
|
tool_call_id=tool_call_id,
|
|
)
|
|
]
|
|
},
|
|
)
|
|
],
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
tool_calls = [
|
|
{"args": {}, "id": "1", "name": "transfer_to_alice", "type": "tool_call"},
|
|
{"args": {}, "id": "2", "name": "transfer_to_bob", "type": "tool_call"},
|
|
]
|
|
|
|
result = ToolNode([transfer_to_alice, transfer_to_bob]).invoke(
|
|
[AIMessage("", tool_calls=tool_calls)],
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
assert result == [
|
|
Command(
|
|
goto=[
|
|
Send(
|
|
"alice",
|
|
{
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Alice",
|
|
name="transfer_to_alice",
|
|
tool_call_id="1",
|
|
)
|
|
]
|
|
},
|
|
),
|
|
Send(
|
|
"bob",
|
|
{
|
|
"messages": [
|
|
ToolMessage(
|
|
content="Transferred to Bob",
|
|
name="transfer_to_bob",
|
|
tool_call_id="2",
|
|
)
|
|
]
|
|
},
|
|
),
|
|
],
|
|
graph=Command.PARENT,
|
|
)
|
|
]
|
|
|
|
|
|
async def test_tool_node_command_remove_all_messages() -> None:
|
|
|
|
@dec_tool
|
|
def remove_all_messages_tool(tool_call_id: Annotated[str, InjectedToolCallId]):
|
|
"""A tool that removes all messages."""
|
|
return Command(update={"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)]})
|
|
|
|
tool_node = ToolNode([remove_all_messages_tool])
|
|
tool_call = {
|
|
"name": "remove_all_messages_tool",
|
|
"args": {},
|
|
"id": "tool_call_123",
|
|
}
|
|
result = await tool_node.ainvoke(
|
|
{"messages": [AIMessage(content="", tool_calls=[tool_call])]},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
command = result[0]
|
|
assert isinstance(command, Command)
|
|
assert command.update == {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)]}
|
|
|
|
|
|
class _InjectStateSchema(TypedDict):
|
|
messages: list
|
|
foo: str
|
|
|
|
|
|
class _InjectedStatePydanticV2Schema(BaseModel):
|
|
messages: list
|
|
foo: str
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class _InjectedStateDataclassSchema:
|
|
messages: list
|
|
foo: str
|
|
|
|
|
|
_INJECTED_STATE_SCHEMAS = [
|
|
_InjectStateSchema,
|
|
_InjectedStatePydanticV2Schema,
|
|
_InjectedStateDataclassSchema,
|
|
]
|
|
|
|
if sys.version_info < (3, 14):
|
|
|
|
class _InjectedStatePydanticSchema(BaseModelV1):
|
|
messages: list
|
|
foo: str
|
|
|
|
_INJECTED_STATE_SCHEMAS.append(_InjectedStatePydanticSchema)
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
@pytest.mark.parametrize("schema_", _INJECTED_STATE_SCHEMAS)
|
|
def test_tool_node_inject_state(schema_: type[T]) -> None:
|
|
def tool1(some_val: int, state: Annotated[T, InjectedState]) -> str:
|
|
"""Tool 1 docstring."""
|
|
if isinstance(state, dict):
|
|
return state["foo"]
|
|
return state.foo
|
|
|
|
def tool2(some_val: int, state: Annotated[T, InjectedState()]) -> str:
|
|
"""Tool 2 docstring."""
|
|
if isinstance(state, dict):
|
|
return state["foo"]
|
|
return state.foo
|
|
|
|
def tool3(
|
|
some_val: int,
|
|
foo: Annotated[str, InjectedState("foo")],
|
|
msgs: Annotated[list[AnyMessage], InjectedState("messages")],
|
|
) -> str:
|
|
"""Tool 1 docstring."""
|
|
return foo
|
|
|
|
def tool4(
|
|
some_val: int, msgs: Annotated[list[AnyMessage], InjectedState("messages")]
|
|
) -> str:
|
|
"""Tool 1 docstring."""
|
|
return msgs[0].content
|
|
|
|
node = ToolNode([tool1, tool2, tool3, tool4], handle_tool_errors=True)
|
|
for tool_name in ("tool1", "tool2", "tool3"):
|
|
tool_call = {
|
|
"name": tool_name,
|
|
"args": {"some_val": 1},
|
|
"id": "some 0",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("hi?", tool_calls=[tool_call])
|
|
result = node.invoke(
|
|
schema_(messages=[msg], foo="bar"), config=_create_config_with_runtime()
|
|
)
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.content == "bar", f"Failed for tool={tool_name}"
|
|
|
|
if tool_name == "tool3":
|
|
failure_input = None
|
|
with contextlib.suppress(Exception):
|
|
failure_input = schema_(messages=[msg], notfoo="bar")
|
|
if failure_input is not None:
|
|
with pytest.raises(KeyError):
|
|
node.invoke(failure_input, config=_create_config_with_runtime())
|
|
|
|
with pytest.raises(ValueError):
|
|
node.invoke([msg], config=_create_config_with_runtime())
|
|
else:
|
|
failure_input = None
|
|
try:
|
|
failure_input = schema_(messages=[msg], notfoo="bar")
|
|
except Exception:
|
|
# We'd get a validation error from pydantic state and wouldn't make it to the node
|
|
# anyway
|
|
pass
|
|
if failure_input is not None:
|
|
messages_ = node.invoke(
|
|
failure_input, config=_create_config_with_runtime()
|
|
)
|
|
tool_message = messages_["messages"][-1]
|
|
assert "KeyError" in tool_message.content
|
|
tool_message = node.invoke([msg], config=_create_config_with_runtime())[
|
|
-1
|
|
]
|
|
assert "KeyError" in tool_message.content
|
|
|
|
tool_call = {
|
|
"name": "tool4",
|
|
"args": {"some_val": 1},
|
|
"id": "some 0",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("hi?", tool_calls=[tool_call])
|
|
result = node.invoke(
|
|
schema_(messages=[msg], foo=""), config=_create_config_with_runtime()
|
|
)
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.content == "hi?"
|
|
|
|
result = node.invoke([msg], config=_create_config_with_runtime())
|
|
tool_message = result[-1]
|
|
assert tool_message.content == "hi?"
|
|
|
|
|
|
def test_tool_node_inject_store() -> None:
|
|
store = InMemoryStore()
|
|
namespace = ("test",)
|
|
|
|
def tool1(some_val: int, store: Annotated[BaseStore, InjectedStore()]) -> str:
|
|
"""Tool 1 docstring."""
|
|
store_val = store.get(namespace, "test_key").value["foo"]
|
|
return f"Some val: {some_val}, store val: {store_val}"
|
|
|
|
def tool2(some_val: int, store: Annotated[BaseStore, InjectedStore()]) -> str:
|
|
"""Tool 2 docstring."""
|
|
store_val = store.get(namespace, "test_key").value["foo"]
|
|
return f"Some val: {some_val}, store val: {store_val}"
|
|
|
|
def tool3(
|
|
some_val: int,
|
|
bar: Annotated[str, InjectedState("bar")],
|
|
store: Annotated[BaseStore, InjectedStore()],
|
|
) -> str:
|
|
"""Tool 3 docstring."""
|
|
store_val = store.get(namespace, "test_key").value["foo"]
|
|
return f"Some val: {some_val}, store val: {store_val}, state val: {bar}"
|
|
|
|
node = ToolNode([tool1, tool2, tool3], handle_tool_errors=True)
|
|
store.put(namespace, "test_key", {"foo": "bar"})
|
|
|
|
class State(MessagesState):
|
|
bar: str
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("tools", node)
|
|
builder.add_edge(START, "tools")
|
|
graph = builder.compile(store=store)
|
|
|
|
for tool_name in ("tool1", "tool2"):
|
|
tool_call = {
|
|
"name": tool_name,
|
|
"args": {"some_val": 1},
|
|
"id": "some 0",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("hi?", tool_calls=[tool_call])
|
|
node_result = node.invoke(
|
|
{"messages": [msg]}, config=_create_config_with_runtime(store=store)
|
|
)
|
|
graph_result = graph.invoke({"messages": [msg]})
|
|
for result in (node_result, graph_result):
|
|
result["messages"][-1]
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.content == "Some val: 1, store val: bar", (
|
|
f"Failed for tool={tool_name}"
|
|
)
|
|
|
|
tool_call = {
|
|
"name": "tool3",
|
|
"args": {"some_val": 1},
|
|
"id": "some 0",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("hi?", tool_calls=[tool_call])
|
|
node_result = node.invoke(
|
|
{"messages": [msg], "bar": "baz"},
|
|
config=_create_config_with_runtime(store=store),
|
|
)
|
|
graph_result = graph.invoke({"messages": [msg], "bar": "baz"})
|
|
for result in (node_result, graph_result):
|
|
result["messages"][-1]
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.content == "Some val: 1, store val: bar, state val: baz", (
|
|
f"Failed for tool={tool_name}"
|
|
)
|
|
|
|
# test injected store without passing store to compiled graph
|
|
failing_graph = builder.compile()
|
|
with pytest.raises(ValueError):
|
|
failing_graph.invoke({"messages": [msg], "bar": "baz"})
|
|
|
|
|
|
def test_tool_node_ensure_utf8() -> None:
|
|
@dec_tool
|
|
def get_day_list(days: list[str]) -> list[str]:
|
|
"""choose days"""
|
|
return days
|
|
|
|
data = ["星期一", "水曜日", "목요일", "Friday"]
|
|
tools = [get_day_list]
|
|
tool_calls = [ToolCall(name=get_day_list.name, args={"days": data}, id="test_id")]
|
|
outputs: list[ToolMessage] = ToolNode(tools).invoke(
|
|
[AIMessage(content="", tool_calls=tool_calls)],
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert outputs[0].content == json.dumps(data, ensure_ascii=False)
|
|
|
|
|
|
def test_tool_node_messages_key() -> None:
|
|
@dec_tool
|
|
def add(a: int, b: int) -> int:
|
|
"""Adds a and b."""
|
|
return a + b
|
|
|
|
model = FakeToolCallingModel(
|
|
tool_calls=[[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")]]
|
|
)
|
|
|
|
class State(TypedDict):
|
|
subgraph_messages: Annotated[list[AnyMessage], add_messages]
|
|
|
|
def call_model(state: State) -> dict[str, Any]:
|
|
response = model.invoke(state["subgraph_messages"])
|
|
model.tool_calls = []
|
|
return {"subgraph_messages": response}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("agent", call_model)
|
|
builder.add_node("tools", ToolNode([add], messages_key="subgraph_messages"))
|
|
builder.add_conditional_edges(
|
|
"agent", partial(tools_condition, messages_key="subgraph_messages")
|
|
)
|
|
builder.add_edge(START, "agent")
|
|
builder.add_edge("tools", "agent")
|
|
|
|
graph = builder.compile()
|
|
result = graph.invoke({"subgraph_messages": [HumanMessage(content="hi")]})
|
|
assert result["subgraph_messages"] == [
|
|
_AnyIdHumanMessage(content="hi"),
|
|
AIMessage(
|
|
content="hi",
|
|
id="0",
|
|
tool_calls=[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")],
|
|
),
|
|
_AnyIdToolMessage(content="3", name=add.name, tool_call_id="test_id"),
|
|
AIMessage(content="hi-hi-3", id="1"),
|
|
]
|
|
|
|
|
|
def test_tool_node_stream_writer() -> None:
|
|
@dec_tool
|
|
def streaming_tool(x: int) -> str:
|
|
"""Do something with writer."""
|
|
my_writer = get_stream_writer()
|
|
for value in ["foo", "bar", "baz"]:
|
|
my_writer({"custom_tool_value": value})
|
|
|
|
return x
|
|
|
|
tool_node = ToolNode([streaming_tool])
|
|
graph = (
|
|
StateGraph(MessagesState)
|
|
.add_node("tools", tool_node)
|
|
.add_edge(START, "tools")
|
|
.compile()
|
|
)
|
|
|
|
tool_call = {
|
|
"name": "streaming_tool",
|
|
"args": {"x": 1},
|
|
"id": "1",
|
|
"type": "tool_call",
|
|
}
|
|
inputs = {
|
|
"messages": [AIMessage("", tool_calls=[tool_call])],
|
|
}
|
|
|
|
assert list(graph.stream(inputs, stream_mode="custom")) == [
|
|
{"custom_tool_value": "foo"},
|
|
{"custom_tool_value": "bar"},
|
|
{"custom_tool_value": "baz"},
|
|
]
|
|
assert list(graph.stream(inputs, stream_mode=["custom", "updates"])) == [
|
|
("custom", {"custom_tool_value": "foo"}),
|
|
("custom", {"custom_tool_value": "bar"}),
|
|
("custom", {"custom_tool_value": "baz"}),
|
|
(
|
|
"updates",
|
|
{
|
|
"tools": {
|
|
"messages": [
|
|
_AnyIdToolMessage(
|
|
content="1",
|
|
name="streaming_tool",
|
|
tool_call_id="1",
|
|
),
|
|
],
|
|
},
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
def test_tool_call_request_setattr_deprecation_warning():
|
|
"""Test that ToolCallRequest raises a deprecation warning on direct attribute modification."""
|
|
|
|
# Create a mock ToolCall
|
|
tool_call = {"name": "test", "args": {"a": 1}, "id": "call_1", "type": "tool_call"}
|
|
|
|
# Create a ToolCallRequest
|
|
request = ToolCallRequest(
|
|
tool_call=tool_call,
|
|
tool=None,
|
|
state={"messages": []},
|
|
runtime=None,
|
|
)
|
|
|
|
# Test 1: Direct attribute assignment should raise deprecation warning but still work
|
|
with pytest.warns(DeprecationWarning, match="deprecated.*override"):
|
|
request.tool_call = {"name": "other", "args": {}, "id": "call_2"}
|
|
|
|
# Verify the attribute was actually modified
|
|
assert request.tool_call == {"name": "other", "args": {}, "id": "call_2"}
|
|
|
|
# Reset for further tests
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
request.tool_call = tool_call
|
|
|
|
# Test 2: override method should work without warnings
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warnings.simplefilter("always")
|
|
new_tool_call = {
|
|
"name": "new_tool",
|
|
"args": {"b": 2},
|
|
"id": "call_3",
|
|
"type": "tool_call",
|
|
}
|
|
new_request = request.override(tool_call=new_tool_call)
|
|
|
|
# Verify no warning was raised
|
|
assert len(w) == 0
|
|
|
|
# Verify original is unchanged
|
|
assert request.tool_call == tool_call
|
|
|
|
# Verify new request has updated values
|
|
assert new_request.tool_call == new_tool_call
|
|
|
|
# Test 3: Initialization should not trigger warning
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warnings.simplefilter("always")
|
|
ToolCallRequest(
|
|
tool_call=tool_call,
|
|
tool=None,
|
|
state={"messages": []},
|
|
runtime=None,
|
|
)
|
|
# Verify no warning was raised during initialization
|
|
assert len(w) == 0
|
|
|
|
|
|
async def test_tool_node_inject_async_all_types_signature_only() -> None:
|
|
"""Test all injection types without @tool decorator."""
|
|
store = InMemoryStore()
|
|
namespace = ("test",)
|
|
store.put(namespace, "test_key", {"store_data": "from_store"})
|
|
|
|
class TestState(TypedDict):
|
|
messages: list
|
|
foo: str
|
|
bar: int
|
|
|
|
async def comprehensive_async_tool(
|
|
x: int,
|
|
whole_state: Annotated[TestState, InjectedState],
|
|
foo_field: Annotated[str, InjectedState("foo")],
|
|
store: Annotated[BaseStore, InjectedStore()],
|
|
runtime: ToolRuntime,
|
|
) -> str:
|
|
"""Async tool that uses all injection types."""
|
|
bar_from_whole = whole_state["bar"]
|
|
foo_value = foo_field
|
|
store_val = store.get(namespace, "test_key").value["store_data"]
|
|
foo_from_runtime = runtime.state["foo"]
|
|
tool_call_id = runtime.tool_call_id
|
|
|
|
return (
|
|
f"x={x}, "
|
|
f"bar_from_whole={bar_from_whole}, "
|
|
f"foo_field={foo_value}, "
|
|
f"store={store_val}, "
|
|
f"foo_from_runtime={foo_from_runtime}, "
|
|
f"tool_call_id={tool_call_id}"
|
|
)
|
|
|
|
node = ToolNode([comprehensive_async_tool], handle_tool_errors=True)
|
|
tool_call = {
|
|
"name": "comprehensive_async_tool",
|
|
"args": {"x": 42},
|
|
"id": "test_call_123",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("hi?", tool_calls=[tool_call])
|
|
|
|
config = _create_config_with_runtime(store=store)
|
|
result = await node.ainvoke(
|
|
{"messages": [msg], "foo": "foo_value", "bar": 99}, config=config
|
|
)
|
|
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.content == (
|
|
"x=42, "
|
|
"bar_from_whole=99, "
|
|
"foo_field=foo_value, "
|
|
"store=from_store, "
|
|
"foo_from_runtime=foo_value, "
|
|
"tool_call_id=test_call_123"
|
|
)
|
|
|
|
|
|
async def test_tool_node_inject_async_all_types_with_decorator() -> None:
|
|
"""Test all injection types with @tool decorator."""
|
|
store = InMemoryStore()
|
|
namespace = ("test",)
|
|
store.put(namespace, "test_key", {"store_data": "from_store"})
|
|
|
|
class TestState(TypedDict):
|
|
messages: list
|
|
foo: str
|
|
bar: int
|
|
|
|
@dec_tool
|
|
async def comprehensive_async_tool(
|
|
x: int,
|
|
whole_state: Annotated[TestState, InjectedState],
|
|
foo_field: Annotated[str, InjectedState("foo")],
|
|
store: Annotated[BaseStore, InjectedStore()],
|
|
runtime: ToolRuntime,
|
|
) -> str:
|
|
"""Async tool that uses all injection types."""
|
|
bar_from_whole = whole_state["bar"]
|
|
foo_value = foo_field
|
|
store_val = store.get(namespace, "test_key").value["store_data"]
|
|
foo_from_runtime = runtime.state["foo"]
|
|
tool_call_id = runtime.tool_call_id
|
|
|
|
return (
|
|
f"x={x}, "
|
|
f"bar_from_whole={bar_from_whole}, "
|
|
f"foo_field={foo_value}, "
|
|
f"store={store_val}, "
|
|
f"foo_from_runtime={foo_from_runtime}, "
|
|
f"tool_call_id={tool_call_id}"
|
|
)
|
|
|
|
node = ToolNode([comprehensive_async_tool], handle_tool_errors=True)
|
|
tool_call = {
|
|
"name": "comprehensive_async_tool",
|
|
"args": {"x": 42},
|
|
"id": "test_call_456",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("hi?", tool_calls=[tool_call])
|
|
|
|
config = _create_config_with_runtime(store=store)
|
|
result = await node.ainvoke(
|
|
{"messages": [msg], "foo": "foo_value", "bar": 99}, config=config
|
|
)
|
|
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.content == (
|
|
"x=42, "
|
|
"bar_from_whole=99, "
|
|
"foo_field=foo_value, "
|
|
"store=from_store, "
|
|
"foo_from_runtime=foo_value, "
|
|
"tool_call_id=test_call_456"
|
|
)
|
|
|
|
|
|
async def test_tool_node_inject_async_all_types_with_schema() -> None:
|
|
"""Test all injection types with explicit schema."""
|
|
store = InMemoryStore()
|
|
namespace = ("test",)
|
|
store.put(namespace, "test_key", {"store_data": "from_store"})
|
|
|
|
class TestState(TypedDict):
|
|
messages: list
|
|
foo: str
|
|
bar: int
|
|
|
|
class ComprehensiveToolSchema(BaseModel):
|
|
model_config = {"arbitrary_types_allowed": True}
|
|
x: int
|
|
whole_state: Annotated[TestState, InjectedState]
|
|
foo_field: Annotated[str, InjectedState("foo")]
|
|
store: Annotated[BaseStore, InjectedStore()]
|
|
runtime: ToolRuntime
|
|
|
|
@dec_tool(args_schema=ComprehensiveToolSchema)
|
|
async def comprehensive_async_tool(
|
|
x: int,
|
|
whole_state: Annotated[TestState, InjectedState],
|
|
foo_field: Annotated[str, InjectedState("foo")],
|
|
store: Annotated[BaseStore, InjectedStore()],
|
|
runtime: ToolRuntime,
|
|
) -> str:
|
|
"""Async tool that uses all injection types."""
|
|
bar_from_whole = whole_state["bar"]
|
|
foo_value = foo_field
|
|
store_val = store.get(namespace, "test_key").value["store_data"]
|
|
foo_from_runtime = runtime.state["foo"]
|
|
tool_call_id = runtime.tool_call_id
|
|
|
|
return (
|
|
f"x={x}, "
|
|
f"bar_from_whole={bar_from_whole}, "
|
|
f"foo_field={foo_value}, "
|
|
f"store={store_val}, "
|
|
f"foo_from_runtime={foo_from_runtime}, "
|
|
f"tool_call_id={tool_call_id}"
|
|
)
|
|
|
|
node = ToolNode([comprehensive_async_tool], handle_tool_errors=True)
|
|
tool_call = {
|
|
"name": "comprehensive_async_tool",
|
|
"args": {"x": 42},
|
|
"id": "test_call_789",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("hi?", tool_calls=[tool_call])
|
|
|
|
config = _create_config_with_runtime(store=store)
|
|
result = await node.ainvoke(
|
|
{"messages": [msg], "foo": "foo_value", "bar": 99}, config=config
|
|
)
|
|
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.content == (
|
|
"x=42, "
|
|
"bar_from_whole=99, "
|
|
"foo_field=foo_value, "
|
|
"store=from_store, "
|
|
"foo_from_runtime=foo_value, "
|
|
"tool_call_id=test_call_789"
|
|
)
|
|
|
|
|
|
async def test_tool_node_tool_runtime_generic() -> None:
|
|
"""Test that ToolRuntime with generic type arguments is correctly injected."""
|
|
|
|
@dataclasses.dataclass
|
|
class MyContext:
|
|
some_info: str
|
|
|
|
@dec_tool
|
|
def get_info(rt: ToolRuntime[MyContext]):
|
|
"""This tool returns info from context."""
|
|
return rt.context.some_info
|
|
|
|
# Create a mock runtime with context
|
|
mock_runtime = _create_mock_runtime()
|
|
mock_runtime.context = MyContext(some_info="test_info")
|
|
|
|
config = {"configurable": {"__pregel_runtime": mock_runtime}}
|
|
|
|
result = await ToolNode([get_info]).ainvoke(
|
|
{
|
|
"messages": [
|
|
AIMessage(
|
|
"call tool",
|
|
tool_calls=[
|
|
{
|
|
"name": "get_info",
|
|
"args": {},
|
|
"id": "call_1",
|
|
}
|
|
],
|
|
)
|
|
]
|
|
},
|
|
config=config,
|
|
)
|
|
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.type == "tool"
|
|
assert tool_message.content == "test_info"
|
|
assert tool_message.tool_call_id == "call_1"
|
|
|
|
|
|
def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call() -> None:
|
|
"""Test that ToolRuntime is injected for dynamically registered tools.
|
|
|
|
Regression test for https://github.com/langchain-ai/langchain/issues/35305.
|
|
When a tool is dynamically provided via wrap_tool_call (not registered at
|
|
ToolNode init time), ToolRuntime should still be injected into the tool.
|
|
"""
|
|
|
|
@dec_tool
|
|
def static_tool(x: int) -> str:
|
|
"""A static tool registered at init."""
|
|
return f"static: {x}"
|
|
|
|
@dec_tool
|
|
def dynamic_tool_with_runtime(x: int, runtime: ToolRuntime) -> str:
|
|
"""A dynamic tool that needs ToolRuntime injection."""
|
|
return f"dynamic: x={x}, tool_call_id={runtime.tool_call_id}"
|
|
|
|
def wrap_tool_call(request, execute):
|
|
"""Middleware that swaps in a dynamic tool."""
|
|
if request.tool_call["name"] == "dynamic_tool_with_runtime":
|
|
# Override tool to the dynamic one (not registered at init)
|
|
new_request = request.override(tool=dynamic_tool_with_runtime)
|
|
return execute(new_request)
|
|
return execute(request)
|
|
|
|
# ToolNode only knows about static_tool at init time
|
|
tool_node = ToolNode(
|
|
[static_tool],
|
|
wrap_tool_call=wrap_tool_call,
|
|
)
|
|
|
|
# Verify the dynamic tool is NOT in the tool node's registered tools
|
|
assert "dynamic_tool_with_runtime" not in tool_node.tools_by_name
|
|
|
|
# Call the dynamic tool
|
|
tool_call = {
|
|
"name": "dynamic_tool_with_runtime",
|
|
"args": {"x": 42},
|
|
"id": "call_dynamic_1",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("", tool_calls=[tool_call])
|
|
result = tool_node.invoke(
|
|
{"messages": [msg]},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
# ToolRuntime should be injected and the tool should execute successfully
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.content == "dynamic: x=42, tool_call_id=call_dynamic_1"
|
|
assert tool_message.tool_call_id == "call_dynamic_1"
|
|
|
|
|
|
async def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call_async() -> None:
|
|
"""Test that ToolRuntime is injected for dynamically registered tools (async).
|
|
|
|
Async version of the regression test for
|
|
https://github.com/langchain-ai/langchain/issues/35305.
|
|
"""
|
|
|
|
@dec_tool
|
|
def static_tool(x: int) -> str:
|
|
"""A static tool registered at init."""
|
|
return f"static: {x}"
|
|
|
|
@dec_tool
|
|
async def dynamic_tool_with_runtime(x: int, runtime: ToolRuntime) -> str:
|
|
"""A dynamic async tool that needs ToolRuntime injection."""
|
|
return f"dynamic: x={x}, tool_call_id={runtime.tool_call_id}"
|
|
|
|
async def awrap_tool_call(request, execute):
|
|
"""Async middleware that swaps in a dynamic tool."""
|
|
if request.tool_call["name"] == "dynamic_tool_with_runtime":
|
|
new_request = request.override(tool=dynamic_tool_with_runtime)
|
|
return await execute(new_request)
|
|
return await execute(request)
|
|
|
|
# ToolNode only knows about static_tool at init time
|
|
tool_node = ToolNode(
|
|
[static_tool],
|
|
awrap_tool_call=awrap_tool_call,
|
|
)
|
|
|
|
# Verify the dynamic tool is NOT in the tool node's registered tools
|
|
assert "dynamic_tool_with_runtime" not in tool_node.tools_by_name
|
|
|
|
# Call the dynamic tool
|
|
tool_call = {
|
|
"name": "dynamic_tool_with_runtime",
|
|
"args": {"x": 42},
|
|
"id": "call_dynamic_2",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("", tool_calls=[tool_call])
|
|
result = await tool_node.ainvoke(
|
|
{"messages": [msg]},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
# ToolRuntime should be injected and the tool should execute successfully
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.content == "dynamic: x=42, tool_call_id=call_dynamic_2"
|
|
assert tool_message.tool_call_id == "call_dynamic_2"
|
|
|
|
|
|
def test_tool_runtime_defaults_tools_to_empty_list() -> None:
|
|
runtime = ToolRuntime(
|
|
state={},
|
|
context=None,
|
|
config={},
|
|
stream_writer=lambda *args, **kwargs: None,
|
|
tool_call_id=None,
|
|
store=None,
|
|
)
|
|
|
|
assert runtime.tools == []
|
|
|
|
|
|
def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
|
|
"""Test that execution_info, server_info, and tools are forwarded from Runtime to ToolRuntime."""
|
|
|
|
exec_info = ExecutionInfo(
|
|
thread_id="t-1",
|
|
checkpoint_id="cp-1",
|
|
checkpoint_ns="",
|
|
task_id="tk-1",
|
|
run_id="r-1",
|
|
)
|
|
server_info = ServerInfo(assistant_id="asst-1", graph_id="graph-1")
|
|
|
|
mock_runtime = Mock()
|
|
mock_runtime.store = None
|
|
mock_runtime.context = None
|
|
mock_runtime.stream_writer = lambda *args, **kwargs: None
|
|
mock_runtime.execution_info = exec_info
|
|
mock_runtime.server_info = server_info
|
|
|
|
captured: dict = {}
|
|
|
|
@dec_tool
|
|
def info_tool(x: int, runtime: ToolRuntime) -> str:
|
|
"""Tool that captures runtime info."""
|
|
captured["execution_info"] = runtime.execution_info
|
|
captured["server_info"] = runtime.server_info
|
|
captured["tools"] = runtime.tools
|
|
return "ok"
|
|
|
|
@dec_tool
|
|
def other_tool(y: int) -> str:
|
|
"""Another tool available to the runtime."""
|
|
return str(y)
|
|
|
|
node = ToolNode([info_tool, other_tool])
|
|
tool_call = {
|
|
"name": "info_tool",
|
|
"args": {"x": 1},
|
|
"id": "call-1",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("", tool_calls=[tool_call])
|
|
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
|
|
result = node.invoke({"messages": [msg]}, config=config)
|
|
|
|
assert result["messages"][-1].content == "ok"
|
|
assert captured["execution_info"] is exec_info
|
|
assert captured["execution_info"].thread_id == "t-1"
|
|
assert captured["execution_info"].task_id == "tk-1"
|
|
assert captured["server_info"] is server_info
|
|
assert captured["server_info"].assistant_id == "asst-1"
|
|
assert [tool.name for tool in captured["tools"]] == ["info_tool", "other_tool"]
|
|
|
|
|
|
async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async() -> (
|
|
None
|
|
):
|
|
"""Test that execution_info, server_info, and tools are forwarded in async path."""
|
|
|
|
exec_info = ExecutionInfo(
|
|
thread_id="t-2",
|
|
checkpoint_id="cp-2",
|
|
checkpoint_ns="",
|
|
task_id="tk-2",
|
|
run_id="r-2",
|
|
)
|
|
server_info = ServerInfo(assistant_id="asst-2", graph_id="graph-2")
|
|
|
|
mock_runtime = Mock()
|
|
mock_runtime.store = None
|
|
mock_runtime.context = None
|
|
mock_runtime.stream_writer = lambda *args, **kwargs: None
|
|
mock_runtime.execution_info = exec_info
|
|
mock_runtime.server_info = server_info
|
|
|
|
captured: dict = {}
|
|
|
|
@dec_tool
|
|
async def info_tool_async(x: int, runtime: ToolRuntime) -> str:
|
|
"""Async tool that captures runtime info."""
|
|
captured["execution_info"] = runtime.execution_info
|
|
captured["server_info"] = runtime.server_info
|
|
captured["tools"] = runtime.tools
|
|
return "ok"
|
|
|
|
@dec_tool
|
|
async def other_tool_async(y: int) -> str:
|
|
"""Another async tool available to the runtime."""
|
|
return str(y)
|
|
|
|
node = ToolNode([info_tool_async, other_tool_async])
|
|
tool_call = {
|
|
"name": "info_tool_async",
|
|
"args": {"x": 1},
|
|
"id": "call-2",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("", tool_calls=[tool_call])
|
|
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
|
|
result = await node.ainvoke({"messages": [msg]}, config=config)
|
|
|
|
assert result["messages"][-1].content == "ok"
|
|
assert captured["execution_info"] is exec_info
|
|
assert captured["execution_info"].thread_id == "t-2"
|
|
assert captured["server_info"] is server_info
|
|
assert captured["server_info"].graph_id == "graph-2"
|
|
assert [tool.name for tool in captured["tools"]] == [
|
|
"info_tool_async",
|
|
"other_tool_async",
|
|
]
|
|
|
|
|
|
# --- InjectedToolArg security tests ---
|
|
|
|
|
|
def test_tool_node_strips_plain_injected_tool_arg() -> None:
|
|
"""Plain InjectedToolArg values supplied by the LLM should be stripped."""
|
|
|
|
@dec_tool
|
|
def read_secret(
|
|
query: str,
|
|
auth: Annotated[dict, InjectedToolArg()],
|
|
) -> str:
|
|
"""Return secret data based on auth role."""
|
|
if auth.get("role") == "admin":
|
|
return "ADMIN_SECRET"
|
|
return "PUBLIC_DATA"
|
|
|
|
node = ToolNode([read_secret], handle_tool_errors=True)
|
|
|
|
# LLM tries to supply the hidden 'auth' field
|
|
tool_call = {
|
|
"name": "read_secret",
|
|
"args": {"query": "hello", "auth": {"role": "admin"}},
|
|
"id": "call-1",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("", tool_calls=[tool_call])
|
|
result = node.invoke({"messages": [msg]}, config=_create_config_with_runtime())
|
|
tool_message = result["messages"][-1]
|
|
# auth should have been stripped, so tool should fail (missing required arg)
|
|
assert "ADMIN_SECRET" not in tool_message.content
|
|
|
|
|
|
def test_tool_node_strips_custom_injected_tool_arg_subclass() -> None:
|
|
"""Custom InjectedToolArg subclasses should also be stripped."""
|
|
|
|
class InjectedAuth(InjectedToolArg):
|
|
pass
|
|
|
|
@dec_tool
|
|
def read_secret(
|
|
query: str,
|
|
auth: Annotated[dict, InjectedAuth()],
|
|
) -> str:
|
|
"""Return secret data based on auth role."""
|
|
if auth.get("role") == "admin":
|
|
return "ADMIN_SECRET"
|
|
return "PUBLIC_DATA"
|
|
|
|
node = ToolNode([read_secret], handle_tool_errors=True)
|
|
|
|
tool_call = {
|
|
"name": "read_secret",
|
|
"args": {"query": "hello", "auth": {"role": "admin"}},
|
|
"id": "call-1",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("", tool_calls=[tool_call])
|
|
result = node.invoke({"messages": [msg]}, config=_create_config_with_runtime())
|
|
tool_message = result["messages"][-1]
|
|
assert "ADMIN_SECRET" not in tool_message.content
|
|
|
|
|
|
def test_tool_node_injected_state_overwrites_llm_value() -> None:
|
|
"""InjectedState should use graph state, not LLM-supplied values."""
|
|
|
|
@dec_tool
|
|
def read_secret(
|
|
query: str,
|
|
auth: Annotated[dict, InjectedState("auth")],
|
|
) -> str:
|
|
"""Return secret data based on auth from graph state."""
|
|
if auth.get("role") == "admin":
|
|
return "ADMIN_SECRET"
|
|
return "PUBLIC_DATA"
|
|
|
|
node = ToolNode([read_secret])
|
|
|
|
# LLM tries to supply auth as admin
|
|
tool_call = {
|
|
"name": "read_secret",
|
|
"args": {"query": "hello", "auth": {"role": "admin"}},
|
|
"id": "call-1",
|
|
"type": "tool_call",
|
|
}
|
|
msg = AIMessage("", tool_calls=[tool_call])
|
|
|
|
# Graph state has auth as viewer
|
|
result = node.invoke(
|
|
{"messages": [msg], "auth": {"role": "viewer"}},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
tool_message = result["messages"][-1]
|
|
assert tool_message.content == "PUBLIC_DATA"
|
|
|
|
|
|
class _ReturningTool(BaseTool):
|
|
"""A tool that returns a configured value verbatim."""
|
|
|
|
name: str = "list_tool"
|
|
description: str = "Returns a configured value"
|
|
return_value: Any = None
|
|
|
|
def _run(self, **kwargs: Any) -> Any:
|
|
return self.return_value
|
|
|
|
async def _arun(self, **kwargs: Any) -> Any:
|
|
return self.return_value
|
|
|
|
|
|
def _list_tool_call(outer_id: str = "call-1") -> dict[str, Any]:
|
|
return {"name": "list_tool", "args": {}, "id": outer_id, "type": "tool_call"}
|
|
|
|
|
|
def _invoke_returning(
|
|
return_value: Any,
|
|
*,
|
|
outer_id: str = "call-1",
|
|
handle_tool_errors: bool = True,
|
|
) -> Any:
|
|
node = ToolNode(
|
|
[_ReturningTool(return_value=return_value)],
|
|
handle_tool_errors=handle_tool_errors,
|
|
)
|
|
return node.invoke(
|
|
{"messages": [AIMessage("", tool_calls=[_list_tool_call(outer_id)])]},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
|
|
|
|
def test_tool_node_list_return_command_and_tool_message() -> None:
|
|
"""Valid: tool returns [Command(update={...}), ToolMessage(...)]."""
|
|
outer_id = "call-1"
|
|
result = _invoke_returning(
|
|
[
|
|
Command(update={"foo": "bar"}),
|
|
ToolMessage(content="done", tool_call_id=outer_id),
|
|
]
|
|
)
|
|
assert isinstance(result, list)
|
|
commands = [r for r in result if isinstance(r, Command)]
|
|
assert len(commands) == 1
|
|
assert commands[0].update == {"foo": "bar"}
|
|
non_commands = [r for r in result if not isinstance(r, Command)]
|
|
assert len(non_commands) == 1
|
|
assert isinstance(non_commands[0], dict)
|
|
msgs = non_commands[0]["messages"]
|
|
assert len(msgs) == 1
|
|
assert isinstance(msgs[0], ToolMessage)
|
|
assert msgs[0].content == "done"
|
|
assert msgs[0].tool_call_id == outer_id
|
|
|
|
|
|
def test_tool_node_list_return_nested_terminator() -> None:
|
|
"""Valid: terminator nested inside Command.update['messages']."""
|
|
outer_id = "call-1"
|
|
result = _invoke_returning(
|
|
[
|
|
Command(update={"foo": "bar"}),
|
|
Command(
|
|
update={
|
|
"messages": [ToolMessage(content="done", tool_call_id=outer_id)]
|
|
}
|
|
),
|
|
]
|
|
)
|
|
assert isinstance(result, list)
|
|
commands = [r for r in result if isinstance(r, Command)]
|
|
assert len(commands) == 2
|
|
updates = [c.update for c in commands]
|
|
assert {"foo": "bar"} in updates
|
|
msgs_update = next(u for u in updates if "messages" in (u or {}))
|
|
assert any(
|
|
isinstance(m, ToolMessage) and m.tool_call_id == outer_id
|
|
for m in msgs_update["messages"]
|
|
)
|
|
|
|
|
|
def test_tool_node_list_return_parent_goto_with_terminator() -> None:
|
|
"""Valid: [Command(graph=PARENT, goto=[Send(...)]), ToolMessage(...)]."""
|
|
outer_id = "call-1"
|
|
result = _invoke_returning(
|
|
[
|
|
Command(graph=Command.PARENT, goto=[Send("child", {})]),
|
|
ToolMessage(content="ok", tool_call_id=outer_id),
|
|
]
|
|
)
|
|
assert isinstance(result, list)
|
|
parent_cmds = [
|
|
r for r in result if isinstance(r, Command) and r.graph is Command.PARENT
|
|
]
|
|
assert len(parent_cmds) == 1
|
|
assert isinstance(parent_cmds[0].goto, list)
|
|
assert any(isinstance(s, Send) for s in parent_cmds[0].goto)
|
|
non_commands = [r for r in result if not isinstance(r, Command)]
|
|
assert len(non_commands) == 1
|
|
|
|
|
|
def test_tool_node_list_return_no_terminator_raises() -> None:
|
|
"""Invalid: list with no terminating ToolMessage."""
|
|
with pytest.raises(ValueError, match="0 messages bound to tool_call_id"):
|
|
_invoke_returning([Command(update={"foo": "bar"})], handle_tool_errors=False)
|
|
|
|
|
|
def test_tool_node_list_return_multiple_terminators_raises() -> None:
|
|
"""Invalid: list with two terminating ToolMessages."""
|
|
outer_id = "call-1"
|
|
with pytest.raises(ValueError, match="2 messages bound to tool_call_id"):
|
|
_invoke_returning(
|
|
[
|
|
ToolMessage(content="a", tool_call_id=outer_id),
|
|
ToolMessage(content="b", tool_call_id=outer_id),
|
|
],
|
|
handle_tool_errors=False,
|
|
)
|
|
|
|
|
|
def test_tool_node_list_return_validation_error_handled() -> None:
|
|
"""handle_tool_errors=True converts validation errors to an error ToolMessage."""
|
|
result = _invoke_returning([Command(update={"foo": "bar"})])
|
|
assert isinstance(result, dict)
|
|
msg = result["messages"][0]
|
|
assert isinstance(msg, ToolMessage)
|
|
assert msg.status == "error"
|
|
assert "0 messages bound to tool_call_id" in msg.content
|
|
|
|
|
|
async def test_tool_node_list_return_async_smoke() -> None:
|
|
"""Async path parallels sync for the happy case."""
|
|
outer_id = "call-1"
|
|
node = ToolNode(
|
|
[
|
|
_ReturningTool(
|
|
return_value=[
|
|
Command(update={"foo": "bar"}),
|
|
ToolMessage(content="done", tool_call_id=outer_id),
|
|
]
|
|
)
|
|
]
|
|
)
|
|
result = await node.ainvoke(
|
|
{"messages": [AIMessage("", tool_calls=[_list_tool_call(outer_id)])]},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert isinstance(result, list)
|
|
commands = [r for r in result if isinstance(r, Command)]
|
|
assert len(commands) == 1 and commands[0].update == {"foo": "bar"}
|
|
|
|
|
|
def test_tool_node_list_return_mixed_with_regular_tool() -> None:
|
|
"""List-returning tool and a regular tool dispatched from the same AIMessage."""
|
|
list_tool_id = "call-list"
|
|
regular_tool_id = "call-regular"
|
|
list_tool = _ReturningTool(
|
|
return_value=[
|
|
Command(update={"foo": "bar"}),
|
|
ToolMessage(content="list done", tool_call_id=list_tool_id),
|
|
]
|
|
)
|
|
|
|
def regular_tool(x: int) -> str:
|
|
"""A normal tool."""
|
|
return f"regular: {x}"
|
|
|
|
tool_calls = [
|
|
{"name": "list_tool", "args": {}, "id": list_tool_id, "type": "tool_call"},
|
|
{
|
|
"name": "regular_tool",
|
|
"args": {"x": 7},
|
|
"id": regular_tool_id,
|
|
"type": "tool_call",
|
|
},
|
|
]
|
|
node = ToolNode([list_tool, regular_tool])
|
|
result = node.invoke(
|
|
{"messages": [AIMessage("", tool_calls=tool_calls)]},
|
|
config=_create_config_with_runtime(),
|
|
)
|
|
assert isinstance(result, list)
|
|
commands = [r for r in result if isinstance(r, Command)]
|
|
assert len(commands) == 1
|
|
assert commands[0].update == {"foo": "bar"}
|
|
all_msgs = [m for r in result if isinstance(r, dict) for m in r["messages"]]
|
|
tool_call_ids = {m.tool_call_id for m in all_msgs}
|
|
assert list_tool_id in tool_call_ids
|
|
assert regular_tool_id in tool_call_ids
|