discord: add /tagme command.

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

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

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

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

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

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)

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

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