feat: conversation with bot

This commit is contained in:
pinks 2023-09-11 01:59:33 +02:00
parent af9d31f75c
commit 3984ca337c
6 changed files with 54 additions and 20 deletions

View File

@ -29,8 +29,8 @@ You can put these in `.env` file or pass them as environment variables.
- [x] Changing params, parsing png info in request - [x] Changing params, parsing png info in request
- [x] Cancelling jobs by deleting message - [x] Cancelling jobs by deleting message
- [x] Multiple parallel workers - [x] Multiple parallel workers
- [ ] Replying to another text message to copy prompt and generate - [x] Replying to another text message to copy prompt and generate
- [ ] Replying to bot message, conversation in DMs - [x] Replying to bot message, conversation in DMs
- [ ] Replying to png message to extract png info nad generate - [ ] Replying to png message to extract png info nad generate
- [ ] Banning tags - [ ] Banning tags
- [ ] Img2Img + Upscale - [ ] Img2Img + Upscale

View File

@ -2,7 +2,7 @@ import { Grammy, GrammyAutoQuote, GrammyParseMode, Log } from "../deps.ts";
import { formatUserChat } from "../utils.ts"; import { formatUserChat } from "../utils.ts";
import { session, SessionFlavor } from "./session.ts"; import { session, SessionFlavor } from "./session.ts";
import { queueCommand } from "./queueCommand.ts"; import { queueCommand } from "./queueCommand.ts";
import { txt2imgCommand } from "./txt2imgCommand.ts"; import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts";
export const logger = () => Log.getLogger(); export const logger = () => Log.getLogger();
@ -62,6 +62,7 @@ bot.api.setMyCommands([
bot.command("start", (ctx) => ctx.reply("Hello! Use the /txt2img command to generate an image")); bot.command("start", (ctx) => ctx.reply("Hello! Use the /txt2img command to generate an image"));
bot.command("txt2img", txt2imgCommand); bot.command("txt2img", txt2imgCommand);
bot.use(txt2imgQuestion.middleware() as any);
bot.command("queue", queueCommand); bot.command("queue", queueCommand);

View File

@ -1,3 +1,4 @@
import { db } from "../db/db.ts";
import { Grammy, GrammyKvStorage } from "../deps.ts"; import { Grammy, GrammyKvStorage } from "../deps.ts";
import { SdApi, SdTxt2ImgRequest } from "../sd.ts"; import { SdApi, SdTxt2ImgRequest } from "../sd.ts";
@ -33,9 +34,7 @@ export interface UserData {
params?: Partial<SdTxt2ImgRequest>; params?: Partial<SdTxt2ImgRequest>;
} }
const globalDb = await Deno.openKv("./app.db"); const globalDbAdapter = new GrammyKvStorage.DenoKVAdapter<GlobalData>(db);
const globalDbAdapter = new GrammyKvStorage.DenoKVAdapter<GlobalData>(globalDb);
const getDefaultGlobalData = (): GlobalData => ({ const getDefaultGlobalData = (): GlobalData => ({
adminUsernames: Deno.env.get("TG_ADMIN_USERS")?.split(",") ?? [], adminUsernames: Deno.env.get("TG_ADMIN_USERS")?.split(",") ?? [],

View File

@ -1,39 +1,72 @@
import { Grammy } from "../deps.ts"; import { Grammy, GrammyStatelessQ } from "../deps.ts";
import { formatUserChat } from "../utils.ts"; import { formatUserChat } from "../utils.ts";
import { jobStore } from "../db/jobStore.ts"; import { jobStore } from "../db/jobStore.ts";
import { parsePngInfo } from "../sd.ts"; import { parsePngInfo } from "../sd.ts";
import { Context, logger } from "./mod.ts"; import { Context, logger } from "./mod.ts";
export const txt2imgQuestion = new GrammyStatelessQ.StatelessQuestion(
"txt2img",
async (ctx) => {
if (!ctx.message.text) return;
await txt2img(ctx as any, ctx.message.text, false);
},
);
export async function txt2imgCommand(ctx: Grammy.CommandContext<Context>) { export async function txt2imgCommand(ctx: Grammy.CommandContext<Context>) {
if (!ctx.from?.id) { await txt2img(ctx, ctx.match, true);
return ctx.reply("I don't know who you are"); }
async function txt2img(ctx: Context, match: string, includeRepliedTo: boolean): Promise<void> {
if (!ctx.message?.from?.id) {
return void ctx.reply("I don't know who you are");
} }
const config = ctx.session.global; const config = ctx.session.global;
if (config.pausedReason != null) { if (config.pausedReason != null) {
return ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`); return void ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`);
} }
const jobs = await jobStore.getBy("status.type", "waiting"); const jobs = await jobStore.getBy("status.type", "waiting");
if (jobs.length >= config.maxJobs) { if (jobs.length >= config.maxJobs) {
return ctx.reply( return void ctx.reply(
`The queue is full. Try again later. (Max queue size: ${config.maxJobs})`, `The queue is full. Try again later. (Max queue size: ${config.maxJobs})`,
); );
} }
const userJobs = jobs.filter((job) => job.value.request.from.id === ctx.from?.id);
const userJobs = jobs.filter((job) => job.value.request.from.id === ctx.message?.from?.id);
if (userJobs.length >= config.maxUserJobs) { if (userJobs.length >= config.maxUserJobs) {
return ctx.reply( return void ctx.reply(
`You already have ${config.maxUserJobs} jobs in queue. Try again later.`, `You already have ${config.maxUserJobs} jobs in queue. Try again later.`,
); );
} }
const params = parsePngInfo(ctx.match);
if (!params.prompt) { let params = parsePngInfo(match);
return ctx.reply("Please describe what you want to see after the command"); const repliedToMsg = ctx.message.reply_to_message;
const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
if (includeRepliedTo && repliedToText) {
const originalParams = parsePngInfo(repliedToText);
params = {
...originalParams,
...params,
prompt: [originalParams.prompt, params.prompt].filter(Boolean).join("\n"),
};
} }
const reply = await ctx.reply("Accepted. You are now in queue."); if (!params.prompt) {
return void ctx.reply(
"Please tell me what you want to see." +
txt2imgQuestion.messageSuffixMarkdown(),
{ reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown" },
);
}
const replyMessage = await ctx.reply("Accepted. You are now in queue.");
await jobStore.create({ await jobStore.create({
params, params,
request: ctx.message, request: ctx.message,
reply, reply: replyMessage,
status: { type: "waiting" }, status: { type: "waiting" },
}); });
logger().debug(`Job enqueued for ${formatUserChat(ctx)}`);
logger().debug(`Job enqueued for ${formatUserChat(ctx.message)}`);
} }

View File

@ -4,7 +4,7 @@ import { db } from "./db.ts";
export interface JobSchema { export interface JobSchema {
params: Partial<SdTxt2ImgRequest>; params: Partial<SdTxt2ImgRequest>;
request: GrammyTypes.Message.TextMessage & { from: GrammyTypes.User }; request: GrammyTypes.Message & { from: GrammyTypes.User };
reply?: GrammyTypes.Message.TextMessage; reply?: GrammyTypes.Message.TextMessage;
status: status:
| { type: "waiting" } | { type: "waiting" }

View File

@ -11,6 +11,7 @@ export * as GrammyTypes from "https://deno.land/x/grammy_types@v3.2.0/mod.ts";
export * as GrammyAutoQuote from "https://deno.land/x/grammy_autoquote@v1.1.2/mod.ts"; export * as GrammyAutoQuote from "https://deno.land/x/grammy_autoquote@v1.1.2/mod.ts";
export * as GrammyParseMode from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts"; export * as GrammyParseMode from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts";
export * as GrammyKvStorage from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/mod.ts"; export * as GrammyKvStorage from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/mod.ts";
export * as GrammyStatelessQ from "npm:@grammyjs/stateless-question";
export * as FileType from "npm:file-type@18.5.0"; export * as FileType from "npm:file-type@18.5.0";
// @deno-types="./types/png-chunks-extract.d.ts" // @deno-types="./types/png-chunks-extract.d.ts"
export * as PngChunksExtract from "npm:png-chunks-extract@1.0.0"; export * as PngChunksExtract from "npm:png-chunks-extract@1.0.0";