telitairos_bot/
lib.rs

1//! # A fully functional AI Powered Telegram bot
2//!
3//! ## Supported commands
4//! You can either:
5//! - `/ask` for a specified question.
6//! - `/mediate` to read the last N messages of a chat group and mitigate an argument.
7//!
8//! ## Environment variables needed
9//!
10//! ```bash
11//! - TELOXIDE_TOKEN= "/* Your Telegram Bot API Key */"
12//! - OPENAI_API_KEY= "/* Your OpenAI API Key */"
13//! - OPENAI_ORG_ID= "/* Your OpenAI Organization ID */"
14//! ```
15//!
16//! ## Basic usage
17//!
18//! You need to specify the personality of the bot as well as its criteria when mitigating an argument.
19//! A size for the context of the N last messages of the chat group is also needed.
20//!
21//! For a detailed example go to [TelitairoBot]
22//!
23//! ### Example
24//!
25//! TelitairoBot struct implements Default trait, so you can start a bot with a generic personality
26//! by just doing this:
27//!
28//! ```
29//! # use telitairos_bot::TelitairoBot;
30//! let telitairo_bot: TelitairoBot = Default::default();
31//! ```
32//!
33//! But if you want to set your own Bot's personality you can use the `new()` function like this:
34//!
35//! ```
36//! # use telitairos_bot::TelitairoBot;
37//! /* #[tokio::main] */
38//! /*async*/ fn main() {
39//! let telitairo_bot = TelitairoBot::new(
40//!     String::from("personality"),
41//!     String::from("mediation criteria"),
42//!     200 /* buffer size */,
43//! );
44//!
45//! /* telitairo_bot.dispatch().await; */
46//! # }
47//! ```
48//!
49mod bot;
50mod gpt;
51mod types;
52
53use crate::bot::*;
54use std::collections::HashMap;
55use std::sync::Arc;
56use teloxide::dispatching::{HandlerExt, UpdateFilterExt};
57use teloxide::{dptree, Bot};
58use tokio::sync::RwLock;
59
60/// Defines the bot behavior
61#[derive(Clone, Debug, PartialEq)]
62pub struct TelitairoBot {
63    /// String to define the bot personality, a descriptive short prompt.
64    ///
65    /// # Example
66    /// ```bash
67    /// "You are a virtual assistant with a touch of acid humour and you love potatoes"
68    /// ```
69    pub personality: String,
70    /// String to define the bot action when `/mediate` command is sent. descriptive short prompt.
71    ///
72    /// # Example
73    /// ```bash
74    /// "Take the messages, search for possible discussions and choose one side"
75    /// ```
76    pub mediate_query: String,
77    /// Buffer capacity, size of the context for the `/mediate` command.
78    ///
79    /// Bot will assert if buffer_size = 0
80    pub buffer_size: usize,
81}
82
83impl TelitairoBot {
84    /// Creates a new bot with the selected parameters
85    ///
86    /// Creation will assert if strings are empty or if buffer_size = 0
87    pub fn new(personality: String, mediate_query: String, buffer_size: usize) -> TelitairoBot {
88        assert!(buffer_size > 0 && !personality.is_empty() && !mediate_query.is_empty());
89
90        TelitairoBot {
91            personality,
92            mediate_query,
93            buffer_size,
94        }
95    }
96
97    /// Open a listening for incoming messages and commands
98    pub async fn dispatch(&self) {
99        let bot = Bot::from_env();
100        let buffer_store: types::BufferStore = Arc::new(RwLock::new(HashMap::new()));
101
102        let handler = Update::filter_message()
103            .branch(
104                dptree::entry()
105                    .filter_command::<ai::AiCommand>()
106                    .endpoint(ai::handle_ai_commands),
107            )
108            .branch(
109                dptree::entry()
110                    .filter_command::<admin::AdminCommand>()
111                    .endpoint(admin::handle_admin_commands),
112            )
113            .branch(dptree::endpoint(handle_messages));
114
115        Dispatcher::builder(bot, handler)
116            .dependencies(dptree::deps![buffer_store, self.clone()])
117            .default_handler(|update| async move {
118                log::warn!("Unhandled update: {:#?}", update);
119            })
120            .error_handler(LoggingErrorHandler::with_custom_text(
121                "An error occurred in the dispatcher",
122            ))
123            .enable_ctrlc_handler()
124            .build()
125            .dispatch()
126            .await
127    }
128}
129
130/// Default attributes for a generic assistant bot
131impl Default for TelitairoBot {
132    fn default() -> Self {
133        TelitairoBot {
134            personality: types::DEFAULT_PERSONALITY.to_string(),
135            mediate_query: types::DEFAULT_MEDIATION_QUERY.to_string(),
136            buffer_size: types::DEFAULT_BUFFER_SIZE,
137        }
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn test_default_and_new() {
147        let telitairo_default = TelitairoBot::default();
148        let telitairo_partially_default = TelitairoBot {
149            buffer_size: 200,
150            ..Default::default()
151        };
152        let telitairo_new = TelitairoBot::new(
153            String::from(
154                "You are a virtual assistant with a touch of acid humour and you love potatoes",
155            ),
156            String::from("Take the messages, search for possible discussions and choose one side"),
157            200,
158        );
159
160        assert_eq!(telitairo_default, telitairo_new);
161        assert_eq!(telitairo_partially_default, telitairo_new);
162    }
163}