Coverage for apps / ai / api_quotas.py: 97%

1""" 

2AI quota management API endpoints. 

3""" 

4 

5from ninja import Router, Schema, Status 

6 

7from apps.core.auth import AdminAuth, SessionAuth 

8from apps.core.models import AppSettings 

9 

10from .services.quota import get_usage, _next_midnight_utc_iso, ALL_FEATURES, FEATURE_LIMIT_FIELDS 

11 

12router = Router(tags=["ai"]) 

13 

14 

15# Schemas 

16 

17 

18class QuotaLimitsIn(Schema): 

19 remix: int 

20 remix_suggestions: int 

21 scale: int 

22 tips: int 

23 discover: int 

24 timer: int 

25 

26 

27class QuotaLimitsOut(Schema): 

28 remix: int 

29 remix_suggestions: int 

30 scale: int 

31 tips: int 

32 discover: int 

33 timer: int 

34 

35 

36class QuotaResponse(Schema): 

37 limits: QuotaLimitsOut 

38 usage: QuotaLimitsOut 

39 unlimited: bool 

40 resets_at: str 

41 

42 

43class ErrorOut(Schema): 

44 error: str 

45 message: str 

46 action: str | None = None 

47 

48 

49# Helpers 

50 

51 

52def _build_quota_response(profile, app_settings): 

53 """Build a QuotaResponse dict for the given profile.""" 

54 limits = {feature: getattr(app_settings, FEATURE_LIMIT_FIELDS[feature]) for feature in ALL_FEATURES} 

55 usage = get_usage(profile.pk) 

56 unlimited = profile.unlimited_ai or (profile.user and profile.user.is_staff) 

57 return { 

58 "limits": limits, 

59 "usage": usage, 

60 "unlimited": unlimited, 

61 "resets_at": _next_midnight_utc_iso(), 

62 } 

63 

64 

65# Endpoints 

66 

67 

68@router.get("/quotas", response={200: QuotaResponse, 404: ErrorOut}, auth=SessionAuth()) 

69def get_quotas(request): 

70 """Get current AI quota limits, usage, and reset time for the requesting user.""" 

71 from django.conf import settings as django_settings 

72 

73 if getattr(django_settings, "AUTH_MODE", "home") == "home": 

74 return Status(404, {"error": "not_found", "message": "Quotas are not available in home mode"}) 

75 

76 profile = request.auth 

77 app_settings = AppSettings.get() 

78 return _build_quota_response(profile, app_settings) 

79 

80 

81@router.put("/quotas", response={200: QuotaResponse, 404: ErrorOut}, auth=AdminAuth()) 

82def update_quotas(request, data: QuotaLimitsIn): 

83 """Update AI quota limits (admin only).""" 

84 from django.conf import settings as django_settings 

85 

86 if getattr(django_settings, "AUTH_MODE", "home") == "home": 

87 return Status(404, {"error": "not_found", "message": "Quotas are not available in home mode"}) 

88 

89 app_settings = AppSettings.get() 

90 for feature in ALL_FEATURES: 

91 setattr(app_settings, FEATURE_LIMIT_FIELDS[feature], getattr(data, feature)) 

92 app_settings.save() 

93 

94 profile = request.auth 

95 return _build_quota_response(profile, app_settings)