pikadick/commands/
iqdb.rs

1use crate::{
2    checks::ENABLED_CHECK,
3    client_data::{
4        CacheStatsBuilder,
5        CacheStatsProvider,
6    },
7    util::{
8        LoadingReaction,
9        TimedCache,
10        TimedCacheEntry,
11    },
12    ClientDataKey,
13};
14use anyhow::Context as _;
15use serenity::{
16    builder::{
17        CreateEmbed,
18        CreateMessage,
19    },
20    framework::standard::{
21        macros::command,
22        Args,
23        CommandResult,
24    },
25    model::prelude::*,
26    prelude::*,
27};
28use std::sync::Arc;
29use tracing::error;
30
31#[derive(Clone, Debug)]
32pub struct IqdbClient {
33    client: iqdb::Client,
34    search_cache: TimedCache<String, iqdb::SearchResults>,
35}
36
37impl IqdbClient {
38    pub fn new() -> Self {
39        Self {
40            client: iqdb::Client::new(),
41            search_cache: TimedCache::new(),
42        }
43    }
44
45    /// Search for an image, with caching
46    pub async fn search(
47        &self,
48        query: &str,
49    ) -> anyhow::Result<Arc<TimedCacheEntry<iqdb::SearchResults>>> {
50        if let Some(entry) = self.search_cache.get_if_fresh(query) {
51            return Ok(entry);
52        }
53
54        let search_results = self
55            .client
56            .search(query)
57            .await
58            .context("failed to search for image")?;
59
60        self.search_cache
61            .insert(String::from(query), search_results);
62
63        self.search_cache
64            .get_if_fresh(query)
65            .context("cache data expired")
66    }
67}
68
69impl Default for IqdbClient {
70    fn default() -> Self {
71        Self::new()
72    }
73}
74
75impl CacheStatsProvider for IqdbClient {
76    fn publish_cache_stats(&self, cache_stats_builder: &mut CacheStatsBuilder) {
77        cache_stats_builder.publish_stat("iqdb", "search_cache", self.search_cache.len() as f32);
78    }
79}
80
81#[command]
82#[description("Search IQDB for an image at a url")]
83#[usage("<img_url>")]
84#[example("https://konachan.com/image/5982d8946ae503351e960f097f84cd90/Konachan.com%20-%20330136%20animal%20nobody%20original%20signed%20yutaka_kana.jpg")]
85#[checks(Enabled)]
86#[min_args(1)]
87#[max_args(1)]
88#[bucket("default")]
89async fn iqdb(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
90    let data_lock = ctx.data.read().await;
91    let client_data = data_lock
92        .get::<ClientDataKey>()
93        .expect("missing client data");
94    let client = client_data.iqdb_client.clone();
95    drop(data_lock);
96
97    let query = args.trimmed().current().expect("missing query");
98
99    let mut loading = LoadingReaction::new(ctx.http.clone(), msg);
100
101    match client
102        .search(query)
103        .await
104        .context("failed to search for image")
105    {
106        Ok(data) => {
107            let data = data.data();
108            match data.best_match.as_ref() {
109                Some(data) => {
110                    let embed_builder = CreateEmbed::new()
111                        .title("IQDB Best Match")
112                        .image(data.image_url.as_str())
113                        .url(data.url.as_str())
114                        .description(data.url.as_str());
115                    let message_builder = CreateMessage::new().embed(embed_builder);
116                    msg.channel_id
117                        .send_message(&ctx.http, message_builder)
118                        .await?;
119
120                    loading.send_ok();
121                }
122                None => {
123                    msg.channel_id
124                        .say(&ctx.http, format!("No results on iqdb for \"{query}\"",))
125                        .await?;
126                }
127            }
128        }
129        Err(error) => {
130            error!("{error:?}");
131            msg.channel_id.say(&ctx.http, format!("{error:?}")).await?;
132        }
133    }
134
135    Ok(())
136}
pikadick/commands/iqdb.rs

pikadick/commands/
iqdb.rs
