Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py: 88%

1import asyncio 

2import json 

3from typing import Annotated, Final, Literal 

4 

5from fastapi import APIRouter, Depends, HTTPException 

6from pydantic import BaseModel, Field, ValidationError, model_validator 

7 

8from litellm._uuid import uuid4 

9from litellm.proxy._types import LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth 

10from litellm.proxy.auth.user_api_key_auth import user_api_key_auth 

11from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerRepository 

12 

13router: Final = APIRouter() 

14 

15USER_BANNER_MAX_MESSAGE_LENGTH: Final = 4000 

16 

17UserBannerSeverity = Literal["info", "warning", "error"] 

18 

19 

20class UserBannerUpdate(BaseModel): 

21 enabled: bool = Field( 

22 default=False, 

23 description="If true, the banner is shown to all authenticated dashboard users.", 

24 ) 

25 message: str = Field( 

26 default="", 

27 max_length=USER_BANNER_MAX_MESSAGE_LENGTH, 

28 description="Banner text shown to dashboard users. Markdown is supported.", 

29 ) 

30 severity: UserBannerSeverity = Field( 

31 default="info", 

32 description="Visual style of the banner.", 

33 ) 

34 

35 @model_validator(mode="after") 

36 def _require_message_when_enabled(self) -> "UserBannerUpdate": 

37 if self.enabled and not self.message.strip(): 

38 raise ValueError("message must be non-empty when the banner is enabled") 

39 return self 

40 

41 

42class UserBanner(UserBannerUpdate): 

43 revision: str = Field( 

44 default="", 

45 description=( 

46 "Server-stamped opaque publish identity; a fresh value is generated on every " 

47 "update so clients re-surface dismissed banners on republish." 

48 ), 

49 ) 

50 

51 

52class UpdateUserBannerResponse(BaseModel): 

53 message: str 

54 banner: UserBanner 

55 

56 

57def parse_user_banner(raw_settings: object) -> UserBanner: 

58 if raw_settings is None: 

59 return UserBanner() 

60 try: 

61 parsed: Final = json.loads(raw_settings) if isinstance(raw_settings, str) else raw_settings 

62 return UserBanner.model_validate(parsed) 

63 except (json.JSONDecodeError, ValidationError): 

64 return UserBanner() 

65 

66 

67@router.get( 

68 "/get/user_banner", 

69 tags=["UI Settings"], # mutable-ok: FastAPI's route decorator only accepts a list 

70 dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI's route decorator only accepts a list 

71 response_model=UserBanner, 

72) 

73async def get_user_banner() -> UserBanner: 

74 """ 

75 Get the admin-published dashboard banner. 

76 Readable by any authenticated user; rendered on every dashboard page. 

77 """ 

78 from litellm.proxy.proxy_server import prisma_client 

79 

80 if prisma_client is None: 80 ↛ 81line 80 didn't jump to line 81 because the condition on line 80 was never true

81 return UserBanner() 

82 

83 raw_settings: Final = await UserBannerRepository(prisma_client).get_raw_settings() 

84 return parse_user_banner(raw_settings) 

85 

86 

87@router.patch( 

88 "/update/user_banner", 

89 tags=["UI Settings"], # mutable-ok: FastAPI's route decorator only accepts a list 

90 response_model=UpdateUserBannerResponse, 

91) 

92async def update_user_banner( 

93 banner_update: UserBannerUpdate, 

94 user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], 

95) -> UpdateUserBannerResponse: 

96 """ 

97 Publish, edit, or unpublish the dashboard banner. 

98 Only proxy admins are allowed to modify it. 

99 """ 

100 from litellm.proxy.proxy_server import create_config_audit_log, prisma_client 

101 

102 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 102 ↛ 103line 102 didn't jump to line 103 because the condition on line 102 was never true

103 raise HTTPException(status_code=403, detail="Only proxy admins can update the user banner.") 

104 

105 if prisma_client is None: 105 ↛ 106line 105 didn't jump to line 106 because the condition on line 105 was never true

106 raise HTTPException(status_code=500, detail="Database not connected. Please connect a database.") 

107 

108 repository: Final = UserBannerRepository(prisma_client) 

109 before: Final = parse_user_banner(await repository.get_raw_settings()) 

110 banner: Final = UserBanner( 

111 enabled=banner_update.enabled, 

112 message=banner_update.message, 

113 severity=banner_update.severity, 

114 revision=uuid4().hex, 

115 ) 

116 

117 await repository.upsert_settings(json.dumps(banner.model_dump())) 

118 

119 asyncio.create_task( 

120 create_config_audit_log( 

121 param_name=USER_BANNER_ROW_ID, 

122 action="updated", 

123 before_value=before.model_dump(), 

124 after_value=banner.model_dump(), 

125 user_api_key_dict=user_api_key_dict, 

126 table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME, 

127 ) 

128 ) 

129 

130 return UpdateUserBannerResponse(message="User banner updated successfully", banner=banner)