langgraph/examples/streaming-events-from-within-tools.ipynb

How to stream events from within a tool

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.

Setup

In [1]:

%%capture --no-stderr
%pip install -U langgraph langchain-openai

In [2]:

import getpass
import os


def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")


_set_env("OPENAI_API_KEY")

OPENAI_API_KEY:  ········

Define graph and tools

We'll use a prebuilt ReAct agent for this guide

In [3]:

from langchain_core.callbacks import Callbacks
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool

from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI

ASYNC IN PYTHON<=3.10

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.

This is a common reason why you may fail to see events being emitted from custom runnables or tools.

In [4]:

@tool
async def get_items(
    place: str, callbacks: Callbacks
) -> str:  # <--- Accept callbacks (Python <= 3.10)
    """Use this tool to look up which items are in the given place."""
    template = ChatPromptTemplate.from_messages(
        [
            (
                "human",
                "Can you tell me what kind of items i might find in the following place: '{place}'. "
                "List at least 3 such items separating them by a comma. And include a brief description of each item..",
            )
        ]
    )
    chain = template | llm.with_config(
        {
            "run_name": "Get Items LLM",
            "tags": ["tool_llm"],
            "callbacks": callbacks,  # <-- Propagate callbacks (Python <= 3.10)
        }
    )
    chunks = [chunk async for chunk in chain.astream({"place": place})]
    return "".join(chunk.content for chunk in chunks)

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

In [5]:

llm = ChatOpenAI(model_name="gpt-3.5-turbo")
tools = [get_items]
agent = create_react_agent(llm, tools=tools)

Stream events from the graph

In [6]:

async for event in agent.astream_events(
    {"messages": [("human", "what items are on the shelf?")]}, version="v2"
):
    tags = event.get("tags", [])
    if event["event"] == "on_chat_model_stream" and "tool_llm" in tags:
        print(event["data"]["chunk"].content, end="", flush=True)

/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.
  warn_beta(

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.

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.

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.

Let's inspect the last event to get the final list of messages from the agent

In [7]:

final_messages = event["data"]["output"]["messages"]

In [8]:

for message in final_messages:
    message.pretty_print()

================================ Human Message =================================

what items are on the shelf?
================================== Ai Message ==================================
Tool Calls:
  get_items (call_5CAMZ3asoLsZm9ocMbCOWxYQ)
 Call ID: call_5CAMZ3asoLsZm9ocMbCOWxYQ
  Args:
    place: shelf
================================= Tool Message =================================
Name: get_items

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.

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.

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.
================================== Ai Message ==================================

The items on the shelf are:
1. Books
2. Picture frames
3. Candles

You can see that the content of the ToolMessage is the same as the output we streamed above