{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
   "metadata": {},
   "source": [
    "# How to view and update graph state\n",
    "\n",
    "Once you start [checkpointing](./persistence.ipynb) your graphs, you can easily **get** or **update** the state of the agent at any point in time. This permits a few things:\n",
    "\n",
    "1. You can surface a state during an interrupt to a user to let them accept an action.\n",
    "2. You can **rewind** the graph to reproduce or avoid issues.\n",
    "3. You can **modify** the state to embed your agent into a larger system, or to let the user better control its actions.\n",
    "\n",
    "The key methods used for this functionality are:\n",
    "\n",
    "- [get_state](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.graph.CompiledGraph.get_state): fetch the values from the target config\n",
    "- [update_state](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.graph.CompiledGraph.update_state): apply the given values to the target state\n",
    "\n",
    "**Note:** this requires passing in a checkpointer.\n",
    "\n",
    "Below is a quick example.\n",
    "\n",
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Note:</p>\n",
    "    <p>\n",
    "        In this how-to, we will create our agent from scratch to be transparent (but verbose). You can accomplish similar functionality using the <code>create_react_agent(model, tools=tool, checkpointer=checkpointer)</code> (<a href=\"https://langchain-ai.github.io/langgraph/reference/prebuilt/#create_react_agent\">API doc</a>) constructor. This may be more appropriate if you are used to LangChain’s <a href=\"https://python.langchain.com/v0.2/docs/how_to/agent_executor/#concepts\">AgentExecutor</a> class.\n",
    "    </p>\n",
    "</div>    "
   ]
  },
  {
   "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 LLM we will use) and Tavily (the search tool we will use)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c903a1cf-2977-4e2d-ad7d-8b3946821d89",
   "metadata": {},
   "outputs": [],
   "source": ["import getpass\nimport os\n\n\ndef _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",
   "id": "f0ed46a8-effe-4596-b0e1-a6a29ee16f5c",
   "metadata": {},
   "source": [
    "Optionally, we can set API key for [LangSmith tracing](https://smith.langchain.com/), which will give us best-in-class observability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "95e25aec-7c9f-4a63-b143-225d0e9a79c3",
   "metadata": {},
   "outputs": [],
   "source": ["os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n_set_env(\"LANGCHAIN_API_KEY\")"]
  },
  {
   "cell_type": "markdown",
   "id": "e36f89e5",
   "metadata": {},
   "source": [
    "## Set up the State\n",
    "\n",
    "The state is the interface for all the nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f5319e01",
   "metadata": {},
   "outputs": [],
   "source": ["from typing import Annotated\n\nfrom typing_extensions import TypedDict\n\nfrom langgraph.graph.message import add_messages\n\n# `add_messages`` essentially does this\n# (with more robust handling)\n# def add_messages(left: list, right: list):\n#     return left + right\n\n\nclass State(TypedDict):\n    messages: Annotated[list, add_messages]"]
  },
  {
   "cell_type": "markdown",
   "id": "21ac643b-cb06-4724-a80c-2862ba4773f1",
   "metadata": {},
   "source": [
    "## Set up the tools\n",
    "\n",
    "We will first define the tools we want to use.\n",
    "For this simple example, we will use create a placeholder search engine.\n",
    "However, it is really easy to create your own tools - see documentation [here](https://python.langchain.com/v0.2/docs/how_to/custom_tools) on how to do that.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d7ef57dd-5d6e-4ad3-9377-a92201c1310e",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_core.tools import tool\n\n\n@tool\ndef search(query: str):\n    \"\"\"Call to surf the web.\"\"\"\n    # This is a placeholder for the actual implementation\n    return [\"The weather is cloudy with a chance of meatballs.\"]\n\n\ntools = [search]"]
  },
  {
   "cell_type": "markdown",
   "id": "01885785-b71a-44d1-b1d6-7b5b14d53b58",
   "metadata": {},
   "source": [
    "We can now wrap these tools in a simple ToolNode.\n",
    "This is a prebuilt node that extracts tool calls from the most recent AIMessage, executes them, and returns a ToolMessage with the results.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5cf3331e-ccb3-41c8-aeb9-a840a94d41e7",
   "metadata": {},
   "outputs": [],
   "source": ["from langgraph.prebuilt import ToolNode\n\ntool_node = ToolNode(tools)"]
  },
  {
   "cell_type": "markdown",
   "id": "5497ed70-fce3-47f1-9cad-46f912bad6a5",
   "metadata": {},
   "source": [
    "## Set up the model\n",
    "\n",
    "Now we need to load the chat model we want to use.\n",
    "Importantly, this should satisfy two criteria:\n",
    "\n",
    "1. It should work with messages. We will represent all agent state in the form of messages, so it needs to be able to work well with them.\n",
    "2. It should work with OpenAI function calling. This means it should either be an OpenAI model or a model that exposes a similar interface.\n",
    "\n",
    "Note: these model requirements are not requirements for using LangGraph - they are just requirements for this one example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "892b54b9-75f0-4804-9ed0-88b5e5532989",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_openai import ChatOpenAI\n\nmodel = ChatOpenAI(temperature=0)"]
  },
  {
   "cell_type": "markdown",
   "id": "a77995c0-bae2-4cee-a036-8688a90f05b9",
   "metadata": {},
   "source": [
    "\n",
    "After we've done this, we should make sure the model knows that it has these tools available to call.\n",
    "We can do this using the `.bind_tools()` method, common to many of LangChain's chat models.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cd3cbae5-d92c-4559-a4aa-44721b80d107",
   "metadata": {},
   "outputs": [],
   "source": ["model = model.bind_tools(tools)"]
  },
  {
   "cell_type": "markdown",
   "id": "e03c5094-9297-4d19-a04e-3eedc75cefb4",
   "metadata": {},
   "source": [
    "## Define the nodes\n",
    "\n",
    "We now need to define a few different nodes in our graph.\n",
    "In `langgraph`, a node can be either a function or a [runnable](https://python.langchain.com/v0.2/docs/concepts/#langchain-expression-language-lcel).\n",
    "There are two main nodes we need for this:\n",
    "\n",
    "1. The agent: responsible for deciding what (if any) actions to take.\n",
    "2. A function to invoke tools: if the agent decides to take an action, this node will then execute that action.\n",
    "\n",
    "We will also need to define some edges.\n",
    "Some of these edges may be conditional.\n",
    "The reason they are conditional is that based on the output of a node, one of several paths may be taken.\n",
    "The path that is taken is not known until that node is run (the LLM decides).\n",
    "\n",
    "1. Conditional Edge: after the agent is called, we should either:\n",
    "   a. If the agent said to take an action, then the function to invoke tools should be called\n",
    "   b. If the agent said that it was finished, then it should finish\n",
    "2. Normal Edge: after the tools are invoked, it should always go back to the agent to decide what to do next\n",
    "\n",
    "Let's define the nodes, as well as a function to decide how what conditional edge to take."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3b541bb9-900c-40d0-964d-7b5dfee30667",
   "metadata": {},
   "outputs": [],
   "source": ["from typing import Literal\n\n\n# Define the function that determines whether to continue or not\ndef should_continue(state: State) -> Literal[\"continue\", \"end\"]:\n    last_message = state[\"messages\"][-1]\n    # If there is no function call, then we finish\n    if not last_message.tool_calls:\n        return \"end\"\n    # Otherwise if there is, we continue\n    else:\n        return \"continue\""]
  },
  {
   "cell_type": "markdown",
   "id": "ffd6e892-946c-4899-8cc0-7c9291c1f73b",
   "metadata": {},
   "source": [
    "## Define the graph\n",
    "\n",
    "We can now put it all together and define the graph!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "812b4e70-4956-4415-8880-db48b3dcbad2",
   "metadata": {},
   "outputs": [],
   "source": ["from langgraph.graph import END, StateGraph, START\n\n# Define a new graph\nworkflow = StateGraph(State)\n\n\n# Define the two nodes we will cycle between\ndef call_model(state: State) -> State:\n    return {\"messages\": model.invoke(state[\"messages\"])}\n\n\nworkflow.add_node(\"agent\", call_model)\nworkflow.add_node(\"action\", tool_node)\n\n# Set the entrypoint as `agent`\n# This means that this node is the first one called\nworkflow.add_edge(START, \"agent\")\n\n# We now add a conditional edge\nworkflow.add_conditional_edges(\n    # First, we define the start node. We use `agent`.\n    # This means these are the edges taken after the `agent` node is called.\n    \"agent\",\n    # Next, we pass in the function that will determine which node is called next.\n    should_continue,\n    # Finally we pass in a mapping.\n    # The keys are strings, and the values are other nodes.\n    # END is a special node marking that the graph should finish.\n    # What will happen is we will call `should_continue`, and then the output of that\n    # will be matched against the keys in this mapping.\n    # Based on which one it matches, that node will then be called.\n    {\n        # If `tools`, then we call the tool node.\n        \"continue\": \"action\",\n        # Otherwise we finish.\n        \"end\": END,\n    },\n)\n\n# We now add a normal edge from `tools` to `agent`.\n# This means that after `tools` is called, `agent` node is called next.\nworkflow.add_edge(\"action\", \"agent\")"]
  },
  {
   "cell_type": "markdown",
   "id": "bc9c8536-f90b-44fa-958d-5df016c66d8f",
   "metadata": {},
   "source": [
    "**Persistence**\n",
    "\n",
    "To add in persistence, we pass in a checkpoint when compiling the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6845ed6a-d155-4105-9160-28849877248b",
   "metadata": {},
   "outputs": [],
   "source": ["from langgraph.checkpoint.memory import MemorySaver\n\nmemory = MemorySaver()"]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "79d29875-8aa8-434c-9f20-1c58346a6249",
   "metadata": {},
   "outputs": [],
   "source": ["# Finally, we compile it!\n# This compiles it into a LangChain Runnable,\n# meaning you can use it as you would any other runnable\napp = workflow.compile(checkpointer=memory)"]
  },
  {
   "cell_type": "markdown",
   "id": "e8aff75b-563e-42b1-969b-742201514fc3",
   "metadata": {},
   "source": [
    "## Preview the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c9ab60eb-679b-4eef-9e64-5ffbf3dffc70",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 12716 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": ["from IPython.display import Image, display\n\ntry:\n    display(Image(app.get_graph().draw_mermaid_png()))\nexcept Exception:\n    # This requires some extra dependencies and is optional\n    pass"]
  },
  {
   "cell_type": "markdown",
   "id": "2a1b56c5-bd61-4192-8bdb-458a1e9f0159",
   "metadata": {},
   "source": [
    "## Interacting with the Agent\n",
    "\n",
    "We can now interact with the agent. Between interactions you can get and update state.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "cfd140f0-a5a6-4697-8115-322242f197b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "hi! I'm bob\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Hello Bob! How can I assist you today?\n"
     ]
    }
   ],
   "source": ["from langchain_core.messages import HumanMessage\n\nconfig = {\"configurable\": {\"thread_id\": \"2\"}}\ninput_message = HumanMessage(content=\"hi! I'm bob\")\nfor event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n    event[\"messages\"][-1].pretty_print()"]
  },
  {
   "cell_type": "markdown",
   "id": "9230c0a7-bab8-4327-a9b4-f21e7a28d32a",
   "metadata": {},
   "source": [
    "See LangSmith example run here https://smith.langchain.com/public/01c1d61c-6943-4db1-8afe-5366f083caf3/r\n",
    "\n",
    "Here you can see the \"agent\" node ran, and then \"should_continue\" returned \"end\" so the graph stopped execution there."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c161d511-0dc3-4824-b89e-0dbcf9a7a759",
   "metadata": {},
   "source": [
    "Let's now get the current state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "23433643-d35d-4df4-80fe-a3002323cd4f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content=\"hi! I'm bob\", id='cd7df241-189c-46a6-b822-69fcfafd8ad4'),\n",
       "  AIMessage(content='Hello Bob! How can I assist you today?', response_metadata={'finish_reason': 'stop', 'logprobs': None, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'token_usage': {'completion_tokens': 11, 'prompt_tokens': 54, 'total_tokens': 65}}, id='run-cc3e7ee7-208e-446e-80cb-0349fe75319b-0')]}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["app.get_state(config).values"]
  },
  {
   "cell_type": "markdown",
   "id": "b894d739-d8b1-485b-841a-5cd997c111f9",
   "metadata": {},
   "source": [
    "The current state is the two messages we've seen above, 1. the HumanMessage we sent in, 2. the AIMessage we got back from the model.\n",
    "\n",
    "The `next` values are empty since the graph has terminated (transitioned to the `__end__`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e2c2531f-6dda-444b-b2f5-bbd607251776",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "()"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["app.get_state(config).next"]
  },
  {
   "cell_type": "markdown",
   "id": "98a953a5-7c04-4210-a6ea-bd42e8abd81c",
   "metadata": {},
   "source": [
    "The graph got to the end without interruptions, so the list of next nodes is empty."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f117f21-cabc-4c6a-b89a-e673085bc067",
   "metadata": {},
   "source": [
    "### Let's get it to execute a tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "73eb35a7-b705-4d7b-9e4b-28f7e2130358",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what is the weather in sf currently\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  search (call_UVPlm7YZ0xksC2VsYsPxN5ag)\n",
      " Call ID: call_UVPlm7YZ0xksC2VsYsPxN5ag\n",
      "  Args:\n",
      "    query: weather in San Francisco\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: search\n",
      "\n",
      "[\"The weather is cloudy with a chance of meatballs.\"]\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The weather in San Francisco is currently cloudy with a chance of meatballs.\n"
     ]
    }
   ],
   "source": ["config = {\"configurable\": {\"thread_id\": \"2\"}}\ninput_message = HumanMessage(content=\"what is the weather in sf currently\")\nfor event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n    event[\"messages\"][-1].pretty_print()"]
  },
  {
   "cell_type": "markdown",
   "id": "c0ebc4b8-b2cc-4394-bb5a-8964bafa8d2d",
   "metadata": {},
   "source": [
    "See LangSmith example run here https://smith.langchain.com/public/c33c04c5-f1f2-4977-9d7d-c48f28be7be2/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18d3141b-9f9a-4ad8-b475-400674848b16",
   "metadata": {},
   "source": [
    "We can see it planned the tool execution (ie the \"agent\" node), then \"should_continue\" edge returned \"continue\" so we proceeded to \"action\" node, which executed the tool, and then \"agent\" node emitted the final response, which made \"should_continue\" edge return \"end\". Let's see how we can have more control over this."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c5d69cf-6de5-4a52-976a-3799aa39d6d5",
   "metadata": {},
   "source": [
    "### Pause before tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06aaa787-32fe-404b-abeb-5b2ed94309fd",
   "metadata": {},
   "source": [
    "If you notice below, we now will add `interrupt_before=[\"action\"]` - this means that before any actions are taken we pause. This is a great moment to allow the user to correct and update the state! This is very useful when you want to have a human-in-the-loop to validate (and potentially change) the action to take. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5a68afc0-606f-4294-a872-b2b563be0d69",
   "metadata": {},
   "outputs": [],
   "source": ["app_w_interrupt = workflow.compile(checkpointer=memory, interrupt_before=[\"action\"])"]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "08ae8246-11d5-40e1-8567-361e5bef8917",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what is the weather in sf currently\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  search (call_sxtKypVZlFrjzdOFYiCh8kin)\n",
      " Call ID: call_sxtKypVZlFrjzdOFYiCh8kin\n",
      "  Args:\n",
      "    query: weather in San Francisco\n"
     ]
    }
   ],
   "source": ["config = {\"configurable\": {\"thread_id\": \"4\"}}\ninput_message = HumanMessage(content=\"what is the weather in sf currently\")\nfor event in app_w_interrupt.stream(\n    {\"messages\": [input_message]}, config, stream_mode=\"values\"\n):\n    event[\"messages\"][-1].pretty_print()"]
  },
  {
   "cell_type": "markdown",
   "id": "c53352c8-dcf7-48b5-940c-f75a279cb42c",
   "metadata": {},
   "source": [
    "See LangSmith example run here https://smith.langchain.com/public/22402055-a50e-4d82-8b3e-733c9d752bc5/r\n",
    "This time it executed the \"agent\" node same as before, and you can see in the LangSmith trace that \"should_continue\" returned \"continue\", but it paused execution per our setting above."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d9fa4f0-fb7d-47a7-b77e-464e8b619273",
   "metadata": {},
   "source": [
    "Notice that this time, the `next` value is populated with `action`. That means that if we resume the graph, it will start at the `action` node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5a53df78-6c25-4176-b049-02a02a713771",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('action',)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["current_values = app_w_interrupt.get_state(config)\ncurrent_values.next"]
  },
  {
   "cell_type": "markdown",
   "id": "a901e2d1-dcf4-470c-88c0-2e4c91307f69",
   "metadata": {},
   "source": [
    "Because we asked to interrupt the graph before getting to the action node, the next node to execute, if we were to resume, would be the \"action\" node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9bad3f92-a3b9-4370-a1f5-0b78d1b67602",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'name': 'search',\n",
       "  'args': {'query': 'weather in San Francisco'},\n",
       "  'id': 'call_sxtKypVZlFrjzdOFYiCh8kin'}]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["current_values.values[\"messages\"][-1].tool_calls"]
  },
  {
   "cell_type": "markdown",
   "id": "b2b7af00-035c-44e6-a76a-a4b83340b32b",
   "metadata": {},
   "source": [
    "Let's update the search string before proceeding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "060e2e33-1f6a-40ef-850e-161b308986fb",
   "metadata": {},
   "outputs": [],
   "source": ["current_values.values[\"messages\"][-1].tool_calls[0][\"args\"][\n    \"query\"\n] = \"weather in San Francisco today\""]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "3210d003-e20e-47fa-b4f4-402fed7409f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'configurable': {'thread_id': '4',\n",
       "  'thread_ts': '2024-05-07T17:30:25.205012+00:00'}}"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["app_w_interrupt.update_state(config, current_values.values)"]
  },
  {
   "cell_type": "markdown",
   "id": "0a21b87a-319c-4a93-9ac1-a3c6ad3f5400",
   "metadata": {},
   "source": [
    "This actually produces a LangSmith run too! See it here https://smith.langchain.com/public/9d86718b-333e-4175-bec0-9a64cdd01dc3/r\n",
    "\n",
    "This is a shorter run that allows you to inspect the edges that reacted to the state update, you can see \"should_continue\" returned \"continue\" as before, given this is still a function call."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3613f445-cf46-46a0-b7c1-e83a02aae96f",
   "metadata": {},
   "source": [
    "The current state now reflects our updated search query!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4a7dfaab-bfa7-47c7-8b93-2b06b3560809",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='what is the weather in sf currently', id='7e198f29-a371-49d5-86df-7e0e0b5a9144'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'function': {'arguments': '{\"query\":\"weather in San Francisco\"}', 'name': 'search'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin', 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'logprobs': None, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'token_usage': {'completion_tokens': 16, 'prompt_tokens': 56, 'total_tokens': 72}}, id='run-0e6d8103-a92e-461d-aa99-8a68f4c99366-0', tool_calls=[{'name': 'search', 'args': {'query': 'weather in San Francisco today'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin'}])]}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["app_w_interrupt.get_state(config).values"]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "77d73f33-93cb-41db-b6bb-81697569fdb0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('action',)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["app_w_interrupt.get_state(config).next"]
  },
  {
   "cell_type": "markdown",
   "id": "27304085-4e44-48bc-afa5-d7ab6da4bb8c",
   "metadata": {},
   "source": [
    "If we start the agent again it will pick up from the state we updated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "273d56a8-f40f-4a51-a27f-7c6bb2bda0ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'messages': [ToolMessage(content='[\"The weather is cloudy with a chance of meatballs.\"]', name='search', id='9dce802a-9811-491f-a1d5-ace400fbcba0', tool_call_id='call_sxtKypVZlFrjzdOFYiCh8kin')]}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'messages': AIMessage(content='The weather in San Francisco is currently cloudy with a chance of meatballs.', response_metadata={'token_usage': {'completion_tokens': 16, 'prompt_tokens': 92, 'total_tokens': 108}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-269cd84d-c5ba-438b-9abf-1d1389bed733-0')}\n"
     ]
    }
   ],
   "source": ["for event in app_w_interrupt.stream(None, config):\n    for v in event.values():\n        print(v)"]
  },
  {
   "cell_type": "markdown",
   "id": "14d288a5-1027-4a12-898c-1c6bfcc19bb6",
   "metadata": {},
   "source": [
    "See this run in LangSmith here https://smith.langchain.com/public/8262c0f9-0701-4d73-95f6-2a32f6d3f96a/r\n",
    "\n",
    "This continues where we left off, with \"action\" node, followed by \"agent\" node, which terminates the execution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c38c505-6cee-427f-9dcd-493a2ade7ebb",
   "metadata": {},
   "source": [
    "## Checking history\n",
    "\n",
    "Let's browse the history of this thread, from newest to oldest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "8578a66d-6489-4e03-8c23-fd0530278455",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "StateSnapshot(values={'messages': [HumanMessage(content='what is the weather in sf currently', id='7e198f29-a371-49d5-86df-7e0e0b5a9144'), AIMessage(content='', additional_kwargs={'tool_calls': [{'function': {'arguments': '{\"query\":\"weather in San Francisco\"}', 'name': 'search'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin', 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'logprobs': None, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'token_usage': {'completion_tokens': 16, 'prompt_tokens': 56, 'total_tokens': 72}}, id='run-0e6d8103-a92e-461d-aa99-8a68f4c99366-0', tool_calls=[{'name': 'search', 'args': {'query': 'weather in San Francisco today'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin'}]), ToolMessage(content='[\"The weather is cloudy with a chance of meatballs.\"]', name='search', id='9dce802a-9811-491f-a1d5-ace400fbcba0', tool_call_id='call_sxtKypVZlFrjzdOFYiCh8kin'), AIMessage(content='The weather in San Francisco is currently cloudy with a chance of meatballs.', response_metadata={'finish_reason': 'stop', 'logprobs': None, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'token_usage': {'completion_tokens': 16, 'prompt_tokens': 92, 'total_tokens': 108}}, id='run-269cd84d-c5ba-438b-9abf-1d1389bed733-0')]}, next=(), config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:25.872512+00:00'}}, metadata={'source': 'loop', 'step': 4}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:25.228389+00:00'}})\n",
      "--\n",
      "StateSnapshot(values={'messages': [HumanMessage(content='what is the weather in sf currently', id='7e198f29-a371-49d5-86df-7e0e0b5a9144'), AIMessage(content='', additional_kwargs={'tool_calls': [{'function': {'arguments': '{\"query\":\"weather in San Francisco\"}', 'name': 'search'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin', 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'logprobs': None, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'token_usage': {'completion_tokens': 16, 'prompt_tokens': 56, 'total_tokens': 72}}, id='run-0e6d8103-a92e-461d-aa99-8a68f4c99366-0', tool_calls=[{'name': 'search', 'args': {'query': 'weather in San Francisco today'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin'}]), ToolMessage(content='[\"The weather is cloudy with a chance of meatballs.\"]', name='search', id='9dce802a-9811-491f-a1d5-ace400fbcba0', tool_call_id='call_sxtKypVZlFrjzdOFYiCh8kin')]}, next=('agent',), config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:25.228389+00:00'}}, metadata={'source': 'loop', 'step': 3}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:25.205012+00:00'}})\n",
      "--\n",
      "StateSnapshot(values={'messages': [HumanMessage(content='what is the weather in sf currently', id='7e198f29-a371-49d5-86df-7e0e0b5a9144'), AIMessage(content='', additional_kwargs={'tool_calls': [{'function': {'arguments': '{\"query\":\"weather in San Francisco\"}', 'name': 'search'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin', 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'logprobs': None, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'token_usage': {'completion_tokens': 16, 'prompt_tokens': 56, 'total_tokens': 72}}, id='run-0e6d8103-a92e-461d-aa99-8a68f4c99366-0', tool_calls=[{'name': 'search', 'args': {'query': 'weather in San Francisco today'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin'}])]}, next=('action',), config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:25.205012+00:00'}}, metadata={'source': 'update', 'step': 2}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:25.186985+00:00'}})\n",
      "--\n",
      "StateSnapshot(values={'messages': [HumanMessage(content='what is the weather in sf currently', id='7e198f29-a371-49d5-86df-7e0e0b5a9144'), AIMessage(content='', additional_kwargs={'tool_calls': [{'function': {'arguments': '{\"query\":\"weather in San Francisco\"}', 'name': 'search'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin', 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'logprobs': None, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'token_usage': {'completion_tokens': 16, 'prompt_tokens': 56, 'total_tokens': 72}}, id='run-0e6d8103-a92e-461d-aa99-8a68f4c99366-0', tool_calls=[{'name': 'search', 'args': {'query': 'weather in San Francisco'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin'}])]}, next=('action',), config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:25.186985+00:00'}}, metadata={'source': 'loop', 'step': 1}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:24.675950+00:00'}})\n",
      "--\n",
      "StateSnapshot(values={'messages': [HumanMessage(content='what is the weather in sf currently', id='7e198f29-a371-49d5-86df-7e0e0b5a9144')]}, next=('agent',), config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:24.675950+00:00'}}, metadata={'source': 'loop', 'step': 0}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:24.672976+00:00'}})\n",
      "--\n",
      "StateSnapshot(values={'messages': []}, next=('__start__',), config={'configurable': {'thread_id': '4', 'thread_ts': '2024-05-07T17:30:24.672976+00:00'}}, metadata={'source': 'input', 'step': -1}, parent_config=None)\n",
      "--\n"
     ]
    }
   ],
   "source": ["for state in app_w_interrupt.get_state_history(config):\n    print(state)\n    print(\"--\")\n    if len(state.values[\"messages\"]) == 2:\n        to_replay = state"]
  },
  {
   "cell_type": "markdown",
   "id": "0ec41c37-7c09-4cc7-8475-bf373fe66584",
   "metadata": {},
   "source": [
    "We can go back to any of these states and restart the agent from there!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "21e7fc18-6fd9-4e11-a84b-e0325c9640c8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='what is the weather in sf currently', id='7e198f29-a371-49d5-86df-7e0e0b5a9144'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'function': {'arguments': '{\"query\":\"weather in San Francisco\"}', 'name': 'search'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin', 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'logprobs': None, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'token_usage': {'completion_tokens': 16, 'prompt_tokens': 56, 'total_tokens': 72}}, id='run-0e6d8103-a92e-461d-aa99-8a68f4c99366-0', tool_calls=[{'name': 'search', 'args': {'query': 'weather in San Francisco'}, 'id': 'call_sxtKypVZlFrjzdOFYiCh8kin'}])]}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["to_replay.values"]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "d4b01634-0041-4632-8d1f-5464580e54f5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('action',)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["to_replay.next"]
  },
  {
   "cell_type": "markdown",
   "id": "29da43ea-9295-43e2-b164-0eb28d96749c",
   "metadata": {},
   "source": [
    "### Replay a past state\n",
    "\n",
    "To replay from this place we just need to pass its config back to the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "e986f94f-706f-4b6f-b3c4-f95483b9e9b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'messages': [ToolMessage(content='[\"The weather is cloudy with a chance of meatballs.\"]', name='search', id='71e6f2b9-46cf-4629-a0e2-fda37da9a3bb', tool_call_id='call_sxtKypVZlFrjzdOFYiCh8kin')]}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'messages': AIMessage(content='The weather in San Francisco is currently cloudy with a chance of meatballs.', response_metadata={'token_usage': {'completion_tokens': 16, 'prompt_tokens': 91, 'total_tokens': 107}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-d2ed2496-271f-4353-8f9c-3fb3157a4f63-0')}\n"
     ]
    }
   ],
   "source": ["for event in app_w_interrupt.stream(None, to_replay.config):\n    for v in event.values():\n        print(v)"]
  },
  {
   "cell_type": "markdown",
   "id": "8b4b81dd-1108-4ab7-b680-3f641b3efa70",
   "metadata": {},
   "source": [
    "See this run in LangSmith here https://smith.langchain.com/public/f26e9e1d-16df-48ae-98f7-c823d6942bf7/r\n",
    "\n",
    "This is similar to the previous run, this time with the original search query, instead of our modified one. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59910951-fae1-4475-8511-f622439b590d",
   "metadata": {},
   "source": [
    "### Branch off a past state\n",
    "\n",
    "Using LangGraph's checkpointing, you can do more than just replay past states. You can branch off previous locations to let the agent explore alternate trajectories or to let a user \"version control\" changes in a workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "b084f141-5800-487b-b115-d2e58421b963",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_core.messages import AIMessage\n\nbranch_config = app_w_interrupt.update_state(\n    to_replay.config,\n    {\n        \"messages\": [\n            AIMessage(content=\"All done here!\", id=to_replay.values[\"messages\"][-1].id)\n        ]\n    },\n)"]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1a7cfcd4-289e-419e-8b49-dfaef4f88641",
   "metadata": {},
   "outputs": [],
   "source": ["branch_state = app_w_interrupt.get_state(branch_config)"]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "5198f9c1-d2d4-458a-993d-3caa55810b1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='what is the weather in sf currently', id='7e198f29-a371-49d5-86df-7e0e0b5a9144'),\n",
       "  AIMessage(content='All done here!', id='run-0e6d8103-a92e-461d-aa99-8a68f4c99366-0')]}"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["branch_state.values"]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "5d89d55d-db84-4c2d-828b-64a29a69947b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "()"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["branch_state.next"]
  },
  {
   "cell_type": "markdown",
   "id": "cc168c90-a374-4280-a9a6-8bc232dbb006",
   "metadata": {},
   "source": [
    "You can see the snapshot was updated and now correctly reflects that there is no next step.\n",
    "\n",
    "You can see this in LangSmith update run here https://smith.langchain.com/public/65104717-6eda-4a0f-93c1-4755c6f929ed/r\n",
    "\n",
    "This shows the \"should_continue\" edge now reacting to this replaced message, and now changing the outcome to \"end\" which finishes the computation."
   ]
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
