Store messages into table, required docker container to create table

This commit is contained in:
Benjamin Sherriff
2023-07-06 22:44:48 -04:00
parent 9a8587e4b8
commit 4a366d6237
11 changed files with 115 additions and 51 deletions

View File

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

View File

@@ -1,6 +1,6 @@
[package] [package]
name = "siren" name = "siren"
version = "0.2.1" version = "0.2.2"
edition = "2021" edition = "2021"
authors = ["Ben Sherriff <hello@bensherriff.com>"] authors = ["Ben Sherriff <hello@bensherriff.com>"]
repository = "https://github.com/bensherriff/siren" repository = "https://github.com/bensherriff/siren"

View File

@@ -17,5 +17,5 @@ FROM debian:bullseye-slim as runtime
WORKDIR /siren WORKDIR /siren
COPY --from=builder /siren/target/release/siren /usr/local/bin/siren COPY --from=builder /siren/target/release/siren /usr/local/bin/siren
COPY --from=packages /packages /usr/bin COPY --from=packages /packages /usr/bin
ADD migrations ./migrations/ # ADD migrations ./migrations/
CMD ["siren"] CMD ["siren"]

View File

@@ -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. [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. 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>
sudo apt install libopus-dev <summary>Unix Installation</summary>
sudo apt install ffmpeg
sudo apt apt install youtube-dl ```
``` 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). - 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](). </details>
``` <details>
sudo apt install libpq5 <summary>Mac Installation</summary>
sudo apt install libpq-dev
``` ```
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) Begin the application with `cargo run` (note the database must still be running)
- `docker compose up -d db` - `docker compose up -d db`

View File

@@ -27,6 +27,7 @@ services:
POSTGRES_DB: ${POSTGRES_DB} POSTGRES_DB: ${POSTGRES_DB}
volumes: volumes:
- ./data:/var/lib/postgresql/data - ./data:/var/lib/postgresql/data
- ./sql/create_tables.sql:/docker-entrypoint-initdb.d/create_tables.sql
ports: ports:
- "5432:5432" - "5432:5432"
restart: unless-stopped restart: unless-stopped

12
sql/create_tables.sql Normal file
View File

@@ -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
);

View File

@@ -2,6 +2,8 @@
use std::error::Error; use std::error::Error;
use std::fmt; use std::fmt;
use diesel::{prelude::*, PgConnection, insert_into};
use diesel::r2d2::{Pool, ConnectionManager};
use log::{error, debug, trace}; use log::{error, debug, trace};
use serde::{Serialize, Deserialize}; use serde::{Serialize, Deserialize};
@@ -9,10 +11,12 @@ use serde_json::Value;
use serenity::model::channel::Message; use serenity::model::channel::Message;
use serenity::prelude::*; use serenity::prelude::*;
use crate::database::models::NewMessageDB;
pub struct OAI { pub struct OAI {
pub client: reqwest::Client, pub client: reqwest::Client,
pub base_url: String, pub base_url: String,
pub max_attempts: u64, pub max_attempts: i64,
pub token: String pub token: String
} }
@@ -29,7 +33,7 @@ struct ChatCompletionRequest {
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
n: Option<f64>, n: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
max_tokens: Option<u64>, max_tokens: Option<i64>,
/// Value between -2.0 and 2.0 /// Value between -2.0 and 2.0
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
presence_penalty: Option<f64>, presence_penalty: Option<f64>,
@@ -62,7 +66,7 @@ enum Role {
struct ChatCompletionResponse { struct ChatCompletionResponse {
id: String, id: String,
object: String, object: String,
created: u64, created: i64,
model: String, model: String,
usage: Usage, usage: Usage,
choices: Vec<Choice> choices: Vec<Choice>
@@ -70,21 +74,21 @@ struct ChatCompletionResponse {
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
struct Usage { struct Usage {
prompt_tokens: u64, prompt_tokens: i64,
completion_tokens: u64, completion_tokens: i64,
total_tokens: u64 total_tokens: i64
} }
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
struct Choice { struct Choice {
message: ChatCompletionMessage, message: ChatCompletionMessage,
finish_reason: String, finish_reason: String,
index: u64 index: i64
} }
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
struct ResponseError { struct ResponseError {
code: Option<u64>, code: Option<i64>,
message: Option<String>, message: Option<String>,
param: Option<String>, param: Option<String>,
#[serde(rename = "type")] #[serde(rename = "type")]
@@ -114,7 +118,7 @@ impl OAI {
async fn get_request(&self, request: ChatCompletionRequest) -> Result<ChatCompletionResponse, OAIError> { async fn get_request(&self, request: ChatCompletionRequest) -> Result<ChatCompletionResponse, OAIError> {
let uri = format!("{}/chat/completions", self.base_url); let uri = format!("{}/chat/completions", self.base_url);
let body = serde_json::to_string(&request).unwrap(); 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 let value = match match self.client
.post(&uri) .post(&uri)
.bearer_auth(&self.token) .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::<OAIResponseEvent>(value) { // let response = match serde_json::from_value::<OAIResponseEvent>(value) {
// Ok(r) => { // Ok(r) => {
@@ -159,8 +163,8 @@ impl OAI {
} }
} }
pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) { pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI, pool: &Pool<ConnectionManager<PgConnection>>) {
let bot_mention = format!("<@{}>", ctx.cache.current_user_id().0); let bot_mention: String = format!("<@{}>", ctx.cache.current_user_id().0);
let parsed_content = msg.content.replace(bot_mention.as_str(), ""); let parsed_content = msg.content.replace(bot_mention.as_str(), "");
debug!("Generating response for message: {}", msg.content); debug!("Generating response for message: {}", msg.content);
@@ -171,12 +175,14 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
}, },
ChatCompletionMessage { ChatCompletionMessage {
role: "user".to_string(), role: "user".to_string(),
content: parsed_content content: parsed_content.to_string()
}, },
]; ];
let model = "gpt-3.5-turbo".to_string();
let request = ChatCompletionRequest { let request = ChatCompletionRequest {
model: "gpt-3.5-turbo".to_string(), model: model.to_string(),
messages, messages,
temperature: None, temperature: None,
top_p: None, top_p: None,
@@ -190,7 +196,23 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
Ok(r) => { Ok(r) => {
debug!("Received response from OpenAI"); debug!("Received response from OpenAI");
if !r.choices.is_empty() { 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 { } else {
"No reply received".to_string() "No reply received".to_string()
} }
@@ -200,6 +222,8 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
err.message err.message
} }
}; };
debug!("Sending response: \"{}\"", response);
if let Err(why) = msg.channel_id.say(&ctx.http, response).await { if let Err(why) = msg.channel_id.say(&ctx.http, response).await {
error!("Cannot send message: {}", why); error!("Cannot send message: {}", why);
} }

View File

@@ -1,5 +1,5 @@
use std::env; use std::env;
use std::path::Path; // use std::path::Path;
use diesel::r2d2::{Pool, ConnectionManager}; use diesel::r2d2::{Pool, ConnectionManager};
use diesel::pg::PgConnection; use diesel::pg::PgConnection;
@@ -7,25 +7,39 @@ use diesel::pg::PgConnection;
pub mod models; pub mod models;
pub mod schema; pub mod schema;
pub fn run_migrations(pool: &Pool<ConnectionManager<PgConnection>>) { // pub fn run_migrations(pool: &Pool<ConnectionManager<PgConnection>>) {
let mut connection = pool.get().unwrap(); // let mut connection = pool.get().unwrap();
let migrations_dir = Path::new("./migrations"); // let migrations_dir = Path::new("./migrations");
let migrations = std::fs::read_dir(&migrations_dir).unwrap(); // let migrations = std::fs::read_dir(&migrations_dir).unwrap();
for migration in migrations { // for migration in migrations {
if migration.as_ref().unwrap().file_type().unwrap().is_dir() { // if migration.as_ref().unwrap().file_type().unwrap().is_dir() {
let migration_paths = std::fs::read_dir(&migration.unwrap().path()).unwrap(); // let migration_paths = std::fs::read_dir(&migration.unwrap().path()).unwrap();
for migration_path in migration_paths { // for migration_path in migration_paths {
if migration_path.as_ref().unwrap().file_name().eq_ignore_ascii_case("up.sql") { // if migration_path.as_ref().unwrap().file_name().eq_ignore_ascii_case("up.sql") {
let path = &migration_path.unwrap().path(); // let path = &migration_path.unwrap().path();
let contents = std::fs::read_to_string(path).expect("Unable to read from file"); // let contents = std::fs::read_to_string(path).expect("Unable to read from file");
// connection.build_transaction() // connection.build_transaction()
} // connection.build_transaction()
} // .read_write()
} // .run(|conn| {
} // // let read_attempt = users.select(name).load::<String>(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<ConnectionManager<PgConnection>> { pub fn establish_connection() -> Pool<ConnectionManager<PgConnection>> {
let database_user = env::var("POSTGRES_USER").expect("Expected a user in the environment"); let database_user = env::var("POSTGRES_USER").expect("Expected a user in the environment");

View File

@@ -5,7 +5,7 @@ use super::schema::messages;
#[derive(Queryable, Selectable)] #[derive(Queryable, Selectable)]
#[diesel(table_name = messages)] #[diesel(table_name = messages)]
pub struct MessageDB { pub struct MessageDB {
pub id: i64, pub id: String,
pub guild_id: i64, pub guild_id: i64,
pub channel_id: i64, pub channel_id: i64,
pub user_id: i64, pub user_id: i64,
@@ -20,7 +20,7 @@ pub struct MessageDB {
#[derive(Insertable)] #[derive(Insertable)]
#[diesel(table_name = messages)] #[diesel(table_name = messages)]
pub struct NewMessageDB<'a> { pub struct NewMessageDB<'a> {
pub id: i64, pub id: &'a str,
pub guild_id: i64, pub guild_id: i64,
pub channel_id: i64, pub channel_id: i64,
pub user_id: i64, pub user_id: i64,

View File

@@ -1,6 +1,6 @@
diesel::table! { diesel::table! {
messages (id) { messages (id) {
id -> BigInt, id -> Text,
guild_id -> BigInt, guild_id -> BigInt,
channel_id -> BigInt, channel_id -> BigInt,
user_id -> BigInt, user_id -> BigInt,

View File

@@ -36,7 +36,7 @@ impl EventHandler for Handler {
match msg.mentions_me(&ctx.http).await { match msg.mentions_me(&ctx.http).await {
Ok(mentioned) => { Ok(mentioned) => {
if 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) Err(why) => warn!("Could not check mentions: {:?}", why)
@@ -116,7 +116,7 @@ async fn main() {
}; };
let pool = database::establish_connection(); let pool = database::establish_connection();
database::run_migrations(&pool); // database::run_migrations(&pool);
let handler = match env::var("OPENAI_API_KEY") { let handler = match env::var("OPENAI_API_KEY") {
Ok(token) => { Ok(token) => {