{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7240d5b5-9dac-4070-8a9e-2350fb01e0be",
   "metadata": {},
   "source": [
    "# How to share state between threads\n",
    "\n",
    "By default, state in a graph is scoped to that thread.\n",
    "LangGraph also allows you to specify a \"scope\" for a given key/value pair that exists between threads. This can be useful for storing information that is shared between threads. For instance, you may want to store information about a user's preferences expressed in one thread, and then use that information in another thread.\n",
    "\n",
    "In this notebook we will go through an example of how to construct and use such a graph.\n",
    "\n",
    "## Setup\n",
    "\n",
    "First, let's install the required packages and set our API keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3457aadf",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langchain_openai langgraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa2c64a7",
   "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(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51b6817d",
   "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": "c4c550b5-1954-496b-8b9d-800361af17dc",
   "metadata": {},
   "source": [
    "## Create graph\n",
    "\n",
    "In this example we will create a graph that will let us store information about a user's preferences. We will do so by defining a state key that will be scoped to a user_id, and allowing the model to populate this field as it deems fit (by providing the model with a tool to save information about the user).\n",
    "\n",
    "    \n",
    "<div class=\"admonition note\">\n",
    "    <p class=\"admonition-title\">Typing shared state keys</p>\n",
    "    <p style=\"margin-top: 5px;\">\n",
    "        Shared state channels (keys) MUST be dictionaries (see <code>info</code> channel in the AgentState example below)\n",
    "    </p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a7f303d6-612e-4e34-bf36-29d4ed25d802",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph.graph import START, END\n",
    "from langgraph.graph.message import MessagesState\n",
    "from langgraph.graph.state import StateGraph\n",
    "from langgraph.store.memory import MemoryStore\n",
    "from langgraph.managed.shared_value import SharedValue\n",
    "from typing import TypedDict, Annotated, Any\n",
    "import uuid\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "\n",
    "class AgentState(MessagesState):\n",
    "    # We use an info key to track information\n",
    "    # This is scoped to a user_id, so it will be information specific to each user\n",
    "    info: Annotated[dict, SharedValue.on(\"user_id\")]\n",
    "\n",
    "\n",
    "# We will give this as a tool to the agent\n",
    "# This will let the agent call this tool to save a fact\n",
    "class Info(TypedDict):\n",
    "    \"\"\"This tool should be called when you want to save a new fact about the user.\n",
    "    \n",
    "    Attributes:\n",
    "        fact (str): A fact about the user.\n",
    "        topic (str): The topic related the fact is about, i.e. Food, Location, Movies, etc.\n",
    "    \"\"\"\n",
    "    fact: str\n",
    "    topic: str\n",
    "\n",
    "\n",
    "# This is the prompt we give the agent\n",
    "# We will pass known info into the prompt\n",
    "# We will tell it to use the Info tool to save more\n",
    "prompt = \"\"\"You are helpful assistant.\n",
    "\n",
    "Here is what you know about the user:\n",
    "\n",
    "<info>\n",
    "{info}\n",
    "</info>\n",
    "\n",
    "Help out the user. If the user tells you any information about themselves, save the information using the `Info` tool.\n",
    "\n",
    "This means if the user provides any sort of fact about themselves, be it an opinion they have, a fact about themselves, etc. SAVE IT!\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "# We give the model access to the Info tool\n",
    "model = ChatOpenAI().bind_tools([Info])\n",
    "\n",
    "\n",
    "# Our first node - this will call the model\n",
    "def call_model(state):\n",
    "    # We get all facts and assemble them into a string\n",
    "    facts = [d['fact'] for d in state['info'].values()]\n",
    "    info = \"\\n\".join(facts)\n",
    "    # Format system prompt\n",
    "    system_msg = prompt.format(info=info)\n",
    "    # Call model\n",
    "    response = model.invoke([{\"role\": \"system\", \"content\": system_msg}] + state['messages'])\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "# Routing function to decide what to do next\n",
    "# If no tool calls, then we end\n",
    "# If tool calls, then we update memory\n",
    "def route(state):\n",
    "    if len(state['messages'][-1].tool_calls) == 0:\n",
    "        return END\n",
    "    else:\n",
    "        return \"update_memory\"\n",
    "\n",
    "\n",
    "# This function is responsible for updating the memory\n",
    "def update_memory(state):\n",
    "    tool_calls = []\n",
    "    memories = {}\n",
    "    # Each tool call is a new memory to save\n",
    "    for tc in state['messages'][-1].tool_calls:\n",
    "        # We append ToolMessages (to pass back to the LLM)\n",
    "        # This is needed because OpenAI requires each tool call be followed by a ToolMessage\n",
    "        tool_calls.append({\"role\": \"tool\", \"content\": \"Saved!\", \"tool_call_id\": tc['id']})\n",
    "        # We create a new memory from this tool call\n",
    "        memories[str(uuid.uuid4())] = {\"fact\": tc['args']['fact'], \"topic\": tc['args']['topic']}\n",
    "    # Return the messages and memories to update the state with\n",
    "    return {\"messages\": tool_calls, \"info\": memories}\n",
    "\n",
    "\n",
    "# This is the in memory checkpointer we will use\n",
    "# We need this because we want to enable threads (conversations)\n",
    "memory = MemorySaver()\n",
    "\n",
    "# This is the in memory Key Value store\n",
    "# This is needed to save the memories\n",
    "kv = MemoryStore()\n",
    "\n",
    "# Construct this relatively simple graph\n",
    "graph = StateGraph(AgentState)\n",
    "graph.add_node(call_model)\n",
    "graph.add_node(update_memory)\n",
    "graph.add_edge(\"update_memory\", END)\n",
    "graph.add_edge(START, \"call_model\")\n",
    "graph.add_conditional_edges(\"call_model\", route)\n",
    "graph = graph.compile(checkpointer=memory, store=kv)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "552d4e33-556d-4fa5-8094-2a076bc21529",
   "metadata": {},
   "source": [
    "## Run graph on one thread\n",
    "\n",
    "We can now run the graph on one thread and give it some information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "18bd8679-3a73-4033-bfb4-5093ac1f5d7f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'call_model': {'messages': [AIMessage(content='Hello! How can I assist you today?', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 171, 'total_tokens': 181}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-fbbb73a4-7c94-4db1-8761-44ea2fe9feaf-0', usage_metadata={'input_tokens': 171, 'output_tokens': 10, 'total_tokens': 181})]}}\n",
      "{'call_model': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_zMUXZfhOCFYvZg5TwXyBzw16', 'function': {'arguments': '{\"fact\":\"I like pepperoni pizza\",\"topic\":\"Food\"}', 'name': 'Info'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 21, 'prompt_tokens': 193, 'total_tokens': 214}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-7297f9fb-1d3e-480e-b125-ab269f648158-0', tool_calls=[{'name': 'Info', 'args': {'fact': 'I like pepperoni pizza', 'topic': 'Food'}, 'id': 'call_zMUXZfhOCFYvZg5TwXyBzw16', 'type': 'tool_call'}], usage_metadata={'input_tokens': 193, 'output_tokens': 21, 'total_tokens': 214})]}}\n",
      "{'update_memory': {'messages': [{'role': 'tool', 'content': 'Saved!', 'tool_call_id': 'call_zMUXZfhOCFYvZg5TwXyBzw16'}]}}\n",
      "{'call_model': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_GjshujJAeqoTuuBeHCD5YTPQ', 'function': {'arguments': '{\"fact\":\"I just moved to SF\",\"topic\":\"Location\"}', 'name': 'Info'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 21, 'prompt_tokens': 239, 'total_tokens': 260}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-4abea1d6-7ccb-49b4-b805-0e04ebb542e3-0', tool_calls=[{'name': 'Info', 'args': {'fact': 'I just moved to SF', 'topic': 'Location'}, 'id': 'call_GjshujJAeqoTuuBeHCD5YTPQ', 'type': 'tool_call'}], usage_metadata={'input_tokens': 239, 'output_tokens': 21, 'total_tokens': 260})]}}\n",
      "{'update_memory': {'messages': [{'role': 'tool', 'content': 'Saved!', 'tool_call_id': 'call_GjshujJAeqoTuuBeHCD5YTPQ'}]}}\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"1\", \"user_id\": \"1\"}}\n",
    "\n",
    "# First let's just say hi to the AI\n",
    "for update in graph.stream({\"messages\": [{\"role\": \"user\", \"content\": \"hi\"}]}, config, stream_mode=\"updates\"):\n",
    "    print(update)\n",
    "\n",
    "# Let's continue the conversation (by passing the same config) and tell the AI we like pepperoni pizza\n",
    "for update in graph.stream({\"messages\": [{\"role\": \"user\", \"content\": \"i like pepperoni pizza\"}]}, config, stream_mode=\"updates\"):\n",
    "    print(update)\n",
    "\n",
    "# Let's continue the conversation even further (by passing the same config) and tell the AI we live in SF\n",
    "for update in graph.stream({\"messages\": [{\"role\": \"user\", \"content\": \"i also just moved to SF\"}]}, config, stream_mode=\"updates\"):\n",
    "    print(update)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8c416fa-086a-491d-a7d3-57091f6413e3",
   "metadata": {},
   "source": [
    "## Run graph on a different thread\n",
    "\n",
    "We can now run the graph on a different thread and see that it remembers facts about the user (specifically that the user likes pepperoni pizza and lives in SF):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e240f025-ff8b-4d17-beb7-2420c0575dd9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'call_model': {'messages': [AIMessage(content=\"Sure! Since you just moved to San Francisco, how about trying some popular local spots? Here are a few restaurant recommendations in SF:\\n\\n1. Tony's Pizza Napoletana - Known for their delicious pepperoni pizza!\\n2. The Slanted Door - A popular Vietnamese restaurant in the city.\\n3. Zuni Cafe - A classic American restaurant with a great ambiance.\\n4. Tartine Bakery - Perfect for a casual dinner with amazing baked goods.\\n5. State Bird Provisions - A unique dining experience with small plates and a lively atmosphere.\\n\\nFeel free to explore these options and enjoy your dinner! If you need more recommendations or information about a specific cuisine, let me know!\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 138, 'prompt_tokens': 197, 'total_tokens': 335}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-de8ad08c-0810-4bb5-b2e8-d3dc89522f8e-0', usage_metadata={'input_tokens': 197, 'output_tokens': 138, 'total_tokens': 335})]}}\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"2\", \"user_id\": \"1\"}}\n",
    "\n",
    "for update in graph.stream({\"messages\": [{\"role\": \"user\", \"content\": \"where and what should i eat for dinner? Can you list some restaurants?\"}]}, config, stream_mode=\"updates\"):\n",
    "    print(update)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "091995d3",
   "metadata": {},
   "source": [
    "Perfect! The AI recommended restaurants in SF, and included a pizza restaurant at the top of it's list.\n",
    "\n",
    "Notice that the `messages` in this new thread do NOT contain the messages from the previous thread since we didn't store them as shared values across the `user_id`. However, the `info` we saved in the previous thread was saved since we passed in the same `user_id` in this new thread.\n",
    "\n",
    "Let's now run the graph for another user to verify that the preferences of the first user are self contained:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f9bf2c15",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'call_model': {'messages': [AIMessage(content='I can definitely help you with that! To provide you with personalized restaurant recommendations, could you please let me know your location or any specific preferences you have for dinner?', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 34, 'prompt_tokens': 185, 'total_tokens': 219}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-5a483acf-1289-4d7f-b707-97760a8c3620-0', usage_metadata={'input_tokens': 185, 'output_tokens': 34, 'total_tokens': 219})]}}\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"3\", \"user_id\": \"2\"}}\n",
    "\n",
    "for update in graph.stream({\"messages\": [{\"role\": \"user\", \"content\": \"where and what should i eat for dinner? Can you list some restaurants?\"}]}, config, stream_mode=\"updates\"):\n",
    "    print(update)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7086cea",
   "metadata": {},
   "source": [
    "Perfect! The graph has forgotten all of the previous preferences and has to ask the user for it's location and dietary preferences."
   ]
  }
 ],
 "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
}