LiteLLMとLangGraphでPlan-and-ExecuteとReflexionの2パターンを実装してみる

Python 3.13
litellm 1.83.14
langgraph 1.1.10
langchain-litellm 0.6.4
langchain-core 1.3.2

$ uv pip install litellm langgraph langchain-litellm langchain-core
$ export OPENAI_API_KEY="sk-..."
$ export ANTHROPIC_API_KEY="sk-ant-..."

"""Plan-and-Execute エージェントの実装。

ReActが「考える→1ツール呼ぶ」を繰り返すのに対し、Plan-and-Executeは
最初に計画を立てて実行→必要なら再計画、というフローを取る。
"""

from typing import TypedDict

import litellm
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.tools import tool
from langchain_litellm import ChatLiteLLM
from langgraph.graph import END, START, StateGraph

# Anthropic（replanner）はツール無し呼出時に modify_params=True が必要
litellm.modify_params = True

# =============================================
# ツール: Executor が利用する
# =============================================
@tool
def get_population(city: str) -> str:
    """都市の人口を返します。"""
    data = {
        "東京": "約1,400万人",
        "ニューヨーク": "約830万人",
        "ロンドン": "約900万人",
        "上海": "約2,500万人",
    }
    return data.get(city, f"{city}の人口データはありません")

@tool
def calculate(expression: str) -> str:
    """数式を計算します。"""
    allowed = set("0123456789+-*/.() ")
    if not all(c in allowed for c in expression):
        return "エラー: 許可されていない文字"
    try:
        return str(eval(expression))  # noqa: S307
    except Exception as e:
        return f"計算エラー: {e}"

# =============================================
# モデル: ロールごとに使い分ける
# =============================================
# Planner: 計画立案。質の高い思考を求めるので高性能モデル
planner_llm = ChatLiteLLM(model="openai/gpt-5.4")
# Executor: 単純なステップ実行なので安価・高速モデル
executor_llm = ChatLiteLLM(model="openai/gpt-5-mini")
# Replanner: 進捗を見て再計画。判断精度を担保するためClaudeに切替
replanner_llm = ChatLiteLLM(model="anthropic/claude-sonnet-4-6")

# Executor は ReAct エージェントとして実行
executor_agent = create_agent(
    model=executor_llm,
    tools=[get_population, calculate],
    system_prompt="あなたはステップ実行担当。与えられたサブタスクをツールを使って解決し、結果だけ返してください。",
)

# =============================================
# ステート定義
# =============================================
class State(TypedDict):
    task: str
    plan: list[str]  # 残りのステップ
    past_steps: list[tuple[str, str]]  # (step, result)
    response: str

# =============================================
# Planner ノード: 初回の計画を立てる
# =============================================
PLANNER_PROMPT = """\
あなたはタスクを段階的に分解するプランナーです。
ユーザーのタスクを解決するための具体的なステップを 3〜5 個列挙してください。
各ステップは独立して実行可能な粒度にしてください。

出力形式: 1行に1ステップ。番号や記号は付けない。
"""

def planner_node(state: State) -> dict:
    response = planner_llm.invoke(
        [SystemMessage(content=PLANNER_PROMPT), HumanMessage(content=state["task"])]
    )
    plan = [s.strip() for s in str(response.content).strip().split("\n") if s.strip()]
    print(f"\n[Plan] {len(plan)} ステップ:")
    for i, step in enumerate(plan, 1):
        print(f"  {i}. {step}")
    return {"plan": plan, "past_steps": []}

# =============================================
# Executor ノード: 計画の先頭ステップを実行する
# =============================================
def executor_node(state: State) -> dict:
    if not state["plan"]:
        return {}
    current_step = state["plan"][0]
    result = executor_agent.invoke({"messages": [HumanMessage(content=current_step)]})
    answer = str(result["messages"][-1].content)
    print(f"  [Execute] {current_step} -> {answer[:60]}")
    return {
        "plan": state["plan"][1:],
        "past_steps": state["past_steps"] + [(current_step, answer)],
    }

# =============================================
# Replanner ノード: 完了/再計画を判断する
# =============================================
REPLANNER_PROMPT = """\
あなたは進捗を見て次の行動を決める判断者です。

これまでの実行結果を踏まえて、以下のいずれかを返してください:
- 「FINISH」 + 改行 + ユーザーへの最終回答  （タスク完了の場合）
- 「CONTINUE」 + 改行 + 残りのステップを箇条書き  （まだ実行が必要な場合）

判断基準:
- 残りの計画ステップが空でも、結果を統合して回答すれば良いだけならFINISH
- 元のタスクに対する答えが既に出ているならFINISH
- 抜けがある、追加調査が必要ならCONTINUE
"""

def replanner_node(state: State) -> dict:
    history = "\n".join(
        f"- {step}\n  -> {result}" for step, result in state["past_steps"]
    )
    remaining = "\n".join(f"- {s}" for s in state["plan"]) or "(なし)"
    response = replanner_llm.invoke(
        [
            SystemMessage(content=REPLANNER_PROMPT),
            HumanMessage(
                content=f"タスク: {state['task']}\n\n実行済み:\n{history}\n\n残り計画:\n{remaining}"
            ),
        ]
    )
    text = str(response.content).strip()
    if text.upper().startswith("FINISH"):
        body = text.split("\n", 1)[1] if "\n" in text else ""
        print("\n[Replanner] FINISH")
        return {"response": body.strip(), "plan": []}

    # CONTINUE: 残りステップを更新
    lines = text.split("\n")[1:]
    new_plan = [line.lstrip("- *・").strip() for line in lines if line.strip()]
    print(f"\n[Replanner] CONTINUE - 残り {len(new_plan)} ステップ")
    return {"plan": new_plan}

# =============================================
# 条件分岐: Replanner 後に終了か継続か
# =============================================
def should_end(state: State) -> str:
    if state.get("response"):
        return "end"
    return "continue"

# =============================================
# グラフ構築
# =============================================
graph = StateGraph(State)
graph.add_node("planner", planner_node)
graph.add_node("executor", executor_node)
graph.add_node("replanner", replanner_node)

graph.add_edge(START, "planner")
graph.add_edge("planner", "executor")
graph.add_edge("executor", "replanner")
graph.add_conditional_edges(
    "replanner",
    should_end,
    {"continue": "executor", "end": END},
)

app = graph.compile()

# =============================================
# 実行
# =============================================
result = app.invoke(
    {
        "task": "東京とニューヨークの人口の合計と差を計算してください",
        "plan": [],
        "past_steps": [],
        "response": "",
    }
)

print("\n=== 最終回答 ===")
print(result["response"])

$ python plan_execute.py

[Plan] 4 ステップ:
  1. 東京とニューヨークの最新の人口データを確認する
  2. 東京とニューヨークで比較対象となる範囲（都市本体か都市圏か）を揃える
  3. 確認した人口値を使って合計を計算する
  4. 同じ人口値を使って差を計算する
  [Execute] 東京とニューヨークの最新の人口データを確認する -> Tokyoの人口データはありません
New Yorkの人口データはありません

[Replanner] FINISH

=== 最終回答 ===
## 東京とニューヨークの人口：合計・差の計算

ツールによる自動取得ができなかったため、一般に広く参照されている統計データをもとに回答します。

### 使用データ（直近の推計値）

| 対象 | 東京 | ニューヨーク |
|------|------|-------------|
| **都市本体**（市区レベル） | 約 **960万人**（東京23区） | 約 **820万人**（NYC市） |
| **都市圏**（Greater area） | 約 **3,700万人** | 約 **2,000万人** |

### 計算結果（都市圏）

| 計算 | 結果 |
|------|------|
| **合計** | 3,700万 ＋ 2,000万 ＝ **約5,700万人** |
| **差** | 3,700万 － 2,000万 ＝ **約1,700万人**（東京圏が多い） |
（以下略）

"""Reflexion パターン: 自己批判ループによる出力品質向上。

生成 -> 批評 -> 再生成 を繰り返す。エッセイや要約など、品質が
反復で改善するタスクに向く。
"""

from typing import TypedDict

import litellm
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_litellm import ChatLiteLLM
from langgraph.graph import END, START, StateGraph

# Anthropic（reflector）はツール無し呼出時に modify_params=True が必要
litellm.modify_params = True

# Generator: ドラフトを書く（多数回呼ばれるので安価モデル）
generator_llm = ChatLiteLLM(model="openai/gpt-5-mini")
# Reflector: 批評する。鋭い指摘が欲しいのでClaudeに切替
reflector_llm = ChatLiteLLM(model="anthropic/claude-sonnet-4-6")

MAX_ITERATIONS = 3

class State(TypedDict):
    task: str
    draft: str
    critique: str
    iteration: int
    is_acceptable: bool

GENERATOR_PROMPT = """\
あなたは執筆者です。与えられたタスクに対する文章を作成してください。
以前の批評がある場合は、それを反映して改善してください。
"""

REFLECTOR_PROMPT = """\
あなたは厳しいレビュアーです。提示された文章を以下の観点で評価してください:
- タスクへの適合性
- 論理性・構成
- 簡潔さ・読みやすさ

最後に「ACCEPT」または「REVISE」のいずれかを必ず1行で返してください。
REVISEの場合は具体的な改善点を箇条書きで列挙してください。
"""

def generator_node(state: State) -> dict:
    """ドラフトを生成（または批評を反映して再生成）。"""
    user_content = f"タスク: {state['task']}"
    if state["critique"]:
        user_content += f"\n\n前回の批評:\n{state['critique']}\n\n上記を反映して書き直してください。"
    response = generator_llm.invoke(
        [SystemMessage(content=GENERATOR_PROMPT), HumanMessage(content=user_content)]
    )
    print(f"\n[Generator#{state['iteration'] + 1}] ドラフト生成")
    return {
        "draft": str(response.content),
        "iteration": state["iteration"] + 1,
    }

def reflector_node(state: State) -> dict:
    """ドラフトを批評し、ACCEPT/REVISE を判定する。"""
    response = reflector_llm.invoke(
        [
            SystemMessage(content=REFLECTOR_PROMPT),
            HumanMessage(
                content=f"タスク: {state['task']}\n\nドラフト:\n{state['draft']}"
            ),
        ]
    )
    text = str(response.content)
    is_acceptable = (
        "ACCEPT" in text.upper().splitlines()[-1] if text.strip() else False
    )
    if "ACCEPT" in text.upper() and "REVISE" not in text.upper():
        is_acceptable = True
    print(f"  [Reflector] {'ACCEPT' if is_acceptable else 'REVISE'}")
    return {"critique": text, "is_acceptable": is_acceptable}

def should_continue(state: State) -> str:
    if state["is_acceptable"]:
        return "end"
    if state["iteration"] >= MAX_ITERATIONS:
        return "end"
    return "regenerate"

graph = StateGraph(State)
graph.add_node("generator", generator_node)
graph.add_node("reflector", reflector_node)

graph.add_edge(START, "generator")
graph.add_edge("generator", "reflector")
graph.add_conditional_edges(
    "reflector",
    should_continue,
    {"regenerate": "generator", "end": END},
)

app = graph.compile()

result = app.invoke(
    {
        "task": "AIエージェントの今後の展望を、技術者でない読者向けに3行で説明する文章を書いてください",
        "draft": "",
        "critique": "",
        "iteration": 0,
        "is_acceptable": False,
    }
)

print("\n=== 最終ドラフト ===")
print(result["draft"])
print(f"\n反復回数: {result['iteration']}")

$ python reflexion.py

[Generator#1] ドラフト生成
  [Reflector] REVISE

[Generator#2] ドラフト生成
  [Reflector] REVISE

[Generator#3] ドラフト生成
  [Reflector] REVISE

=== 最終ドラフト ===
AIエージェント（自律的なデジタルアシスタント）は、日々の作業を自動化して私たちの暮らしをもっと便利にしてくれます。
その結果、人はより創造的な仕事や重要な判断に集中できるようになります。
一方で、使い方のルールや人間の責任について、みんなでしっかり考えて決めていく必要があります。

反復回数: 3