{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
   "metadata": {},
   "source": [
    "# How to pass runtime values to tools\n",
    "\n",
    "Sometimes, you want to let a tool-calling LLM populate a *subset* of the tool functions' arguments and provide the other values for the other arguments at runtime. If you're using LangChain-style [tools](https://python.langchain.com/docs/concepts/#tools), an easy way to handle this is by annotating function parameters with [InjectedArg](https://python.langchain.com/docs/how_to/tool_runtime/). This annotation excludes that parameter from being shown to the LLM.\n",
    "\n",
    "In LangGraph applications you might want to pass the graph state or [shared memory](https://langchain-ai.github.io/langgraph/how-tos/cross-thread-persistence/) (store) to the tools at runtime. This type of stateful tools is useful when a tool's output is affected by past agent steps (e.g. if you're using a sub-agent as a tool, and want to pass the message history in to the sub-agent), or when a tool's input needs to be validated given context from past agent steps.\n",
    "\n",
    "In this guide we'll demonstrate how to do so using LangGraph's prebuilt [ToolNode](https://langchain-ai.github.io/langgraph/how-tos/tool-calling/).\n",
    "\n",
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Prerequisites</p>\n",
    "    <p>\n",
    "        This guide targets **LangChain tool calling** assumes familiarity with the following:\n",
    "        <ul>\n",
    "            <li>\n",
    "                <a href=\"https://python.langchain.com/docs/concepts/#tools\">\n",
    "                    Tools\n",
    "                </a>\n",
    "            </li>\n",
    "            <li>\n",
    "                <a href=\"https://langchain-ai.github.io/langgraph/concepts/low_level/#state\">\n",
    "                    State\n",
    "                </a>\n",
    "            </li>\n",
    "            <li>\n",
    "                <a href=\"https://langchain-ai.github.io/langgraph/concepts/agentic_concepts/#tool-calling-agent\">\n",
    "                    Tool-calling\n",
    "                </a>\n",
    "            </li>\n",
    "        </ul>\n",
    "        You can still use tool calling in LangGraph using your provider SDK without losing any of LangGraph's core features.\n",
    "    </p>\n",
    "</div> \n",
    "\n",
    "The core technique the examples below is to **annotate** a parameter as \"injected\", meaning it will be injected by your program and should not be seen or populated by the LLM. Let the following codesnippet serve as a tl;dr:\n",
    "\n",
    "```python\n",
    "from typing import Annotated\n",
    "\n",
    "from langchain_core.runnables import RunnableConfig\n",
    "from langchain_core.tools import InjectedToolArg\n",
    "from langgraph.store.base import BaseStore\n",
    "\n",
    "from langgraph.prebuilt import InjectedState, InjectedStore\n",
    "\n",
    "\n",
    "# Can be sync or async; @tool decorator not required\n",
    "async def my_tool(\n",
    "    # These arguments are populated by the LLM\n",
    "    some_arg: str,\n",
    "    another_arg: float,\n",
    "    # The config: RunnableConfig is always available in LangChain calls\n",
    "    # This is not exposed to the LLM\n",
    "    config: RunnableConfig,\n",
    "    # The following three are specific to the prebuilt ToolNode\n",
    "    # (and `create_react_agent` by extension). If you are invoking the\n",
    "    # tool on its own (in your own node), then you would need to provide these yourself.\n",
    "    store: Annotated[BaseStore, InjectedStore],\n",
    "    # This passes in the full state.\n",
    "    state: Annotated[State, InjectedState],\n",
    "    # You can also inject single fields from your state if you\n",
    "    messages: Annotated[list, InjectedState(\"messages\")]\n",
    "    # The following is not compatible with create_react_agent or ToolNode\n",
    "    # You can also exclude other arguments from being shown to the model.\n",
    "    # These must be provided manually and are useful if you call the tools/functions in your own node\n",
    "    # some_other_arg=Annotated[\"MyPrivateClass\", InjectedToolArg],\n",
    "):\n",
    "    \"\"\"Call my_tool to have an impact on the real world.\n",
    "\n",
    "    Args:\n",
    "        some_arg: a very important argument\n",
    "        another_arg: another argument the LLM will provide\n",
    "    \"\"\" # The docstring becomes the description for your tool and is passed to the model\n",
    "    print(some_arg, another_arg, config, store, state, messages)\n",
    "    # Config, some_other_rag, store, and state  are all \"hidden\" from\n",
    "    # LangChain models when passed to bind_tools or with_structured_output\n",
    "    return \"... some response\"\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c205242",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "7cbd446a-808f-4394-be92-d45ab818953c",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First we need to install the packages required"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "af4ce0ba-7596-4e5f-8bf8-0b0bd6e62833",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install --quiet -U langgraph langchain-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0abe11f4-62ed-4dc4-8875-3db21e260d1d",
   "metadata": {},
   "source": [
    "Next, we need to set API keys for OpenAI (the chat model we will use)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c903a1cf-2977-4e2d-ad7d-8b3946821d89",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OPENAI_API_KEY:  ········\n"
     ]
    }
   ],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "\n",
    "_set_env(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0ed46a8-effe-4596-b0e1-a6a29ee16f5c",
   "metadata": {},
   "source": [
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Set up <a href=\"https://smith.langchain.com\">LangSmith</a> for LangGraph development</p>\n",
    "    <p style=\"padding-top: 5px;\">\n",
    "        Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started <a href=\"https://docs.smith.langchain.com\">here</a>. \n",
    "    </p>\n",
    "</div>    "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1a5908ab-7c1f-4832-85ae-511536b89b9f",
   "metadata": {},
   "source": [
    "## Pass graph state to tools\n",
    "\n",
    "Let's first take a look at how to give our tools access to the graph state. We'll need to define our graph state:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a2ded72f-1450-4295-a879-fd213ecfd4b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "# this is the state schema used by the prebuilt create_react_agent we'll be using below\n",
    "from langgraph.prebuilt.chat_agent_executor import AgentState\n",
    "from langchain_core.documents import Document\n",
    "\n",
    "\n",
    "class State(AgentState):\n",
    "    docs: List[str]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75d2d28c-0da3-46c6-9c67-979a71b5f517",
   "metadata": {},
   "source": [
    "### Define 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 `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": "markdown",
   "id": "2f7e2c8d",
   "metadata": {},
   "source": [
    "<div class=\"admonition note\">\n",
    "    <p class=\"admonition-title\">Using Pydantic with LangChain</p>\n",
    "    <p>\n",
    "        This notebook uses Pydantic v2 <code>BaseModel</code>, which requires <code>langchain-core >= 0.3</code>. Using <code>langchain-core < 0.3</code> will result in errors due to mixing of Pydantic v1 and v2 <code>BaseModels</code>.\n",
    "    </p>\n",
    "</div>  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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.messages import ToolMessage\n",
    "from langchain_core.tools import tool\n",
    "from langgraph.prebuilt import InjectedState\n",
    "\n",
    "\n",
    "@tool\n",
    "def get_context(question: str, state: Annotated[dict, InjectedState]):\n",
    "    \"\"\"Get relevant context for answering the question.\"\"\"\n",
    "    return \"\\n\\n\".join(doc for doc in state[\"docs\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c2d0de0-0f3e-4bbe-b0b6-cc0f70b11993",
   "metadata": {},
   "source": [
    "If we look at the input schemas for these tools, we'll see that `state` is still listed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1092929b-c939-4b2a-9f9c-e725b0e34af2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'description': 'Get relevant context for answering the question.',\n",
       " 'properties': {'question': {'title': 'Question', 'type': 'string'},\n",
       "  'state': {'title': 'State', 'type': 'object'}},\n",
       " 'required': ['question', 'state'],\n",
       " 'title': 'get_context',\n",
       " 'type': 'object'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_context.get_input_schema().schema()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e346a26e-e00b-48e5-82c5-c930ea6084a4",
   "metadata": {},
   "source": [
    "But if we look at the tool call schema, which is what is passed to the model for tool-calling, `state` has been removed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3912bb51-3107-4335-a659-021c5d89fb37",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'description': 'Get relevant context for answering the question.',\n",
       " 'properties': {'question': {'title': 'Question', 'type': 'string'}},\n",
       " 'required': ['question'],\n",
       " 'title': 'get_context',\n",
       " 'type': 'object'}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_context.tool_call_schema.schema()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e8b9211-93d0-4ad5-aa7a-9c09099c53ff",
   "metadata": {},
   "source": [
    "### Define the graph\n",
    "\n",
    "In this example we will be using a [prebuilt ReAct agent](https://langchain-ai.github.io/langgraph/how-tos/create-react-agent/). We'll first need to define our model and a tool-calling node ([ToolNode](https://langchain-ai.github.io/langgraph/reference/prebuilt/#langgraph.prebuilt.tool_node.ToolNode)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "be341be4-bdc3-4e78-9bc1-7486da27fd7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.prebuilt import ToolNode, create_react_agent\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o\", temperature=0)\n",
    "tools = [get_context]\n",
    "\n",
    "# ToolNode will automatically take care of injecting state into tools\n",
    "tool_node = ToolNode(tools)\n",
    "\n",
    "checkpointer = MemorySaver()\n",
    "graph = create_react_agent(model, tools, state_schema=State, checkpointer=checkpointer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "547c3931-3dae-4281-ad4e-4b51305594d4",
   "metadata": {},
   "source": [
    "### Use it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c0858273-10f2-45c4-a922-b11321ac3fae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what's the latest news about FooBar\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  get_context (call_UkqfR7z2cLJQjhatUpDeEa5H)\n",
      " Call ID: call_UkqfR7z2cLJQjhatUpDeEa5H\n",
      "  Args:\n",
      "    question: latest news about FooBar\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: get_context\n",
      "\n",
      "FooBar company just raised 1 Billion dollars!\n",
      "\n",
      "FooBar company was founded in 2019\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The latest news about FooBar is that the company has just raised 1 billion dollars.\n"
     ]
    }
   ],
   "source": [
    "docs = [\n",
    "    \"FooBar company just raised 1 Billion dollars!\",\n",
    "    \"FooBar company was founded in 2019\",\n",
    "]\n",
    "\n",
    "inputs = {\n",
    "    \"messages\": [{\"type\": \"user\", \"content\": \"what's the latest news about FooBar\"}],\n",
    "    \"docs\": docs,\n",
    "}\n",
    "config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "for chunk in graph.stream(inputs, config, stream_mode=\"values\"):\n",
    "    chunk[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09dd6d7d-d8a8-48f9-a30d-fba79fb7cb1e",
   "metadata": {},
   "source": [
    "## Pass shared memory (store) to the graph\n",
    "\n",
    "You might also want to give tools access to memory that is shared across multiple conversations or users. We can do it by passing LangGraph [Store](https://langchain-ai.github.io/langgraph/how-tos/cross-thread-persistence/) to the tools using a different annotation -- `InjectedStore`.\n",
    "\n",
    "Let's modify our example to save the documents in an in-memory store and retrieve them using `get_context` tool. We'll also make the documents accessible based on a user ID, so that some documents are only visible to certain users. The tool will then use the `user_id` provided in the [config](https://langchain-ai.github.io/langgraph/how-tos/pass-config-to-tools/) to retrieve a correct set of documents."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41644e7e-bda2-4c0d-95da-99572554c02d",
   "metadata": {},
   "source": [
    "<div class=\"admonition note\">\n",
    "    <p class=\"admonition-title\">Note</p>\n",
    "    <list>\n",
    "        <li>\n",
    "        Support for <code>Store</code> API and <code>InjectedStore</code> used in this notebook was added in LangGraph <code>v0.2.34</code>.\n",
    "        </li>\n",
    "        <li>\n",
    "        <code>InjectedStore</code> annotation requires <code>langchain-core >= 0.3.8</code>\n",
    "        </li>\n",
    "    <list>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "759186ed-5506-4e9a-80c4-0a2405ff58de",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.store.memory import InMemoryStore\n",
    "\n",
    "doc_store = InMemoryStore()\n",
    "\n",
    "namespace = (\"documents\", \"1\")  # user ID\n",
    "doc_store.put(\n",
    "    namespace, \"doc_0\", {\"doc\": \"FooBar company just raised 1 Billion dollars!\"}\n",
    ")\n",
    "namespace = (\"documents\", \"2\")  # user ID\n",
    "doc_store.put(namespace, \"doc_1\", {\"doc\": \"FooBar company was founded in 2019\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1a6e4e4-3256-4d42-aa5f-de337bfa6a97",
   "metadata": {},
   "source": [
    "### Define the tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "91ab7e2b-7df4-41ea-8b20-e2ba950aacb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.store.base import BaseStore\n",
    "from langchain_core.runnables import RunnableConfig\n",
    "from langgraph.prebuilt import InjectedStore\n",
    "\n",
    "\n",
    "@tool\n",
    "def get_context(\n",
    "    question: str,\n",
    "    config: RunnableConfig,\n",
    "    store: Annotated[BaseStore, InjectedStore()],\n",
    ") -> Tuple[str, List[Document]]:\n",
    "    \"\"\"Get relevant context for answering the question.\"\"\"\n",
    "    user_id = config.get(\"configurable\", {}).get(\"user_id\")\n",
    "    docs = [item.value[\"doc\"] for item in store.search((\"documents\", user_id))]\n",
    "    return \"\\n\\n\".join(doc for doc in docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcd29b33-b647-4e11-8942-2d47f129095b",
   "metadata": {},
   "source": [
    "We can also verify that the tool-calling model will ignore `store` arg of `get_context` tool:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "07bdc124-06bb-4c25-a18e-700ab7aa6521",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'description': 'Get relevant context for answering the question.',\n",
       " 'properties': {'question': {'title': 'Question', 'type': 'string'}},\n",
       " 'required': ['question'],\n",
       " 'title': 'get_context',\n",
       " 'type': 'object'}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_context.tool_call_schema.schema()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10622c58-db67-476b-81a3-1b964f5b471f",
   "metadata": {},
   "source": [
    "### Define the graph\n",
    "\n",
    "Let's update our ReAct agent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "10c70125-a105-4103-89d2-2e93b401080c",
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [get_context]\n",
    "\n",
    "# ToolNode will automatically take care of injecting Store into tools\n",
    "tool_node = ToolNode(tools)\n",
    "\n",
    "checkpointer = MemorySaver()\n",
    "# NOTE: we need to pass our store to `create_react_agent` to make sure our graph is aware of it\n",
    "graph = create_react_agent(model, tools, checkpointer=checkpointer, store=doc_store)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35ba8adc-c706-48a6-992f-24c03c4cee46",
   "metadata": {},
   "source": [
    "### Use it!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c155ef73-91fe-459d-a3af-7ed254c19e23",
   "metadata": {},
   "source": [
    "Let's try running our graph with a `\"user_id\"` in the config."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d683986f-faf4-4724-a13f-bac39ea9bafe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what's the latest news about FooBar\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  get_context (call_ocyHBpGgF3LPFOgRKURBfkGG)\n",
      " Call ID: call_ocyHBpGgF3LPFOgRKURBfkGG\n",
      "  Args:\n",
      "    question: latest news about FooBar\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: get_context\n",
      "\n",
      "FooBar company just raised 1 Billion dollars!\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The latest news about FooBar is that the company has just raised 1 billion dollars.\n"
     ]
    }
   ],
   "source": [
    "messages = [{\"type\": \"user\", \"content\": \"what's the latest news about FooBar\"}]\n",
    "config = {\"configurable\": {\"thread_id\": \"1\", \"user_id\": \"1\"}}\n",
    "for chunk in graph.stream({\"messages\": messages}, config, stream_mode=\"values\"):\n",
    "    chunk[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fac8ead-a4fa-488f-835d-21872c67ec72",
   "metadata": {},
   "source": [
    "We can see that the tool only retrieved the correct document for user \"1\" when looking up the information in the store. Let's now try it again for a different user:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2e3fd1e2-cc19-4023-8ffa-b0fc13da9e09",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what's the latest news about FooBar\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  get_context (call_zxO9KVlL8UxFQUMb8ETeHNvs)\n",
      " Call ID: call_zxO9KVlL8UxFQUMb8ETeHNvs\n",
      "  Args:\n",
      "    question: latest news about FooBar\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: get_context\n",
      "\n",
      "FooBar company was founded in 2019\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "FooBar company was founded in 2019. If you need more specific or recent news, please let me know!\n"
     ]
    }
   ],
   "source": [
    "messages = [{\"type\": \"user\", \"content\": \"what's the latest news about FooBar\"}]\n",
    "config = {\"configurable\": {\"thread_id\": \"2\", \"user_id\": \"2\"}}\n",
    "for chunk in graph.stream({\"messages\": messages}, config, stream_mode=\"values\"):\n",
    "    chunk[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "196fd8ec-38e7-4300-88f6-9f8774bba313",
   "metadata": {},
   "source": [
    "We can see that the tool pulled in a different document this time."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}