Merge pull request #6 from bensherriff/develop

Added Open AI basic functionality

.env.TEMPLATE

@@ -3,3 +3,4 @@ RUST_LOG=warn,siren=info
 POSTGRES_USER=siren
 POSTGRES_PASSWORD=
 POSTGRES_DB=siren
+OPENAI_API_KEY=

.gitignore

@@ -2,6 +2,7 @@
 target/
 .idea/
 .vscode/
+**/Cargo.lock
 
 audio/
 logs/

.version

@@ -1 +1 @@
-SIREN_VERSION=0.2.1
+SIREN_VERSION=0.2.2

Cargo.toml

@@ -1,7 +1,11 @@
 [package]
 name = "siren"
-version = "0.2.1"
+version = "0.2.2"
 edition = "2021"
+authors = ["Ben Sherriff <hello@bensherriff.com>"]
+repository = "https://github.com/bensherriff/siren"
+readme = "README.md"
+license = "GPL-3.0-or-later"
 
 [dependencies]
 dotenv = "0.15.0"
@@ -30,3 +34,8 @@ features = ["derive"]
 version = "0.11.18"
 default-features = false
 features = ["json", "rustls-tls"]
+
+[dependencies.diesel]
+version = "2.1.0"
+default-features = false
+features = ["postgres", "32-column-tables", "serde_json", "r2d2", "with-deprecated"]

Dockerfile

@@ -1,12 +1,13 @@
 FROM rust:1.67 as builder
 WORKDIR /siren
 RUN apt-get update && apt-get install -y cmake && apt-get auto-remove -y
-COPY . .
+ADD src ./src/
+ADD Cargo.toml ./
 RUN cargo build --release --bin siren
 
 FROM debian:bullseye-slim as packages
 WORKDIR /packages
-RUN apt-get update && apt-get install -y curl tar xz-utils
+RUN apt-get update && apt-get install -y libopus-dev libpq5 libpq-dev curl tar xz-utils
 RUN curl -L https://github.com/yt-dlp/yt-dlp/releases/latest/download/yt-dlp_linux > yt-dlp && \
   chmod +x yt-dlp
 RUN curl -L https://github.com/yt-dlp/FFmpeg-Builds/releases/download/latest/ffmpeg-master-latest-linux64-gpl.tar.xz > ffmpeg.tar.xz && \
@@ -14,7 +15,7 @@ RUN curl -L https://github.com/yt-dlp/FFmpeg-Builds/releases/download/latest/ffm
 
 FROM debian:bullseye-slim as runtime
 WORKDIR /siren
-RUN apt-get update && apt-get install -y libopus-dev ffmpeg youtube-dl
 COPY --from=builder /siren/target/release/siren /usr/local/bin/siren
 COPY --from=packages /packages /usr/bin
+# ADD migrations ./migrations/
 CMD ["siren"]

Makefile

@@ -5,6 +5,8 @@ include .version
 export $(shell sed 's/=.*//' .env)
 export $(shell sed 's/=.*//' .version)
 
+SIREN_IMAGES = $(shell docker images 'siren' -a -q)
+
 .PHONY: help build test up down exec clean
 
 build:
@@ -23,4 +25,4 @@ exec:
 	docker exec -it siren bash
 
 clean:
-	docker rmi siren
+	docker rmi $(SIREN_IMAGES)

README.md

@@ -18,7 +18,7 @@ Example Invite:
 ```
 https://discord.com/api/oauth2/authorize?client_id=<CLIENT_ID>&permissions=40671259392832&scope=bot%20applications.commands
 ```
- - The CLIENT_ID can be found in the General Information tab on the Discord Developer Portal for your application, under `Application ID`
+The CLIENT_ID can be found in the General Information tab on the Discord Developer Portal for your application, under `Application ID`
 
 1. Copy `.env.TEMPLATE` to `.env` and fill out the fields
 2. Build the [Docker](https://www.docker.com/) application with `make build`
@@ -27,16 +27,35 @@ https://discord.com/api/oauth2/authorize?client_id=<CLIENT_ID>&permissions=40671
 ## Contributing
 [Rust](https://www.rust-lang.org/) must be installed to run locally. See [serenity-rs/serenity](https://github.com/serenity-rs/serenity) for more information about Rust Discord API Library. 
 
-Furthermore, the following packages must be installed for [serenity-rs/songbird](https://github.com/serenity-rs/songbird). View the repository for additional installation and setup information on other operating systems.
-```
-sudo apt install libopus-dev
-sudo apt install ffmpeg
-sudo apt apt install youtube-dl
-```
+The following packages must be installed for [serenity-rs/songbird](https://github.com/serenity-rs/songbird). View the repository for additional installation and setup information on other operating systems.
+<details>
+  <summary>Unix Installation</summary>
 
-Potentially requires [yt-dlp](https://github.com/yt-dlp/yt-dlp#installation) and [yt-dlp FFmpeg Static Auto-Builds](https://github.com/yt-dlp/FFmpeg-Builds).
+  ```
+  sudo apt install libopus-dev
+  sudo apt install ffmpeg
+  sudo apt apt install youtube-dl
+  # PostgreSQL Headers
+  sudo apt install libpq5
+  sudo apt install libpq-dev
+  ```
 
-Begin the application with `cargo run`
+  - Potentially requires [yt-dlp](https://github.com/yt-dlp/yt-dlp#installation) and [yt-dlp FFmpeg Static Auto-Builds](https://github.com/yt-dlp/FFmpeg-Builds).
+
+</details>
+<details>
+  <summary>Mac Installation</summary>
+
+  ```
+  brew install opus
+  brew install ffmpeg
+  brew install youtube-dl
+  brew install postgresql
+  ```
+</details>
+
+Begin the application with `cargo run` (note the database must still be running)
+  - `docker compose up -d db`
 
 The application can also be tested from within a Docker container:
 ```

docker-compose.yml

@@ -9,10 +9,11 @@ services:
       dockerfile: ./Dockerfile
       args:
         - VERSION=${SIREN_VERSION}
-    #volumes:
-    #  - ./app:/siren
+    volumes:
+     - ./app:/siren
     environment:
       DISCORD_TOKEN: ${DISCORD_TOKEN}
+      RUST_LOG: ${RUST_LOG}
       DATABASE_URL: postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@db/${POSTGRES_DB}
     depends_on:
       - db
@@ -24,8 +25,8 @@ services:
       POSTGRES_USER: ${POSTGRES_USER}
       POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
       POSTGRES_DB: ${POSTGRES_DB}
-    #volumes:
-    #  - ./data:/var/lib/postgresql/data
+    volumes:
+     - ./data:/var/lib/postgresql/data
     ports:
       - "5432:5432"
     restart: unless-stopped

migrations/create_messages/down.sql

@@ -0,0 +1 @@
+DROP TABLE messages

migrations/create_messages/up.sql

@@ -0,0 +1,12 @@
+CREATE TABLE IF NOT EXISTS messages (
+  id TEXT PRIMARY KEY NOT NULL,
+  guild_id BIGINT NOT NULL,
+  channel_id BIGINT NOT NULL,
+  user_id BIGINT NOT NULL,
+  created BIGINT NOT NULL,
+  model TEXT NOT NULL,
+  request TEXT NOT NULL,
+  response TEXT NOT NULL,
+  request_tags TEXT[] NOT NULL,
+  response_tags TEXT[] NOT NULL
+)

src/commands/mod.rs

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

src/commands/oai.rs

@@ -0,0 +1,242 @@
+
+use std::error::Error;
+use std::fmt;
+
+use diesel::{prelude::*, PgConnection, insert_into};
+use diesel::r2d2::{Pool, ConnectionManager};
+use log::{error, debug, trace};
+
+use serde::{Serialize, Deserialize};
+use serde_json::Value;
+use serenity::model::channel::Message;
+use serenity::prelude::*;
+
+use crate::database::models::NewMessageDB;
+
+pub struct OAI {
+  pub client: reqwest::Client,
+  pub base_url: String,
+  pub max_attempts: i64,
+  pub token: String,
+  pub max_context_questions: i64
+}
+
+#[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<i64>,
+  /// 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: i64,
+  model: String,
+  usage: Usage,
+  choices: Vec<Choice>
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct Usage {
+  prompt_tokens: i64,
+  completion_tokens: i64,
+  total_tokens: i64
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct Choice {
+  message: ChatCompletionMessage,
+  finish_reason: String,
+  index: i64
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct ResponseError {
+  code: Option<i64>,
+  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, body);
+
+    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)
+        })
+      };
+
+    trace!("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, pool: &Pool<ConnectionManager<PgConnection>>) {
+  let bot_mention: String = 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 instructions = vec![
+    ChatCompletionMessage {
+      role: "system".to_string(),
+      content: "You are a Discord user named Siren.".to_string()
+    },
+    ChatCompletionMessage {
+      role: "system".to_string(),
+      content: "Siren is an expert on Dungeons and Dragons.".to_string()
+    },
+  ];
+
+  let user_request = ChatCompletionMessage {
+    role: "user".to_string(),
+    content: parsed_content.to_string()
+  };
+
+  let mut messages: Vec<ChatCompletionMessage> = vec![];
+  messages.extend(instructions);
+  // TODO: Get previous messages
+  messages.push(user_request);
+
+  let model = "gpt-3.5-turbo".to_string();
+
+  let request = ChatCompletionRequest {
+    model: model.to_string(),
+    messages,
+    temperature: Some(0.5),
+    top_p: None,
+    n: None,
+    max_tokens: Some(1000),
+    presence_penalty: Some(0.6),
+    frequency_penalty: Some(0.0),
+    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() {
+        let mut connection = pool.get().unwrap();
+        let res = r.choices[0].message.content.clone();
+        if let Err(err) = insert_into(crate::database::schema::messages::table).values(NewMessageDB {
+          id: &r.id,
+          guild_id: msg.guild_id.unwrap().0 as i64,
+          channel_id: msg.channel_id.0 as i64,
+          user_id: msg.author.id.0 as i64,
+          created: r.created,
+          model: &model,
+          request: &parsed_content,
+          response: &res,
+          request_tags: vec![],
+          response_tags: vec![],
+        }).execute(&mut connection) {
+          error!("Could not insert message into database: {}", err);
+        }
+        res
+      } else {
+        "No reply received".to_string()
+      }
+    }
+    Err(err) => {
+      error!("Could not get response from OpenAI: {}", err.message);
+      err.message
+    }
+  };
+  debug!("Sending response: \"{}\"", response);
+
+  if let Err(why) = msg.channel_id.say(&ctx.http, response).await {
+    error!("Cannot send message: {}", why);
+  }
+}

src/database/mod.rs

@@ -0,0 +1,44 @@
+use std::env;
+use std::path::Path;
+
+use diesel::RunQueryDsl;
+use diesel::r2d2::{Pool, ConnectionManager};
+use diesel::pg::PgConnection;
+use log::{error, info};
+
+pub mod models;
+pub mod schema;
+
+pub fn run_migrations(pool: &Pool<ConnectionManager<PgConnection>>) {
+  let mut connection = pool.get().unwrap();
+  let migrations_dir = Path::new("./migrations");
+  let migrations = std::fs::read_dir(&migrations_dir).unwrap();
+  
+  for migration in migrations {
+    if migration.as_ref().unwrap().file_type().unwrap().is_dir() {
+      let migration_paths = std::fs::read_dir(&migration.unwrap().path()).unwrap();
+
+      for migration_path in migration_paths {
+        if migration_path.as_ref().unwrap().file_name().eq_ignore_ascii_case("up.sql") {
+          let path = &migration_path.unwrap().path();
+          let contents = std::fs::read_to_string(path).expect("Unable to read from file");
+          if let Err(err) = diesel::sql_query(&contents).execute(&mut connection) {
+            error!("Could not run migration: {}", err);
+          } else {
+            info!("Successfully ran migration: {}", path.display());
+          }
+        }
+      }
+    }
+  }
+}
+
+pub fn establish_connection() -> Pool<ConnectionManager<PgConnection>> {
+  let database_user = env::var("POSTGRES_USER").expect("Expected a user in the environment");
+  let database_password = env::var("POSTGRES_PASSWORD").expect("Expected a password in the environment");
+  let database_name = env::var("POSTGRES_DB").expect("Expected a database name in the environment");
+
+  let database_url = format!("postgres://{}:{}@localhost/{}", database_user, database_password, database_name);
+  let manager = ConnectionManager::<PgConnection>::new(database_url);
+  Pool::builder().build(manager).expect("Failed to create pool.")
+}

src/database/models.rs

@@ -0,0 +1,33 @@
+use diesel::prelude::*;
+
+use super::schema::messages;
+
+#[derive(Queryable, Selectable)]
+#[diesel(table_name = messages)]
+pub struct MessageDB {
+  pub id: String,
+  pub guild_id: i64,
+  pub channel_id: i64,
+  pub user_id: i64,
+  pub created: i64,
+  pub model: String,
+  pub request: String,
+  pub response: String,
+  pub request_tags: Vec<String>,
+  pub response_tags: Vec<String>,
+}
+
+#[derive(Insertable)]
+#[diesel(table_name = messages)]
+pub struct NewMessageDB<'a> {
+  pub id: &'a str,
+  pub guild_id: i64,
+  pub channel_id: i64,
+  pub user_id: i64,
+  pub created: i64,
+  pub model: &'a str,
+  pub request: &'a str,
+  pub response: &'a str,
+  pub request_tags: Vec<&'a str>,
+  pub response_tags: Vec<&'a str>,
+}

src/database/schema.rs

@@ -0,0 +1,14 @@
+diesel::table! {
+  messages (id) {
+    id -> Text,
+    guild_id -> BigInt,
+    channel_id -> BigInt,
+    user_id -> BigInt,
+    created -> BigInt,
+    model -> Text,
+    request -> Text,
+    response -> Text,
+    request_tags -> Array<Text>,
+    response_tags -> Array<Text>,
+  }
+}

src/main.rs

@@ -2,21 +2,51 @@ use std::collections::HashSet;
 use std::env;
 
 use commands::audio::create_response;
+use diesel::r2d2::{Pool, ConnectionManager};
+use diesel::pg::PgConnection;
+
 use dotenv::dotenv;
 use log::{error, warn, info};
 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;
+mod database;
+struct Handler {
+  // Open AI Config
+  oai: Option<commands::oai::OAI>,
+  pool: Pool<ConnectionManager<PgConnection>>
+}
 
 #[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, &self.pool).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 +93,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();
@@ -85,15 +115,27 @@ async fn main() {
     Err(why) => panic!("Could not access application info: {:?}", why)
   };
 
-  let framework = StandardFramework::new()
-    .configure(|c| c
-      .owners(owners)
-      .prefix("!")
-    );
+  let pool = database::establish_connection();
+  database::run_migrations(&pool);
+  
+  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 , max_context_questions: 10 }),
+        pool
+      }
+    }
+    Err(err) => {
+      warn!("Could not load OpenAI token: {}", err);
+      Handler { oai: None, pool }
+    }
+  };
 
   let mut client = Client::builder(token, intents)
-    .event_handler(Handler)
-    .framework(framework)
+    .event_handler(handler)
+    .framework(StandardFramework::new()
+    .configure(|c| c.owners(owners)))
     .register_songbird()
     .await
     .expect("Error creating client");