From c194de0c9bc830477053a2442a0490957d683b00 Mon Sep 17 00:00:00 2001 From: artie Date: Tue, 11 Feb 2025 18:41:10 +0100 Subject: [PATCH] fix dm commands, add url support for ocr --- package.json | 2 +- src/client.ts | 9 ++- src/commands/language/translate.ts | 8 +-- src/commands/language/translateMenu.ts | 3 + src/commands/ocr/ocr.ts | 88 ++++++++++++++++++-------- src/commands/ocr/ocrMenu.ts | 14 ++-- src/commands/ocr/ocrTranslate.ts | 21 +++--- src/commands/ocr/ocrTranslateMenu.ts | 14 ++-- src/scripts/sandbox.ts | 6 ++ src/scripts/test.ts | 3 - src/utils/constants.ts | 3 + src/utils/functions.ts | 29 +++++++++ 12 files changed, 139 insertions(+), 61 deletions(-) create mode 100644 src/scripts/sandbox.ts delete mode 100644 src/scripts/test.ts diff --git a/package.json b/package.json index 573618b..c1337e9 100644 --- a/package.json +++ b/package.json @@ -7,7 +7,7 @@ "dev": "bun --watch src/index.ts", "start": "bun run src/index.ts", "sync": "bun run src/scripts/sync.ts", - "test": "bun run src/scripts/test.ts", + "sandbox": "bun run src/scripts/sandbox.ts", "deploy": "src/scripts/deploy.sh" }, "dependencies": { diff --git a/src/client.ts b/src/client.ts index 5ac0554..7e50acd 100644 --- a/src/client.ts +++ b/src/client.ts @@ -1,4 +1,10 @@ -import { Client, Collection, GatewayIntentBits, REST } from "discord.js"; +import { + Client, + Collection, + GatewayIntentBits, + Partials, + REST, +} from "discord.js"; import { env } from "./env"; import { ActivityType, API } from "@discordjs/core"; import type { Command } from "./types/command"; @@ -28,6 +34,7 @@ export class ArtemisClient extends Client { presence: { activities: [{ name: "🩷", type: ActivityType.Custom }], }, + partials: [Partials.Channel], }); const rest = new REST().setToken(env.DISCORD_TOKEN); diff --git a/src/commands/language/translate.ts b/src/commands/language/translate.ts index fff6120..3428aa4 100644 --- a/src/commands/language/translate.ts +++ b/src/commands/language/translate.ts @@ -1,6 +1,6 @@ import { - Attachment, AutocompleteInteraction, + hyperlink, SlashCommandBuilder, type InteractionEditReplyOptions, } from "discord.js"; @@ -39,7 +39,7 @@ export async function translateImpl( text: string, source: string | null, target: string, - attachment?: Attachment + imageUrl?: string ): Promise { const { text: translatedText, @@ -56,6 +56,7 @@ export async function translateImpl( if (translatedText.length > 4096) { return { + content: imageUrl ? `${hyperlink("Image", imageUrl)}\n\n` : undefined, files: [ { name: `${displaySource}-${displayTarget}.txt`, @@ -63,7 +64,6 @@ export async function translateImpl( `--- From ${displaySource} to ${displayTarget} ---\n${translatedText}` ), }, - ...(attachment ? [attachment] : []), ], }; } @@ -74,6 +74,7 @@ export async function translateImpl( title: `From ${displaySource} to ${displayTarget}`, description: translatedText, color: 0x0f2b46, + ...(imageUrl ? { image: { url: imageUrl } } : {}), author: { name: "DeepL", icon_url: "https://www.google.com/s2/favicons?domain=deepl.com&sz=64", @@ -83,7 +84,6 @@ export async function translateImpl( }, }, ], - files: attachment ? [attachment] : [], }; } diff --git a/src/commands/language/translateMenu.ts b/src/commands/language/translateMenu.ts index 8072041..bf6a263 100644 --- a/src/commands/language/translateMenu.ts +++ b/src/commands/language/translateMenu.ts @@ -1,6 +1,7 @@ import { ApplicationCommandType, ContextMenuCommandBuilder } from "discord.js"; import { defineCommand } from ".."; import { translateImpl } from "./translate"; +import { abort } from "../../utils/error"; export default defineCommand({ data: new ContextMenuCommandBuilder() @@ -11,6 +12,8 @@ export default defineCommand({ if (!interaction.isMessageContextMenuCommand()) return; const text = interaction.targetMessage.content; + if (!text) abort("No text to translate"); + await interaction.deferReply(); const payload = await translateImpl(text, null, "en-US"); diff --git a/src/commands/ocr/ocr.ts b/src/commands/ocr/ocr.ts index 4d56621..6b74575 100644 --- a/src/commands/ocr/ocr.ts +++ b/src/commands/ocr/ocr.ts @@ -1,6 +1,6 @@ import { - Attachment, codeBlock, + hyperlink, inlineCode, SlashCommandBuilder, type InteractionEditReplyOptions, @@ -10,41 +10,73 @@ import { downloadFile } from "../../utils/http"; import { abort } from "../../utils/error"; import { yandexOcr } from "../../utils/yandex"; import sharp from "sharp"; +import { getImageFromAttachmentOrString, run } from "../../utils/functions"; export function buildOcrPayload( text: string, detected_lang: string, - attachment?: Attachment + imageUrl?: string ): InteractionEditReplyOptions { - const languageName = - new Intl.DisplayNames(["en"], { type: "language" }).of(detected_lang) ?? - "unknown"; + const languageName = run(() => { + try { + return ( + new Intl.DisplayNames(["en"], { type: "language" }).of(detected_lang) ?? + "unknown" + ); + } catch { + return "unknown"; + } + }); - const content = `Detected language: ${inlineCode(languageName)}\n${codeBlock( - text - )}`; + const content = + `Detected language: ${inlineCode(languageName)}\n${codeBlock(text)}` + + (imageUrl ? `\n${hyperlink("Image", imageUrl)}` : ""); - if (content.length > 2000) { + if (content.length > 4096) { return { - content: `Detected language: ${inlineCode(languageName)}`, + content: + `Detected language: ${inlineCode(languageName)}` + + (imageUrl ? `\n${hyperlink("Image", imageUrl)}` : ""), files: [ { name: "ocr.txt", attachment: Buffer.from(text), }, - ...(attachment ? [attachment] : []), ], }; } return { - content, - files: attachment ? [attachment] : [], + embeds: [ + { + description: codeBlock(text), + color: 0xffdb4d, + fields: [ + { + name: "Detected language", + value: inlineCode(languageName), + }, + ], + ...(imageUrl ? { image: { url: imageUrl } } : {}), + author: { + name: "Yandex", + icon_url: + "https://www.google.com/s2/favicons?domain=yandex.com&sz=64", + }, + }, + ], }; } -export async function ocrImpl(attachment: Attachment) { - const { data, type } = await downloadFile(attachment.url); +export async function ocrImpl(url: string) { + const { data, type } = await run(async () => { + try { + return await downloadFile(url); + } catch { + abort("Failed to download the image"); + } + }); + if (!type?.mime.startsWith("image/")) { abort("The file must be an image!"); } @@ -54,7 +86,12 @@ export async function ocrImpl(attachment: Attachment) { .jpeg({ quality: 90 }) .toBuffer(); - return yandexOcr(compressed, type.mime); + const result = await yandexOcr(compressed, type.mime); + if (!result.text) { + result.text = "No text detected"; + } + + return result; } export default defineCommand({ @@ -62,25 +99,24 @@ export default defineCommand({ .setName("ocr") .setDescription("OCR an image using Yandex") .addAttachmentOption((option) => - option - .setName("image") - .setDescription("The image to OCR") - .setRequired(true) + option.setName("image").setDescription("The image to OCR") + ) + .addStringOption((option) => + option.setName("url").setDescription("The image URL to OCR") ), async execute(interaction) { - const attachment = interaction.options.getAttachment("image", true); - if (!attachment.contentType?.startsWith("image/")) { - abort("The file must be an image!"); - } + const attachment = interaction.options.getAttachment("image"); + const url = interaction.options.getString("url"); + const imageUrl = getImageFromAttachmentOrString(attachment, url); await interaction.deferReply(); - const result = await ocrImpl(attachment); + const result = await ocrImpl(imageUrl); const payload = buildOcrPayload( result.text, result.detected_lang, - attachment + imageUrl ); await interaction.editReply(payload); }, diff --git a/src/commands/ocr/ocrMenu.ts b/src/commands/ocr/ocrMenu.ts index 7960751..a500577 100644 --- a/src/commands/ocr/ocrMenu.ts +++ b/src/commands/ocr/ocrMenu.ts @@ -1,7 +1,7 @@ import { ApplicationCommandType, ContextMenuCommandBuilder } from "discord.js"; import { defineCommand } from ".."; -import { abort } from "../../utils/error"; import { buildOcrPayload, ocrImpl } from "./ocr"; +import { getImageFromAttachmentOrString } from "../../utils/functions"; export default defineCommand({ data: new ContextMenuCommandBuilder() @@ -12,16 +12,14 @@ export default defineCommand({ if (!interaction.isMessageContextMenuCommand()) return; const attachment = interaction.targetMessage.attachments.first(); - if (!attachment) { - abort("No attachment found"); - } - if (!attachment.contentType?.startsWith("image/")) { - abort("The file must be an image!"); - } + const imageUrl = getImageFromAttachmentOrString( + attachment, + interaction.targetMessage.content + ); await interaction.deferReply(); - const result = await ocrImpl(attachment); + const result = await ocrImpl(imageUrl); const payload = buildOcrPayload(result.text, result.detected_lang); await interaction.editReply(payload); }, diff --git a/src/commands/ocr/ocrTranslate.ts b/src/commands/ocr/ocrTranslate.ts index 0a1adeb..93e5f44 100644 --- a/src/commands/ocr/ocrTranslate.ts +++ b/src/commands/ocr/ocrTranslate.ts @@ -10,6 +10,7 @@ import { translateImpl, } from "../language/translate"; import { ocrImpl } from "./ocr"; +import { getImageFromAttachmentOrString } from "../../utils/functions"; export default defineCommand({ data: new SlashCommandBuilder() @@ -18,10 +19,10 @@ export default defineCommand({ "OCR an image using Yandex and translate the result using DeepL" ) .addAttachmentOption((option) => - option - .setName("image") - .setDescription("The image to OCR") - .setRequired(true) + option.setName("image").setDescription("The image to OCR") + ) + .addStringOption((option) => + option.setName("url").setDescription("The image URL to OCR") ) .addStringOption((option) => option @@ -39,13 +40,13 @@ export default defineCommand({ autocomplete: translateAutocompleteImpl, async execute(interaction) { - const attachment = interaction.options.getAttachment("image", true); + const attachment = interaction.options.getAttachment("image"); + const url = interaction.options.getString("url"); + const source = interaction.options.getString("source") ?? null; const target = interaction.options.getString("target") ?? "en-US"; - if (!attachment.contentType?.startsWith("image/")) { - abort("The file must be an image!"); - } + const imageUrl = getImageFromAttachmentOrString(attachment, url); await interaction.deferReply(); @@ -56,8 +57,8 @@ export default defineCommand({ abort("Target language not supported"); } - const { text } = await ocrImpl(attachment); - const payload = await translateImpl(text, source, target, attachment); + const { text } = await ocrImpl(imageUrl); + const payload = await translateImpl(text, source, target, imageUrl); await interaction.editReply(payload); }, }); diff --git a/src/commands/ocr/ocrTranslateMenu.ts b/src/commands/ocr/ocrTranslateMenu.ts index 07204ba..c00d66e 100644 --- a/src/commands/ocr/ocrTranslateMenu.ts +++ b/src/commands/ocr/ocrTranslateMenu.ts @@ -1,8 +1,8 @@ import { ApplicationCommandType, ContextMenuCommandBuilder } from "discord.js"; import { defineCommand } from ".."; -import { abort } from "../../utils/error"; import { translateImpl } from "../language/translate"; import { ocrImpl } from "./ocr"; +import { getImageFromAttachmentOrString } from "../../utils/functions"; export default defineCommand({ data: new ContextMenuCommandBuilder() @@ -13,16 +13,14 @@ export default defineCommand({ if (!interaction.isMessageContextMenuCommand()) return; const attachment = interaction.targetMessage.attachments.first(); - if (!attachment) { - abort("No attachment found"); - } - if (!attachment.contentType?.startsWith("image/")) { - abort("The file must be an image!"); - } + const imageUrl = getImageFromAttachmentOrString( + attachment, + interaction.targetMessage.content + ); await interaction.deferReply(); - const { text } = await ocrImpl(attachment); + const { text } = await ocrImpl(imageUrl); const payload = await translateImpl(text, null, "en-US"); await interaction.editReply(payload); }, diff --git a/src/scripts/sandbox.ts b/src/scripts/sandbox.ts new file mode 100644 index 0000000..5ac37b4 --- /dev/null +++ b/src/scripts/sandbox.ts @@ -0,0 +1,6 @@ +import { findFirstUrl } from "../utils/functions"; + +const text = + "I ended up doing https://cdn.discordapp.com/attachments/338689901111541760/1338884762785288253/image.png?ex=67acb51a&is=67ab639a&hm=92ed060dabffccf9544157da2922bce79386eca444a182db92f19d833d66fba6&b because yah"; + +console.log(findFirstUrl(text)); diff --git a/src/scripts/test.ts b/src/scripts/test.ts deleted file mode 100644 index f39b39f..0000000 --- a/src/scripts/test.ts +++ /dev/null @@ -1,3 +0,0 @@ -import { getWikipediaEditions } from "../utils/wikipedia"; - -console.log(await getWikipediaEditions()); diff --git a/src/utils/constants.ts b/src/utils/constants.ts index 642b11e..eb12fe2 100644 --- a/src/utils/constants.ts +++ b/src/utils/constants.ts @@ -8,6 +8,9 @@ export const USER_AGENT = `artemis (discord.js ${version})`; export const FAKE_USER_AGENT = "Mozilla/5.0 (X11; Linux x86_64; rv:135.0) Gecko/20100101 Firefox/135.0"; +export const URL_REGEX = + /\bhttps?:\/\/(www\.)?[-a-zA-Z0-9@:%._+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9()@:%_+.~#?&//=]*)\b/gi; + export const GOOD_BOT_EMOJIS = [ "<:teehee:825098257742299136>", "<:teehee2:825098258741067787>", diff --git a/src/utils/functions.ts b/src/utils/functions.ts index c5ab982..04f681d 100644 --- a/src/utils/functions.ts +++ b/src/utils/functions.ts @@ -1,6 +1,9 @@ import * as cheerio from "cheerio"; import { execa } from "execa"; import { customAlphabet } from "nanoid"; +import { URL_REGEX } from "./constants"; +import type { Attachment } from "discord.js"; +import { abort } from "./error"; export const nanoid = customAlphabet("1234567890abcdef"); export const shell = execa({ reject: false }); @@ -63,3 +66,29 @@ export function lazy(cb: () => T) { export function trim(str: string, maxLength: number) { return str.length > maxLength ? str.slice(0, maxLength) + "…" : str; } + +export function findUrls(text: string) { + return text.match(URL_REGEX) ?? []; +} + +export function findFirstUrl(text: string) { + return findUrls(text)[0]; +} + +export function getImageFromAttachmentOrString( + attachment?: Attachment | null, + str?: string | null +) { + if (attachment) { + if (!attachment.contentType?.startsWith("image/")) { + abort("The file must be an image!"); + } + return attachment.url; + } else if (str) { + const match = findFirstUrl(str); + if (!match) abort("The URL is invalid!"); + return match; + } else { + abort("You must provide an image or an image URL!"); + } +}