Bedrock Agentsでユーザー確認機能が追加されたので試してみた

## はじめに
ドキュメントを眺めていたら9/11にしれっとアップデートが記載されていたので息抜きがてらに検証してみました。

https://docs.aws.amazon.com/bedrock/latest/userguide/agents-userconfirmation.html

## ユーザー確認とは
Agentsがアクショングループを呼び出す前に、ユーザーに実行承認を求める機能になります。公式ドキュメント上では悪意あるプロンプトインジェクションからアプリケーションを保護出来るというメリットが記載されていました。

また、確認したところデフォルトでは無効になっており、各アクショングループ上で有効化する必要があるのでユースケースに応じて使い分けていけそうです。

## コンソールで有効化して試してみた
私が確認した時点ではユーザー確認機能を有効化するためにはアクショングループの種類が「**[Define with function details]()**」である必要がありました。

![スクリーンショット_2024-09-15_3_21_18.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/354022/9d95b4e0-093c-8c4e-171d-131c9167d848.png)

**Define with function details**を選択することで、「**Action group function**」内にて「**Enable confirmation of action group function - オプション**」が表示され、有効/無効を選択できるようになります。

![スクリーンショット_2024-09-15_3_22_00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/354022/d424ca5b-29d0-27f0-d85e-3b96f27c50a5.png)

**Define with function details**の使い方は@hayao_kさんの以下の記事が参考になるかと思います。※記事中にも記載にあるように準備する際にはLambda関数のResponse形式にも注意が必要です。

https://qiita.com/hayao_k/items/57cf2e7950c04e88895a

[#Bedrock開発入門](https://www.sbcr.jp/product/4815626440/)で使用していた、AWSブログ取得のLambda関数をアクショングループとして紐づけ、ユーザー確認機能を有効化にしてコンソール上で試してみました。

<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/354022/27533fb0-fd3f-751a-fed4-05fea239c27a.png" width="350" alt="スクリーンショット">

AWSの最新情報について指示を投げかけたところ、「**search-aws-blogs**」のアクショングループを実行してよいかの承認を求められることが確認できましたので**Comfirm**」をクリックします。

<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/354022/122674cc-289a-2507-815f-dafd58b4896d.png" width="350" alt="スクリーンショット">

アクショングループが実行されて執筆当時のAWSブログから最新のものが無事取得できました。

## Streamlitから試してみた。

※Streamlitを実行するとAWSプロフィールやリージョン、エージェントID、エージェントエイリアスIDが求められるので自身の環境に合わせて入力してください。

**InvokeAgent**からのレスポンスに**returnControl**のイベントがあり、**invocationInputs**が含まれている場合にユーザー確認が求められます。

```response.json
HTTP/1.1 200
x-amzn-bedrock-agent-content-type: contentType
x-amz-bedrock-agent-session-id: sessionId
Content-type: application/json

{
   "chunk": { 
      ...
   },
   ...
   "returnControl": { 
      "invocationId": "string",
      "invocationInputs": [ 
         { ... }
      ]
   },
   "trace": { 
      "agentAliasId": "string",
      "agentId": "string",
      "agentVersion": "string",
      "sessionId": "string",
      "trace": { ... }
   },
}
```

また、ユーザー確認リクエストには**sessionState**に**invocationId**と**returnControlInvocationResults**を含め、更に**[functionResult](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_FunctionResult.html)**にはactionGroup、function、そしてresponseBody**に***confirmationState**としてユーザー確認の結果(CONFIRM | DENY)を一緒に送る必要があります。

```request.json
POST /agents/agentId/agentAliases/agentAliasId/sessions/sessionId/text HTTP/1.1
Content-type: application/json

{
   "enableTrace": boolean,
   "endSession": boolean,
   "inputText": "string",
   "sessionState": { 
      "invocationId": "string",
      "promptSessionAttributes": { 
         "string" : "string" 
      },
      "returnControlInvocationResults": [ 
         { ... }
      ],
      "sessionAttributes": { 
         "string" : "string" 
      }
   }
}
```

### コード

:::note info

```python
import uuid
import boto3
import streamlit as st
import json

st.sidebar.title("Configuration")
profile_name = st.sidebar.text_input("AWS Profile Name", "")
region_name = st.sidebar.text_input("AWS Region", "ap-northeast-1")
AGENT_ID = st.sidebar.text_input("エージェントID")
AGENT_ALIAS_ID = st.sidebar.text_input("エージェントエイリアスID")

# セッションIDの生成と保存
if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())
session_id = st.session_state.session_id

# Bedrockクライアントの初期化
@st.cache_resource
def get_bedrock_client(profile_name, region_name):
    session = boto3.Session(profile_name=profile_name)
    return session.client('bedrock-agent-runtime', region_name=region_name)

bedrock_client = get_bedrock_client(profile_name, region_name)


st.title("ユーザー確認デモ")

# セッション状態の初期化
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "pending_confirmation" not in st.session_state:
    st.session_state.pending_confirmation = None
if "agent_session_state" not in st.session_state:
    st.session_state.agent_session_state = {
        "promptSessionAttributes": {},
        "sessionAttributes": {}
    }

# Bedrock Agentを呼び出す関数
def invoke_bedrock_agent(input_text, agent_session_state=None):
    invoke_params = {
        "inputText": input_text,
        "agentId": AGENT_ID,
        "agentAliasId": AGENT_ALIAS_ID,
        "sessionId": session_id,
        "enableTrace": True,
    }
    if agent_session_state:
        invoke_params["sessionState"] = agent_session_state

    return bedrock_client.invoke_agent(**invoke_params)

# Bedrock Agentからのレスポンスを処理する関数
def process_agent_response(response):

    full_response = ""
    requires_confirmation = False

    for event in response.get("completion", []):
        if "chunk" in event:
            chunk = event["chunk"]["bytes"].decode("utf-8")
            full_response += chunk
        
        if "returnControl" in event:
            return_control = event["returnControl"]
            if "invocationInputs" in return_control:
                st.session_state.pending_confirmation = return_control
                requires_confirmation = True

    return full_response, requires_confirmation

# ユーザー確認処理
if st.session_state.pending_confirmation:
    for input_item in st.session_state.pending_confirmation["invocationInputs"]:
        action_group = input_item["functionInvocationInput"]["actionGroup"]
        st.write(f"{action_group}の実行を承認しますか？")
    
    # 確認ボタンと拒否ボタンの表示
    col1, col2 = st.columns(2)
    with col1:
        if st.button("確認"):
            user_decision = "CONFIRM"
    with col2:
        if st.button("拒否"):
            user_decision = "DENY"
    
    if 'user_decision' in locals():
        if user_decision == "CONFIRM":
            # ユーザーが確認した場合の処理
            invocation_id = st.session_state.pending_confirmation.get("invocationId")
            function = input_item["functionInvocationInput"]["function"]

            updated_session_state = {
                "invocationId": invocation_id,
                "returnControlInvocationResults": [
                    {
                        "functionResult": {
                            "actionGroup": action_group,
                            "function": function,
                            "responseBody": {
                                "TEXT": { 
                                    "body": json.dumps({
                                        "confirmationState": user_decision
                                    })
                                }
                            }
                        }
                    }
                ],
            }
            st.session_state.agent_session_state = updated_session_state
        
            with st.spinner("処理中..."):
                response = invoke_bedrock_agent("", updated_session_state)
                full_response, requires_confirmation = process_agent_response(response)
                
                st.write(f"{full_response}")
                
                if not requires_confirmation:
                    # 確認が不要な場合、チャット履歴を更新し、セッション状態をリセット
                    st.session_state.chat_history.append({"role": "assistant", "content": full_response})
                    st.session_state.pending_confirmation = None
                    st.session_state.agent_session_state = {
                        "promptSessionAttributes": {},
                        "sessionAttributes": {}
                    }
                    st.success("処理が完了しました。")

        elif user_decision == "DENY":
            # ユーザーが拒否した場合、セッション状態をクリアして再実行
            st.session_state.clear()
            st.experimental_rerun()

# ユーザー入力の処理
if user_input := st.chat_input("何でも聞いてください"):
    # ユーザー入力をチャット履歴に追加
    st.session_state.chat_history.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    with st.chat_message("assistant"):
        response_placeholder = st.empty()
        full_response = ""

        # Bedrock Agentを呼び出し、レスポンスを処理
        response = invoke_bedrock_agent(user_input, st.session_state.agent_session_state)
        full_response, requires_confirmation = process_agent_response(response)
        
        if not requires_confirmation:
            # 確認が不要な場合、レスポンスを表示してチャット履歴を更新
            response_placeholder.markdown(full_response)
            st.session_state.chat_history.append({"role": "assistant", "content": full_response})
        else:
            st.rerun()
```
:::

![スクリーンショット_2024-09-15_3_46_53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/354022/e0ef3e0e-0dd2-5336-516b-73a6afd464a8.png)

### 実行結果動画

<iframe width="560" height="315" src="https://www.youtube.com/embed/iVK8BGvkcQ4?si=rCW84VQ9SMvZ-pBK" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" referrerpolicy="strict-origin-when-cross-origin" allowfullscreen></iframe>

ユーザーの指示に応じてAgentsがアクショングループの実行確認を求める->「確認」ボタンがクリック->AWSブログ取得のアクショングループが実行->実行結果表示の一連の流れを確認することができました。

## おわりに
アクショングループに紐づけるLambda関数の処理内容によっては、ワンクッションおいてユーザー確認を必ず挟ませたい、そんなユースケースがある場合に使えそうなので覚えておいて損はないかなという印象でした。

今回はユーザ確認でDenyを選択した場合の処理の実装は省きましたが、コンソール上で確認したところLLMが持つ一般的な知識を用いて回答が出力されました。エージェント向けの指示プロンプト次第ではその場合の分岐も制御できるかもしれませんので、また改めて試してみようと思います。