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. boxMCP

    Claude

    boxMCPサーバーを使ってみた

    こんにちは、中村です。最近なにかと話題のMCPサーバーですが…

  2. ChatGPT

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

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

  3. KARTE

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

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

  4. Generative AI

    VScode拡張機能「Cline」を利用してMCPのツール利用する方法

    こんにちは、エクスチュアの石原です。今回はVScodeの拡張…

  5. Python

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

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

  6. Google Cloud Platform

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

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

カテゴリ
最近の記事
  1. dbt Projects on SnowflakeをTASK…
  2. AWS発のAIエージェントIDE「Kiro」を使用した仕様駆…
  3. AWS発のAIエージェントIDE「Kiro」を使用した仕様駆…
  4. TableauとSnowflakeを接続する方法
  5. 【dbts25】Snowflake×PostgreSQLのニ…
  1. Snowflake

    SnowPro Associate: Platform 合格体験記
  2. KARTE

    CXツール「KARTE」ってどんなツール?主な機能や魅力をざっくり紹介!
  3. ヒートマップ

    ムーブメント・ヒートマップ
  4. ObservePoint

    ObservePointの「Journey」とは?
  5. IT用語集

    テーブル(Table)・カラム(Column)・ロー(Row)・フィールド(Fi…
PAGE TOP