diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index 86084c0bb..425ee38b0 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -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, + ) diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index e4b05e5ad..eaf42106a 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -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, []) diff --git a/libs/prebuilt/tests/test_tool_node.py b/libs/prebuilt/tests/test_tool_node.py index c6e143f4f..00845d398 100644 --- a/libs/prebuilt/tests/test_tool_node.py +++ b/libs/prebuilt/tests/test_tool_node.py @@ -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" + )