Pydantic AI Graph Beta API：定義步驟與連線

在使用 Pydantic Graph 的 Beta API 時，我們不再需要繼承類別來建立節點，而是使用 GraphBuilder 配合 Python 的裝飾器與函數來完成。這種宣告式的寫法讓程式碼看起來更加簡潔、直覺。

在這篇文章中，我們將透過一個簡單的計數器範例，學習如何定義狀態、步驟 (Steps)，以及如何把它們連線 (Edges) 起來。

1. 建立 Builder 與狀態

首先，我們需要定義流程中共用的狀態 (State)，並實例化一個 GraphBuilder。在建立 Builder 時，你需要指定傳遞狀態的型別，以及最終輸出的型別。

from dataclasses import dataclass
from pydantic_graph.beta import GraphBuilder, StepContext

# 定義一個簡單的狀態物件，用來追蹤計數器的數值
@dataclass
class CounterState:
    value: int = 0

# 建立一個 GraphBuilder 實例
# 宣告：共用狀態為 CounterState，流程最終輸出一個 int 數值
g = GraphBuilder(state_type=CounterState, output_type=int)

2. 定義步驟 (Steps)

在 Beta API 中，節點被稱為「步驟 (Steps)」。要定義一個步驟，你只需要寫一個非同步函數 (Async Function)，並加上 @g.step 裝飾器。

每個步驟函數都會接收一個 StepContext 參數，它包含了三個重要的屬性：

ctx.state：你可以讀取或修改的共用狀態。
ctx.inputs：從上一個節點傳遞過來的輸入資料。
ctx.deps：外部注入的依賴資源（如果有設定的話）。

# 定義第一個步驟：將計數器加 1
@g.step
async def increment(ctx: StepContext[CounterState, None, None]) -> int:
    # 存取並修改狀態
    ctx.state.value += 1
    # 這裡回傳的值，將會變成下一個節點的 ctx.inputs
    return ctx.state.value

# 定義第二個步驟：將收到的輸入值乘以 2
@g.step
async def double_it(ctx: StepContext[CounterState, None, int]) -> int:
    # ctx.inputs 就是上一個節點回傳的數字
    return ctx.inputs * 2

3. 連線節點 (Edges)

定義好所有的步驟後，我們需要告訴 GraphBuilder 流程該怎麼走。這就是「連線 (Edges)」的工作。我們使用 g.add() 與 g.edge_from() 來畫出這張流程圖。

所有的 Graph 都有一個內建的起點 g.start_node 和終點 g.end_node。

# 定義流程走向
g.add(
    # 從起點走到 increment 步驟
    g.edge_from(g.start_node).to(increment),

    # 從 increment 步驟走到 double_it 步驟
    g.edge_from(increment).to(double_it),

    # 從 double_it 步驟走向終點
    g.edge_from(double_it).to(g.end_node),
)

4. 建置並執行流程

當流程圖畫好後，我們呼叫 g.build() 把它鎖定成一個可執行的 Graph。接著就可以開始運行了！

import asyncio

async def main():
    # 建置 Graph
    graph = g.build()

    # 準備初始狀態
    state = CounterState()

    # 執行整個流程，直到遇到終點
    result = await graph.run(state=state)

    print(f'最終回傳結果: {result}')     # 輸出: 2
    print(f'最終狀態的數值: {state.value}') # 輸出: 1

if __name__ == "__main__":
    asyncio.run(main())

進階技巧：手動逐步執行 (Step-by-Step Execution)

在除錯或製作互動式 UI 時，你可能不希望 graph.run() 一口氣把整個流程跑完。Beta API 提供了 graph.iter() 方法，讓你能夠「一步一步」地手動控制流程的推進。

async def step_by_step_demo():
    graph = g.build()
    state = CounterState()

    print(f'初始狀態: {state.value}')

    # 使用 iter() 開啟逐步執行的上下文
    async with graph.iter(state=state) as graph_run:
        # 使用 async for 迴圈，每次推進一個步驟
        async for event in graph_run:
            print(f'當前狀態={state.value} | 準備執行的事件={event}')

        # 當迴圈結束時，代表流程走到了終點
        if graph_run.output is not None:
            print(f'流程完成，最終輸出: {graph_run.output}')

# 這個功能讓你可以在每個步驟之間安插自訂的邏輯，例如暫停等待使用者確認！

產生視覺化流程圖

Beta API 內建了強大的視覺化功能，你可以隨時將畫好的流程圖輸出成 Mermaid 語法，這對於寫文件或除錯非常有幫助。

graph = g.build()

# 產生 Mermaid 語法字串，並設定流程圖為由左至右 (LR)
mermaid_diagram = graph.render(title='我的計數器流程', direction='LR')
print(mermaid_diagram)

產生的內容可以直接貼在支援 Mermaid 的 Markdown 編輯器中，它看起來會像是：

---
title: 我的計數器流程
---
stateDiagram-v2
  direction LR
  increment
  double_it
  [*] --> increment
  increment --> double_it
  double_it --> [*]

透過 Builder Pattern，我們將節點邏輯與流程控制解耦了。在下一篇文章中，我們將探討 Beta API 最強大的特色：如何輕易地設定平行執行與匯總結果。