{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "562ddb82",
   "metadata": {},
   "source": [
    "# Streaming Tokens\n",
    "\n",
    "In this example, we will stream tokens from the language model powering an\n",
    "agent. We will use a ReAct agent as an example. The tl;dr is to use\n",
    "[streamEvents](https://js.langchain.com/v0.2/docs/how_to/chat_streaming/#stream-events)\n",
    "([API Ref](https://api.js.langchain.com/classes/langchain_core_runnables.Runnable.html#streamEvents)).\n",
    "\n",
    "<div class=\"admonition info\">\n",
    "    <p class=\"admonition-title\">Note</p>\n",
    "    <p>\n",
    "      If you are using a version of `@langchain/core` < 0.2.3, when calling chat models or LLMs you need to call `await model.stream()` within your nodes to get token-by-token streaming events, and aggregate final outputs if needed to update the graph state. In later versions of `@langchain/core`, this occurs automatically, and you can call `await model.invoke()`.\n",
    "\n",
    "      For more on how to upgrade `@langchain/core`, check out [the instructions here](https://js.langchain.com/v0.2/docs/how_to/installation/#installing-integration-packages).\n",
    "    </p>\n",
    "\n",
    "</div>\n",
    "\n",
    "This how-to guide closely follows the others in this directory, showing how to\n",
    "incorporate the functionality into a prototypical agent in LangGraph.\n",
    "\n",
    "This works for\n",
    "[StateGraph](https://langchain-ai.github.io/langgraphjs/reference/classes/index.StateGraph.html)\n",
    "and all its subclasses, such as\n",
    "[MessageGraph](https://langchain-ai.github.io/langgraphjs/reference/classes/index.MessageGraph.html).\n",
    "\n",
    "<div class=\"admonition info\">\n",
    "    <p class=\"admonition-title\">Streaming Support</p>\n",
    "    <p>\n",
    "        Token streaming is supported by many, but not all chat models. Check to see if your LLM integration supports token streaming <a href=\"https://js.langchain.com/v0.2/docs/integrations/chat/\">here (doc)</a>. Note that some integrations may support _general_ token streaming but lack support for streaming tool calls.\n",
    "    </p>\n",
    "</div>\n",
    "\n",
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Note</p>\n",
    "    <p>\n",
    "        In this how-to, we will create our agent from scratch to be transparent (but verbose). You can accomplish similar functionality using the <code>createReactAgent(model, tools=tool)</code> (<a href=\"https://langchain-ai.github.io/langgraphjs/reference/functions/prebuilt.createReactAgent.html\">API doc</a>) constructor. This may be more appropriate if you are used to LangChain's <a href=\"https://js.langchain.com/v0.2/docs/how_to/agent_executor\">AgentExecutor</a> class.\n",
    "    </p>\n",
    "</div>\n",
    "\n",
    "## Setup\n",
    "\n",
    "This guide will use OpenAI's GPT-4o model. We will optionally set our API key\n",
    "for [LangSmith tracing](https://smith.langchain.com/), which will give us\n",
    "best-in-class observability.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8e76833b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "// process.env.OPENAI_API_KEY = \"sk_...\";\n",
    "\n",
    "// Optional, add tracing in LangSmith\n",
    "// process.env.LANGCHAIN_API_KEY = \"ls__...\";\n",
    "process.env.LANGCHAIN_CALLBACKS_BACKGROUND = \"true\";\n",
    "process.env.LANGCHAIN_TRACING_V2 = \"true\";\n",
    "process.env.LANGCHAIN_PROJECT = \"Stream Tokens: LangGraphJS\";"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab95dc97",
   "metadata": {},
   "source": [
    "## Define the state\n",
    "\n",
    "The state is the interface for all of the nodes in our graph.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1648124b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import { BaseMessage } from \"@langchain/core/messages\";\n",
    "import { StateGraphArgs } from \"@langchain/langgraph\";\n",
    "\n",
    "interface IState {\n",
    "  messages: BaseMessage[];\n",
    "}\n",
    "\n",
    "// This defines the agent state\n",
    "const graphState: StateGraphArgs<IState>[\"channels\"] = {\n",
    "  messages: {\n",
    "    value: (x: BaseMessage[], y: BaseMessage[]) => x.concat(y),\n",
    "    default: () => [],\n",
    "  },\n",
    "};"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da50fbd8",
   "metadata": {},
   "source": [
    "## Set up the tools\n",
    "\n",
    "We will first define the tools we want to use. For this simple example, we will\n",
    "use create a placeholder search engine. However, it is really easy to create\n",
    "your own tools - see documentation\n",
    "[here](https://js.langchain.com/v0.2/docs/how_to/custom_tools) on how to do\n",
    "that.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a8f1ae1c",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[32m\"Cold, with a low of 3℃\"\u001b[39m"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import { DynamicStructuredTool } from \"@langchain/core/tools\";\n",
    "import { z } from \"zod\";\n",
    "\n",
    "const searchTool = new DynamicStructuredTool({\n",
    "  name: \"search\",\n",
    "  description:\n",
    "    \"Use to surf the web, fetch current information, check the weather, and retrieve other information.\",\n",
    "  schema: z.object({\n",
    "    query: z.string().describe(\"The query to use in your search.\"),\n",
    "  }),\n",
    "  func: async ({ query: _query }: { query: string }) => {\n",
    "    // This is a placeholder for the actual implementation\n",
    "    return \"Cold, with a low of 3℃\";\n",
    "  },\n",
    "});\n",
    "\n",
    "await searchTool.invoke({ query: \"What's the weather like?\" });\n",
    "\n",
    "const tools = [searchTool];"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19b27cb3",
   "metadata": {},
   "source": [
    "We can now wrap these tools in a simple\n",
    "[ToolNode](https://langchain-ai.github.io/langgraphjs/reference/classes/prebuilt.ToolNode.html).\n",
    "This object will actually run the tools (functions) whenever they are invoked by\n",
    "our LLM.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f02278b1",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import { ToolNode } from \"@langchain/langgraph/prebuilt\";\n",
    "\n",
    "const toolNode = new ToolNode<{ messages: BaseMessage[] }>(tools);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd55ee5a",
   "metadata": {},
   "source": [
    "## Set up the model\n",
    "\n",
    "Now we will load the\n",
    "[chat model](https://js.langchain.com/v0.2/docs/concepts/#chat-models).\n",
    "\n",
    "1. It should work with messages. We will represent all agent state in the form\n",
    "   of messages, so it needs to be able to work well with them.\n",
    "2. It should work with\n",
    "   [tool calling](https://js.langchain.com/v0.2/docs/how_to/tool_calling/#passing-tools-to-llms),\n",
    "   meaning it can return function arguments in its response.\n",
    "\n",
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Note</p>\n",
    "    <p>\n",
    "        These model requirements are not general requirements for using LangGraph - they are just requirements for this one example.\n",
    "    </p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9c7210e7",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import { ChatOpenAI } from \"@langchain/openai\";\n",
    "\n",
    "const model = new ChatOpenAI({ model: \"gpt-4o\" });"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73e59248",
   "metadata": {},
   "source": [
    "After we've done this, we should make sure the model knows that it has these\n",
    "tools available to call. We can do this by calling\n",
    "[bindTools](https://v01.api.js.langchain.com/classes/langchain_core_language_models_chat_models.BaseChatModel.html#bindTools)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b4ff23ee",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "const boundModel = model.bindTools(tools);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbe67356",
   "metadata": {},
   "source": [
    "## Define the graph\n",
    "\n",
    "We can now put it all together. Time travel requires a checkpointer to save the\n",
    "state - otherwise you wouldn't have anything go `get` or `update`. We will use\n",
    "the\n",
    "[MemorySaver](https://langchain-ai.github.io/langgraphjs/reference/classes/index.MemorySaver.html),\n",
    "which \"saves\" checkpoints in-memory.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0ba603bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import { END, START, StateGraph } from \"@langchain/langgraph\";\n",
    "import { AIMessage } from \"@langchain/core/messages\";\n",
    "import { RunnableConfig } from \"@langchain/core/runnables\";\n",
    "\n",
    "const routeMessage = (state: IState) => {\n",
    "  const { messages } = state;\n",
    "  const lastMessage = messages[messages.length - 1] as AIMessage;\n",
    "  // If no tools are called, we can finish (respond to the user)\n",
    "  if (!lastMessage?.tool_calls?.length) {\n",
    "    return END;\n",
    "  }\n",
    "  // Otherwise if there is, we continue and call the tools\n",
    "  return \"tools\";\n",
    "};\n",
    "\n",
    "const callModel = async (\n",
    "  state: IState,\n",
    "  config?: RunnableConfig,\n",
    ") => {\n",
    "  // For versions of @langchain/core < 0.2.3, you must call `.stream()`\n",
    "  // and aggregate the message from chunks instead of calling `.invoke()`.\n",
    "  const { messages } = state;\n",
    "  const responseMessage = await boundModel.invoke(messages, config);\n",
    "  return { messages: [responseMessage] };\n",
    "};\n",
    "\n",
    "const workflow = new StateGraph<IState>({\n",
    "  channels: graphState,\n",
    "})\n",
    "  .addNode(\"agent\", callModel)\n",
    "  .addNode(\"tools\", toolNode)\n",
    "  .addEdge(START, \"agent\")\n",
    "  .addConditionalEdges(\"agent\", routeMessage)\n",
    "  .addEdge(\"tools\", \"agent\");\n",
    "\n",
    "const graph = workflow.compile();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1ab3ad3",
   "metadata": {},
   "source": [
    "## Call streamEvents\n",
    "\n",
    "We can now interact with the agent. Between interactions you can get and update\n",
    "state.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cbcf7c39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Hello\n",
      " Jo\n",
      "!\n",
      " How\n",
      " can\n",
      " I\n",
      " assist\n",
      " you\n",
      " today\n",
      "?\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import { ChatGenerationChunk } from \"@langchain/core/outputs\";\n",
    "\n",
    "let config = { configurable: { thread_id: \"conversation-num-1\" } };\n",
    "let inputs = { messages: [[\"user\", \"Hi I'm Jo.\"]] };\n",
    "\n",
    "for await (\n",
    "  const event of await graph.streamEvents(inputs, {\n",
    "    ...config,\n",
    "    streamMode: \"values\",\n",
    "    version: \"v1\",\n",
    "  })\n",
    ") {\n",
    "  if (event.event === \"on_llm_stream\") {\n",
    "    let chunk: ChatGenerationChunk = event.data?.chunk;\n",
    "    let msg = chunk.message as AIMessageChunk;\n",
    "    if (msg.tool_call_chunks && msg.tool_call_chunks.length > 0) {\n",
    "      console.log(msg.tool_call_chunks);\n",
    "    } else {\n",
    "      console.log(msg.content);\n",
    "    }\n",
    "  }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "055aacad",
   "metadata": {},
   "source": [
    "## How to stream tool calls\n",
    "\n",
    "Many providers support token-level streaming of tool invocations. To get the\n",
    "partially populated results, you can access the message chunks'\n",
    "`tool_call_chunks` property."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c704d23c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[\n",
      "  {\n",
      "    name: \"search\",\n",
      "    args: \"\",\n",
      "    id: \"call_zesYhE7nGptBKnmj5yD78yrB\",\n",
      "    index: 0\n",
      "  }\n",
      "]\n",
      "[ { name: undefined, args: '{\"', id: undefined, index: 0 } ]\n",
      "[ { name: undefined, args: \"query\", id: undefined, index: 0 } ]\n",
      "[ { name: undefined, args: '\":\"', id: undefined, index: 0 } ]\n",
      "[ { name: undefined, args: \"current\", id: undefined, index: 0 } ]\n",
      "[ { name: undefined, args: \" weather\", id: undefined, index: 0 } ]\n",
      "[ { name: undefined, args: '\"}', id: undefined, index: 0 } ]\n",
      "\n",
      "\n",
      "The\n",
      " weather\n",
      " today\n",
      " is\n",
      " quite\n",
      " cold\n",
      ",\n",
      " with\n",
      " temperatures\n",
      " dropping\n",
      " to\n",
      " a\n",
      " low\n",
      " of\n",
      " \n",
      "3\n",
      "℃\n",
      ".\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for await (\n",
    "  const event of await graph.streamEvents(\n",
    "    { messages: [[\"user\", \"What's the weather like today?\"]] },\n",
    "    {\n",
    "      ...config,\n",
    "      streamMode: \"values\",\n",
    "      version: \"v1\",\n",
    "    },\n",
    "  )\n",
    ") {\n",
    "  if (event.event === \"on_llm_stream\") {\n",
    "    let chunk: ChatGenerationChunk = event.data?.chunk;\n",
    "    let msg = chunk.message as AIMessageChunk;\n",
    "    if (msg.tool_call_chunks && msg.tool_call_chunks.length > 0) {\n",
    "      console.log(msg.tool_call_chunks);\n",
    "    } else {\n",
    "      console.log(msg.content);\n",
    "    }\n",
    "  }\n",
    "}"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TypeScript",
   "language": "typescript",
   "name": "tslab"
  },
  "language_info": {
   "file_extension": ".ts",
   "mimetype": "text/x.typescript",
   "name": "typescript",
   "nb_converter": "script",
   "pygments_lexer": "typescript",
   "version": "5.3.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}