用 LangGraph 寫 LeetCode 解題機器人 Agent

LangGraph 三部曲：本文是是範例篇：用 LangGraph 解 LeetCode，剖析實戰使用 LangGraph 有哪些要注意的

本文假設你對 LangGraph 有一定了解。如果不熟請先看 入門篇 與 進階篇

思路 #

最一開始的思路是

• 要能讀取題目跟題目已有的測資
• 請 LLM 分析題目，生成程式
• 請 LLM 分析題目，生出驗證測資 – 但可信度不知
• 無論哪種測資，都要能拿來測試 AI 生出的解法
• 如果解法錯誤，要能反省更正
• 人類要能介入並給意見

有些是單純工程問題，有些是 LLM prompt 問題，有些是流程問題

第一版隨手畫的流程圖如下

思路轉換成 LangGraph 的框架 #

從手繪流程圖轉到 LangGraph 的 graph 表面上看起來一樣：定義 node function, 定義 edge / conditional edge, 以及定義 state

注意到手寫流程圖的菱形 Decision，對應到 LangGraph 這張是用 node 而非全然是 edge。後面詳解

「LangGraph 就是執行流程圖，所以有了流程圖就簡單了吧！」 – 對，也不對

class AgentState(TypedDict):
  # Chat history of solution-generating LLM
  main_coding_llm_messages: Annotated[list[BaseMessage], add_messages]

  # leetcode problem definition and example tests
  problem: LeetCodeProblem

  # Code generated by LLM
  main_code: str
  ai_test_code: str

  # Number of generations
  ai_test_code_revision: int
  main_code_revision: int
  main_code_revision_limit: int

  # Flags
  is_main_code_good: bool
  is_ai_test_code_good: bool
  skip_ai_test_code: bool

  # Test result
  last_test_result: CodeExecutionResult
  last_test_type: TestType

  # Human in the loop: feedback to give to main LLM
  has_human_interference: bool
  human_comment_on_main_code: str
  human_test_code: Annotated[str, lambda a, b: f"{a}\n{b}"]

「這麼長!?」所以我才先縮起來。從現在開始請務必打開原始碼對照

實作 #

大部分是照著流程圖，實作每個點跟線的 function。其中對應的 prompt 在 prompt.py 。底下列一些值得提的

Node 中的邏輯判斷 #

就算 LLM 做同件事，因為情境的不同，呼叫 LLM 的方式也可能不同

例如「產生解法程式」是第一次呢？還是驗證過且失敗了？會用不同的 prompt 給他錯誤訊息去反省

若想在同一個 node 處理，就要在裡面做判斷；判斷的方式是靠「變數」也就是 state，讓過去執行的 node 寫入關鍵的 state

例 node_generate_main_code() 181行

if (not state["main_code"]) or (not state["last_test_result"]):
  # 單純根據題目叫 LLM 產生解法
else:
  # 額外把 LLM 過去自省的對話與錯誤訊息給他
  history_messages = state["main_coding_llm_messages"]
  last_test_result = state["last_test_result"]
  ...
  response = main_coding_llm.invoke(history_messages + [message])

是否該拆成不同 node？ 這是可以考慮的，就像思考一個 python function 該做多少事

總之，State 含有某些 node 生成後，需要給其他 node 知道的資訊

# Code generated by LLM
main_code: str
ai_test_code: str
# Test result
last_test_result: CodeExecutionResult
last_test_type: TestType

Edge 的邏輯判斷 #

無論是 if-else 或 for-loop 流程，在 LangGraph 的實作都是 Conditional Edge (switch-case)

一個有趣的決定就是：conditional edge 裡面要做多複雜的事情？

以我的程式碼來說，edge 只看 state 做單純的判斷；主要計算在 node 裡面，把特性（注意，不見得是「決定」）放到 state

例 to_regenerate_main_code() 298行

def to_regenerate_main_code(state: AgentState) -> str:
  revision = state.get("main_code_revision") or 0
  revision_limit = state.get("main_code_revision_limit", 4)
  if revision >= revision_limit:
    return "no"

  if state.get("is_main_code_good", False):
    return "no"
  else:
    return "yes"

上面的 edge function 做了兩種判斷

1. 流程 / Flag 型：單純次數檢定，用來實踐 for-loop
2. 商業邏輯：例如解法測試結果

Conditional edge 就這麼單純。其中 is_main_code_good 是其他 node 決定，像是 228行 在驗證程式碼

def node_test_main_with_examples(state: AgentState):
  ...
  code_result = execute_code(
    combine_test_with_code(main_code, example_test_code)
  )
  return {
    "is_main_code_good": not code_result.has_error,
    ...
  }

我沒有在 edge 裡面去測試解法（而是在 node 做）。我也沒有在 node 裡直接說出要 skip_code_generation (而只是說 is_main_code_good，交由 edge 去判斷)

不能說這一定好，反而也凸顯了 LangGraph 隨便寫都還是有可能會動的特點（就跟寫普通 python 程式一樣）

總之，State 裡也有這兩類的值

# Number of generations
ai_test_code_revision: int
main_code_revision: int
main_code_revision_limit: int
# Flags
is_main_code_good: bool
is_ai_test_code_good: bool
skip_ai_test_code: bool

有記憶的對話 #

為了讓 LLM 保有對話記錄，根據之前的問答 + 自省做出比較好的解法答案，用了之前 進階篇的聊天機器人的例子

1. State 裡面放聊天歷史，list 的每個元素就是一筆訊息 (human/ai 交錯)，並且用 LangGraph 提供的 add_messages reducer, 概念上是 list 相加 (71行)

from langgraph.graph.message import add_messages

class AgentState(TypedDict):
  # Chat history of solution-generating LLM
  main_coding_llm_messages: Annotated[list[BaseMessage], add_messages]

2. 在 node 裡，呼叫 LLM 之前先從 state 拿到聊天歷史 list ; Node 的最後回傳「一輪」回答 – Human message 是對於問題的輸入、要求，而 AI message 是 LLM 輸出的程式解法等等 (198行)

  history_messages = state["main_coding_llm_messages"]

  message = HumanMessage(
    content=REGEN_BY_COMMENT_USER_PROMPT.format(
      human_comment=human_comment
    ))

  response = main_coding_llm.invoke(history_messages + [message])

  return {
    "main_coding_llm_messages": [message, response],
    ...
  }

Human-in-the-loop #

「如果前面一切順利，結果把 AI 的解法放到 LeetCode 上測試不過，怎麼辦？」這就是為什麼要 Human-in-the-loop

就像我們能與 ChatGPT 對話，不斷修正問題，「互動」是最重要的。因此 LeetCode 機器人流程圖的最後一步是

1. 給出 AI 生成的解法
2. 詢問使用者覺得滿不滿意，無限期等使用者的答案
3. 使用者可以人工檢查，submit 到 Leetcode 等等，最後回饋（例如多加測資，給文字意見）

Human-in-the-loop 的實作在 進階篇講了很多。這裡是用 interrupt_after=["give_answer"]，也就是

• give_answer 這個 node 只是給答案，沒做別的
• graph 執行到這會先結束，等使用者輸入。 等待使用者這件事不是在 node 裡面做，是在外面 python 這一層做
• 使用者的輸入由 update_state() 放到 state 裡
• 無論如何，都 resume 接續著 graph 的執行
• graph 用 conditional edge 判斷 state 裡面的使用者輸入，判斷要不要讓 graph 走到 END

原始碼可以從整個 solve_leetcode() 開始看 – 他呼叫 graph 執行，是在外層不是在 graph 內部的

流程概念是 do-while：第一次先讓 LLM 生程式，之後看使用者要不要進行下一輪生成

相對應的 State 如下，其中人類測資希望能不斷累積，所以用 reducer

# Human in the loop: feedback to give to main LLM
has_human_interference: bool
human_comment_on_main_code: str
human_test_code: Annotated[str, lambda a, b: f"{a}\n{b}"]

AI 程式執行沙盒 #

這跟 LangGraph, LLM 一點關係都沒有，單純是「你要怎麼執行來路不明的程式而不被傷害？」 否則你電腦的檔案就不見了

我倒沒有特別講究，單純弄個 docker container 當沙盒，裡面放 API server 把程式碼文字 POST 進去，回傳 stdout/stderr。電腦資源用 docker 控制。非常陽春，無法輸入 stdin / 檔案，也無法 on-the-fly 裝 package

原始碼在 https://github.com/unclefomotw/code-executor

哪裡可以做更好？ #

單元假設驗證 #

LeetCode 機器人看起來很唬人，實際有個致命缺點：「用 AI 生測試資料，比 AI 生的程式還不可信」 – 這點倒是跟 AlphaCodium 的觀察相反

如果根據這個觀察，這流程圖壓根就是不妥當的，而這件事情應該要早點發現，用更多資料去驗證，如果為真的話去修正

所以這是軟體開發的問題：如何更 agile 去調整 – 而 LangGraph 以單元（一個個 node function）開發來說並不難做到

（題外話，對於 AI 寫程式有興趣的，可以看 AlphaCodium 論文的第四章， 概念是先拿到一堆一定是對的測試資料，確定 AI 程式會通過這些測資，然後讓 AI 一筆一筆生測資，每次的測資如果不過，先假設 AI 解法程式是錯的，改寫後確定正確的測資一定要通過，剩下就是看新的這一筆有沒有通過）

Mock #

一些情境會想做單元測試，或是在 end-to-end 上 mock 某一些 LLM 呼叫的 output，為了

• Deterministic, 方便實驗偵錯
• 省錢 $$$

我在這裡並沒有做，而是仗著 gpt-3.5 便宜，但無論是不是用 LangGraph 甚至是不是用框架，有個能夠動態「錄影+播放」的 cache 或 mock 是非常實用的

（例如以前沒呼叫過就呼叫 LLM，把 output 記錄下來，之後的呼叫可選擇性地直接回傳同一個 output – 可能是根據 input，也可能是人工切換）

測試版本: LangGraph 0.0.65 (+LangChain 0.2.3) 抱怨一下，他們版本真的跑太快了