{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to pass config to tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may need to pass values to a tool that are only known at runtime. For example, the tool logic may require using the ID of the user who made the request.\n",
    "\n",
    "Most of the time, such values should not be controlled by the LLM. In fact, allowing the LLM to control the user ID may lead to a security risk.\n",
    "\n",
    "Instead, the LLM should only control the parameters of the tool that are meant to be controlled by the LLM, while other parameters (such as user ID) should be fixed by the application logic.\n",
    "\n",
    "To pass run time information, we will use tools that leverage the LangChain Runnable interface. The standard runnables methods (invoke, batch, stream etc.) accept a 2nd argument which is a RunnableConfig. RunnableConfig has a few standard fields, but allows users to use other fields for run time information.\n",
    "\n",
    "Here, we will show how to set up a simple agent that has access to three tools for saving, reading, and deleting a list of the user's favorite pets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install --quiet -U langgraph langchain_anthropic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "ANTHROPIC_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(\"ANTHROPIC_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define tools and model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from langchain_core.tools import tool\n",
    "from langchain_core.runnables.config import RunnableConfig\n",
    "\n",
    "from langgraph.prebuilt import ToolNode\n",
    "\n",
    "user_to_pets = {}\n",
    "\n",
    "\n",
    "@tool(parse_docstring=True)\n",
    "def update_favorite_pets(\n",
    "    # NOTE: config arg does not need to be added to docstring, as we don't want it to be included in the function signature attached to the LLM\n",
    "    pets: List[str],\n",
    "    config: RunnableConfig,\n",
    ") -> None:\n",
    "    \"\"\"Add the list of favorite pets.\n",
    "\n",
    "    Args:\n",
    "        pets: List of favorite pets to set.\n",
    "    \"\"\"\n",
    "    user_id = config.get(\"configurable\", {}).get(\"user_id\")\n",
    "    user_to_pets[user_id] = pets\n",
    "\n",
    "\n",
    "@tool\n",
    "def delete_favorite_pets(config: RunnableConfig) -> None:\n",
    "    \"\"\"Delete the list of favorite pets.\"\"\"\n",
    "    user_id = config.get(\"configurable\", {}).get(\"user_id\")\n",
    "    if user_id in user_to_pets:\n",
    "        del user_to_pets[user_id]\n",
    "\n",
    "\n",
    "@tool\n",
    "def list_favorite_pets(config: RunnableConfig) -> None:\n",
    "    \"\"\"List favorite pets if any.\"\"\"\n",
    "    user_id = config.get(\"configurable\", {}).get(\"user_id\")\n",
    "    return \", \".join(user_to_pets.get(user_id, []))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [update_favorite_pets, delete_favorite_pets, list_favorite_pets]\n",
    "tool_node = ToolNode(tools)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll be using a small chat model from Anthropic in our example. To use chat models with tool calling, we need to first ensure that the model is aware of the available tools. We do this by calling `.bind_tools` method on `ChatAnthropic` moodel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_anthropic import ChatAnthropic\n",
    "from langgraph.graph import StateGraph, MessagesState\n",
    "from langgraph.prebuilt import ToolNode\n",
    "\n",
    "\n",
    "model_with_tools = ChatAnthropic(\n",
    "    model=\"claude-3-haiku-20240307\", temperature=0\n",
    ").bind_tools(tools)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ReAct Agent"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's set up a graph implementation of the [ReAct agent](https://langchain-ai.github.io/langgraph/concepts/agentic_concepts/#react-agent). This agent takes some query as input, then repeatedly call tools until it has enough information to resolve the query. We'll be using prebuilt `ToolNode` and the Anthropic model with tools we just defined"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "\n",
    "from langgraph.graph import StateGraph, MessagesState\n",
    "\n",
    "\n",
    "def should_continue(state: MessagesState) -> Literal[\"tools\", \"__end__\"]:\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "    if last_message.tool_calls:\n",
    "        return \"tools\"\n",
    "    return \"__end__\"\n",
    "\n",
    "\n",
    "def call_model(state: MessagesState):\n",
    "    messages = state[\"messages\"]\n",
    "    response = model_with_tools.invoke(messages)\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "workflow = StateGraph(MessagesState)\n",
    "\n",
    "# Define the two nodes we will cycle between\n",
    "workflow.add_node(\"agent\", call_model)\n",
    "workflow.add_node(\"tools\", tool_node)\n",
    "\n",
    "workflow.add_edge(\"__start__\", \"agent\")\n",
    "workflow.add_conditional_edges(\n",
    "    \"agent\",\n",
    "    should_continue,\n",
    ")\n",
    "workflow.add_edge(\"tools\", \"agent\")\n",
    "\n",
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 10320 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "try:\n",
    "    display(Image(app.get_graph().draw_mermaid_png()))\n",
    "except Exception:\n",
    "    # This requires some extra dependencies and is optional\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User information prior to run: {}\n",
      "Output from node 'agent':\n",
      "---\n",
      "{'messages': [AIMessage(content=[{'text': \"Okay, let's update your favorite pets:\", 'type': 'text'}, {'id': 'toolu_01LQK6fgtAyEo3xBfzg1fSuv', 'input': {'pets': ['cats', 'dogs']}, 'name': 'update_favorite_pets', 'type': 'tool_use'}], response_metadata={'id': 'msg_014bUFindzuzqqGmNVPX67zH', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'tool_use', 'stop_sequence': None, 'usage': {'input_tokens': 438, 'output_tokens': 70}}, id='run-2c77cfe0-ba1f-4cd5-922c-614330368ca3-0', tool_calls=[{'name': 'update_favorite_pets', 'args': {'pets': ['cats', 'dogs']}, 'id': 'toolu_01LQK6fgtAyEo3xBfzg1fSuv', 'type': 'tool_call'}], usage_metadata={'input_tokens': 438, 'output_tokens': 70, 'total_tokens': 508})]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node 'tools':\n",
      "---\n",
      "{'messages': [ToolMessage(content='null', name='update_favorite_pets', tool_call_id='toolu_01LQK6fgtAyEo3xBfzg1fSuv')]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node 'agent':\n",
      "---\n",
      "{'messages': [AIMessage(content='Your favorite pets have been updated to cats and dogs.', response_metadata={'id': 'msg_01JyfYdPiFHEPyE5PGeBXxqu', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 521, 'output_tokens': 15}}, id='run-c78b8fce-9358-4823-ac6c-896714860af2-0', usage_metadata={'input_tokens': 521, 'output_tokens': 15, 'total_tokens': 536})]}\n",
      "\n",
      "---\n",
      "\n",
      "User information after the run: {'123': ['cats', 'dogs']}\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "user_to_pets.clear()  # Clear the state\n",
    "\n",
    "print(f\"User information prior to run: {user_to_pets}\")\n",
    "\n",
    "inputs = {\"messages\": [HumanMessage(content=\"my favorite pets are cats and dogs\")]}\n",
    "for output in app.stream(inputs, {\"configurable\": {\"user_id\": \"123\"}}):\n",
    "    # stream() yields dictionaries with output keyed by node name\n",
    "    for key, value in output.items():\n",
    "        print(f\"Output from node '{key}':\")\n",
    "        print(\"---\")\n",
    "        print(value)\n",
    "    print(\"\\n---\\n\")\n",
    "\n",
    "print(f\"User information after the run: {user_to_pets}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User information prior to run: {'123': ['cats', 'dogs']}\n",
      "Output from node 'agent':\n",
      "---\n",
      "{'messages': [AIMessage(content=[{'id': 'toolu_01EsSgrDZ8aRZsg9y7ngroiu', 'input': {}, 'name': 'list_favorite_pets', 'type': 'tool_use'}], response_metadata={'id': 'msg_01Dp1VYH5RssYbReL6KzPfNM', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'tool_use', 'stop_sequence': None, 'usage': {'input_tokens': 437, 'output_tokens': 38}}, id='run-c620472c-ac52-488a-90f5-141fa65f1ce9-0', tool_calls=[{'name': 'list_favorite_pets', 'args': {}, 'id': 'toolu_01EsSgrDZ8aRZsg9y7ngroiu', 'type': 'tool_call'}], usage_metadata={'input_tokens': 437, 'output_tokens': 38, 'total_tokens': 475})]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node 'tools':\n",
      "---\n",
      "{'messages': [ToolMessage(content='cats, dogs', name='list_favorite_pets', tool_call_id='toolu_01EsSgrDZ8aRZsg9y7ngroiu')]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node 'agent':\n",
      "---\n",
      "{'messages': [AIMessage(content='Based on the output, your favorite pets are cats and dogs.', response_metadata={'id': 'msg_017heQczfgTMCzAo5qcYdYWW', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 490, 'output_tokens': 17}}, id='run-c0cb9626-61a1-4151-b194-be0e7d655a8d-0', usage_metadata={'input_tokens': 490, 'output_tokens': 17, 'total_tokens': 507})]}\n",
      "\n",
      "---\n",
      "\n",
      "User information after the run: {'123': ['cats', 'dogs']}\n"
     ]
    }
   ],
   "source": [
    "print(f\"User information prior to run: {user_to_pets}\")\n",
    "\n",
    "\n",
    "inputs = {\"messages\": [HumanMessage(content=\"what are my favorite pets?\")]}\n",
    "for output in app.stream(inputs, {\"configurable\": {\"user_id\": \"123\"}}):\n",
    "    # stream() yields dictionaries with output keyed by node name\n",
    "    for key, value in output.items():\n",
    "        print(f\"Output from node '{key}':\")\n",
    "        print(\"---\")\n",
    "        print(value)\n",
    "    print(\"\\n---\\n\")\n",
    "\n",
    "\n",
    "print(f\"User information after the run: {user_to_pets}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User information prior to run: {'123': ['cats', 'dogs']}\n",
      "Output from node 'agent':\n",
      "---\n",
      "{'messages': [AIMessage(content=[{'id': 'toolu_01EcVWNpWQnoRuRtXXbndeWn', 'input': {}, 'name': 'delete_favorite_pets', 'type': 'tool_use'}], response_metadata={'id': 'msg_01PfMPkCHuV1UvcCKdqT5jXH', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'tool_use', 'stop_sequence': None, 'usage': {'input_tokens': 441, 'output_tokens': 38}}, id='run-eeac69b6-812e-4630-ba6f-9b22b672493b-0', tool_calls=[{'name': 'delete_favorite_pets', 'args': {}, 'id': 'toolu_01EcVWNpWQnoRuRtXXbndeWn', 'type': 'tool_call'}], usage_metadata={'input_tokens': 441, 'output_tokens': 38, 'total_tokens': 479})]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node 'tools':\n",
      "---\n",
      "{'messages': [ToolMessage(content='null', name='delete_favorite_pets', tool_call_id='toolu_01EcVWNpWQnoRuRtXXbndeWn')]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node 'agent':\n",
      "---\n",
      "{'messages': [AIMessage(content='I have deleted the information about your favorite pets. The list of favorite pets has been cleared.', response_metadata={'id': 'msg_01PvNPmzfgSvGdWQp6ATWs6Q', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 492, 'output_tokens': 23}}, id='run-0cb06dbb-7d6c-4aa1-9ba7-685b0de62e06-0', usage_metadata={'input_tokens': 492, 'output_tokens': 23, 'total_tokens': 515})]}\n",
      "\n",
      "---\n",
      "\n",
      "User information prior to run: {}\n"
     ]
    }
   ],
   "source": [
    "print(f\"User information prior to run: {user_to_pets}\")\n",
    "\n",
    "\n",
    "inputs = {\n",
    "    \"messages\": [\n",
    "        HumanMessage(content=\"please forget what i told you about my favorite animals\")\n",
    "    ]\n",
    "}\n",
    "for output in app.stream(inputs, {\"configurable\": {\"user_id\": \"123\"}}):\n",
    "    # stream() yields dictionaries with output keyed by node name\n",
    "    for key, value in output.items():\n",
    "        print(f\"Output from node '{key}':\")\n",
    "        print(\"---\")\n",
    "        print(value)\n",
    "    print(\"\\n---\\n\")\n",
    "\n",
    "\n",
    "print(f\"User information prior to run: {user_to_pets}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langgraph",
   "language": "python",
   "name": "langgraph"
  },
  "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": 4
}
