🦜

LangGraphのMermaid出力機能とその活用事例

2024/08/30に公開

はじめに

こんにちは。PharmaX でエンジニアをしている諸岡(@hakoten)です。

この記事では、LLMアプリケーション開発をサポートするグラフ管理ツールである、LangGraphのMermaid出力機能の活用方法をご紹介します。

少しニッチな内容となりますが、既にLangGraphを使っている方や、これから使うことを検討している方にとって、参考になれば幸いです。

LangGraphとは

LangGraphは、LangChainのツール群の一つで、LLMエージェントのステップをグラフ化し、状態管理を行うためのツールです。

LangChainは、大規模言語モデル(LLM)を活用したアプリケーション開発を支援するフレームワークですが、LangGraphはその中でも特にグラフ管理に特化しており、LLMの実行とは独立して使用することができます。

詳しくは、過去に書いた記事もぜひご覧ください。

https://zenn.dev/pharmax/articles/8796b892eed183

環境

この記事執筆時点では、以下のバージョンで実施しています。
とくにLangChain周りは非常に開発速度が早いため、現在の最新バージョンを合わせてご確認ください

  • langgraph: 0.2.14
  • Python: 3.10.12

LangGraphのMermaid出力機能

LangGraphは、作成したグラフをMermaidフォーマットで出力することができます。

グラフの出力方法

例えば、以下のようなLangGraphのコードがあるとします。

from typing import Annotated
from typing_extensions import TypedDict
from langchain_core.runnables import RunnableConfig
from langgraph.graph import StateGraph

class State(TypedDict):
    value: str

def node(state: State, config: RunnableConfig):
    return {"value": "1"}

def node2(state: State, config: RunnableConfig):
    return {"value": "2"}

graph_builder = StateGraph(State)

graph_builder.add_node("node", node)
graph_builder.add_node("node2", node2)
graph_builder.add_edge("node", "node2")
graph_builder.set_entry_point("node")
graph_builder.set_finish_point("node2")

graph = graph_builder.compile()

このグラフをMermaidで出力するには、コンパイル後のオブジェクトに draw_mermaid メソッドを呼び出すだけです。

print(graph.get_graph().draw_mermaid())

(出力されるMermaidのフロー図)

%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
	__start__([__start__]):::first
	node(node)
	node2(node2)
	__end__([__end__]):::last
	__start__ --> node;
	node --> node2;
	node2 --> __end__;
	classDef default fill:#f2f0ff,line-height:1.2
	classDef first fill-opacity:0
	classDef last fill:#bfb6fc

また、MermaidのグラフはPNG形式の画像としても出力することが可能です。

from IPython.display import Image, display

display(Image(graph.get_graph().draw_mermaid_png()))

グラフのスタイルをカスタマイズ

Mermaidで出力されるグラフのスタイルは、draw_mermaid メソッドを使用してある程度カスタマイズすることが可能です。もちろん、Mermaid自体がコードベースのツールなので、出力されたコードを直接編集してカスタマイズすることもできます。

from langchain_core.runnables.graph_mermaid import CurveStyle, NodeStyles

mermaid = graph.get_graph().draw_mermaid(
        with_styles=True,
        curve_style=CurveStyle.LINEAR,
        node_colors=NodeStyles(
          default="fill:#f2f0ff,line-height:1.2",
          first="fill-opacity:0",
          last="fill:#bfb6fc"
        ),
        wrap_label_n_words=9
      )
print(mermaid)

draw_mermaid メソッドには、以下のような引数を指定できます。スタイルを調整したい場合には、これらを適宜変更してください。

引数名 デフォルト値 説明
with_styles True スタイルを含むかどうかを指定します。
curve_style CurveStyle.LINEAR エッジのスタイルを指定します。
node_colors NodeStyles() ノードの色を指定できます。
wrap_label_n_words 9 ノードラベルを折り返す単語数を指定します。

Mermaid出力の活用方法

PharmaXでは、LangGraphのMermaid出力機能をどのように活用しているかをご紹介します。

スナップショットテスト

まず、最初の活用方法として、スナップショットテストに利用しています。

スナップショットテストとは、フロントエンドでよく使用されるテスト手法の一つです。ローカルに特定のUIの状態をスナップショットとして保存し、その後のUIと比較することで、予期せぬ変更が発生していないかを確認することを目的としています。

LangGraphでも同様に、「意図しない形でグラフ構造が変更されていないか?」を確認するために、現在のグラフのスナップショットと draw_mermaid メソッドで生成したグラフを比較するテストを行っています。

以下は、弊社のプロダクト「YOJO」で使用されている業務フローの一部を抜粋したグラフの例です。

以下は、テストのサンプルコードです。

def test_compile():
    """
    グラフ構造が正しいことを確認する為のスナップショットテスト
    """
    graph = RuleBaseCheckGraph().compile()
    mermaid_output = graph.get_graph().draw_mermaid()
    expected_mermaid = """
    %%{init: {'flowchart': {'curve': 'linear'}}}%%
    graph TD;
        __start__([__start__]):::first
        _________7_______(次回決済予定日まで7日以下かどうか)
        ___________(要返信に上がっているか)
        ______2_______(前回の会話が2日以内かどうか)
        _______________________(ルールベースのルールが全て満たされているか判定)
        __end__([__end__]):::last
        __start__ --> ______2_______;
        __start__ --> _________7_______;
        __start__ --> ___________;
        _______________________ --> __end__;
        ______2_______ --> _______________________;
        _________7_______ --> _______________________;
        ___________ --> _______________________;
        classDef default fill:#f2f0ff,line-height:1.2
        classDef first fill-opacity:0
        classDef last fill:#bfb6fc
    """
    assert mermaid_output == expected_mermaid

RuleBaseCheckGraph().compile() は独自の処理ですが、内部ではStateGraphをコンパイルして、CompiledGraphを返すメソッドです。

このように、コンパイルされたグラフの状態と前回までのスナップショットを比較して、変更がないことを確認するテストを実施しています。

このテストには、スナップショットテストの目的の他に、グラフのコンパイル自体をテストできる というメリットもあります。コンパイルのテストを行うことで、実際にグラフを実行する前にノードやエッジの定義に誤りがないか確認することができます。

開発メンバー向けのドキュメント

もう一つの目的として、業務フローのコード実装であるMermaidを開発ドキュメントとしてチームメンバーに共有するために利用しています。

前述のスナップショットテストで作成されたMermaidのコードは、そのままコード上に残っているため、新しく改修作業を行う際や、類似の機能を開発する際に、そのMermaidコードを見ることで、現在の実装フローがすぐに把握できます。

このように、実装されたコードフローを可視化できることは、LangGraphを使用する上での大きな魅力の一つであり、Mermaid出力は非常に役立っています。

現在はスナップショットとしてコードに残しているだけですが、NotionやGitHubなどに自動でアップロードしてドキュメント化するなど、開発効率が向上する工夫は他にもありそうです。

Mermaid出力を行う上での注意事項

弊社でMermaid出力を運用する際にいくつかの注意点があったため、ここでご紹介します。

Mermaidで使用できない文字がある

LangGraphのノード名には任意の文字列を付けることができますが、Mermaidでは正しく動作しない文字が存在するため注意が必要です。例えば、コロン(:) などがその例です。

例えば次のようなグラフです。

from typing_extensions import TypedDict
from langgraph.graph import StateGraph

class State(TypedDict):
    value: str

graph_builder = StateGraph(State)

graph_builder.add_node("A: Node a", lambda: {"value": "1"})
graph_builder.add_node("B: Node b", lambda: {"value": "2"})
graph_builder.add_edge("A: Node a", "B: Node b")
graph_builder.set_entry_point("A: Node a")
graph_builder.set_finish_point("B: Node b")

graph = graph_builder.compile()

この例では、ノード名("A: Node a")にコロンを使用しています。このグラフをMermaidで出力すると、次のようになります。

%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
	__start__([__start__]):::first
	A__Node_a( Node a)
	B__Node_b( Node b)
	__end__([__end__]):::last
	A__Node_a --> B__Node_b;
	B__Node_b --> __end__;
	__start__ --> A__Node_a;
	classDef default fill:#f2f0ff,line-height:1.2
	classDef first fill-opacity:0
	classDef last fill:#bfb6fc

このように、「A: Node a」の「A」の部分が消えてしまっています。

さらに、グラフをネスト構造にするサブグラフを使用した場合、Mermaid上では subgraph として表現されますが、このときサブグラフ名に全角コロン(:)などを使うと、Mermaidのシンタックスエラーが発生してしまいます。

ノード名に特殊な文字を使用することはそれほど多くないかもしれませんが、LangGraph上で正しく動作していても、Mermaid上で問題が生じるケースがあるため、ノード名に記号を使用する際は注意が必要です。

条件付き分岐があるグラフ

特定の条件でノードから分岐が発生する場合(add_conditional_edgesを使う場合)には、Mermaidが正しく出力できるようにグラフを定義する必要があります。

例として次のようなグラフで考えます。

このグラフは、node_cは終了時に、「node_aに戻る」 or 「終了する(__end__へ遷移)」の2つの分岐を持っています。

LangGraphは、Mermaid出力時に実行(invoke)は行わず、コンパイル(compile)のみを行うため、出力経路は静的に解析されます。この点を考慮して、静的解析で正しい経路を出力できるように実装する必要があります。

静的解析の方法は大きく2つあり、その方法について以下で説明します。

add_conditional_edges メソッドの戻り値を使う

前述のグラフのLangGraphのコードは以下のようになっています。

グラフのコード
from langgraph.graph import StateGraph, END
from operator import add
from typing_extensions import TypedDict
from typing import Annotated, Literal

class State(TypedDict):
    path: Annotated[list[str], add]

graph_builder = StateGraph(State)

def routing() -> Literal["node_a", "__end__"]:
    return "node_a"

graph_builder.add_node("node_a", lambda: { "path": ["node_a"] })
graph_builder.add_node("node_b", lambda: { "path": ["node_b"] })
graph_builder.add_node("node_c", lambda: { "path": ["node_"] })

graph_builder.set_entry_point("node_a")
graph_builder.add_edge("node_a", "node_b")
graph_builder.add_edge("node_b", "node_c")
graph_builder.set_finish_point("node_c")

graph_builder.add_conditional_edges("node_c", routing)

graph = graph_builder.compile()
...
def routing() -> Literal["node_a", "__end__"]:
    return "node_a"
...
graph_builder.add_conditional_edges("node_c", routing)
...

このように、add_conditional_edges メソッドで渡すルーティング関数の戻り値には、特定のリテラル値を表現する型である「Literal」を使用することができます。

LangGraphでは、add_conditional_edges を使って分岐するグラフの経路を、関数の戻り値に定義されたリテラル値から取得します。

このグラフでは、Literal["node_a", "__end__"] という定義から、node_c 終了時に node_a または __end__ への分岐が発生することを戻り値から判断し、この経路情報をもとにMermaidが出力されます。

では、add_conditional_edges で定義する関数に Literal を使用しなかった場合はどうなるでしょうか?

次のコードでは、routing 関数の戻り値を Literal ではなく str に変更しています。

...
def routing() -> str:
    return "node_a"
...
graph_builder.add_conditional_edges("node_c", routing)
...
全てのコード
from langgraph.graph import StateGraph, END
from operator import add
from typing_extensions import TypedDict
from typing import Annotated, Literal

class State(TypedDict):
    path: Annotated[list[str], add]

graph_builder = StateGraph(State)

def routing() -> str:
    return "node_a"

graph_builder.add_node("node_a", lambda: { "path": ["node_a"] })
graph_builder.add_node("node_b", lambda: { "path": ["node_b"] })
graph_builder.add_node("node_c", lambda: { "path": ["node_"] })

graph_builder.set_entry_point("node_a")
graph_builder.add_edge("node_a", "node_b")
graph_builder.add_edge("node_b", "node_c")
graph_builder.set_finish_point("node_c")

graph_builder.add_conditional_edges("node_c", routing)

graph = graph_builder.compile()

Mermaidで出力されるグラフは次のようになります。

このように、経路が特定できないため、node_c からすべてのノードへの分岐が出力され、正しいグラフとして表示されていないことが分かります。

path_mapを指定する

前述のように、add_conditional_edges の戻り値に Literal を使用しない場合、正しい経路が出力されないことがあります。

しかし、PharmaXではノード名を定数値で管理しているため、Literal には定数を定義できず、add_conditional_edges のルーティング関数の戻り値として str を使用しています。

このような場合は、add_conditional_edgespath_map 引数を追加することで、問題を解決することができます。

...
# ルーティングはstrを返す
def routing() -> str:
    return "node_a"
...
graph_builder.add_conditional_edges(
    source="node_c", 
    path=routing,
    # path_map引数に、分岐対象のノード名を指定する
    path_map=['node_a', END]
)
...
全てのコード
from langgraph.graph import StateGraph, END
from operator import add
from typing_extensions import TypedDict
from typing import Annotated, Literal

class State(TypedDict):
    path: Annotated[list[str], add]

graph_builder = StateGraph(State)

def routing() -> str:
    return "node_a"

graph_builder.add_node("node_a", lambda: { "path": ["node_a"] })
graph_builder.add_node("node_b", lambda: { "path": ["node_b"] })
graph_builder.add_node("node_c", lambda: { "path": ["node_"] })

graph_builder.set_entry_point("node_a")
graph_builder.add_edge("node_a", "node_b")
graph_builder.add_edge("node_b", "node_c")
graph_builder.set_finish_point("node_c")

graph_builder.add_conditional_edges(
    source="node_c", 
    path=routing,
    # path_map引数に、分岐対象のノード名を指定する
    path_map=['node_a', END]
)

graph = graph_builder.compile()

このように、add_conditional_edges で対象のノード名を指定することで、Mermaidで正しい分岐を出力することができます。

LangGraphの path_map のコード

実際に、LangGraphの add_conditional_edges メソッドは以下のように実装されています。

path_map 引数が指定されない場合は、path(指定された関数やメソッド)の戻り値の型を取得し、Literal 型であれば、その引数を path_map としてセットしていることがわかります。

https://github.com/langchain-ai/langgraph/blob/a8758661bc4533fd4719860b0ed0687cdf402f04/libs/langgraph/langgraph/graph/graph.py#L248-L259

おわりに

今回は、LangGraphのMermaid出力機能について、その基本的な使い方とPharmaXでの具体的な活用事例を紹介しました。LangGraphを使用することで、業務フローやグラフの構造を視覚化し、より効率的に開発やデバッグを行えるようになりました。特にMermaidの出力機能は、スナップショットテストやドキュメント化において非常に役立っています。

PharmaXでは、AIやLLMに関連する技術の活用を積極的に進めています。もし、この記事が興味を引いた方や、LangGraphの活用に関心がある方は、ぜひ私のXアカウント(@hakoten)やコメントで気軽にお声がけください。PharmaXのエンジニアチームで一緒に働けることを楽しみにしています。

まずはカジュアルにお話できることを楽しみにしています!

PharmaXテックブログ

Discussion