LangChainでChatGPTのWebみたいに返信を徐々に表示する

ChatGPT の API を呼び出すと結果がまとめて返ってくるけど、ChatGPT の WebUI のように１文字１文字徐々に表示することもできます。LangChain を使った場合のコードはこちら。

chat.py

from typing import Any, Dict, List
from langchain.chat_models import ChatOpenAI
from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import (
    LLMResult,
    HumanMessage,
)
import openai
import os


openai.api_key = os.environ.get("OPENAI_API_KEY", "")


class MyCustomCallbackHandler(BaseCallbackHandler):
    def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> None:
        print("START")

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        print("\nEND")

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        print(token, end="", flush=True)

    def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None:
        pass


chat = ChatOpenAI(streaming=True, callbacks=[MyCustomCallbackHandler()], temperature=0)
chat([HumanMessage(content="炭酸ソーダについての詞を書いてください。")])

streaming=True と callbacks=[MyCustomCallbackHandler()] を指定し、 BaseCallbackHandler を継承したクラスを定義してください。