From 0cf3a64d66f1d43222a5827c4a6288df249182a7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 28 Jan 2025 09:48:33 -0800 Subject: [PATCH] Do not inject args in RunnableCallable if arg already exists (#3185) Co-authored-by: Eugene Yurtsev --- .../langgraph/langgraph/prebuilt/tool_node.py | 33 +-- libs/langgraph/langgraph/utils/runnable.py | 81 ++++--- libs/langgraph/tests/test_pregel.py | 19 ++ libs/langgraph/tests/test_pregel_async.py | 20 ++ libs/langgraph/tests/test_runnable.py | 202 +++++++++++++++++- libs/langgraph/tests/test_utils.py | 5 +- 6 files changed, 310 insertions(+), 50 deletions(-) diff --git a/libs/langgraph/langgraph/prebuilt/tool_node.py b/libs/langgraph/langgraph/prebuilt/tool_node.py index 30a4a54e3..581a89997 100644 --- a/libs/langgraph/langgraph/prebuilt/tool_node.py +++ b/libs/langgraph/langgraph/prebuilt/tool_node.py @@ -27,7 +27,6 @@ from langchain_core.runnables.config import ( get_config_list, get_executor_for_config, ) -from langchain_core.runnables.utils import Input from langchain_core.tools import BaseTool, InjectedToolArg from langchain_core.tools import tool as create_tool from langchain_core.tools.base import get_all_basemodel_annotations @@ -210,7 +209,7 @@ class ToolNode(RunnableCallable): ], config: RunnableConfig, *, - store: BaseStore, + store: Optional[BaseStore], ) -> Any: tool_calls, input_type = self._parse_input(input, store) config_list = get_config_list(config, len(tool_calls)) @@ -220,12 +219,14 @@ class ToolNode(RunnableCallable): *executor.map(self._run_one, tool_calls, input_types, config_list) ] - # preserve existing behavior for non-command tool outputs for backwards compatibility + # preserve existing behavior for non-command tool outputs for backwards + # compatibility if not any(isinstance(output, Command) for output in outputs): # TypedDict, pydantic, dataclass, etc. should all be able to load from dict return outputs if input_type == "list" else {self.messages_key: outputs} - # LangGraph will automatically handle list of Command and non-command node updates + # LangGraph will automatically handle list of Command and non-command node + # updates combined_outputs: list[ Command | list[ToolMessage] | dict[str, list[ToolMessage]] ] = [] @@ -238,20 +239,6 @@ class ToolNode(RunnableCallable): ) return combined_outputs - def invoke( - self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any - ) -> Any: - if "store" not in kwargs: - kwargs["store"] = None - return super().invoke(input, config, **kwargs) - - async def ainvoke( - self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any - ) -> Any: - if "store" not in kwargs: - kwargs["store"] = None - return await super().ainvoke(input, config, **kwargs) - async def _afunc( self, input: Union[ @@ -261,7 +248,7 @@ class ToolNode(RunnableCallable): ], config: RunnableConfig, *, - store: BaseStore, + store: Optional[BaseStore], ) -> Any: tool_calls, input_type = self._parse_input(input, store) outputs = await asyncio.gather( @@ -404,7 +391,7 @@ class ToolNode(RunnableCallable): dict[str, Any], BaseModel, ], - store: BaseStore, + store: Optional[BaseStore], ) -> Tuple[list[ToolCall], Literal["list", "dict"]]: if isinstance(input, list): input_type = "list" @@ -484,7 +471,9 @@ class ToolNode(RunnableCallable): } return tool_call - def _inject_store(self, tool_call: ToolCall, store: BaseStore) -> ToolCall: + def _inject_store( + self, tool_call: ToolCall, store: Optional[BaseStore] + ) -> ToolCall: store_arg = self.tool_to_store_arg[tool_call["name"]] if not store_arg: return tool_call @@ -509,7 +498,7 @@ class ToolNode(RunnableCallable): dict[str, Any], BaseModel, ], - store: BaseStore, + store: Optional[BaseStore], ) -> ToolCall: if tool_call["name"] not in self.tools_by_name: return tool_call diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index 8416efc72..f3f4c4b79 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -14,6 +14,7 @@ from typing import ( Iterator, Optional, Sequence, + Tuple, Union, cast, ) @@ -66,9 +67,10 @@ class StrEnum(str, enum.Enum): # Special type to denote any type is accepted ANY_TYPE = object() - ASYNCIO_ACCEPTS_CONTEXT = sys.version_info >= (3, 11) +# List of keyword arguments that can be injected at runtime from the config object. +# A named argument may appear multiple times if it appears with distinct types. KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = ( ( sys.intern("writer"), @@ -77,11 +79,31 @@ KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = ( lambda _: None, ), ( + # Covers store that is not optional (will raise an error if a store + # cannot be injected). sys.intern("store"), - (BaseStore, "BaseStore", inspect.Parameter.empty), + ( + BaseStore, + "BaseStore", + inspect.Parameter.empty, + ), CONFIG_KEY_STORE, inspect.Parameter.empty, ), + ( + # Covers store that is optional. Will set to None if not found in config. + sys.intern("store"), + ( + Optional[BaseStore], + # Best effort to catch some forward references. + # This will not work for cases like `"Union[None, BaseStore]"`, + # we'll need to re-write logic to use get_type_hints() + # to resolve forward references. + "Optional[BaseStore]", + ), + CONFIG_KEY_STORE, + None, + ), ( sys.intern("previous"), (ANY_TYPE,), @@ -149,15 +171,24 @@ class RunnableCallable(Runnable): params = inspect.signature(cast(Callable, func or afunc)).parameters self.func_accepts_config = "config" in params - self.func_accepts: dict[str, bool] = {} - for kw, typ, _, _ in KWARGS_CONFIG_KEYS: + # Mapping from kwarg name to (config key, default value) to be used. + # The default value is used if the config key is not found in the config. + self.func_accepts: dict[str, Tuple[str, Any]] = {} + + for kw, typ, config_key, default in KWARGS_CONFIG_KEYS: p = params.get(kw) - if typ == (ANY_TYPE,): - self.func_accepts[kw] = p is not None and p.kind in VALID_KINDS - else: - self.func_accepts[kw] = ( - p is not None and p.annotation in typ and p.kind in VALID_KINDS - ) + + if p is None or p.kind not in VALID_KINDS: + # If parameter is not found or is not a valid kind, skip + continue + + if typ != (ANY_TYPE,) and p.annotation not in typ: + # A specific type is required, but the function annotation does + # not match the expected type. + continue + + # If the kwarg is accepted by the function, store the default value + self.func_accepts[kw] = (config_key, default) def __repr__(self) -> str: repr_args = { @@ -187,20 +218,22 @@ class RunnableCallable(Runnable): if self.func_accepts_config: kwargs["config"] = config _conf = config[CONF] - for kw, _, config_key, default_value in KWARGS_CONFIG_KEYS: - if not self.func_accepts[kw]: + + for kw, (config_key, default_value) in self.func_accepts.items(): + # If the kwarg is already set, use the set value + if kw in kwargs: continue if ( - default_value is inspect.Parameter.empty - and kw not in kwargs - and config_key not in _conf + # If the kwarg is requested, but isn't in the config AND has no + # default value, raise an error + config_key not in _conf and default_value is inspect.Parameter.empty ): raise ValueError( f"Missing required config key '{config_key}' for '{self.name}'." ) - elif kwargs.get(kw) is None: - kwargs[kw] = _conf.get(config_key, default_value) + + kwargs[kw] = _conf.get(config_key, default_value) context = copy_context() if self.trace: @@ -244,20 +277,20 @@ class RunnableCallable(Runnable): if self.func_accepts_config: kwargs["config"] = config _conf = config[CONF] - for kw, _, config_key, default_value in KWARGS_CONFIG_KEYS: - if not self.func_accepts[kw]: + for kw, (config_key, default_value) in self.func_accepts.items(): + # If the kwarg has already been set, use the set value + if kw in kwargs: continue if ( - default_value is inspect.Parameter.empty - and kw not in kwargs - and config_key not in _conf + # If the kwarg is requested, but isn't in the config AND has no + # default value, raise an error + config_key not in _conf and default_value is inspect.Parameter.empty ): raise ValueError( f"Missing required config key '{config_key}' for '{self.name}'." ) - elif kwargs.get(kw) is None: - kwargs[kw] = _conf.get(config_key, default_value) + kwargs[kw] = _conf.get(config_key, default_value) context = copy_context() if self.trace: callback_manager = get_async_callback_manager_for_config(config, self.tags) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 58ac09b65..4aca6edf2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6220,6 +6220,25 @@ def test_entrypoint_with_return_and_save() -> None: assert previous_ == ["hello", "goodbye"] +def test_overriding_injectable_args_with_tasks() -> None: + """Test overriding injectable args in tasks.""" + from langgraph.store.memory import InMemoryStore + + @task + def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None: + assert store is value + assert writer is value + + @entrypoint(store=InMemoryStore()) + def main(inputs, store: BaseStore) -> str: + assert store is not None + foo(store=None, writer=None, value=None).result() + foo(store="hello", writer="hello", value="hello").result() + return "OK" + + assert main.invoke({}) == "OK" + + def test_named_tasks_functional() -> None: class Foo: def foo(self, value: str) -> dict: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 932330dc7..ac387b015 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7446,3 +7446,23 @@ async def test_named_tasks_functional() -> None: {"qux": "foo|bar|baz|custom_baz|qux"}, {"workflow": "foo|bar|baz|custom_baz|qux"}, ] + + +@NEEDS_CONTEXTVARS +async def test_overriding_injectable_args_with_async_task() -> None: + """Test overriding injectable args in tasks.""" + from langgraph.store.memory import InMemoryStore + + @task + async def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None: + assert store is value + assert writer is value + + @entrypoint(store=InMemoryStore()) + async def main(inputs, store: BaseStore) -> str: + assert store is not None + await foo(store=None, writer=None, value=None) + await foo(store="hello", writer="hello", value="hello") + return "OK" + + assert await main.ainvoke({}) == "OK" diff --git a/libs/langgraph/tests/test_runnable.py b/libs/langgraph/tests/test_runnable.py index 6cb4b5e9e..0a81be368 100644 --- a/libs/langgraph/tests/test_runnable.py +++ b/libs/langgraph/tests/test_runnable.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any +from typing import Any, Optional import pytest @@ -43,8 +43,15 @@ def test_runnable_callable_func_accepts(): expected_writer = {"with_writer": True, "awith_writer": True} for name, runnable in runnables.items(): - assert runnable.func_accepts["writer"] == expected_writer.get(name, False) - assert runnable.func_accepts["store"] == expected_store.get(name, False) + if expected_writer.get(name, False): + assert "writer" in runnable.func_accepts + else: + assert "writer" not in runnable.func_accepts + + if expected_store.get(name, False): + assert "store" in runnable.func_accepts + else: + assert "store" not in runnable.func_accepts async def test_runnable_callable_basic(): @@ -63,3 +70,192 @@ async def test_runnable_callable_basic(): # Test asynchronous ainvoke result_async = await runnable_async.ainvoke("test") assert result_async == "test" + + +def test_runnable_callable_injectable_arguments() -> None: + """Test injectable arguments for RunnableCallable. + + This test verifies that injectable arguments like BaseStore work correctly. + It tests: + - Optional store injection + - Required store injection + - Store injection via config + - Store injection override behavior + - Store value injection and validation + """ + + # Test Optional[BaseStore] annotation. + def func_optional_store(inputs: Any, store: Optional[BaseStore]) -> str: + """Test function that accepts an optional store parameter.""" + assert store is None + return "success" + + assert RunnableCallable(func_optional_store).invoke({"x": "1"}) == "success" + + # Test BaseStore annotation + def func_required_store(inputs: Any, store: BaseStore) -> str: + """Test function that requires a store parameter.""" + assert store is None + return "success" + + with pytest.raises(ValueError): + # Should fail b/c store is not Optional and config is not populated with store. + assert RunnableCallable(func_required_store).invoke({}) == "success" + + # Manually provide store + assert RunnableCallable(func_required_store).invoke({}, store=None) == "success" + + # Specify a value for store in the config + assert ( + RunnableCallable(func_required_store).invoke( + {}, config={"configurable": {"__pregel_store": None}} + ) + == "success" + ) + + # Specify a value for store in config, but override with None + assert ( + RunnableCallable(func_optional_store).invoke( + {"x": "1"}, + store=None, + config={"configurable": {"__pregel_store": "foobar"}}, + ) + == "success" + ) + + # Set of tests where we verify that 'foobar' is injected as the store value. + def func_required_store_v2(inputs: Any, store: BaseStore) -> str: + """Test function that requires a store parameter and validates its value. + + The store value is expected to be 'foobar' when injected. + """ + assert store == "foobar" + return "success" + + assert ( + RunnableCallable(func_required_store_v2).invoke( + {}, config={"configurable": {"__pregel_store": "foobar"}} + ) + == "success" + ) + + assert RunnableCallable(func_required_store_v2).invoke( + # And manual override takes precedence. + {}, + store="foobar", + config={"configurable": {"__pregel_store": "barbar"}}, + ) + + +async def test_runnable_callable_injectable_arguments_async() -> None: + """Test injectable arguments for async RunnableCallable. + + This test verifies that injectable arguments like BaseStore work correctly + in the async context. It tests: + - Optional store injection + - Required store injection + - Store injection via config + - Store injection override behavior + """ + + # Test Optional[BaseStore] annotation. + def func_optional_store(inputs: Any, store: Optional[BaseStore]) -> str: + """Test function that accepts an optional store parameter.""" + assert store is None + return "success" + + async def afunc_optional_store(inputs: Any, store: Optional[BaseStore]) -> str: + """Async version of func_optional_store.""" + assert store is None + return "success" + + assert ( + await RunnableCallable( + func=func_optional_store, afunc=afunc_optional_store + ).ainvoke({"x": "1"}) + == "success" + ) + + # Test BaseStore annotation + def func_required_store(inputs: Any, store: BaseStore) -> str: + """Test function that requires a store parameter.""" + assert store is None + return "success" + + async def afunc_required_store(inputs: Any, store: BaseStore) -> str: + """Async version of func_required_store.""" + assert store is None + return "success" + + with pytest.raises(ValueError): + # Should fail b/c store is not Optional and config is not populated with store. + assert ( + await RunnableCallable( + func=func_required_store, afunc=afunc_required_store + ).ainvoke({}) + == "success" + ) + + # Manually provide store + assert ( + await RunnableCallable( + func=func_required_store, afunc=afunc_required_store + ).ainvoke({}, store=None) + == "success" + ) + + # Specify a value for store in the config + assert ( + await RunnableCallable( + func=func_required_store, afunc=afunc_required_store + ).ainvoke({}, config={"configurable": {"__pregel_store": None}}) + == "success" + ) + + # Specify a value for store in config, but override with None + assert ( + await RunnableCallable( + func=func_optional_store, afunc=afunc_optional_store + ).ainvoke( + {"x": "1"}, + store=None, + config={"configurable": {"__pregel_store": "foobar"}}, + ) + == "success" + ) + + # Set of tests where we verify that 'foobar' is injected as the store value. + def func_required_store_v2(inputs: Any, store: BaseStore) -> str: + """Test function that requires a store parameter with specific value. + + The store parameter is expected to be 'foobar' when injected. + """ + assert store == "foobar" + return "success" + + async def afunc_required_store_v2(inputs: Any, store: BaseStore) -> str: + """Async version of func_required_store_v2. + + The store parameter is expected to be 'foobar' when injected. + """ + assert store == "foobar" + return "success" + + assert ( + await RunnableCallable( + func=func_required_store_v2, afunc=afunc_required_store_v2 + ).ainvoke({}, config={"configurable": {"__pregel_store": "foobar"}}) + == "success" + ) + + assert ( + await RunnableCallable( + func=func_required_store_v2, afunc=afunc_required_store_v2 + ).ainvoke( + # And manual override takes precedence. + {}, + store="foobar", + config={"configurable": {"__pregel_store": "barbar"}}, + ) + == "success" + ) diff --git a/libs/langgraph/tests/test_utils.py b/libs/langgraph/tests/test_utils.py index 8c9e304c9..3549bd574 100644 --- a/libs/langgraph/tests/test_utils.py +++ b/libs/langgraph/tests/test_utils.py @@ -26,7 +26,10 @@ from langgraph.utils.fields import ( get_enhanced_type_hints, get_field_default, ) -from langgraph.utils.runnable import is_async_callable, is_async_generator +from langgraph.utils.runnable import ( + is_async_callable, + is_async_generator, +) pytestmark = pytest.mark.anyio