Generative AI

ChainlitでのOAuth認証にスコープを追加する方法

こんにちは!

ChainlitというPythonでチャットアプリを簡単につくれるライブラリをご存知でしょうか。ChainlitではChat機能はもちろんとして、OAuthを使用したログイン機能を簡単に実装できます。

しかし、ログインした人の権限を使用してBigQueryとか叩きたいな~と思った時に、デフォルトでは必要最低限の権限(スコープ)のみ与えるようになっているため権限がなく実行できません。

今回はその権限を追加するために少し試行錯誤したので備忘録としてブログにしました。
※Google OAuth前提となりますので違うプロバイダの場合は適宜変更してください。

実際に権限(スコープ)を設定しているコードの箇所の確認

Chainlitでログイン機能を実装するには以下のようなコードを書き、環境変数を設定するだけで簡単に実装できます。

環境変数としてはOAUTH_GOOGLE_CLIENT_IDとOAUTH_GOOGLE_CLIENT_SECRETを設定しています。

from typing import Dict, Optional
import chainlit as cl


@cl.oauth_callback
def oauth_callback(
  provider_id: str,
  token: str,
  raw_user_data: Dict[str, str],
  default_user: cl.User,
) -> Optional[cl.User]:
  if provider_id == "google":
    if raw_user_data["hd"] == "example.org":
      return default_user
  return None

これだけじゃどこで権限を設定しているかわからないのでデコレータの実装をみてみましょう。

@trace
def oauth_callback(
    func: Callable[
        [str, str, Dict[str, str], User, Optional[str]], Awaitable[Optional[User]]
    ],
) -> Callable:
    """
    Framework agnostic decorator to authenticate the user via oauth

    Args:
        func (Callable[[str, str, Dict[str, str], User, Optional[str]], Awaitable[Optional[User]]]): The authentication callback to execute.

    Example:
        @cl.oauth_callback
        async def oauth_callback(provider_id: str, token: str, raw_user_data: Dict[str, str], default_app_user: User, id_token: Optional[str]) -> Optional[User]:

    Returns:
        Callable[[str, str, Dict[str, str], User, Optional[str]], Awaitable[Optional[User]]]: The decorated authentication callback.
    """

    if len(get_configured_oauth_providers()) == 0:
        raise ValueError(
            "You must set the environment variable for at least one oauth provider to use oauth authentication."
        )

    config.code.oauth_callback = wrap_user_function(func)
    return func

まだわからないですね。get_configured_oauth_providers()が怪しいのでみてみます。

def get_configured_oauth_providers():
    return [p.id for p in providers if p.is_configured()]

近づいてきた気がします。今度はprovidersが何か見てみましょう。

providers = [
    GithubOAuthProvider(),
    GoogleOAuthProvider(),
    AzureADOAuthProvider(),
    AzureADHybridOAuthProvider(),
    OktaOAuthProvider(),
    Auth0OAuthProvider(),
    DescopeOAuthProvider(),
    AWSCognitoOAuthProvider(),
    GitlabOAuthProvider(),
    KeycloakOAuthProvider(),
    GenericOAuthProvider(),
]

それっぽいのがありました!今回はGoogleOAuthを使用しているのでGoogleOAuthProvider()の中をみましょう。

class GoogleOAuthProvider(OAuthProvider):
    id = "google"
    env = ["OAUTH_GOOGLE_CLIENT_ID", "OAUTH_GOOGLE_CLIENT_SECRET"]
    authorize_url = "https://accounts.google.com/o/oauth2/v2/auth"

    def __init__(self):
        self.client_id = os.environ.get("OAUTH_GOOGLE_CLIENT_ID")
        self.client_secret = os.environ.get("OAUTH_GOOGLE_CLIENT_SECRET")
        self.authorize_params = {
            "scope": "https://www.googleapis.com/auth/userinfo.profile https://www.googleapis.com/auth/userinfo.email",
            "response_type": "code",
            "access_type": "offline",
        }

        if prompt := self.get_prompt():
            self.authorize_params["prompt"] = prompt

    async def get_token(self, code: str, url: str):
        payload = {
            "client_id": self.client_id,
            "client_secret": self.client_secret,
            "code": code,
            "grant_type": "authorization_code",
            "redirect_uri": url,
        }
        async with httpx.AsyncClient() as client:
            response = await client.post(
                "https://oauth2.googleapis.com/token",
                data=payload,
            )
            response.raise_for_status()
            json = response.json()
            token = json.get("access_token")
            if not token:
                raise httpx.HTTPStatusError(
                    "Failed to get the access token",
                    request=response.request,
                    response=response,
                )
            return token

    async def get_user_info(self, token: str):
        async with httpx.AsyncClient() as client:
            response = await client.get(
                "https://www.googleapis.com/userinfo/v2/me",
                headers={"Authorization": f"Bearer {token}"},
            )
            response.raise_for_status()
            google_user = response.json()
            user = User(
                identifier=google_user["email"],
                metadata={"image": google_user["picture"], "provider": "google"},
            )
            return (google_user, user)

コンストラクタの中のauthorize_paramsにscopeを設定している箇所がありますね。そこに任意のスコープを追加すればその権限がログイン時に付与されます。しかし引数や環境変数で設定できそうにありません。直接ライブラリのコードを書き換えるのはやりたくないので今回はモンキーパッチをあてることにしました。

下記のコードをアプリケーションコードの上部に書いておけば動くと思います。

from chainlit import oauth_providers
from chainlit.oauth_providers import GoogleOAuthProvider

def setup_google_oauth():
    """
    GoogleOAuthProviderをカスタマイズするための関数
    """
    # 任意のスコープの設定
    scope = (
        "https://www.googleapis.com/auth/userinfo.profile"
        "https://www.googleapis.com/auth/userinfo.email"
        "https://www.googleapis.com/auth/bigquery" # 追加
    )

    for provider in oauth_providers.providers:
        if isinstance(provider, GoogleOAuthProvider):
            provider.authorize_params.update(scope=scope)
            print("✅ GoogleOAuthProvider の authorize_params を更新しました!")


setup_google_oauth()

簡単に説明をするとchainlit.oauth_providersのprovidersに該当のインスタンスが入っているのでそのインスタンスのauthorize_paramsを新しいscopeにupdateしています。

そうすると冒頭のログイン機能のためのコードのtoken引数にアクセストークンが入るのでログインユーザーのメタデータに追加します。

@cl.oauth_callback
def oauth_callback(
  provider_id: str,
  token: str,
  raw_user_data: Dict[str, str],
  default_user: cl.User,
) -> Optional[cl.User]:
  if provider_id == "google":
    if raw_user_data["hd"] == "example.org":
      default_app_user.metadata["access_token"] = token # 追加
      return default_user
  return None

実際にスコープが含まれているかはログイン時に取得したアクセストークンを使用して下記のようなコードを書けば確認できます。

import requests

access_token = ACCESS_TOKEN
url = f"https://oauth2.googleapis.com/tokeninfo?access_token={access_token}"

response = requests.get(url)

print(response.content)

あとはアクセストークンを使って下記のようなコードを書けば、ログインしたユーザーの権限でBigQueryが使えるようになります!

from google.oauth2.credentials import Credentials
from google.cloud import bigquery

creds = Credentials(
    token=access_token,
    token_uri="https://oauth2.googleapis.com/token",
    client_id=OAUTH_GOOGLE_CLIENT_ID,
    client_secret=OAUTH_GOOGLE_CLIENT_SECRET,
    scopes="https://www.googleapis.com/auth/bigquery"
)

bq_client = bigquery.Client(credentials=creds, project="YOUR_PROJECT_ID")

おわりに

Chainlitを使っている人に少しでもこの記事が参考になれば幸いです。

エクスチュアはマーケティングテクノロジーを実践的に利用することで企業のマーケティング活動を支援しています。
ツールの活用にお困りの方はお気軽にお問い合わせください。

ピックアップ記事

  1. 最速で理解したい人のためのIT用語集

関連記事

  1. ChatGPT

    ChatGPT ProからClaude3 Proへ移行した話

    こんにちは、エクスチュアの黒岩です。突然ですが、皆様はどんな…

  2. Python

    【完全版】MacでSeleniumを環境構築から実行まで 〜Python&Chrome〜

    Seleniumって何?Selenium(セレニウム)とは、Webア…

  3. Python

    わかりやすいPyTorch入門②(ニューラルネットワークによる分類)

    ニューラルネットワークを使ってワインの種類を分類する今回はsciki…

  4. Python

    Streamlit in Snowflakeによるダッシュボード作成

    こんにちは、エクスチュアの石原です。前回に引き続き、Stre…

  5. 未分類

    databricksのnotebookを使ってみよう その2

    こんにちは。エクスチュアの松村です。…

  6. Python

    わかりやすいPyTorch入門③(手書き数字認識と精度の向上)

    手書き数字認識今回は前回に続きニューラルネットワークを扱います。デ…

カテゴリ

最近の記事

  1. ChainlitでのOAuth認証にスコープを追加する方法
  2. Snowflake無料トライアルの始め方
  3. TROCCO入門
  4. コンポーザブルCDPにおけるSnowflakeのマルチモーダ…
  5. boxMCPサーバーを使ってみた
  1. Google Analytics

    Google Analytics 導入方法
  2. Google Analytics

    Google Analytics4 を知る。
  3. Google Cloud Platform

    Google Compute Engine: 一定時間経過したらタスクを強制終了…
  4. ブログ

    カスタマージャーニー作成
  5. Adobe Analytics

    Adobe Analytics: BigQueryにロードしたデータフィードをD…
PAGE TOP