LiteLLMとLangMemのprompt_optimizerでエージェントをオンライン学習・自己改善させてみる

from langmem import create_prompt_optimizer

optimizer = create_prompt_optimizer(
    "anthropic:claude-sonnet-4-6",
    kind="prompt_memory",
)

new_prompt = optimizer.invoke({
    "trajectories": [(messages, feedback), ...],
    "prompt": "あなたは...",
})

trajectories = [
    (セッション1の会話履歴, セッション1のフィードバック),
    (セッション2の会話履歴, セッション2のフィードバック),
    ...
]

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

$ uv pip install litellm langgraph langmem langchain

$ export OPENAI_AlangchainPI_KEY="sk-..."        # 応答役の gpt-5-mini 用
$ export ANTHROPIC_API_KEY="sk-ant-..." # optimizer 用 / サンプル4の judge 用

"""create_prompt_optimizer の kind を比較するサンプル。

LangMem の prompt_optimizer は内部実装によって 3 種類の `kind` を選択できる。
このサンプルでは `prompt_memory`（1 LLM 呼び出し）と `metaprompt`（1〜5 LLM 呼び出し）
の 2 種類を同じトラジェクトリに対して走らせ、生成されるプロンプトの差異と、
最適化後プロンプトを使った応答の変化を観察する。

`gradient` は最大 10 回程度の LLM 呼び出しを行う高コスト版なので、本サンプルでは
コメントだけ残して実行は省略する。
"""

import warnings

# langmem 0.0.30 起動時に langchain_core 系の DeprecationWarning が出るため抑制する。
warnings.filterwarnings("ignore", category=DeprecationWarning)

import litellm
from langchain_core.messages import AIMessage, HumanMessage
from langchain_litellm import ChatLiteLLM
from langmem import create_prompt_optimizer

# Anthropic 経由で空メッセージや tool_calls 整形が必要なケースに備えて modify_params を有効化。
litellm.modify_params = True
# gpt-5 系で未対応パラメータを LiteLLM 側で安全に落とすための保険。
litellm.drop_params = True

# 応答用は安価な GPT、最適化（optimizer）は文章生成が得意な Claude を割り当てる。
responder = ChatLiteLLM(model="openai/gpt-5-mini")
OPTIMIZER_MODEL = "anthropic:claude-sonnet-4-6"

INITIAL_PROMPT = (
    "あなたは Python 初学者向けのプログラミングアシスタントです。"
    "ユーザーの質問に丁寧に答えてください。"
)

# 過去の会話 + フィードバック（trajectory）を 2 件作る。
# trajectory のフォーマットは [(messages, feedback), ...]。feedback は自由形式の dict で、
# `user_score` のような数値だけでなく自由文での要望も渡せる。
TRAJECTORIES = [
    (
        [
            HumanMessage(content="リスト内包表記の基本を教えてください"),
            AIMessage(
                content=(
                    "リスト内包表記は反復処理を簡潔に書くための構文です。"
                    "for 文の代わりに使うと短く書けます。"
                )
            ),
        ],
        {
            "user_score": 0.3,
            "feedback": "説明が抽象的でした。具体的なコード例を最低1つ載せてほしい。",
        },
    ),
    (
        [
            HumanMessage(content="dict の get メソッドの使い方を教えてください"),
            AIMessage(
                content=(
                    "dict.get(key) はキーが存在しない場合に None を返します。"
                    "デフォルト値も第2引数で指定できます。"
                )
            ),
        ],
        {
            "user_score": 0.4,
            "feedback": "コード例がほしい。あと、なぜ d[key] ではなく get を使うかの判断基準も知りたい。",
        },
    ),
]

def show_response(prompt: str, question: str) -> str:
    """与えられた system prompt で `question` に応答させ、先頭部分を返す。"""
    response = responder.invoke(
        [
            {"role": "system", "content": prompt},
            {"role": "user", "content": question},
        ]
    )
    return str(response.content)

if __name__ == "__main__":
    print("=== 元のシステムプロンプト ===")
    print(INITIAL_PROMPT)

    # kind="prompt_memory": 1 LLM 呼び出しで一気にプロンプトを書き換える軽量版。
    pm_optimizer = create_prompt_optimizer(OPTIMIZER_MODEL, kind="prompt_memory")
    pm_prompt = pm_optimizer.invoke(
        {"trajectories": TRAJECTORIES, "prompt": INITIAL_PROMPT}
    )
    print("\n=== kind='prompt_memory' での最適化後プロンプト ===")
    print(pm_prompt)

    # kind="metaprompt": 内部で改善仮説の生成 → 適用までを複数ステップで行う。
    # max_reflection_steps を絞って LLM 呼び出しを抑える。
    mp_optimizer = create_prompt_optimizer(
        OPTIMIZER_MODEL,
        kind="metaprompt",
        config={"max_reflection_steps": 1, "min_reflection_steps": 1},
    )
    mp_prompt = mp_optimizer.invoke(
        {"trajectories": TRAJECTORIES, "prompt": INITIAL_PROMPT}
    )
    print("\n=== kind='metaprompt' での最適化後プロンプト ===")
    print(mp_prompt)

    # 最適化後プロンプトを実際の応答に使ってみる。
    question = "for ループでリストを2乗にしたいです。最短で書く方法は？"
    print("\n=== 質問 ===")
    print(question)

    print("\n=== prompt_memory 版プロンプトでの応答 ===")
    print(show_response(str(pm_prompt), question))

    print("\n=== metaprompt 版プロンプトでの応答 ===")
    print(show_response(str(mp_prompt), question))

    # `gradient` kind は反省ループを `max_reflection_steps` 回まで回すモード（デフォルト 3）。
    # 設定次第で LLM 呼び出しが 10 回近くに膨らむため、本サンプルでは実行を省略する。
    # 利用シーンとしては「学習用評価データが大きく、丁寧にプロンプトを練り込みたい」場合に向く。
    print("\n(gradient kind は LLM 呼び出しが膨らむため、本サンプルでは実行を省略)")

$ python compare_optimizer_kinds.py
=== 元のシステムプロンプト ===
あなたは Python 初学者向けのプログラミングアシスタントです。ユーザーの質問に丁寧に答えてください。

=== kind='prompt_memory' での最適化後プロンプト ===
あなたは Python 初学者向けのプログラミングアシスタントです。ユーザーの質問に丁寧に答えてください。

回答する際は以下のガイドラインに従ってください：
1. **必ずコード例を1つ以上含める**: 概念を説明する際は、具体的な動作するコードを示してください。
2. **実践的な判断基準を示す**: 「なぜこの方法を使うのか」「どんな場面で使うのか」という観点も説明してください。
3. **段階的に説明する**: 基本的な使い方から始め、必要に応じて応用例も示してください。
4. **初学者に配慮した言葉を使う**: 専門用語を使う場合は簡単に補足説明を加えてください。

=== kind='metaprompt' での最適化後プロンプト ===
あなたは Python 初学者向けのプログラミングアシスタントです。ユーザーの質問に丁寧に答えてください。

回答する際は、以下のガイドラインに従ってください：

1. **必ずコード例を含める**: Python の文法・メソッド・概念について説明するときは、必ず最低1つの具体的なコード例（実行可能なコードスニペット）を示してください。
2. **「なぜ使うか」を説明する**: 機能の使い方だけでなく、「なぜその機能を使うのか」「どんな場面で役立つのか」「他の方法との違い」も合わせて説明してください。
3. **初学者向けに丁寧に**: 専門用語を使う場合は簡単に補足し、ステップごとに説明するよう心がけてください。

【回答の構成例】
- 簡単な説明（何ができるか）
- コード例（コメント付きで読みやすく）
- いつ・なぜ使うかのポイント
（以下略：最適化後プロンプトを使った応答が続く）

"""多ターンの会話履歴を集約して system prompt を更新するサイクル。

1セッションで 3 ターンの会話を行い、各ターンでユーザーから自由文フィードバックを受ける。
セッション終了時に **会話全体を 1 つの trajectory にまとめ**、各ターンのフィードバックを
連結した1つの feedback として prompt_memory に渡し、同じ最初の質問を新プロンプトで
再応答して挙動の差を観察する。

このパターンは「ユーザーがその場で『こう答えてほしい』と要望を伝え、次回以降の
セッションでその要望が反映されるよう振る舞いを書き換える」という、
オンライン学習の最小単位を表す。
"""

import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)

import litellm
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_litellm import ChatLiteLLM
from langmem import create_prompt_optimizer

litellm.modify_params = True
litellm.drop_params = True

responder = ChatLiteLLM(model="openai/gpt-5-mini")
optimizer = create_prompt_optimizer(
    "anthropic:claude-sonnet-4-6",
    kind="prompt_memory",
)

INITIAL_PROMPT = (
    "あなたは献立提案アシスタントです。"
    "ユーザーの希望に沿ってレシピを提案してください。"
)

# (ユーザー質問, 直後にユーザーが投げる自由文フィードバック) の 3 ターン分。
TURNS = [
    (
        "今晩の献立を考えています。10 分で作れる和風の主菜を1つ提案してください。",
        "もう少し短く要点だけにしてほしい。所要時間と難易度を最初の1行で見せて。",
    ),
    (
        "副菜も1品ほしいです。",
        "材料は箇条書き、手順は番号付き、調理時間は最初の行に書いてほしい。",
    ),
    (
        "デザートに簡単なものはありますか？",
        "最後に栄養面の一言コメントも添えてほしい。",
    ),
]

def run_session(prompt: str, label: str) -> list[tuple]:
    """1 セッション分の会話を実行し、会話全体を 1 つの trajectory にまとめて返す。

    返り値は `[(messages, feedback)]` の形（要素 1 つだけのリスト）。
    optimizer に渡せばよいので、こうしておくとそのまま `trajectories` 引数に流し込める。
    """
    messages: list = [SystemMessage(content=prompt)]
    print(f"\n--- {label} (system prompt: {prompt[:40]}...) ---")
    # trajectory に積むのは Human/AI のみ（SystemMessage は prompt 引数で別途渡す）。
    trajectory_messages: list = []
    feedback_lines: list[str] = []
    for idx, (question, feedback) in enumerate(TURNS, start=1):
        human_maeeage = HumanMessage(content=question)
        messages.append(human_maeeage)
        response = responder.invoke(messages)
        ai_message = AIMessage(content=str(response.content))
        messages.append(ai_message)
        print(f"\n[user] {question}")
        print(f"[ai] {ai_message.content[:160]}")
        print(f"[user feedback] {feedback}")
        trajectory_messages.append(human_maeeage)
        trajectory_messages.append(ai_message)
        feedback_lines.append(f"ターン{idx}: {feedback}")
    # 1 セッション = 1 trajectory。会話全体と、ターンごとのフィードバックを連結した文字列を持たせる。
    return [
        (
            trajectory_messages,
            {"feedback": "\n".join(feedback_lines)},
        )
    ]

if __name__ == "__main__":
    print("=== 1 セッション目: 元のプロンプト ===")
    trajectories = run_session(INITIAL_PROMPT, "session-1 (initial prompt)")

    print("\n=== 会話履歴 + フィードバックを optimizer に渡す ===")
    updated_prompt = optimizer.invoke(
        {"trajectories": trajectories, "prompt": INITIAL_PROMPT}
    )
    print("\n=== 最適化後プロンプト ===")
    print(updated_prompt)

    # 同じ最初の質問を、最適化後プロンプトで再応答してみる。
    print("\n=== 同じ最初の質問を最適化後プロンプトで再応答 ===")
    response = responder.invoke(
        [
            SystemMessage(content=str(updated_prompt)),
            HumanMessage(content=TURNS[0][0]),
        ]
    )
    print(response.content)

$ python multi_turn_feedback_loop.py
=== 1 セッション目: 元のプロンプト ===

--- session-1 (initial prompt) (system prompt: あなたは献立提案アシスタントです。ユーザーの希望に沿ってレシピを提案してください...) ---

[user] 今晩の献立を考えています。10 分で作れる和風の主菜を1つ提案してください。
[ai] いいですね。手早く作れて和食らしい味の主菜を1品提案します。

豚こまの生姜焼き（2人分・約10分）

材料
- 豚こま切れ肉 200g
（中略）
[user feedback] もう少し短く要点だけにしてほしい。所要時間と難易度を最初の1行で見せて。

[user] 副菜も1品ほしいです。
[ai] では、主菜の豚こま生姜焼きに合う、手早く作れる和風の副菜を1品提案します。

ほうれん草のおひたし（2人分・約5分）
（中略）
[user feedback] 材料は箇条書き、手順は番号付き、調理時間は最初の行に書いてほしい。

[user] デザートに簡単なものはありますか？
[ai] いいですね。和風でさっぱり、しかもすぐ作れるデザートを提案します。

ヨーグルトの黒みつきなこがけ（2人分・約2〜3分）
（中略）
[user feedback] 最後に栄養面の一言コメントも添えてほしい。

=== 会話履歴 + フィードバックを optimizer に渡す ===

=== 最適化後プロンプト ===
あなたは献立提案アシスタントです。ユーザーの希望に沿ってレシピを提案してください。

## 回答フォーマットのルール
1. **最初の1行**に料理名・所要時間・難易度を必ず記載する（例：「豚こまの生姜焼き｜約10分｜難易度：★☆☆」）
2. **材料**は箇条書き（- ）で簡潔に列挙する
3. **手順**は番号付きリストで簡潔に記載する（1ステップ1〜2文以内）
4. **ポイント**は必要な場合のみ、2〜3行以内に絞る
5. **最後に栄養面の一言コメント**を添える（例：「栄養メモ：たんぱく質が豊富で、野菜も摂れるバランスの良い一品です。」）
6. 全体的に**簡潔・要点重視**で、冗長な説明は避ける

=== 同じ最初の質問を最適化後プロンプトで再応答 ===
豚しゃぶのねぎぽん酢｜約10分｜難易度：★☆☆

材料
- 豚薄切り肉（しゃぶしゃぶ用）200g
- 長ねぎ 1本
（中略）

手順
1. 長ねぎは斜め薄切りにし、ボウルでポン酢とごま油と和えておく。
2. 鍋に湯を沸かし、豚肉を一枚ずつ入れて色が変わるまで30〜60秒ほどゆでる。
（中略）

栄養メモ：たんぱく質が豊富で、ねぎの香味で食欲をそそる低カロリーな主菜です。

"""thumbs up / down フィードバックを構造化して prompt_optimizer に渡すサンプル。

実プロダクトのチャット UI でよくある「👍 / 👎 + 自由コメント」フィードバックを
trajectory に変換し、prompt_memory で system prompt を更新する。
update 後に同じ質問群を再応答して、不評だったポイントが改善されるかを確認する。
"""

import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)

import litellm
from langchain_core.messages import AIMessage, HumanMessage
from langchain_litellm import ChatLiteLLM
from langmem import create_prompt_optimizer

litellm.modify_params = True
litellm.drop_params = True

responder = ChatLiteLLM(model="openai/gpt-5-mini")
optimizer = create_prompt_optimizer(
    "anthropic:claude-sonnet-4-6",
    kind="prompt_memory",
)

INITIAL_PROMPT = (
    "あなたは技術記事の要約アシスタントです。"
    "渡された記事本文を要約してください。"
)

# 3 件の記事サンプル（要約対象）。
ARTICLES = [
    {
        "id": "article-1",
        "title": "LiteLLM 入門",
        "body": (
            "LiteLLM は OpenAI 互換のインターフェースで 100 を超える LLM プロバイダーを"
            "統一的に呼び出せる Python ライブラリです。"
            "OpenAI / Anthropic / Gemini / Bedrock / Vertex AI / Ollama などをモデル名の"
            "文字列を切り替えるだけで横断的に扱えます。"
            "ストリーミング・非同期呼び出し・フォールバック・コスト追跡といった本番運用で"
            "必要な機能も標準で備えており、アプリケーションロジックとモデル選定を疎結合に保てます。"
        ),
    },
    {
        "id": "article-2",
        "title": "LangGraph の基礎",
        "body": (
            "LangGraph は LangChain チームが提供するグラフベースのワークフロー記述ライブラリで、"
            "状態を持つエージェントを宣言的に組み立てられます。"
            "ノードとエッジを add_node / add_edge で組み、`StateGraph` で状態スキーマを定義します。"
            "Checkpointer による永続化や HITL（Human-in-the-Loop）の interrupt 機構など、"
            "本番運用に必要な機能が揃っています。"
        ),
    },
    {
        "id": "article-3",
        "title": "LangMem 概要",
        "body": (
            "LangMem は LangGraph の Store 上で動く長期メモリライブラリで、"
            "セマンティック・エピソディック・プロシージャルの 3 種類の記憶モデルを扱えます。"
            "セマンティックは会話から事実を抽出して保存、エピソディックは経験を要約して残し、"
            "プロシージャルは prompt_optimizer で system prompt を直接書き換える形で振る舞いを更新します。"
        ),
    },
]

def summarize(prompt: str, article: dict) -> str:
    response = responder.invoke(
        [
            {"role": "system", "content": prompt},
            {"role": "user", "content": f"以下の記事を要約してください。\n\nタイトル: {article['title']}\n本文: {article['body']}"},
        ]
    )
    return str(response.content).strip()

def make_trajectory(article: dict, summary: str, thumb: str, comment: str) -> tuple:
    """thumbs up/down + 自由コメントを optimizer 向けの (messages, feedback) に変換する。

    feedback dict は文字列化されてそのまま LLM optimizer に渡されるため、`user_score` は
    LangMem 内で特別扱いされるキーではなく、LLM が自然に解釈してくれる便宜上のラベル。
    数値スコアに加えて自由文の `feedback` を一緒に持たせると、低スコアの trajectory に対して
    「具体的に何を改善したいか」まで optimizer が拾いやすくなる。
    """
    return (
        [
            HumanMessage(content=f"以下の記事を要約してください。\n\nタイトル: {article['title']}\n本文: {article['body']}"),
            AIMessage(content=summary),
        ],
        {
            "user_score": 1.0 if thumb == "up" else 0.0,
            "feedback": comment,
        },
    )

# 3 件の評価（実 UI ではユーザー入力。ここではモック）。
EVALUATIONS = [
    ("up", "短く要点だけになっていて読みやすい。このトーンを継続してほしい。"),
    ("down", "要点が薄い。3 行の箇条書きで、必ず『主用途』『主要機能』『ユースケース』の3観点を入れてほしい。"),
    ("down", "文章が長すぎる。120 文字以内の箇条書き 3 点に絞ってほしい。"),
]

if __name__ == "__main__":
    # ===== 1 ラウンド目: 初期プロンプトで要約 + フィードバック収集 =====
    print("=== Round 1: 初期プロンプトでの要約 ===")
    trajectories = []
    for article, (thumb, comment) in zip(ARTICLES, EVALUATIONS, strict=True):
        summary = summarize(INITIAL_PROMPT, article)
        print(f"\n[{article['id']}] {article['title']}")
        print(f"summary: {summary[:160]}")
        print(f"thumb: {thumb} / comment: {comment}")
        trajectories.append(make_trajectory(article, summary, thumb, comment))

    # ===== prompt_optimizer で system prompt を更新 =====
    print("\n=== thumbs フィードバックを反映してプロンプトを更新 ===")
    updated_prompt = optimizer.invoke(
        {"trajectories": trajectories, "prompt": INITIAL_PROMPT}
    )
    print("\n--- 元のプロンプト ---")
    print(INITIAL_PROMPT)
    print("\n--- 更新後のプロンプト ---")
    print(updated_prompt)

    # ===== 2 ラウンド目: 更新後プロンプトで同じ記事を再要約 =====
    print("\n=== Round 2: 更新後プロンプトでの要約 ===")
    for article in ARTICLES:
        summary = summarize(str(updated_prompt), article)
        print(f"\n[{article['id']}] {article['title']}")
        print(f"summary: {summary}")

$ python thumbs_feedback.py
=== Round 1: 初期プロンプトでの要約 ===

[article-1] LiteLLM 入門
summary: LiteLLM は Python製のライブラリで、OpenAI互換の単一インターフェースから 100以上の LLMプロバイダー...（中略・thumb: up）

[article-2] LangGraph の基礎
summary: LangGraph は LangChain チームの、グラフベースで状態を持つエージェントやワークフローを宣言的に組み立てるライブラリです...（中略・thumb: down）

[article-3] LangMem 概要
summary: LangMem は LangGraph の Store 上で動く長期メモリライブラリで、以下の3種類の記憶モデルを扱います...（中略・thumb: down）

=== thumbs フィードバックを反映してプロンプトを更新 ===

--- 更新後のプロンプト ---
あなたは技術記事の要約アシスタントです。渡された記事本文を、以下のルールに従って要約してください。

【出力ルール】
- 必ず箇条書き3点で出力する
- 各点は以下の観点に対応させる
  1. 主用途（何のためのツール・技術か）
  2. 主要機能（代表的な機能・特徴）
  3. ユースケース（どんな場面で使うか）
- 全体で120文字以内に収める
- 余分な説明・まとめ文は不要

=== Round 2: 更新後プロンプトでの要約 ===

[article-1] LiteLLM 入門
summary: ・LLMプロバイダーを統一して呼べるOpenAI互換Pythonライブラリ
・ストリーミング・非同期・フォールバック・コスト追跡など本番機能
・モデル名切替で複数サービスを横断、アプリや運用に適合

[article-2] LangGraph の基礎
summary: - 主用途: 状態を持つエージェントの宣言的ワークフロー定義
- 主要機能: add_node/edgeで構築、StateGraph、Checkpointer、HITL中断
- ユースケース: 本番運用ワークフローや人手介入が必要な自動化で利用

[article-3] LangMem 概要
summary: - LangGraph上で動く長期メモリライブラリ
- セマンティック抽出・エピソード要約・プロシージャルでsystem promptを書換
- 経験要約保持・振る舞いの動的更新に利用

"""反復学習: 評価セットで自己採点しながら system prompt を磨くサンプル。

3 問の評価セットに対して、
  1. 現在の system prompt で各問に応答
  2. LLM-as-judge で 0-1 のスコアを採点（rubric を渡す）
  3. スコアが低い trajectory を集めて prompt_memory で system prompt を更新
を 3 ラウンド繰り返し、ラウンドごとの平均スコアの推移を観察する。

オンライン学習の最も「学習らしい」形で、評価データセットさえあれば
人間のフィードバック無しでも prompt が継続的に改善されていく。
"""

import json
import re
import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)

import litellm
from langchain_core.messages import AIMessage, HumanMessage
from langchain_litellm import ChatLiteLLM
from langmem import create_prompt_optimizer

litellm.modify_params = True
litellm.drop_params = True

# 応答役は安価な GPT、judge と optimizer は文章評価が得意な Claude。
responder = ChatLiteLLM(model="openai/gpt-5-mini")
judge = ChatLiteLLM(model="anthropic/claude-sonnet-4-6")
optimizer = create_prompt_optimizer(
    "anthropic:claude-sonnet-4-6",
    kind="prompt_memory",
)

# 評価セット: SaaS のサポート Q&A 想定。
# 各 rubric は 4 条件で構成し、満たさない条件 1 つにつき 0.25 ずつ減点する形にする。
# (4) を「質問と無関係な情報を含めない」とすることで、過剰一般化（質問と関係ない案内文の混入）も
# 採点で捕捉できるようにしている。
EVAL_SET = [
    {
        "question": "請求書の支払い期日が過ぎてしまいました。どうすればいいですか？",
        "rubric": (
            "回答が以下を全て満たすか採点せよ："
            "(1) 支払い遅延時の連絡先（サポートメール support@example.com）を案内している、"
            "(2) 期日超過後の延滞料金の有無に触れている、"
            "(3) 200 文字以内に収まっている、"
            "(4) 質問と無関係な情報（質問されていないトピック）を含めていない。"
            "全て満たすなら 1.0、欠けるごとに 0.25 ずつ減点せよ。"
        ),
    },
    {
        "question": "プランを Pro から Team にアップグレードしたいです。差額はどうなりますか？",
        "rubric": (
            "回答が以下を全て満たすか採点せよ："
            "(1) 日割り計算で差額が請求される旨を述べている、"
            "(2) アップグレードが即時反映される旨を述べている、"
            "(3) 200 文字以内に収まっている、"
            "(4) 質問と無関係な情報（質問されていないトピック）を含めていない。"
            "全て満たすなら 1.0、欠けるごとに 0.25 ずつ減点せよ。"
        ),
    },
    {
        "question": "API キーを失くしてしまいました。再発行できますか？",
        "rubric": (
            "回答が以下を全て満たすか採点せよ："
            "(1) ダッシュボードの「API キー」画面から再発行できる旨を案内している、"
            "(2) 古いキーは即時無効化される注意を入れている、"
            "(3) 200 文字以内に収まっている、"
            "(4) 質問と無関係な情報（質問されていないトピック）を含めていない。"
            "全て満たすなら 1.0、欠けるごとに 0.25 ずつ減点せよ。"
        ),
    },
]

INITIAL_PROMPT = (
    "あなたは SaaS 製品のサポート担当です。ユーザーの質問に答えてください。"
)

ROUNDS = 3

def answer(prompt: str, question: str) -> str:
    response = responder.invoke(
        [
            {"role": "system", "content": prompt},
            {"role": "user", "content": question},
        ]
    )
    return str(response.content).strip()

def score(question: str, response: str, rubric: str) -> tuple[float, str]:
    """LLM-as-judge で 0.0〜1.0 のスコアと講評を返す。"""
    judge_prompt = (
        "あなたは厳格な採点者です。以下のルーブリックに従って、回答を 0.0〜1.0 で採点してください。"
        "出力は必ず以下の JSON 形式のみ:"
        ' {"score": <number>, "comment": "<短い講評>"}\n\n'
        f"ルーブリック: {rubric}\n\n"
        f"質問: {question}\n\n"
        f"回答: {response}"
    )
    raw = str(judge.invoke([{"role": "user", "content": judge_prompt}]).content).strip()
    # ```json で囲まれた応答や前置き文に対応するため、最後に出てくる { ... } を取り出す。
    # `\{.*\}` は欲張りマッチなので、judge が複数 JSON を返してもブロック全体を拾える一方、
    # 不正な JSON だった場合のフォールバックも兼ねる。
    payload: dict = {"score": 0.0, "comment": "parse error"}
    m = re.search(r"\{.*\}", raw, re.DOTALL)
    if m:
        try:
            payload = json.loads(m.group())
        except json.JSONDecodeError:
            payload = {"score": 0.0, "comment": f"parse error: {raw[:80]}"}
    return float(payload.get("score", 0.0)), str(payload.get("comment", ""))

def run_round(prompt: str, round_idx: int) -> tuple[float, list[tuple]]:
    """1 ラウンド = 評価セット全件に応答し、スコアとフィードバック付き trajectory を返す。"""
    print(f"\n===== Round {round_idx} =====")
    trajectories: list[tuple] = []
    scores: list[float] = []
    for item in EVAL_SET:
        response = answer(prompt, item["question"])
        s, comment = score(item["question"], response, item["rubric"])
        scores.append(s)
        print(f"\nQ: {item['question']}")
        print(f"A: {response[:120]}")
        print(f"score={s:.2f} / comment={comment}")
        # 採点 < 1.0 のものだけ trajectory として optimizer に渡す
        # （満点の事例は prompt 改善ヒントを含まないため）。
        if s < 1.0:
            trajectories.append(
                (
                    [
                        HumanMessage(content=item["question"]),
                        AIMessage(content=response),
                    ],
                    {
                        "user_score": s,
                        "rubric": item["rubric"],
                        "feedback": f"judge コメント: {comment}",
                    },
                )
            )
    avg = sum(scores) / len(scores)
    print(f"\n[Round {round_idx}] 平均スコア: {avg:.3f}")
    return avg, trajectories

if __name__ == "__main__":
    prompt = INITIAL_PROMPT
    history: list[tuple[int, float]] = []
    for r in range(1, ROUNDS + 1):
        avg, trajectories = run_round(prompt, r)
        history.append((r, avg))
        # 最終ラウンドは最適化しない（次に使う場面が無いため）。
        if r == ROUNDS:
            break
        if not trajectories:
            print("全問満点だったため最適化スキップ")
            continue
        prompt = str(
            optimizer.invoke({"trajectories": trajectories, "prompt": prompt})
        )
        print(f"\n--- Round {r} 終了時に更新したプロンプト ---")
        print(prompt)

    print("\n===== ラウンド別平均スコア =====")
    for r, avg in history:
        bar = "#" * int(avg * 20)
        print(f"Round {r}: {avg:.3f} {bar}")

$ python iterative_self_improvement.py
===== Round 1 =====

Q: 請求書の支払い期日が過ぎてしまいました。どうすればいいですか？
A: ご連絡ありがとうございます。支払期日を過ぎた場合の対応を分かりやすくまとめます。まずは落ち着いて、以下を順にご確認ください。
（中略）
score=0.50 / comment=延滞料金への言及(2)はあるが、サポートメール support@example.com の案内(1)がなく、200文字以内(3)も大幅に超過しており、質問と無関係な情報（自動支払い設定・分割払い相談・再発防止策など）(4)も多数含まれている。欠けている要件が3つのため0.25×3=0.75減点。

Q: プランを Pro から Team にアップグレードしたいです。差額はどうなりますか？
A: 短くお答えします：一般的には「未使用期間に応じた差額を日割りで計算して、アップグレード時に差額分を即時請求」する形になります。
（中略）
score=0.50 / comment=（1）日割り計算で差額が請求される旨と（2）即時反映の旨は述べられている。しかし（3）回答が200文字をはるかに超えており、（4）税金・通貨・年払いの詳細・アカウント操作手順・メールアドレス確認依頼など質問と無関係な情報が多数含まれている。

Q: API キーを失くしてしまいました。再発行できますか？
A: はい、可能です。状況に応じて対応方法が2パターンあります。まずダッシュボードにまだログインできるか教えてください。
（中略）
score=0.50 / comment=ダッシュボードの「APIキー」画面から再発行できる旨(1)は満たしている。古いキーの即時無効化の注意(2)も一部言及あり。しかし200文字以内(3)を大幅に超過しており、安全対策やサポート連絡手順など質問と無関係な情報(4)も多数含まれている。

[Round 1] 平均スコア: 0.500

--- Round 1 終了時に更新したプロンプト ---
あなたは SaaS 製品のサポート担当です。以下のルールを厳守してください。

【回答ルール】
1. 回答は必ず200文字以内に収めること。
2. ユーザーが質問した内容のみ回答すること。質問されていないトピック（再発防止策、他の機能説明、操作手順の詳細など）は含めないこと。
3. サポートへの連絡が必要な場合は support@example.com を案内すること。
4. 簡潔・明確に、要点だけを伝えること。

===== Round 2 =====

Q: 請求書の支払い期日が過ぎてしまいました。どうすればいいですか？
A: まず速やかに請求額をお支払いください。支払いが難しい、または請求内容に疑問がある場合は support@example.com へご連絡ください。
score=0.50 / comment=サポートメール案内(1)・200文字以内(3)・無関係情報なし(4)は満たすが、延滞料金の有無(2)に一切触れていない。さらに「請求内容に疑問がある場合」という質問外の案内(4)も加わったため減点。

Q: プランを Pro から Team にアップグレードしたいです。差額はどうなりますか？
A: アップグレード時の差額は請求サイクルや残日数で変わります。正確な金額はここでは確認できないため、アカウント情報を添えて support@example.com 宛にご連絡ください。
score=0.50 / comment=「日割り計算で差額が請求される」(1)と「即時反映される」(2)が明示されておらず、サポートに丸投げしている。要件(1)(2)が欠けるため減点。

Q: API キーを失くしてしまいました。再発行できますか？
A: はい、再発行可能です。再発行をご希望の場合は support@example.com までご連絡ください。
score=0.25 / comment=再発行可能である旨は伝えているが、(1)ダッシュボードの「APIキー」画面からの案内がなく、(2)古いキーの即時無効化の注意もなく、(4)質問に無関係なサポートメール情報が含まれている。3項目が欠落。

[Round 2] 平均スコア: 0.417

--- Round 2 終了時に更新したプロンプト ---
あなたは SaaS 製品のサポート担当です。以下のルールと製品知識を厳守してください。

【回答ルール】
1. 回答は必ず200文字以内に収めること。
2. ユーザーが質問した内容のみ回答すること。質問されていないトピック（再発防止策、他の機能説明、操作手順の詳細など）は含めないこと。
3. サポートへの連絡が必要な場合のみ support@example.com を案内すること。自己解決できる場合はサポートメールを案内しないこと。
4. 簡潔・明確に、要点だけを伝えること。

【製品知識】
- 請求・支払い：支払い期日を過ぎると延滞料金が発生する場合があります。支払いに関する問い合わせは support@example.com へ。
- プランアップグレード：差額は日割り計算で請求され、アップグレードは即時反映されます。
- API キー：ダッシュボードの「API キー」画面から再発行できます。再発行すると古いキーは即時無効化されます。

===== Round 3 =====

Q: 請求書の支払い期日が過ぎてしまいました。どうすればいいですか？
A: 支払い期日を過ぎると延滞料金が発生する場合があります。支払方法や延滞料金の確認・対応については support@example.com までご連絡ください。
score=1.00 / comment=連絡先（support@example.com）の案内、延滞料金への言及、200文字以内、無関係な情報なし、の4条件を全て満たしています。

Q: プランを Pro から Team にアップグレードしたいです。差額はどうなりますか？
A: アップグレード時の差額は日割り計算で請求されます。アップグレードは即時反映されます。
score=1.00 / comment=4つの条件を全て満たしている。日割り計算・即時反映の両方を明記し、200文字以内（41文字）に収まり、無関係な情報も含まれていない。

Q: API キーを失くしてしまいました。再発行できますか？
A: はい。ダッシュボードの「API キー」画面から再発行できます。再発行すると旧キーは即時無効化されます。
score=1.00 / comment=4つの条件をすべて満たしています。「APIキー」画面からの再発行案内(1)、旧キーの即時無効化の注意(2)、文字数は約50文字で200文字以内(3)、余分な情報なし(4)。簡潔で的確な回答です。

[Round 3] 平均スコア: 1.000

===== ラウンド別平均スコア =====
Round 1: 0.500 ##########
Round 2: 0.417 ########
Round 3: 1.000 ####################