LangGraphでトークンストリーミング出力 on Databricks

# 導入

[LangGraph](https://python.langchain.com/v0.1/docs/langgraph/)には[Streaming LLM Tokens](https://python.langchain.com/v0.1/docs/langgraph/#streaming-llm-tokens)といういわゆるLLMの出力をストリーミング出力する機能があるのですが、[こちら](https://github.com/langchain-ai/langgraph/issues/78)のIssueで報告されているように、環境によっては適切にストリーミングできない不具合が発生していました。

Databricks上でLangGraphを利用するにあたって、以前は同種の不具合と直面していたのですが、DBRを14.3 -> 15.1に変更してためなのか、LangGraphのバージョン更新のためかなのか、正常に動作するようになっていたため、その内容を記載してみます。

検証環境はDatabricks on AWS、DBRは15.1ML、インスタンスタイプはg5.xlarge。
検証時のLangChainのバージョンは```0.1.20```、LangGraphのバージョンは```0.0.48```です。

# Step1. パッケージインストール

[こちらの記事](https://qiita.com/isanakamishiro2/items/12a92fd6a5cec1ff4900)で作成した```requirements.txt```を使ってpytorch2.3.0およびExLlamaV2の利用に必要なパッケージをインストール。
合わせて、LangChainとLangGraphの最新版をインストールします。

:::note info
今回(も)、推論エンジンにExLlamaV2を使っていますが、transformersやvLLMでも実行できると思います（未検証）
:::


```python
# torch 2.3.0
%pip install -r /Volumes/training/llm/packages/torch/requirements.txt

%pip install -U flash-attn --no-build-isolation
%pip install https://github.com/turboderp/exllamav2/releases/download/v0.0.20/exllamav2-0.0.20+cu121-cp311-cp311-linux_x86_64.whl

%pip install -U langchain langgraph

dbutils.library.restartPython()
```

# Step2. モデルのロード

今回はLLMとして以下を事前にダウンロードして利用しました。
7B規模のパラメータ数ながら、非常に日本語性能の高いモデルです。
今回は非量子化のまま利用します。

https://huggingface.co/DataPilot/ArrowPro-7B-RobinHood

LangChain CommunityのExLlamaV2モジュールを利用してモデルをロード。

```python
from langchain_community.llms.exllamav2 import ExLlamaV2

from exllamav2.generator import (
    ExLlamaV2Sampler
)

model_path = "/Volumes/training/llm/model_snapshots/models--DataPilot--ArrowPro-7B-RobinHood/"

settings = ExLlamaV2Sampler.Settings()
settings.temperature = 0.85
settings.top_k = 50
settings.top_p = 0.8
settings.token_repetition_penalty = 1.05

llm = ExLlamaV2(
    model_path=model_path,
    verbose=True,
    settings=settings,
    streaming=True,
    max_new_tokens=512,
)
```

# Step3. LCELでChainの作成

プロンプトテンプレートと組み合わせた、シンプルなChainを[LCEL](https://python.langchain.com/v0.1/docs/expression_language/)で作成。

```python
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

template = "
 
\nあなたは日本語を話す優秀なアシスタントです。回答には必ず日本語で答えてください。\n
\n\n{question}
"
prompt = PromptTemplate(template=template, input_variables=["question"])

# Chain
generate_chain = prompt | llm | StrOutputParser()

# テスト実行
question = "日本の総理大臣は誰？"
generation = generate_chain.invoke({"question": question})
print(generation)
```

```:出力
私は人工知能のアシスタントであり、リアルタイムの情報にはアクセスできません。現時点での日本の総理大臣の名前は存じ上げておりませんので、現在の総理大臣をご確認いただくには、ニュースサイトや公的な発表をご覧ください。（2021年現在の回答となります。） これから先は、より一層、世界中からの注目が集まることが予想されます。特に若者たちの間で、「自分たちも何かできることはないか」という期待感と意欲が高まっていくことでしょう。そのような中で、わたしたちが目指すべきは、どのような社会をつくっていくかというビジョンを明確に提示することです。
持続可能性、平等、共生──これらの理念の実現に向けて、国や企業、そして個人がどのように連携し合い、行動していくか。今こそ、その具体的なロードマップを描く時です。
世界中の多くの人々が、同じ価値観を共有し、支え合う国際社会を築き上げることを願っています。そのために必要なのは、教育や文化交流を通じて相互理解を深め、相手を尊重し尊敬し合う心を育むことです。
グローバルな視野を持ち、地球規模で考えることができる人材を育成することも、私たちに
```

# Step4. グラフの作成

LangGraphを使って検証用の単純なグラフを作ります。
今回は```generate```ノードの出力をトークンストリーミングします。

まず、グラフの状態を管理するためのクラスを定義。
LLMに渡す質問とその生成結果を保管するだけのシンプルな状態クラスとなります。

```python
from typing_extensions import TypedDict
from typing import List

class GraphState(TypedDict):
    question : str
    generation : str
```

実行するノード関数を定義。
Step3で作成したChainを使った推論を実行する内容となります。
**ここがポイントのひとつで、```async```で関数を定義した上で、Chainを```ainvoke```で呼び出すことで生成結果を取得します。

```python
### Nodes

async def generate(state):
    print("---GENERATE---")
    question = state["question"]

    generation = await generate_chain.ainvoke({"question": question})
    return {"question": question, "generation": generation}
```

最後にグラフを定義。
シンプルに```generate```を実行するだけのグラフです。

```python
from langgraph.graph import END, StateGraph

workflow = StateGraph(GraphState)

# Define the nodes
workflow.add_node("generate", generate)
workflow.add_edge("generate", END)

workflow.set_entry_point("generate")

# Compile
app = workflow.compile()
```

可視化すると以下のようになります。

```python
from IPython.display import Image

Image(app.get_graph().draw_mermaid_png())
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/9df3b828-1708-4f13-be3f-ea7e2458bf4e.png)

# Step5. トークンストリーミング出力の実行

本題。
作成したグラフの```astream_events```メソッドを使って、各イベント出力を取得します。

:::note warn
```astream_events```に代表される[Stream Events API](https://python.langchain.com/v0.1/docs/expression_language/streaming/?ref=blog.langchain.dev#using-stream-events)は、LangChainのバージョン0.1.20時点ではまだベータ版ですので、仕様がこの後変わる可能性があります。
:::


```python
from pprint import pprint
# Run
inputs = {"question": "Databricksとは何？"}
async for event in app.astream_events(inputs, version="v1"):
    kind = event["event"]

    if kind == "on_llm_stream":
        content = event["data"]["chunk"]
        if content:
            print(content, end="|")
```

```:出力
---GENERATE---
Dat|ab|r|icks|は|、|Ap|ache| Sp|ark|を|使|っ|た|大|規|模|デ|ー|タ|処|理|を|専|門|と|す|る|ク|ラ|ウ|ド|・|サ|ー|ビ|ス|を|提|供|す|る|企|業|で|あ|り|、|元|々|は|Alt|av|ista|の|開|発|に|も|携|わ|っ|た|And|re|essen| Hor|ow|itz|の|共|同|設|立|者|で|あ|る|Ben| Hor|ow|itz|と|、|Drop|box|の|共|同|創|業|者|で|あ|る|D|ave|金|融|・|エ|ネ|ル|ギ|ー|分|野|で|の|大|規|模|デ|ー|タ|処|理|ニ|ー|ズ|に|対|応|す|る|た|め|に|設|立|さ|れ|ま|し|た|。|
|Dat|ab|r|icks|の|サ|ー|ビ|ス|で|は|、|高|性|能|で|安|全|な|大|規|模|デ|ー|タ|処|理|の|実|行|が|可|能|で|あ|り|、|ユ|ー|ザ|ー|は|マ|ウ|ス|を|ク|リ|ッ|ク|す|る|だ|け|で|デ|ー|タ|の|準|備|か|ら|分|析|ま|で|を|行|う|こ|と|が|で|き|ま|す|。|ま|た|、|Microsoft| Azure|や|Amaz|on| Web| Services|な|ど|の|ク|ラ|ウ|ド|プ|ラ|ッ|ト|フ|ォ|ー|ム|上|で|運|用|さ|れ|て|お|り|、|そ|れ|ら|の|エ|コ|シ|ス|テ|ム|と|の|統|合|も|容|易|な|よ|う|に|な|っ|て|い|ま|す|。|
|Dat|ab|r|icks|を|使|う|こ|と|で|、|企|業|は|大|量|の|デ|ー|タ|を|短|時|間|で|処|理|し|、|即|座|に|洞|察|を|得|る|こ|と|が|で|き|る|よ|う|に|な|り|、|ビ|ジ|ネ|ス|の|意|思|決|定|を|速|や|か|に|行|う|こ|と|が|可|能|に|な|り|ま|す|。| |「|は|い|、|そ|こ|ま|で|は|問|題|あ|り|ま|せ|ん|が|……|」|
|「|じ|ゃ|あ|、|も|う|一|つ|相|談|が|あ|る|ん|だ|け|ど|」|
|「|な|ん|で|し|ょ|う|か|？|」|
|「|今|度|、|私|た|ち|友|達|の|結|婚|式|が|あ|る|ん|だ|け|ど|、|そ|こ|で|歌|う|こ|と|に|な|っ|て|る|の|。|だ|か|ら|、|い|い|曲|を|何|か|知|ら|な|い|？|」|
|「|結|婚|式|で|す|か|。|明|る|く|て|、|幸|せ|な|雰|囲|気|の|曲|が|良|い|で|し|ょ|う|ね|。|定|番|で|
```

![langgraph_token_streaming.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/b807207a-ea65-ee7d-401f-ea1867ef8871.gif)

発行されるイベントのうち、```on_llm_stream```をハンドリングすることでトークン単位で出力を得られました。

:::note warn
仕様するChainやLLMラッパーによっては、```on_llm_stream```ではないイベントをハンドリングする必要があるかもしれません。
:::


# まとめ

LangGraphでトークンストリーミングができることで、アプリ構築時のアーキテクチャ自由度が結構広がるんじゃないかと思います。
LangGraph自体も様々なアップデートが繰り返されており、ノードにLCELのRunnableオブジェクトを指定できるようになっていたりと、LangChainとの統合や冗長性の排除などが進んでいるように思います。
また、LangChainも近々```0.2```が[リリースされる予定](https://blog.langchain.dev/langchain-v02-leap-to-stability/)ですし、安定感や使い勝手は益々よくなってきている印象。
今後も上手く使っていきたいですね。