{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "15c4bd28",
   "metadata": {},
   "source": [
    "# How to stream from the final node"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "964686a6-8fed-4360-84d2-958c48186008",
   "metadata": {},
   "source": [
    "A common use case is streaming from an agent is to stream LLM tokens from inside the final node. This guide demonstrates how you can do this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c04a3f8e-0bc9-430b-85db-3edfa026d2cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n%pip install -U langgraph langchain-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c87e4a47-4099-4d1a-907c-a99fa857165a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "OPENAI_API_KEY:  ········\n"
     ]
    }
   ],
   "source": [
    "import getpass\nimport os\n\n\ndef _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": "17f994ca-28e7-4379-a1c9-8c1682773b5f",
   "metadata": {},
   "source": [
    "## Define model and tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1d51c35c-dbf2-4c01-932d-c5d308ea37d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\nfrom langchain_community.tools.tavily_search import TavilySearchResults\nfrom langchain_core.runnables import ConfigurableField\nfrom langchain_core.tools import tool\nfrom langchain_openai import ChatOpenAI\nfrom langgraph.prebuilt import create_react_agent\nfrom langgraph.prebuilt import ToolNode\n\n\n@tool\ndef 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\ntools = [get_weather]\nmodel = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)\nfinal_model = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)\n\nmodel = model.bind_tools(tools)\n# NOTE: this is where we're adding a tag that we'll be using later to filter the outputs of the final node\nfinal_model = final_model.with_config(tags=[\"final_node\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0af37212-e592-484d-9194-35d53fa79678",
   "metadata": {},
   "outputs": [],
   "source": [
    "tool_node = ToolNode(tools=tools)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ac9d4f5b-655a-48f3-b514-a4a0815714a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import TypedDict, Annotated\n\nfrom langgraph.graph import END, StateGraph, START\nfrom langgraph.graph.message import MessagesState\nfrom langchain_core.messages import BaseMessage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9acef997-5dd6-4108-baf1-c4d6be3e4999",
   "metadata": {},
   "source": [
    "## Define graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3948c6b8-0317-4001-b699-32b25306a023",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import SystemMessage, HumanMessage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2efe9fb4-c6c2-4171-becd-d45bbf899209",
   "metadata": {},
   "outputs": [],
   "source": [
    "def should_continue(state: MessagesState) -> Literal[\"tools\", \"final\"]:\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "    # If the LLM makes a tool call, then we route to the \"tools\" node\n",
    "    if last_message.tool_calls:\n",
    "        return \"tools\"\n",
    "    # Otherwise, we stop (reply to the user)\n",
    "    return \"final\"\n",
    "\n",
    "\n",
    "def call_model(state: MessagesState):\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",
    "def call_final_model(state: MessagesState):\n",
    "    messages = state[\"messages\"]\n",
    "    last_ai_message = messages[-1]\n",
    "    response = final_model.invoke(\n",
    "        [\n",
    "            SystemMessage(\"Rewrite this in the voice of Al Roker\"),\n",
    "            HumanMessage(last_ai_message.content),\n",
    "        ]\n",
    "    )\n",
    "    # overwrite the last AI message from the agent\n",
    "    response.id = last_ai_message.id\n",
    "    return {\"messages\": [response]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b1a9a981-8629-4d25-a0e1-d666c3968b30",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = StateGraph(MessagesState)\n",
    "\n",
    "workflow.add_node(\"agent\", call_model)\n",
    "workflow.add_node(\"tools\", tool_node)\n",
    "# add a separate final node\n",
    "workflow.add_node(\"final\", call_final_model)\n",
    "\n",
    "workflow.add_edge(START, \"agent\")\n",
    "workflow.add_conditional_edges(\n",
    "    \"agent\",\n",
    "    should_continue,\n",
    ")\n",
    "\n",
    "workflow.add_edge(\"tools\", \"agent\")\n",
    "workflow.add_edge(\"final\", END)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a7b0251f-dcee-49d6-8133-af50d4a55e22",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f8b77e74-17e9-4fee-a164-4637013b55ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "98adcef2-1dec-4503-99cc-48613cdb7a85",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 12768 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Image(app.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "521adaef-dd2f-46d6-8f6a-5cc1d6e0aefc",
   "metadata": {},
   "source": [
    "## Stream outputs from the final node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a37c3a5f-5a43-46db-940e-c583df776520",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/vadymbarda/.virtualenvs/langgraph/lib/python3.12/site-packages/langchain_core/_api/beta_decorator.py:87: LangChainBetaWarning: This API is in beta and may change in the future.\n",
      "  warn_beta(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Well| folks|,| looks| like| we|'ve| got| some| cloudy| skies| in| the| Big| Apple| today|.| So| grab| your| umbrella| just| in| case|,| and| don|'t| let| those| clouds| rain| on| your| parade|!|"
     ]
    }
   ],
   "source": [
    "inputs = {\"messages\": [(\"human\", \"what's the weather in nyc?\")]}\nasync for event in app.astream_events(inputs, version=\"v2\"):\n    kind = event[\"event\"]\n    tags = event.get(\"tags\", [])\n    if kind == \"on_chat_model_stream\" and \"final_node\" in tags:\n        data = event[\"data\"]\n        if data[\"chunk\"].content:\n            # Empty content in the context of OpenAI or Anthropic usually means\n            # that the model is asking for a tool to be invoked.\n            # So we only print non-empty content\n            print(data[\"chunk\"].content, end=\"|\")"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
