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. Generative AI

    Snowflake の Copilot が優秀すぎる件について

    こんにちは、喜田です。いままでSnowflakeのライトユー…

  2. KARTE

    KARTE 成果の出るアイデアを考える_ツールを活用できていないと感じたら

    エクスチュアの林です。今回はKARTEを活用していらっしゃる…

  3. Claude

    boxMCPサーバーを使ってみた

    はじめにはじめまして、今年度入社した中村です。最近なにかと話…

  4. ChatGPT

    Open Interpreter+VScode+Dockerで生成AIによるコード開発環境構築(Wi…

    はじめにこんにちは、エクスチュアの石原です。皆さん、…

  5. Google Tag Manager

    【GA4/GTM】dataLayerを活用しよう

    はじめにこんにちは、エクスチュアの岩川です。GA4の…

  6. Python

    Streamlit in SnowflakeによるStreamlitアプリケーション作成

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

カテゴリ

最近の記事

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

    Tableau:分かりやすいLOD-概要編
  2. Linux

    Linux Tips: コマンドラインでJSONをCSVに変換する
  3. Adobe Analytics

    訪問回数 とは-Adobe Analyticsの指標説明
  4. Adobe Analytics

    Adobe Analytics: DatafeedをGoogle BigQuer…
  5. Adobe Analytics

    AdobeAnalytics: s.Util.getQueryParam で複数…
PAGE TOP