Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/auth/team_grants.py: 61%

1"""Project a team row (plus the caller's membership in it) onto the ``team_*`` fields of ``UserAPIKeyAuth``. 

2 

3The virtual-key path gets these fields for free from the combined-view SQL join. Every other auth path 

4starts from a ``LiteLLM_TeamTable`` object instead and has to copy them over by hand, which is how JWT 

5callers kept losing grants (aliases, permissions, limits) one field at a time. Build the badge through 

6``team_grants`` and the two paths cannot drift. 

7""" 

8 

9from collections.abc import Mapping, Sequence 

10from types import MappingProxyType 

11from typing import Annotated, Final 

12 

13from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError 

14from pydantic.main import IncEx 

15from typing_extensions import ReadOnly, TypedDict 

16 

17from litellm.proxy._types import ( 

18 LiteLLM_ObjectPermissionTable, 

19 LiteLLM_TeamMembership, 

20 LiteLLM_TeamTable, 

21 Member, 

22) 

23 

24_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str]) 

25_JSON_COLUMNS: Final[Mapping[str, IncEx | bool]] = MappingProxyType( 

26 {"metadata": True, "litellm_model_table": MappingProxyType({"model_aliases": True})} 

27) 

28 

29 

30def _decode_model_aliases(value: object) -> object: 

31 """``LiteLLM_ModelTable.model_aliases`` is typed ``str | dict``; writers hand Prisma ``json.dumps(...)``, so take both.""" 

32 if not isinstance(value, str): 

33 return value 

34 try: 

35 return _MODEL_ALIASES_ADAPTER.validate_json(value) 

36 except ValidationError: 

37 return None 

38 

39 

40class TeamModelAliasTable(BaseModel): 

41 model_config = ConfigDict(protected_namespaces=()) 

42 

43 model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None 

44 

45 

46class _TeamJsonColumns(BaseModel): 

47 """The two loosely typed columns on ``LiteLLM_TeamTable``, re-read with the shape the badge needs.""" 

48 

49 metadata: Mapping[str, object] | None = None 

50 litellm_model_table: TeamModelAliasTable | None = None 

51 

52 

53class TeamGrants(TypedDict, total=False): 

54 """Keyword arguments for ``UserAPIKeyAuth``. Empty when the caller has no team, so the model's own defaults apply.""" 

55 

56 team_alias: ReadOnly[str | None] 

57 team_tpm_limit: ReadOnly[int | None] 

58 team_rpm_limit: ReadOnly[int | None] 

59 team_tpd_limit: ReadOnly[int | None] 

60 team_max_budget: ReadOnly[float | None] 

61 team_soft_budget: ReadOnly[float | None] 

62 team_model_max_budget: ReadOnly[dict[str, object] | None] 

63 team_spend: ReadOnly[float | None] 

64 team_models: ReadOnly[Sequence[str]] 

65 team_blocked: ReadOnly[bool] 

66 team_metadata: ReadOnly[Mapping[str, object] | None] 

67 team_model_aliases: ReadOnly[Mapping[str, str] | None] 

68 team_object_permission_id: ReadOnly[str | None] 

69 team_object_permission: ReadOnly[LiteLLM_ObjectPermissionTable | None] 

70 team_member: ReadOnly[Member | None] 

71 team_member_spend: ReadOnly[float | None] 

72 team_member_tpm_limit: ReadOnly[int | None] 

73 team_member_rpm_limit: ReadOnly[int | None] 

74 

75 

76def _json_columns(team_object: LiteLLM_TeamTable) -> _TeamJsonColumns: 

77 try: 

78 return _TeamJsonColumns.model_validate(team_object.model_dump(include=_JSON_COLUMNS)) 

79 except ValidationError: 

80 return _TeamJsonColumns() 

81 

82 

83def team_model_aliases(team_object: LiteLLM_TeamTable | None) -> Mapping[str, str] | None: 

84 if team_object is None: 

85 return None 

86 alias_table: Final = _json_columns(team_object).litellm_model_table 

87 return alias_table.model_aliases if alias_table is not None else None 

88 

89 

90def team_grants( 

91 team_object: LiteLLM_TeamTable | None, 

92 team_membership: LiteLLM_TeamMembership | None, 

93 user_id: str | None, 

94) -> TeamGrants: 

95 if team_object is None: 

96 return TeamGrants() 

97 json_columns: Final = _json_columns(team_object) 

98 return TeamGrants( 

99 team_alias=team_object.team_alias, 

100 team_tpm_limit=team_object.tpm_limit, 

101 team_rpm_limit=team_object.rpm_limit, 

102 team_tpd_limit=team_object.tpd_limit, 

103 team_max_budget=team_object.max_budget, 

104 team_soft_budget=team_object.soft_budget, 

105 team_model_max_budget=team_object.model_max_budget, 

106 team_spend=team_object.spend, 

107 team_models=tuple(team_object.models), 

108 team_blocked=team_object.blocked, 

109 team_metadata=json_columns.metadata, 

110 team_model_aliases=( 

111 json_columns.litellm_model_table.model_aliases if json_columns.litellm_model_table is not None else None 

112 ), 

113 team_object_permission_id=team_object.object_permission_id, 

114 team_object_permission=team_object.object_permission, 

115 team_member=next( 

116 (m for m in team_object.members_with_roles if user_id is not None and m.user_id == user_id), 

117 None, 

118 ), 

119 team_member_spend=team_membership.spend if team_membership is not None else None, 

120 team_member_tpm_limit=( 

121 team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None 

122 ), 

123 team_member_rpm_limit=( 

124 team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None 

125 ), 

126 )