diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/litellm/proxy/custom_sso.py')
-rw-r--r-- | .venv/lib/python3.12/site-packages/litellm/proxy/custom_sso.py | 44 |
1 files changed, 44 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/litellm/proxy/custom_sso.py b/.venv/lib/python3.12/site-packages/litellm/proxy/custom_sso.py new file mode 100644 index 00000000..210e9eea --- /dev/null +++ b/.venv/lib/python3.12/site-packages/litellm/proxy/custom_sso.py @@ -0,0 +1,44 @@ +""" +Example Custom SSO Handler + +Use this if you want to run custom code after litellm has retrieved information from your IDP (Identity Provider). + +Flow: +- User lands on Admin UI +- LiteLLM redirects user to your SSO provider +- Your SSO provider redirects user back to LiteLLM +- LiteLLM has retrieved user information from your IDP +- Your custom SSO handler is called and returns an object of type SSOUserDefinedValues +- User signed in to UI +""" + +from fastapi_sso.sso.base import OpenID + +from litellm.proxy._types import LitellmUserRoles, SSOUserDefinedValues +from litellm.proxy.management_endpoints.internal_user_endpoints import user_info + + +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}" + ) + + # 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=[], + user_id=userIDPInfo.id, + user_email=userIDPInfo.email, + user_role=LitellmUserRoles.INTERNAL_USER.value, + max_budget=10, + budget_duration="1d", + ) + except Exception: + raise Exception("Failed custom auth") |