{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b23ced4e-dc29-43be-9f94-0c36bb181b8a",
   "metadata": {},
   "source": [
    "# How to stream events from within a tool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7044eeb8-4074-4f9c-8a62-962488744557",
   "metadata": {},
   "source": [
    "If your LangGraph graph needs to use tools that call LLMs (or any other LangChain `Runnable` objects -- other graphs, LCEL chains, retrievers, etc.), you might want to stream events from the underlying `Runnable`. This guide shows how you can do that."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a37f60af-43ea-4aa6-847a-df8cc47065f5",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "47f79af8-58d8-4a48-8d9a-88823d88701f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langgraph langchain-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0cf6b41d-7fcb-40b6-9a72-229cdd00a094",
   "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": "markdown",
   "id": "e3d02ebb-c2e1-4ef7-b187-810d55139317",
   "metadata": {},
   "source": [
    "## Define graph and tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d74a1760-a063-4d05-8c6f-9d16bc31fa82",
   "metadata": {},
   "source": [
    "We'll use a prebuilt ReAct agent for this guide"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "083757a9-26d7-481e-8f3d-3e34bcba154b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.callbacks import Callbacks\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.tools import tool\n",
    "\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "from langchain_openai import ChatOpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9378fd4a-69e4-49e2-b34c-a98a0505ea35",
   "metadata": {},
   "source": [
    "<div class=\"admonition warning\">\n",
    "    <p class=\"admonition-title\">ASYNC IN PYTHON<=3.10</p>\n",
    "    <p>\n",
    "Any Langchain RunnableLambda, a RunnableGenerator, or Tool that invokes other runnables and is running async in python<=3.10, will have to propagate callbacks to child objects manually. This is because LangChain cannot automatically propagate callbacks to child objects in this case.\n",
    "    \n",
    "This is a common reason why you may fail to see events being emitted from custom runnables or tools.\n",
    "    </p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2cb38dd9-74d8-456d-9e39-4655f2bf3f37",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool\n",
    "async def get_items(\n",
    "    place: str, callbacks: Callbacks\n",
    ") -> str:  # <--- Accept callbacks (Python <= 3.10)\n",
    "    \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n",
    "    template = ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            (\n",
    "                \"human\",\n",
    "                \"Can you tell me what kind of items i might find in the following place: '{place}'. \"\n",
    "                \"List at least 3 such items separating them by a comma. And include a brief description of each item..\",\n",
    "            )\n",
    "        ]\n",
    "    )\n",
    "    chain = template | llm.with_config(\n",
    "        {\n",
    "            \"run_name\": \"Get Items LLM\",\n",
    "            \"tags\": [\"tool_llm\"],\n",
    "            \"callbacks\": callbacks,  # <-- Propagate callbacks (Python <= 3.10)\n",
    "        }\n",
    "    )\n",
    "    chunks = [chunk async for chunk in chain.astream({\"place\": place})]\n",
    "    return \"\".join(chunk.content for chunk in chunks)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17279b8a-049d-483d-af63-8a875098e71f",
   "metadata": {},
   "source": [
    "We're adding a custom tag (`tool_llm`) to our LLM runnable within the tool. This will allow us to filter events that we'll stream from the compiled graph (`agent`) Runnable below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7254310e-7016-45f7-9795-6d52a1160086",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\")\n",
    "tools = [get_items]\n",
    "agent = create_react_agent(llm, tools=tools)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7d88960-a66b-4699-adee-c12d40b4318a",
   "metadata": {},
   "source": [
    "## Stream events from the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "31fe94ab-80de-4729-843e-5a0fe1bb52c0",
   "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": [
      "1. Books - A collection of written or printed works bound together with covers. They can be fiction or non-fiction and come in various genres.\n",
      "\n",
      "2. Picture frames - A decorative border for a photograph or artwork, typically made of wood, metal, or plastic. Picture frames are used to display and protect a picture or painting.\n",
      "\n",
      "3. Candles - A cylinder of wax with a central wick that is lit to produce light or fragrance. Candles are often used for decoration, ambiance, or religious ceremonies."
     ]
    }
   ],
   "source": [
    "async for event in agent.astream_events(\n",
    "    {\"messages\": [(\"human\", \"what items are on the shelf?\")]}, version=\"v2\"\n",
    "):\n",
    "    tags = event.get(\"tags\", [])\n",
    "    if event[\"event\"] == \"on_chat_model_stream\" and \"tool_llm\" in tags:\n",
    "        print(event[\"data\"][\"chunk\"].content, end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebd8902e-935b-4724-8b5d-551b7674fd34",
   "metadata": {},
   "source": [
    "Let's inspect the last event to get the final list of messages from the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ca382c1f-b1c7-4c8a-bd9b-7a873b891b3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_messages = event[\"data\"][\"output\"][\"messages\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3fa7d768-5a84-475a-950e-fd351a44841b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what items are on the shelf?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  get_items (call_5CAMZ3asoLsZm9ocMbCOWxYQ)\n",
      " Call ID: call_5CAMZ3asoLsZm9ocMbCOWxYQ\n",
      "  Args:\n",
      "    place: shelf\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: get_items\n",
      "\n",
      "1. Books - A collection of written or printed works bound together with covers. They can be fiction or non-fiction and come in various genres.\n",
      "\n",
      "2. Picture frames - A decorative border for a photograph or artwork, typically made of wood, metal, or plastic. Picture frames are used to display and protect a picture or painting.\n",
      "\n",
      "3. Candles - A cylinder of wax with a central wick that is lit to produce light or fragrance. Candles are often used for decoration, ambiance, or religious ceremonies.\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The items on the shelf are:\n",
      "1. Books\n",
      "2. Picture frames\n",
      "3. Candles\n"
     ]
    }
   ],
   "source": [
    "for message in final_messages:\n",
    "    message.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7f9457c-5665-4cd5-9a99-d54c84270616",
   "metadata": {},
   "source": [
    "You can see that the content of the `ToolMessage` is the same as the output we streamed above"
   ]
  }
 ],
 "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
}