Refactor discord bot's event handler and message processing

Cargo.toml

@@ -19,3 +19,4 @@ env_logger = "0.11.2"
 config = "0.14.0"
 url = "2.5.0"
 reqwest = "0.12.2"
+thiserror = "1.0.0"

script/run.sh

@@ -1,3 +1,3 @@
-export DISCORD_BOT_TOKEN="your_discord_bot_token"
+export DISCORD_TOKEN="your_discord_bot_token"
 export OPENAI_API_KEY="your_openai_api_key"
-./bootstrap
+./bootstrap

src/bin/bootstrap.rs

@@ -7,6 +7,7 @@ use std::path::Path;
 #[tokio::main]
 async fn main() {
     // Initialize the logger
+
     env_logger::init();
 
     // Load environment variables

src/discord/mod.rs

@@ -1,3 +1,103 @@
-mod bot;
-pub use bot::Error;
-pub use bot::Handler;
+use chatgpt::client::ChatGPT;
+use chatgpt::config::{ChatGPTEngine, ModelConfiguration, ModelConfigurationBuilder};
+use chatgpt::types::ResponseChunk;
+use futures::StreamExt;
+use serenity::async_trait;
+use serenity::builder::EditMessage;
+use serenity::client::{Context, EventHandler};
+use serenity::model::channel::Message;
+use std::time::Duration;
+use thiserror::Error;
+use tokio::time::interval;
+
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error(transparent)]
+    ModelConfigurationBuilderError(#[from] chatgpt::config::ModelConfigurationBuilderError),
+
+    #[error(transparent)]
+    ChatGPT(#[from] chatgpt::err::Error),
+}
+
+fn init_gpt_client(api_key: &str) -> Result<ChatGPT, Error> {
+    let config: ModelConfiguration = ModelConfigurationBuilder::default()
+        .engine(ChatGPTEngine::Gpt4)
+        .timeout(Duration::from_secs(500))
+        .build()?;
+
+    let client = ChatGPT::new_with_config(api_key, config)?;
+    Ok(client)
+}
+
+// The Handler struct is the event handler for the bot.
+pub struct Handler {
+    pub gpt_client: ChatGPT,
+}
+
+impl Handler {
+    pub fn new(api_key: &str) -> Result<Self, Error> {
+        let gpt_client = init_gpt_client(api_key)?;
+        Ok(Self { gpt_client })
+    }
+}
+
+#[async_trait]
+impl EventHandler for Handler {
+    async fn message<'a>(&'a self, ctx: Context, msg: Message) {
+        if msg.author.bot || !msg.content.starts_with('.') {
+            return;
+        }
+
+        let prompt = &msg.content[1..];
+        let mut stream = match self.gpt_client.send_message_streaming(prompt).await {
+            Ok(stream) => stream,
+            Err(e) => {
+                log::error!("Error sending message to ChatGPT: {:?}", e);
+                return;
+            }
+        };
+
+        let processing_message = match msg.channel_id.say(&ctx.http, "Processing...").await {
+            Ok(message) => message,
+            Err(why) => {
+                log::error!("Error sending message: {:?}", why);
+                return;
+            }
+        };
+
+        let mut result = String::new();
+        let mut interval = interval(Duration::from_millis(900));
+
+        loop {
+            tokio::select! {
+                chunk = stream.next() => {
+                    if let Some(chunk) = chunk {
+                        if let ResponseChunk::Content { delta, response_index: _ } = chunk {
+                                result.push_str(&delta);
+                        }
+                    } else {
+                        // Stream has ended, break the loop
+                        break;
+                    }
+                },
+                _ = interval.tick() => {
+                    if !result.is_empty() {
+                        let edit = EditMessage::default().content(&result);
+                        if let Err(why) = msg.channel_id.edit_message(&ctx.http, processing_message.id, edit).await {
+                            log::error!("Error editing message: {:?}", why);
+                        }
+                    }
+                }
+            }
+        }
+
+        // Ensure any remaining content is also sent
+        if !result.is_empty() {
+            let edit = EditMessage::default().content(&result);
+            let _ = msg
+                .channel_id
+                .edit_message(&ctx.http, processing_message.id, edit)
+                .await;
+        }
+    }
+}

src/lib.rs

@@ -1,3 +1,2 @@
 mod discord;
-pub use crate::discord::Error;
 pub use crate::discord::Handler;