상태 관리 — 스텝들이 값을 나누어 쓰는 법

상태 관리 — 스텝들이 값을 나누어 쓰는 법

여러 스텝이 공유해야 하는 값, 예컨대 지금까지 센 횟수 같은 걸 어디에 둬야 할까요? 워크플로 실행마다 Context가 있고, 그 컨텍스트마다 상태 저장소(state store) 가 있어요. 실행 중에 스텝들이 서로 나눠 써야 하는 값이나, 같은 컨텍스트를 재사용·복원해 다음 실행으로 의도적으로 넘길 값을 여기에 담아요.

출처: 공식문서

상태에 무엇을 둬야 하지 말아야 하나

상태는 무거운 클라이언트나 인덱스, 파일 핸들 같은 런타임 의존성의 자리가 아니에요. 그런 건 Resource에 두세요. 상태는 워크플로를 스냅샷하거나 재개할 때 직렬화할 의향이 있는 데이터여야 해요.

기본적으로 워크플로는 타입이 없는 상태 저장소를 초기화하고, ctx.store로 읽고 쓸 수 있어요.

from llama_index.core.workflow import Workflow, Context, step
from llama_index.core.workflow.events import StartEvent, StopEvent

class MyWorkflow(Workflow):
    @step
    async def my_step(self, ctx: Context, ev: StartEvent) -> StopEvent:
        current_count = await ctx.store.get("count", default=0)
        current_count += 1
        await ctx.store.set("count", current_count)
        return StopEvent(result=current_count)

상태 잠그기: edit_state()

여러 스텝이 동시에 상태를 건드리는 상황이라면 안전하게 읽고-수정하고-쓰는 작업을 원자적으로 해야 해요. 그럴 때는 edit_state()로 상태를 잠그면 됩니다.

@step
async def my_step(self, ctx: Context, ev: StartEvent) -> StopEvent:
    async with ctx.store.edit_state() as state:
        current_count = state.get("count", 0)
        state["count"] = current_count + 1
        result = state["count"]
    return StopEvent(result=result)

블록이 실행되는 동안에는 다른 스텝이 상태를 수정할 수 없어요. 블록은 작게 유지하세요. 읽고-수정하고-쓰는 작업은 안에 두고, 느린 LLM·네트워크 호출은 블록 바깥에 두는 게 좋아요.

타입드 상태: Pydantic 모델로

워크플로 상태의 형태가 정해져 있다면 Pydantic 모델을 쓰면 돼요. 타입 검사와 자동 초기화 등을 얻을 수 있어요. 모든 필드에 기본값을 가진 Pydantic 모델을 쓰면 Context가 상태를 자동으로 초기화해 줍니다.

from pydantic import BaseModel, Field

class CounterState(BaseModel):
    count: int = Field(default=0)

그리고 워크플로의 상태를 이 모델로 타입을 붙여주면 edit_state() 안에서 필드 단위로 원자적 갱신이 가능해요.

from llama_index.core.workflow import Workflow, Context, step
from llama_index.core.workflow.events import StartEvent, StopEvent

class MyWorkflow(Workflow):
    @step
    async def start(self, ctx: Context[CounterState], ev: StartEvent) -> StopEvent:
        async with ctx.store.edit_state() as state:
            state.count += 1
        return StopEvent(result="Done!")

타입드 상태를 필드 하나씩 다루고 싶다면 ctx.store.get("count")/ctx.store.set("count", ...)를 쓰거나 ctx.store.get_state()로 전체 상태를 가져올 수도 있어요.

직렬화 가능한 상태

커스텀 값을 저장한다면 Pydantic으로 직렬화 가능하게 만들거나, Context.to_dict()와 Context.from_dict()를 호출할 때 커스텀 시리얼라이저를 제공해야 해요. 스냅샷·재개 여정을 염두에 두고 상태 설계를 하라는 뜻이에요.

더 알아보기