{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3631f2b9-aa79-472e-a9d6-9125a90ee704",
   "metadata": {},
   "source": [
    "# How to stream state updates of your graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "858c7499-0c92-40a9-bd95-e5a5a5817e92",
   "metadata": {},
   "source": [
    "LangGraph supports multiple streaming modes. The main ones are:\n",
    "\n",
    "- `values`: This streaming mode streams back values of the graph. This is the **full state of the graph** after each node is called.\n",
    "- `updates`: This streaming mode streams back updates to the graph. This is the **update to the state of the graph** after each node is called.\n",
    "\n",
    "This guide covers `stream_mode=\"updates\"`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c2f84f1-0751-4779-97d4-5cbb286093b7",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "323db423-b644-40bd-9c2d-976a53f602f7",
   "metadata": {},
   "source": [
    "We'll be using a simple ReAct agent for this guide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6b4285e4-7434-4971-bde0-aabceef8ee7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langgraph langchain-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f7f9f24a-e3d0-422b-8924-47950b2facd6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "OPENAI_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(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "85cf2e23-29f2-40cc-b302-5377b3b49da9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "from langchain_core.runnables import ConfigurableField\n",
    "from langchain_core.tools import tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "\n",
    "@tool\n",
    "def get_weather(city: Literal[\"nyc\", \"sf\"]):\n",
    "    \"\"\"Use this to get weather information.\"\"\"\n",
    "    if city == \"nyc\":\n",
    "        return \"It might be cloudy in nyc\"\n",
    "    elif city == \"sf\":\n",
    "        return \"It's always sunny in sf\"\n",
    "    else:\n",
    "        raise AssertionError(\"Unknown city\")\n",
    "\n",
    "\n",
    "tools = [get_weather]\n",
    "\n",
    "model = ChatOpenAI(model_name=\"gpt-4o\", temperature=0)\n",
    "graph = create_react_agent(model, tools)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "956db549-5207-4be1-a823-78311738e3f8",
   "metadata": {},
   "source": [
    "## Stream updates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e9e9ffb0-2cd5-466f-b70b-b6ed51b852d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Receiving update from node: 'agent'\n",
      "{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_kc6cvcEkTAUGRlSHrP4PK9fn', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_3e7d703517', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-cd68b3a0-86c3-4afa-9649-1b962a0dd062-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_kc6cvcEkTAUGRlSHrP4PK9fn'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71})]}\n",
      "\n",
      "\n",
      "\n",
      "Receiving update from node: 'tools'\n",
      "{'messages': [ToolMessage(content=\"It's always sunny in sf\", name='get_weather', tool_call_id='call_kc6cvcEkTAUGRlSHrP4PK9fn')]}\n",
      "\n",
      "\n",
      "\n",
      "Receiving update from node: 'agent'\n",
      "{'messages': [AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_3e7d703517', 'finish_reason': 'stop', 'logprobs': None}, id='run-009d83c4-b874-4acc-9494-20aba43132b9-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})]}\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "inputs = {\"messages\": [(\"human\", \"what's the weather in sf\")]}\n",
    "async for chunk in graph.astream(inputs, stream_mode=\"updates\"):\n",
    "    for node, values in chunk.items():\n",
    "        print(f\"Receiving update from node: '{node}'\")\n",
    "        print(values)\n",
    "        print(\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cc57240-243c-4d9a-a845-cb55ef973a59",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langgraph",
   "language": "python",
   "name": "langgraph"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}