跳到主要内容

SSO 登录事件钩子（自定义处理程序）

如果您想在用户使用 SSO 登录 LiteLLM 界面后运行自己的代码，请使用此功能。

工作原理​

  • 用户访问管理界面
  • LiteLLM 将用户重定向到您的 SSO 提供商
  • 您的 SSO 提供商将用户重定向回 LiteLLM
  • LiteLLM 已从您的 IDP 获取用户信息
  • 调用您的自定义 SSO 处理程序，并返回一个类型为 SSOUserDefinedValues 的对象
  • 用户登录界面成功

用法​

1. 创建自定义 SSO 处理程序文件。​

确保响应类型遵循 SSOUserDefinedValues pydantic 对象。这用于将用户登录到管理界面。

from fastapi import Request
from fastapi_sso.sso.base import OpenID

from litellm.proxy._types import LitellmUserRoles, SSOUserDefinedValues
from litellm.proxy.management_endpoints.internal_user_endpoints import (
    new_user,
    user_info,
)
from litellm.proxy.management_endpoints.team_endpoints import add_new_member


async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues:
    try:
        print("inside custom sso handler")  # noqa
        print(f"userIDPInfo: {userIDPInfo}")  # noqa

        if userIDPInfo.id is None:
            raise ValueError(
                f"No ID found for user. userIDPInfo.id is None {userIDPInfo}"
            )
        

        #################################################
        # Run you custom code / logic here
        # check if user exists in litellm proxy DB
        _user_info = await user_info(user_id=userIDPInfo.id)
        print("_user_info from litellm DB ", _user_info)  # noqa
        #################################################

        return SSOUserDefinedValues(
            models=[],                                      # models user has access to
            user_id=userIDPInfo.id,                         # user id to use in the LiteLLM DB
            user_email=userIDPInfo.email,                   # user email to use in the LiteLLM DB
            user_role=LitellmUserRoles.INTERNAL_USER.value, # role to use for the user 
            max_budget=0.01,                                # Max budget for this UI login Session
            budget_duration="1d",                           # Duration of the budget for this UI login Session, 1d, 2d, 30d ...
        )
    except Exception as e:
        raise Exception("Failed custom auth")

2. 传递文件路径（相对于 config.yaml）​

将文件路径传递给 config.yaml

例如，如果它们都在同一目录下 - ./config.yaml 和 ./custom_sso.py，如下所示

model_list: 
  - model_name: "openai-model"
    litellm_params: 
      model: "gpt-3.5-turbo"

litellm_settings:
  drop_params: True
  set_verbose: True

general_settings:
  custom_sso: custom_sso.custom_sso_handler

3. 启动代理​

$ litellm --config /path/to/config.yaml