Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/custom_sso.py: 0%

1""" 

2Example Custom SSO Handler 

3 

4Use this if you want to run custom code after litellm has retrieved information from your IDP (Identity Provider). 

5 

6Flow: 

7- User lands on Admin UI 

8- LiteLLM redirects user to your SSO provider 

9- Your SSO provider redirects user back to LiteLLM 

10- LiteLLM has retrieved user information from your IDP 

11- Your custom SSO handler is called and returns an object of type SSOUserDefinedValues 

12- User signed in to UI 

13""" 

14 

15from typing import Final 

16 

17from fastapi_sso.sso.base import OpenID 

18 

19from litellm.proxy import proxy_server 

20from litellm.proxy._types import LitellmUserRoles, SSOUserDefinedValues 

21 

22 

23async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: 

24 try: 

25 if userIDPInfo.id is None: 

26 raise ValueError(f"No ID found for user. userIDPInfo.id is None {userIDPInfo}") 

27 

28 # Access extra fields from the IDP response (requires GENERIC_USER_EXTRA_ATTRIBUTES env var) 

29 # Example: Set GENERIC_USER_EXTRA_ATTRIBUTES="group,NTID,domain" to capture these fields 

30 # extra_fields = getattr(userIDPInfo, 'extra_fields', None) or {} 

31 # user_groups = extra_fields.get("group", []) 

32 

33 # check if user exists in litellm proxy DB 

34 if proxy_server.prisma_client is not None: 

35 _user_info: Final = await proxy_server.prisma_client.get_data(user_id=userIDPInfo.id) 

36 

37 return SSOUserDefinedValues( 

38 models=[], 

39 user_id=userIDPInfo.id, 

40 user_email=userIDPInfo.email, 

41 user_role=LitellmUserRoles.INTERNAL_USER.value, 

42 max_budget=10, 

43 budget_duration="1d", 

44 ) 

45 except Exception: 

46 raise Exception("Failed custom auth")