izzilis/src/main.rs

161 lines
4.4 KiB
Rust

use std::{convert::Infallible, error::Error, sync::Arc, time::Duration};
use async_std::io::stdin;
use futures_timer::Delay;
use mammut::{
apps::{AppBuilder, Scopes},
Mastodon, Registration,
};
use rand::Rng;
use telegram_bot::{Api, ChatId, ChatRef, ToChatRef};
use crate::{
publish::MastodonPublisher,
selection::{telegram::get_chat_ref, SelectorExt, TelegramSelector},
};
use futures::{channel::mpsc::channel, sink::unfold, SinkExt, StreamExt, TryStreamExt};
use model::SampleModelExt;
mod config;
mod model;
mod publish;
mod selection;
const CONFIG_PATH: &str = "bot_config.json";
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
let mut cfg = match config::Config::from(CONFIG_PATH) {
Ok(cfg) => Ok(cfg),
Err(_) => {
println!(
"Failed reading config at [{}], writing default",
CONFIG_PATH
);
config::Config::default().save(CONFIG_PATH)?;
Err(anyhow::anyhow!("Invalid configuration"))
}
}?;
let app = AppBuilder {
client_name: "izzilis",
redirect_uris: "urn:ietf:wg:oauth:2.0:oob",
scopes: Scopes::Write,
website: None,
};
let mut registration = Registration::new(cfg.fediverse_base_url.clone());
registration.register(app)?;
let mastodon = if let Some(data) = &cfg.fediverse_token {
Mastodon::from_data(data.clone())
} else {
let url = registration.authorise()?;
println!("{}", url);
let mut buffer = String::new();
stdin().read_line(&mut buffer).await?;
let fedi = registration.create_access_token(buffer)?;
cfg.fediverse_token = Some(fedi.data.clone());
cfg.save(CONFIG_PATH)?;
fedi
};
let publisher = MastodonPublisher::new(mastodon);
let api = Arc::new(Api::new(cfg.bot_token.clone()));
let chat = get_chat_ref(
api.clone(),
if ChatId::new(0) == cfg.chat_ref {
None
} else {
Some(cfg.chat_ref.clone().to_chat_ref())
},
)
.await
.expect("could not get chat ref");
let chat_ref = chat.lock().await.clone();
if let ChatRef::Id(id) = chat_ref {
cfg.chat_ref = id;
cfg.save(CONFIG_PATH)?;
}
let (sender, receiver) = channel(cfg.post_buffer as usize);
let cfg_clone = cfg.clone();
let mut model = TelegramSelector::new(
api,
chat,
Box::pin(unfold((), move |_, chat_ref| {
let mut cfg_clone = cfg_clone.clone();
async move {
if let ChatRef::Id(id) = &chat_ref {
cfg_clone.chat_ref = id.clone();
let _ = cfg_clone.save(CONFIG_PATH);
}
Ok::<_, Infallible>(())
}
})),
)
.filter(
model::GPTSampleModel::new(
cfg.python_path.clone(),
cfg.gpt_code_path.clone(),
vec![
"generate_unconditional_samples.py".to_string(),
"--model_name".to_string(),
cfg.model_name.clone(),
"--temperature".to_string(),
cfg.temperature.clone(),
"--top_k".to_string(),
cfg.top_k.clone(),
"--nsamples".to_string(),
"1".to_string(),
],
)
.into_stream()
.try_filter(|message| {
let not_empty = !message.is_empty();
async move { not_empty }
}),
)
.map_err(|e| Box::new(e) as Box<dyn Error>);
tokio::spawn(async move {
sender
.sink_map_err(|e| Box::new(e) as Box<dyn Error>)
.send_all(&mut model)
.await
.expect("Broken buffer");
});
publisher
.sink_map_err(|e| Box::new(e) as Box<dyn Error>)
.send_all(
&mut receiver
.then(|item| {
let interval_seconds = cfg.interval_seconds.clone();
Box::pin(async move {
Delay::new(Duration::from_secs(
rand::thread_rng()
.gen_range(interval_seconds.min..=interval_seconds.max),
))
.await;
item
})
})
.map(Ok),
)
.await?;
return Ok(());
}