diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index b2bcaf5ca..06d2102a1 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -348,34 +348,34 @@ class ToolNode(RunnableCallable): for tool_ in tools: if not isinstance(tool_, BaseTool): tool_ = create_tool(tool_) - + # Check for deprecated annotation usage and emit warnings # We need to check for annotations directly, not just the presence of state/store args # because _get_state_args returns both reserved keywords and annotations reserved_args = _get_reserved_keyword_args(tool_) - + # Check for InjectedState and InjectedStore annotations full_schema = tool_.get_input_schema() has_injected_state = False has_injected_store = False - + for name, type_ in get_all_basemodel_annotations(full_schema).items(): type_args = get_args(type_) - + # Check for InjectedState (can be class or instance) for type_arg in type_args: if _is_injection(type_arg, InjectedState): - if 'state' not in reserved_args: + if "state" not in reserved_args: has_injected_state = True break - + # Check for InjectedStore (can be class or instance) for type_arg in type_args: if _is_injection(type_arg, InjectedStore): - if 'runtime' not in reserved_args: + if "runtime" not in reserved_args: has_injected_store = True break - + # Emit deprecation warnings if has_injected_state: warnings.warn( @@ -385,9 +385,9 @@ class ToolNode(RunnableCallable): f"def {tool_.name}(..., state: Annotated[dict, InjectedState]). " f"The annotation-based approach will be removed in a future version.", DeprecationWarning, - stacklevel=2 + stacklevel=2, ) - + if has_injected_store: warnings.warn( f"Tool '{tool_.name}' uses deprecated InjectedStore annotation. " @@ -395,14 +395,14 @@ class ToolNode(RunnableCallable): f"Example: def {tool_.name}(..., runtime) and access store via runtime.store. " f"The annotation-based approach will be removed in a future version.", DeprecationWarning, - stacklevel=2 + stacklevel=2, ) - + # Check for reserved keywords and wrap the tool if needed # (already computed above, no need to recompute) if reserved_args: tool_ = _wrap_tool_with_reserved_keywords(tool_, reserved_args) - + 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_) @@ -812,7 +812,9 @@ class ToolNode(RunnableCallable): tool_call_with_state = self._inject_state(tool_call_copy, input) tool_call_with_store = self._inject_store(tool_call_with_state, store) if config: - tool_call_with_runtime = self._inject_runtime(tool_call_with_store, store, config) + tool_call_with_runtime = self._inject_runtime( + tool_call_with_store, store, config + ) return tool_call_with_runtime return tool_call_with_store @@ -961,10 +963,10 @@ class InjectedState(InjectedToolArg): .. deprecated:: 0.2.0 Use reserved keyword 'state' instead of InjectedState annotation. The annotation-based approach will be removed in a future version. - + Instead of: def tool(x: int, state: Annotated[dict, InjectedState]) -> str: - + Use: def tool(x: int, state) -> str: @@ -1045,10 +1047,10 @@ class InjectedStore(InjectedToolArg): .. deprecated:: 0.2.0 Use reserved keyword 'runtime' instead of InjectedStore annotation. The annotation-based approach will be removed in a future version. - + Instead of: def tool(x: int, store: Annotated[BaseStore, InjectedStore()]) -> str: - + Use: def tool(x: int, runtime) -> str: # Access store via runtime.store @@ -1167,15 +1169,15 @@ def _get_reserved_keyword_args(tool: BaseTool) -> dict[str, str]: Keys are parameter names, values are either 'state' or 'runtime'. """ reserved_args = {} - + # Get the underlying function from the tool - if hasattr(tool, 'func'): + if hasattr(tool, "func"): func = tool.func - elif hasattr(tool, '_run'): + elif hasattr(tool, "_run"): func = tool._run else: return reserved_args - + # Inspect the function signature try: sig = inspect.signature(func) @@ -1183,20 +1185,24 @@ def _get_reserved_keyword_args(tool: BaseTool) -> dict[str, str]: # Check for reserved keywords only if they don't have injection annotations # Parameters with InjectedState or InjectedStore annotations should not be # considered reserved keywords (they use the old annotation-based approach) - if param_name == 'state' and param.annotation == inspect.Parameter.empty: + if param_name == "state" and param.annotation == inspect.Parameter.empty: # Only consider 'state' as reserved if it has no annotation - reserved_args['state'] = 'state' - elif param_name == 'runtime' and param.annotation == inspect.Parameter.empty: + reserved_args["state"] = "state" + elif ( + param_name == "runtime" and param.annotation == inspect.Parameter.empty + ): # Only consider 'runtime' as reserved if it has no annotation - reserved_args['runtime'] = 'runtime' + reserved_args["runtime"] = "runtime" except (ValueError, TypeError): # If we can't inspect the signature, return empty pass - + return reserved_args -def _wrap_tool_with_reserved_keywords(tool: BaseTool, reserved_args: dict[str, str]) -> BaseTool: +def _wrap_tool_with_reserved_keywords( + tool: BaseTool, reserved_args: dict[str, str] +) -> BaseTool: """Wrap a tool to exclude reserved keyword parameters from its schema. This function creates a wrapper around tools that use reserved keywords @@ -1212,21 +1218,22 @@ def _wrap_tool_with_reserved_keywords(tool: BaseTool, reserved_args: dict[str, s """ # Get the original schema original_schema = tool.get_input_schema() - + # Create a new schema class that excludes reserved keywords # We need to dynamically create a new Pydantic model with filtered fields from pydantic import create_model - from pydantic.fields import FieldInfo - + # Get the original fields - if hasattr(original_schema, 'model_fields'): + if hasattr(original_schema, "model_fields"): # Pydantic v2 original_fields = original_schema.model_fields filtered_fields = {} for field_name, field_info in original_fields.items(): if field_name not in reserved_args: # Create a tuple for create_model: (type, field_info) - field_type = field_info.annotation if hasattr(field_info, 'annotation') else Any + field_type = ( + field_info.annotation if hasattr(field_info, "annotation") else Any + ) filtered_fields[field_name] = (field_type, field_info) else: # Pydantic v1 (backward compatibility) @@ -1235,42 +1242,50 @@ def _wrap_tool_with_reserved_keywords(tool: BaseTool, reserved_args: dict[str, s for field_name, field_info in original_fields.items(): if field_name not in reserved_args: filtered_fields[field_name] = (field_info.type_, field_info.field_info) - + # Create the filtered schema model FilteredSchema = create_model( - f"{original_schema.__name__}Filtered", - __base__=BaseModel, - **filtered_fields + f"{original_schema.__name__}Filtered", __base__=BaseModel, **filtered_fields ) - + # Create a wrapper tool with the filtered schema class WrappedTool(type(tool)): """Tool wrapper that excludes reserved keywords from schema.""" - - def get_input_schema(self, config: Optional[RunnableConfig] = None) -> Type[BaseModel]: + + def get_input_schema( + self, config: Optional[RunnableConfig] = None + ) -> Type[BaseModel]: """Return the filtered schema without reserved keywords.""" return FilteredSchema - + # Create the wrapped tool instance wrapped = WrappedTool( name=tool.name, description=tool.description, - func=tool.func if hasattr(tool, 'func') else None, + func=tool.func if hasattr(tool, "func") else None, args_schema=FilteredSchema, # Set the filtered schema ) - + # Copy over other attributes - for attr in ['return_direct', 'verbose', 'callbacks', 'tags', 'metadata', - 'handle_tool_error', 'handle_validation_error', 'response_format']: + for attr in [ + "return_direct", + "verbose", + "callbacks", + "tags", + "metadata", + "handle_tool_error", + "handle_validation_error", + "response_format", + ]: if hasattr(tool, attr): setattr(wrapped, attr, getattr(tool, attr)) - + # Ensure the wrapped tool still has access to the original run method - if hasattr(tool, '_run'): + if hasattr(tool, "_run"): wrapped._run = tool._run - if hasattr(tool, '_arun'): + if hasattr(tool, "_arun"): wrapped._arun = tool._arun - + return wrapped @@ -1289,19 +1304,19 @@ def _get_state_args(tool: BaseTool) -> dict[str, Optional[str]]: name is None, the entire state should be injected for that argument. """ tool_args_to_state_fields: dict = {} - + # First check for reserved keywords reserved_args = _get_reserved_keyword_args(tool) - if 'state' in reserved_args: - tool_args_to_state_fields['state'] = None - + if "state" in reserved_args: + tool_args_to_state_fields["state"] = None + # Then check for annotation-based injection (backward compatibility) full_schema = tool.get_input_schema() for name, type_ in get_all_basemodel_annotations(full_schema).items(): # Skip if already handled by reserved keyword if name in tool_args_to_state_fields: continue - + injections = [ type_arg for type_arg in get_args(type_) @@ -1377,29 +1392,4 @@ def _get_runtime_arg(tool: BaseTool) -> Optional[str]: The string 'runtime' if the tool has a runtime parameter, or None otherwise. """ reserved_args = _get_reserved_keyword_args(tool) - return 'runtime' if 'runtime' in reserved_args else None - - - - - - - - - - - - - - - - - - - - - - - - - + return "runtime" if "runtime" in reserved_args else None diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index a4781b3ab..c29d021a2 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -912,6 +912,7 @@ def test_tool_node_inject_store() -> None: def test_tool_node_inject_state_reserved_keyword() -> None: """Test that tools can use 'state' as a reserved keyword parameter.""" + def tool1(some_val: int, state) -> str: """Tool 1 with reserved keyword 'state'.""" if isinstance(state, dict): @@ -935,15 +936,19 @@ def test_tool_node_inject_state_reserved_keyword() -> None: # Test with dict state node = ToolNode([tool1, tool2, tool3]) - + # Verify that 'state' is excluded from tool schemas for tool in [tool1, tool2, tool3]: schema = node.tools_by_name[tool.__name__].get_input_schema() - if hasattr(schema, 'model_fields'): - assert "state" not in schema.model_fields, f"'state' should be excluded from {tool.__name__} schema" + if hasattr(schema, "model_fields"): + assert "state" not in schema.model_fields, ( + f"'state' should be excluded from {tool.__name__} schema" + ) else: - assert "state" not in schema.__fields__, f"'state' should be excluded from {tool.__name__} schema" - + assert "state" not in schema.__fields__, ( + f"'state' should be excluded from {tool.__name__} schema" + ) + for tool_name in ("tool1", "tool2"): tool_call = { "name": tool_name, @@ -957,8 +962,10 @@ def test_tool_node_inject_state_reserved_keyword() -> None: if tool_name == "tool1": assert tool_message.content == "bar", f"Failed for tool={tool_name}" else: - assert tool_message.content == "val: 1, foo: bar", f"Failed for tool={tool_name}" - + assert tool_message.content == "val: 1, foo: bar", ( + f"Failed for tool={tool_name}" + ) + # Test tool3 with additional parameter tool_call = { "name": "tool3", @@ -987,15 +994,19 @@ def test_tool_node_inject_state_reserved_keyword() -> None: result = node_pydantic.invoke(State(messages=[msg], foo="baz")) tool_message = result["messages"][-1] if tool_name == "tool1": - assert tool_message.content == "baz", f"Failed for tool={tool_name} with Pydantic state" + assert tool_message.content == "baz", ( + f"Failed for tool={tool_name} with Pydantic state" + ) else: - assert tool_message.content == "val: 2, foo: baz", f"Failed for tool={tool_name} with Pydantic state" + assert tool_message.content == "val: 2, foo: baz", ( + f"Failed for tool={tool_name} with Pydantic state" + ) def test_tool_node_inject_runtime_reserved_keyword() -> None: """Test that tools can use 'runtime' as a reserved keyword parameter.""" from langgraph.runtime import Runtime - + def tool1(some_val: int, runtime) -> str: """Tool 1 with reserved keyword 'runtime'.""" assert isinstance(runtime, Runtime) @@ -1007,16 +1018,16 @@ def test_tool_node_inject_runtime_reserved_keyword() -> None: store = InMemoryStore() store.put(("test",), "test_key", {"foo": "bar"}) - + node = ToolNode([tool1]) - + # Verify that 'runtime' is excluded from tool schemas schema = node.tools_by_name[tool1.__name__].get_input_schema() - if hasattr(schema, 'model_fields'): + if hasattr(schema, "model_fields"): assert "runtime" not in schema.model_fields else: assert "runtime" not in schema.__fields__ - + # Test with store tool_call = { "name": "tool1", @@ -2282,7 +2293,3 @@ def test_create_react_agent_inject_vars_with_post_model_hook( AIMessage("hi-hi-6", id="1"), ] assert result["foo"] == 2 - - - - diff --git a/libs/prebuilt/tests/test_reserved_keywords.py b/libs/prebuilt/tests/test_reserved_keywords.py index 27820e329..a86274184 100644 --- a/libs/prebuilt/tests/test_reserved_keywords.py +++ b/libs/prebuilt/tests/test_reserved_keywords.py @@ -1,16 +1,18 @@ """Test reserved keywords for tool injection.""" + from typing import Annotated + from langchain_core.messages import AIMessage -from langgraph.prebuilt import ToolNode, InjectedState, InjectedStore -from langgraph.store.memory import InMemoryStore + +from langgraph.prebuilt import InjectedState, InjectedStore, ToolNode from langgraph.store.base import BaseStore -from langgraph.graph import MessagesState +from langgraph.store.memory import InMemoryStore def test_tool_node_inject_runtime_reserved_keyword() -> None: """Test that tools can use 'runtime' as a reserved keyword parameter.""" from langgraph.runtime import Runtime - + def tool1(some_val: int, runtime) -> str: """Tool 1 with reserved keyword 'runtime'.""" assert isinstance(runtime, Runtime), "runtime should be a Runtime instance" @@ -26,7 +28,9 @@ def test_tool_node_inject_runtime_reserved_keyword() -> None: assert isinstance(runtime, Runtime), "runtime should be a Runtime instance" # Access context from runtime if runtime.context: - return f"val: {some_val}, context: {runtime.context.get('user_id', 'unknown')}" + return ( + f"val: {some_val}, context: {runtime.context.get('user_id', 'unknown')}" + ) return f"val: {some_val}, no context" def tool3(x: int, y: str, runtime) -> str: @@ -38,17 +42,21 @@ def test_tool_node_inject_runtime_reserved_keyword() -> None: store = InMemoryStore() store.put(("test",), "test_key", {"foo": "bar"}) - + node = ToolNode([tool1, tool2, tool3]) - + # Verify that 'runtime' is excluded from tool schemas for tool in [tool1, tool2, tool3]: schema = node.tools_by_name[tool.__name__].get_input_schema() - if hasattr(schema, 'model_fields'): - assert "runtime" not in schema.model_fields, f"'runtime' should be excluded from {tool.__name__} schema" + if hasattr(schema, "model_fields"): + assert "runtime" not in schema.model_fields, ( + f"'runtime' should be excluded from {tool.__name__} schema" + ) else: - assert "runtime" not in schema.__fields__, f"'runtime' should be excluded from {tool.__name__} schema" - + assert "runtime" not in schema.__fields__, ( + f"'runtime' should be excluded from {tool.__name__} schema" + ) + # Test with store tool_call = { "name": "tool1", @@ -60,11 +68,12 @@ def test_tool_node_inject_runtime_reserved_keyword() -> None: result = node.invoke({"messages": [msg]}, store=store) tool_message = result["messages"][-1] assert tool_message.content == "val: 1, store: bar" - + # Test with context from langchain_core.runnables import RunnableConfig + config = RunnableConfig(configurable={"context": {"user_id": "test_user"}}) - + tool_call = { "name": "tool2", "args": {"some_val": 2}, @@ -75,7 +84,7 @@ def test_tool_node_inject_runtime_reserved_keyword() -> None: result = node.invoke({"messages": [msg]}, config=config) tool_message = result["messages"][-1] assert tool_message.content == "val: 2, context: test_user" - + # Test with both store and context tool_call = { "name": "tool3", @@ -92,28 +101,28 @@ def test_tool_node_inject_runtime_reserved_keyword() -> None: def test_tool_node_mixed_injection_styles() -> None: """Test that tools can mix reserved keywords and annotations.""" from langgraph.runtime import Runtime - + def tool1(some_val: int, state) -> str: """Tool with reserved keyword 'state'.""" if isinstance(state, dict): return f"reserved state: {state['foo']}" else: return f"reserved state: {getattr(state, 'foo')}" - + def tool2(some_val: int, state: Annotated[dict, InjectedState]) -> str: """Tool with annotation-based state injection.""" return f"annotated state: {state['foo']}" - + def tool3(some_val: int, runtime) -> str: """Tool with reserved keyword 'runtime'.""" assert isinstance(runtime, Runtime) return f"reserved runtime: {runtime.context.get('user_id', 'none') if runtime.context else 'none'}" - + def tool4(some_val: int, store: Annotated[BaseStore, InjectedStore()]) -> str: """Tool with annotation-based store injection.""" store_val = store.get(("test",), "test_key") return f"annotated store: {store_val.value['foo'] if store_val else 'none'}" - + def tool5(x: int, state, runtime) -> str: """Tool with both reserved keywords.""" assert isinstance(runtime, Runtime) @@ -121,12 +130,12 @@ def test_tool_node_mixed_injection_styles() -> None: return f"both: state={state['foo']}, runtime={runtime.context.get('user_id', 'none') if runtime.context else 'none'}" else: return f"both: state={getattr(state, 'foo')}, runtime={runtime.context.get('user_id', 'none') if runtime.context else 'none'}" - + store = InMemoryStore() store.put(("test",), "test_key", {"foo": "bar"}) - + node = ToolNode([tool1, tool2, tool3, tool4, tool5]) - + # Verify schemas exclude injected parameters for tool_name, expected_excluded in [ ("tool1", ["state"]), @@ -136,16 +145,19 @@ def test_tool_node_mixed_injection_styles() -> None: ("tool5", ["state", "runtime"]), ]: schema = node.tools_by_name[tool_name].get_input_schema() - if hasattr(schema, 'model_fields'): + if hasattr(schema, "model_fields"): fields = schema.model_fields else: fields = schema.__fields__ for param in expected_excluded: - assert param not in fields, f"'{param}' should be excluded from {tool_name} schema" - + assert param not in fields, ( + f"'{param}' should be excluded from {tool_name} schema" + ) + from langchain_core.runnables import RunnableConfig + config = RunnableConfig(configurable={"context": {"user_id": "test_user"}}) - + # Test each tool test_cases = [ ("tool1", {"some_val": 1}, "reserved state: baz"), @@ -154,7 +166,7 @@ def test_tool_node_mixed_injection_styles() -> None: ("tool4", {"some_val": 4}, "annotated store: bar"), ("tool5", {"x": 5}, "both: state=baz, runtime=test_user"), ] - + for tool_name, args, expected in test_cases: tool_call = { "name": tool_name, @@ -163,6 +175,10 @@ def test_tool_node_mixed_injection_styles() -> None: "type": "tool_call", } msg = AIMessage("test", tool_calls=[tool_call]) - result = node.invoke({"messages": [msg], "foo": "baz"}, store=store, config=config) + result = node.invoke( + {"messages": [msg], "foo": "baz"}, store=store, config=config + ) tool_message = result["messages"][-1] - assert tool_message.content == expected, f"Failed for {tool_name}: got {tool_message.content}, expected {expected}" + assert tool_message.content == expected, ( + f"Failed for {tool_name}: got {tool_message.content}, expected {expected}" + )