用 LangGraph 寫 LeetCode 解題機器人 Agent
目錄
LangGraph 三部曲:本文是是範例篇:用 LangGraph 解 LeetCode,剖析實戰使用 LangGraph 有哪些要注意的
本文假設你對 LangGraph 有一定了解。如果不熟請先看 入門篇 與 進階篇
思路 #
最一開始的思路是
- 要能讀取題目跟題目已有的測資
- 請 LLM 分析題目,生成程式
- 請 LLM 分析題目,生出驗證測資 – 但可信度不知
- 無論哪種測資,都要能拿來測試 AI 生出的解法
- 如果解法錯誤,要能反省更正
- 人類要能介入並給意見
有些是單純工程問題,有些是 LLM prompt 問題,有些是流程問題
第一版隨手畫的流程圖如下
![Human flow chart for leetcode solver](./leetcode-draft-flow.webp)
思路轉換成 LangGraph 的框架 #
從手繪流程圖轉到 LangGraph 的 graph 表面上看起來一樣:定義 node function, 定義 edge / conditional edge, 以及定義 state
注意到手寫流程圖的菱形 Decision,對應到 LangGraph 這張是用 node 而非全然是 edge。後面詳解
![From human flow chart to LangGraph](./leetcode-langgraph-flow-sbs.webp)
「LangGraph 就是執行流程圖,所以有了流程圖就簡單了吧!」 – 對,也不對
>> 想被嚇?點開看 State 宣告 <<
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 做了兩種判斷
- 流程 / Flag 型:單純次數檢定,用來實踐 for-loop
- 商業邏輯:例如解法測試結果
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 保有對話記錄,根據之前的問答 + 自省做出比較好的解法答案,用了之前 進階篇的聊天機器人的例子
- 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]
- 在 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 機器人流程圖的最後一步是
- 給出 AI 生成的解法
- 詢問使用者覺得滿不滿意,無限期等使用者的答案
- 使用者可以人工檢查,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
![Explanation of main code invoking LangGraph graph](./main-code-explanation.webp)
原始碼可以從整個 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) 抱怨一下,他們版本真的跑太快了