""" User repository for database operations on LiteLLM_UserTable. """ import json from typing import Any, Dict, List, Optional, Type from litellm.models.user import LiteLLM_UserTable from litellm.repositories.base_repository import BaseRepository class UserRepository(BaseRepository[LiteLLM_UserTable]): """Repository for user database operations.""" @property def table(self) -> Any: return self.prisma_client.db.litellm_usertable @property def model_class(self) -> Type[LiteLLM_UserTable]: return LiteLLM_UserTable def _to_model(self, record: Any) -> Optional[LiteLLM_UserTable]: """Convert a database record to a User model.""" if record is None: return None data = record.dict() if hasattr(record, "dict") else dict(record) json_fields = ["metadata", "model_spend", "model_max_budget"] for field in json_fields: if isinstance(data.get(field), str): data[field] = json.loads(data[field]) return LiteLLM_UserTable(**data) async def find_by_id( self, user_id: str, id_field: str = "user_id" ) -> Optional[LiteLLM_UserTable]: return await super().find_by_id(user_id, id_field) async def find_by_email(self, user_email: str) -> Optional[LiteLLM_UserTable]: """Find a user by email.""" records = await self.table.find_many(where={"user_email": user_email}) if records: return self._to_model(records[0]) return None async def find_by_sso_id(self, sso_user_id: str) -> Optional[LiteLLM_UserTable]: """Find a user by SSO ID.""" record = await self.table.find_unique(where={"sso_user_id": sso_user_id}) return self._to_model(record) async def find_by_organization_id( self, organization_id: str ) -> List[LiteLLM_UserTable]: """Find all users in an organization.""" records = await self.table.find_many(where={"organization_id": organization_id}) return self._to_model_list(records) async def find_by_team_id(self, team_id: str) -> List[LiteLLM_UserTable]: """Find all users in a team.""" records = await self.table.find_many(where={"teams": {"has": team_id}}) return self._to_model_list(records) async def create_user( self, user_id: str, user_alias: Optional[str] = None, team_id: Optional[str] = None, sso_user_id: Optional[str] = None, organization_id: Optional[str] = None, password: Optional[str] = None, teams: Optional[List[str]] = None, user_role: Optional[str] = None, max_budget: Optional[float] = None, user_email: Optional[str] = None, models: Optional[List[str]] = None, metadata: Optional[Dict[str, Any]] = None, max_parallel_requests: Optional[int] = None, tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, budget_duration: Optional[str] = None, allowed_cache_controls: Optional[List[str]] = None, policies: Optional[List[str]] = None, object_permission_id: Optional[str] = None, ) -> LiteLLM_UserTable: """Create a new user.""" data: Dict[str, Any] = {"user_id": user_id} if user_alias is not None: data["user_alias"] = user_alias if team_id is not None: data["team_id"] = team_id if sso_user_id is not None: data["sso_user_id"] = sso_user_id if organization_id is not None: data["organization_id"] = organization_id if password is not None: data["password"] = password if teams is not None: data["teams"] = teams if user_role is not None: data["user_role"] = user_role if max_budget is not None: data["max_budget"] = max_budget if user_email is not None: data["user_email"] = user_email if models is not None: data["models"] = models if metadata is not None: data["metadata"] = json.dumps(metadata) if max_parallel_requests is not None: data["max_parallel_requests"] = max_parallel_requests if tpm_limit is not None: data["tpm_limit"] = tpm_limit if rpm_limit is not None: data["rpm_limit"] = rpm_limit if budget_duration is not None: data["budget_duration"] = budget_duration if allowed_cache_controls is not None: data["allowed_cache_controls"] = allowed_cache_controls if policies is not None: data["policies"] = policies if object_permission_id is not None: data["object_permission_id"] = object_permission_id return await self.create(data) async def update_user( self, user_id: str, user_alias: Optional[str] = None, team_id: Optional[str] = None, sso_user_id: Optional[str] = None, organization_id: Optional[str] = None, password: Optional[str] = None, teams: Optional[List[str]] = None, user_role: Optional[str] = None, max_budget: Optional[float] = None, user_email: Optional[str] = None, models: Optional[List[str]] = None, metadata: Optional[Dict[str, Any]] = None, max_parallel_requests: Optional[int] = None, tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, budget_duration: Optional[str] = None, allowed_cache_controls: Optional[List[str]] = None, policies: Optional[List[str]] = None, object_permission_id: Optional[str] = None, ) -> Optional[LiteLLM_UserTable]: """Update a user.""" data: Dict[str, Any] = {} if user_alias is not None: data["user_alias"] = user_alias if team_id is not None: data["team_id"] = team_id if sso_user_id is not None: data["sso_user_id"] = sso_user_id if organization_id is not None: data["organization_id"] = organization_id if password is not None: data["password"] = password if teams is not None: data["teams"] = teams if user_role is not None: data["user_role"] = user_role if max_budget is not None: data["max_budget"] = max_budget if user_email is not None: data["user_email"] = user_email if models is not None: data["models"] = models if metadata is not None: data["metadata"] = json.dumps(metadata) if max_parallel_requests is not None: data["max_parallel_requests"] = max_parallel_requests if tpm_limit is not None: data["tpm_limit"] = tpm_limit if rpm_limit is not None: data["rpm_limit"] = rpm_limit if budget_duration is not None: data["budget_duration"] = budget_duration if allowed_cache_controls is not None: data["allowed_cache_controls"] = allowed_cache_controls if policies is not None: data["policies"] = policies if object_permission_id is not None: data["object_permission_id"] = object_permission_id return await self.update(user_id, data, id_field="user_id") async def delete_user(self, user_id: str) -> Optional[LiteLLM_UserTable]: """Delete a user.""" return await self.delete(user_id, id_field="user_id") async def update_spend( self, user_id: str, spend: float ) -> Optional[LiteLLM_UserTable]: """Update user spend.""" return await self.update(user_id, {"spend": spend}, id_field="user_id") async def add_to_team( self, user_id: str, team_id: str ) -> Optional[LiteLLM_UserTable]: """Add a user to a team using atomic array push operation.""" if not await self.exists(user_id, id_field="user_id"): return None record = await self.table.update( where={"user_id": user_id}, data={"teams": {"push": team_id}}, ) return self._to_model(record) async def remove_from_team( self, user_id: str, team_id: str ) -> Optional[LiteLLM_UserTable]: """Remove a user from a team. Note: Prisma doesn't support atomic array removal, so we use a read-modify-write pattern here. For high-concurrency scenarios, consider using raw SQL with array_remove(). """ user = await self.find_by_id(user_id) if user is None: return None teams = [t for t in user.teams if t != team_id] return await self.update(user_id, {"teams": teams}, id_field="user_id")