[アップデート]Amazon Bedrock AgentCore Runtime の双方向ストリーミングを試してみた

wss://bedrock-agentcore.<region>.amazonaws.com/runtimes/<EncodedRuntimeArn>/ws

# プロジェクトの作成と初期化
uv init agentcore-websocket-strands
cd agentcore-websocket-strands

# 必要なパッケージのインストール
uv add bedrock-agentcore strands-agents bedrock-agentcore-starter-toolkit

import asyncio
import json
from bedrock_agentcore import BedrockAgentCoreApp
from strands import Agent
from strands.models.bedrock import BedrockModel
from starlette.websockets import WebSocketDisconnect

app = BedrockAgentCoreApp()

# Bedrock モデルの設定（Claude Haiku 4.5）
model = BedrockModel(
    model_id="us.anthropic.claude-haiku-4-5-20251001-v1:0",
    region_name="us-west-2"
)

# Strands エージェントの初期化
agent = Agent(
    model=model,
    system_prompt="""あなたは親切で知識豊富なAIアシスタントです。
ユーザーの質問に対して、簡潔かつ分かりやすく回答してください。
日本語で回答してください。""",
    callback_handler=None
)

@app.websocket
async def websocket_handler(websocket, context):
    """割り込み対応の WebSocket ハンドラー"""
    await websocket.accept()
    print("WebSocket 接続を受け入れました")

    # 現在実行中のストリーミングタスク
    current_task: asyncio.Task | None = None

    async def stream_response(user_message: str):
        """ストリーミングレスポンスを送信（キャンセル可能）"""
        full_response = ""
        try:
            async for event in agent.stream_async(user_message):
                if "data" in event:
                    chunk = event["data"]
                    full_response += chunk
                    await websocket.send_json({
                        "type": "chunk",
                        "data": chunk
                    })

            # 正常完了
            await websocket.send_json({
                "type": "complete",
                "fullResponse": full_response
            })
            print("レスポンス完了")

        except asyncio.CancelledError:
            # 割り込みによるキャンセル
            await websocket.send_json({
                "type": "interrupted",
                "partialResponse": full_response
            })
            print("割り込みでキャンセルされました")
            raise  # CancelledError は必ず再送出

    try:
        while True:
            # クライアントからのメッセージを受信
            raw_data = await websocket.receive_text()
            data = json.loads(raw_data)
            user_message = data.get("inputText", "")

            if not user_message:
                await websocket.send_json({"error": "inputText が見つかりません"})
                continue

            print(f"受信: {user_message}")

            # 既存のタスクがあればキャンセル（割り込み）
            if current_task and not current_task.done():
                print("割り込み: 現在の回答をキャンセルします")
                current_task.cancel()
                try:
                    await current_task  # キャンセル完了を待機
                except asyncio.CancelledError:
                    pass  # 期待通りの例外

            # 新しいストリーミングタスクをバックグラウンドで開始
            current_task = asyncio.create_task(stream_response(user_message))

    except WebSocketDisconnect:
        print("クライアントが切断しました")
    except Exception as e:
        print(f"エラー: {e}")
        try:
            await websocket.send_json({"error": str(e)})
        except Exception:
            pass
    finally:
        # クリーンアップ: 残存タスクをキャンセル
        if current_task and not current_task.done():
            current_task.cancel()
            try:
                await current_task
            except asyncio.CancelledError:
                pass

if __name__ == "__main__":
    app.run(log_level="info")

uv run websocket_strands_agent.py

INFO:     Started server process [xxxxx]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)

import asyncio
import websockets
import json

async def test_interruption():
    """割り込みテスト用クライアント"""
    uri = "ws://localhost:8080/ws"

    async with websockets.connect(uri) as websocket:

        async def receive_responses():
            """レスポンス受信ループ"""
            while True:
                response = await websocket.recv()
                data = json.loads(response)

                if data.get("type") == "chunk":
                    print(data["data"], end="", flush=True)
                elif data.get("type") == "complete":
                    print("\n\n[完了]")
                elif data.get("type") == "interrupted":
                    print(f"\n\n[割り込み中断] 途中経過: {data['partialResponse'][:50]}...")
                elif "error" in data:
                    print(f"\n[エラー] {data['error']}")

        # 受信タスクをバックグラウンドで開始
        receive_task = asyncio.create_task(receive_responses())

        try:
            # 最初の質問を送信（長い回答が返ってくる質問）
            print("送信: 日本の歴史について詳しく教えてください")
            print("-" * 50)
            await websocket.send(json.dumps({
                "inputText": "日本の歴史について詳しく教えてください"
            }))

            # 2秒後に割り込み
            await asyncio.sleep(2)

            print("\n" + "=" * 50)
            print(">>> 割り込み送信: やっぱりPythonについて教えて")
            print("=" * 50 + "\n")
            await websocket.send(json.dumps({
                "inputText": "やっぱりPythonについて教えて"
            }))

            # 回答完了を待つ（サンプルのため固定時間で待機）
            await asyncio.sleep(15)

        finally:
            receive_task.cancel()
            try:
                await receive_task
            except asyncio.CancelledError:
                pass

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

uv run websocket_client_local.py

送信: 日本の歴史について詳しく教えてください
--------------------------------------------------
# 日本の歴史について

日本の歴史を主要な時代別にご説明します。

## 主な時代区分

**古代**
- 縄文時代（約12,000年前～）：狩猟採集社会
- 弥生時代（紀元前3世紀～3世紀）：稲作農業の開始
- 古墳時代～飛鳥時代：国家形
==================================================
>>> 割り込み送信: やっぱりPythonについて教えて
==================================================

[割り込み中断] 途中経過: # 日本の歴史について

日本の歴史を主要な時代別にご説明します。

## 主な時代区分

**古代...

uv run agentcore configure -e websocket_strands_agent.py

uv run agentcore launch

arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/websocket_strands_agent-abc123

export AGENT_ARN="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/websocket_strands_agent-abc123"

import asyncio
import json
import os

import websockets
from bedrock_agentcore.runtime import AgentCoreRuntimeClient

async def main():
    # 環境変数から ARN を取得
    runtime_arn = os.getenv("AGENT_ARN")
    if not runtime_arn:
        raise ValueError("AGENT_ARN 環境変数が設定されていません")

    # AgentCoreRuntimeClient の初期化
    client = AgentCoreRuntimeClient(region="us-west-2")

    # 接続URL、ヘッダーを作成
    ws_url, headers = client.generate_ws_connection(
        runtime_arn=runtime_arn, session_id=None
    )

    print("AgentCore Runtime に接続します...\n")

    try:
        async with websockets.connect(ws_url, additional_headers=headers) as ws:

            async def receive_responses():
                """レスポンス受信ループ"""
                while True:
                    response = await ws.recv()
                    data = json.loads(response)

                    if data.get("type") == "chunk":
                        print(data["data"], end="", flush=True)
                    elif data.get("type") == "complete":
                        print("\n\n[完了]")
                    elif data.get("type") == "interrupted":
                        print(
                            f"\n\n[割り込み中断] 途中経過: {data['partialResponse'][:50]}..."
                        )
                    elif "error" in data:
                        print(f"\n[エラー] {data['error']}")

            # 受信タスクをバックグラウンドで開始
            receive_task = asyncio.create_task(receive_responses())

            try:
                # 最初の質問を送信
                print("送信: AWSの主要サービスについて詳しく教えてください")
                print("-" * 50)
                await ws.send(
                    json.dumps(
                        {"inputText": "AWSの主要サービスについて詳しく教えてください"}
                    )
                )

                # 2秒後に割り込み
                await asyncio.sleep(2)

                print("\n" + "=" * 50)
                print(">>> 割り込み送信: やっぱりEC2だけ教えて")
                print("=" * 50 + "\n")
                await ws.send(json.dumps({"inputText": "やっぱりEC2だけ教えて"}))

                # 回答完了を待つ（サンプルのため固定時間で待機）
                await asyncio.sleep(15)

            finally:
                receive_task.cancel()
                try:
                    await receive_task
                except asyncio.CancelledError:
                    pass

    except websockets.exceptions.InvalidStatus as e:
        print(f"WebSocket ハンドシェイクに失敗しました: {e.response.status_code}")
    except Exception as e:
        print(f"接続エラー: {e}")

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

uv run websocket_client_deployed.py

uv run websocket_client_deployed.py
AgentCore Runtime に接続します...

送信: AWSの主要サービスについて詳しく教えてください
--------------------------------------------------
# AWSの主要サービス

AWSには200以上のサービスがありますが、主要なものを以下に整理しました。

## **コンピューティング**
- **EC2** - 仮想サーバー
- **Lambda** - サーバーレスコンピューティング
- **ECS/EKS** - コンテナオーケストレーション

## **ストレージ**
- **S3** - オブジェクトストレージ（最も利用されるサービス）
- **EBS** - ブロックストレージ
- **EFS** - ファイルストレージ

## **データベース**
- **RDS** - 関係型データベース（MySQL、PostgreSQL等
==================================================
>>> 割り込み送信: やっぱりEC2だけ教えて
==================================================

）
- **DynamoDB** - NoSQLデータベース
- **ElastiCache** - キャッシュサービス

## **ネットワーク・

[割り込み中断] 途中経過: # AWSの主要サービス

AWSには200以上のサービスがありますが、主要なものを以下に整理しまし...
# AWS EC2について

## EC2とは
**Elastic Compute Cloud** の略で、AWSのクラウド上で仮想サーバーを利用できるサービスです。

## 主な特徴

| 項目 | 内容 |
|------|------|
| **スケーラビリティ** | 必要に応じてサーバー数を増減できる |
| **料金体系** | 使用した分だけ支払う（従量課金） |
| **OS選択** | Linux、Windows等から選択可能 |
| **インスタンスタイプ** | 用途に応じて性能を選べる |

## よく使われるインスタンスタイプ
- **t2/t3** - 小～中規模アプリケーション
- **m5/m6** - 汎用サーバー
- **c5/c6** - 計算集約型処理
- **g4/p3** - GPU利用

## 基本的な流れ
1. インスタンスタイプを選択
2. インスタンスを起動
3. セキュリティグループ設定
4. アプリケーションをデプロイ
5. 必要に応じてスケール

シンプルに始められるので、クラウド初心者にもおすすめです。

何か具体的に知りたい点はありますか？

[完了]