ローカル環境のStreamlitにAmazon Cognitoで認証をかける

みなさま、Bedrockしてますか？
Streamlitを使ったBedrockのデモアプリが出来上がったら、いろんな人に使ってもらいたいですよね。

AWS上で構築する場合は、CognitoとALBを組み合わせると簡単にログイン機能が実現できます。

https://docs.aws.amazon.com/ja_jp/elasticloadbalancing/latest/application/listener-authenticate-users.html

https://repost.aws/ja/knowledge-center/cognito-user-pool-alb-authentication

ただ、デモとして作ったものをちょっと動かすために、EC2を構築するのは少し面倒ですし、インターネット公開は抵抗があるかもしれません。

* EC2を起動するほどではないのでローカル環境で動作させたい
* 利用者を限定したいので認証機能はほしい

こんな要件を満たせないか調査したところ、良い方法が見つかりましたのでご紹介します。


## アーキテクチャ

`Oauth2 Proxy`と`Amazon Cognito`を使います。

Oauth2 Proxyは認証機能を提供するリバースプロキシサーバーです。未ログイン状態でアクセスするとログイン画面に遷移させ、ログイン状態では、バックエンドサーバーにリクエストを転送します。

https://github.com/oauth2-proxy/oauth2-proxy

Oauth2 ProxyとCognitoのユーザープール機能を連携させることで、ログイン機能を実現します。

https://aws.amazon.com/jp/cognito/

更にCognitoのアイデンティティプール機能で、ログインの際に取得したIDトークンからAWSの認証情報（アクセスキー、シークレットキー、セッショントークン）を取得することができますので、この認証情報を使ってBedrockへアクセスします。

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/4b2f475d-b4c1-e672-2653-4efce687e211.png)

ログインしたユーザー自身がAWSへのアクセス権限を持つので、 **サーバー側プログラムは認証情報（アクセスキー、シークレットキー）を持っていなくてもよい構成** となります。

**アクセスキー、シークレットキーのベタ書きから開放されます！**

素敵でしょ？


:::note
ちょうどこの投稿を検証している最中に同様のブログが公開されています。こちらもご参照いただければと思います。

Nginx + OAuth2 Proxy + StreamlitでGoogleログイン後にStreamlitにアクセスする環境をローカルコンテナ環境で作ってみた
https://dev.classmethod.jp/articles/nginx-oauth2proxy-streamlit/
:::


Amplify Authと同等の仕組みを通常のWebアプリで実現した構成ですので、Streamlit以外にも応用できると思います。

## 手順

1. Cognitoユーザープールを作成する
1. Cognitoアイデンティティプールを作成する
1. OAuth2 ProxyとStreamlitの連携設定を行う
1. StreamlitにAWS認証情報取得ロジックを追加する

Streamlitのアプリは以下のようなものがある前提で進めます。

<details><summary>app.py</summary>

```python:app.py
import json

import boto3

import streamlit as st


def stream(stream: dict):
    for event in stream.get("body"):
        chunk = json.loads(event["chunk"]["bytes"])
        if "delta" in chunk and "text" in chunk["delta"]:
            yield chunk["delta"]["text"]


st.title("Bedrock チャット")


if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"][0]["text"])

if prompt := st.chat_input("何でも聞いてください。"):
    st.session_state.messages.append(
        {
            "role": "user",
            "content": [{"type": "text", "text": prompt}],
        }
    )

    with st.chat_message("user"):
        st.markdown(prompt)

    client = boto3.client("bedrock-runtime")

    response = client.invoke_model_with_response_stream(
        modelId="anthropic.claude-3-haiku-20240307-v1:0",
        body=json.dumps(
            {
                "anthropic_version": "bedrock-2023-05-31",
                "max_tokens": 300,
                "system": "あなたは優秀なAIボットです",
                "messages": st.session_state.messages,
            }
        ),
    )

    with st.chat_message("assistant"):
        result = st.write_stream(stream(response))

    st.session_state.messages.append(
        {
            "role": "assistant",
            "content": [{"type": "text", "text": result}],
        }
    )
```

</details>

:::note
今回は手元の環境で動作させることを前提とし、Webアプリのエンドポイントは`http://localhost:8080`としています。
ホスト名やIPアドレス、ポート番号を変更する場合は適宜読み替えてください。
:::

## Cognitoユーザープールを作成する

:::note
手順がわかってる人向け

* ホストされた UIを有効にする
* クライアントシークレットを生成する
* コールバックに`http://localhost:8080/oauth2/callback`を、サインアウトに`http://localhost:8080/oauth2/sign_out`を入力
:::

* 後で必要になる情報

    | 項目 |
    | --- |
    | リージョン |
    | ユーザープールID |
    | クライアントID |
    | クライアントシークレット |
    | ホストされた UIのドメインプレフィックス |

↓↓↓操作手順は長くなったので折りたたんでいます↓↓↓

<details><summary>マネジメントコンソールでの手順</summary>

1. マネジメントコンソールで、[Cognito管理画面](https://ap-northeast-1.console.aws.amazon.com/cognito/v2/home)にアクセスする。
    `ユーザープールを作成`ボタンをクリック

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/63529b6e-ff2c-4b24-3bea-06c4a74907e0.png)

    サインインオプションを一つ選び`次へ`ボタンをクリック。（今回はEメールを選択）

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/39a10721-930e-5fec-3c66-39173b5583cc.png)

1. セキュリティ要件を設定画面
    お好きな設定をして`次へ`ボタンをクリック。

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/10d0bc7a-5f88-4b4b-8993-3e9b07b802b0.png)

1. サインアップエクスペリエンスを設定画面
    お好きな設定をして`次へ`ボタンをクリック。

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/5fad9ceb-189f-de63-149d-5bc97335e643.png)

1. メッセージ配信を設定画面
    お好きな設定をして`次へ`ボタンをクリック。

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/c6dd51b4-fa47-0ace-617f-e576f4f0cb8f.png)

1. アプリケーションを統合画面
    * `ユーザープール名`を入力する
    * `Cognito のホストされた UI を使用`にチェックを入れる
    * ドメインタイプは`Cognito ドメインを使用する`を選び、ドメインプレフィックスを入力する

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/04dab04e-8650-5e24-c17a-f9fbc2d09eb4.png)

    * `アプリケーションタイプ`は`パブリッククライアント`を選択する
    * `アプリケーションクライアント名`を入力する
    * `クライアントシークレット`は`クライアントのシークレットを生成する`を選択する
    * `許可されているコールバック URL`に`http://localhost:8080/oauth2/callback`と入力

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/69b2eacd-2866-21e6-5da7-aec1ddf586c2.png)

    * 高度なアプリケーションクライアントの設定の中の一番下にある`許可されているサインアウト URL`に`http://localhost:8080/oauth2/sign_out`を追加

    `次へ`をクリック

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/c870dff5-c6b4-8e61-f9ee-202d8b6c34a5.png)


1. 確認および作成画面
    `ユーザープールを作成`ボタンをクリック

1. 作成できました。`詳細を表示`をクリック

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/fe4d8beb-5bf2-dcf9-a3fa-5e067ac41268.png)

1. ユーザープールIDをメモする

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/36fc7044-73fd-432e-87f8-b3ce3fdb62e2.png)

1. アプリケーションの統合タブの最下部にあるアプリケーションクライアントをクリックする

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/fce96200-9aec-7c5f-e21d-597bcf90d805.png)

    クライアントID、クライアントシークレットをメモする

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/071bda94-f18e-adb6-3c5a-a300c0be7dd1.png)

1. ユーザープールの設定画面に戻り、ユーザータブを選択、`ユーザーを作成`をクリック

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/27d0ab4f-5c93-4f65-9f0d-6e5ae2d23e15.png)

    ユーザーを作成する
    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/2802399b-cdb0-08e2-3818-389a37b3fe3f.png)

これでCognitoユーザープールの作成は完了です。

</details>


## Cognitoアイデンティティプールを作成する

:::note
手順がわかってる人向け

* 認証されたアクセスとして先ほど作成したCognitoユーザープールを選択する
* 作成されるIAMロールにBedrockへのアクセス権限を付与する
:::

* 後で必要になる情報

    | 項目 |
    | --- |
    | IDプールのID |

↓↓↓操作手順は長くなったので折りたたんでいます↓↓↓

<details><summary>マネジメントコンソールでの手順</summary>

1. Cognito管理画面の左メニューの`IDプール`を選択し、`ID プールを作成`ボタンをクリック。

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/9f8a4432-20cf-ed27-efdc-54e8aac87dc8.png)

1. ID プールの信頼を設定画面
    * ユーザーアクセスの`認証されたアクセス`にチェックを入れる
    * 認証された ID ソースの`Amazon Cognito ユーザープール`にチェックを入れる

    `次へ`ボタンをクリック

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/9e811d50-a0c8-bf8b-aede-25c8d0c015a3.png)

1. 許可を設定画面
    * `新しいIAMロールを作成`を選択し、IAMロール名を入力（streamlit-idpool-roleとしました）

    `次へ`ボタンをクリック

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/cad8c9a1-675d-831d-e538-52caedfd5ad9.png)

1. ID プロバイダーを接続画面

    * ユーザープールの詳細欄の`ユーザープール ID`と`アプリクライアント ID`を先程作成したものを選択
    * その他はデフォルトのまま

    `次へ`ボタンをクリック
    
    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/68c9c31c-a073-eba4-9114-1630184c433d.png)


1. プロパティを設定画面

    * IDプール名を入力

    `次へ`ボタンをクリック

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/a4b8d345-60e9-eb24-18cf-48aab142aa81.png)

1. 確認および作成画面

    内容を確認し、`ID プールを作成`ボタンをクリック

1. 作成できました。`IDプールのID`をメモします。
    画面上部の青いバーの`ロールの表示`をクリック

1. 作成したIAMロールの画面が開きます。

    Bedrockへの許可を追加するので、`許可を追加`メニューの`ポリシーをアタッチ`をクリック

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/9924ab37-09f0-be9c-ee80-fc1d361313bf.png)

1. `BedrockFullAccess`にチェックを入れ、`許可を追加`ボタンをクリック 

    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/96a26ba9-02b1-abaf-aa98-0720ffc1a8ed.png)

これでCognitoアイデンティティプールの作成は完了です。

</details>


## OAuth2 ProxyとStreamlitの連携設定を行う

今回は、Docker Composeを使って環境を構築します。

作成するディレクトリー構成は以下のとおりです。

```shell:tree
.
├── .env
├── app
│   ├── Dockerfile
│   ├── app.py
│   └── requirements.txt
└── docker-compose.yaml

1 directory, 5 files
```

### OAuth2 Proxy

* docker-compose.yaml
    oauth2-proxyは公開されているコンテナイメージをそのまま使用します。
    
    ```yaml:docker-compose.yaml （一部抜粋）
    services:
      oauth2-proxy:
        image: quay.io/oauth2-proxy/oauth2-proxy:latest
        ports:
          - 8080:4180
        env_file:
          - ./.env
    ```

* .env
    設定はすべて環境変数で行います。環境変数は`.env`ファイルに記述します。
    
    ```toml:.env
    COGNITO_REGION=xxxxx # Cognitoのリージョン
    COGNITO_IDENTITY_ID=xxxxx # CognitoのIDプールのID
    COGNITO_USERPOOL_ID=xxxxx # CognitoのユーザープールのID
    COGNITO_CLIENT_ID=xxxxx # CognitoのクライアントID
    COGNITO_CLIENT_SECRET=xxxxx # Cognitoのクライアントシークレット
    COGNITO_HOSTED_UI_DOMAIN_PREFIX=xxxxx # CognitoのユーザープールのID
    
    REDIRECT_URL=http://localhost:8080/oauth2/callback
    SIGNOUT_URL=http://localhost:8080/oauth2/sign_out
    
    OAUTH2_PROXY_PROVIDER=oidc
    OAUTH2_PROXY_REDIRECT_URL=${REDIRECT_URL}
    OAUTH2_PROXY_OIDC_ISSUER_URL=https://cognito-idp.${COGNITO_REGION}.amazonaws.com/${COGNITO_USERPOOL_ID}
    OAUTH2_PROXY_UPSTREAMS=http://streamlit:8501
    OAUTH2_PROXY_EMAIL_DOMAINS=*
    
    OAUTH2_PROXY_CLIENT_ID=${COGNITO_CLIENT_ID}
    OAUTH2_PROXY_CLIENT_SECRET=${COGNITO_CLIENT_SECRET}
    
    OAUTH2_PROXY_PASS_ACCESS_TOKEN=true
    OAUTH2_PROXY_PASS_AUTHORIZATION_HEADER=true
    
    OAUTH2_PROXY_COOKIE_SECRET=""
    OAUTH2_PROXY_COOKIE_SECURE=false
    
    OAUTH2_PROXY_SKIP_PROVIDER_BUTTON=true
    OAUTH2_PROXY_HTTP_ADDRESS=0.0.0.0:4180
    OAUTH2_PROXY_SCOPE="openid email phone"
    ```

    値の変更が必要な項目はこちら
    
    | 項目 | 設定する値 |
    | --- | --- |
    | COGNITO_REGION | Cognitoのリージョン |
    | COGNITO_IDENTITY_ID | CognitoのIDプールのID |
    | COGNITO_USERPOOL_ID | CognitoのユーザープールのID |
    | COGNITO_CLIENT_ID | CognitoのクライアントID |
    | COGNITO_CLIENT_SECRET | Cognitoのクライアントシークレット |
    | COGNITO_HOSTED_UI_DOMAIN_PREFIX | CognitoのユーザープールのID |
    | OAUTH2_PROXY_COOKIE_SECRET | クッキーシークレット（※１） |
    
    （※１）[ドキュメント](https://oauth2-proxy.github.io/oauth2-proxy/configuration/overview/)を参考に作成した値をセットします。
    
    ```python:Python3の場合
    python3 -c 'import os,base64; print(base64.urlsafe_b64encode(os.urandom(32)).decode())'
    ```
    
    ```shell:Bashの場合
    dd if=/dev/urandom bs=32 count=1 2>/dev/null | base64 | tr -d -- '\n' | tr -- '+/' '-_' ; echo
    ```

:::note warn
今回の検証はHTTP通信で行いましたので、`OAUTH2_PROXY_COOKIE_SECURE`（CookieのSecure属性）は`false`としました。OAuth2 ProxyにはTLSの設定も可能ですので、こちらも一度検討ください。

TLS Configuration
https://oauth2-proxy.github.io/oauth2-proxy/configuration/tls
:::

### appディレクトリ（Streamlitアプリ）

* app/Dockerfile

    `python:3.12-slim`をベースイメージに、StreamlitとBoto3をインストールしたコンテナイメージを作成します。
    app.pyはイメージには含めず、Volumeマウントする方式としました。

    ```docker:app/Dockerfile
    FROM python:3.12-slim
    
    WORKDIR /opt/streamlit
    
    COPY requirements.txt .
    
    RUN pip install -r requirements.txt
    
    ENTRYPOINT [ "streamlit" ]
    CMD [ "run", "app.py" ]
    
    ```

* docker-compose.yaml

    環境変数はOAuth2 Proxyと共用するようにしました。

    ```yaml:docker-compose.yaml （全体）
    services:
      oauth2-proxy:
        image: quay.io/oauth2-proxy/oauth2-proxy:latest
        ports:
          - 8080:4180
        env_file:
          - ./.env
      streamlit:
          build:
            context: ./app
            dockerfile: Dockerfile
          volumes:
            - ./app/app.py:/opt/streamlit/app.py
          env_file:
            - ./.env
    ```

* app/requirements.txt
    必要なライブラリーをインストール

    ```text:app/requirements.txt
    boto3
    streamlit
    ```

* app/app.py

    あとで修正しますが一旦はもともとのソースのままです。

    <details><summary>app.py</summary>

    ```python:app/app.py
    import json
    
    import boto3
    
    import streamlit as st
    
    
    def stream(stream: dict):
        for event in stream.get("body"):
            chunk = json.loads(event["chunk"]["bytes"])
            if "delta" in chunk and "text" in chunk["delta"]:
                yield chunk["delta"]["text"]
    
    
    st.title("Bedrock チャット")
    
    
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"][0]["text"])
    
    if prompt := st.chat_input("何でも聞いてください。"):
        st.session_state.messages.append(
            {
                "role": "user",
                "content": [{"type": "text", "text": prompt}],
            }
        )
    
        with st.chat_message("user"):
            st.markdown(prompt)
    
        client = boto3.client("bedrock-runtime")
    
        response = client.invoke_model_with_response_stream(
            modelId="anthropic.claude-3-haiku-20240307-v1:0",
            body=json.dumps(
                {
                    "anthropic_version": "bedrock-2023-05-31",
                    "max_tokens": 300,
                    "system": "あなたは優秀なAIボットです",
                    "messages": st.session_state.messages,
                }
            ),
        )
    
        with st.chat_message("assistant"):
            result = st.write_stream(stream(response))
    
        st.session_state.messages.append(
            {
                "role": "assistant",
                "content": [{"type": "text", "text": result}],
            }
        )
    ```
    
    </details>

### Docker Composeビルド＆起動

```shell
docker compose build
```

```shell
docker compose up
```

## 動作確認

ブラウザで`http://localhost:8080`へアクセス。リダイレクトされ、Cognito のホストされた UIのログイン画面に遷移します。

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/03d0dcd7-4a5d-d468-a931-7cf0c0986b60.png)

ログインに成功すると、Streamlitの画面に遷移します。

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/93c8560c-bf00-6a12-6ac0-787064d5094f.png)

まだ、アプリ側での認証情報取得処理を入れていないので、Bedrockへのアクセスはエラーとなります。

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/b0033f4a-6f8b-e809-a4fe-e5a8b218a794.png)


## StreamlitにAWS認証情報取得ロジックを追加する

それでは最後の手順です。Streamlitアプリの中で、AWSの認証情報を取得する処理を追加します。

まず、環境変数から値を取得します。

```diff_python:app/app.py
  import json
+ import os
  
  import boto3
  
  import streamlit as st
  
  
+ # 環境変数を取得
+ cognito_region = os.environ["COGNITO_REGION"]
+ cognito_identity_id = os.environ["COGNITO_IDENTITY_ID"]
+ cognito_userpool_id = os.environ["COGNITO_USERPOOL_ID"]
+ cognito_client_id = os.environ["COGNITO_CLIENT_ID"]
+ cognito_hosted_ui_domain = os.environ["COGNITO_HOSTED_UI_DOMAIN_PREFIX"]
+ 
+ logout_uri = os.environ["SIGNOUT_URL"]
```

次に、StreamlitでHTTPヘッダーを取得するためのメソッドをインポートします。

```diff_python:app/app.py
  import json
  import os
  
  import boto3
  
  import streamlit as st
+ from streamlit.web.server.websocket_headers import _get_websocket_headers
```

HTTPヘッダーの`Authorization`からIDトークンを取得します。
CognitoIdentityの`get_id`と`get_credentials_for_identity`を使って、AWSの認証情報を取得します。
（IdentityPoolIdとIdentityIdの違いはいまいちわかりませんが、この手順が必要なようです。）

```diff_python:app/app.py
+ def get_credentials():
+     headers = _get_websocket_headers()
+ 
+     authorization = headers["Authorization"]
+     id_token = authorization.replace("Bearer ", "")
+ 
+     client = boto3.client("cognito-identity", region_name=cognito_region)
+ 
+     cognito_idp_name = (
+         f"cognito-idp.{cognito_region}.amazonaws.com/{cognito_userpool_id}"
+     )
+ 
+     response = client.get_id(
+         IdentityPoolId=cognito_identity_id, Logins={cognito_idp_name: id_token}
+     )
+ 
+     identity_id = response["IdentityId"]
+ 
+     response = client.get_credentials_for_identity(
+         IdentityId=identity_id,
+         Logins={cognito_idp_name: id_token},
+     )
+ 
+     return response["Credentials"]
```

あとは`get_credentials()`の値を使ってboto3.clientを作る処理の追加と、呼び出し元の変更を行います。

```diff_python:app/app.py
+ def create_bedrock_runtime_client():
+     credentials = get_credentials()
+     client = boto3.Session(
+         aws_access_key_id=credentials["AccessKeyId"],
+         aws_secret_access_key=credentials["SecretKey"],
+         aws_session_token=credentials["SessionToken"],
+         region_name="us-east-1",
+     ).client("bedrock-runtime")
+     return client
```

```diff_python:app/app.py
  if prompt := st.chat_input("何でも聞いてください。"):
      st.session_state.messages.append(
          {
              "role": "user",
              "content": [{"type": "text", "text": prompt}],
          }
      )
  
      with st.chat_message("user"):
          st.markdown(prompt)
  
-     client = boto3.client("bedrock-runtime")
+     client = create_bedrock_runtime_client()
  
      response = client.invoke_model_with_response_stream(
          modelId="anthropic.claude-3-haiku-20240307-v1:0",
          body=json.dumps(
              {
                  "anthropic_version": "bedrock-2023-05-31",
                  "max_tokens": 300,
                  "system": "あなたは優秀なAIボットです",
                  "messages": st.session_state.messages,
              }
          ),
      )
  
      with st.chat_message("assistant"):
          result = st.write_stream(stream(response))
  
      st.session_state.messages.append(
          {
              "role": "assistant",
              "content": [{"type": "text", "text": result}],
          }
      )
```

最後にサインアウトボタンを追加しましょう。

```diff_python
+ def sign_out_url():
+     return f"https://{cognito_hosted_ui_domain}.auth.{cognito_region}.amazoncognito.com/logout?client_id={cognito_client_id}&logout_uri={logout_uri}"
```

```diff_python
  st.title("Bedrock チャット")
  
+ st.link_button("Sign out", sign_out_url())
  
  
  if "messages" not in st.session_state:
      st.session_state.messages = []
```

:::note
サインアウト処理は、Cognito のホストされた UIの`/logout`エンドポイントに`client_id`と`logout_uri`（`http://localhost:8080/oauth2/sign_out`）をパラメーターとして付与しアクセスします。処理の流れの詳細は以下のような感じです

1. `/logout`エンドポイントのアクセスによりCognitoからサインアウトする
1. `logout_uri`にリダイレクトされ、クッキー情報がクリアされる
1. `http://localhost:8080`にリダイレクトされる
1. 未ログイン状態で`http://localhost:8080`にアクセスするので、Cognito のホストされた UIのログイン画面にリダイレクトされる
:::

StreamlitのPythonスクリプトの修正は完了です。

<details><summary>app.py（修正後全体）</summary>

```python:app/app.py
import json
import os

import boto3

import streamlit as st
from streamlit.web.server.websocket_headers import _get_websocket_headers


# 環境変数を取得
cognito_region = os.environ["COGNITO_REGION"]
cognito_identity_id = os.environ["COGNITO_IDENTITY_ID"]
cognito_userpool_id = os.environ["COGNITO_USERPOOL_ID"]
cognito_client_id = os.environ["COGNITO_CLIENT_ID"]
cognito_hosted_ui_domain = os.environ["COGNITO_HOSTED_UI_DOMAIN_PREFIX"]

logout_uri = os.environ["SIGNOUT_URL"]


def sign_out_url():
    return f"https://{cognito_hosted_ui_domain}.auth.{cognito_region}.amazoncognito.com/logout?client_id={cognito_client_id}&logout_uri={logout_uri}"


def get_credentials():
    headers = _get_websocket_headers()

    authorization = headers["Authorization"]
    id_token = authorization.replace("Bearer ", "")

    client = boto3.client("cognito-identity", region_name=cognito_region)

    cognito_idp_name = (
        f"cognito-idp.{cognito_region}.amazonaws.com/{cognito_userpool_id}"
    )

    response = client.get_id(
        IdentityPoolId=cognito_identity_id, Logins={cognito_idp_name: id_token}
    )

    identity_id = response["IdentityId"]

    response = client.get_credentials_for_identity(
        IdentityId=identity_id,
        Logins={cognito_idp_name: id_token},
    )

    return response["Credentials"]


def create_bedrock_runtime_client():
    credentials = get_credentials()
    client = boto3.Session(
        aws_access_key_id=credentials["AccessKeyId"],
        aws_secret_access_key=credentials["SecretKey"],
        aws_session_token=credentials["SessionToken"],
        region_name="us-east-1",
    ).client("bedrock-runtime")
    return client


def stream(stream: dict):
    for event in stream.get("body"):
        chunk = json.loads(event["chunk"]["bytes"])
        if "delta" in chunk and "text" in chunk["delta"]:
            yield chunk["delta"]["text"]


st.title("Bedrock チャット")

st.link_button("Sign out", sign_out_url())

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"][0]["text"])

if prompt := st.chat_input("何でも聞いてください。"):
    st.session_state.messages.append(
        {
            "role": "user",
            "content": [{"type": "text", "text": prompt}],
        }
    )

    with st.chat_message("user"):
        st.markdown(prompt)

    client = create_bedrock_runtime_client()

    response = client.invoke_model_with_response_stream(
        modelId="anthropic.claude-3-haiku-20240307-v1:0",
        body=json.dumps(
            {
                "anthropic_version": "bedrock-2023-05-31",
                "max_tokens": 300,
                "system": "あなたは優秀なAIボットです",
                "messages": st.session_state.messages,
            }
        ),
    )

    with st.chat_message("assistant"):
        result = st.write_stream(stream(response))

    st.session_state.messages.append(
        {
            "role": "assistant",
            "content": [{"type": "text", "text": result}],
        }
    )
```

</details>

---

これで当初の目的は達成しました。

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/4a388010-7b57-3828-148e-948e9e854618.png)

ソースコードはGitHubにて公開しています。

https://github.com/moritalous/streamlit-cognito-auth/