Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_helpers/auto_router_permissions.py: 40%
161 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1from collections.abc import Mapping
2from dataclasses import dataclass
3from datetime import datetime
4from types import MappingProxyType
5from typing import TYPE_CHECKING, Final, Literal
7from fastapi import HTTPException
8from pydantic import BaseModel, ConfigDict, Field, ValidationError
9from typing_extensions import ReadOnly, TypedDict
11from litellm.models.organization import LiteLLM_OrganizationTable
12from litellm.models.project import LiteLLM_ProjectTable
13from litellm.proxy._types import (
14 UI_TEAM_ID,
15 CommonProxyErrors,
16 KeyManagementRoutes,
17 LiteLLM_TeamMembership,
18 LiteLLM_TeamTable,
19 LitellmUserRoles,
20 UserAPIKeyAuth,
21)
22from litellm.proxy.auth.auth_checks import (
23 _check_team_member_model_access, # pyright: ignore[reportPrivateUsage] # shared membership authorization owner
24 can_key_call_model,
25 can_org_access_model,
26 can_project_access_model,
27 can_team_access_model,
28)
29from litellm.proxy.auth.team_grants import team_model_aliases
30from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
31from litellm.repositories.organization_repository import OrganizationRepository
32from litellm.repositories.prisma_protocols import DatabaseClient
33from litellm.repositories.project_repository import ProjectRepository
34from litellm.repositories.table_repositories import TeamMembershipRepository
35from litellm.router import Router
36from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model, strategy_router_dependencies
37from litellm.types.management_endpoints.auto_router_endpoints import RequestComplexityRouterConfig
38from litellm.types.router import Deployment, updateDeployment
40if TYPE_CHECKING: 40 ↛ 41line 40 didn't jump to line 41 because the condition on line 40 was never true
41 from prisma import types as prisma_types
44class _MemberRouterThinking(BaseModel):
45 model_config = ConfigDict(extra="forbid", frozen=True)
47 type: Literal["enabled", "disabled", "adaptive"]
48 budget_tokens: int | None = Field(default=None, gt=0, le=1_000_000)
51class _MemberRouterGenerationParams(BaseModel):
52 model_config = ConfigDict(extra="forbid", frozen=True)
54 reasoning_effort: str | None = None
55 thinking: _MemberRouterThinking | None = None
56 verbosity: Literal["low", "medium", "high"] | None = None
57 max_tokens: int | None = Field(default=None, gt=0, le=1_000_000)
58 max_completion_tokens: int | None = Field(default=None, gt=0, le=1_000_000)
59 max_output_tokens: int | None = Field(default=None, gt=0, le=1_000_000)
60 temperature: float | None = Field(default=None, ge=0, le=2, allow_inf_nan=False)
61 top_p: float | None = Field(default=None, ge=0, le=1, allow_inf_nan=False)
62 frequency_penalty: float | None = Field(default=None, ge=-2, le=2, allow_inf_nan=False)
63 presence_penalty: float | None = Field(default=None, ge=-2, le=2, allow_inf_nan=False)
64 seed: int | None = None
65 stop: str | tuple[str, ...] | None = None
68class _MemberJevClassifierConfig(BaseModel):
69 """The Jev classifier settings a team member may set. Credentials stay the proxy's own: a member-chosen
70 api_base would receive the proxy's TYPESAFE_API_KEY, and a member-chosen api_key would be sent from the proxy."""
72 model_config = ConfigDict(extra="forbid")
74 model: str
75 api_key: None = None
76 api_base: None = None
77 timeout_ms: int
78 instructions: str | None = None
79 circuit_breaker_enabled: bool
80 circuit_breaker_cooldown_seconds: float
83class _MemberComplexityRouterConfig(RequestComplexityRouterConfig):
84 model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
87class _RouterConfigSource(BaseModel):
88 model: str | None = None
89 complexity_router_config: Mapping[str, object] | None = None
92class _MembershipKey(TypedDict):
93 user_id: ReadOnly[str]
94 team_id: ReadOnly[str]
97class _MembershipWhere(TypedDict):
98 user_id_team_id: ReadOnly[_MembershipKey]
101@dataclass(frozen=True, slots=True)
102class MemberAutoRouterDependencyObjects:
103 membership: LiteLLM_TeamMembership | None
104 organization: LiteLLM_OrganizationTable | None
105 project: LiteLLM_ProjectTable | None
108def authorize_member_auto_router_team(
109 *, user_api_key_dict: UserAPIKeyAuth, team: LiteLLM_TeamTable, premium_user: bool
110) -> None:
111 if not premium_user:
112 raise HTTPException(status_code=403, detail=CommonProxyErrors.not_premium_user.value)
113 if (
114 user_api_key_dict.user_role
115 not in (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.TEAM, LitellmUserRoles.ORG_ADMIN)
116 or not user_api_key_dict.user_id
117 or not any(member.user_id == user_api_key_dict.user_id for member in team.members_with_roles)
118 or user_api_key_dict.team_id not in (None, UI_TEAM_ID, team.team_id)
119 or team.blocked
120 or KeyManagementRoutes.AUTO_ROUTER_MANAGE.value not in (team.team_member_permissions or ())
121 ):
122 raise HTTPException(status_code=403, detail="This team does not allow you to manage your own auto routers.")
125def validate_member_auto_router_config(config: Mapping[str, object]) -> RequestComplexityRouterConfig:
126 try:
127 validated: Final = _MemberComplexityRouterConfig.model_validate(config)
128 for entries in validated.tier_model_configs.values():
129 for entry in entries:
130 _MemberRouterGenerationParams.model_validate(entry.litellm_params)
131 if validated.jev_classifier_config is not None:
132 _MemberJevClassifierConfig.model_validate(validated.jev_classifier_config.model_dump())
133 return validated
134 except ValidationError as exc:
135 location: Final = ".".join(str(part) for part in exc.errors()[0]["loc"])
136 raise HTTPException(status_code=400, detail=f"Invalid member auto-router configuration at {location}.") from exc
139async def authorize_member_auto_router_dependencies(
140 *,
141 config: RequestComplexityRouterConfig,
142 default_model: str | None,
143 user_api_key_dict: UserAPIKeyAuth,
144 team: LiteLLM_TeamTable,
145 prisma_client: DatabaseClient | None,
146 llm_router: Router,
147 dependency_objects: MemberAutoRouterDependencyObjects | None = None,
148) -> None:
149 from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
151 if team.blocked:
152 raise HTTPException(status_code=403, detail="This auto router's team is blocked.")
153 aliases: Final = team_model_aliases(team)
154 alias_dict: Final = (
155 dict(aliases) if aliases is not None else None # mutable-ok: auth model and helpers require dict
156 )
157 scoped_actor: Final = user_api_key_dict.model_copy(
158 update=MappingProxyType({"team_id": team.team_id, "team_models": team.models, "team_model_aliases": alias_dict})
159 )
160 objects: Final = (
161 dependency_objects
162 if dependency_objects is not None
163 else await _load_member_auto_router_dependency_objects(
164 user_api_key_dict=scoped_actor, team=team, prisma_client=prisma_client
165 )
166 )
167 if team.organization_id and objects.organization is None:
168 raise HTTPException(status_code=403, detail="The auto router's organization is unavailable.")
169 if scoped_actor.project_id and (
170 objects.project is None or objects.project.team_id != team.team_id or objects.project.blocked
171 ):
172 raise HTTPException(status_code=403, detail="The auto router's project is unavailable.")
173 dependencies: Final = strategy_router_dependencies(
174 MappingProxyType(
175 {
176 "model": "auto_router/complexity_router",
177 "complexity_router_config": config.model_dump(exclude_none=True),
178 "complexity_router_default_model": default_model,
179 }
180 )
181 )
182 for dependency, model, deployments in (
183 (
184 dependency,
185 dependency.model_name,
186 llm_router.get_model_list(model_name=dependency.model_name, team_id=team.team_id),
187 )
188 for dependency in dependencies
189 ):
190 if dependency.role != "evaluation" and (
191 not deployments
192 or any(
193 classify_strategy_router_model(
194 _RouterConfigSource.model_validate(deployment["litellm_params"]).model or ""
195 )
196 is not None
197 for deployment in deployments
198 )
199 ):
200 raise HTTPException(status_code=400, detail=f"Auto-router target {model!r} must be a configured model.")
201 await can_team_access_model(
202 model=model,
203 team_object=team,
204 llm_router=llm_router,
205 team_model_aliases=alias_dict,
206 prisma_client=prisma_client,
207 )
208 await can_key_call_model(
209 model=model,
210 llm_model_list=None,
211 valid_token=scoped_actor,
212 llm_router=llm_router,
213 prisma_client=prisma_client,
214 )
215 await _check_team_member_model_access(
216 model=model,
217 team_object=team,
218 valid_token=scoped_actor,
219 llm_router=llm_router,
220 prisma_client=None,
221 user_api_key_cache=user_api_key_cache,
222 proxy_logging_obj=proxy_logging_obj,
223 team_membership=objects.membership,
224 team_membership_loaded=True,
225 )
226 if objects.organization is not None:
227 can_org_access_model(model=model, org_object=objects.organization, llm_router=llm_router)
228 if objects.project is not None:
229 can_project_access_model(model=model, project_object=objects.project, llm_router=llm_router)
232async def _load_member_auto_router_dependency_objects(
233 *, user_api_key_dict: UserAPIKeyAuth, team: LiteLLM_TeamTable, prisma_client: DatabaseClient | None
234) -> MemberAutoRouterDependencyObjects:
235 if prisma_client is None:
236 raise HTTPException(status_code=503, detail="Cannot verify auto-router model access without a database")
237 membership_where: Final[_MembershipWhere] = {
238 "user_id_team_id": {"user_id": user_api_key_dict.user_id or "", "team_id": team.team_id}
239 }
240 membership_include: Final[prisma_types.LiteLLM_TeamMembershipInclude] = {"litellm_budget_table": True}
241 membership_row: Final = (
242 await TeamMembershipRepository(prisma_client).table.find_unique(
243 where=membership_where, include=membership_include
244 )
245 if user_api_key_dict.user_id
246 else None
247 )
248 membership: Final = (
249 LiteLLM_TeamMembership.model_validate(membership_row.model_dump()) if membership_row is not None else None
250 )
251 organization: Final = (
252 await OrganizationRepository(prisma_client).find_by_id(team.organization_id) if team.organization_id else None
253 )
254 if team.organization_id and organization is None:
255 raise HTTPException(status_code=403, detail="The auto router's organization is unavailable.")
256 project: Final = (
257 await ProjectRepository(prisma_client).find_by_id(user_api_key_dict.project_id)
258 if user_api_key_dict.project_id
259 else None
260 )
261 return MemberAutoRouterDependencyObjects(membership=membership, organization=organization, project=project)
264class StoredAutoRouterIdentity(BaseModel):
265 created_by: str | None = None
266 updated_at: datetime | None = None
269@dataclass(frozen=True, slots=True)
270class MemberAutoRouterWrite:
271 actor: UserAPIKeyAuth
272 team_id: str
273 model_id: str | None
274 public_name: str
275 updated_at: datetime | None
276 config: RequestComplexityRouterConfig
277 default_model: str | None
280async def authorize_member_auto_router_write(
281 *,
282 incoming: Deployment | updateDeployment,
283 existing: Deployment | None,
284 user_api_key_dict: UserAPIKeyAuth,
285 team: LiteLLM_TeamTable,
286 premium_user: bool,
287 prisma_client: DatabaseClient,
288 llm_router: Router,
289) -> MemberAutoRouterWrite:
290 authorize_member_auto_router_team(user_api_key_dict=user_api_key_dict, team=team, premium_user=premium_user)
291 stored: Final = StoredAutoRouterIdentity.model_validate(existing.model_dump()) if existing is not None else None
292 if stored is not None and stored.created_by != user_api_key_dict.user_id:
293 raise HTTPException(status_code=403, detail="Team members can update only their own auto routers.")
294 params: Final = incoming.litellm_params
295 if params is None or incoming.model_fields_set - frozenset({"model_name", "litellm_params", "model_info"}):
296 raise HTTPException(status_code=403, detail="Team members may change only auto-router configuration.")
297 if params.model_fields_set - frozenset({"model", "complexity_router_config", "complexity_router_default_model"}):
298 raise HTTPException(status_code=403, detail="Team members may change only auto-router configuration.")
299 info: Final = incoming.model_info
300 if info is not None and (
301 info.model_fields_set - frozenset({"id", "team_id"})
302 or info.team_id not in (None, team.team_id)
303 or (existing is not None and "id" in info.model_fields_set and info.id != existing.model_info.id)
304 ):
305 raise HTTPException(
306 status_code=403, detail="Team members cannot change model ownership or administrative settings."
307 )
308 existing_model: Final = (
309 decrypt_value_helper(existing.litellm_params.model, key="model", return_original_value=True)
310 if existing is not None
311 else None
312 )
313 effective_model: Final = params.model or existing_model
314 if (
315 not isinstance(effective_model, str)
316 or classify_strategy_router_model(effective_model) != "complexity"
317 or (existing is not None and effective_model != existing_model)
318 ):
319 raise HTTPException(status_code=403, detail="Team members may manage only complexity auto routers.")
320 public_name: Final = (
321 existing.model_info.team_public_model_name or existing.model_name
322 if existing is not None
323 else incoming.model_name
324 )
325 if (
326 not public_name
327 or public_name != public_name.strip()
328 or any(character in public_name for character in "*?[]")
329 or public_name.startswith("model_name_")
330 ):
331 raise HTTPException(
332 status_code=400, detail="Choose a non-empty auto-router name without wildcards or internal prefixes."
333 )
334 if existing is not None and incoming.model_name not in (None, public_name, existing.model_name):
335 raise HTTPException(status_code=403, detail="Team members cannot rename an auto router.")
336 supplied_config: Final = _RouterConfigSource.model_validate(params.model_dump()).complexity_router_config
337 raw_config: Final = (
338 supplied_config
339 if supplied_config is not None
340 else _RouterConfigSource.model_validate(existing.litellm_params.model_dump()).complexity_router_config
341 if existing is not None
342 else None
343 )
344 if raw_config is None:
345 raise HTTPException(status_code=400, detail="A complexity_router_config is required.")
346 config: Final = validate_member_auto_router_config(raw_config)
347 stored_default: Final = existing.litellm_params.complexity_router_default_model if existing is not None else None
348 default_model: Final = (
349 params.complexity_router_default_model
350 if params.complexity_router_default_model is not None
351 else decrypt_value_helper(stored_default, key="complexity_router_default_model", return_original_value=True)
352 if stored_default is not None
353 else None
354 )
355 await authorize_member_auto_router_dependencies(
356 config=config,
357 default_model=default_model,
358 user_api_key_dict=user_api_key_dict,
359 team=team,
360 prisma_client=prisma_client,
361 llm_router=llm_router,
362 )
363 return MemberAutoRouterWrite(
364 actor=user_api_key_dict,
365 team_id=team.team_id,
366 model_id=existing.model_info.id if existing is not None else None,
367 public_name=public_name,
368 updated_at=stored.updated_at if stored is not None else None,
369 config=config,
370 default_model=default_model,
371 )