forked from pinks/eris
1
0
Fork 0
eris/app/userStatsStore.ts

123 lines
3.5 KiB
TypeScript
Raw Permalink Normal View History

2023-10-09 19:03:31 +00:00
import { minutesToMilliseconds } from "date-fns";
import { Store } from "indexed_kv";
2023-10-15 19:13:38 +00:00
import { info } from "std/log/mod.ts";
import { Static, t } from "elysia";
2023-10-09 19:03:31 +00:00
import { db } from "./db.ts";
import { generationStore } from "./generationStore.ts";
import { kvMemoize } from "../utils/kvMemoize.ts";
2023-10-27 19:49:46 +00:00
import { sortBy } from "std/collections/sort_by.ts";
2023-10-09 19:03:31 +00:00
export const userStatsSchema = t.Object({
userId: t.Number(),
imageCount: t.Number(),
stepCount: t.Number(),
pixelCount: t.Number(),
pixelStepCount: t.Number(),
tagCountMap: t.Record(t.String(), t.Number()),
timestamp: t.Number(),
});
2023-10-09 19:03:31 +00:00
export type UserStats = Static<typeof userStatsSchema>;
2023-10-09 19:03:31 +00:00
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;
2023-10-13 11:47:57 +00:00
let stepCount = 0;
2023-10-09 19:03:31 +00:00
let pixelCount = 0;
2023-10-13 11:47:57 +00:00
let pixelStepCount = 0;
2023-10-27 19:49:46 +00:00
const tagCountMap = new Map<string, number>();
2023-10-09 19:03:31 +00:00
2023-10-15 19:13:38 +00:00
info(`Calculating user stats for ${userId}`);
2023-10-09 19:03:31 +00:00
for await (
const generation of generationStore.listBy("fromId", { value: userId })
) {
imageCount++;
2023-10-13 11:47:57 +00:00
stepCount += generation.value.info?.steps ?? 0;
2023-10-09 19:03:31 +00:00
pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0);
2023-10-13 11:47:57 +00:00
pixelStepCount += (generation.value.info?.width ?? 0) *
(generation.value.info?.height ?? 0) *
(generation.value.info?.steps ?? 0);
2023-10-10 16:21:25 +00:00
const tags = generation.value.info?.prompt
// split on punctuation and newlines
.split(/[,;.]\s+|\n/)
// remove `:weight` syntax
2023-10-27 19:49:46 +00:00
.map((tag) => tag.replace(/:[\d\.]+/g, ""))
2023-10-10 16:21:25 +00:00
// remove `(tag)` and `[tag]` syntax
.map((tag) => tag.replace(/[()[\]]/g, " "))
// collapse multiple whitespace to one
.map((tag) => tag.replace(/\s+/g, " "))
// trim whitespace
2023-10-09 19:03:31 +00:00
.map((tag) => tag.trim())
2023-10-10 16:21:25 +00:00
// remove empty tags
2023-10-09 19:03:31 +00:00
.filter((tag) => tag.length > 0)
2023-10-10 16:21:25 +00:00
// lowercase tags
.map((tag) => tag.toLowerCase()) ??
// default to empty array
[];
2023-10-09 19:03:31 +00:00
for (const tag of tags) {
2023-10-27 19:49:46 +00:00
const count = tagCountMap.get(tag) ?? 0;
tagCountMap.set(tag, count + 1);
2023-10-09 19:03:31 +00:00
}
}
2023-10-27 19:49:46 +00:00
const tagCountObj = Object.fromEntries(
sortBy(
Array.from(tagCountMap.entries()),
([_tag, count]) => -count,
).filter(([_tag, count]) => count >= 3),
);
2023-10-09 19:03:31 +00:00
return {
userId,
imageCount,
2023-10-13 11:47:57 +00:00
stepCount,
2023-10-09 19:03:31 +00:00
pixelCount,
2023-10-13 11:47:57 +00:00
pixelStepCount,
2023-10-27 19:49:46 +00:00
tagCountMap: tagCountObj,
2023-10-09 19:03:31 +00:00
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);
},
},
},
);