From 4a366d6237b63c50ba6325f78780f6f9f7d8515b Mon Sep 17 00:00:00 2001 From: Benjamin Sherriff Date: Thu, 6 Jul 2023 22:44:48 -0400 Subject: [PATCH] Store messages into table, required docker container to create table --- .version | 2 +- Cargo.toml | 2 +- Dockerfile | 2 +- README.md | 33 ++++++++++++++++++-------- docker-compose.yml | 1 + sql/create_tables.sql | 12 ++++++++++ src/commands/oai.rs | 54 ++++++++++++++++++++++++++++++------------ src/database/mod.rs | 50 ++++++++++++++++++++++++-------------- src/database/models.rs | 4 ++-- src/database/schema.rs | 2 +- src/main.rs | 4 ++-- 11 files changed, 115 insertions(+), 51 deletions(-) create mode 100644 sql/create_tables.sql diff --git a/.version b/.version index 6fcdfad..cdca5d4 100644 --- a/.version +++ b/.version @@ -1 +1 @@ -SIREN_VERSION=0.2.1 \ No newline at end of file +SIREN_VERSION=0.2.2 \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index 58dd4bb..bd3ba7e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "siren" -version = "0.2.1" +version = "0.2.2" edition = "2021" authors = ["Ben Sherriff "] repository = "https://github.com/bensherriff/siren" diff --git a/Dockerfile b/Dockerfile index 3a5db82..3f9b54d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -17,5 +17,5 @@ FROM debian:bullseye-slim as runtime WORKDIR /siren COPY --from=builder /siren/target/release/siren /usr/local/bin/siren COPY --from=packages /packages /usr/bin -ADD migrations ./migrations/ +# ADD migrations ./migrations/ CMD ["siren"] diff --git a/README.md b/README.md index 1c4f331..10d4b3b 100644 --- a/README.md +++ b/README.md @@ -28,18 +28,31 @@ The CLIENT_ID can be found in the General Information tab on the Discord Develop [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. 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 -``` +
+ Unix Installation + + ``` + 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 + ``` + - 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). -The following packages must be installed for [PostgreSQL](). -``` -sudo apt install libpq5 -sudo apt install libpq-dev -``` +
+
+ Mac Installation + + ``` + brew install opus + brew install ffmpeg + brew install youtube-dl + brew install postgresql + ``` +
Begin the application with `cargo run` (note the database must still be running) - `docker compose up -d db` diff --git a/docker-compose.yml b/docker-compose.yml index fdfc898..1cfd21a 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -27,6 +27,7 @@ services: POSTGRES_DB: ${POSTGRES_DB} volumes: - ./data:/var/lib/postgresql/data + - ./sql/create_tables.sql:/docker-entrypoint-initdb.d/create_tables.sql ports: - "5432:5432" restart: unless-stopped diff --git a/sql/create_tables.sql b/sql/create_tables.sql new file mode 100644 index 0000000..f876ba8 --- /dev/null +++ b/sql/create_tables.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 +); \ No newline at end of file diff --git a/src/commands/oai.rs b/src/commands/oai.rs index 0980300..366a9c3 100644 --- a/src/commands/oai.rs +++ b/src/commands/oai.rs @@ -2,6 +2,8 @@ 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}; @@ -9,10 +11,12 @@ 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: u64, + pub max_attempts: i64, pub token: String } @@ -29,7 +33,7 @@ struct ChatCompletionRequest { #[serde(skip_serializing_if = "Option::is_none")] n: Option, #[serde(skip_serializing_if = "Option::is_none")] - max_tokens: Option, + max_tokens: Option, /// Value between -2.0 and 2.0 #[serde(skip_serializing_if = "Option::is_none")] presence_penalty: Option, @@ -62,7 +66,7 @@ enum Role { struct ChatCompletionResponse { id: String, object: String, - created: u64, + created: i64, model: String, usage: Usage, choices: Vec @@ -70,21 +74,21 @@ struct ChatCompletionResponse { #[derive(Debug, Clone, Serialize, Deserialize)] struct Usage { - prompt_tokens: u64, - completion_tokens: u64, - total_tokens: u64 + prompt_tokens: i64, + completion_tokens: i64, + total_tokens: i64 } #[derive(Debug, Clone, Serialize, Deserialize)] struct Choice { message: ChatCompletionMessage, finish_reason: String, - index: u64 + index: i64 } #[derive(Debug, Clone, Serialize, Deserialize)] struct ResponseError { - code: Option, + code: Option, message: Option, param: Option, #[serde(rename = "type")] @@ -114,7 +118,7 @@ impl OAI { async fn get_request(&self, request: ChatCompletionRequest) -> Result { let uri = format!("{}/chat/completions", self.base_url); let body = serde_json::to_string(&request).unwrap(); - trace!("Sending request to {}", uri); + trace!("Sending request to {}: {}", uri, body); let value = match match self.client .post(&uri) .bearer_auth(&self.token) @@ -135,7 +139,7 @@ impl OAI { }) }; - debug!("Received response from OpenAI: {:?}", value); + trace!("Received response from OpenAI: {:?}", value); // let response = match serde_json::from_value::(value) { // Ok(r) => { @@ -159,8 +163,8 @@ impl OAI { } } -pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) { - let bot_mention = format!("<@{}>", ctx.cache.current_user_id().0); +pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI, pool: &Pool>) { + 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); @@ -171,12 +175,14 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) { }, ChatCompletionMessage { role: "user".to_string(), - content: parsed_content + content: parsed_content.to_string() }, ]; + let model = "gpt-3.5-turbo".to_string(); + let request = ChatCompletionRequest { - model: "gpt-3.5-turbo".to_string(), + model: model.to_string(), messages, temperature: None, top_p: None, @@ -190,7 +196,23 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) { Ok(r) => { debug!("Received response from OpenAI"); if !r.choices.is_empty() { - r.choices[0].message.content.clone() + 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() } @@ -200,6 +222,8 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) { err.message } }; + debug!("Sending response: \"{}\"", response); + if let Err(why) = msg.channel_id.say(&ctx.http, response).await { error!("Cannot send message: {}", why); } diff --git a/src/database/mod.rs b/src/database/mod.rs index 1223358..2d5019b 100644 --- a/src/database/mod.rs +++ b/src/database/mod.rs @@ -1,5 +1,5 @@ use std::env; -use std::path::Path; +// use std::path::Path; use diesel::r2d2::{Pool, ConnectionManager}; use diesel::pg::PgConnection; @@ -7,25 +7,39 @@ use diesel::pg::PgConnection; pub mod models; pub mod schema; -pub fn run_migrations(pool: &Pool>) { - let mut connection = pool.get().unwrap(); - let migrations_dir = Path::new("./migrations"); - let migrations = std::fs::read_dir(&migrations_dir).unwrap(); +// pub fn run_migrations(pool: &Pool>) { +// 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 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"); - // connection.build_transaction() - } - } - } - } -} +// 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"); +// connection.build_transaction() +// connection.build_transaction() +// .read_write() +// .run(|conn| { +// // let read_attempt = users.select(name).load::(conn); +// // assert!(read_attempt.is_ok()); + +// // let write_attempt = diesel::insert_into(users) +// // .values(name.eq("Ruby")) +// // .execute(conn); +// // assert!(write_attempt.is_ok()); +// // diesel::migration::CREATE_MIGRATIONS_TABLE + +// // Ok(()) +// }); +// } +// } +// } +// } +// } pub fn establish_connection() -> Pool> { let database_user = env::var("POSTGRES_USER").expect("Expected a user in the environment"); diff --git a/src/database/models.rs b/src/database/models.rs index 56947bb..93d9bab 100644 --- a/src/database/models.rs +++ b/src/database/models.rs @@ -5,7 +5,7 @@ use super::schema::messages; #[derive(Queryable, Selectable)] #[diesel(table_name = messages)] pub struct MessageDB { - pub id: i64, + pub id: String, pub guild_id: i64, pub channel_id: i64, pub user_id: i64, @@ -20,7 +20,7 @@ pub struct MessageDB { #[derive(Insertable)] #[diesel(table_name = messages)] pub struct NewMessageDB<'a> { - pub id: i64, + pub id: &'a str, pub guild_id: i64, pub channel_id: i64, pub user_id: i64, diff --git a/src/database/schema.rs b/src/database/schema.rs index 47676b5..825d574 100644 --- a/src/database/schema.rs +++ b/src/database/schema.rs @@ -1,6 +1,6 @@ diesel::table! { messages (id) { - id -> BigInt, + id -> Text, guild_id -> BigInt, channel_id -> BigInt, user_id -> BigInt, diff --git a/src/main.rs b/src/main.rs index e3a23d6..da98df9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -36,7 +36,7 @@ impl EventHandler for Handler { match msg.mentions_me(&ctx.http).await { Ok(mentioned) => { if mentioned { - commands::oai::generate_response(&ctx, &msg, oai).await; + commands::oai::generate_response(&ctx, &msg, oai, &self.pool).await; } } Err(why) => warn!("Could not check mentions: {:?}", why) @@ -116,7 +116,7 @@ async fn main() { }; let pool = database::establish_connection(); - database::run_migrations(&pool); + // database::run_migrations(&pool); let handler = match env::var("OPENAI_API_KEY") { Ok(token) => {