WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-13 21:27:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix: refactor injection logic to respect function signatures (#6468)

Browse Source 
## overview

The main purpose of this is to respect tool signatures that request
injected args (like `ToolRuntime`) even when the explicitly specified
`args_schema` does not.

Ex in the following example, we should still inject `runtime` despite
its absence in `ArgsSchema`

```py
class ArgsSchema(BaseModel):
    some_arg: int = Field(...)

@tool(args_schema=ArgsSchema)
def my_tool(some_arg: int, runtime: ToolRuntime): ...
```

This is accompanied by
https://github.com/langchain-ai/langchain/pull/34051 which has tests
that pass w/ this change. This tests injection w/ `create_agent` (more
end to end than tests added in
https://github.com/langchain-ai/langchain/pull/33999.

This unblocks the injection of `ToolRuntime` into MCP tools which is
exciting bc that exposes tool call id and state, which we previously
were unable to do.

## other benefits

* Cleaner code structure w/ more helpful docs about injected args.
* Nice perf boost, we're no longer inspecting the annotations of a
tool's schema 3 different times to detect store, state, and runtime
injections.

## additional notes

1. I could see a world where we want more of this logic to reside on the
tools themselves, but tools don't now about LG specific injection types
(like `ToolRuntime`, hence having this logic here for now).
2. We could separately add validation for the case where something is
specified in `args_schema` and not in the function signature (probably
at the tool level though).
This commit is contained in:
Sydney Runkle
2025-11-20 11:37:54 -05:00
committed by GitHub
parent 6d20a0b9c7
commit 2164b7daa3
3 changed files with 383 additions and 229 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/prebuilt/langgraph/prebuilt/tool_node.py
+195 -213
View File
@@ -499,9 +499,7 @@ def _infer_handled_types(handler: Callable[..., str]) -> tuple[type[Exception],
def _filter_validation_errors(
validation_error: ValidationError,
tool_to_state_args: dict[str, str | None],
tool_to_store_arg: str | None,
tool_to_runtime_arg: str | None,
injected_args: _InjectedArgs | None,
) -> list[ErrorDetails]:
"""Filter validation errors to only include LLM-controlled arguments.
@@ -516,25 +514,28 @@ def _filter_validation_errors(
Args:
validation_error: The Pydantic ValidationError raised during tool invocation.
tool_to_state_args: Mapping of state argument names to state field names.
tool_to_store_arg: Name of the store argument, if any.
tool_to_runtime_arg: Name of the runtime argument, if any.
injected_args: The _InjectedArgs structure containing all injected arguments,
or None if there are no injected arguments.
Returns:
List of ErrorDetails containing only errors for LLM-controlled arguments,
with system-injected argument values removed from the input field.
"""
injected_args = set(tool_to_state_args.keys())
if tool_to_store_arg:
injected_args.add(tool_to_store_arg)
if tool_to_runtime_arg:
injected_args.add(tool_to_runtime_arg)
# Collect all injected argument names
injected_arg_names: set[str] = set()
if injected_args:
if injected_args.state:
injected_arg_names.update(injected_args.state.keys())
if injected_args.store:
injected_arg_names.add(injected_args.store)
if injected_args.runtime:
injected_arg_names.add(injected_args.runtime)
filtered_errors: list[ErrorDetails] = []
for error in validation_error.errors():
# Check if error location contains any injected argument
# error['loc'] is a tuple like ('field_name',) or ('field_name', 'nested_field')
if error["loc"] and error["loc"][0] not in injected_args:
if error["loc"] and error["loc"][0] not in injected_arg_names:
# Create a copy of the error dict to avoid mutating the original
error_copy: dict[str, Any] = {**error}
@@ -542,7 +543,7 @@ def _filter_validation_errors(
if isinstance(error_copy.get("input"), dict):
input_dict = error_copy["input"]
input_copy = {
k: v for k, v in input_dict.items() if k not in injected_args
k: v for k, v in input_dict.items() if k not in injected_arg_names
}
error_copy["input"] = input_copy
@@ -552,6 +553,60 @@ def _filter_validation_errors(
return filtered_errors
@dataclass
class _InjectedArgs:
"""Internal structure for tracking injected arguments for a tool.
This data structure is built once during ToolNode initialization by analyzing
the tool's signature and args schema, then reused during execution for efficient
injection without repeated reflection.
The structure maps from tool parameter names to their injection sources, enabling
the ToolNode to know exactly which arguments need to be injected and where to
get their values from.
Attributes:
state: Mapping from tool parameter names to state field names for injection.
Keys are tool parameter names, values are either:
- str: Name of the state field to extract and inject
- None: Inject the entire state object
Empty dict if no state injection is needed.
store: Name of the tool parameter where the store should be injected,
or None if no store injection is needed.
runtime: Name of the tool parameter where the runtime should be injected,
or None if no runtime injection is needed.
Example:
For a tool with signature:
```python
def my_tool(
x: int,
messages: Annotated[list, InjectedState("messages")],
full_state: Annotated[dict, InjectedState()],
store: Annotated[BaseStore, InjectedStore()],
runtime: ToolRuntime,
) -> str:
...
```
The resulting `_InjectedArgs` would be:
```python
_InjectedArgs(
state={
"messages": "messages", # Extract state["messages"]
"full_state": None, # Inject entire state
},
store="store", # Inject into "store" parameter
runtime="runtime", # Inject into "runtime" parameter
)
```
"""
state: dict[str, str | None]
store: str | None
runtime: str | None
class ToolNode(RunnableCallable):
"""A node for executing tools in LangGraph workflows.
@@ -696,9 +751,7 @@ class ToolNode(RunnableCallable):
"""
super().__init__(self._func, self._afunc, name=name, tags=tags, trace=False)
self._tools_by_name: dict[str, BaseTool] = {}
self._tool_to_state_args: dict[str, dict[str, str | None]] = {}
self._tool_to_store_arg: dict[str, str | None] = {}
self._tool_to_runtime_arg: dict[str, str | None] = {}
self._injected_args: dict[str, _InjectedArgs] = {}
self._handle_tool_errors = handle_tool_errors
self._messages_key = messages_key
self._wrap_tool_call = wrap_tool_call
@@ -709,9 +762,8 @@ class ToolNode(RunnableCallable):
else:
tool_ = tool
self._tools_by_name[tool_.name] = tool_
self._tool_to_state_args[tool_.name] = _get_state_args(tool_)
self._tool_to_store_arg[tool_.name] = _get_store_arg(tool_)
self._tool_to_runtime_arg[tool_.name] = _get_runtime_arg(tool_)
# Build injected args mapping once during initialization in a single pass
self._injected_args[tool_.name] = _get_all_injected_args(tool_)
@property
def tools_by_name(self) -> dict[str, BaseTool]:
@@ -864,12 +916,8 @@ class ToolNode(RunnableCallable):
response = tool.invoke(call_args, config)
except ValidationError as exc:
# Filter out errors for injected arguments
filtered_errors = _filter_validation_errors(
exc,
self._tool_to_state_args.get(call["name"], {}),
self._tool_to_store_arg.get(call["name"]),
self._tool_to_runtime_arg.get(call["name"]),
)
injected = self._injected_args.get(call["name"])
filtered_errors = _filter_validation_errors(exc, injected)
# Use original call["args"] without injected values for error reporting
raise ToolInvocationError(
call["name"], exc, call["args"], filtered_errors
@@ -1021,12 +1069,8 @@ class ToolNode(RunnableCallable):
response = await tool.ainvoke(call_args, config)
except ValidationError as exc:
# Filter out errors for injected arguments
filtered_errors = _filter_validation_errors(
exc,
self._tool_to_state_args.get(call["name"], {}),
self._tool_to_store_arg.get(call["name"]),
self._tool_to_runtime_arg.get(call["name"]),
)
injected = self._injected_args.get(call["name"])
filtered_errors = _filter_validation_errors(exc, injected)
# Use original call["args"] without injected values for error reporting
raise ToolInvocationError(
call["name"], exc, call["args"], filtered_errors
@@ -1219,86 +1263,6 @@ class ToolNode(RunnableCallable):
return input["state"]
return input
def _inject_state(
self,
tool_call: ToolCall,
state: list[AnyMessage] | dict[str, Any] | BaseModel,
) -> ToolCall:
state_args = self._tool_to_state_args[tool_call["name"]]
if state_args and isinstance(state, list):
required_fields = list(state_args.values())
if (
len(required_fields) == 1 and required_fields[0] == self._messages_key
) or required_fields[0] is None:
state = {self._messages_key: state}
else:
err_msg = (
f"Invalid input to ToolNode. Tool {tool_call['name']} requires "
f"graph state dict as input."
)
if any(state_field for state_field in state_args.values()):
required_fields_str = ", ".join(f for f in required_fields if f)
err_msg += f" State should contain fields {required_fields_str}."
raise ValueError(err_msg)
if isinstance(state, dict):
tool_state_args = {
tool_arg: state[state_field] if state_field else state
for tool_arg, state_field in state_args.items()
}
else:
tool_state_args = {
tool_arg: getattr(state, state_field) if state_field else state
for tool_arg, state_field in state_args.items()
}
tool_call["args"] = {
**tool_call["args"],
**tool_state_args,
}
return tool_call
def _inject_store(self, tool_call: ToolCall, store: BaseStore | None) -> ToolCall:
store_arg = self._tool_to_store_arg[tool_call["name"]]
if not store_arg:
return tool_call
if store is None:
msg = (
"Cannot inject store into tools with InjectedStore annotations - "
"please compile your graph with a store."
)
raise ValueError(msg)
tool_call["args"] = {
**tool_call["args"],
store_arg: store,
}
return tool_call
def _inject_runtime(
self, tool_call: ToolCall, tool_runtime: ToolRuntime
) -> ToolCall:
"""Inject ToolRuntime into tool call arguments.
Args:
tool_call: The tool call to inject runtime into.
tool_runtime: The ToolRuntime instance to inject.
Returns:
The tool call with runtime injected if needed.
"""
runtime_arg = self._tool_to_runtime_arg.get(tool_call["name"])
if not runtime_arg:
return tool_call
tool_call["args"] = {
**tool_call["args"],
runtime_arg: tool_runtime,
}
return tool_call
def _inject_tool_args(
self,
tool_call: ToolCall,
@@ -1337,12 +1301,64 @@ class ToolNode(RunnableCallable):
if tool_call["name"] not in self.tools_by_name:
return tool_call
injected = self._injected_args.get(tool_call["name"])
if not injected:
return tool_call
tool_call_copy: ToolCall = copy(tool_call)
tool_call_with_state = self._inject_state(tool_call_copy, tool_runtime.state)
tool_call_with_store = self._inject_store(
tool_call_with_state, tool_runtime.store
)
return self._inject_runtime(tool_call_with_store, tool_runtime)
injected_args = {}
# Inject state
if injected.state:
state = tool_runtime.state
# Handle list state by converting to dict
if isinstance(state, list):
required_fields = list(injected.state.values())
if (
len(required_fields) == 1
and required_fields[0] == self._messages_key
) or required_fields[0] is None:
state = {self._messages_key: state}
else:
err_msg = (
f"Invalid input to ToolNode. Tool {tool_call['name']} requires "
f"graph state dict as input."
)
if any(state_field for state_field in injected.state.values()):
required_fields_str = ", ".join(f for f in required_fields if f)
err_msg += (
f" State should contain fields {required_fields_str}."
)
raise ValueError(err_msg)
# Extract state values
if isinstance(state, dict):
for tool_arg, state_field in injected.state.items():
injected_args[tool_arg] = (
state[state_field] if state_field else state
)
else:
for tool_arg, state_field in injected.state.items():
injected_args[tool_arg] = (
getattr(state, state_field) if state_field else state
)
# Inject store
if injected.store:
if tool_runtime.store is None:
msg = (
"Cannot inject store into tools with InjectedStore annotations - "
"please compile your graph with a store."
)
raise ValueError(msg)
injected_args[injected.store] = tool_runtime.store
# Inject runtime
if injected.runtime:
injected_args[injected.runtime] = tool_runtime
tool_call_copy["args"] = {**tool_call_copy["args"], **injected_args}
return tool_call_copy
def _validate_tool_command(
self,
@@ -1735,120 +1751,86 @@ def _is_injection(
return False
def _get_state_args(tool: BaseTool) -> dict[str, str | None]:
"""Extract state injection mappings from tool annotations.
This function analyzes a tool's input schema to identify arguments that should
be injected with graph state. It processes InjectedState annotations to build
a mapping of tool argument names to state field names.
def _get_injection_from_type(
type_: Any, injection_type: type[InjectedState | InjectedStore | ToolRuntime]
) -> Any | None:
"""Extract injection instance from a type annotation.
Args:
tool: The tool to analyze for state injection requirements.
type_: The type annotation to check.
injection_type: The injection type to look for.
Returns:
A dictionary mapping tool argument names to state field names. If a field
name is None, the entire state should be injected for that argument.
The injection instance if found, True if injection marker found without instance, None otherwise.
"""
full_schema = tool.get_input_schema()
tool_args_to_state_fields: dict = {}
type_args = get_args(type_)
matches = [arg for arg in type_args if _is_injection(arg, injection_type)]
for name, type_ in get_all_basemodel_annotations(full_schema).items():
injections = [
type_arg
for type_arg in get_args(type_)
if _is_injection(type_arg, InjectedState)
]
if len(injections) > 1:
msg = (
"A tool argument should not be annotated with InjectedState more than "
f"once. Received arg {name} with annotations {injections}."
)
raise ValueError(msg)
if len(injections) == 1:
injection = injections[0]
if isinstance(injection, InjectedState) and injection.field:
tool_args_to_state_fields[name] = injection.field
else:
tool_args_to_state_fields[name] = None
else:
pass
return tool_args_to_state_fields
if len(matches) > 1:
msg = (
f"A tool argument should not be annotated with {injection_type.__name__} "
f"more than once. Found: {matches}"
)
raise ValueError(msg)
def _get_store_arg(tool: BaseTool) -> str | None:
"""Extract store injection argument from tool annotations.
This function analyzes a tool's input schema to identify the argument that
should be injected with the graph store. Only one store argument is supported
per tool.
Args:
tool: The tool to analyze for store injection requirements.
Returns:
The name of the argument that should receive the store injection, or None
if no store injection is required.
Raises:
ValueError: If a tool argument has multiple InjectedStore annotations.
"""
full_schema = tool.get_input_schema()
for name, type_ in get_all_basemodel_annotations(full_schema).items():
injections = [
type_arg
for type_arg in get_args(type_)
if _is_injection(type_arg, InjectedStore)
]
if len(injections) > 1:
msg = (
"A tool argument should not be annotated with InjectedStore more than "
f"once. Received arg {name} with annotations {injections}."
)
raise ValueError(msg)
if len(injections) == 1:
return name
if len(matches) == 1:
return matches[0]
elif _is_injection(type_, injection_type):
return True
return None
def _get_runtime_arg(tool: BaseTool) -> str | None:
"""Extract runtime injection argument from tool annotations.
def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
"""Extract all injected arguments from tool in a single pass.
This function analyzes a tool's input schema to identify the argument that
should be injected with the ToolRuntime instance. Only one runtime argument
is supported per tool.
This function analyzes both the tool's input schema and function signature
to identify all arguments that should be injected (state, store, runtime).
Args:
tool: The tool to analyze for runtime injection requirements.
tool: The tool to analyze for injection requirements.
Returns:
The name of the argument that should receive the runtime injection, or None
if no runtime injection is required.
Raises:
ValueError: If a tool argument has multiple ToolRuntime annotations.
_InjectedArgs structure containing all detected injections.
"""
# Get annotations from both schema and function signature
full_schema = tool.get_input_schema()
for name, type_ in get_all_basemodel_annotations(full_schema).items():
# Check if the parameter name is "runtime" (regardless of type)
schema_annotations = get_all_basemodel_annotations(full_schema)
func = getattr(tool, "func", None) or getattr(tool, "coroutine", None)
func_annotations = get_type_hints(func, include_extras=True) if func else {}
# Combine both annotation sources, preferring schema annotations
# In the future, we might want to add more restrictions here...
all_annotations = {**func_annotations, **schema_annotations}
# Track injected args
state_args: dict[str, str | None] = {}
store_arg: str | None = None
runtime_arg: str | None = None
for name, type_ in all_annotations.items():
# Check for runtime (special case: parameter named "runtime")
if name == "runtime":
return name
# Check if the type itself is ToolRuntime (direct usage)
if _is_injection(type_, ToolRuntime):
return name
# Check if ToolRuntime is in Annotated args
injections = [
type_arg
for type_arg in get_args(type_)
if _is_injection(type_arg, ToolRuntime)
]
if len(injections) > 1:
msg = (
"A tool argument should not be annotated with ToolRuntime more than "
f"once. Received arg {name} with annotations {injections}."
)
raise ValueError(msg)
if len(injections) == 1:
return name
runtime_arg = name
return None
# Check for InjectedState
if state_inj := _get_injection_from_type(type_, InjectedState):
if isinstance(state_inj, InjectedState) and state_inj.field:
state_args[name] = state_inj.field
else:
state_args[name] = None
# Check for InjectedStore
if _get_injection_from_type(type_, InjectedStore):
store_arg = name
# Check for ToolRuntime
if _get_injection_from_type(type_, ToolRuntime):
runtime_arg = name
return _InjectedArgs(
state=state_args,
store=store_arg,
runtime=runtime_arg,
)
libs/prebuilt/tests/test_react_agent.py
-16
View File
@@ -53,7 +53,6 @@ from langgraph.prebuilt.chat_agent_executor import (
from langgraph.prebuilt.tool_node import (
InjectedState,
InjectedStore,
_get_state_args,
_infer_handled_types,
)
from tests.any_str import AnyStr
@@ -1084,21 +1083,6 @@ async def test_return_direct(version: str) -> None:
]
def test__get_state_args() -> None:
class Schema1(BaseModel):
a: Annotated[str, InjectedState]
class Schema2(Schema1):
b: Annotated[int, InjectedState("bar")]
@dec_tool(args_schema=Schema2)
def foo(a: str, b: int) -> float:
"""return"""
return 0.0
assert _get_state_args(foo) == {"a": None, "b": "bar"}
def test_inspect_react() -> None:
model = FakeToolCallingModel(tool_calls=[])
agent = create_react_agent(model, [])
libs/prebuilt/tests/test_tool_node.py
+188
View File
@@ -42,6 +42,7 @@ from langgraph.prebuilt import (
from langgraph.prebuilt.tool_node import (
TOOL_CALL_ERROR_TEMPLATE,
ToolInvocationError,
ToolRuntime,
tools_condition,
)
@@ -1672,3 +1673,190 @@ def test_tool_call_request_setattr_deprecation_warning():
)
# 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"
)