From 33d2a2dbc32c3d359484214e3ad05cd8463dcf98 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 23 Sep 2024 17:00:16 -0700 Subject: [PATCH] store sso provider in user metadata --- litellm/proxy/management_endpoints/ui_sso.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 7d21b201531..073ea8466be 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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 = {"user_identity_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