refactor session module

This commit is contained in:
pinks 2023-09-06 02:53:34 +02:00
parent 5722238c06
commit 20b8cb52f2
6 changed files with 132 additions and 106 deletions

111
bot.ts
View File

@ -1,84 +1,27 @@
import { import { autoQuote, bold, Bot, Context, hydrateReply, ParseModeFlavor } from "./deps.ts";
autoQuote, import { fmt, formatOrdinal } from "./intl.ts";
autoRetry,
bold,
Bot,
Context,
DenoKVAdapter,
fmt,
hydrateReply,
ParseModeFlavor,
session,
SessionFlavor,
} from "./deps.ts";
import { fmtArray, formatOrdinal } from "./intl.ts";
import { queue } from "./queue.ts"; import { queue } from "./queue.ts";
import { SdRequest } from "./sd.ts"; import { mySession, MySessionFlavor } from "./session.ts";
type AppContext = ParseModeFlavor<Context> & SessionFlavor<SessionData>; export type MyContext = ParseModeFlavor<Context> & MySessionFlavor;
export const bot = new Bot<MyContext>(Deno.env.get("TG_BOT_TOKEN") ?? "");
interface SessionData {
global: {
adminUsernames: string[];
pausedReason: string | null;
sdApiUrl: string;
maxUserJobs: number;
maxJobs: number;
defaultParams?: Partial<SdRequest>;
};
user: {
steps: number;
detail: number;
batchSize: number;
};
}
export const bot = new Bot<AppContext>(Deno.env.get("TG_BOT_TOKEN") ?? "");
bot.use(autoQuote); bot.use(autoQuote);
bot.use(hydrateReply); bot.use(hydrateReply);
bot.api.config.use(autoRetry({ maxRetryAttempts: 5, maxDelaySeconds: 60 })); bot.use(mySession);
const db = await Deno.openKv("./app.db"); // Automatically retry bot requests if we get a 429 error
bot.api.config.use(async (prev, method, payload, signal) => {
const getDefaultGlobalSession = (): SessionData["global"] => ({ let remainingAttempts = 5;
adminUsernames: (Deno.env.get("ADMIN_USERNAMES") ?? "").split(",").filter(Boolean), while (true) {
pausedReason: null, const result = await prev(method, payload, signal);
sdApiUrl: Deno.env.get("SD_API_URL") ?? "http://127.0.0.1:7860/", if (result.ok) return result;
maxUserJobs: 3, if (result.error_code !== 429 || remainingAttempts <= 0) return result;
maxJobs: 20, remainingAttempts -= 1;
defaultParams: { const retryAfterMs = (result.parameters?.retry_after ?? 30) * 1000;
batch_size: 1, await new Promise((resolve) => setTimeout(resolve, retryAfterMs));
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",
},
}); });
bot.use(session<SessionData, AppContext>({
type: "multi",
global: {
getSessionKey: () => "global",
initial: getDefaultGlobalSession,
storage: new DenoKVAdapter(db),
},
user: {
initial: () => ({
steps: 20,
detail: 8,
batchSize: 2,
}),
},
}));
export async function getGlobalSession(): Promise<SessionData["global"]> {
const entry = await db.get<SessionData["global"]>(["sessions", "global"]);
return entry.value ?? getDefaultGlobalSession();
}
bot.api.setMyShortDescription("I can generate furry images from text"); bot.api.setMyShortDescription("I can generate furry images from text");
bot.api.setMyDescription( bot.api.setMyDescription(
"I can generate furry images from text. Send /txt2img to generate an image.", "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"); if (queue.length === 0) return ctx.reply("Queue is empty");
return ctx.replyFmt( return ctx.replyFmt(
fmt`Current queue:\n\n${ fmt`Current queue:\n\n${
fmtArray( queue.map((job, index) =>
queue.map((job, index) => fmt`${bold(index + 1)}. ${bold(job.userName)} in ${bold(job.chatName)}\n`
fmt`${bold(index + 1)}. ${bold(job.userName)} in ${bold(job.chatName)}`
),
"\n",
) )
}`, }`,
); );
@ -272,14 +212,13 @@ bot.command("setsdparam", (ctx) => {
bot.command("sdparams", (ctx) => { bot.command("sdparams", (ctx) => {
if (!ctx.from?.username) return; if (!ctx.from?.username) return;
const config = ctx.session.global; const config = ctx.session.global;
return ctx.replyFmt(fmt`Current config:\n\n${ return ctx.replyFmt(
fmtArray( fmt`Current config:\n\n${
Object.entries(config.defaultParams ?? {}).map(([key, value]) => Object.entries(config.defaultParams ?? {}).map(([key, value]) =>
fmt`${bold(key)} = ${String(value)}` fmt`${bold(key)} = ${String(value)}\n`
), )
"\n", }`,
) );
}`);
}); });
bot.catch((err) => { bot.catch((err) => {

View File

@ -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_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_parse_mode@1.7.1/mod.ts";
export * from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/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";

View File

44
intl.ts
View File

@ -8,26 +8,36 @@ export function formatOrdinal(n: number) {
return `${n}th`; return `${n}th`;
} }
type DeepArray<T> = Array<T | DeepArray<T>>;
type StringLikes = DeepArray<FormattedString | string | number | null | undefined>;
/** /**
* 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 * @see https://deno.land/x/grammy_parse_mode@1.7.1/format.ts?source=#L182
*/ */
export function fmtArray( export const fmt = (
stringLikes: FormattedString[], rawStringParts: TemplateStringsArray | StringLikes,
separator = "", ...stringLikes: StringLikes
): FormattedString { ): FormattedString => {
let text = ""; let text = "";
const entities: ConstructorParameters<typeof FormattedString>[1] = []; const entities: ConstructorParameters<typeof FormattedString>[1][] = [];
for (let i = 0; i < stringLikes.length; i++) {
const stringLike = stringLikes[i]; const length = Math.max(rawStringParts.length, stringLikes.length);
entities.push( for (let i = 0; i < length; i++) {
...stringLike.entities.map((e) => ({ for (let stringLike of [rawStringParts[i], stringLikes[i]]) {
...e, if (Array.isArray(stringLike)) {
offset: e.offset + text.length, stringLike = fmt(stringLike);
})), }
); if (stringLike instanceof FormattedString) {
text += stringLike.toString(); entities.push(
if (i < stringLikes.length - 1) text += separator; ...stringLike.entities.map((e) => ({
...e,
offset: e.offset + text.length,
})),
);
}
if (stringLike != null) text += stringLike.toString();
}
} }
return new FormattedString(text, entities); return new FormattedString(text, entities);
} };

View File

@ -1,5 +1,6 @@
import { InputFile, InputMediaBuilder } from "./deps.ts"; 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 { formatOrdinal } from "./intl.ts";
import { SdProgressResponse, SdRequest, txt2img } from "./sd.ts"; import { SdProgressResponse, SdRequest, txt2img } from "./sd.ts";
import { extFromMimeType, mimeTypeFromBase64 } from "./mimeType.ts"; import { extFromMimeType, mimeTypeFromBase64 } from "./mimeType.ts";

78
session.ts Normal file
View File

@ -0,0 +1,78 @@
import { Context, DenoKVAdapter, session, SessionFlavor } from "./deps.ts";
import { SdRequest } from "./sd.ts";
export type MySessionFlavor = SessionFlavor<SessionData>;
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<SdRequest>;
}
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<GlobalData>(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<SessionData, Context & MySessionFlavor>({
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<GlobalData> {
const data = await globalDbAdapter.read("global");
return data ?? getDefaultGlobalData();
}