Added basic Open AI chat

src/commands/mod.rs

@@ -1,2 +1,4 @@
+pub mod audio;
+pub mod help;
+pub mod oai;
 pub mod ping;
-pub mod audio;

src/commands/oai.rs

@@ -0,0 +1,206 @@
+
+use std::error::Error;
+use std::fmt;
+
+use log::{error, debug, trace};
+
+use serde::{Serialize, Deserialize};
+use serde_json::Value;
+use serenity::model::channel::Message;
+use serenity::prelude::*;
+
+pub struct OAI {
+  pub client: reqwest::Client,
+  pub base_url: String,
+  pub max_attempts: u64,
+  pub token: String
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct ChatCompletionRequest {
+  model: String,
+  messages: Vec<ChatCompletionMessage>,
+  /// Value between 0 and 2
+  #[serde(skip_serializing_if = "Option::is_none")]
+  temperature: Option<f64>,
+  /// Value between 0 and 1
+  #[serde(skip_serializing_if = "Option::is_none")]
+  top_p: Option<f64>,
+  #[serde(skip_serializing_if = "Option::is_none")]
+  n: Option<f64>,
+  #[serde(skip_serializing_if = "Option::is_none")]
+  max_tokens: Option<u64>,
+  /// Value between -2.0 and 2.0
+  #[serde(skip_serializing_if = "Option::is_none")]
+  presence_penalty: Option<f64>,
+  /// Value between -2.0 and 2.0
+  #[serde(skip_serializing_if = "Option::is_none")]
+  frequency_penalty: Option<f64>,
+  #[serde(skip_serializing_if = "Option::is_none")]
+  user: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct ChatCompletionMessage {
+  role: String,
+  content: String
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+enum Role {
+  #[serde(rename = "system")]
+  SYSTEM,
+  #[serde(rename = "user")]
+  USER,
+  #[serde(rename = "assistant")]
+  ASSISTANT,
+  #[serde(rename = "function")]
+  FUNCTION
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct ChatCompletionResponse {
+  id: String,
+  object: String,
+  created: u64,
+  model: String,
+  usage: Usage,
+  choices: Vec<Choice>
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct Usage {
+  prompt_tokens: u64,
+  completion_tokens: u64,
+  total_tokens: u64
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct Choice {
+  message: ChatCompletionMessage,
+  finish_reason: String,
+  index: u64
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct ResponseError {
+  code: Option<u64>,
+  message: Option<String>,
+  param: Option<String>,
+  #[serde(rename = "type")]
+  error_type: Option<String>
+}
+
+#[derive(Debug)]
+struct OAIError {
+  pub message: String
+}
+
+impl fmt::Display for OAIError {
+  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+    write!(f, "OAIError: {}", self.message)
+  }
+}
+
+impl Error for OAIError {}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+enum ResponseEvent {
+  ChatCompletionResponse(ChatCompletionResponse),
+  ResponseError(ResponseError)
+}
+
+impl OAI {
+  async fn get_request(&self, request: ChatCompletionRequest) -> Result<ChatCompletionResponse, OAIError> {
+    let uri = format!("{}/chat/completions", self.base_url);
+    let body = serde_json::to_string(&request).unwrap();
+    trace!("Sending request to {}", uri);
+    let value = match match self.client
+      .post(&uri)
+      .bearer_auth(&self.token)
+      .header("Content-Type", "application/json".to_string())
+      .body(body)
+      .send()
+      .await {
+        Ok(r) => r,
+        Err(err) => return Err(OAIError {
+          message: format!("Could not send request to OpenAI: {}", err),
+        })
+      }
+      .json::<Value>()
+      .await {
+        Ok(r) => r,
+        Err(err) => return Err(OAIError {
+          message: format!("Could not read response from OpenAI: {}", err)
+        })
+      };
+
+    debug!("Received response from OpenAI: {:?}", value);
+    
+    // let response = match serde_json::from_value::<OAIResponseEvent>(value) {
+    //   Ok(r) => {
+    //     match r {
+    //       OAIResponseEvent::OAIResponse(r) => r,
+    //       OAIResponseEvent::OAIError(e) => return Err(OAIError { message: e.message.unwrap_or("Unknown error".to_string()) })
+    //     }
+    //   },
+    //   Err(err) => return Err(OAIError {
+    //     message: format!("Could not parse response from OpenAI: {}", err)
+    //   })
+    // };
+    let response = match serde_json::from_value::<ChatCompletionResponse>(value) {
+      Ok(r) => r,
+      Err(err) => return Err(OAIError {
+        message: format!("Could not parse response from OpenAI: {}", err)
+      })
+    };
+
+    Ok(response)
+  }
+}
+
+pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
+  let bot_mention = format!("<@{}>", ctx.cache.current_user_id().0);
+  let parsed_content = msg.content.replace(bot_mention.as_str(), "");
+  debug!("Generating response for message: {}", msg.content);
+
+  let messages = vec![
+    ChatCompletionMessage {
+      role: "system".to_string(),
+      content: "You are a helpful Discord bot named Siren.".to_string()
+    },
+    ChatCompletionMessage {
+      role: "user".to_string(),
+      content: parsed_content
+    },
+  ];
+
+  let request = ChatCompletionRequest {
+    model: "gpt-3.5-turbo".to_string(),
+    messages,
+    temperature: None,
+    top_p: None,
+    n: None,
+    max_tokens: Some(2000),
+    presence_penalty: None,
+    frequency_penalty: None,
+    user: Some(msg.author.name.clone())
+  };
+  let response = match oai.get_request(request).await {
+    Ok(r) => {
+      debug!("Received response from OpenAI");
+      if !r.choices.is_empty() {
+        r.choices[0].message.content.clone()
+      } else {
+        "No reply received".to_string()
+      }
+    }
+    Err(err) => {
+      error!("Could not get response from OpenAI: {}", err.message);
+      err.message
+    }
+  };
+  if let Err(why) = msg.channel_id.say(&ctx.http, response).await {
+    error!("Cannot send message: {}", why);
+  }
+}

src/main.rs

@@ -8,15 +8,40 @@ use serenity::async_trait;
 use serenity::framework::StandardFramework;
 use serenity::model::application::interaction::Interaction;
 use serenity::model::gateway::Ready;
+use serenity::model::channel::Message;
 use serenity::http::Http;
 use serenity::prelude::*;
 use songbird::SerenityInit;
 
 mod commands;
-struct Handler;
+struct Handler {
+  // Open AI Config
+  oai: Option<commands::oai::OAI>
+}
 
 #[async_trait]
 impl EventHandler for Handler {
+  async fn message(&self, ctx: Context, msg: Message) {
+    // Ignore messages from bots
+    if msg.author.bot {
+      return;
+    }
+    match &self.oai {
+      Some(oai) => {
+        match msg.mentions_me(&ctx.http).await {
+          Ok(mentioned) => {
+            if mentioned {
+              commands::oai::generate_response(&ctx, &msg, oai).await;
+            }
+          }
+          Err(why) => warn!("Could not check mentions: {:?}", why)
+        };
+      }
+      None => {}
+    }
+
+  }
+
   async fn interaction_create(&self, ctx: Context, interaction: Interaction) {
     if let Interaction::ApplicationCommand(command) = interaction {
       match command.data.name.as_str() {
@@ -63,7 +88,7 @@ impl EventHandler for Handler {
 #[tokio::main]
 async fn main() {
   dotenv().ok();
-  env_logger::init();
+  env_logger::init_from_env(env_logger::Env::default().filter_or("RUST_LOG", "warn,siren=info"));
 
   let token: String = env::var("DISCORD_TOKEN").expect("Expected a token in the environment");
   let intents: GatewayIntents = GatewayIntents::all();
@@ -90,9 +115,22 @@ async fn main() {
       .owners(owners)
       .prefix("!")
     );
+  
+  let handler = match env::var("OPENAI_API_KEY") {
+    Ok(token) => {
+      info!("Loaded OpenAI token");
+      Handler {
+        oai: Some(commands::oai::OAI { client: reqwest::Client::new(), base_url: "https://api.openai.com/v1".to_string(), max_attempts: 5, token })
+      }
+    }
+    Err(err) => {
+      warn!("Could not load OpenAI token: {}", err);
+      Handler { oai: None }
+    }
+  };
 
   let mut client = Client::builder(token, intents)
-    .event_handler(Handler)
+    .event_handler(handler)
     .framework(framework)
     .register_songbird()
     .await