[langgraph] 상태(state)에 pydantic 모델을 사용하는 방법

LangGraph 공식문서를 번역한 내용입니다. 필요한 경우 부연 설명을 추가하였고 이해하기 쉽게 예제를 일부 변경하였습니다. 문제가 되면 삭제하겠습니다.

https://langchain-ai.github.io/langgraph/how-tos/state-model/

StateGraph는 초기화 시 state_schema 인자를 받아들여 그래프의 노드들이 접근하고 업데이트할 수 있는 상태의 "형태"를 지정한다.

예시에서는 보통 상태 스키마로 Python-native TypedDict를 사용하지만 (MessageGraph의 경우 리스트를 사용), state_schema는 어떤 타입도 될 수 있다.

여기서는 Pydantic BaseModel을 state_schema로 사용하여 입력에 대한 런타임 검증을 추가하는 방법을 살펴보자.

준비

우선, 필요한 패키지를 설치하자.

pip install langgraph

입력값 검증

from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict

from pydantic import BaseModel


class OverallState(BaseModel):
    a: str


def node(state: OverallState):
    return {"a": "goodbye"}


builder = StateGraph(OverallState)
builder.add_node(node) 
builder.add_edge(START, "node")  
builder.add_edge("node", END)  
graph = builder.compile()

result = graph.invoke({"a": "hello"})
print(result)

{'a': 'goodbye'}

잘못된 입력으로 그래프를 호출하기

try:
    graph.invoke({"a": 123})  # string 이어야 한다
except Exception as e:
    print("An exception was raised because `a` is an integer rather than a string.")
    print(e)

An exception was raised because `a` is an integer rather than a string.
1 validation error for OverallState
a
  Input should be a valid string [type=string_type, input_value=123, input_type=int]
    For further information visit https://errors.pydantic.dev/2.9/v/string_type

다중 노드

런타임 검증은 다중 노드 그래프에서도 동작한다. 아래 예시에서 bad_node는 a를 정수로 업데이트한다.

입력에서 런타임 검증이 발생하기 때문에, 검증 오류는 bad_node가 상태를 업데이트할 때가 아니라, ok_node가 호출될 때 발생한다(상태가 스키마와 일치하지 않게 업데이트된 경우).

from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict

from pydantic import BaseModel


class OverallState(BaseModel):
    a: str


def bad_node(state: OverallState):
    return {
        "a": 123  # Invalid
    }


def ok_node(state: OverallState):
    return {"a": "goodbye"}


builder = StateGraph(OverallState)
builder.add_node(bad_node)
builder.add_node(ok_node)
builder.add_edge(START, "bad_node")
builder.add_edge("bad_node", "ok_node")
builder.add_edge("ok_node", END)
graph = builder.compile()

try:
    graph.invoke({"a": "hello"})
except Exception as e:
    print("An exception was raised because bad_node sets `a` to an integer.")
    print(e)

An exception was raised because bad_node sets `a` to an integer.
1 validation error for OverallState
a
  Input should be a valid string [type=string_type, input_value=123, input_type=int]
    For further information visit https://errors.pydantic.dev/2.9/v/string_type

LangGraph 참고 자료

저작자표시

준비

입력값 검증

다중 노드

LangGraph 참고 자료

티스토리툴바