|
| 1 | +use std::cell::LazyCell; |
| 2 | + |
| 3 | +use extism_pdk::FnResult; |
| 4 | +use extism_pdk::FromBytes; |
| 5 | +use extism_pdk::HttpRequest; |
| 6 | +use extism_pdk::Json; |
| 7 | +use extism_pdk::ToBytes; |
| 8 | +use extism_pdk::config; |
| 9 | +use extism_pdk::http; |
| 10 | +use extism_pdk::plugin_fn; |
| 11 | +use extism_pdk::var; |
| 12 | +use lemmy_api_common::comment::CommentInsertForm; |
| 13 | +use lemmy_api_common::language::Language as LemmyLanguage; |
| 14 | +use lemmy_api_common::language::LanguageId; |
| 15 | +use lemmy_api_common::plugin::PluginMetadata; |
| 16 | +use lemmy_api_common::post::PostInsertForm; |
| 17 | +use lemmy_api_common::site::GetSiteResponse; |
| 18 | +use lingua::Language; |
| 19 | +use lingua::LanguageDetector; |
| 20 | +use lingua::LanguageDetectorBuilder; |
| 21 | +use serde::Deserialize; |
| 22 | +use serde::Serialize; |
| 23 | + |
| 24 | +// Returns info about the plugin which gets included in /api/v4/site |
| 25 | +#[plugin_fn] |
| 26 | +pub fn metadata() -> FnResult<Json<PluginMetadata>> { |
| 27 | + // initialize the detector because it takes a long time (~5s) |
| 28 | + LazyCell::<LanguageDetector>::force(&DETECTOR); |
| 29 | + |
| 30 | + Ok(Json(PluginMetadata::new( |
| 31 | + "Lingua", |
| 32 | + "https://github.com/LemmyNet/lemmy-plugins/", |
| 33 | + "Automatic language tagging for posts and comments", |
| 34 | + ))) |
| 35 | +} |
| 36 | + |
| 37 | +// Usage: https://docs.rs/lingua/1.7.2/lingua/index.html |
| 38 | +// There are various optimizations available, which could be exposed as plugin settings |
| 39 | +const DETECTOR: LazyCell<LanguageDetector> = |
| 40 | + LazyCell::new(|| LanguageDetectorBuilder::from_all_languages().build()); |
| 41 | + |
| 42 | +#[plugin_fn] |
| 43 | +pub fn local_post_before_create( |
| 44 | + Json(mut form): Json<PostInsertForm>, |
| 45 | +) -> FnResult<Json<PostInsertForm>> { |
| 46 | + let content = format!("{} {}", form.name, form.body.clone().unwrap_or_default()); |
| 47 | + detect_language(content, &mut form.language_id)?; |
| 48 | + Ok(Json(form)) |
| 49 | +} |
| 50 | + |
| 51 | +#[plugin_fn] |
| 52 | +pub fn local_comment_before_create( |
| 53 | + Json(mut form): Json<CommentInsertForm>, |
| 54 | +) -> FnResult<Json<CommentInsertForm>> { |
| 55 | + detect_language(form.content.clone(), &mut form.language_id)?; |
| 56 | + Ok(Json(form)) |
| 57 | +} |
| 58 | + |
| 59 | +#[plugin_fn] |
| 60 | +pub fn federated_post_before_receive( |
| 61 | + Json(mut form): Json<PostInsertForm>, |
| 62 | +) -> FnResult<Json<PostInsertForm>> { |
| 63 | + let content = format!("{} {}", form.name, form.body.clone().unwrap_or_default()); |
| 64 | + detect_language(content, &mut form.language_id)?; |
| 65 | + Ok(Json(form)) |
| 66 | +} |
| 67 | + |
| 68 | +#[plugin_fn] |
| 69 | +pub fn federated_comment_before_receive( |
| 70 | + Json(mut form): Json<CommentInsertForm>, |
| 71 | +) -> FnResult<Json<CommentInsertForm>> { |
| 72 | + detect_language(form.content.clone(), &mut form.language_id)?; |
| 73 | + Ok(Json(form)) |
| 74 | +} |
| 75 | + |
| 76 | +fn detect_language(content: String, language_id: &mut Option<LanguageId>) -> FnResult<()> { |
| 77 | + if language_id.is_none() { |
| 78 | + let detected_language: Option<Language> = DETECTOR.detect_language_of(content); |
| 79 | + |
| 80 | + if let Some(detected_language) = detected_language { |
| 81 | + let all_langs = all_languages()?; |
| 82 | + let lang = all_langs |
| 83 | + .iter() |
| 84 | + .find(|l| l.code == detected_language.iso_code_639_1().to_string()); |
| 85 | + *language_id = lang.map(|l| l.id); |
| 86 | + } |
| 87 | + } |
| 88 | + Ok(()) |
| 89 | +} |
| 90 | + |
| 91 | +#[derive(Deserialize, Serialize, FromBytes, ToBytes)] |
| 92 | +#[encoding(Json)] |
| 93 | +struct AllLanguages(Vec<LemmyLanguage>); |
| 94 | + |
| 95 | +fn all_languages() -> FnResult<Vec<LemmyLanguage>> { |
| 96 | + const KEY: &str = "all_languages"; |
| 97 | + let langs = var::get::<AllLanguages>(KEY)?; |
| 98 | + if let Some(langs) = langs { |
| 99 | + Ok(langs.0) |
| 100 | + } else { |
| 101 | + let lemmy_url = config::get("lemmy_url")?.unwrap(); |
| 102 | + let req = HttpRequest { |
| 103 | + url: format!("{lemmy_url}api/v4/site"), |
| 104 | + headers: Default::default(), |
| 105 | + method: Some("GET".to_string()), |
| 106 | + }; |
| 107 | + let site: GetSiteResponse = http::request::<()>(&req, None)?.json()?; |
| 108 | + let langs = site.all_languages; |
| 109 | + var::set(KEY, AllLanguages(langs.clone()))?; |
| 110 | + Ok(langs) |
| 111 | + } |
| 112 | +} |
0 commit comments