diff --git a/home-mixer/candidate_hydrators/quoted_post_text_hydrator.rs b/home-mixer/candidate_hydrators/quoted_post_text_hydrator.rs index 7c6041bd..d188ed0b 100644 --- a/home-mixer/candidate_hydrators/quoted_post_text_hydrator.rs +++ b/home-mixer/candidate_hydrators/quoted_post_text_hydrator.rs @@ -23,31 +23,40 @@ impl Hydrator for QuotedPostTextHydrator { _query: &ScoredPostsQuery, candidates: &[PostCandidate], ) -> Vec> { - let quoted_ids: Vec = candidates + let fetch_ids: Vec = candidates .iter() - .filter_map(|c| c.quoted_tweet_id) + .flat_map(|c| { + c.quoted_tweet_id + .into_iter() + .chain(c.retweeted_tweet_id) + .chain(c.ancestors.iter().copied()) + }) .collect::>() .into_iter() .collect(); - let quoted_core = if quoted_ids.is_empty() { + let core = if fetch_ids.is_empty() { HashMap::new() } else { - self.tes_client.get_tweet_core_datas(quoted_ids).await + self.tes_client.get_tweet_core_datas(fetch_ids).await }; candidates .iter() .map(|candidate| { Ok(PostCandidate { - quoted_tweet_text: candidate.quoted_tweet_id.and_then(|id| { - match quoted_core.get(&id) { - Some(Ok(Some(data))) if !data.text.is_empty() => { - Some(data.text.clone()) - } - _ => None, - } - }), + quoted_tweet_text: candidate + .quoted_tweet_id + .and_then(|id| text_from_core(&core, id)), + retweeted_tweet_text: candidate + .retweeted_tweet_id + .and_then(|id| text_from_core(&core, id)), + ancestor_texts: candidate + .ancestors + .iter() + .copied() + .filter_map(|id| text_from_core(&core, id).map(|text| (id, text))) + .collect(), ..Default::default() }) }) @@ -56,6 +65,18 @@ impl Hydrator for QuotedPostTextHydrator { fn update(&self, candidate: &mut PostCandidate, hydrated: PostCandidate) { candidate.quoted_tweet_text = hydrated.quoted_tweet_text; + candidate.retweeted_tweet_text = hydrated.retweeted_tweet_text; + candidate.ancestor_texts = hydrated.ancestor_texts; + } +} + +fn text_from_core( + core: &HashMap, E>>, + id: u64, +) -> Option { + match core.get(&id) { + Some(Ok(Some(data))) if !data.text.is_empty() => Some(data.text.clone()), + _ => None, } } @@ -102,5 +123,78 @@ mod tests { assert_eq!(with_quote.quoted_tweet_text.as_deref(), Some("quoted text")); assert_eq!(without_quote.quoted_tweet_text, None); + assert!(with_quote.ancestor_texts.is_empty()); + } + + #[tokio::test] + async fn fills_ancestor_texts_without_changing_ancestor_ids() { + let mut core_data = HashMap::new(); + core_data.insert( + 20, + Some(PureCoreData { + text: "parent spam".to_string(), + ..Default::default() + }), + ); + core_data.insert( + 10, + Some(PureCoreData { + text: "root text".to_string(), + ..Default::default() + }), + ); + let client = Arc::new(MockTESClient { + core_data, + ..Default::default() + }); + let hydrator = QuotedPostTextHydrator::new(client as Arc); + + let mut reply = PostCandidate { + tweet_id: 30, + ancestors: vec![20, 10], + ..Default::default() + }; + + let hydrated = hydrator + .hydrate(&ScoredPostsQuery::default(), &[reply.clone()]) + .await; + hydrator.update(&mut reply, hydrated[0].clone().unwrap()); + + assert_eq!(reply.ancestors, vec![20, 10]); + assert_eq!(reply.ancestor_texts.get(&20).map(String::as_str), Some("parent spam")); + assert_eq!(reply.ancestor_texts.get(&10).map(String::as_str), Some("root text")); + assert_eq!(reply.quoted_tweet_text, None); + assert_eq!(reply.retweeted_tweet_text, None); + } + + #[tokio::test] + async fn fills_retweeted_tweet_text() { + let mut core_data = HashMap::new(); + core_data.insert( + 50, + Some(PureCoreData { + text: "original spam".to_string(), + ..Default::default() + }), + ); + let client = Arc::new(MockTESClient { + core_data, + ..Default::default() + }); + let hydrator = QuotedPostTextHydrator::new(client as Arc); + + let mut rt = PostCandidate { + tweet_id: 51, + retweeted_tweet_id: Some(50), + ..Default::default() + }; + + let hydrated = hydrator + .hydrate(&ScoredPostsQuery::default(), &[rt.clone()]) + .await; + hydrator.update(&mut rt, hydrated[0].clone().unwrap()); + + assert_eq!(rt.retweeted_tweet_text.as_deref(), Some("original spam")); + assert_eq!(rt.quoted_tweet_text, None); } } diff --git a/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs index e8c9c397..4803d289 100644 --- a/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs +++ b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs @@ -12,6 +12,7 @@ use crate::candidate_hydrators::language_code_hydrator::LanguageCodeHydrator; use crate::candidate_hydrators::media_info_hydrator::MediaInfoHydrator; use crate::candidate_hydrators::mutual_follow_jaccard_hydrator::MutualFollowJaccardHydrator; use crate::candidate_hydrators::quote_hydrator::QuoteHydrator; +use crate::candidate_hydrators::quoted_post_text_hydrator::QuotedPostTextHydrator; use crate::candidate_hydrators::semantic_id_hydrator::SemanticIdHydrator; use crate::candidate_hydrators::subscription_hydrator::SubscriptionHydrator; use crate::candidate_hydrators::topic_feedback_context_hydrator::TopicFeedbackContextHydrator; @@ -331,6 +332,7 @@ impl PhoenixCandidatePipeline { }), Box::new(core_data_hydrator), Box::new(QuoteHydrator::new(tes_client.clone(), socialgraph_client.clone()).await), + Box::new(QuotedPostTextHydrator::new(tes_client.clone())), Box::new(MediaInfoHydrator::new(media_info_cache_client).await), Box::new(SubscriptionHydrator::new(tes_client.clone()).await), Box::new(GizmoduckCandidateHydrator::new(gizmoduck_client).await), diff --git a/home-mixer/filters/following_viewer_muted_keyword_filter.rs b/home-mixer/filters/following_viewer_muted_keyword_filter.rs index b5ad3769..98d1171a 100644 --- a/home-mixer/filters/following_viewer_muted_keyword_filter.rs +++ b/home-mixer/filters/following_viewer_muted_keyword_filter.rs @@ -59,6 +59,7 @@ fn candidate_matches( ) -> bool { std::iter::once(candidate.tweet_text.as_str()) .chain(candidate.quoted_tweet_text.as_deref()) + .chain(candidate.retweeted_tweet_text.as_deref()) .chain(candidate.ancestor_texts.values().map(String::as_str)) .filter(|text| !text.is_empty()) .any(|text| matcher.matches(&tokenizer.tokenize(text))) diff --git a/home-mixer/filters/viewer_muted_keyword_filter.rs b/home-mixer/filters/viewer_muted_keyword_filter.rs index 482a8182..1d6aec6b 100644 --- a/home-mixer/filters/viewer_muted_keyword_filter.rs +++ b/home-mixer/filters/viewer_muted_keyword_filter.rs @@ -43,8 +43,7 @@ impl Filter for ViewerMutedKeywordFilter { let mut removed = Vec::new(); for candidate in candidates { - let tweet_text_token_sequence = tokenizer.tokenize(&candidate.tweet_text); - if matcher.matches(&tweet_text_token_sequence) { + if candidate_matches(&candidate, &tokenizer, &matcher) { removed.push(candidate); } else { kept.push(candidate); @@ -56,6 +55,19 @@ impl Filter for ViewerMutedKeywordFilter { } } +fn candidate_matches( + candidate: &PostCandidate, + tokenizer: &TweetTokenizer, + matcher: &MatchTweetGroup, +) -> bool { + std::iter::once(candidate.tweet_text.as_str()) + .chain(candidate.quoted_tweet_text.as_deref()) + .chain(candidate.retweeted_tweet_text.as_deref()) + .chain(candidate.ancestor_texts.values().map(String::as_str)) + .filter(|text| !text.is_empty()) + .any(|text| matcher.matches(&tokenizer.tokenize(text))) +} + #[cfg(test)] mod tests { use super::*; @@ -315,4 +327,83 @@ mod tests { assert_eq!(result.kept[0].tweet_id, 4); assert_eq!(result.removed.len(), 3); } + + #[tokio::test(flavor = "multi_thread")] + async fn drops_quote_whose_quoted_text_matches_muted_keyword() { + let filter = ViewerMutedKeywordFilter::new(); + let query = create_test_query(vec!["spam".to_string()]); + + let quote = PostCandidate { + tweet_id: 1, + tweet_text: "sharing this".to_string(), + quoted_tweet_text: Some("this is spam content".to_string()), + author_id: 12345, + ..Default::default() + }; + let clean = create_test_candidate(2, "sharing this"); + + let result = filter.filter(&query, vec![quote, clean]); + + assert_eq!(result.kept.len(), 1); + assert_eq!(result.kept[0].tweet_id, 2); + assert_eq!(result.removed.len(), 1); + assert_eq!(result.removed[0].tweet_id, 1); + } + + #[tokio::test(flavor = "multi_thread")] + async fn keeps_quote_when_quoted_text_does_not_match() { + let filter = ViewerMutedKeywordFilter::new(); + let query = create_test_query(vec!["spam".to_string()]); + + let quote = PostCandidate { + tweet_id: 1, + tweet_text: "sharing this".to_string(), + quoted_tweet_text: Some("ordinary news".to_string()), + author_id: 12345, + ..Default::default() + }; + + let result = filter.filter(&query, vec![quote]); + + assert_eq!(result.kept.len(), 1); + assert!(result.removed.is_empty()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn drops_reply_when_ancestor_text_matches_muted_keyword() { + let filter = ViewerMutedKeywordFilter::new(); + let query = create_test_query(vec!["spam".to_string()]); + + let mut reply = create_test_candidate(1, "ok"); + reply.ancestor_texts.insert(9, "this is spam content".to_string()); + + let result = filter.filter(&query, vec![reply, create_test_candidate(2, "ok")]); + + assert_eq!(result.kept.len(), 1); + assert_eq!(result.kept[0].tweet_id, 2); + assert_eq!(result.removed.len(), 1); + assert_eq!(result.removed[0].tweet_id, 1); + } + + #[tokio::test(flavor = "multi_thread")] + async fn drops_retweet_whose_original_text_matches_muted_keyword() { + let filter = ViewerMutedKeywordFilter::new(); + let query = create_test_query(vec!["spam".to_string()]); + + let rt = PostCandidate { + tweet_id: 1, + tweet_text: String::new(), + retweeted_tweet_id: Some(50), + retweeted_tweet_text: Some("this is spam content".to_string()), + author_id: 12345, + ..Default::default() + }; + + let result = filter.filter(&query, vec![rt, create_test_candidate(2, "ok")]); + + assert_eq!(result.kept.len(), 1); + assert_eq!(result.kept[0].tweet_id, 2); + assert_eq!(result.removed.len(), 1); + assert_eq!(result.removed[0].tweet_id, 1); + } } diff --git a/home-mixer/models/candidate.rs b/home-mixer/models/candidate.rs index f94d4a4b..3c86a3b5 100644 --- a/home-mixer/models/candidate.rs +++ b/home-mixer/models/candidate.rs @@ -39,6 +39,8 @@ pub struct PostCandidate { pub ancestor_users: Vec, pub ancestor_texts: HashMap, pub quoted_tweet_text: Option, + #[serde(default)] + pub retweeted_tweet_text: Option, pub min_video_duration_ms: Option, pub quoted_video_duration_ms: Option, pub author_followers_count: Option,