如何从特定节点流式传输LLM令牌

前提条件

本指南假设您熟悉以下内容:

• 流式传输
• 聊天模型

当流式传输LLM令牌时，一个常见的用例是仅从特定节点流式传输它们。为此，您可以使用stream_mode="messages"并根据流式传输元数据中的langgraph_node字段过滤输出：

from langgraph.graph import StateGraph
from langchain_openai import ChatOpenAI

model = ChatOpenAI()

def node_a(state: State):
    model.invoke(...)
    ...

def node_b(state: State):
    model.invoke(...)
    ...

graph = (
    StateGraph(State)
    .add_node(node_a)
    .add_node(node_b)
    ...
    .compile()

for msg, metadata in graph.stream(
    inputs,
    # 高亮下一行
    stream_mode="messages"
):
    # 从'node_a'流式传输
    # 高亮下一行
    if metadata["langgraph_node"] == "node_a":
        print(msg)

从特定LLM调用流式传输

如果您需要将流式传输的LLM令牌过滤到特定的LLM调用，请查看此指南

环境搭建

首先我们需要安装所需的包

%%capture --no-stderr
%pip install --quiet -U langgraph langchain_openai

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")

示例

from typing import TypedDict
from langgraph.graph import START, StateGraph, MessagesState
from langchain_openai import ChatOpenAI

model = ChatOpenAI(model="gpt-4o-mini")


class State(TypedDict):
    topic: str
    joke: str
    poem: str


def write_joke(state: State):
    topic = state["topic"]
    joke_response = model.invoke(
        [{"role": "user", "content": f"Write a joke about {topic}"}]
    )
    return {"joke": joke_response.content}


def write_poem(state: State):
    topic = state["topic"]
    poem_response = model.invoke(
        [{"role": "user", "content": f"Write a short poem about {topic}"}]
    )
    return {"poem": poem_response.content}


graph = (
    StateGraph(State)
    .add_node(write_joke)
    .add_node(write_poem)
    # write both the joke and the poem concurrently
    .add_edge(START, "write_joke")
    .add_edge(START, "write_poem")
    .compile()
)

API Reference: START | StateGraph | ChatOpenAI

for msg, metadata in graph.stream(
    {"topic": "cats"},
    stream_mode="messages",
):
    if msg.content and metadata["langgraph_node"] == "write_poem":
        print(msg.content, end="|", flush=True)

In| shadows| soft|,| they| quietly| creep|,|  
|Wh|isk|ered| wonders|,| in| dreams| they| leap|.|  
|With| eyes| like| lantern|s|,| bright| and| wide|,|  
|Myst|eries| linger| where| they| reside|.|  

|P|aws| that| pat|ter| on| silent| floors|,|  
|Cur|led| in| sun|be|ams|,| they| seek| out| more|.|  
|A| flick| of| a| tail|,| a| leap|,| a| p|ounce|,|  
|In| their| playful| world|,| we| can't| help| but| bounce|.|  

|Guard|ians| of| secrets|,| with| gentle| grace|,|  
|Each| little| me|ow|,| a| warm| embrace|.|  
|Oh|,| the| joy| that| they| bring|,| so| pure| and| true|,|  
|In| the| heart| of| a| cat|,| there's| magic| anew|.|  |

Comments