WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-27 20:15:00 +02:00
Code Issues Packages Projects Releases Wiki Activity

langgraph: add support for store in ToolNode (#1968)

Browse Source
This commit is contained in:
Vadym Barda
2024-10-02 17:39:25 +00:00
committed by GitHub
parent 74b36adb18
commit 7efd3c726e
5 changed files with 271 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/prebuilt/__init__.py
+7 -1
View File
@@ -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",
]
libs/langgraph/langgraph/prebuilt/tool_node.py
+157 -21
View File
@@ -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
libs/langgraph/langgraph/utils/runnable.py
+18 -14
View File
@@ -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)
libs/langgraph/poetry.lock
Generated
+7 -7
View File
@@ -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]
libs/langgraph/tests/test_prebuilt.py
+82 -1
View File
@@ -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]: