サブシステム間のAPIのやり取りをシーケンス図にまとめる

Streaming
For LLMs, streaming has become an increasingly popular feature. The idea is to rapidly return tokens as an LLM is generating them, rather than waiting for a full response to be created before returning anything.

Streaming is actually very easy to implement for simple use-cases, but it can get complicated when we start including things like Agents which have their own logic running which can block our attempts at streaming. Fortunately, we can make it work — it just requires a little extra effort.

We'll start easy by implementing streaming to the terminal for LLMs, but by the end of the notebook we'll be handling the more complex task of streaming via FastAPI for Agents.

First, let's install all of the libraries we'll be using.

!pip install -qU \
    openai==0.28.0 \
    langchain==0.0.301 \
    fastapi==0.103.1 \
    "uvicorn[standard]"==0.23.2
LLM Streaming to Stdout
The simplest form of streaming is to simply "print" the tokens as they're generated. To set this up we need to initialize an LLM (one that supports streaming, not all do) with two specific parameters:

streaming=True, to enable streaming
callbacks=[SomeCallBackHere()], where we pass a LangChain callback class (or list containing multiple).
The streaming parameter is self-explanatory. The callbacks parameter and callback classes less so — essentially they act as little bits of code that do something as each token from our LLM is generated. As mentioned, the simplest form of streaming is to print the tokens as they're being generated, like with the StreamingStdOutCallbackHandler.

import os
from langchain.chat_models import ChatOpenAI
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY") or "YOUR_API_KEY"

llm = ChatOpenAI(
    openai_api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
    model_name="gpt-3.5-turbo",
    streaming=True,  # ! important
    callbacks=[StreamingStdOutCallbackHandler()]  # ! important
)
Now if we run the LLM we'll see the response being streamed.

from langchain.schema import HumanMessage

# create messages to be passed to chat LLM
messages = [HumanMessage(content="tell me a long story")]

llm(messages)
That was surprisingly easy, but things begin to get much more complicated as soon as we begin using agents. Let's first initialize an agent.

from langchain.memory import ConversationBufferWindowMemory
from langchain.agents import load_tools, AgentType, initialize_agent

# initialize conversational memory
memory = ConversationBufferWindowMemory(
    memory_key="chat_history",
    k=5,
    return_messages=True,
    output_key="output"
)

# create a single tool to see how it impacts streaming
tools = load_tools(["llm-math"], llm=llm)

# initialize the agent
agent = initialize_agent(
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
    tools=tools,
    llm=llm,
    memory=memory,
    verbose=True,
    max_iterations=3,
    early_stopping_method="generate",
    return_intermediate_steps=False
)
We already added our StreamingStdOutCallbackHandler to the agent as we initialized the agent with the same llm as we created with that callback. So let's see what we get when running the agent.

prompt = "Hello, how are you?"

agent(prompt)
Not bad, but we do now have the issue of streaming the entire output from the LLM. Because we're using an agent, the LLM is instructed to output the JSON format we can see here so that the agent logic can handle tool usage, multiple "thinking" steps, and so on. For example, if we ask a math question we'll see this:

agent("what is the square root of 71?")
It's interesting to see during development but we'll want to clean this streaming up a little in any actual use-case. For that we can go with two approaches — either we build a custom callback handler, or use a purpose built callback handler from LangChain (as usual, LangChain has something for everything). Let's first try LangChain's purpose-built FinalStreamingStdOutCallbackHandler.

We will overwrite the existing callbacks attribute found here:

agent.agent.llm_chain.llm
With the new callback handler:

from langchain.callbacks.streaming_stdout_final_only import (
    FinalStreamingStdOutCallbackHandler,
)

agent.agent.llm_chain.llm.callbacks = [
    FinalStreamingStdOutCallbackHandler(
        answer_prefix_tokens=["Final", "Answer"]
    )
]
Let's try it:

agent("what is the square root of 71?")
Not quite there, we should really clean up the answer_prefix_tokens argument but it is hard to get right. It's generally easier to use a custom callback handler like so:

import sys

class CallbackHandler(StreamingStdOutCallbackHandler):
    def __init__(self):
        self.content: str = ""
        self.final_answer: bool = False

    def on_llm_new_token(self, token: str, **kwargs: any) -> None:
        self.content += token
        if "Final Answer" in self.content:
            # now we're in the final answer section, but don't print yet
            self.final_answer = True
            self.content = ""
        if self.final_answer:
            if '"action_input": "' in self.content:
                if token not in ["}"]:
                    sys.stdout.write(token)  # equal to `print(token, end="")`
                    sys.stdout.flush()

agent.agent.llm_chain.llm.callbacks = [CallbackHandler()]
Let's try again:

agent("what is the square root of 71?")
agent.agent.llm_chain.llm
It isn't perfect, but this is getting better. Now, in most scenarios we're unlikely to simply be printing output to a terminal or notebook. When we want to do something more complex like stream this data through another API, we need to do things differently.

Using FastAPI with Agents
In most cases we'll be placing our LLMs, Agents, etc behind something like an API. Let's add that into the mix and see how we can implement streaming for agents with FastAPI.

First, we'll create a simple main.py script to contain our FastAPI logic. You can find it in the same GitHub repo location as this notebook (here's a link).

To run the API, navigate to the directory and run uvicorn main:app --reload. Once complete, you can confirm it is running by looking for the 🤙 status in the next cell output:

import requests

res = requests.get("http://localhost:8000/health")
res.json()
{'status': '🤙'}
res = requests.get("http://localhost:8000/chat",
    json={"text": "hello there!"}
)
res
<Response [200]>
res.json()
{'input': 'hello there!',
 'chat_history': [],
 'output': 'Hello! How can I assist you today?'}
Unlike with our StdOut streaming, we now need to send our tokens to a generator function that feeds those tokens to FastAPI via a StreamingResponse object. To handle this we need to use async code, otherwise our generator will not begin emitting anything until after generation is already complete.

The Queue is accessed by our callback handler, as as each token is generated, it puts the token into the queue. Our generator function asyncronously checks for new tokens being added to the queue. As soon as the generator sees a token has been added, it gets the token and yields it to our StreamingResponse.

To see it in action, we'll define a stream requests function called get_stream:

def get_stream(query: str):
    s = requests.Session()
    with s.get(
        "http://localhost:8000/chat",
        stream=True,
        json={"text": query}
    ) as r:
        for line in r.iter_content():
            print(line.decode("utf-8"), end="")
get_stream("hi there!")
 "Hello! How can I assist you today?"

このチュートリアルのシステム構成は？

処理フローをmermaidで

participantは、ユーザー、フロントエンド、バックエンド、LLMとしてください。

participantは、ユーザー、フロントエンド、バックエンド、LLMとしたまま
処理内容として使用しているAPIがわかるように記載してください、

ユーザーからFastAPIへのリクエスト送信: ユーザーがFastAPIにリクエストを送信します。
FastAPIからエージェントへのリクエスト受信: FastAPIがリクエストを受け取り、エージェントに渡します。
エージェントからLLMへの設定: エージェントがLLMにストリーミングパラメータ（streaming=Trueおよびcallbacks=[...]）を設定します。
LLMからコールバックハンドラへのトークン生成のリアルタイム監視: LLMがトークンをリアルタイムで生成し、コールバックハンドラに通知します。
コールバックハンドラによるトークンのリアルタイム処理: コールバックハンドラが生成された各トークンをリアルタイムで処理します。
コールバックハンドラからFastAPIへのストリーミング応答の構築: コールバックハンドラがストリーミング応答を構築し、FastAPIに送信します。
FastAPIからユーザーへのストリーミング応答送信: FastAPIが構築されたストリーミング応答をユーザーに送信します。