From 7efd3c726e104d4bf0796e0808a60597078789bf Mon Sep 17 00:00:00 2001 From: Vadym Barda Date: Wed, 2 Oct 2024 13:39:25 -0400 Subject: [PATCH] langgraph: add support for store in ToolNode (#1968) --- libs/langgraph/langgraph/prebuilt/__init__.py | 8 +- .../langgraph/langgraph/prebuilt/tool_node.py | 178 +++++++++++++++--- libs/langgraph/langgraph/utils/runnable.py | 32 ++-- libs/langgraph/poetry.lock | 14 +- libs/langgraph/tests/test_prebuilt.py | 83 +++++++- 5 files changed, 271 insertions(+), 44 deletions(-) diff --git a/libs/langgraph/langgraph/prebuilt/__init__.py b/libs/langgraph/langgraph/prebuilt/__init__.py index 0554b8036..671804258 100644 --- a/libs/langgraph/langgraph/prebuilt/__init__.py +++ b/libs/langgraph/langgraph/prebuilt/__init__.py @@ -2,7 +2,12 @@ from langgraph.prebuilt.chat_agent_executor import create_react_agent from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation -from langgraph.prebuilt.tool_node import InjectedState, ToolNode, tools_condition +from langgraph.prebuilt.tool_node import ( + InjectedState, + InjectedStore, + ToolNode, + tools_condition, +) from langgraph.prebuilt.tool_validator import ValidationNode __all__ = [ @@ -13,4 +18,5 @@ __all__ = [ "tools_condition", "ValidationNode", "InjectedState", + "InjectedStore", ] diff --git a/libs/langgraph/langgraph/prebuilt/tool_node.py b/libs/langgraph/langgraph/prebuilt/tool_node.py index be87c0f0f..d01ce82db 100644 --- a/libs/langgraph/langgraph/prebuilt/tool_node.py +++ b/libs/langgraph/langgraph/prebuilt/tool_node.py @@ -13,6 +13,7 @@ from typing import ( Optional, Sequence, Tuple, + Type, Union, cast, ) @@ -28,10 +29,12 @@ 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 typing_extensions import Annotated, get_args, get_origin +from langgraph.store.base import BaseStore from langgraph.utils.runnable import RunnableCallable if TYPE_CHECKING: @@ -102,11 +105,15 @@ class ToolNode(RunnableCallable): ) -> None: 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, Optional[str]]] = {} + self.tool_to_store_arg: Dict[str, Optional[str]] = {} self.handle_tool_errors = handle_tool_errors for tool_ in tools: if not isinstance(tool_, BaseTool): tool_ = cast(BaseTool, create_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_) def _func( self, @@ -116,14 +123,30 @@ class ToolNode(RunnableCallable): BaseModel, ], config: RunnableConfig, + *, + store: BaseStore, ) -> Any: - tool_calls, output_type = self._parse_input(input) + tool_calls, output_type = self._parse_input(input, store) config_list = get_config_list(config, len(tool_calls)) with get_executor_for_config(config) as executor: outputs = [*executor.map(self._run_one, tool_calls, config_list)] # TypedDict, pydantic, dataclass, etc. should all be able to load from dict return outputs if output_type == "list" else {"messages": 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[ @@ -132,8 +155,10 @@ class ToolNode(RunnableCallable): BaseModel, ], config: RunnableConfig, + *, + store: BaseStore, ) -> Any: - tool_calls, output_type = self._parse_input(input) + tool_calls, output_type = self._parse_input(input, store) outputs = await asyncio.gather( *(self._arun_one(call, config) for call in tool_calls) ) @@ -184,6 +209,7 @@ class ToolNode(RunnableCallable): dict[str, Any], BaseModel, ], + store: BaseStore, ) -> Tuple[List[ToolCall], Literal["list", "dict"]]: if isinstance(input, list): output_type = "list" @@ -201,7 +227,9 @@ class ToolNode(RunnableCallable): if not isinstance(message, AIMessage): raise ValueError("Last message is not an AIMessage") - tool_calls = [self._inject_state(call, input) for call in message.tool_calls] + tool_calls = [ + self._inject_tool_args(call, input, store) for call in message.tool_calls + ] return tool_calls, output_type def _validate_tool_call(self, call: ToolCall) -> Optional[ToolMessage]: @@ -223,9 +251,7 @@ class ToolNode(RunnableCallable): BaseModel, ], ) -> ToolCall: - if tool_call["name"] not in self.tools_by_name: - return tool_call - state_args = _get_state_args(self.tools_by_name[tool_call["name"]]) + state_args = self.tool_to_state_args[tool_call["name"]] if state_args and isinstance(input, list): required_fields = list(state_args.values()) if ( @@ -255,12 +281,46 @@ class ToolNode(RunnableCallable): for tool_arg, state_field in state_args.items() } - tool_call_copy: ToolCall = copy(tool_call) - tool_call_copy["args"] = { - **tool_call_copy["args"], + tool_call["args"] = { + **tool_call["args"], **tool_state_args, } - return tool_call_copy + return tool_call + + def _inject_store(self, tool_call: ToolCall, store: BaseStore) -> ToolCall: + store_arg = self.tool_to_store_arg[tool_call["name"]] + if not store_arg: + return tool_call + + if store is None: + raise ValueError( + "Cannot inject store into tools with InjectedStore annotations - " + "please compile your graph with a store." + ) + + tool_call["args"] = { + **tool_call["args"], + store_arg: store, + } + return tool_call + + def _inject_tool_args( + self, + tool_call: ToolCall, + input: Union[ + list[AnyMessage], + dict[str, Any], + BaseModel, + ], + store: BaseStore, + ) -> ToolCall: + if tool_call["name"] not in self.tools_by_name: + return tool_call + + tool_call_copy: ToolCall = copy(tool_call) + tool_call_with_state = self._inject_state(tool_call_copy, input) + tool_call_with_store = self._inject_store(tool_call_with_state, store) + return tool_call_with_store def tools_condition( @@ -391,23 +451,78 @@ class InjectedState(InjectedToolArg): self.field = field +class InjectedStore(InjectedToolArg): + """Annotation for a Tool arg that is meant to be populated with LangGraph store. + + Any Tool argument annotated with InjectedStore will be hidden from a tool-calling + model, so that the model doesn't attempt to generate the argument. If using + ToolNode, the appropriate store field will be automatically injected into + the model-generated tool args. Note: if a graph is compiled with a store object, + the store will be automatically propagated to the tools with InjectedStore args + when using ToolNode. + + Example: + ```python + from typing import Any + from typing_extensions import Annotated + + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + + from langgraph.store.memory import InMemoryStore + from langgraph.prebuilt import InjectedStore, ToolNode + + store = InMemoryStore() + store.put(("values",), "foo", {"bar": 2}) + + @tool + def store_tool(x: int, my_store: Annotated[Any, InjectedStore()]) -> str: + '''Do something with store.''' + stored_value = my_store.get(("values",), "foo").value["bar"] + return stored_value + x + + node = ToolNode([store_tool]) + + tool_call = {"name": "store_tool", "args": {"x": 1}, "id": "1", "type": "tool_call"} + state = { + "messages": [AIMessage("", tool_calls=[tool_call])], + } + + node.invoke(state, store=store) + ``` + + ```pycon + { + "messages": [ + ToolMessage(content='3', name='store_tool', tool_call_id='1'), + ] + } + ``` + """ # noqa: E501 + + +def _is_injection( + type_arg: Any, injection_type: Union[Type[InjectedState], Type[InjectedStore]] +) -> bool: + if isinstance(type_arg, injection_type) or ( + isinstance(type_arg, type) and issubclass(type_arg, injection_type) + ): + return True + origin_ = get_origin(type_arg) + if origin_ is Union or origin_ is Annotated: + return any(_is_injection(ta, injection_type) for ta in get_args(type_arg)) + return False + + def _get_state_args(tool: BaseTool) -> Dict[str, Optional[str]]: full_schema = tool.get_input_schema() tool_args_to_state_fields: Dict = {} - def _is_injection(type_arg: Any) -> bool: - if isinstance(type_arg, InjectedState) or ( - isinstance(type_arg, type) and issubclass(type_arg, InjectedState) - ): - return True - origin_ = get_origin(type_arg) - if origin_ is Union or origin_ is Annotated: - return any(_is_injection(ta) for ta in get_args(type_arg)) - return False - for name, type_ in full_schema.__annotations__.items(): injections = [ - type_arg for type_arg in get_args(type_) if _is_injection(type_arg) + type_arg + for type_arg in get_args(type_) + if _is_injection(type_arg, InjectedState) ] if len(injections) > 1: raise ValueError( @@ -423,3 +538,24 @@ def _get_state_args(tool: BaseTool) -> Dict[str, Optional[str]]: else: pass return tool_args_to_state_fields + + +def _get_store_arg(tool: BaseTool) -> Optional[str]: + full_schema = tool.get_input_schema() + for name, type_ in full_schema.__annotations__.items(): + injections = [ + type_arg + for type_arg in get_args(type_) + if _is_injection(type_arg, InjectedStore) + ] + if len(injections) > 1: + ValueError( + "A tool argument should not be annotated with InjectedStore more than " + f"once. Received arg {name} with annotations {injections}." + ) + elif len(injections) == 1: + return name + else: + pass + + return None diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index 2545eb49e..376a6f3ee 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -150,13 +150,15 @@ class RunnableCallable(Runnable): kwargs["config"] = config _conf = config[CONF] for kw, _, ck, defv in KWARGS_CONFIG_KEYS: - if self.func_accepts[kw]: - if defv is inspect.Parameter.empty and ck not in _conf: - raise ValueError( - f"Missing required config key '{ck}' for '{self.name}'." - ) - else: - kwargs[kw] = _conf.get(ck, defv) + if not self.func_accepts[kw]: + continue + + if defv is inspect.Parameter.empty and kw not in kwargs and ck not in _conf: + raise ValueError( + f"Missing required config key '{ck}' for '{self.name}'." + ) + elif kwargs.get(kw) is None: + kwargs[kw] = _conf.get(ck, defv) context = copy_context() if self.trace: callback_manager = get_callback_manager_for_config(config, self.tags) @@ -195,13 +197,15 @@ class RunnableCallable(Runnable): kwargs["config"] = config _conf = config[CONF] for kw, _, ck, defv in KWARGS_CONFIG_KEYS: - if self.func_accepts[kw]: - if defv is inspect.Parameter.empty and ck not in _conf: - raise ValueError( - f"Missing required config key '{ck}' for '{self.name}'." - ) - else: - kwargs[kw] = _conf.get(ck, defv) + if not self.func_accepts[kw]: + continue + + if defv is inspect.Parameter.empty and kw not in kwargs and ck not in _conf: + raise ValueError( + f"Missing required config key '{ck}' for '{self.name}'." + ) + elif kwargs.get(kw) is None: + kwargs[kw] = _conf.get(ck, defv) context = copy_context() if self.trace: callback_manager = get_async_callback_manager_for_config(config, self.tags) diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index b32d88589..56b859257 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -1215,18 +1215,18 @@ files = [ [[package]] name = "langchain-core" -version = "0.3.0" +version = "0.3.8" description = "Building applications with LLMs through composability" optional = false python-versions = "<4.0,>=3.9" files = [ - {file = "langchain_core-0.3.0-py3-none-any.whl", hash = "sha256:bee6dae2366d037ef0c5b87401fed14b5497cad26f97724e8c9ca7bc9239e847"}, - {file = "langchain_core-0.3.0.tar.gz", hash = "sha256:1249149ea3ba24c9c761011483c14091573a5eb1a773aa0db9c8ad155dd4a69d"}, + {file = "langchain_core-0.3.8-py3-none-any.whl", hash = "sha256:07015f7b1d9f52eefe05130e8cafe4dcbdbbf72a8411c9edafe38422e4d11b5c"}, + {file = "langchain_core-0.3.8.tar.gz", hash = "sha256:7485904f7082f1df880d5ae470a488161616132f30d99f556a1877901fffd1cb"}, ] [package.dependencies] jsonpatch = ">=1.33,<2.0" -langsmith = ">=0.1.117,<0.2.0" +langsmith = ">=0.1.125,<0.2.0" packaging = ">=23.2,<25" pydantic = [ {version = ">=2.5.2,<3.0.0", markers = "python_full_version < \"3.12.4\""}, @@ -1291,13 +1291,13 @@ url = "../checkpoint-sqlite" [[package]] name = "langsmith" -version = "0.1.120" +version = "0.1.129" description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform." optional = false python-versions = "<4.0,>=3.8.1" files = [ - {file = "langsmith-0.1.120-py3-none-any.whl", hash = "sha256:54d2785e301646c0988e0a69ebe4d976488c87b41928b358cb153b6ddd8db62b"}, - {file = "langsmith-0.1.120.tar.gz", hash = "sha256:25499ca187b41bd89d784b272b97a8d76f60e0e21bdf20336e8a2aa6a9b23ac9"}, + {file = "langsmith-0.1.129-py3-none-any.whl", hash = "sha256:31393fbbb17d6be5b99b9b22d530450094fab23c6c37281a6a6efb2143d05347"}, + {file = "langsmith-0.1.129.tar.gz", hash = "sha256:6c3ba66471bef41b9f87da247cc0b493268b3f54656f73648a256a205261b6a0"}, ] [package.dependencies] diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 42ab73957..6fc7f0d35 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -35,11 +35,15 @@ from pydantic.v1 import BaseModel as BaseModelV1 from typing_extensions import TypedDict from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.graph import START, MessagesState, StateGraph from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent -from langgraph.prebuilt.tool_node import InjectedState +from langgraph.prebuilt.tool_node import InjectedState, InjectedStore +from langgraph.store.base import BaseStore +from langgraph.store.memory import InMemoryStore from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, ALL_CHECKPOINTERS_SYNC, + IS_LANGCHAIN_CORE_030_OR_GREATER, awith_checkpointer, ) from tests.messages import _AnyIdHumanMessage @@ -692,6 +696,83 @@ def test_tool_node_inject_state(schema_: Type[T]) -> None: assert tool_message.content == "hi?" +@pytest.mark.skipif( + not IS_LANGCHAIN_CORE_030_OR_GREATER, + reason="Langchain core 0.3.0 or greater is required", +) +def test_tool_node_inject_store() -> None: + store = InMemoryStore() + namespace = ("test",) + + def tool1(some_val: int, store: Annotated[BaseStore, InjectedStore()]) -> str: + """Tool 1 docstring.""" + store_val = store.get(namespace, "test_key").value["foo"] + return f"Some val: {some_val}, store val: {store_val}" + + def tool2(some_val: int, store: Annotated[BaseStore, InjectedStore()]) -> str: + """Tool 2 docstring.""" + store_val = store.get(namespace, "test_key").value["foo"] + return f"Some val: {some_val}, store val: {store_val}" + + def tool3( + some_val: int, + bar: Annotated[str, InjectedState("bar")], + store: Annotated[BaseStore, InjectedStore()], + ) -> str: + """Tool 3 docstring.""" + store_val = store.get(namespace, "test_key").value["foo"] + return f"Some val: {some_val}, store val: {store_val}, state val: {bar}" + + node = ToolNode([tool1, tool2, tool3], handle_tool_errors=True) + store.put(namespace, "test_key", {"foo": "bar"}) + + class State(MessagesState): + bar: str + + builder = StateGraph(State) + builder.add_node("tools", node) + builder.add_edge(START, "tools") + graph = builder.compile(store=store) + + for tool_name in ("tool1", "tool2"): + tool_call = { + "name": tool_name, + "args": {"some_val": 1}, + "id": "some 0", + "type": "tool_call", + } + msg = AIMessage("hi?", tool_calls=[tool_call]) + node_result = node.invoke({"messages": [msg]}, store=store) + graph_result = graph.invoke({"messages": [msg]}) + for result in (node_result, graph_result): + result["messages"][-1] + tool_message = result["messages"][-1] + assert ( + tool_message.content == "Some val: 1, store val: bar" + ), f"Failed for tool={tool_name}" + + tool_call = { + "name": "tool3", + "args": {"some_val": 1}, + "id": "some 0", + "type": "tool_call", + } + msg = AIMessage("hi?", tool_calls=[tool_call]) + node_result = node.invoke({"messages": [msg], "bar": "baz"}, store=store) + graph_result = graph.invoke({"messages": [msg], "bar": "baz"}) + for result in (node_result, graph_result): + result["messages"][-1] + tool_message = result["messages"][-1] + assert ( + tool_message.content == "Some val: 1, store val: bar, state val: baz" + ), f"Failed for tool={tool_name}" + + # test injected store without passing store to compiled graph + failing_graph = builder.compile() + with pytest.raises(ValueError): + failing_graph.invoke({"messages": [msg], "bar": "baz"}) + + def test_tool_node_ensure_utf8() -> None: @dec_tool def get_day_list(days: list[str]) -> list[str]: