{ "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 ensure_config, get_callback_manager_for_config\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\": {\n", " \"place\": {\n", " \"type\": \"string\"\n", " }\n", " },\n", " \"required\": [\n", " \"place\"\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,\n", " model=\"gpt-3.5-turbo\",\n", " tools=[tool],\n", " 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(content=\"\", additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]})\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\": {\"name\": tool_call_function_name, \"arguments\": tool_call_function_arguments},\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", "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", "# define mapping to look up functions when running tools\n", "function_name_to_function = {\n", " \"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 {\n", " \"messages\": [tool_message]\n", " }" ] }, { "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", "class State(TypedDict):\n", " messages: Annotated[list, operator.add]\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", "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({\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"):\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 }