{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3631f2b9-aa79-472e-a9d6-9125a90ee704",
   "metadata": {},
   "source": [
    "# How to stream full state 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=\"values\"`."
   ]
  },
  {
   "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": "ef5a3ec6-0cd0-4541-ab1b-d63ede22720e",
   "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": "002a715b-e0be-4e89-8d42-f0098882586b",
   "metadata": {},
   "source": [
    "## Stream values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e9e9ffb0-2cd5-466f-b70b-b6ed51b852d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what's the weather in sf\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  get_weather (call_61VvIzqVGtyxcXi0z6knZkjZ)\n",
      " Call ID: call_61VvIzqVGtyxcXi0z6knZkjZ\n",
      "  Args:\n",
      "    city: sf\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: get_weather\n",
      "\n",
      "It's always sunny in sf\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The weather in San Francisco is currently sunny.\n"
     ]
    }
   ],
   "source": [
    "inputs = {\"messages\": [(\"human\", \"what's the weather in sf\")]}\n",
    "async for chunk in graph.astream(inputs, stream_mode=\"values\"):\n",
    "    chunk[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d73de237-bf45-4fa7-93ef-6dae7eacffc0",
   "metadata": {},
   "source": [
    "If we want to just get the final result, we can use the same method and just keep track of the last value we received"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c122bf15-a489-47bf-b482-a744a54e2cc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {\"messages\": [(\"human\", \"what's the weather in sf\")]}\n",
    "async for chunk in graph.astream(inputs, stream_mode=\"values\"):\n",
    "    final_result = chunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "316022e5-4c65-48e4-9878-8d94a2425ed4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content=\"what's the weather in sf\", id='54b39b6f-054b-4306-980b-86905e48a6bc'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_avoKnK8reERzTUSxrN9cgFxY', '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_5e6c71d4a8', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-f2f43c89-2c96-45f4-975c-2d0f22d0d2d1-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_avoKnK8reERzTUSxrN9cgFxY'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}),\n",
       "  ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='fc18a798-c7b2-4f73-84fa-8ffdffb6ddcb', tool_call_id='call_avoKnK8reERzTUSxrN9cgFxY'),\n",
       "  AIMessage(content='The weather in San Francisco is currently sunny. Enjoy the sunshine!', response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 84, 'total_tokens': 98}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_5e6c71d4a8', 'finish_reason': 'stop', 'logprobs': None}, id='run-21418147-da8e-4738-a076-239377397c40-0', usage_metadata={'input_tokens': 84, 'output_tokens': 14, 'total_tokens': 98})]}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "final_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0f64ebbe-535c-4b35-a95f-0a7490cfed90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The weather in San Francisco is currently sunny. Enjoy the sunshine!\n"
     ]
    }
   ],
   "source": [
    "final_result[\"messages\"][-1].pretty_print()"
   ]
  }
 ],
 "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
}