From 610b6cc78cf682e4dbc7c155fbda04227136ebf0 Mon Sep 17 00:00:00 2001 From: Bagatur <22008038+baskaryan@users.noreply.github.com> Date: Fri, 19 Jul 2024 20:07:21 -0700 Subject: [PATCH] langgraph[patch]: InjectedState annotation (#1067) Add annotated for injecting state vars into a Tool --- docs/docs/reference/prebuilt.md | 11 +- examples/pass-run-time-values-to-tools.ipynb | 105 +++++------- libs/langgraph/langgraph/prebuilt/__init__.py | 3 +- .../langgraph/langgraph/prebuilt/tool_node.py | 159 ++++++++++++++++-- libs/langgraph/poetry.lock | 2 +- libs/langgraph/pyproject.toml | 2 +- libs/langgraph/tests/test_prebuilt.py | 85 ++++++++-- 7 files changed, 277 insertions(+), 90 deletions(-) diff --git a/docs/docs/reference/prebuilt.md b/docs/docs/reference/prebuilt.md index fc276aea0..4bfc9b314 100644 --- a/docs/docs/reference/prebuilt.md +++ b/docs/docs/reference/prebuilt.md @@ -55,4 +55,13 @@ from langgraph.prebuilt import tools_condition from langgraph.prebuilt import ValidationNode ``` -::: langgraph.prebuilt.ValidationNode \ No newline at end of file +::: langgraph.prebuilt.ValidationNode + +## InjectedState + +```python +from langgraph.prebuilt import InjectedState +``` + +::: langgraph.prebuilt.InjectedState + handler: python diff --git a/examples/pass-run-time-values-to-tools.ipynb b/examples/pass-run-time-values-to-tools.ipynb index 48d575aa1..3c74bd06d 100644 --- a/examples/pass-run-time-values-to-tools.ipynb +++ b/examples/pass-run-time-values-to-tools.ipynb @@ -86,30 +86,30 @@ "source": [ "## Defining the tools\n", "\n", - "We'll want our tool to take graph state as an input, but we don't want the model to try to generate this input when calling the tool. We can use the `InjectedToolArg` annotation to mark `state` as being injected at runtime. Any argument annotated with `InjectedToolArg` will not be generated by the model.\n", + "We'll want our tool to take graph state as an input, but we don't want the model to try to generate this input when calling the tool. We can use the `InjectedState` annotation to mark arguments as required graph state (or some field of graph state. These arguments will not be generated by the model. When using `ToolNode`, graph state will automatically be passed in to the relevant tools and arguments.\n", "\n", "In this example we'll create a tool that returns Documents and then another tool that actually cites the Documents that justify a claim." ] }, { "cell_type": "code", - "execution_count": 63, + "execution_count": 6, "id": "1d36e782-80f4-4334-b7d7-ee4c79864480", "metadata": {}, "outputs": [], "source": [ "from typing import List, Tuple\n", + "from typing_extensions import Annotated\n", "\n", "from langchain_core.documents import Document\n", "from langchain_core.pydantic_v1 import BaseModel\n", - "from langchain_core.tools import InjectedToolArg, tool\n", - "from typing_extensions import Annotated\n", + "from langchain_core.tools import tool\n", + "\n", + "from langgraph.prebuilt import InjectedState\n", "\n", "\n", "@tool(parse_docstring=True, response_format=\"content_and_artifact\")\n", - "def get_context(\n", - " question: List[str], state: Annotated[dict, InjectedToolArg]\n", - ") -> Tuple[str, List[Document]]:\n", + "def get_context(question: List[str]) -> Tuple[str, List[Document]]:\n", " \"\"\"Get context on the question.\n", "\n", " Args:\n", @@ -136,7 +136,7 @@ "\n", "@tool(parse_docstring=True, response_format=\"content_and_artifact\")\n", "def cite_context_sources(\n", - " claim: str, state: Annotated[dict, InjectedToolArg]\n", + " claim: str, state: Annotated[dict, InjectedState]\n", ") -> Tuple[str, List[Document]]:\n", " \"\"\"Cite which source a claim was based on.\n", "\n", @@ -175,31 +175,30 @@ }, { "cell_type": "code", - "execution_count": 64, + "execution_count": 9, "id": "1092929b-c939-4b2a-9f9c-e725b0e34af2", "metadata": {}, "outputs": [ { "data": { "text/plain": [ - "{'title': 'get_contextSchema',\n", - " 'description': 'Get context on the question.',\n", + "{'title': 'cite_context_sourcesSchema',\n", + " 'description': 'Cite which source a claim was based on.',\n", " 'type': 'object',\n", - " 'properties': {'question': {'title': 'Question',\n", - " 'description': 'The user question',\n", - " 'type': 'array',\n", - " 'items': {'type': 'string'}},\n", + " 'properties': {'claim': {'title': 'Claim',\n", + " 'description': 'The claim that was made.',\n", + " 'type': 'string'},\n", " 'state': {'title': 'State', 'type': 'object'}},\n", - " 'required': ['question', 'state']}" + " 'required': ['claim', 'state']}" ] }, - "execution_count": 64, + "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "get_context.get_input_schema().schema()" + "cite_context_sources.get_input_schema().schema()" ] }, { @@ -212,30 +211,29 @@ }, { "cell_type": "code", - "execution_count": 65, + "execution_count": 11, "id": "3912bb51-3107-4335-a659-021c5d89fb37", "metadata": {}, "outputs": [ { "data": { "text/plain": [ - "{'title': 'get_context',\n", - " 'description': 'Get context on the question.',\n", + "{'title': 'cite_context_sources',\n", + " 'description': 'Cite which source a claim was based on.',\n", " 'type': 'object',\n", - " 'properties': {'question': {'title': 'Question',\n", - " 'description': 'The user question',\n", - " 'type': 'array',\n", - " 'items': {'type': 'string'}}},\n", - " 'required': ['question']}" + " 'properties': {'claim': {'title': 'Claim',\n", + " 'description': 'The claim that was made.',\n", + " 'type': 'string'}},\n", + " 'required': ['claim']}" ] }, - "execution_count": 65, + "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "get_context.tool_call_schema.schema()" + "cite_context_sources.tool_call_schema.schema()" ] }, { @@ -258,7 +256,7 @@ }, { "cell_type": "code", - "execution_count": 66, + "execution_count": 12, "id": "ea793afa-2eab-4901-910d-6eed90cd6564", "metadata": {}, "outputs": [], @@ -302,7 +300,7 @@ }, { "cell_type": "code", - "execution_count": 67, + "execution_count": 18, "id": "3b541bb9-900c-40d0-964d-7b5dfee30667", "metadata": {}, "outputs": [], @@ -312,7 +310,7 @@ "from langchain_core.messages import ToolMessage\n", "from langchain_openai import ChatOpenAI\n", "\n", - "from langgraph.prebuilt import ToolExecutor, ToolInvocation\n", + "from langgraph.prebuilt import ToolNode\n", "\n", "model = ChatOpenAI(model=\"gpt-4o\", temperature=0)\n", "\n", @@ -330,8 +328,6 @@ "\n", "\n", "tools = [get_context, cite_context_sources]\n", - "tool_map = {tool_.name: tool_ for tool_ in tools}\n", - "\n", "\n", "# Define the function that calls the model\n", "def call_model(state, config):\n", @@ -342,25 +338,8 @@ " return {\"messages\": [response]}\n", "\n", "\n", - "# Helper function for adding state to each tool call's arguments\n", - "def inject_state(message, state):\n", - " tool_calls = []\n", - " for tool_call in message.tool_calls:\n", - " tool_call_copy = deepcopy(tool_call)\n", - " tool_call_copy[\"args\"][\"state\"] = state\n", - " tool_calls.append(tool_call_copy)\n", - " return tool_calls\n", - "\n", - "\n", - "# Define the function to execute tools\n", - "def call_tool(state, config):\n", - " messages = state[\"messages\"]\n", - " last_message = messages[-1]\n", - " tool_messages = []\n", - " for tool_call in inject_state(last_message, state):\n", - " tool_messages.append(tool_map[tool_call[\"name\"]].invoke(tool_call, config))\n", - " # We return a list, because this will get added to the existing list\n", - " return {\"messages\": tool_messages}" + "# ToolNode will automatically take care of injecting state into tools\n", + "tool_node = ToolNode(tools)" ] }, { @@ -375,7 +354,7 @@ }, { "cell_type": "code", - "execution_count": 68, + "execution_count": 19, "id": "813ae66c-3b58-4283-a02a-36da72a2ab90", "metadata": {}, "outputs": [], @@ -387,7 +366,7 @@ "\n", "# Define the two nodes we will cycle between\n", "workflow.add_node(\"agent\", call_model)\n", - "workflow.add_node(\"action\", call_tool)\n", + "workflow.add_node(\"action\", tool_node)\n", "\n", "# Set the entrypoint as `agent`\n", "# This means that this node is the first one called\n", @@ -426,7 +405,7 @@ }, { "cell_type": "code", - "execution_count": 69, + "execution_count": 20, "id": "a8afd6ef", "metadata": {}, "outputs": [ @@ -464,7 +443,7 @@ }, { "cell_type": "code", - "execution_count": 70, + "execution_count": 21, "id": "8edb04b9-40b6-46f1-a7a8-4b2d8aba7752", "metadata": {}, "outputs": [ @@ -474,19 +453,19 @@ "text": [ "Output from node 'agent':\n", "---\n", - "{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_aFUFt3TdazRnmD3FTZfxFAgL', 'function': {'arguments': '{\"question\":[\"what\\'s the latest news about FooBar\"]}', 'name': 'get_context'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 22, 'prompt_tokens': 87, 'total_tokens': 109}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-adf99f00-a903-49f2-b0c3-37b84b9b801f-0', tool_calls=[{'name': 'get_context', 'args': {'question': [\"what's the latest news about FooBar\"]}, 'id': 'call_aFUFt3TdazRnmD3FTZfxFAgL', 'type': 'tool_call'}], usage_metadata={'input_tokens': 87, 'output_tokens': 22, 'total_tokens': 109})]}\n", + "{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_BidVTw5NiW2wp8Ez7m8dDoHI', 'function': {'arguments': '{\"question\":[\"latest news about FooBar\"]}', 'name': 'get_context'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 19, 'prompt_tokens': 87, 'total_tokens': 106}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-fcac1b73-563e-4f4c-b1b0-626f55d377be-0', tool_calls=[{'name': 'get_context', 'args': {'question': ['latest news about FooBar']}, 'id': 'call_BidVTw5NiW2wp8Ez7m8dDoHI', 'type': 'tool_call'}], usage_metadata={'input_tokens': 87, 'output_tokens': 19, 'total_tokens': 106})]}\n", "\n", "---\n", "\n", "Output from node 'action':\n", "---\n", - "{'messages': [ToolMessage(content=\"FooBar company just raised 1 Billion dollars!\\n\\nFooBar company is now only hiring AI's\\n\\nFooBar company was founded in 2019\\n\\nFooBar company makes friendly robots\", name='get_context', tool_call_id='call_aFUFt3TdazRnmD3FTZfxFAgL', artifact=[Document(metadata={'source': 'twitter'}, page_content='FooBar company just raised 1 Billion dollars!'), Document(metadata={'source': 'twitter'}, page_content=\"FooBar company is now only hiring AI's\"), Document(metadata={'source': 'wikipedia'}, page_content='FooBar company was founded in 2019'), Document(metadata={'source': 'wikipedia'}, page_content='FooBar company makes friendly robots')])]}\n", + "{'messages': [ToolMessage(content=\"FooBar company just raised 1 Billion dollars!\\n\\nFooBar company is now only hiring AI's\\n\\nFooBar company was founded in 2019\\n\\nFooBar company makes friendly robots\", name='get_context', tool_call_id='call_BidVTw5NiW2wp8Ez7m8dDoHI', artifact=[Document(metadata={'source': 'twitter'}, page_content='FooBar company just raised 1 Billion dollars!'), Document(metadata={'source': 'twitter'}, page_content=\"FooBar company is now only hiring AI's\"), Document(metadata={'source': 'wikipedia'}, page_content='FooBar company was founded in 2019'), Document(metadata={'source': 'wikipedia'}, page_content='FooBar company makes friendly robots')])]}\n", "\n", "---\n", "\n", "Output from node 'agent':\n", "---\n", - "{'messages': [AIMessage(content='The latest news about FooBar is that the company just raised 1 billion dollars!', response_metadata={'token_usage': {'completion_tokens': 18, 'prompt_tokens': 153, 'total_tokens': 171}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'stop', 'logprobs': None}, id='run-c229a397-fda3-415b-a188-1416fd5f21b7-0', usage_metadata={'input_tokens': 153, 'output_tokens': 18, 'total_tokens': 171})]}\n", + "{'messages': [AIMessage(content='The latest news about FooBar is that the company has just raised 1 billion dollars!', response_metadata={'token_usage': {'completion_tokens': 19, 'prompt_tokens': 150, 'total_tokens': 169}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'stop', 'logprobs': None}, id='run-a8407471-7715-4c16-bd46-c29e5751e882-0', usage_metadata={'input_tokens': 150, 'output_tokens': 19, 'total_tokens': 169})]}\n", "\n", "---\n", "\n" @@ -509,7 +488,7 @@ }, { "cell_type": "code", - "execution_count": 71, + "execution_count": 22, "id": "4a2128ed-e23f-4f25-a026-0c6590f01a1c", "metadata": {}, "outputs": [ @@ -519,19 +498,19 @@ "text": [ "Output from node 'agent':\n", "---\n", - "{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_qqB4kucZnVhrZ5mJSH1dF8Lb', 'function': {'arguments': '{\"claim\":\"The latest news about FooBar is that the company just raised 1 billion dollars!\"}', 'name': 'cite_context_sources'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 32, 'prompt_tokens': 185, 'total_tokens': 217}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-686d4706-81c9-4ca0-8f09-d9af02f4ad7f-0', tool_calls=[{'name': 'cite_context_sources', 'args': {'claim': 'The latest news about FooBar is that the company just raised 1 billion dollars!'}, 'id': 'call_qqB4kucZnVhrZ5mJSH1dF8Lb', 'type': 'tool_call'}], usage_metadata={'input_tokens': 185, 'output_tokens': 32, 'total_tokens': 217})]}\n", + "{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_EB0zaQypXMqEUzaqwflUr0zH', 'function': {'arguments': '{\"claim\":\"FooBar company just raised 1 Billion dollars!\"}', 'name': 'cite_context_sources'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 25, 'prompt_tokens': 183, 'total_tokens': 208}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_c4e5b6fa31', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b4952777-e2b3-4448-be87-200e6e80981b-0', tool_calls=[{'name': 'cite_context_sources', 'args': {'claim': 'FooBar company just raised 1 Billion dollars!'}, 'id': 'call_EB0zaQypXMqEUzaqwflUr0zH', 'type': 'tool_call'}], usage_metadata={'input_tokens': 183, 'output_tokens': 25, 'total_tokens': 208})]}\n", "\n", "---\n", "\n", "Output from node 'action':\n", "---\n", - "{'messages': [ToolMessage(content='twitter', name='cite_context_sources', tool_call_id='call_qqB4kucZnVhrZ5mJSH1dF8Lb', artifact=[Document(metadata={'source': 'twitter'}, page_content='FooBar company just raised 1 Billion dollars!')])]}\n", + "{'messages': [ToolMessage(content='twitter', name='cite_context_sources', tool_call_id='call_EB0zaQypXMqEUzaqwflUr0zH', artifact=[Document(metadata={'source': 'twitter'}, page_content='FooBar company just raised 1 Billion dollars!')])]}\n", "\n", "---\n", "\n", "Output from node 'agent':\n", "---\n", - "{'messages': [AIMessage(content='The information about FooBar raising 1 billion dollars came from Twitter.', response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 227, 'total_tokens': 242}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_18cc0f1fa0', 'finish_reason': 'stop', 'logprobs': None}, id='run-343ad465-9a62-4d72-91bf-ab29c4fe8781-0', usage_metadata={'input_tokens': 227, 'output_tokens': 15, 'total_tokens': 242})]}\n", + "{'messages': [AIMessage(content='The information that FooBar company just raised 1 billion dollars comes from Twitter.', response_metadata={'token_usage': {'completion_tokens': 17, 'prompt_tokens': 218, 'total_tokens': 235}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_400f27fa1f', 'finish_reason': 'stop', 'logprobs': None}, id='run-a0dede05-dadd-46f6-8654-746520d4cef8-0', usage_metadata={'input_tokens': 218, 'output_tokens': 17, 'total_tokens': 235})]}\n", "\n", "---\n", "\n" diff --git a/libs/langgraph/langgraph/prebuilt/__init__.py b/libs/langgraph/langgraph/prebuilt/__init__.py index b05e0e660..615de4c7b 100644 --- a/libs/langgraph/langgraph/prebuilt/__init__.py +++ b/libs/langgraph/langgraph/prebuilt/__init__.py @@ -3,7 +3,7 @@ from langgraph.prebuilt import chat_agent_executor from langgraph.prebuilt.agent_executor import create_agent_executor from langgraph.prebuilt.chat_agent_executor import create_react_agent from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation -from langgraph.prebuilt.tool_node import ToolNode, tools_condition +from langgraph.prebuilt.tool_node import InjectedState, ToolNode, tools_condition from langgraph.prebuilt.tool_validator import ValidationNode __all__ = [ @@ -15,4 +15,5 @@ __all__ = [ "ToolNode", "tools_condition", "ValidationNode", + "InjectedState", ] diff --git a/libs/langgraph/langgraph/prebuilt/tool_node.py b/libs/langgraph/langgraph/prebuilt/tool_node.py index fc70d261b..72c0d37e2 100644 --- a/libs/langgraph/langgraph/prebuilt/tool_node.py +++ b/libs/langgraph/langgraph/prebuilt/tool_node.py @@ -1,11 +1,24 @@ import asyncio -from typing import Any, Callable, Dict, Literal, Optional, Sequence, Tuple, Union, cast +from copy import copy +from typing import ( + Any, + Callable, + Dict, + List, + Literal, + Optional, + Sequence, + Tuple, + Union, + cast, +) from langchain_core.messages import AIMessage, AnyMessage, ToolCall, ToolMessage from langchain_core.runnables import RunnableConfig from langchain_core.runnables.config import get_config_list, get_executor_for_config -from langchain_core.tools import BaseTool +from langchain_core.tools import BaseTool, InjectedToolArg from langchain_core.tools import tool as create_tool +from typing_extensions import get_args from langgraph.utils import RunnableCallable @@ -60,18 +73,18 @@ class ToolNode(RunnableCallable): def _func( self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig ) -> Any: - message, output_type = self._parse_input(input) - config_list = get_config_list(config, len(message.tool_calls)) + tool_calls, output_type = self._parse_input(input) + config_list = get_config_list(config, len(tool_calls)) with get_executor_for_config(config) as executor: - outputs = [*executor.map(self._run_one, message.tool_calls, config_list)] + outputs = [*executor.map(self._run_one, tool_calls, config_list)] return outputs if output_type == "list" else {"messages": outputs} async def _afunc( self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig ) -> Any: - message, output_type = self._parse_input(input) + tool_calls, output_type = self._parse_input(input) outputs = await asyncio.gather( - *(self._arun_one(call, config) for call in message.tool_calls) + *(self._arun_one(call, config) for call in tool_calls) ) return outputs if output_type == "list" else {"messages": outputs} @@ -102,7 +115,7 @@ class ToolNode(RunnableCallable): def _parse_input( self, input: Union[list[AnyMessage], dict[str, Any]] - ) -> Tuple[AIMessage, Literal["list", "dict"]]: + ) -> Tuple[List[ToolCall], Literal["list", "dict"]]: if isinstance(input, list): output_type = "list" message: AnyMessage = input[-1] @@ -114,8 +127,12 @@ class ToolNode(RunnableCallable): if not isinstance(message, AIMessage): raise ValueError("Last message is not an AIMessage") - else: - return cast(AIMessage, message), output_type + + tool_calls = [ + self._inject_state(call, input) + for call in cast(AIMessage, message).tool_calls + ] + return tool_calls, output_type def _validate_tool_call(self, call: ToolCall) -> Optional[ToolMessage]: if (requested_tool := call["name"]) not in self.tools_by_name: @@ -127,6 +144,39 @@ class ToolNode(RunnableCallable): else: return None + def _inject_state( + self, tool_call: ToolCall, input: Union[list[AnyMessage], dict[str, Any]] + ) -> 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"]]) + if state_args and not isinstance(input, dict): + required_fields = list(state_args.values()) + if ( + len(required_fields) == 1 + and required_fields[0] == "messages" + or required_fields[0] is None + ): + input = {"messages": input} + else: + err_msg = ( + f"Invalid input to ToolNode. Tool {tool_call['name']} requires " + f"graph state dict as input." + ) + if any(state_field for state_field in state_args.values()): + required_fields_str = ", ".join(f for f in required_fields if f) + err_msg += f" State should contain fields {required_fields_str}." + raise ValueError(err_msg) + tool_call_copy: ToolCall = copy(tool_call) + tool_call_copy["args"] = { + **tool_call_copy["args"], + **{ + tool_arg: cast(dict, input)[state_field] if state_field else input + for tool_arg, state_field in state_args.items() + }, + } + return tool_call_copy + def tools_condition( state: Union[list[AnyMessage], dict[str, Any]], @@ -183,3 +233,92 @@ def tools_condition( if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0: return "tools" return "__end__" + + +class InjectedState(InjectedToolArg): + """Annotation for a Tool arg that is meant to be populated with the graph state. + + Any Tool argument annotated with InjectedState will be hidden from a tool-calling + model, so that the model doesn't attempt to generate the argument. If using + ToolNode, the appropriate graph state field will be automatically injected into + the model-generated tool args. + + Args: + field: The key from state to insert. If None, the entire state is expected to + be passed in. + + Example: + ```python + from typing import List + from typing_extensions import Annotated, TypedDict + + from langchain_core.messages import BaseMessage, AIMessage + from langchain_core.tools import tool + + from langgraph.prebuilt import InjectedState, ToolNode + + + class AgentState(TypedDict): + messages: List[BaseMessage] + foo: str + + @tool + def state_tool(x: int, state: Annotated[dict, InjectedState]) -> str: + '''Do something with state.''' + if len(state["messages"]) > 2: + return state["foo"] + str(x) + else: + return "not enough messages" + + @tool + def foo_tool(x: int, foo: Annotated[str, InjectedState("foo")]) -> str: + '''Do something else with state.''' + return foo + str(x + 1) + + node = ToolNode([state_tool, foo_tool]) + + tool_call1 = {"name": "state_tool", "args": {"x": 1}, "id": "1", "type": "tool_call"} + tool_call2 = {"name": "foo_tool", "args": {"x": 1}, "id": "2", "type": "tool_call"} + state = { + "messages": [AIMessage("", tool_calls=[tool_call1, tool_call2])], + "foo": "bar", + } + node.invoke(state) + ``` + + ```pycon + [ + ToolMessage(content='not enough messages', name='state_tool', tool_call_id='1'), + ToolMessage(content='bar2', name='foo_tool', tool_call_id='2') + ] + ``` + """ # noqa: E501 + + def __init__(self, field: Optional[str] = None) -> None: + self.field = field + + +def _get_state_args(tool: BaseTool) -> Dict[str, Optional[str]]: + full_schema = tool.get_input_schema() + tool_args_to_state_fields: Dict = {} + for name, type_ in full_schema.__annotations__.items(): + injections = [ + type_arg + for type_arg in get_args(type_) + if isinstance(type_arg, InjectedState) + or (isinstance(type_arg, type) and issubclass(type_arg, InjectedState)) + ] + if len(injections) > 1: + raise ValueError( + "A tool argument should not be annotated with InjectedState more than " + f"once. Received arg {name} with annotations {injections}." + ) + elif len(injections) == 1: + injection = injections[0] + if isinstance(injection, InjectedState) and injection.field: + tool_args_to_state_fields[name] = injection.field + else: + tool_args_to_state_fields[name] = None + else: + pass + return tool_args_to_state_fields diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index 3bee7a407..94075d612 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -4165,4 +4165,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "5fb6190a1b01d0cd351ea9a0023c8c8d6acf4fe831101ab87f9fc41308f74b74" +content-hash = "0d877d3879473de43aca1e1d36a8f420ff3f4b140807cb5cc24935d9114947be" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 7b4ab2078..593e659c8 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/langgraph" [tool.poetry.dependencies] python = ">=3.9.0,<4.0" -langchain-core = ">=0.2.19,<0.3" +langchain-core = ">=0.2.22,<0.3" [tool.poetry.group.dev.dependencies] diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index a0a8436e5..e21766cd0 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -1,15 +1,11 @@ -from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union +from typing import Annotated, Any, Callable, Dict, List, Optional, Sequence, Type, Union import pytest -from langchain_core.callbacks import ( - CallbackManagerForLLMRun, -) -from langchain_core.language_models import ( - BaseChatModel, - LanguageModelInput, -) +from langchain_core.callbacks import CallbackManagerForLLMRun +from langchain_core.language_models import BaseChatModel, LanguageModelInput from langchain_core.messages import ( AIMessage, + AnyMessage, BaseMessage, HumanMessage, SystemMessage, @@ -23,11 +19,8 @@ from langchain_core.tools import tool as dec_tool from pydantic import BaseModel as BaseModelV2 from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.prebuilt import ( - ToolNode, - ValidationNode, - create_react_agent, -) +from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent +from langgraph.prebuilt.tool_node import InjectedState from tests.any_str import AnyStr from tests.memory_assert import MemorySaverAssertImmutable @@ -453,3 +446,69 @@ async def test_validation_node(tool_schema: Any, use_message_key: bool): if use_message_key: result_sync = result_sync["messages"] check_results(result_sync) + + +def test_tool_node_inject_state() -> None: + def tool1(some_val: int, state: Annotated[dict, InjectedState]) -> str: + """Tool 1 docstring.""" + return state["foo"] + + def tool2(some_val: int, state: Annotated[dict, InjectedState()]) -> str: + """Tool 1 docstring.""" + return state["foo"] + + def tool3( + some_val: int, + foo: Annotated[str, InjectedState("foo")], + msgs: Annotated[List[AnyMessage], InjectedState("messages")], + ) -> str: + """Tool 1 docstring.""" + return foo + + def tool4( + some_val: int, msgs: Annotated[List[AnyMessage], InjectedState("messages")] + ) -> str: + """Tool 1 docstring.""" + return msgs[0].content + + node = ToolNode([tool1, tool2, tool3, tool4]) + for tool_name in ("tool1", "tool2", "tool3"): + tool_call = { + "name": tool_name, + "args": {"some_val": 1}, + "id": "some 0", + "type": "tool_call", + } + msg = AIMessage("hi?", tool_calls=[tool_call]) + result = node.invoke({"messages": [msg], "foo": "bar"}) + tool_message = result["messages"][-1] + assert tool_message.content == "bar" + + if tool_name == "tool3": + with pytest.raises(KeyError): + node.invoke({"messages": [msg], "notfoo": "bar"}) + + with pytest.raises(ValueError): + node.invoke([msg]) + else: + tool_message = node.invoke({"messages": [msg], "notfoo": "bar"})[ + "messages" + ][-1] + assert "KeyError" in tool_message.content + tool_message = node.invoke([msg])[-1] + assert "KeyError" in tool_message.content + + tool_call = { + "name": "tool4", + "args": {"some_val": 1}, + "id": "some 0", + "type": "tool_call", + } + msg = AIMessage("hi?", tool_calls=[tool_call]) + result = node.invoke({"messages": [msg]}) + tool_message = result["messages"][-1] + assert tool_message.content == "hi?" + + result = node.invoke([msg]) + tool_message = result[-1] + assert tool_message.content == "hi?"