{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b23ced4e-dc29-43be-9f94-0c36bb181b8a",
   "metadata": {},
   "source": [
    "# How to stream LLM tokens (without LangChain LLMs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7044eeb8-4074-4f9c-8a62-962488744557",
   "metadata": {},
   "source": [
    "In this example we will stream tokens from the language model powering an agent. We'll be using OpenAI client library directly, without using LangChain chat models. We will also use a ReAct agent as an example."
   ]
  },
  {
   "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 openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0cf6b41d-7fcb-40b6-9a72-229cdd00a094",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "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 model, tools and graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ba684f1-d46b-42e4-95cf-9685209a5992",
   "metadata": {},
   "source": [
    "### Define a node that will call OpenAI API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d59234f9-173e-469d-a725-c13e0979663e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import AsyncOpenAI\n",
    "from langchain_core.language_models.chat_models import ChatGenerationChunk\n",
    "from langchain_core.messages import AIMessageChunk\n",
    "from langchain_core.runnables.config import (\n",
    "    ensure_config,\n",
    "    get_callback_manager_for_config,\n",
    ")\n",
    "\n",
    "openai_client = AsyncOpenAI()\n",
    "# define tool schema for openai tool calling\n",
    "\n",
    "tool = {\n",
    "    \"type\": \"function\",\n",
    "    \"function\": {\n",
    "        \"name\": \"get_items\",\n",
    "        \"description\": \"Use this tool to look up which items are in the given place.\",\n",
    "        \"parameters\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\"place\": {\"type\": \"string\"}},\n",
    "            \"required\": [\"place\"],\n",
    "        },\n",
    "    },\n",
    "}\n",
    "\n",
    "\n",
    "async def call_model(state, config=None):\n",
    "    config = ensure_config(config | {\"tags\": [\"agent_llm\"]})\n",
    "    callback_manager = get_callback_manager_for_config(config)\n",
    "    messages = state[\"messages\"]\n",
    "\n",
    "    llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n",
    "    response = await openai_client.chat.completions.create(\n",
    "        messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n",
    "    )\n",
    "\n",
    "    response_content = \"\"\n",
    "    role = None\n",
    "\n",
    "    tool_call_id = None\n",
    "    tool_call_function_name = None\n",
    "    tool_call_function_arguments = \"\"\n",
    "    async for chunk in response:\n",
    "        delta = chunk.choices[0].delta\n",
    "        if delta.role is not None:\n",
    "            role = delta.role\n",
    "\n",
    "        if delta.content:\n",
    "            response_content += delta.content\n",
    "            llm_run_manager.on_llm_new_token(delta.content)\n",
    "\n",
    "        if delta.tool_calls:\n",
    "            # note: for simplicity we're only handling a single tool call here\n",
    "            if delta.tool_calls[0].function.name is not None:\n",
    "                tool_call_function_name = delta.tool_calls[0].function.name\n",
    "                tool_call_id = delta.tool_calls[0].id\n",
    "\n",
    "            # note: we're wrapping the tools calls in ChatGenerationChunk so that the events from .astream_events in the graph can render tool calls correctly\n",
    "            tool_call_chunk = ChatGenerationChunk(\n",
    "                message=AIMessageChunk(\n",
    "                    content=\"\",\n",
    "                    additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]},\n",
    "                )\n",
    "            )\n",
    "            llm_run_manager.on_llm_new_token(\"\", chunk=tool_call_chunk)\n",
    "            tool_call_function_arguments += delta.tool_calls[0].function.arguments\n",
    "\n",
    "    if tool_call_function_name is not None:\n",
    "        tool_calls = [\n",
    "            {\n",
    "                \"id\": tool_call_id,\n",
    "                \"function\": {\n",
    "                    \"name\": tool_call_function_name,\n",
    "                    \"arguments\": tool_call_function_arguments,\n",
    "                },\n",
    "                \"type\": \"function\",\n",
    "            }\n",
    "        ]\n",
    "    else:\n",
    "        tool_calls = None\n",
    "\n",
    "    response_message = {\n",
    "        \"role\": role,\n",
    "        \"content\": response_content,\n",
    "        \"tool_calls\": tool_calls,\n",
    "    }\n",
    "    return {\"messages\": [response_message]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a3877e8-8ace-40d5-ad04-cbf21c6f3250",
   "metadata": {},
   "source": [
    "### Define our tools and a tool-calling node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b756ea32",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "\n",
    "async def get_items(place: str) -> str:\n",
    "    \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n",
    "    if \"bed\" in place:  # For under the bed\n",
    "        return \"socks, shoes and dust bunnies\"\n",
    "    if \"shelf\" in place:  # For 'shelf'\n",
    "        return \"books, penciles and pictures\"\n",
    "    else:  # if the agent decides to ask about a different place\n",
    "        return \"cat snacks\"\n",
    "\n",
    "\n",
    "# define mapping to look up functions when running tools\n",
    "function_name_to_function = {\"get_items\": get_items}\n",
    "\n",
    "\n",
    "async def call_tools(state):\n",
    "    messages = state[\"messages\"]\n",
    "\n",
    "    tool_call = messages[-1][\"tool_calls\"][0]\n",
    "    function_name = tool_call[\"function\"][\"name\"]\n",
    "    function_arguments = tool_call[\"function\"][\"arguments\"]\n",
    "    arguments = json.loads(function_arguments)\n",
    "\n",
    "    function_response = await function_name_to_function[function_name](**arguments)\n",
    "    tool_message = {\n",
    "        \"tool_call_id\": tool_call[\"id\"],\n",
    "        \"role\": \"tool\",\n",
    "        \"name\": function_name,\n",
    "        \"content\": function_response,\n",
    "    }\n",
    "    return {\"messages\": [tool_message]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6685898c-9a1c-4803-a492-bd70574ebe38",
   "metadata": {},
   "source": [
    "### Define our graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "228260be-1f9a-4195-80e0-9604f8a5dba6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated, TypedDict, Literal\n",
    "\n",
    "from langgraph.graph import StateGraph, END\n",
    "\n",
    "\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[list, operator.add]\n",
    "\n",
    "\n",
    "def should_continue(state) -> Literal[\"tools\", END]:\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "    if last_message[\"tool_calls\"]:\n",
    "        return \"tools\"\n",
    "    return END\n",
    "\n",
    "\n",
    "workflow = StateGraph(State)\n",
    "workflow.set_entry_point(\"model\")\n",
    "workflow.add_node(\"model\", call_model)  # i.e. our \"agent\"\n",
    "workflow.add_node(\"tools\", call_tools)\n",
    "workflow.add_conditional_edges(\"model\", should_continue)\n",
    "workflow.add_edge(\"tools\", \"model\")\n",
    "graph = workflow.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d046e2ef-f208-4831-ab31-203b2e75a49a",
   "metadata": {},
   "source": [
    "## Stream tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "45c96a79-4147-42e3-89fd-d942b2b49f6c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/vadymbarda/.virtualenvs/langgraph/lib/python3.11/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": [
      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_xUcx3IPa8GREPOpjHVj5k9Wx', 'function': {'arguments': '', 'name': 'get_items'}, 'type': 'function'}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': 'get_items', 'args': '', 'id': 'call_xUcx3IPa8GREPOpjHVj5k9Wx', 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': 'get_items', 'args': '', 'id': 'call_xUcx3IPa8GREPOpjHVj5k9Wx', 'index': 0}]}\n",
      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': '{\"', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [{'name': '', 'args': {}, 'id': None}], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': '{\"', 'id': None, 'index': 0}]}\n",
      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': 'place', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': None, 'args': 'place', 'id': None, 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': 'place', 'id': None, 'index': 0}]}\n",
      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': '\":\"', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': None, 'args': '\":\"', 'id': None, 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': '\":\"', 'id': None, 'index': 0}]}\n",
      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': 'bed', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': None, 'args': 'bed', 'id': None, 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': 'bed', 'id': None, 'index': 0}]}\n",
      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': 'room', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': None, 'args': 'room', 'id': None, 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': 'room', 'id': None, 'index': 0}]}\n",
      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': '\"}', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': None, 'args': '\"}', 'id': None, 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': '\"}', 'id': None, 'index': 0}]}\n",
      "LLM token {'content': 'In', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' the', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' bedroom', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ',', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' you', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' have', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' socks', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ',', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' shoes', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ',', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' and', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' some', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' dust', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' b', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': 'unn', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': 'ies', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': '.', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' Is', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' there', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' anything', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' else', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' you', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' would', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' like', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' to', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': ' know', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
      "LLM token {'content': '?', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n"
     ]
    }
   ],
   "source": [
    "async for event in graph.astream_events(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"\n",
    "):\n",
    "    tags = event.get(\"tags\", [])\n",
    "    if event[\"event\"] == \"on_chat_model_stream\" and \"agent_llm\" in tags:\n",
    "        print(\"LLM token\", event[\"data\"][\"chunk\"].dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adb0f7bc-6e51-478e-bd32-8f72df072d6c",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
