AmazonBedrockのpromptFlowをSDKを使っていじってみた

このチュートリアルでは、Amazon Bedrock Agent FlowとそのエイリアスをPythonで操作するためのライブラリの実装と実践的な使い方を説明します。

## 前提条件

このライブラリを使用する前に、以下の準備が必要です：

1. **Flowの作成**: Amazon Bedrock Consoleで事前にAgentとFlowを作成しておく必要があります。
2. **Flowのバージョンの作成**: Flowの変更を保存する際に新しいバージョンが作成されます。
3. **エイリアスの登録**: 作成したFlowのバージョンに対してエイリアスを登録する必要があります。

これらの作業はAWS Management ConsoleまたはAWS CLIを使用して行います。

## 1. ライブラリのコード

以下のコードを`bedrock_agent_library.py`として保存します。

```python
import boto3
from botocore.exceptions import ClientError
from datetime import datetime

class BedrockAgentFlow:
    def __init__(self, region='us-east-1'):
        self.region = region
        self.client_runtime = boto3.client("bedrock-agent-runtime", region_name=self.region)
        self.client_agent = boto3.client('bedrock-agent', region_name=self.region)

    def invoke_flow(self, message, flow_identifier, flow_alias_identifier):
        response = self.client_runtime.invoke_flow(
            flowIdentifier=flow_identifier,
            flowAliasIdentifier=flow_alias_identifier,
            inputs=[
                {
                    "content": {
                        "document": {
                            "message": message
                        }
                    },
                    "nodeName": "FlowInputNode",
                    "nodeOutputName": "document",
                }
            ],
        )
        
        print("Full response:", response)
        
        event_stream = response["responseStream"]
        generated_text = ""
        for event in event_stream:
            print("Event:", event)
            if 'flowOutputEvent' in event:
                output_event = event['flowOutputEvent']
                if 'content' in output_event and 'document' in output_event['content']:
                    generated_text += output_event['content']['document']
        
        if not generated_text:
            print("Warning: No text was generated from the event stream.")
        
        return generated_text

class BedrockAgentFlowAlias:
    def __init__(self, region='us-east-1'):
        self.region = region
        self.client = boto3.client('bedrock-agent', region_name=self.region)

    def list_aliases(self, flow_identifier):
        try:
            response = self.client.list_flow_aliases(
                flowIdentifier=flow_identifier
            )
            
            aliases = response.get('flowAliasSummaries', [])
            for alias in aliases:
                print(f"Alias Name: {alias['name']}")
                print(f"Alias Identifier: {alias['id']}")
                print(f"ARN: {alias['arn']}")
                print(f"Flow ID: {alias['flowId']}")
                print(f"Created At: {alias['createdAt']}")
                print(f"Updated At: {alias['updatedAt']}")
                if 'description' in alias:
                    print(f"Description: {alias['description']}")
                if 'routingConfiguration' in alias and alias['routingConfiguration']:
                    print(f"Flow Version: {alias['routingConfiguration'][0]['flowVersion']}")
                print("---")
            
            return aliases
        except Exception as e:
            print(f"エラーが発生しました: {str(e)}")
            return None

    def get_alias_by_version(self, flow_identifier, version) -> dict:
        if not isinstance(version, int) or version <= 0:
            raise ValueError("バージョンは正の整数である必要があります。")
        
        if not flow_identifier or not isinstance(flow_identifier, str):
            raise ValueError("有効なflow_identifierを指定してください。")

        try:
            response = self.client.list_flow_aliases(flowIdentifier=flow_identifier)
            aliases = response.get('flowAliasSummaries', [])
            
            for alias in aliases:
                if 'routingConfiguration' in alias and alias['routingConfiguration']:
                    if int(alias['routingConfiguration'][0]['flowVersion']) == version:
                        return {
                            'name': alias['name'],
                            'id': alias['id'],
                            'arn': alias['arn'],
                            'flowId': alias['flowId'],
                            'createdAt': alias['createdAt'],
                            'updatedAt': alias['updatedAt'],
                            'description': alias.get('description'),
                            'flowVersion': alias['routingConfiguration'][0]['flowVersion']
                        }
            
            print(f"バージョン {version} のエイリアスが見つかりませんでした。")
            return None

        except ClientError as e:
            error_code = e.response['Error']['Code']
            error_message = e.response['Error']['Message']
            print(f"エラーが発生しました: {error_code} - {error_message}")
            return None
        except Exception as e:
            print(f"予期せぬエラーが発生しました: {str(e)}")
            return None
```

## 2. セットアップ

必要なライブラリをインストールします。

```
pip install boto3
```

AWSの認証情報を設定します（通常は`~/.aws/credentials`ファイルか環境変数で行います）。

## 3. ライブラリの使用

以下は、ライブラリを使用して特定のバージョンのエイリアスを取得し、そのエイリアスを使用してFlowを呼び出す完全な例です。

```python
from bedrock_agent_library import BedrockAgentFlow, BedrockAgentFlowAlias

# FlowとAliasの設定
flow_identifier = "YOUR_FLOW_ID"  # コンソールで作成したFlowのID
version = 6  # 使用したいFlowのバージョン

# BedrockAgentFlowAliasのインスタンスを作成
agent_flow_alias = BedrockAgentFlowAlias()

# 特定のバージョンのエイリアスを取得
alias_info = agent_flow_alias.get_alias_by_version(flow_identifier, version)

if alias_info:
    print("取得されたエイリアス情報:")
    for key, value in alias_info.items():
        print(f"{key}: {value}")
    
    # エイリアスIDを取得
    flow_alias_identifier = alias_info['id']
    
    # BedrockAgentFlowのインスタンスを作成
    agent_flow = BedrockAgentFlow()

    # Flowを呼び出す
    message = "Hello, Bedrock Agent!"
    result = agent_flow.invoke_flow(message, flow_identifier, flow_alias_identifier)
    print("生成されたテキスト:", result)
else:
    print(f"バージョン {version} のエイリアスが見つかりませんでした。")
```

## 主な更新点

1. `BedrockAgentFlow`クラスの`invoke_flow`メソッドが更新され、生成されたテキストを返すようになりました。
2. `invoke_flow`メソッド内でイベントストリームの処理が改善され、生成されたテキストを抽出するようになりました。
3. 呼び出し例が更新され、生成されたテキストを表示するようになりました。

## 注意事項

- `YOUR_FLOW_ID`は、実際にコンソールで作成したFlowのIDに置き換えてください。
- `version`は、使用したいFlowのバージョンに合わせて変更してください。
- FlowやエイリアスがAWS上に存在しない場合、エラーが発生します。
- AWS認証情報が正しく設定されていることを確認してください。
- Flowの作成、バージョニング、エイリアスの登録はこのライブラリでは行えません。これらはAWS Management ConsoleまたはAWS CLIを使用して行う必要があります。

このワークフローでは、まず`get_alias_by_version`メソッドを使用して特定のバージョンのエイリアス情報を取得し、その後で取得したエイリアスIDを使用してFlowを呼び出しています。これにより、特定のバージョンのFlowを確実に使用することができます。また、Flowから生成されたテキストを取得し、表示することができるようになりました。