From 1a7a108d473879a9598bb81c739e3503b60f83a8 Mon Sep 17 00:00:00 2001 From: evazion Date: Fri, 19 Mar 2021 04:44:22 -0500 Subject: [PATCH] discord: add /tagme command. --- Gemfile | 1 + Gemfile.lock | 2 + app/logical/danbooru/http.rb | 25 ++++- app/logical/deep_danbooru_client.rb | 28 ++++++ app/logical/discord_api_client.rb | 46 +++++++-- app/logical/discord_slash_command.rb | 42 +++++++- .../discord_slash_command/tagme_command.rb | 97 +++++++++++++++++++ app/models/tag.rb | 4 + 8 files changed, 231 insertions(+), 14 deletions(-) create mode 100644 app/logical/deep_danbooru_client.rb create mode 100644 app/logical/discord_slash_command/tagme_command.rb diff --git a/Gemfile b/Gemfile index 3bc6cd97e..794ec1974 100644 --- a/Gemfile +++ b/Gemfile @@ -50,6 +50,7 @@ gem 'hsluv' gem 'google-cloud-bigquery', require: "google/cloud/bigquery" gem 'google-cloud-storage', require: "google/cloud/storage" gem 'ed25519' +gem 'terminal-table' group :production, :staging do gem 'unicorn', :platforms => :ruby diff --git a/Gemfile.lock b/Gemfile.lock index 57be2ab40..19be8ae30 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -446,6 +446,7 @@ GEM dante (>= 0.2.0) multi_json (~> 1.0) stripe (> 5, < 6) + terminal-table (1.6.0) thor (1.1.0) tzinfo (2.0.4) concurrent-ruby (~> 1.0) @@ -556,6 +557,7 @@ DEPENDENCIES streamio-ffmpeg stripe stripe-ruby-mock + terminal-table tzinfo-data unicorn unicorn-worker-killer diff --git a/app/logical/danbooru/http.rb b/app/logical/danbooru/http.rb index bbf30cafd..a792b1492 100644 --- a/app/logical/danbooru/http.rb +++ b/app/logical/danbooru/http.rb @@ -11,8 +11,9 @@ require "danbooru/http/unpolish_cloudflare" module Danbooru class Http - class DownloadError < StandardError; end - class FileTooLargeError < StandardError; end + class Error < StandardError; end + class DownloadError < Error; end + class FileTooLargeError < Error; end DEFAULT_TIMEOUT = 10 MAX_REDIRECTS = 5 @@ -43,7 +44,7 @@ module Danbooru end def put(url, **options) - request(:get, url, **options) + request(:put, url, **options) end def post(url, **options) @@ -54,6 +55,14 @@ module Danbooru request(:delete, url, **options) end + def get!(url, **options) + request!(:get, url, **options) + end + + def post!(url, **options) + request!(:post, url, **options) + end + def follow(*args) dup.tap { |o| o.http = o.http.follow(*args) } end @@ -136,6 +145,16 @@ module Danbooru fake_response(599, "") end + def request!(method, url, **options) + response = request(method, url, **options) + + if response.status.in?(200..399) + response + else + raise Error, "#{method.upcase} #{url} failed (HTTP #{response.status})" + end + end + def fake_response(status, body) ::HTTP::Response.new(status: status, version: "1.1", body: ::HTTP::Response::Body.new(body)) end diff --git a/app/logical/deep_danbooru_client.rb b/app/logical/deep_danbooru_client.rb new file mode 100644 index 000000000..c1c256f9c --- /dev/null +++ b/app/logical/deep_danbooru_client.rb @@ -0,0 +1,28 @@ +class DeepDanbooruClient + attr_reader :http + + def initialize(http: Danbooru::Http.new) + @http = http + end + + def tags!(file) + html = post!("/upload", form: { + file: HTTP::FormData::File.new(file) + }).parse + + tags = html.css("tbody tr").map do |row| + tag_name = row.css("td:first-child").text + confidence = row.css("td:last-child").text + + # If tag_name is "rating:safe", then make a mock tag. + tag = Tag.find_by_name_or_alias(tag_name) || Tag.new(name: tag_name).freeze + [tag, confidence.to_f] + end.to_h + + tags + end + + def post!(url, **options) + http.post!("http://dev.kanotype.net:8003/deepdanbooru/#{url}", **options) + end +end diff --git a/app/logical/discord_api_client.rb b/app/logical/discord_api_client.rb index 6b4035031..3ecc59f26 100644 --- a/app/logical/discord_api_client.rb +++ b/app/logical/discord_api_client.rb @@ -3,6 +3,10 @@ class DiscordApiClient BASE_URL = "https://discord.com/api/v8" + # https://discord.com/developers/docs/resources/webhook#execute-webhook + # Actual limit is 2000; use 1950 for headroom. + MAX_MESSAGE_LENGTH = 1950 + attr_reader :application_id, :bot_token, :http def initialize(application_id: Danbooru.config.discord_application_client_id, bot_token: Danbooru.config.discord_bot_token, http: Danbooru::Http.new) @@ -11,6 +15,9 @@ class DiscordApiClient @http = http end + # https://discord.com/developers/docs/interactions/slash-commands#registering-a-command + # https://discord.com/developers/docs/interactions/slash-commands#create-global-application-command + # https://discord.com/developers/docs/interactions/slash-commands#create-guild-application-command def register_slash_command(name:, description:, options: [], guild_id: nil) json = { name: name, @@ -25,34 +32,55 @@ class DiscordApiClient end end + # https://discord.com/developers/docs/interactions/slash-commands#create-followup-message + def create_followup_message(interaction_token, allowed_mentions: { parse: [] }, **options) + post("/webhooks/#{application_id}/#{interaction_token}", { + allowed_mentions: allowed_mentions, + **options + }) + end + + # https://discord.com/developers/docs/resources/channel#get-channel def get_channel(channel_id, **options) get("/channels/#{channel_id}", **options) end - def me(**options) + def get_channel_messages(channel_id, limit: 50, **options) + get("/channels/#{channel_id}/messages", params: { limit: limit }, **options) + end + + # https://discord.com/developers/docs/resources/channel#trigger-typing-indicator + def trigger_typing_indicator(channel_id) + post("/channels/#{channel_id}/typing") + end + + # https://discord.com/developers/docs/resources/user#get-current-user + def get_current_user(**options) get("/users/@me", **options) end def get(url, cache: nil, **options) if cache - client.cache(cache).get("#{BASE_URL}/#{url}").parse + client.cache(cache).get("#{BASE_URL}/#{url}", **options).parse else - client.get("#{BASE_URL}/#{url}").parse + client.get("#{BASE_URL}/#{url}", **options).parse end end - def post(url, data) + def post(url, data = {}) client.post("#{BASE_URL}/#{url}", json: data).parse end - def client - headers = { + def http_headers + { "User-Agent": "#{Danbooru.config.canonical_app_name} (#{Danbooru.config.source_code_url}, 1.0)", "Authorization": "Bot #{bot_token}" } - - http.headers(headers) end - memoize :client + def client + http.headers(http_headers) + end + + memoize :client, :http_headers end diff --git a/app/logical/discord_slash_command.rb b/app/logical/discord_slash_command.rb index 589174b47..e8e35bffb 100644 --- a/app/logical/discord_slash_command.rb +++ b/app/logical/discord_slash_command.rb @@ -7,10 +7,18 @@ class DiscordSlashCommand ApplicationCommand = 2 end + # https://discord.com/developers/docs/interactions/slash-commands#interaction-response-interactionresponsetype + module InteractionResponseType + Pong = 1 + ChannelMessageWithSource = 4 + DeferredChannelMessageWithSource = 5 + end + # https://discord.com/developers/docs/interactions/slash-commands#applicationcommandoptiontype module ApplicationCommandOptionType String = 3 Integer = 4 + Boolean = 5 end # The name of the slash command. @@ -48,7 +56,7 @@ class DiscordSlashCommand # https://discord.com/developers/docs/interactions/slash-commands#responding-to-an-interaction # https://discord.com/developers/docs/interactions/slash-commands#interaction-response - def respond_with(content = nil, type: 4, posts: [], **options) + def respond_with(content = nil, type: InteractionResponseType::ChannelMessageWithSource, posts: [], **options) if posts.present? embeds = posts.map { |post| DiscordSlashCommand::PostEmbed.new(post, self).to_h } options[:embeds] = embeds @@ -63,6 +71,35 @@ class DiscordSlashCommand } end + # Post a response to the command later, after it's ready. + # https://discord.com/developers/docs/interactions/slash-commands#interaction-response + def respond_later(&block) + create_deferred_followup(&block) + trigger_typing_indicator + respond_with(type: InteractionResponseType::DeferredChannelMessageWithSource) + end + + def create_deferred_followup(&block) + Thread.new do + content = block.call + create_followup_message(content) + rescue StandardError => e + create_followup_message("`Error: #{e.message}`") + end + end + + def get_channel_messages(**options) + discord.get_channel_messages(data[:channel_id], **options) + end + + def trigger_typing_indicator + discord.trigger_typing_indicator(data[:channel_id]) + end + + def create_followup_message(content) + discord.create_followup_message(data[:token], content: content) + end + def channel discord.get_channel(data[:channel_id], cache: 1.minute) end @@ -84,6 +121,7 @@ class DiscordSlashCommand count: DiscordSlashCommand::CountCommand, posts: DiscordSlashCommand::PostsCommand, random: DiscordSlashCommand::RandomCommand, + tagme: DiscordSlashCommand::TagmeCommand, time: DiscordSlashCommand::TimeCommand, wiki: DiscordSlashCommand::WikiCommand, } @@ -101,7 +139,7 @@ class DiscordSlashCommand case data[:type] when InteractionType::Ping - { type: InteractionType::Ping } + { type: InteractionResponseType::Pong } when InteractionType::ApplicationCommand name = data.dig(:data, :name) klass = slash_commands.fetch(name&.to_sym) diff --git a/app/logical/discord_slash_command/tagme_command.rb b/app/logical/discord_slash_command/tagme_command.rb new file mode 100644 index 000000000..23b974d66 --- /dev/null +++ b/app/logical/discord_slash_command/tagme_command.rb @@ -0,0 +1,97 @@ +class DiscordSlashCommand + class TagmeCommand < DiscordSlashCommand + self.name = "tagme" + self.description = "Automatically tag an image" + self.options = [{ + name: "url", + description: "The URL of the image to tag", + required: false, + type: ApplicationCommandOptionType::String + }, { + name: "table", + description: "Format the output as a table", + required: false, + type: ApplicationCommandOptionType::Boolean + }] + + def call + table = params.fetch(:table, false) + + # Use the given URL, if present, or the last message with an attachment, if not. + if params[:url].present? + respond_later { tagme(params[:url], table: table) } + elsif result = get_last_message_with_url + message, url = result + respond_later { tagme(url, table: table) } + else + respond_with("No image found. Post an image or provide a URL.") + end + end + + def tagme(url, table: false) + tags = get_tags(url) + + if table + build_tag_table(tags) + else + build_tag_list(tags) + end + end + + def get_last_message_with_url(limit: 10) + messages = get_channel_messages(limit: limit) + + messages.each do |message| + if message["attachments"].present? + url = message["attachments"].first["url"] + # else + # url = message["content"].scan(%r!https?://[^ ]+!i).first + end + + return [message, url] if url.present? + end + + nil + end + + def get_tags(url, size: 500, minimum_confidence: 0.5) + response, file = http.download_media(url) + preview = file.preview(size, size) + tags = deep_danbooru.tags!(preview).to_a + tags = tags.reject { |tag, confidence| confidence < minimum_confidence } + tags = tags.sort_by { |tag, confidence| [tag.general? ? 1 : 0, tag.name] }.to_h + tags + end + + def build_tag_table(tags) + table = Terminal::Table.new + table.headings = ["Tag", "Count", "Confidence"] + + tags.each do |tag, confidence| + table << [tag.name, tag.post_count, "%.f%%" % (100 * confidence)] + break if table.to_s.size >= DiscordApiClient::MAX_MESSAGE_LENGTH + end + + "```\n#{table}\n```" + end + + def build_tag_list(tags) + msg = "" + + tags.keys.each do |tag| + msg << "[#{tag.name}](#{Routes.posts_url(tags: tag.name)}) " + break if msg.size >= DiscordApiClient::MAX_MESSAGE_LENGTH + end + + msg + end + + def http + @http ||= Danbooru::Http.timeout(15) + end + + def deep_danbooru + @deep_danbooru ||= DeepDanbooruClient.new(http: http) + end + end +end diff --git a/app/models/tag.rb b/app/models/tag.rb index dee7997ae..f587f5c75 100644 --- a/app/models/tag.rb +++ b/app/models/tag.rb @@ -260,6 +260,10 @@ class Tag < ApplicationRecord where("regexp_replace(tags.name, ?, '\\1', 'g') LIKE ?", ABBREVIATION_REGEXP.source, abbrev.to_escaped_for_sql_like) end + def find_by_name_or_alias(name) + find_by_name(TagAlias.to_aliased(normalize_name(name))) + end + def find_by_abbreviation(abbrev) abbreviation_matches(abbrev.escape_wildcards).order(post_count: :desc).first end