LoginSignup
17
13

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@moritalous

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

Posted at

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

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

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

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

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

アーキテクチャ

Oauth2 ProxyとAmazon Cognitoを使います。

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

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

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

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

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

素敵でしょ？

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

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

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

手順

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

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

app.py
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}],
        }
    )

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

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

手順がわかってる人向け

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

  • 後で必要になる情報

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

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

マネジメントコンソールでの手順

  1. マネジメントコンソールで、Cognito管理画面にアクセスする。
    ユーザープールを作成ボタンをクリック

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

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

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

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

  5. アプリケーションを統合画面

    • ユーザープール名を入力する
    • Cognito のホストされた UI を使用にチェックを入れる
    • ドメインタイプはCognito ドメインを使用するを選び、ドメインプレフィックスを入力する

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

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

    次へをクリック

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

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

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

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

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

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

    ユーザーを作成する

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

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

手順がわかってる人向け

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

  • 後で必要になる情報

    項目
    IDプールのID

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

マネジメントコンソールでの手順

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

  2. ID プールの信頼を設定画面

    • ユーザーアクセスの認証されたアクセスにチェックを入れる
    • 認証された ID ソースのAmazon Cognito ユーザープールにチェックを入れる

    次へボタンをクリック

  3. 許可を設定画面

    • 新しいIAMロールを作成を選択し、IAMロール名を入力（streamlit-idpool-roleとしました）

    次へボタンをクリック

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

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

    次へボタンをクリック

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

    • IDプール名を入力

    次へボタンをクリック

  6. 確認および作成画面

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

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

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

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

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

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

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

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

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

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

1 directory, 5 files

OAuth2 Proxy

  • docker-compose.yaml
    oauth2-proxyは公開されているコンテナイメージをそのまま使用します。

    docker-compose.yaml （一部抜粋）
    services:
  oauth2-proxy:
    image: quay.io/oauth2-proxy/oauth2-proxy:latest
    ports:
      - 8080:4180
    env_file:
      - ./.env

  • .env
    設定はすべて環境変数で行います。環境変数は.envファイルに記述します。

    .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 クッキーシークレット（※１）

    （※１）ドキュメントを参考に作成した値をセットします。

    Python3の場合
    python3 -c 'import os,base64; print(base64.urlsafe_b64encode(os.urandom(32)).decode())'
    Bashの場合
    dd if=/dev/urandom bs=32 count=1 2>/dev/null | base64 | tr -d -- '\n' | tr -- '+/' '-_' ; echo

今回の検証は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マウントする方式としました。

    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と共用するようにしました。

    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
    必要なライブラリーをインストール

    app/requirements.txt
    boto3
streamlit

  • app/app.py

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

    app.py
    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}],
        }
    )

Docker Composeビルド＆起動

docker compose build

docker compose up

動作確認

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

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

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

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

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

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

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ヘッダーを取得するためのメソッドをインポートします。

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の違いはいまいちわかりませんが、この手順が必要なようです。）

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を作る処理の追加と、呼び出し元の変更を行います。

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
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}],
          }
      )

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

+ 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}"

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

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

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

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

app.py（修正後全体）
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}],
        }
    )

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

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

17
13
1

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
17
13

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?