feat: stats api

api/mod.ts

@@ -14,7 +14,7 @@ export async function serveUi() {
           aliasMap: {
             "/utils/*": "../utils/",
           },
-          quiet: true,
+          log: (_request, response) => response.status >= 400,
         }),
     }));
 

api/serveApi.ts

@@ -1,14 +1,19 @@
-import { createPathFilter } from "t_rest/server";
+import { createLoggerMiddleware, createPathFilter } from "t_rest/server";
 import { jobsRoute } from "./jobsRoute.ts";
-import { sessionsRoute } from "./sessionsRoute.ts";
-import { usersRoute } from "./usersRoute.ts";
 import { paramsRoute } from "./paramsRoute.ts";
+import { sessionsRoute } from "./sessionsRoute.ts";
+import { statsRoute } from "./statsRoute.ts";
+import { usersRoute } from "./usersRoute.ts";
 
-export const serveApi = createPathFilter({
-  "jobs": jobsRoute,
-  "sessions": sessionsRoute,
-  "users": usersRoute,
-  "settings/params": paramsRoute,
-});
+export const serveApi = createLoggerMiddleware(
+  createPathFilter({
+    "jobs": jobsRoute,
+    "sessions": sessionsRoute,
+    "users": usersRoute,
+    "settings/params": paramsRoute,
+    "stats": statsRoute,
+  }),
+  { filterStatus: (status) => status >= 400 },
+);
 
 export type ApiHandler = typeof serveApi;

api/statsRoute.ts

@@ -0,0 +1,46 @@
+import { createEndpoint, createMethodFilter, createPathFilter } from "t_rest/server";
+import { liveGlobalStats } from "../app/globalStatsStore.ts";
+import { getDailyStats } from "../app/dailyStatsStore.ts";
+
+export const statsRoute = createPathFilter({
+  "global": createMethodFilter({
+    GET: createEndpoint(
+      { query: null, body: null },
+      async () => {
+        return {
+          status: 200,
+          body: {
+            type: "application/json",
+            data: liveGlobalStats,
+          },
+        };
+      },
+    ),
+  }),
+  "daily/{year}/{month}/{day}": createMethodFilter({
+    GET: createEndpoint(
+      { query: null, body: null },
+      async ({ params }) => {
+        const year = Number(params.year);
+        const month = Number(params.month);
+        const day = Number(params.day);
+        const minDate = new Date("2023-01-01");
+        const maxDate = new Date();
+        const date = new Date(Date.UTC(year, month - 1, day));
+        if (date < minDate || date > maxDate) {
+          return {
+            status: 404,
+            body: { type: "text/plain", data: "Not found" },
+          };
+        }
+        return {
+          status: 200,
+          body: {
+            type: "application/json",
+            data: await getDailyStats(year, month, day),
+          },
+        };
+      },
+    ),
+  }),
+});

app/dailyStatsStore.ts

@@ -0,0 +1,65 @@
+import { UTCDateMini } from "@date-fns/utc";
+import { hoursToMilliseconds, isSameDay, minutesToMilliseconds } from "date-fns";
+import { getLogger } from "std/log/mod.ts";
+import { JsonSchema, jsonType } from "t_rest/server";
+import { db } from "./db.ts";
+import { generationStore } from "./generationStore.ts";
+import { kvMemoize } from "./kvMemoize.ts";
+
+const logger = () => getLogger();
+
+export const dailyStatsSchema = {
+  type: "object",
+  properties: {
+    userIds: { type: "array", items: { type: "number" } },
+    imageCount: { type: "number" },
+    pixelCount: { type: "number" },
+    timestamp: { type: "number" },
+  },
+  required: ["userIds", "imageCount", "pixelCount", "timestamp"],
+} as const satisfies JsonSchema;
+
+export type DailyStats = jsonType<typeof dailyStatsSchema>;
+
+export const getDailyStats = kvMemoize(
+  db,
+  ["dailyStats"],
+  async (year: number, month: number, day: number): Promise<DailyStats> => {
+    const userIdSet = new Set<number>();
+    let imageCount = 0;
+    let pixelCount = 0;
+
+    const after = new Date(Date.UTC(year, month - 1, day));
+    const before = new Date(Date.UTC(year, month - 1, day + 1));
+
+    logger().info(`Calculating daily stats for ${year}-${month}-${day}`);
+
+    for await (
+      const generation of generationStore.listAll({ after, before })
+    ) {
+      userIdSet.add(generation.value.from.id);
+      imageCount++;
+      pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0);
+    }
+
+    return {
+      userIds: [...userIdSet],
+      imageCount,
+      pixelCount,
+      timestamp: Date.now(),
+    };
+  },
+  {
+    // expire in 1 minute if was calculated on the same day, otherwise 7-14 days.
+    expireIn: (result, year, month, day) => {
+      const requestDate = new UTCDateMini(year, month - 1, day);
+      const calculatedDate = new UTCDateMini(result.timestamp);
+      return isSameDay(requestDate, calculatedDate)
+        ? minutesToMilliseconds(1)
+        : hoursToMilliseconds(24 * 7 + Math.random() * 24 * 7);
+    },
+    // should cache if the stats are non-zero
+    shouldCache: (result) =>
+      result.userIds.length > 0 || result.imageCount > 0 || result.pixelCount > 0,
+  },
+);

app/globalStatsStore.ts

@@ -0,0 +1,56 @@
+import { addDays } from "date-fns";
+import { JsonSchema, jsonType } from "t_rest/server";
+import { decodeTime } from "ulid";
+import { getDailyStats } from "./dailyStatsStore.ts";
+import { generationStore } from "./generationStore.ts";
+
+export const globalStatsSchema = {
+  type: "object",
+  properties: {
+    userIds: { type: "array", items: { type: "number" } },
+    imageCount: { type: "number" },
+    pixelCount: { type: "number" },
+    timestamp: { type: "number" },
+  },
+  required: ["userIds", "imageCount", "pixelCount", "timestamp"],
+} as const satisfies JsonSchema;
+
+export type GlobalStats = jsonType<typeof globalStatsSchema>;
+
+export const liveGlobalStats: GlobalStats = await getGlobalStats();
+
+export async function getGlobalStats(): Promise<GlobalStats> {
+  // find the year/month/day of the first generation
+  const startDate = await generationStore.getAll({}, { limit: 1 })
+    .then((generations) => generations[0]?.id)
+    .then((generationId) => generationId ? new Date(decodeTime(generationId)) : new Date());
+
+  // iterate to today and sum up stats
+  const userIdSet = new Set<number>();
+  let imageCount = 0;
+  let pixelCount = 0;
+
+  const tomorrow = addDays(new Date(), 1);
+
+  for (
+    let date = startDate;
+    date < tomorrow;
+    date = addDays(date, 1)
+  ) {
+    const dailyStats = await getDailyStats(
+      date.getUTCFullYear(),
+      date.getUTCMonth() + 1,
+      date.getUTCDate(),
+    );
+    for (const userId of dailyStats.userIds) userIdSet.add(userId);
+    imageCount += dailyStats.imageCount;
+    pixelCount += dailyStats.pixelCount;
+  }
+
+  return {
+    userIds: [...userIdSet],
+    imageCount,
+    pixelCount,
+    timestamp: Date.now(),
+  };
+}

app/kvMemoize.ts

@@ -0,0 +1,45 @@
+/**
+ * Memoizes the function result in KV storage.
+ */
+export function kvMemoize<A extends Deno.KvKey, R>(
+  db: Deno.Kv,
+  key: Deno.KvKey,
+  fn: (...args: A) => Promise<R>,
+  options?: {
+    expireIn?: number | ((result: R, ...args: A) => number);
+    shouldRecalculate?: (result: R, ...args: A) => boolean;
+    shouldCache?: (result: R, ...args: A) => boolean;
+    override?: {
+      set: (key: Deno.KvKey, args: A, value: R, options: { expireIn?: number }) => Promise<void>;
+      get: (key: Deno.KvKey, args: A) => Promise<R | undefined>;
+    };
+  },
+): (...args: A) => Promise<R> {
+  return async (...args) => {
+    const cachedResult = options?.override?.get
+      ? await options.override.get(key, args)
+      : (await db.get<R>([...key, ...args])).value;
+
+    if (cachedResult != null) {
+      if (!options?.shouldRecalculate?.(cachedResult, ...args)) {
+        return cachedResult;
+      }
+    }
+
+    const result = await fn(...args);
+
+    const expireIn = typeof options?.expireIn === "function"
+      ? options.expireIn(result, ...args)
+      : options?.expireIn;
+
+    if (options?.shouldCache?.(result, ...args) ?? (result != null)) {
+      if (options?.override?.set) {
+        await options.override.set(key, args, result, { expireIn });
+      } else {
+        await db.set([...key, ...args], result, { expireIn });
+      }
+    }
+
+    return result;
+  };
+}

app/uploadQueue.ts

@@ -9,6 +9,7 @@ import { bot } from "../bot/mod.ts";
 import { formatUserChat } from "../utils/formatUserChat.ts";
 import { db, fs } from "./db.ts";
 import { generationStore, SdGenerationInfo } from "./generationStore.ts";
+import { liveGlobalStats } from "./globalStatsStore.ts";
 
 const logger = () => getLogger();
 
@@ -109,6 +110,15 @@ export async function processUploadQueue() {
       info: state.info,
     });
 
+    // update live stats
+    {
+      liveGlobalStats.imageCount++;
+      liveGlobalStats.pixelCount += state.info.width * state.info.height;
+      const userIdSet = new Set(liveGlobalStats.userIds);
+      userIdSet.add(state.from.id);
+      liveGlobalStats.userIds = [...userIdSet];
+    }
+
     // delete the status message
     await bot.api.deleteMessage(state.replyMessage.chat.id, state.replyMessage.message_id)
       .catch(() => undefined);

app/userStatsStore.ts

@@ -0,0 +1,100 @@
+import { minutesToMilliseconds } from "date-fns";
+import { Store } from "indexed_kv";
+import { getLogger } from "std/log/mod.ts";
+import { JsonSchema, jsonType } from "t_rest/server";
+import { db } from "./db.ts";
+import { generationStore } from "./generationStore.ts";
+import { kvMemoize } from "./kvMemoize.ts";
+
+const logger = () => getLogger();
+
+export const userStatsSchema = {
+  type: "object",
+  properties: {
+    userId: { type: "number" },
+    imageCount: { type: "number" },
+    pixelCount: { type: "number" },
+    tagsCount: {
+      type: "object",
+      additionalProperties: { type: "number" },
+    },
+    timestamp: { type: "number" },
+  },
+  required: ["userId", "imageCount", "pixelCount", "tagsCount", "timestamp"],
+} as const satisfies JsonSchema;
+
+export type UserStats = jsonType<typeof userStatsSchema>;
+
+type UserStatsIndices = {
+  userId: number;
+  imageCount: number;
+  pixelCount: number;
+};
+
+const userStatsStore = new Store<UserStats, UserStatsIndices>(
+  db,
+  "userStats",
+  {
+    indices: {
+      userId: { getValue: (item) => item.userId },
+      imageCount: { getValue: (item) => item.imageCount },
+      pixelCount: { getValue: (item) => item.pixelCount },
+    },
+  },
+);
+
+export const getUserStats = kvMemoize(
+  db,
+  ["userStats"],
+  async (userId: number): Promise<UserStats> => {
+    let imageCount = 0;
+    let pixelCount = 0;
+    const tagsCount: Record<string, number> = {};
+
+    logger().info(`Calculating user stats for ${userId}`);
+
+    for await (
+      const generation of generationStore.listBy("fromId", { value: userId })
+    ) {
+      imageCount++;
+      pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0);
+      const tags = generation.value.info?.prompt.split(/[,;.\n]/)
+        .map((tag) => tag.trim())
+        .filter((tag) => tag.length > 0)
+        .map((tag) => tag.toLowerCase())
+        .map((tag) => tag.replace(/[()[\]]/, ""))
+        .map((tag) => tag.replace(/:[\d.]/g, ""))
+        .map((tag) => tag.replace(/ +/g, " ")) ?? [];
+      for (const tag of tags) {
+        tagsCount[tag] = (tagsCount[tag] ?? 0) + 1;
+      }
+    }
+
+    return {
+      userId,
+      imageCount,
+      pixelCount,
+      tagsCount,
+      timestamp: Date.now(),
+    };
+  },
+  {
+    // expire in random time between 5-10 minutes
+    expireIn: () => minutesToMilliseconds(5 + Math.random() * 5),
+    // override default set/get behavior to use userStatsStore
+    override: {
+      get: async (_key, [userId]) => {
+        const items = await userStatsStore.getBy("userId", { value: userId }, { reverse: true });
+        return items[0]?.value;
+      },
+      set: async (_key, [userId], value, options) => {
+        // delete old stats
+        for await (const item of userStatsStore.listBy("userId", { value: userId })) {
+          await item.delete();
+        }
+        // set new stats
+        await userStatsStore.create(value, options);
+      },
+    },
+  },
+);

deno.json

@@ -19,10 +19,10 @@
 
     "async": "https://deno.land/x/async@v2.0.2/mod.ts",
     "ulid": "https://deno.land/x/ulid@v0.3.0/mod.ts",
-    "indexed_kv": "https://deno.land/x/indexed_kv@v0.4.0/mod.ts",
+    "indexed_kv": "https://deno.land/x/indexed_kv@v0.5.0/mod.ts",
     "kvmq": "https://deno.land/x/kvmq@v0.3.0/mod.ts",
     "kvfs": "https://deno.land/x/kvfs@v0.1.0/mod.ts",
-    "serve_spa": "https://deno.land/x/serve_spa@v0.1.0/mod.ts",
+    "serve_spa": "https://deno.land/x/serve_spa@v0.2.0/mod.ts",
     "reroute": "https://deno.land/x/reroute@v0.1.0/mod.ts",
     "grammy": "https://lib.deno.dev/x/grammy@1/mod.ts",
     "grammy_types": "https://lib.deno.dev/x/grammy_types@3/mod.ts",
@@ -30,13 +30,15 @@
     "grammy_parse_mode": "https://lib.deno.dev/x/grammy_parse_mode@1/mod.ts",
     "grammy_stateless_question": "https://lib.deno.dev/x/grammy_stateless_question_alpha@3/mod.ts",
     "grammy_files": "https://lib.deno.dev/x/grammy_files@1/mod.ts",
+    "date-fns": "https://cdn.skypack.dev/date-fns@2.30.0?dts",
+    "@date-fns/utc": "https://cdn.skypack.dev/@date-fns/utc@1.1.0?dts",
     "file_type": "https://esm.sh/file-type@18.5.0",
     "png_chunks_extract": "https://esm.sh/png-chunks-extract@1.0.0",
     "png_chunk_text": "https://esm.sh/png-chunk-text@1.0.0",
     "openapi_fetch": "https://esm.sh/openapi-fetch@0.7.6",
-    "t_rest/server": "https://esm.sh/ty-rest@0.3.1/server?dev",
+    "t_rest/server": "https://esm.sh/ty-rest@0.3.2/server?dev",
 
-    "t_rest/client": "https://esm.sh/ty-rest@0.3.1/client?dev",
+    "t_rest/client": "https://esm.sh/ty-rest@0.3.2/client?dev",
     "react": "https://esm.sh/react@18.2.0?dev",
     "react-dom/client": "https://esm.sh/react-dom@18.2.0/client?dev",
     "react-router-dom": "https://esm.sh/react-router-dom@6.16.0?dev",