{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
   "metadata": {},
   "source": [
    "# How to delete messages\n",
    "\n",
    "One of the common states for a graph is a list of messages. Usually you only add messages to that state. However, sometimes you may want to remove messages (either by directly modifying the state or as part of the graph). To do that, you can use the `RemoveMessage` modifier. In this guide, we will cover how to do that.\n",
    "\n",
    "The key idea is that each state key has a `reducer` key. This key specifies how to combine updates to the state. The default `MessagesState` has a messages key, and the reducer for that key accepts these `RemoveMessage` modifiers. That reducer then uses these `RemoveMessage` to delete messages from the key.\n",
    "\n",
    "So note that just because your graph state has a key that is a list of messages, it doesn't mean that that this `RemoveMessage` modifier will work. You also have to have a `reducer` defined that knows how to work with this.\n",
    "\n",
    "**NOTE**: Many models expect certain rules around lists of messages. For example, some expect them to start with a `user` message, others expect all messages with tool calls to be followed by a tool message. **When deleting messages, you will want to make sure you don't violate these rules.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cbd446a-808f-4394-be92-d45ab818953c",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, let's build a simple graph that uses messages. Note that it's using the `MessagesState` which has the required `reducer`."
   ]
  },
  {
   "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_anthropic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0abe11f4-62ed-4dc4-8875-3db21e260d1d",
   "metadata": {},
   "source": [
    "Next, we need to set API keys for Anthropic (the LLM we will use)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c903a1cf-2977-4e2d-ad7d-8b3946821d89",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "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",
   "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>    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4767ef1c-a7cf-41f8-a301-558988cb7ac5",
   "metadata": {},
   "source": [
    "## Build the agent\n",
    "Let's now build a simple ReAct style agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "378899a9-3b9a-4748-95b6-eb00e0828677",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "\n",
    "from langchain_anthropic import ChatAnthropic\n",
    "from langchain_core.tools import tool\n",
    "\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import MessagesState, StateGraph, START, END\n",
    "from langgraph.prebuilt import ToolNode\n",
    "\n",
    "memory = MemorySaver()\n",
    "\n",
    "\n",
    "@tool\n",
    "def search(query: str):\n",
    "    \"\"\"Call to surf the web.\"\"\"\n",
    "    # This is a placeholder for the actual implementation\n",
    "    # Don't let the LLM know this though 😊\n",
    "    return \"It's sunny in San Francisco, but you better look out if you're a Gemini 😈.\"\n",
    "\n",
    "\n",
    "tools = [search]\n",
    "tool_node = ToolNode(tools)\n",
    "model = ChatAnthropic(model_name=\"claude-3-haiku-20240307\")\n",
    "bound_model = model.bind_tools(tools)\n",
    "\n",
    "\n",
    "def should_continue(state: MessagesState):\n",
    "    \"\"\"Return the next node to execute.\"\"\"\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",
    "    return \"action\"\n",
    "\n",
    "\n",
    "# Define the function that calls the model\n",
    "def call_model(state: MessagesState):\n",
    "    response = model.invoke(state[\"messages\"])\n",
    "    # We return a list, because this will get added to the existing list\n",
    "    return {\"messages\": response}\n",
    "\n",
    "\n",
    "# Define a new graph\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(\"action\", tool_node)\n",
    "\n",
    "# Set the entrypoint as `agent`\n",
    "# This means that this node is the first one called\n",
    "workflow.add_edge(START, \"agent\")\n",
    "\n",
    "# We now add a conditional edge\n",
    "workflow.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",
    "    # Next, we pass in the path map - all the possible nodes this edge could go to\n",
    "    [\"action\", END],\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.\n",
    "workflow.add_edge(\"action\", \"agent\")\n",
    "\n",
    "# Finally, we compile it!\n",
    "# This compiles it into a LangChain Runnable,\n",
    "# meaning you can use it as you would any other runnable\n",
    "app = workflow.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "57b27553-21be-43e5-ac48-d1d0a3aa0dca",
   "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",
      "It's nice to meet you, Bob! I'm an AI assistant created by Anthropic. I'm here to help out with any questions or tasks you might have. Please let me know if there's anything I can assist you with.\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what's my name?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "You said your name is Bob.\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
    "input_message = HumanMessage(content=\"hi! I'm bob\")\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
    "    event[\"messages\"][-1].pretty_print()\n",
    "\n",
    "\n",
    "input_message = HumanMessage(content=\"what's my name?\")\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
    "    event[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fb0de5b-30ec-42d4-813a-7ad63fe1c367",
   "metadata": {},
   "source": [
    "## Manually deleting messages\n",
    "\n",
    "First, we will cover how to manually delete messages. Let's take a look at the current state of the thread:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8a850529-d038-48f7-b5a2-8d4d2923f83a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content=\"hi! I'm bob\", additional_kwargs={}, response_metadata={}, id='db576005-3a60-4b3b-8925-dc602ac1c571'),\n",
       " AIMessage(content=\"It's nice to meet you, Bob! I'm an AI assistant created by Anthropic. I'm here to help out with any questions or tasks you might have. Please let me know if there's anything I can assist you with.\", additional_kwargs={}, response_metadata={'id': 'msg_01BKAnYxmoC6bQ9PpCuHk8ZT', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 52}}, id='run-3a60c536-b207-4c56-98f3-03f94d49a9e4-0', usage_metadata={'input_tokens': 12, 'output_tokens': 52, 'total_tokens': 64}),\n",
       " HumanMessage(content=\"what's my name?\", additional_kwargs={}, response_metadata={}, id='2088c465-400b-430b-ad80-fad47dc1f2d6'),\n",
       " AIMessage(content='You said your name is Bob.', additional_kwargs={}, response_metadata={'id': 'msg_013UWTLTzwZi81vke8mMQ2KP', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 72, 'output_tokens': 10}}, id='run-3a6883be-0c52-4938-af98-e9e7476659eb-0', usage_metadata={'input_tokens': 72, 'output_tokens': 10, 'total_tokens': 82})]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages = app.get_state(config).values[\"messages\"]\n",
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81be8a0a-1e94-4302-bd84-d1b72e3c501c",
   "metadata": {},
   "source": [
    "We can call `update_state` and pass in the id of the first message. This will delete that message."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "df1a0970-7e64-4170-beef-2855d10eef42",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'configurable': {'thread_id': '2',\n",
       "  'checkpoint_ns': '',\n",
       "  'checkpoint_id': '1ef75157-f251-6a2a-8005-82a86a6593a0'}}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import RemoveMessage\n",
    "\n",
    "app.update_state(config, {\"messages\": RemoveMessage(id=messages[0].id)})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c9127ae-0d42-42b8-957f-ea69a5da555f",
   "metadata": {},
   "source": [
    "If we now look at the messages, we can verify that the first one was deleted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8bfe4ffa-e170-43bc-aec4-6e36ac620931",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[AIMessage(content=\"It's nice to meet you, Bob! I'm Claude, an AI assistant created by Anthropic. How can I assist you today?\", response_metadata={'id': 'msg_01XPSAenmSqK8rX2WgPZHfz7', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 32}}, id='run-1c69af09-adb1-412d-9010-2456e5a555fb-0', usage_metadata={'input_tokens': 12, 'output_tokens': 32, 'total_tokens': 44}),\n",
       " HumanMessage(content=\"what's my name?\", id='f3c71afe-8ce2-4ed0-991e-65021f03b0a5'),\n",
       " AIMessage(content='Your name is Bob, as you introduced yourself at the beginning of our conversation.', response_metadata={'id': 'msg_01BPZdwsjuMAbC1YAkqawXaF', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 52, 'output_tokens': 19}}, id='run-b2eb9137-2f4e-446f-95f5-3d5f621a2cf8-0', usage_metadata={'input_tokens': 52, 'output_tokens': 19, 'total_tokens': 71})]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages = app.get_state(config).values[\"messages\"]\n",
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef129a75-4cad-44d7-b532-eb37b0553c0c",
   "metadata": {},
   "source": [
    "## Programmatically deleting messages\n",
    "\n",
    "We can also delete messages programmatically from inside the graph. Here we'll modify the graph to delete any old messages (longer than 3 messages ago) at the end of a graph run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bb22ede0-e153-4fd0-a4c0-f9af2f7663b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import RemoveMessage\n",
    "from langgraph.graph import END\n",
    "\n",
    "\n",
    "def delete_messages(state):\n",
    "    messages = state[\"messages\"]\n",
    "    if len(messages) > 3:\n",
    "        return {\"messages\": [RemoveMessage(id=m.id) for m in messages[:-3]]}\n",
    "\n",
    "\n",
    "# We need to modify the logic to call delete_messages rather than end right away\n",
    "def should_continue(state: MessagesState) -> Literal[\"action\", \"delete_messages\"]:\n",
    "    \"\"\"Return the next node to execute.\"\"\"\n",
    "    last_message = state[\"messages\"][-1]\n",
    "    # If there is no function call, then we call our delete_messages function\n",
    "    if not last_message.tool_calls:\n",
    "        return \"delete_messages\"\n",
    "    # Otherwise if there is, we continue\n",
    "    return \"action\"\n",
    "\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(MessagesState)\n",
    "workflow.add_node(\"agent\", call_model)\n",
    "workflow.add_node(\"action\", tool_node)\n",
    "\n",
    "# This is our new node we're defining\n",
    "workflow.add_node(delete_messages)\n",
    "\n",
    "\n",
    "workflow.add_edge(START, \"agent\")\n",
    "workflow.add_conditional_edges(\n",
    "    \"agent\",\n",
    "    should_continue,\n",
    ")\n",
    "workflow.add_edge(\"action\", \"agent\")\n",
    "\n",
    "# This is the new edge we're adding: after we delete messages, we finish\n",
    "workflow.add_edge(\"delete_messages\", END)\n",
    "app = workflow.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52cbdef6-7db7-45a2-8194-de4f8929bd1f",
   "metadata": {},
   "source": [
    "We can now try this out. We can call the graph twice and then check the state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3975f34c-c243-40ea-b9d2-424d50a48dc9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('human', \"hi! I'm bob\")]\n",
      "[('human', \"hi! I'm bob\"), ('ai', \"Hello Bob! It's nice to meet you. I'm an AI assistant created by Anthropic. I'm here to help with any questions or tasks you might have. Please let me know how I can assist you.\")]\n",
      "[('human', \"hi! I'm bob\"), ('ai', \"Hello Bob! It's nice to meet you. I'm an AI assistant created by Anthropic. I'm here to help with any questions or tasks you might have. Please let me know how I can assist you.\"), ('human', \"what's my name?\")]\n",
      "[('human', \"hi! I'm bob\"), ('ai', \"Hello Bob! It's nice to meet you. I'm an AI assistant created by Anthropic. I'm here to help with any questions or tasks you might have. Please let me know how I can assist you.\"), ('human', \"what's my name?\"), ('ai', 'You said your name is Bob, so that is the name I have for you.')]\n",
      "[('ai', \"Hello Bob! It's nice to meet you. I'm an AI assistant created by Anthropic. I'm here to help with any questions or tasks you might have. Please let me know how I can assist you.\"), ('human', \"what's my name?\"), ('ai', 'You said your name is Bob, so that is the name I have for you.')]\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"3\"}}\n",
    "input_message = HumanMessage(content=\"hi! I'm bob\")\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
    "    print([(message.type, message.content) for message in event[\"messages\"]])\n",
    "\n",
    "\n",
    "input_message = HumanMessage(content=\"what's my name?\")\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
    "    print([(message.type, message.content) for message in event[\"messages\"]])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67b2fd2a-14a1-4c47-8632-f8cbb0ba1d35",
   "metadata": {},
   "source": [
    "If we now check the state, we should see that it is only three messages long. This is because we just deleted the earlier messages - otherwise it would be four!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a3e15abb-81d8-4072-9f10-61ae0fd61dac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[AIMessage(content=\"Hello Bob! It's nice to meet you. I'm an AI assistant created by Anthropic. I'm here to help with any questions or tasks you might have. Please let me know how I can assist you.\", response_metadata={'id': 'msg_01XPEgPPbcnz5BbGWUDWTmzG', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 48}}, id='run-eded3820-b6a9-4d66-9210-03ca41787ce6-0', usage_metadata={'input_tokens': 12, 'output_tokens': 48, 'total_tokens': 60}),\n",
       " HumanMessage(content=\"what's my name?\", id='a0ea2097-3280-402b-92e1-67177b807ae8'),\n",
       " AIMessage(content='You said your name is Bob, so that is the name I have for you.', response_metadata={'id': 'msg_01JGT62pxhrhN4SykZ57CSjW', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 68, 'output_tokens': 20}}, id='run-ace3519c-81f8-45fe-a777-91f42d48b3a3-0', usage_metadata={'input_tokens': 68, 'output_tokens': 20, 'total_tokens': 88})]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages = app.get_state(config).values[\"messages\"]\n",
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "359cfeae-d43a-46ee-9069-a1cab9a5720a",
   "metadata": {},
   "source": [
    "Remember, when deleting messages you will want to make sure that the remaining message list is still valid. This message list **may actually not be** - this is because it currently starts with an AI message, which some models do not allow."
   ]
  }
 ],
 "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
}