mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
langgraph[patch]: InjectedState annotation (#1067)
Add annotated for injecting state vars into a Tool
This commit is contained in:
Bagatur
@@ -55,4 +55,13 @@ from langgraph.prebuilt import tools_condition
|
||||
from langgraph.prebuilt import ValidationNode
|
||||
```
|
||||
|
||||
::: langgraph.prebuilt.ValidationNode
|
||||
::: langgraph.prebuilt.ValidationNode
|
||||
|
||||
## InjectedState
|
||||
|
||||
```python
|
||||
from langgraph.prebuilt import InjectedState
|
||||
```
|
||||
|
||||
::: langgraph.prebuilt.InjectedState
|
||||
handler: python
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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
|
||||
|
||||
Generated
+1
-1
@@ -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"
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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?"
|
||||
|
||||
Reference in New Issue
Block a user