From 3408e525f8fd6f29f8d1f8dba72a563ff664c838 Mon Sep 17 00:00:00 2001 From: "open-swe[bot]" Date: Fri, 22 Aug 2025 12:57:52 +0000 Subject: [PATCH] Apply patch [skip ci] --- test_reserved_keywords_full.py | 181 +++++++++++++++++++++++++++++++++ 1 file changed, 181 insertions(+) create mode 100644 test_reserved_keywords_full.py diff --git a/test_reserved_keywords_full.py b/test_reserved_keywords_full.py new file mode 100644 index 000000000..547cd3e5a --- /dev/null +++ b/test_reserved_keywords_full.py @@ -0,0 +1,181 @@ +"""Test reserved keywords for tool injection.""" +import pytest +from typing import Annotated, List +from langchain_core.messages import AIMessage, AnyMessage +from langgraph.prebuilt import ToolNode, InjectedState, InjectedStore +from langgraph.store.memory import InMemoryStore +from langgraph.store.base import BaseStore +from langgraph.types import MessagesState + + +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" + # Access store from runtime + if runtime.store: + store_val = runtime.store.get(("test",), "test_key") + if store_val: + return f"val: {some_val}, store: {store_val.value['foo']}" + return f"val: {some_val}, no store" + + def tool2(some_val: int, runtime) -> str: + """Tool 2 with reserved keyword 'runtime'.""" + 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}, no context" + + def tool3(x: int, y: str, runtime) -> str: + """Tool 3 with reserved keyword 'runtime' and other params.""" + assert isinstance(runtime, Runtime), "runtime should be a Runtime instance" + has_store = "yes" if runtime.store else "no" + has_context = "yes" if runtime.context else "no" + return f"x: {x}, y: {y}, store: {has_store}, context: {has_context}" + + 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" + else: + assert "runtime" not in schema.__fields__, f"'runtime' should be excluded from {tool.__name__} schema" + + # Test with store + tool_call = { + "name": "tool1", + "args": {"some_val": 1}, + "id": "some 0", + "type": "tool_call", + } + msg = AIMessage("hi?", tool_calls=[tool_call]) + 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}, + "id": "some 1", + "type": "tool_call", + } + msg = AIMessage("hi?", tool_calls=[tool_call]) + 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", + "args": {"x": 3, "y": "test"}, + "id": "some 2", + "type": "tool_call", + } + msg = AIMessage("hi?", tool_calls=[tool_call]) + result = node.invoke({"messages": [msg]}, store=store, config=config) + tool_message = result["messages"][-1] + assert tool_message.content == "x: 3, y: test, store: yes, context: yes" + + +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) + if isinstance(state, dict): + 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"]), + ("tool2", ["state"]), + ("tool3", ["runtime"]), + ("tool4", ["store"]), + ("tool5", ["state", "runtime"]), + ]: + schema = node.tools_by_name[tool_name].get_input_schema() + 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" + + 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"), + ("tool2", {"some_val": 2}, "annotated state: baz"), + ("tool3", {"some_val": 3}, "reserved runtime: test_user"), + ("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, + "args": args, + "id": f"id_{tool_name}", + "type": "tool_call", + } + msg = AIMessage("test", tool_calls=[tool_call]) + 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}" + + +if __name__ == "__main__": + print("Testing runtime reserved keyword...") + test_tool_node_inject_runtime_reserved_keyword() + print("✓ Runtime reserved keyword test passed!") + + print("\nTesting mixed injection styles...") + test_tool_node_mixed_injection_styles() + print("✓ Mixed injection styles test passed!") + + print("\nAll tests passed!")