Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_helpers/auto_router_availability.py: 91%

1from collections.abc import Mapping, Sequence 

2from copy import deepcopy 

3from dataclasses import dataclass 

4from types import MappingProxyType 

5from typing import Final 

6 

7from pydantic import BaseModel, Json, TypeAdapter, ValidationError 

8 

9from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper 

10from litellm.router_utils.auto_router_model_naming import ( 

11 GATED_AUTO_ROUTER_CAPABILITIES, 

12 capability_limit_violation, 

13 classify_strategy_router_model, 

14 count_capability_routers, 

15 gated_capability_of, 

16) 

17from litellm.router_utils.auto_router_tuning_baseline import ( 

18 is_mutable_tuned_candidate, 

19 mutable_tuned_identities, 

20 tuning_quota_violation, 

21) 

22from litellm.types.management_endpoints.auto_router_endpoints import ( 

23 AutoRouterAllowance, 

24 AutoRouterAvailabilityResponse, 

25) 

26 

27 

28class _CatalogModelInfo(BaseModel): 

29 team_id: str | None = None 

30 

31 

32class _CatalogSource(BaseModel): 

33 model_id: str 

34 created_by: str | None = None 

35 litellm_params: Json[dict[str, object]] | dict[str, object] 

36 model_info: Json[_CatalogModelInfo] | _CatalogModelInfo | None = None 

37 

38 

39@dataclass(frozen=True, slots=True) 

40class AutoRouterCatalogEntry: 

41 model_id: str 

42 team_id: str | None 

43 created_by: str | None 

44 deployment: Mapping[str, object] 

45 

46 

47def _catalog_field(value: object, key: str) -> object: 

48 if not isinstance(value, str): 48 ↛ 49line 48 didn't jump to line 49 because the condition on line 48 was never true

49 return deepcopy(value) 

50 return decrypt_value_helper(value, key=key, exception_type="debug", return_original_value=True) 

51 

52 

53def build_auto_router_catalog(rows: Sequence[object]) -> tuple[AutoRouterCatalogEntry, ...] | None: 

54 try: 

55 sources: Final = TypeAdapter(tuple[_CatalogSource, ...]).validate_python(rows, from_attributes=True) 

56 except ValidationError: 

57 return None 

58 return tuple( 

59 AutoRouterCatalogEntry( 

60 model_id=row.model_id, 

61 team_id=row.model_info.team_id if row.model_info is not None else None, 

62 created_by=row.created_by, 

63 deployment=MappingProxyType( 

64 { 

65 "litellm_params": MappingProxyType( 

66 { 

67 "model": model, 

68 "complexity_router_config": _catalog_field( 

69 row.litellm_params.get("complexity_router_config"), "complexity_router_config" 

70 ), 

71 } 

72 ), 

73 "model_info": MappingProxyType({"id": row.model_id, "db_model": True}), 

74 } 

75 ), 

76 ) 

77 for row in sources 

78 if isinstance(model := _catalog_field(row.litellm_params.get("model"), "model"), str) 

79 and classify_strategy_router_model(model) == "complexity" 

80 ) 

81 

82 

83def auto_router_availability( 

84 *, 

85 others: Sequence[Mapping[str, object]], 

86 existing: Mapping[str, object] | None, 

87 candidate: Mapping[str, object], 

88 baselines: Mapping[str, str] | None, 

89 limit: int | None, 

90) -> AutoRouterAvailabilityResponse: 

91 existing_params: Final = None if existing is None else existing.get("litellm_params") 

92 candidate_params: Final = candidate.get("litellm_params") 

93 owned: Final = gated_capability_of(existing_params) if isinstance(existing_params, Mapping) else None 

94 claimed: Final = gated_capability_of(candidate_params) if isinstance(candidate_params, Mapping) else None 

95 counts: Final = tuple( 

96 (capability, count_capability_routers(others, capability=capability)) 

97 for capability in GATED_AUTO_ROUTER_CAPABILITIES 

98 ) 

99 tuned_count: Final = len(mutable_tuned_identities(others, baselines)) if baselines is not None else 0 

100 allowances: Final = tuple( 

101 AutoRouterAllowance( 

102 key=capability.key, 

103 limit=limit, 

104 remaining=None if limit is None else max(0, limit - held), 

105 used_by_this_router=owned is capability, 

106 ) 

107 for capability, held in counts 

108 ) 

109 capability_error: Final = next( 

110 ( 

111 capability_limit_violation(capability=capability, held=held + 1, limit=limit) 

112 for capability, held in counts 

113 if capability is claimed 

114 ), 

115 None, 

116 ) 

117 tuning_error: Final = ( 

118 tuning_quota_violation(candidate=candidate, others=others, baselines=baselines, limit=limit) 

119 if baselines is not None 

120 else None 

121 ) 

122 capability_labels: Final = { 

123 "heuristic_v2": "Heuristic v2", 

124 "capability": "Capability", 

125 "llm_v2": "Fuse v2", 

126 "tier_or_classifier_prompt": "Custom tiers or classifier instructions", 

127 } 

128 return AutoRouterAvailabilityResponse( 

129 allowances=( 

130 *allowances, 

131 AutoRouterAllowance( 

132 key="heuristic_tuning", 

133 limit=limit, 

134 remaining=None if limit is None or baselines is None else max(0, limit - tuned_count), 

135 available=limit is None or baselines is not None, 

136 used_by_this_router=bool( 

137 existing is not None and baselines is not None and is_mutable_tuned_candidate(existing, baselines) 

138 ), 

139 ), 

140 ), 

141 error=( 

142 f"{capability_labels[claimed.key]} has no available allowance. Choose another option or free an existing allowance." 

143 if capability_error is not None and claimed is not None 

144 else "These scoring rules need an available Rule-based tuning allowance. Check the weights, thresholds, keywords, and custom dimensions in Advanced settings. Model choices do not use this allowance." 

145 if tuning_error is not None 

146 else None 

147 ), 

148 )