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を使っている人に少しでもこの記事が参考になれば幸いです。

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

Snowflake無料トライアルの始め方前のページ

dbt Cloud使ってみた次のページ

ピックアップ記事

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

関連記事

  1. Python

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

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

  2. Data Clean Room

    PostgreSQLによるデータクリーンルームの可能性について

    こんにちは、喜田です。本投稿は PostgreSQL Adv…

  3. Google Cloud Platform

    Vertex AIのベクトル検索によってブログの検索エンジンを作成してみた

    はじめにこんにちは、石原と申します。こちらの記事は前…

  4. 未分類

    Databricksが買収した8080Labのbamboolibをひと足早く使って見る

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

  5. Generative AI

    AIを使ったマーケティングゲームを作ってみた

    こんにちは、石原です。私の所属しているエクスチュア株式会社で…

  6. Google Cloud Platform

    Vertex AI Embeddings for Text によるテキストエンベディングをやってみた…

    こんにちは、石原と申します。自然言語処理(NLP)は近年のA…

カテゴリ
最近の記事
  1. Looker Studioで日付フィルターの開始日・終了日の…
  2. Streamlitでdbt-elementary風ダッシュボ…
  3. Adobe WebSDK FPIDでECIDの復元を検証
  4. dbt Projects on Snowflakeで作成した…
  5. Dataformでtype:’increment…
  1. Adobe Analytics

    Adobe Analytics: SegmentsAPIを使って大量のセグメント…
  2. Adobe Analytics

    ページとは-Adobe Analyticsの指標説明
  3. IT用語集

    KPI、KGIって何?
  4. 赤の女王

    データサイエンス

    進化学に倣う最適化手法「遺伝的アルゴリズム」
  5. IT用語集

    メモリ(Memory)、仮想メモリ(Virtual Memory)って何?
PAGE TOP