From 20b8cb52f244b4abb7a8b0a6d2001be413ccb216 Mon Sep 17 00:00:00 2001 From: pinks Date: Wed, 6 Sep 2023 02:53:34 +0200 Subject: [PATCH] refactor session module --- bot.ts | 111 ++++++++++++---------------------------------------- deps.ts | 2 - fmtArray.ts | 0 intl.ts | 44 +++++++++++++-------- queue.ts | 3 +- session.ts | 78 ++++++++++++++++++++++++++++++++++++ 6 files changed, 132 insertions(+), 106 deletions(-) delete mode 100644 fmtArray.ts create mode 100644 session.ts diff --git a/bot.ts b/bot.ts index cf7b33c..a63fda3 100644 --- a/bot.ts +++ b/bot.ts @@ -1,84 +1,27 @@ -import { - autoQuote, - autoRetry, - bold, - Bot, - Context, - DenoKVAdapter, - fmt, - hydrateReply, - ParseModeFlavor, - session, - SessionFlavor, -} from "./deps.ts"; -import { fmtArray, formatOrdinal } from "./intl.ts"; +import { autoQuote, bold, Bot, Context, hydrateReply, ParseModeFlavor } from "./deps.ts"; +import { fmt, formatOrdinal } from "./intl.ts"; import { queue } from "./queue.ts"; -import { SdRequest } from "./sd.ts"; +import { mySession, MySessionFlavor } from "./session.ts"; -type AppContext = ParseModeFlavor & SessionFlavor; - -interface SessionData { - global: { - adminUsernames: string[]; - pausedReason: string | null; - sdApiUrl: string; - maxUserJobs: number; - maxJobs: number; - defaultParams?: Partial; - }; - user: { - steps: number; - detail: number; - batchSize: number; - }; -} - -export const bot = new Bot(Deno.env.get("TG_BOT_TOKEN") ?? ""); +export type MyContext = ParseModeFlavor & MySessionFlavor; +export const bot = new Bot(Deno.env.get("TG_BOT_TOKEN") ?? ""); bot.use(autoQuote); bot.use(hydrateReply); -bot.api.config.use(autoRetry({ maxRetryAttempts: 5, maxDelaySeconds: 60 })); +bot.use(mySession); -const db = await Deno.openKv("./app.db"); - -const getDefaultGlobalSession = (): SessionData["global"] => ({ - adminUsernames: (Deno.env.get("ADMIN_USERNAMES") ?? "").split(",").filter(Boolean), - pausedReason: null, - sdApiUrl: Deno.env.get("SD_API_URL") ?? "http://127.0.0.1:7860/", - maxUserJobs: 3, - maxJobs: 20, - defaultParams: { - batch_size: 1, - n_iter: 1, - width: 128 * 2, - height: 128 * 3, - steps: 20, - cfg_scale: 9, - send_images: true, - negative_prompt: "boring_e621_fluffyrock_v4 boring_e621_v4", - }, +// Automatically retry bot requests if we get a 429 error +bot.api.config.use(async (prev, method, payload, signal) => { + let remainingAttempts = 5; + while (true) { + const result = await prev(method, payload, signal); + if (result.ok) return result; + if (result.error_code !== 429 || remainingAttempts <= 0) return result; + remainingAttempts -= 1; + const retryAfterMs = (result.parameters?.retry_after ?? 30) * 1000; + await new Promise((resolve) => setTimeout(resolve, retryAfterMs)); + } }); -bot.use(session({ - type: "multi", - global: { - getSessionKey: () => "global", - initial: getDefaultGlobalSession, - storage: new DenoKVAdapter(db), - }, - user: { - initial: () => ({ - steps: 20, - detail: 8, - batchSize: 2, - }), - }, -})); - -export async function getGlobalSession(): Promise { - const entry = await db.get(["sessions", "global"]); - return entry.value ?? getDefaultGlobalSession(); -} - bot.api.setMyShortDescription("I can generate furry images from text"); bot.api.setMyDescription( "I can generate furry images from text. Send /txt2img to generate an image.", @@ -135,11 +78,8 @@ bot.command("queue", (ctx) => { if (queue.length === 0) return ctx.reply("Queue is empty"); return ctx.replyFmt( fmt`Current queue:\n\n${ - fmtArray( - queue.map((job, index) => - fmt`${bold(index + 1)}. ${bold(job.userName)} in ${bold(job.chatName)}` - ), - "\n", + queue.map((job, index) => + fmt`${bold(index + 1)}. ${bold(job.userName)} in ${bold(job.chatName)}\n` ) }`, ); @@ -272,14 +212,13 @@ bot.command("setsdparam", (ctx) => { bot.command("sdparams", (ctx) => { if (!ctx.from?.username) return; const config = ctx.session.global; - return ctx.replyFmt(fmt`Current config:\n\n${ - fmtArray( + return ctx.replyFmt( + fmt`Current config:\n\n${ Object.entries(config.defaultParams ?? {}).map(([key, value]) => - fmt`${bold(key)} = ${String(value)}` - ), - "\n", - ) - }`); + fmt`${bold(key)} = ${String(value)}\n` + ) + }`, + ); }); bot.catch((err) => { diff --git a/deps.ts b/deps.ts index b84b8b8..2522297 100644 --- a/deps.ts +++ b/deps.ts @@ -2,5 +2,3 @@ export * from "https://deno.land/x/grammy@v1.18.1/mod.ts"; export * from "https://deno.land/x/grammy_autoquote@v1.1.2/mod.ts"; export * from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts"; export * from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/mod.ts"; -export { autoRetry } from "https://esm.sh/@grammyjs/auto-retry@1.1.1"; -export * from "https://deno.land/x/zod/mod.ts"; diff --git a/fmtArray.ts b/fmtArray.ts deleted file mode 100644 index e69de29..0000000 diff --git a/intl.ts b/intl.ts index c955dfa..fb943fe 100644 --- a/intl.ts +++ b/intl.ts @@ -8,26 +8,36 @@ export function formatOrdinal(n: number) { return `${n}th`; } +type DeepArray = Array>; +type StringLikes = DeepArray; + /** - * Like `fmt` from `grammy_parse_mode` but accepts an array instead of template string. + * Like `fmt` from `grammy_parse_mode` but additionally accepts arrays. * @see https://deno.land/x/grammy_parse_mode@1.7.1/format.ts?source=#L182 */ -export function fmtArray( - stringLikes: FormattedString[], - separator = "", -): FormattedString { +export const fmt = ( + rawStringParts: TemplateStringsArray | StringLikes, + ...stringLikes: StringLikes +): FormattedString => { let text = ""; - const entities: ConstructorParameters[1] = []; - for (let i = 0; i < stringLikes.length; i++) { - const stringLike = stringLikes[i]; - entities.push( - ...stringLike.entities.map((e) => ({ - ...e, - offset: e.offset + text.length, - })), - ); - text += stringLike.toString(); - if (i < stringLikes.length - 1) text += separator; + const entities: ConstructorParameters[1][] = []; + + const length = Math.max(rawStringParts.length, stringLikes.length); + for (let i = 0; i < length; i++) { + for (let stringLike of [rawStringParts[i], stringLikes[i]]) { + if (Array.isArray(stringLike)) { + stringLike = fmt(stringLike); + } + if (stringLike instanceof FormattedString) { + entities.push( + ...stringLike.entities.map((e) => ({ + ...e, + offset: e.offset + text.length, + })), + ); + } + if (stringLike != null) text += stringLike.toString(); + } } return new FormattedString(text, entities); -} +}; diff --git a/queue.ts b/queue.ts index 57dcbab..bbb0621 100644 --- a/queue.ts +++ b/queue.ts @@ -1,5 +1,6 @@ import { InputFile, InputMediaBuilder } from "./deps.ts"; -import { bot, getGlobalSession } from "./bot.ts"; +import { bot } from "./bot.ts"; +import { getGlobalSession } from "./session.ts"; import { formatOrdinal } from "./intl.ts"; import { SdProgressResponse, SdRequest, txt2img } from "./sd.ts"; import { extFromMimeType, mimeTypeFromBase64 } from "./mimeType.ts"; diff --git a/session.ts b/session.ts new file mode 100644 index 0000000..2223d8f --- /dev/null +++ b/session.ts @@ -0,0 +1,78 @@ +import { Context, DenoKVAdapter, session, SessionFlavor } from "./deps.ts"; +import { SdRequest } from "./sd.ts"; + +export type MySessionFlavor = SessionFlavor; + +export interface SessionData { + global: GlobalData; + chat: ChatData; + user: UserData; +} + +export interface GlobalData { + adminUsernames: string[]; + pausedReason: string | null; + sdApiUrl: string; + maxUserJobs: number; + maxJobs: number; + defaultParams?: Partial; +} + +export interface ChatData { + language: string; +} + +export interface UserData { + steps: number; + detail: number; + batchSize: number; +} + +const globalDb = await Deno.openKv("./app.db"); + +const globalDbAdapter = new DenoKVAdapter(globalDb); + +const getDefaultGlobalData = (): GlobalData => ({ + adminUsernames: (Deno.env.get("ADMIN_USERNAMES") ?? "").split(",").filter(Boolean), + pausedReason: null, + sdApiUrl: Deno.env.get("SD_API_URL") ?? "http://127.0.0.1:7860/", + maxUserJobs: 3, + maxJobs: 20, + defaultParams: { + batch_size: 1, + n_iter: 1, + width: 128 * 2, + height: 128 * 3, + steps: 20, + cfg_scale: 9, + send_images: true, + negative_prompt: "boring_e621_fluffyrock_v4 boring_e621_v4", + }, +}); + +export const mySession = session({ + type: "multi", + global: { + getSessionKey: () => "global", + initial: getDefaultGlobalData, + storage: globalDbAdapter, + }, + chat: { + initial: () => ({ + language: "en", + }), + }, + user: { + getSessionKey: (ctx) => ctx.from?.id.toFixed(), + initial: () => ({ + steps: 20, + detail: 8, + batchSize: 2, + }), + }, +}); + +export async function getGlobalSession(): Promise { + const data = await globalDbAdapter.read("global"); + return data ?? getDefaultGlobalData(); +}