[アップデート] Amazon Bedrock AgentCore RuntimeがStateful MCPサーバーに対応しました

"""
国内旅行プランナー - Stateful MCP Server
Elicitation / Sampling / Progress Notifications のデモ
"""
import asyncio
import logging
from fastmcp import FastMCP, Context
from enum import Enum

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("travel-planner")

mcp = FastMCP("Japan-Travel-Planner")

class TripType(str, Enum):
    SOLO = "ひとり旅"
    COUPLE = "カップル"
    FAMILY = "家族旅行"
    FRIENDS = "友人旅行"

DESTINATIONS = {
    "京都": {
        "name": "京都",
        "transport": 14000,
        "hotel": 12000,
        "highlights": ["伏見稲荷大社", "嵐山竹林", "清水寺"],
        "gourmet": ["湯豆腐", "抹茶スイーツ", "にしんそば"],
    },
    "沖縄": {
        "name": "沖縄",
        "transport": 40000,
        "hotel": 10000,
        "highlights": ["美ら海水族館", "古宇利島", "首里城"],
        "gourmet": ["ソーキそば", "タコライス", "海ぶどう"],
    },
    "北海道": {
        "name": "北海道",
        "transport": 35000,
        "hotel": 9000,
        "highlights": ["富良野ラベンダー畑", "小樽運河", "旭山動物園"],
        "gourmet": ["海鮮丼", "ジンギスカン", "スープカレー"],
    },
    "福岡": {
        "name": "福岡",
        "transport": 22000,
        "hotel": 8000,
        "highlights": ["太宰府天満宮", "中洲屋台", "糸島"],
        "gourmet": ["博多ラーメン", "もつ鍋", "明太子"],
    },
}

@mcp.tool()
async def plan_trip(ctx: Context) -> str:
    """
    国内旅行プランを作成（全MCP機能を使用）:
    1. Elicitation - 旅行の希望をヒアリング
    2. Progress - 交通手段・宿泊先の検索進捗
    3. Sampling - AIによるおすすめ情報の生成
    """
    # ---- Phase 1: Elicitation ----
    logger.info("[Phase 1] Elicitation 開始 - 旅行先の質問")
    dest_result = await ctx.elicit(
        message="どこに行きたいですか？\n選択肢: 京都、沖縄、北海道、福岡",
        response_type=str,
    )
    if dest_result.action != "accept":
        logger.info("[Phase 1] ユーザーがキャンセル（旅行先）")
        return "プラン作成をキャンセルしました。"
    dest_key = dest_result.data.strip()
    dest = DESTINATIONS.get(dest_key, DESTINATIONS["京都"])
    logger.info(f"[Phase 1] 旅行先: {dest['name']}")

    type_result = await ctx.elicit(
        message="旅行のタイプは？\n1. ひとり旅\n2. カップル\n3. 家族旅行\n4. 友人旅行",
        response_type=TripType,
    )
    if type_result.action != "accept":
        logger.info("[Phase 1] ユーザーがキャンセル（旅行タイプ）")
        return "プラン作成をキャンセルしました。"
    trip_type = type_result.data.value if hasattr(type_result.data, "value") else type_result.data
    logger.info(f"[Phase 1] 旅行タイプ: {trip_type}")

    days_result = await ctx.elicit(
        message="何泊しますか？（1〜7泊）",
        response_type=int,
    )
    if days_result.action != "accept":
        return "プラン作成をキャンセルしました。"
    days = max(1, min(7, days_result.data))

    travelers_result = await ctx.elicit(
        message="何人で行きますか？",
        response_type=int,
    )
    if travelers_result.action != "accept":
        return "プラン作成をキャンセルしました。"
    travelers = travelers_result.data
    logger.info(f"[Phase 1] Elicitation 完了 - {dest['name']}/{trip_type}/{days}泊/{travelers}人")

    # ---- Phase 2: Progress Notifications ----
    logger.info("[Phase 2] Progress Notifications 開始 - 検索処理")
    total_steps = 5
    for step in range(1, total_steps + 1):
        await ctx.report_progress(progress=step, total=total_steps)
        await asyncio.sleep(0.4)

    transport_cost = dest["transport"] * travelers
    hotel_cost = dest["hotel"] * days * ((travelers + 1) // 2)
    total_cost = transport_cost + hotel_cost
    logger.info(f"[Phase 2] Progress 完了 - 費用算出: ¥{total_cost:,}")

    # ---- Phase 3: Sampling ----
    logger.info("[Phase 3] Sampling 開始 - クライアントにAI推論を要求")
    ai_tips = f"{dest['name']}を楽しんでください！"
    try:
        response = await ctx.sample(
            messages=(
                f"{dest['name']}への{trip_type}（{travelers}人、{days}泊）のおすすめを"
                f"3つ、60文字以内で簡潔に教えてください。日本語で回答してください。"
            ),
            max_tokens=200,
        )
        if hasattr(response, "text") and response.text:
            ai_tips = response.text
        logger.info(f"[Phase 3] Sampling 完了 - 応答: {ai_tips[:50]}...")
    except Exception as e:
        logger.warning(f"[Phase 3] Sampling 失敗: {e}")
        ai_tips = f"{dest['highlights'][0]}は必見です。{dest['gourmet'][0]}もぜひ！"

    # ---- Final Confirmation ----
    logger.info("[Phase 4] 最終確認 Elicitation")
    confirm = await ctx.elicit(
        message=f"""
========== 旅行プラン概要 ==========
旅行先: {dest['name']}
タイプ: {trip_type}
日程: {days}泊{days + 1}日
人数: {travelers}人

費用概算:
  交通費: ¥{transport_cost:,}
  宿泊費: ¥{hotel_cost:,}（{(travelers + 1) // 2}部屋 × {days}泊）
  合計: ¥{total_cost:,}

この内容で予約しますか？""",
        response_type=["予約する", "キャンセル"],
    )
    if confirm.action != "accept" or confirm.data == "キャンセル":
        logger.info("[Phase 4] ユーザーがキャンセル")
        return "予約をキャンセルしました。検索結果は24時間保存されます。"

    logger.info(f"[Phase 4] 予約確定 - {dest['name']}/{trip_type}/{days}泊/{travelers}人/¥{total_cost:,}")
    highlights_str = "\n".join(f"  - {h}" for h in dest["highlights"])
    gourmet_str = "\n".join(f"  - {g}" for g in dest["gourmet"])
    return f"""
{'=' * 50}
予約が確定しました！
{'=' * 50}
予約番号: TRV-{ctx.session_id[:8].upper()}

旅行先: {dest['name']}
日程: {days}泊{days + 1}日 ／ {travelers}人
タイプ: {trip_type}

交通費: ¥{transport_cost:,}
宿泊費: ¥{hotel_cost:,}（{(travelers + 1) // 2}部屋 × {days}泊）
合計: ¥{total_cost:,}

おすすめスポット:
{highlights_str}

ご当地グルメ:
{gourmet_str}

AIからのおすすめ:
{ai_tips}
{'=' * 50}
"""

if __name__ == "__main__":
    mcp.run(
        transport="streamable-http",
        host="0.0.0.0",
        port=8000,
        stateless_http=False,
    )

"""
国内旅行プランナー - テストクライアント
Elicitation / Sampling / Progress のハンドラを実装

環境変数:
  LOCAL_TEST=true (default) : ローカルサーバーに接続
  LOCAL_TEST=false          : AgentCore Runtime に接続（AGENT_ARN, BEARER_TOKEN 必須）
  USE_BEDROCK=true          : Sampling に Amazon Bedrock を使用
  BEDROCK_MODEL_ID          : Bedrock モデルID（デフォルト: us.anthropic.claude-haiku-4-5-20251001-v1:0）
"""
import asyncio
import os
import sys
import typing
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
from fastmcp.client.elicitation import ElicitResult
from mcp.types import CreateMessageResult, TextContent

# --- Elicitation ハンドラ ---

def _extract_options(response_type) -> list[str] | None:
    """response_type から選択肢リストを取り出す。
    fastmcp は ["予約する", "キャンセル"] を Literal["予約する", "キャンセル"] に変換して渡してくる。
    """
    if isinstance(response_type, list):
        return response_type
    if typing.get_origin(response_type) is typing.Literal:
        args = typing.get_args(response_type)
        if args:
            return list(args)
    return None

def _prompt_choice(options: list[str]) -> str:
    """番号で選択させる。不正入力はリトライ。"""
    for i, opt in enumerate(options, 1):
        print(f"    {i}. {opt}")
    while True:
        raw = input("    番号を選択: ").strip()
        try:
            idx = int(raw)
            if 1 <= idx <= len(options):
                return options[idx - 1]
        except ValueError:
            pass
        print(f"    ※ 1〜{len(options)} の番号を入力してください")

def _prompt_int(label: str = "回答（数字）") -> int:
    """整数を入力させる。不正入力はリトライ。"""
    while True:
        raw = input(f"    {label}: ").strip()
        try:
            return int(raw)
        except ValueError:
            print("    ※ 数字を入力してください")

async def elicit_handler(message, response_type, params, ctx):
    print(f"\n>>> サーバーからの質問: {message}")
    try:
        options = _extract_options(response_type)
        if options:
            response = _prompt_choice(options)
        elif response_type is int:
            response = _prompt_int()
        else:
            response = input("    回答: ").strip()
        return ElicitResult(action="accept", content={"value": response})
    except (KeyboardInterrupt, EOFError):
        print("\n    (キャンセルされました)")
        return ElicitResult(action="decline", content=None)

# --- Sampling ハンドラ ---

def _extract_prompt_text(messages) -> str:
    """SamplingMessage のリストからプロンプトテキストを抽出"""
    if isinstance(messages, str):
        return messages
    if isinstance(messages, list):
        parts = []
        for msg in messages:
            if hasattr(msg, "content") and hasattr(msg.content, "text"):
                parts.append(msg.content.text)
            else:
                parts.append(str(msg))
        return "\n".join(parts)
    return str(messages)

def _invoke_bedrock(prompt: str) -> str:
    """Amazon Bedrock で推論を実行"""
    import boto3
    model_id = os.getenv("BEDROCK_MODEL_ID", "us.anthropic.claude-haiku-4-5-20251001-v1:0")
    client = boto3.client("bedrock-runtime")
    response = client.converse(
        modelId=model_id,
        messages=[{"role": "user", "content": [{"text": prompt}]}],
        inferenceConfig={"maxTokens": 200},
    )
    return response["output"]["message"]["content"][0]["text"]

async def sampling_handler(messages, params, ctx):
    prompt_text = _extract_prompt_text(messages)
    print(f"\n>>> AIサンプリングリクエスト")
    print(f"    プロンプト: {prompt_text[:120]}...")

    use_bedrock = os.getenv("USE_BEDROCK", "false").lower() == "true"
    try:
        if use_bedrock:
            print("    (Bedrock で推論中...)")
            ai_response = _invoke_bedrock(prompt_text)
            print(f"    Bedrock応答: {ai_response}")
        else:
            ai_response = input("    AI応答を入力（Enterでデフォルト）: ").strip()
            if not ai_response:
                ai_response = "1. 人気スポットは朝一で 2. ご当地グルメを堪能 3. 地元の人おすすめの穴場へ"
    except Exception as e:
        print(f"    ※ Sampling エラー: {e}")
        ai_response = "おすすめ情報を取得できませんでした。"

    return CreateMessageResult(
        role="assistant",
        content=TextContent(type="text", text=ai_response),
        model="bedrock" if use_bedrock else "manual",
        stopReason="endTurn",
    )

# --- Progress ハンドラ ---

async def progress_handler(progress, total, message):
    pct = int((progress / total) * 100) if total else 0
    bar = "#" * (pct // 5) + "-" * (20 - pct // 5)
    print(f"\r    進捗: [{bar}] {pct}%", end="", flush=True)
    if progress == total:
        print(" 完了！")

# --- メイン ---

async def main():
    local_test = os.getenv("LOCAL_TEST", "true").lower() == "true"

    if local_test:
        url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8000/mcp"
        headers = {}
    else:
        agent_arn = os.getenv("AGENT_ARN")
        token = os.getenv("BEARER_TOKEN")
        if not agent_arn or not token:
            print("ERROR: AGENT_ARN / BEARER_TOKEN が未設定です")
            sys.exit(1)
        encoded_arn = agent_arn.replace(":", "%3A").replace("/", "%2F")
        endpoint = os.getenv(
            "MCP_ENDPOINT",
            f"https://bedrock-agentcore.{os.getenv('AWS_REGION', 'us-west-2')}.amazonaws.com",
        )
        url = f"{endpoint}/runtimes/{encoded_arn}/invocations?qualifier=DEFAULT"
        headers = {"Authorization": f"Bearer {token}"}

    transport = StreamableHttpTransport(url=url, headers=headers)
    client = Client(
        transport,
        elicitation_handler=elicit_handler,
        sampling_handler=sampling_handler,
        progress_handler=progress_handler,
    )

    try:
        await client.__aenter__()
    except Exception as e:
        print(f"\nERROR: サーバーへの接続に失敗しました: {e}")
        sys.exit(1)

    try:
        print("\nplan_trip ツールのテスト...")
        print("（Elicitation → Progress → Sampling の全フローを実行）\n")
        result = await client.call_tool("plan_trip", {})
        print("\n" + "=" * 60)
        print("結果:")
        print("=" * 60)
        print(result.content[0].text)
    except KeyboardInterrupt:
        print("\n\n中断されました")
    except Exception as e:
        print(f"\nERROR: {e}")
    finally:
        try:
            await client.__aexit__(None, None, None)
        except Exception:
            pass

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