Added basic Open AI chat

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);
+  }
+}