{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
   "metadata": {},
   "source": [
    "# How to add summary of the conversation history\n",
    "\n",
    "One of the most common use cases for persistence is to use it to keep track of conversation history. This is great - it makes it easy to continue conversations. As conversations get longer and longer, however, this conversation history can build up and take up more and more of the context window. This can often be undesirable as it leads to more expensive and longer calls to the LLM, and potentially ones that error. One way to work around that is to create a summary of the conversation to date, and use that with the past N messages. This guide will go through an example of how to do that.\n",
    "\n",
    "This will involve a few steps:\n",
    "- Check if the conversation is too long (can be done by checking number of messages or length of messages)\n",
    "- If yes, the create summary (will need a prompt for this)\n",
    "- Then remove all except the last N messages\n",
    "\n",
    "A big part of this is deleting old messages. For an in depth guide on how to do that, see [this guide](./delete-messages.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cbd446a-808f-4394-be92-d45ab818953c",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, let's set up the packages we're going to want to use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": [],
   "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": [
    "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": "84835fdb-a5f3-4c90-85f3-0e6257650aba",
   "metadata": {},
   "source": [
    "## Build the chatbot\n",
    "\n",
    "Let's now build the chatbot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "378899a9-3b9a-4748-95b6-eb00e0828677",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "\n",
    "from langchain_anthropic import ChatAnthropic\n",
    "from langchain_core.messages import SystemMessage, RemoveMessage\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import MessagesState, StateGraph, START, END\n",
    "\n",
    "memory = MemorySaver()\n",
    "\n",
    "\n",
    "# We will add a `summary` attribute (in addition to `messages` key,\n",
    "# which MessagesState already has)\n",
    "class State(MessagesState):\n",
    "    summary: str\n",
    "\n",
    "\n",
    "# We will use this model for both the conversation and the summarization\n",
    "model = ChatAnthropic(model_name=\"claude-3-haiku-20240307\")\n",
    "\n",
    "\n",
    "# Define the logic to call the model\n",
    "def call_model(state: State):\n",
    "    # If a summary exists, we add this in as a system message\n",
    "    summary = state.get(\"summary\", \"\")\n",
    "    if summary:\n",
    "        system_message = f\"Summary of conversation earlier: {summary}\"\n",
    "        messages = [SystemMessage(content=system_message)] + state[\"messages\"]\n",
    "    else:\n",
    "        messages = state[\"messages\"]\n",
    "    response = model.invoke(messages)\n",
    "    # We return a list, because this will get added to the existing list\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "# We now define the logic for determining whether to end or summarize the conversation\n",
    "def should_continue(state: State) -> Literal[\"summarize_conversation\", END]:\n",
    "    \"\"\"Return the next node to execute.\"\"\"\n",
    "    messages = state[\"messages\"]\n",
    "    # If there are more than six messages, then we summarize the conversation\n",
    "    if len(messages) > 6:\n",
    "        return \"summarize_conversation\"\n",
    "    # Otherwise we can just end\n",
    "    return END\n",
    "\n",
    "\n",
    "def summarize_conversation(state: State):\n",
    "    # First, we summarize the conversation\n",
    "    summary = state.get(\"summary\", \"\")\n",
    "    if summary:\n",
    "        # If a summary already exists, we use a different system prompt\n",
    "        # to summarize it than if one didn't\n",
    "        summary_message = (\n",
    "            f\"This is summary of the conversation to date: {summary}\\n\\n\"\n",
    "            \"Extend the summary by taking into account the new messages above:\"\n",
    "        )\n",
    "    else:\n",
    "        summary_message = \"Create a summary of the conversation above:\"\n",
    "\n",
    "    messages = state[\"messages\"] + [HumanMessage(content=summary_message)]\n",
    "    response = model.invoke(messages)\n",
    "    # We now need to delete messages that we no longer want to show up\n",
    "    # I will delete all but the last two messages, but you can change this\n",
    "    delete_messages = [RemoveMessage(id=m.id) for m in state[\"messages\"][:-2]]\n",
    "    return {\"summary\": response.content, \"messages\": delete_messages}\n",
    "\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(State)\n",
    "\n",
    "# Define the conversation node and the summarize node\n",
    "workflow.add_node(\"conversation\", call_model)\n",
    "workflow.add_node(summarize_conversation)\n",
    "\n",
    "# Set the entrypoint as conversation\n",
    "workflow.add_edge(START, \"conversation\")\n",
    "\n",
    "# We now add a conditional edge\n",
    "workflow.add_conditional_edges(\n",
    "    # First, we define the start node. We use `conversation`.\n",
    "    # This means these are the edges taken after the `conversation` node is called.\n",
    "    \"conversation\",\n",
    "    # Next, we pass in the function that will determine which node is called next.\n",
    "    should_continue,\n",
    ")\n",
    "\n",
    "# We now add a normal edge from `summarize_conversation` to END.\n",
    "# This means that after `summarize_conversation` is called, we end.\n",
    "workflow.add_edge(\"summarize_conversation\", END)\n",
    "\n",
    "# Finally, we compile it!\n",
    "app = workflow.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41c2872e-04b3-4c44-9e03-9e84a5230adf",
   "metadata": {},
   "source": [
    "## Using the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "dc697132-8fa1-4bf5-9722-56a9859331ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_update(update):\n",
    "    for k, v in update.items():\n",
    "        for m in v[\"messages\"]:\n",
    "            m.pretty_print()\n",
    "        if \"summary\" in v:\n",
    "            print(v[\"summary\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "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. How can I help you today?\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what's my name?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Your name is Bob, as you told me at the beginning of our conversation.\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "i like the celtics!\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "That's great, the Celtics are a fun team to follow! Basketball is an exciting sport. Do you have a favorite Celtics player or a favorite moment from a Celtics game you've watched? I'd be happy to discuss the team and the sport with you.\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
    "input_message = HumanMessage(content=\"hi! I'm bob\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)\n",
    "\n",
    "input_message = HumanMessage(content=\"what's my name?\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)\n",
    "\n",
    "input_message = HumanMessage(content=\"i like the celtics!\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9760e219-a7fc-4d81-b4e8-1334c5afc510",
   "metadata": {},
   "source": [
    "We can see that so far no summarization has happened - this is because there are only six messages in the list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "935265a0-d511-475a-8a0d-b3c3cc5e42a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content=\"hi! I'm bob\", id='6534853d-b8a7-44b9-837b-eb7abaf7ebf7'),\n",
       "  AIMessage(content=\"It's nice to meet you, Bob! I'm an AI assistant created by Anthropic. How can I help you today?\", response_metadata={'id': 'msg_015wCFew2vwMQJcpUh2VZ5ah', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 30}}, id='run-0d33008b-1094-4f5e-94ce-293283fc3024-0'),\n",
       "  HumanMessage(content=\"what's my name?\", id='0a4f203a-b95a-42a9-b1c5-bb20f68b3251'),\n",
       "  AIMessage(content='Your name is Bob, as you told me at the beginning of our conversation.', response_metadata={'id': 'msg_01PLp8wg2xDsJbNR9uCtxcGz', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 50, 'output_tokens': 19}}, id='run-3815dd4d-ee0c-4fc2-9889-f6dd40325961-0'),\n",
       "  HumanMessage(content='i like the celtics!', id='ac128172-42d1-4390-b7cc-7bcb2d22ee48'),\n",
       "  AIMessage(content=\"That's great, the Celtics are a fun team to follow! Basketball is an exciting sport. Do you have a favorite Celtics player or a favorite moment from a Celtics game you've watched? I'd be happy to discuss the team and the sport with you.\", response_metadata={'id': 'msg_01CSg5avZEx6CKcZsSvSVXpr', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 78, 'output_tokens': 61}}, id='run-698faa28-0f72-495f-8ebe-e948664d2200-0')]}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "values = app.get_state(config).values\n",
    "values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb40eddb-9a31-4410-a4c0-9762e2d89e56",
   "metadata": {},
   "source": [
    "Now let's send another message in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "048805a4-3d97-4e76-ac45-8d80d4364c46",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "i like how much they win\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "That's understandable, the Celtics have been one of the more successful NBA franchises over the years. Their history of winning championships is very impressive. It's always fun to follow a team that regularly competes for titles. What do you think has been the key to the Celtics' sustained success? Is there a particular era or team that stands out as your favorite?\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "Here is a summary of our conversation so far:\n",
      "\n",
      "- You introduced yourself as Bob and said you like the Boston Celtics basketball team.\n",
      "- I acknowledged that it's nice to meet you, Bob, and noted that you had shared your name earlier in the conversation.\n",
      "- You expressed that you like how much the Celtics win, and I agreed that their history of sustained success and championship pedigree is impressive.\n",
      "- I asked if you have a favorite Celtics player or moment that stands out to you, and invited further discussion about the team and the sport of basketball.\n",
      "- The overall tone has been friendly and conversational, with me trying to engage with your interest in the Celtics by asking follow-up questions.\n"
     ]
    }
   ],
   "source": [
    "input_message = HumanMessage(content=\"i like how much they win\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b196367-6151-4982-9430-3db7373de06e",
   "metadata": {},
   "source": [
    "If we check the state now, we can see that we have a summary of the conversation, as well as the last two messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "09ebb693-4738-4474-a095-6491def5c5f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='i like how much they win', id='bb916ce7-534c-4d48-9f92-e269f9dc4859'),\n",
       "  AIMessage(content=\"That's understandable, the Celtics have been one of the more successful NBA franchises over the years. Their history of winning championships is very impressive. It's always fun to follow a team that regularly competes for titles. What do you think has been the key to the Celtics' sustained success? Is there a particular era or team that stands out as your favorite?\", response_metadata={'id': 'msg_01B7TMagaM8xBnYXLSMwUDAG', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 148, 'output_tokens': 82}}, id='run-c5aa9a8f-7983-4a7f-9c1e-0c0055334ac1-0')],\n",
       " 'summary': \"Here is a summary of our conversation so far:\\n\\n- You introduced yourself as Bob and said you like the Boston Celtics basketball team.\\n- I acknowledged that it's nice to meet you, Bob, and noted that you had shared your name earlier in the conversation.\\n- You expressed that you like how much the Celtics win, and I agreed that their history of sustained success and championship pedigree is impressive.\\n- I asked if you have a favorite Celtics player or moment that stands out to you, and invited further discussion about the team and the sport of basketball.\\n- The overall tone has been friendly and conversational, with me trying to engage with your interest in the Celtics by asking follow-up questions.\"}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "values = app.get_state(config).values\n",
    "values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "966e4177-c0fc-4fd0-a494-dd03f7f2fddb",
   "metadata": {},
   "source": [
    "We can now resume having a conversation! Note that even though we only have the last two messages, we can still ask it questions about things mentioned earlier in the conversation (because we summarized those)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "7094c5ab-66f8-42ff-b1c3-90c8a9468e62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what's my name?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "In our conversation so far, you introduced yourself as Bob. I acknowledged that earlier when you had shared your name.\n"
     ]
    }
   ],
   "source": [
    "input_message = HumanMessage(content=\"what's my name?\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "40e5db8e-9db9-4ac7-9d76-a99fd4034bf3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what NFL team do you think I like?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "I don't actually have any information about what NFL team you might like. In our conversation so far, you've only mentioned that you're a fan of the Boston Celtics basketball team. I don't have any prior knowledge about your preferences for NFL teams. Unless you provide me with that information, I don't have a basis to guess which NFL team you might be a fan of.\n"
     ]
    }
   ],
   "source": [
    "input_message = HumanMessage(content=\"what NFL team do you think I like?\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "0a1a0fda-5309-45f0-9465-9f3dff604d74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "i like the patriots!\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Okay, got it! Thanks for sharing that you're also a fan of the New England Patriots in the NFL. That makes sense, given your interest in other Boston sports teams like the Celtics. The Patriots have also had a very successful run over the past couple of decades, winning multiple Super Bowls. It's fun to follow winning franchises like the Celtics and Patriots. Do you have a favorite Patriots player or moment that stands out to you?\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "Okay, extending the summary with the new information:\n",
      "\n",
      "- You initially introduced yourself as Bob and said you like the Boston Celtics basketball team. \n",
      "- I acknowledged that and we discussed your appreciation for the Celtics' history of winning.\n",
      "- You then asked what your name was, and I reminded you that you had introduced yourself as Bob earlier in the conversation.\n",
      "- You followed up by asking what NFL team I thought you might like, and I explained that I didn't have any prior information about your NFL team preferences.\n",
      "- You then revealed that you are also a fan of the New England Patriots, which made sense given your Celtics fandom.\n",
      "- I responded positively to this new information, noting the Patriots' own impressive success and dynasty over the past couple of decades.\n",
      "- I then asked if you have a particular favorite Patriots player or moment that stands out to you, continuing the friendly, conversational tone.\n",
      "\n",
      "Overall, the discussion has focused on your sports team preferences, with you sharing that you are a fan of both the Celtics and the Patriots. I've tried to engage with your interests and ask follow-up questions to keep the dialogue flowing.\n"
     ]
    }
   ],
   "source": [
    "input_message = HumanMessage(content=\"i like the patriots!\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67d26013-1362-4cee-b135-ab5c3c4eb3d0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
