Coverage for baibu/metrics/rollups.py: 100%

1"""Compute the daily rollups from the platform's own tables. 

2 

3Days are calendar days in the platform time zone (``TIME_ZONE``). Computing 

4a day again replaces what was stored for it, so every function here can be 

5run as often as needed. Only counts are stored. 

6""" 

7 

8from collections import Counter 

9from datetime import date 

10from datetime import datetime 

11from datetime import time 

12from datetime import timedelta 

13 

14from django.conf import settings 

15from django.contrib.auth import get_user_model 

16from django.db import transaction 

17from django.db.models import Aggregate 

18from django.db.models import Count 

19from django.db.models import Exists 

20from django.db.models import F 

21from django.db.models import FloatField 

22from django.db.models import Func 

23from django.db.models import Max 

24from django.db.models import OuterRef 

25from django.db.models import QuerySet 

26from django.utils import timezone 

27 

28from baibu.chat.models import Conversation 

29from baibu.chat.models import Message 

30from baibu.chat.models import Run 

31from baibu.chat.models import ToolInvocation 

32from baibu.submissions.models import Submission 

33from baibu.users.models import ConsentRecord 

34 

35from .definitions import DAILY_METRICS 

36from .models import DailyMetric 

37 

38User = get_user_model() 

39 

40Key = tuple[str, str] # (name, dimension) 

41 

42 

43class PercentileCont(Aggregate): 

44 """PostgreSQL's continuous percentile (interpolated), e.g. 0.5 for the median.""" 

45 

46 function = "percentile_cont" 

47 template = "%(function)s(%(fraction)s) WITHIN GROUP (ORDER BY %(expressions)s)" 

48 output_field = FloatField() 

49 

50 def __init__(self, expression, fraction: float, **extra): 

51 super().__init__(expression, fraction=float(fraction), **extra) 

52 

53 

54class EpochSeconds(Func): 

55 template = "EXTRACT(EPOCH FROM %(expressions)s)" 

56 output_field = FloatField() 

57 

58 

59def day_start(day: date) -> datetime: 

60 """Midnight at the start of ``day`` in the platform time zone.""" 

61 return datetime.combine(day, time.min, tzinfo=timezone.get_default_timezone()) 

62 

63 

64def day_bounds(first: date, last: date | None = None) -> dict: 

65 """Lookup values for ``created_at`` from the start of ``first`` to the end of ``last``.""" 

66 return {"gte": day_start(first), "lt": day_start((last or first) + timedelta(days=1))} 

67 

68 

69def week_start(day: date) -> date: 

70 """The Monday of the week that contains ``day``.""" 

71 return day - timedelta(days=day.weekday()) 

72 

73 

74def yesterday() -> date: 

75 return timezone.localdate() - timedelta(days=1) 

76 

77 

78def _between(field: str, first: date, last: date | None = None) -> dict: 

79 bounds = day_bounds(first, last) 

80 return {f"{field}__gte": bounds["gte"], f"{field}__lt": bounds["lt"]} 

81 

82 

83def active_users(first: date, last: date | None = None, *, users: QuerySet | None = None) -> QuerySet: 

84 """Users who sent a chat message or made a submission between ``first`` and ``last``.""" 

85 chatted = Message.objects.filter( 

86 conversation__user=OuterRef("pk"), role=Message.Role.USER, **_between("created_at", first, last) 

87 ) 

88 contributed = Submission.objects.filter(user=OuterRef("pk"), **_between("created_at", first, last)) 

89 users = User.objects.all() if users is None else users 

90 return users.filter(Exists(chatted) | Exists(contributed)) 

91 

92 

93def _counts_by(queryset: QuerySet, field: str) -> dict[str, int]: 

94 rows = queryset.order_by().values(field).annotate(n=Count("pk")).values_list(field, "n") 

95 # A blank value (for example a conversation without a language) is "unknown", 

96 # because an empty dimension means a plain total. 

97 return {str(key) if key else "unknown": n for key, n in rows} 

98 

99 

100def _consent_at(end: datetime) -> Counter: 

101 """Each user's tier per scope at ``end``, counted as ``scope:tier``.""" 

102 latest = ( 

103 ConsentRecord.objects.filter(granted_at__lt=end) 

104 .order_by("user_id", "scope", "-granted_at", "-created_at") 

105 .distinct("user_id", "scope") 

106 .values_list("scope", "tier", "revoked_at") 

107 ) 

108 counts: Counter = Counter() 

109 for scope, tier, revoked_at in latest: 

110 current = ConsentRecord.Tier.NONE if revoked_at is not None and revoked_at < end else tier 

111 counts[f"{scope}:{current}"] += 1 

112 return counts 

113 

114 

115def compute_day(day: date) -> dict[Key, int]: 

116 """Every daily metric for ``day``. Breakdowns leave out zero values.""" 

117 created = _between("created_at", day) 

118 end = day_bounds(day)["lt"] 

119 values: dict[Key, int] = { 

120 ("users_new", ""): User.objects.filter(**_between("date_joined", day)).count(), 

121 ("users_total", ""): User.objects.filter(date_joined__lt=end).count(), 

122 ("active_users", ""): active_users(day).count(), 

123 ("active_users_7d", ""): active_users(day - timedelta(days=6), day).count(), 

124 ("active_users_30d", ""): active_users(day - timedelta(days=29), day).count(), 

125 ("chat_users", ""): Message.objects.filter(role=Message.Role.USER, **created) 

126 .values("conversation__user_id") 

127 .distinct() 

128 .count(), 

129 ("contributors", ""): Submission.objects.filter(**created).values("user_id").distinct().count(), 

130 ("messages_user", ""): Message.objects.filter(role=Message.Role.USER, **created).count(), 

131 ("messages_assistant", ""): Message.objects.filter(role=Message.Role.ASSISTANT, **created).count(), 

132 } 

133 

134 conversations = Conversation.objects.filter(**created) 

135 values["conversations", ""] = conversations.count() 

136 breakdowns = { 

137 "conversations_by_language": _counts_by(conversations, "language_code"), 

138 "runs": _counts_by(Run.objects.filter(**created), "status"), 

139 "tool_calls": _counts_by(ToolInvocation.objects.filter(**created), "status"), 

140 "submissions_by_status": _counts_by(Submission.objects.filter(**created), "status"), 

141 "submissions_by_language": _counts_by(Submission.objects.filter(**created), "language_code"), 

142 "consent": _consent_at(end), 

143 } 

144 values["submissions", ""] = sum(breakdowns["submissions_by_status"].values()) 

145 for name, counts in breakdowns.items(): 

146 values.update({(name, dimension): n for dimension, n in counts.items() if n}) 

147 

148 reply_seconds = EpochSeconds(F("completed_at") - F("created_at")) 

149 times = Run.objects.filter(status=Run.Status.COMPLETED, completed_at__isnull=False, **created).aggregate( 

150 median=PercentileCont(reply_seconds, 0.5), p90=PercentileCont(reply_seconds, 0.9) 

151 ) 

152 if times["median"] is not None: 

153 values["reply_time_median_ms", ""] = round(times["median"] * 1000) 

154 values["reply_time_p90_ms", ""] = round(times["p90"] * 1000) 

155 return values 

156 

157 

158def compute_retention(day: date) -> dict[tuple[date, str, str], int | None]: 

159 """Retention cells that activity on ``day`` can change. 

160 

161 A cohort is the users who joined in one week (Monday to Sunday). For the 

162 week that contains ``day`` and each of the ``METRICS_RETENTION_WEEKS`` 

163 cohorts before it, this returns the cohort's size and how many of its 

164 users were active in that week. ``None`` marks cells of an empty cohort. 

165 """ 

166 activity_week = week_start(day) 

167 activity_end = activity_week + timedelta(days=6) 

168 cells: dict[tuple[date, str, str], int | None] = {} 

169 for weeks_after in range(settings.METRICS_RETENTION_WEEKS + 1): 

170 cohort_week = activity_week - timedelta(weeks=weeks_after) 

171 cohort = User.objects.filter(**_between("date_joined", cohort_week, cohort_week + timedelta(days=6))) 

172 size = cohort.count() 

173 cells[cohort_week, "retention_cohort_size", ""] = size or None 

174 if weeks_after: 

175 active = active_users(activity_week, activity_end, users=cohort).count() if size else None 

176 cells[cohort_week, "retention_active", str(weeks_after)] = active 

177 return cells 

178 

179 

180@transaction.atomic 

181def store_day(day: date) -> int: 

182 """Compute and store ``day``, replacing earlier values. Returns the rows stored.""" 

183 values = compute_day(day) 

184 DailyMetric.objects.filter(date=day, name__in=DAILY_METRICS).delete() 

185 DailyMetric.objects.bulk_create( 

186 DailyMetric(date=day, name=name, dimension=dimension, value=value) 

187 for (name, dimension), value in values.items() 

188 ) 

189 stored = len(values) 

190 for (cohort_week, name, dimension), value in compute_retention(day).items(): 

191 if value is None: 

192 DailyMetric.objects.filter(date=cohort_week, name=name, dimension=dimension).delete() 

193 else: 

194 DailyMetric.objects.update_or_create( 

195 date=cohort_week, name=name, dimension=dimension, defaults={"value": value} 

196 ) 

197 stored += 1 

198 return stored 

199 

200 

201def store_range(first: date, last: date) -> int: 

202 """Compute and store every day from ``first`` to ``last``. Returns the number of days.""" 

203 day = first 

204 while day <= last: 

205 store_day(day) 

206 day += timedelta(days=1) 

207 return (last - first).days + 1 if last >= first else 0 

208 

209 

210def days_to_catch_up(today: date | None = None) -> tuple[date, date]: 

211 """The days the nightly task computes: from the day after the last stored day to yesterday. 

212 

213 Always includes yesterday, and goes back at most 

214 ``METRICS_CATCH_UP_DAYS`` days, so a stopped worker fills the gap. 

215 """ 

216 last = (today or timezone.localdate()) - timedelta(days=1) 

217 earliest = last - timedelta(days=max(settings.METRICS_CATCH_UP_DAYS, 1) - 1) 

218 latest_stored = DailyMetric.objects.filter(name="active_users", date__lte=last).aggregate(Max("date"))["date__max"] 

219 first = last if latest_stored is None else max(earliest, min(latest_stored + timedelta(days=1), last)) 

220 return first, last