mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-11 04:07:52 +02:00
langgraph[patch]: InjectedState annotation (#1067)
Add annotated for injecting state vars into a Tool
This commit is contained in:
Bagatur
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user