black-stats.ts 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312
  1. import { Database, and, eq, inArray, isNotNull, sql } from "../src/drizzle/index.js"
  2. import { BillingTable, BlackPlans, SubscriptionTable, UsageTable } from "../src/schema/billing.sql.js"
  3. if (process.argv.length < 3) {
  4. console.error("Usage: bun black-stats.ts <plan>")
  5. process.exit(1)
  6. }
  7. const plan = process.argv[2] as (typeof BlackPlans)[number]
  8. if (!BlackPlans.includes(plan)) {
  9. console.error("Usage: bun black-stats.ts <plan>")
  10. process.exit(1)
  11. }
  12. const cutoff = new Date(Date.UTC(2026, 1, 0, 23, 59, 59, 999))
  13. // get workspaces
  14. const workspaces = await Database.use((tx) =>
  15. tx
  16. .select({ workspaceID: BillingTable.workspaceID })
  17. .from(BillingTable)
  18. .where(
  19. and(isNotNull(BillingTable.subscriptionID), sql`JSON_UNQUOTE(JSON_EXTRACT(subscription, '$.plan')) = ${plan}`),
  20. ),
  21. )
  22. if (workspaces.length === 0) throw new Error(`No active Black ${plan} subscriptions found`)
  23. const week = sql<number>`YEARWEEK(${UsageTable.timeCreated}, 3)`
  24. const workspaceIDs = workspaces.map((row) => row.workspaceID)
  25. // Get subscription spend
  26. const spend = await Database.use((tx) =>
  27. tx
  28. .select({
  29. workspaceID: UsageTable.workspaceID,
  30. week,
  31. amount: sql<number>`COALESCE(SUM(${UsageTable.cost}), 0)`,
  32. })
  33. .from(UsageTable)
  34. .where(
  35. and(inArray(UsageTable.workspaceID, workspaceIDs), sql`JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) = 'sub'`),
  36. )
  37. .groupBy(UsageTable.workspaceID, week),
  38. )
  39. // Get pay per use spend
  40. const ppu = await Database.use((tx) =>
  41. tx
  42. .select({
  43. workspaceID: UsageTable.workspaceID,
  44. week,
  45. amount: sql<number>`COALESCE(SUM(${UsageTable.cost}), 0)`,
  46. })
  47. .from(UsageTable)
  48. .where(
  49. and(
  50. inArray(UsageTable.workspaceID, workspaceIDs),
  51. sql`(${UsageTable.enrichment} IS NULL OR JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) != 'sub')`,
  52. ),
  53. )
  54. .groupBy(UsageTable.workspaceID, week),
  55. )
  56. const models = await Database.use((tx) =>
  57. tx
  58. .select({
  59. workspaceID: UsageTable.workspaceID,
  60. model: UsageTable.model,
  61. amount: sql<number>`COALESCE(SUM(${UsageTable.cost}), 0)`,
  62. })
  63. .from(UsageTable)
  64. .where(
  65. and(inArray(UsageTable.workspaceID, workspaceIDs), sql`JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) = 'sub'`),
  66. )
  67. .groupBy(UsageTable.workspaceID, UsageTable.model),
  68. )
  69. const tokens = await Database.use((tx) =>
  70. tx
  71. .select({
  72. workspaceID: UsageTable.workspaceID,
  73. week,
  74. input: sql<number>`COALESCE(SUM(${UsageTable.inputTokens}), 0)`,
  75. cacheRead: sql<number>`COALESCE(SUM(${UsageTable.cacheReadTokens}), 0)`,
  76. output: sql<number>`COALESCE(SUM(${UsageTable.outputTokens}), 0) + COALESCE(SUM(${UsageTable.reasoningTokens}), 0)`,
  77. })
  78. .from(UsageTable)
  79. .where(
  80. and(inArray(UsageTable.workspaceID, workspaceIDs), sql`JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) = 'sub'`),
  81. )
  82. .groupBy(UsageTable.workspaceID, week),
  83. )
  84. const allWeeks = [...spend, ...ppu].map((row) => row.week)
  85. const weeks = [...new Set(allWeeks)].sort((a, b) => a - b)
  86. const spendMap = new Map<string, Map<number, number>>()
  87. const totals = new Map<string, number>()
  88. const ppuMap = new Map<string, Map<number, number>>()
  89. const ppuTotals = new Map<string, number>()
  90. const modelMap = new Map<string, { model: string; amount: number }[]>()
  91. const tokenMap = new Map<string, Map<number, { input: number; cacheRead: number; output: number }>>()
  92. for (const row of spend) {
  93. const workspace = spendMap.get(row.workspaceID) ?? new Map<number, number>()
  94. const total = totals.get(row.workspaceID) ?? 0
  95. const amount = toNumber(row.amount)
  96. workspace.set(row.week, amount)
  97. totals.set(row.workspaceID, total + amount)
  98. spendMap.set(row.workspaceID, workspace)
  99. }
  100. for (const row of ppu) {
  101. const workspace = ppuMap.get(row.workspaceID) ?? new Map<number, number>()
  102. const total = ppuTotals.get(row.workspaceID) ?? 0
  103. const amount = toNumber(row.amount)
  104. workspace.set(row.week, amount)
  105. ppuTotals.set(row.workspaceID, total + amount)
  106. ppuMap.set(row.workspaceID, workspace)
  107. }
  108. for (const row of models) {
  109. const current = modelMap.get(row.workspaceID) ?? []
  110. current.push({ model: row.model, amount: toNumber(row.amount) })
  111. modelMap.set(row.workspaceID, current)
  112. }
  113. for (const row of tokens) {
  114. const workspace = tokenMap.get(row.workspaceID) ?? new Map()
  115. workspace.set(row.week, {
  116. input: toNumber(row.input),
  117. cacheRead: toNumber(row.cacheRead),
  118. output: toNumber(row.output),
  119. })
  120. tokenMap.set(row.workspaceID, workspace)
  121. }
  122. const users = await Database.use((tx) =>
  123. tx
  124. .select({
  125. workspaceID: SubscriptionTable.workspaceID,
  126. subscribed: SubscriptionTable.timeCreated,
  127. subscription: BillingTable.subscription,
  128. })
  129. .from(SubscriptionTable)
  130. .innerJoin(BillingTable, eq(SubscriptionTable.workspaceID, BillingTable.workspaceID))
  131. .where(
  132. and(inArray(SubscriptionTable.workspaceID, workspaceIDs), sql`${SubscriptionTable.timeCreated} <= ${cutoff}`),
  133. ),
  134. )
  135. const counts = new Map<string, number>()
  136. for (const user of users) {
  137. const current = counts.get(user.workspaceID) ?? 0
  138. counts.set(user.workspaceID, current + 1)
  139. }
  140. const rows = users
  141. .map((user) => {
  142. const workspace = spendMap.get(user.workspaceID) ?? new Map<number, number>()
  143. const ppuWorkspace = ppuMap.get(user.workspaceID) ?? new Map<number, number>()
  144. const count = counts.get(user.workspaceID) ?? 1
  145. const amount = (totals.get(user.workspaceID) ?? 0) / count
  146. const ppuAmount = (ppuTotals.get(user.workspaceID) ?? 0) / count
  147. const monthStart = user.subscribed ? startOfMonth(user.subscribed) : null
  148. const modelRows = (modelMap.get(user.workspaceID) ?? []).sort((a, b) => b.amount - a.amount).slice(0, 3)
  149. const modelTotal = totals.get(user.workspaceID) ?? 0
  150. const modelCells = modelRows.map((row) => ({
  151. model: row.model,
  152. percent: modelTotal > 0 ? `${((row.amount / modelTotal) * 100).toFixed(1)}%` : "0.0%",
  153. }))
  154. const modelData = [0, 1, 2].map((index) => modelCells[index] ?? { model: "-", percent: "-" })
  155. const weekly = Object.fromEntries(
  156. weeks.map((item) => {
  157. const value = (workspace.get(item) ?? 0) / count
  158. const beforeMonth = monthStart ? isoWeekStart(item) < monthStart : false
  159. return [formatWeek(item), beforeMonth ? "-" : formatMicroCents(value)]
  160. }),
  161. )
  162. const ppuWeekly = Object.fromEntries(
  163. weeks.map((item) => {
  164. const value = (ppuWorkspace.get(item) ?? 0) / count
  165. const beforeMonth = monthStart ? isoWeekStart(item) < monthStart : false
  166. return [formatWeek(item), beforeMonth ? "-" : formatMicroCents(value)]
  167. }),
  168. )
  169. const tokenWorkspace = tokenMap.get(user.workspaceID) ?? new Map()
  170. const weeklyTokens = Object.fromEntries(
  171. weeks.map((item) => {
  172. const t = tokenWorkspace.get(item) ?? { input: 0, cacheRead: 0, output: 0 }
  173. const beforeMonth = monthStart ? isoWeekStart(item) < monthStart : false
  174. return [
  175. formatWeek(item),
  176. beforeMonth
  177. ? { input: "-", cacheRead: "-", output: "-" }
  178. : {
  179. input: Math.round(t.input / count),
  180. cacheRead: Math.round(t.cacheRead / count),
  181. output: Math.round(t.output / count),
  182. },
  183. ]
  184. }),
  185. )
  186. return {
  187. workspaceID: user.workspaceID,
  188. useBalance: user.subscription?.useBalance ?? false,
  189. subscribed: formatDate(user.subscribed),
  190. subscribedAt: user.subscribed?.getTime() ?? 0,
  191. amount,
  192. ppuAmount,
  193. models: modelData,
  194. weekly,
  195. ppuWeekly,
  196. weeklyTokens,
  197. }
  198. })
  199. .sort((a, b) => a.subscribedAt - b.subscribedAt)
  200. console.log(`Black ${plan} subscribers: ${rows.length}`)
  201. const header = [
  202. "workspaceID",
  203. "subscribed",
  204. "useCredit",
  205. "subTotal",
  206. "ppuTotal",
  207. "model1",
  208. "model1%",
  209. "model2",
  210. "model2%",
  211. "model3",
  212. "model3%",
  213. ...weeks.flatMap((item) => [
  214. formatWeek(item) + " sub",
  215. formatWeek(item) + " ppu",
  216. formatWeek(item) + " input",
  217. formatWeek(item) + " cache",
  218. formatWeek(item) + " output",
  219. ]),
  220. ]
  221. const lines = [header.map(csvCell).join(",")]
  222. for (const row of rows) {
  223. const model1 = row.models[0]
  224. const model2 = row.models[1]
  225. const model3 = row.models[2]
  226. const cells = [
  227. row.workspaceID,
  228. row.subscribed ?? "",
  229. row.useBalance ? "yes" : "no",
  230. formatMicroCents(row.amount),
  231. formatMicroCents(row.ppuAmount),
  232. model1.model,
  233. model1.percent,
  234. model2.model,
  235. model2.percent,
  236. model3.model,
  237. model3.percent,
  238. ...weeks.flatMap((item) => {
  239. const t = row.weeklyTokens[formatWeek(item)] ?? { input: "-", cacheRead: "-", output: "-" }
  240. return [
  241. row.weekly[formatWeek(item)] ?? "",
  242. row.ppuWeekly[formatWeek(item)] ?? "",
  243. String(t.input),
  244. String(t.cacheRead),
  245. String(t.output),
  246. ]
  247. }),
  248. ]
  249. lines.push(cells.map(csvCell).join(","))
  250. }
  251. const output = `${lines.join("\n")}\n`
  252. const file = Bun.file(`black-stats-${plan}.csv`)
  253. await file.write(output)
  254. console.log(`Wrote ${lines.length - 1} rows to ${file.name}`)
  255. const total = rows.reduce((sum, row) => sum + row.amount, 0)
  256. const average = rows.length === 0 ? 0 : total / rows.length
  257. console.log(`Average spending per user: ${formatMicroCents(average)}`)
  258. function formatMicroCents(value: number) {
  259. return `$${(value / 100000000).toFixed(2)}`
  260. }
  261. function formatDate(value: Date | null | undefined) {
  262. if (!value) return null
  263. return value.toISOString().split("T")[0]
  264. }
  265. function formatWeek(value: number) {
  266. return formatDate(isoWeekStart(value)) ?? ""
  267. }
  268. function startOfMonth(value: Date) {
  269. return new Date(Date.UTC(value.getUTCFullYear(), value.getUTCMonth(), 1))
  270. }
  271. function isoWeekStart(value: number) {
  272. const year = Math.floor(value / 100)
  273. const weekNumber = value % 100
  274. const jan4 = new Date(Date.UTC(year, 0, 4))
  275. const day = jan4.getUTCDay() || 7
  276. const weekStart = new Date(Date.UTC(year, 0, 4 - (day - 1)))
  277. weekStart.setUTCDate(weekStart.getUTCDate() + (weekNumber - 1) * 7)
  278. return weekStart
  279. }
  280. function toNumber(value: unknown) {
  281. if (typeof value === "number") return value
  282. if (typeof value === "bigint") return Number(value)
  283. if (typeof value === "string") return Number(value)
  284. return 0
  285. }
  286. function csvCell(value: string | number) {
  287. const text = String(value)
  288. if (!/[",\n]/.test(text)) return text
  289. return `"${text.replace(/"/g, '""')}"`
  290. }