From 4a52a71956a8d46fcb7294ac71734504bb09bcc2 Mon Sep 17 00:00:00 2001 From: S. Solomon Darnell Date: Fri, 28 Mar 2025 21:52:21 -0500 Subject: two version of R2R are here --- .../site-packages/litellm/proxy/custom_sso.py | 44 ++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 .venv/lib/python3.12/site-packages/litellm/proxy/custom_sso.py (limited to '.venv/lib/python3.12/site-packages/litellm/proxy/custom_sso.py') 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") -- cgit v1.2.3