Skip to main content

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

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

工作原理​

  • 用户访问管理 UI
  • LiteLLM 将用户重定向到你的 SSO 提供商
  • SSO 提供商将用户重定向回 LiteLLM
  • LiteLLM 已从你的身份提供者（IDP）获取用户信息
  • 调用你的自定义 SSO 处理程序并返回一个 SSOUserDefinedValues 类型的对象
  • 用户登录到 UI

使用方法​

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

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

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}"
            )
        

        #################################################
        # 在这里运行你的自定义代码 / 逻辑
        # 检查用户是否存在于 litellm 代理数据库中
        _user_info = await user_info(user_id=userIDPInfo.id)
        print("_user_info from litellm DB ", _user_info)  # noqa
        #################################################

        return SSOUserDefinedValues(
            models=[],                                      # 用户可以访问的模型
            user_id=userIDPInfo.id,                         # 用于 LiteLLM 数据库的用户 ID
            user_email=userIDPInfo.email,                   # 用于 LiteLLM 数据库的用户邮箱
            user_role=LitellmUserRoles.INTERNAL_USER.value, # 用户角色
            max_budget=0.01,                                # 此 UI 登录会话的最大预算
            budget_duration="1d",                           # 此 UI 登录会话的预算持续时间，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
优云智算