🔖

LangGraphでSubGraphとPydanticを使った場合のハマりどころ

hakotensan

2024/08/05に公開

はじめに

この記事では、LangGraphでSubGraphと呼ばれるネスト構造を使った際に発生した問題点について、事例を紹介しています。

LangGraphについての詳細な説明は割愛しますので、LangGraphが何か知りたい方は、ぜひ次の記事を御覧ください。

環境

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

langgraph: 0.1.14
Python: 3.10.12

LangGraphのStateにPydanticを使用する

LangGraphでは、グラフを初期化するときに、グラフ全体で扱う状態(State)を定義することができます。LangGraphのチュートリアルでは主にTypedDict（辞書型）を使用していますが、実際にはどんな型でも使用可能です。

一方、Pythonでデータクラスを表現する際に良く使われるライブラリにPydanticがあります。Pydanticは、型ヒントやバリデーション機能を備えたクラスを提供するライブラリで、TypedDictを使うよりも堅牢で便利にデータを作成することができます。

LangGraphでも、グラフのStateにPydanticを使用する方法がHow to Guideで紹介されています。

弊社でも、キー値の存在が不定であるTypedDictよりも、Pydanticを使用する方が望ましいと考え、グラフのStateにはPydanticを利用しています。

しかし、Stateの運用にPydanticを使用すると、いくつか問題点が発生したためその事例を紹介します。

SubGraphとは

SubGraphについて簡単に説明します。
LangGraphのグラフは、あるグラフから別のグラフを呼び出す階層構造にすることが可能です。

例えば次のようなグラフ構造を指します。

実際のコード

from typing import Annotated, TypedDict

from langgraph.graph import StateGraph


def reduce_list(left: list | None, right: list | None) -> list:
    if not left:
        left = []
    if not right:
        right = []
    return left + right


class ChildState(TypedDict):
    path: Annotated[list[str], reduce_list]


class ParentState(TypedDict):
    path: Annotated[list[str], reduce_list]


child_builder = StateGraph(ChildState)

child_builder.add_node('child_start', lambda state: {'path': ['child_start']})
child_builder.add_node('child_node', lambda state: {'path': ['child_node']})
child_builder.add_node('child_end', lambda state: {'path': ['child_end']})

child_builder.set_entry_point('child_start')
child_builder.add_edge('child_start', 'child_node')
child_builder.add_edge('child_node', 'child_end')
child_builder.set_finish_point('child_end')

builder = StateGraph(ParentState)

builder.add_node('parent_start', lambda state: {'path': ['parent_start']})
builder.add_node('parent_node', lambda state: {'path': ['parent_node']})
builder.add_node('child_node', child_builder.compile())
builder.add_node('parent_end', lambda state: {'path': ['parent_end']})

builder.set_entry_point('parent_start')
builder.add_edge('parent_start', 'parent_node')
builder.add_edge('parent_node', 'child_node')
builder.add_edge('child_node', 'parent_end')
builder.set_finish_point('parent_end')
graph = builder.compile()

graph.get_graph(xray=1).draw_mermaid()

このグラフでは、parent_nodeからchild_nodeへとグラフがネストして呼び出されています。

このようなグラフのネスト構造をLangGraphでは SubGraph と表現しています。
SubGraphについての詳細は次の記事も参考にしてください。

事例1: SubGraphのStateのdefault値が適用されない

PydanticのBaseModelでは、インスタンスのデフォルトの値を設定しておくことができますが、LangGraphのSubGraphのStateにBaseModelを使った場合、このdefault値が上手く動作しないという問題があります。

(PydanticのBaseModelでデフォルト値を設定する例)

from langchain_core.pydantic_v1 import BaseModel, Field

class SomeState(BaseModel):
    value: str = Field(default='default value')

# インスタンス時に引数を渡さなくても、初期値を定義しておける
SomeState()
# => SomeState(value='default value')

次に、具体的なコードを示します。

全てのコード

from typing import Annotated, TypedDict

from langchain_core.pydantic_v1 import BaseModel, Field
from langgraph.graph import StateGraph


def reduce_list(left: list | None, right: list | None) -> list:
    if not left:
        left = []
    if not right:
        right = []
    return left + right


class ChildState(BaseModel):
    path: Annotated[list[str], reduce_list]
    value: str = Field(default='default value')


class ParentState(TypedDict):
    path: Annotated[list[str], reduce_list]


child_builder = StateGraph(ChildState)

child_builder.add_node('child_start', lambda state: {'path': ['child_start']})
child_builder.add_node('child_node', lambda state: {'path': ['child_node'], 'value': 'child value'})
child_builder.add_node('child_end', lambda state: {'path': ['child_end']})

child_builder.set_entry_point('child_start')
child_builder.add_edge('child_start', 'child_node')
child_builder.add_edge('child_node', 'child_end')
child_builder.set_finish_point('child_end')

builder = StateGraph(ParentState)

builder.add_node('parent_start', lambda state: {'path': ['parent_start']})
builder.add_node('parent_node', lambda state: {'path': ['parent_node']})
builder.add_node('child_node', child_builder.compile())
builder.add_node('parent_end', lambda state: {'path': ['parent_end']})

builder.set_entry_point('parent_start')
builder.add_edge('parent_start', 'parent_node')
builder.add_edge('parent_node', 'child_node')
builder.add_edge('child_node', 'parent_end')
builder.set_finish_point('parent_end')
graph = builder.compile()

このコードでは、SubGraphのStateにデフォルト値が設定された value フィールドが含まれています。

class ChildState(BaseModel):
    path: Annotated[list[str], reduce_list]
    value: str = Field(default='default value')

SubGraphは親のノードの一部として以下のように実行されます。

...
builder.add_node('child_node', child_builder.compile())
...

このグラフをinvokeすると、以下のエラーが発生します。

ValidationError                           Traceback (most recent call last)
<ipython-input-8-a4d6077ec631> in <cell line: 1>()
----> 1 graph.invoke({'path': []})

15 frames
/usr/local/lib/python3.10/dist-packages/pydantic/v1/main.py in __init__(__pydantic_self__, **data)
    339         values, fields_set, validation_error = validate_model(__pydantic_self__.__class__, data)
    340         if validation_error:
--> 341             raise validation_error
    342         try:
    343             object_setattr(__pydantic_self__, '__dict__', values)

ValidationError: 1 validation error for ChildState
value
  none is not an allowed value (type=type_error.none.not_allowed)

このエラーは、BaseModelを継承したChildStateの初期化時に value プロパティのデフォルト値が設定されていないために発生しています。

これはLangGraphではBaseModelの本来期待する初期化ではない方法でインスタンスが生成されているためだと思われます。

また、SubGraphでは起動時の入力パラメータ（invoke時に渡すinput）を指定することができない という点もポイントで、現時点ではNode指定時にSubGraphのStateの初期値を渡すことはできないため、SubGraphのみで使用されるStateのフィールドを初期化することができません。

回避方法: イニシャライザを使って自分でデフォルト値を設定する

この問題に対して、弊社では以下のようにイニシャライザを自前で定義し、デフォルト値を指定することで対応しています。

class ChildState(BaseModel):
    def __init__(self, **input):
      input['value'] =  input.get('value') or 'default value'
      super().__init__(**input)

    path: Annotated[list[str], reduce_list]
    value: str

事例2: SubGraphのStateを親のStateが持っていないとSubGraph実行時にエラーになる

Pydanticを使ってSubGraphのStateを定義した際、グラフの実行時に次のようなエラーが発生することがあります。

InvalidUpdateError(f"Expected dict, got {input}")

これは、子グラフのStateが保持するフィールドを親グラフのStateが保持していないために発生します。

具体的には、次のようなコードでこのエラーが発生します。こちらは前述のSubGraphの説明で使用したものと同じ構成のグラフですが、StateにはPydanticを使用しています。

全てのコード

from typing import Annotated, TypedDict

from langchain_core.pydantic_v1 import BaseModel
from langgraph.graph import StateGraph

def reduce_list(left: list | None, right: list | None) -> list:
    if not left:
        left = []
    if not right:
        right = []
    return left + right

# PydanticでStateを宣言
class ChildState(BaseModel):
    def __init__(self, **input):
      input['value'] =  input.get('value') or '_get_state_key'
      super().__init__(**input)

    path: Annotated[list[str], reduce_list]
    # 子のグラフにだけ存在するフィールド
    value: str


# PydanticでStateを宣言
class ParentState(BaseModel):
    path: Annotated[list[str], reduce_list]


child_builder = StateGraph(ChildState)

child_builder.add_node('child_start', lambda state: {'path': ['child_start']})
child_builder.add_node('child_node', lambda state: {'path': ['child_node']})
child_builder.add_node('child_end', lambda state: {'path': ['child_end']})

child_builder.set_entry_point('child_start')
child_builder.add_edge('child_start', 'child_node')
child_builder.add_edge('child_node', 'child_end')
child_builder.set_finish_point('child_end')

builder = StateGraph(ParentState)

builder.add_node('parent_start', lambda state: {'path': ['parent_start']})
builder.add_node('parent_node', lambda state: {'path': ['parent_node']})
builder.add_node('child_node', child_builder.compile())
builder.add_node('parent_end', lambda state: {'path': ['parent_end']})

builder.set_entry_point('parent_start')
builder.add_edge('parent_start', 'parent_node')
builder.add_edge('parent_node', 'child_node')
builder.add_edge('child_node', 'parent_end')
builder.set_finish_point('parent_end')
graph = builder.compile()

このグラフではChildStateにのみ存在する value フィールドが含まれています。

...
# PydanticでStateを宣言
class ChildState(BaseModel):
    def __init__(self, **input):
      input['value'] =  input.get('value') or '_get_state_key'
      super().__init__(**input)

    path: Annotated[list[str], reduce_list]
    # 子のグラフにだけ存在するフィールド
    value: str
...

# PydanticでStateを宣言
class ParentState(BaseModel):
    path: Annotated[list[str], reduce_list]
...

このグラフをinvokeすると、以下のエラーが発生します。

graph.invoke({'path': []}, debug=True)

[0:tasks] Starting step 0 with 1 task:
- __start__ -> {'path': []}
[0:writes] Finished step 0 with writes to 1 channel:
- path -> []
[1:tasks] Starting step 1 with 1 task:
- parent_start -> ParentState(path=[])
[1:writes] Finished step 1 with writes to 1 channel:
- path -> ['parent_start']
[2:tasks] Starting step 2 with 1 task:
- parent_node -> ParentState(path=['parent_start'])
[2:writes] Finished step 2 with writes to 1 channel:
- path -> ['parent_node']
[3:tasks] Starting step 3 with 1 task:
- child_node -> ParentState(path=['parent_start', 'parent_node'])

...
lib/python3.12/site-packages/langgraph/graph/state.py", line 353, in _get_state_key
    raise InvalidUpdateError(f"Expected dict, got {input}")

child_nodeの実行前にエラーが発生しました。このエラーが発生しているstate.pyのコードを見ると、次のようになっています。

この処理はchild_node実行時にstateを初期化する処理だと思われます。input引数にはParentStateが、keyには name が渡されてきます。

PydanticのBaseModelは辞書型（dict）ではないため、期待する分岐である elif get_type_hints(type(input)).get(key): の箇所に入らないことがエラーの原因のようです。

つまり、Stateに辞書型以外の型を使う場合、親のStateが子のフィールドを全て持っていないとエラーになるということです。

回避方法

~~弊社では、仕方なく子のStateのフィールドを全て親のフィールドに含めて対応していました。~~

辞書型以外の型が使用された場合でも、キーの存在確認までは行わないように修正されています。

おわりに

以上、LangGraphでPydanticを使った際のハマりどころについて、弊社の事例を紹介させていただきました。

LangChain周りのコードはバージョンの更新が早く、LangGraphも比較的新しいライブラリのため、バグなのか仕様なのかが不明瞭なことがあります。

おかしな挙動が発生した場合は、最新バージョンでの確認やバージョンごとの差分を丁寧に確認することが重要です。

リリースノートを見る限りPydanticへの対応も進んでいるように見えるため、徐々に使いやすくなることが期待できそうです。

PharmaXでは、様々なバックグラウンドを持つエンジニアの採用をお待ちしております。弊社はAI活用にも力を入れていますので、LLM関連の開発に興味がある方もぜひ気軽にお声がけください。

興味をお持ちの場合は、私のXアカウント（@hakoten）や記事のコメントにお気軽にメッセージをいただければと思います。まずはカジュアルにお話できれば嬉しいです！

PharmaXテックブログPublication

PharmaXエンジニアチームのテックブログです。エンジニアメンバーが、PharmaXの事業を通じて得た技術的な知見や、チームマネジメントについての知見を共有します。 PharmaXエンジニアチームやメンバーの雰囲気が分かるような記事は、note（note.com/pharmax）もご覧ください。