Skip to content

Commit aeaea42

Browse files
avoonixavoonix
committed
update sticker import
1 parent f83fda2 commit aeaea42

4 files changed

Lines changed: 102 additions & 46 deletions

File tree

‎fuzzle/src/background_tasks/periodic.rs‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
use std::sync::Arc;
2-
3-
use chrono::Duration;
2+
use std::time::Duration;
43

54
use teloxide::types::UserId;
65
use tokio::time::sleep;
@@ -47,7 +46,7 @@ pub fn start_periodic_tasks(
4746
let services = services_clone.clone();
4847
tokio::spawn(async move {
4948
loop {
50-
sleep(Duration::minutes(15).to_std().expect("no overflow")).await;
49+
sleep(Duration::from_mins(15)).await;
5150
let span = tracing::info_span!("periodic_refetch_stickers");
5251
let bot = bot.clone();
5352
let database = database.clone();
@@ -94,9 +93,9 @@ pub fn start_periodic_tasks(
9493
.instrument(span)
9594
.await;
9695
if result_len == 0 && new_offset != offset && !should_wait_longer {
97-
sleep(Duration::seconds(5).to_std().expect("no overflow")).await;
96+
sleep(Duration::from_secs(5)).await;
9897
} else {
99-
sleep(Duration::minutes(20).to_std().expect("no overflow")).await;
98+
sleep(Duration::from_mins(5)).await;
10099
// TODO: i should really add a config for the automatic tasks
101100
}
102101
offset = new_offset;
@@ -118,15 +117,15 @@ pub fn start_periodic_tasks(
118117
}
119118
.instrument(span)
120119
.await;
121-
sleep(Duration::hours(24).to_std().expect("no overflow")).await;
120+
sleep(Duration::from_hours(24)).await;
122121
}
123122
});
124123

125124
let bot = bot_clone;
126125
let database = database_clone.clone();
127126
tokio::spawn(async move {
128127
loop {
129-
sleep(Duration::days(7).to_std().expect("no overflow")).await;
128+
sleep(Duration::from_hours(7 * 24)).await;
130129
let span = tracing::info_span!("periodic_database_export");
131130
let bot = bot.clone();
132131
let database = database.clone();
@@ -156,7 +155,7 @@ pub fn start_periodic_tasks(
156155
}
157156
.instrument(span)
158157
.await;
159-
sleep(Duration::hours(23).to_std().expect("no overflow")).await;
158+
sleep(Duration::from_hours(23)).await;
160159
}
161160
});
162161

@@ -165,7 +164,7 @@ pub fn start_periodic_tasks(
165164
let config = config_clone.clone();
166165
tokio::spawn(async move {
167166
loop {
168-
sleep(Duration::days(2).to_std().expect("no overflow")).await;
167+
sleep(Duration::from_hours(24 * 2)).await;
169168
let span = tracing::info_span!("periodic_tag_insertion");
170169
// TODO: do daily; also refetch e6 tags
171170
let vector_db = vector_db.clone();
@@ -185,7 +184,7 @@ pub fn start_periodic_tasks(
185184
let vector_db = vector_db_clone.clone();
186185
tokio::spawn(async move {
187186
loop {
188-
sleep(Duration::seconds(10).to_std().expect("no overflow")).await;
187+
sleep(Duration::from_secs(10)).await;
189188
let span = tracing::info_span!("periodic_sticker_file_cleanup");
190189
let database = database.clone();
191190
let vector_db = vector_db.clone();
@@ -195,7 +194,7 @@ pub fn start_periodic_tasks(
195194
}
196195
.instrument(span)
197196
.await;
198-
sleep(Duration::hours(23).to_std().expect("no overflow")).await;
197+
sleep(Duration::from_hours(23)).await;
199198
}
200199
});
201200
}
@@ -306,7 +305,7 @@ async fn fix_missing_tag_implications(
306305
database
307306
.tag_file(&file.id, &vec![implication.clone()], None)
308307
.await?;
309-
sleep(Duration::seconds(1).to_std().expect("no overflow")).await;
308+
sleep(Duration::from_secs(1)).await;
310309
}
311310
}
312311
}

‎fuzzle/src/services/import_service.rs‎

Lines changed: 60 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use crate::{
1919
Histogram, automerge, calculate_color_histogram, calculate_sticker_file_hash,
2020
fetch_sticker_file,
2121
},
22-
util::{Emoji, Required, decode_sticker_set_id, is_wrong_file_id_error},
22+
util::{Emoji, FloatIteratorExt, Required, decode_sticker_set_id, is_wrong_file_id_error},
2323
};
2424

2525
#[derive(Clone)]
@@ -420,32 +420,7 @@ impl ImportService {
420420

421421
// todo: tag animated?
422422
for sticker in stickers_not_in_database_yet {
423-
let result = self
424-
.fetch_sticker_and_save_to_db(sticker.clone(), set.name.clone())
425-
.await;
426-
427-
match result {
428-
Err(InternalError::Teloxide(teloxide::RequestError::Api(api_err)))
429-
if is_wrong_file_id_error(&api_err) =>
430-
{
431-
tracing::warn!("invalid file_id for a sticker, continuing");
432-
}
433-
Err(InternalError::Database(DatabaseError::TryingToInsertRemovedSticker)) => {
434-
tracing::info!("trying to insert removed sticker")
435-
}
436-
Err(other) => {
437-
if other.is_timeout_error() {
438-
tracing::warn!("sticker fetch timed out, continuing");
439-
} else {
440-
return Err(other);
441-
}
442-
}
443-
Ok(()) => {}
444-
}
445-
446-
self.analyze_sticker(sticker.file.unique_id.clone()).await?;
447-
self.possibly_auto_ban_set(&set.name, set.stickers.len())
448-
.await?;
423+
self.import_new_sticker(sticker, set.clone()).await?;
449424
}
450425
for sticker in saved_stickers.clone() {
451426
let Some(s) = set.stickers.iter().find(|s| s.file.unique_id == sticker.id) else {
@@ -457,6 +432,16 @@ impl ImportService {
457432
.update_sticker(sticker.id, s.file.id.clone())
458433
.await?;
459434
}
435+
let files = self.database.get_sticker_files_by_ids(&[sticker.sticker_file_id]).await?;
436+
assert!(files.len() <= 1);
437+
if let Some(file) = files.first() {
438+
if file.thumbnail_file_id.is_none() && s.thumb.is_some() {
439+
tracing::info!(set_id = %sticker.sticker_set_id, "found previously non-imported thumbnail");
440+
self.import_new_sticker(s.clone(), set.clone()).await?;
441+
}
442+
} else {
443+
self.import_new_sticker(s.clone(), set.clone()).await?;
444+
}
460445
}
461446

462447
let deleted_stickers = saved_stickers
@@ -493,6 +478,40 @@ impl ImportService {
493478
Ok(())
494479
}
495480

481+
482+
async fn import_new_sticker(&self, sticker: teloxide::types::Sticker, set: teloxide::types::StickerSet) -> Result<(), InternalError> {
483+
let result = self
484+
.fetch_sticker_and_save_to_db(sticker.clone(), set.name.clone())
485+
.await;
486+
487+
match result {
488+
Err(InternalError::Teloxide(teloxide::RequestError::Api(api_err)))
489+
if is_wrong_file_id_error(&api_err) =>
490+
{
491+
tracing::warn!("invalid file_id for a sticker, continuing");
492+
}
493+
Err(InternalError::Database(DatabaseError::TryingToInsertRemovedSticker)) => {
494+
tracing::info!("trying to insert removed sticker")
495+
}
496+
Err(other) => {
497+
if other.is_timeout_error() {
498+
tracing::warn!("sticker fetch timed out, continuing");
499+
} else {
500+
return Err(other);
501+
}
502+
}
503+
Ok(()) => {}
504+
}
505+
506+
if let Err(err) = self.analyze_sticker(sticker.file.unique_id.clone()).await {
507+
tracing::error!(error = %err, "error during analyze");
508+
}
509+
if let Err(err) = self.possibly_auto_ban_set(&set.name, set.stickers.len()).await {
510+
tracing::error!(error = %err, "error during auto ban");
511+
}
512+
Ok(())
513+
}
514+
496515
#[tracing::instrument(skip(self), err(Debug))]
497516
async fn get_clip_embedding(
498517
&self,
@@ -599,8 +618,17 @@ impl ImportService {
599618
}
600619
let matches = self
601620
.vector_db
602-
.find_banned_stickers_given_vector(clip_vector.clone(), 5, Some(0.6))
621+
.find_banned_stickers_given_vector(clip_vector.clone(), 5, Some(0.7))
603622
.await?;
623+
let Some(worst_score) = matches.iter().map(|m| m.score).fmax() else {
624+
return Ok(false);
625+
};
626+
let best_regular_matches = self
627+
.vector_db
628+
.find_stickers_given_vector(clip_vector.clone(), 100, 0, Some(worst_score))
629+
.await?;
630+
let best_regular_matches_len = best_regular_matches.len();
631+
604632
let mut should_ban = false;
605633
for m in matches {
606634
if let Some(banned_sticker) = self.database.get_banned_sticker(&m.file_hash).await? {
@@ -610,13 +638,11 @@ impl ImportService {
610638

611639
// if the sticker under consideration for a ban matches a banned sticker closer than
612640
// the best match from a different set, also ban
613-
let best_regular_matches = self
614-
.vector_db
615-
.find_stickers_given_vector(clip_vector.clone(), 100, 0, Some(m.score))
616-
.await?; // TODO: call outside loop
617-
let best_regular_matches_len = best_regular_matches.len();
618-
for m in best_regular_matches {
641+
for m in &best_regular_matches {
619642
let score_non_banned_sticker = m.score;
643+
if score_non_banned_sticker < score_banned_sticker {
644+
continue;
645+
}
620646
if let Some(matched_sticker) = self
621647
.database
622648
.get_some_sticker_by_file_id(&m.file_hash)

‎fuzzle/src/util/float_ext.rs‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
pub trait FloatIteratorExt: Iterator + Sized {
2+
fn fmax(self) -> Option<Self::Item> where Self::Item: FloatTraits;
3+
fn fmin(self) -> Option<Self::Item> where Self::Item: FloatTraits;
4+
}
5+
6+
impl<I: Iterator> FloatIteratorExt for I {
7+
fn fmax(self) -> Option<Self::Item> where Self::Item: FloatTraits {
8+
self.filter(|x| !x.is_nan_val()).max_by(|a, b| a.total_compare(b))
9+
}
10+
11+
fn fmin(self) -> Option<Self::Item> where Self::Item: FloatTraits {
12+
self.filter(|x| !x.is_nan_val()).min_by(|a, b| a.total_compare(b))
13+
}
14+
}
15+
16+
pub trait FloatTraits {
17+
fn is_nan_val(&self) -> bool;
18+
fn total_compare(&self, other: &Self) -> std::cmp::Ordering;
19+
}
20+
21+
macro_rules! impl_float_traits {
22+
($($t:ty),*) => {
23+
$(impl FloatTraits for $t {
24+
fn is_nan_val(&self) -> bool { self.is_nan() }
25+
fn total_compare(&self, other: &Self) -> std::cmp::Ordering { self.total_cmp(other) }
26+
})*
27+
};
28+
}
29+
impl_float_traits!(f32, f64);

‎fuzzle/src/util/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ mod bot;
22
mod emoji;
33
mod parsers;
44
mod required;
5+
mod float_ext;
56

67
pub use bot::*;
78
use chrono::{Duration, NaiveDateTime, TimeDelta};
@@ -10,6 +11,7 @@ pub use parsers::*;
1011
use rand::Rng;
1112
use regex::Regex;
1213
pub use required::*;
14+
pub use float_ext::*;
1315

1416
pub fn format_relative_time(time: NaiveDateTime) -> String {
1517
let now = chrono::Utc::now().naive_utc();

0 commit comments

Comments
 (0)