[Feat UI sso] store 'provider' in user metadata (#5856)

* store sso provider in user metadata

* store user metadata

* store user auth_provider in user metadata

* add "metadata" for LiteLLM_UserTable

* fix sso test
This commit is contained in:
Ishaan Jaff 2024-09-23 17:49:36 -07:00 committed by GitHub
parent 922c8ac758
commit 391b107909
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 16 additions and 0 deletions

View file

@ -12,6 +12,7 @@ from typing import List, Optional
from fastapi import APIRouter, Depends, HTTPException, Request, status
from fastapi.responses import RedirectResponse
from fastapi_sso.sso.base import OpenID
import litellm
from litellm._logging import verbose_proxy_logger
@ -498,6 +499,7 @@ async def auth_callback(request: Request):
else:
# user not in DB, insert User into LiteLLM DB
user_role = await insert_sso_user(
result_openid=result,
user_defined_values=user_defined_values,
)
except Exception as e:
@ -575,10 +577,15 @@ async def auth_callback(request: Request):
async def insert_sso_user(
result_openid: Optional[OpenID],
user_defined_values: Optional[SSOUserDefinedValues] = None,
) -> str:
"""
Helper function to create a New User in LiteLLM DB after a successful SSO login
Args:
result_openid (OpenID): User information in OpenID format if the login was successful.
user_defined_values (Optional[SSOUserDefinedValues], optional): LiteLLM SSOValues / fields that were read
"""
verbose_proxy_logger.debug(
f"Inserting SSO user into DB. User values: {user_defined_values}"
@ -610,6 +617,9 @@ async def insert_sso_user(
budget_duration=user_defined_values["budget_duration"],
)
if result_openid:
new_user_request.metadata = {"auth_provider": result_openid.provider}
await new_user(data=new_user_request, user_api_key_dict=UserAPIKeyAuth())
return user_defined_values["user_role"] or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY