bsherriff/siren
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Updated dependencies, fixed migrations usage

Browse Source
This commit is contained in:
Benjamin Sherriff
2023-10-03 10:05:14 -04:00
parent e80ad9680a
commit 6c8a7ceefc
16 changed files with 141 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
.env.TEMPLATE
View File

@@ -1,6 +1,8 @@
DISCORD_TOKEN=
RUST_LOG=warn,siren=info
POSTGRES_USER=siren
POSTGRES_PASSWORD=
POSTGRES_DB=siren
DATABASE_USER=siren
DATABASE_PASSWORD=
DATABASE_NAME=siren
DATABASE_HOST=localhost
DATABASE_PORT=5432
DISCORD_TOKEN=
OPENAI_API_KEY=

2
.version
View File

@@ -1 +1 @@
SIREN_VERSION=0.2.3
SIREN_VERSION=0.2.4

19
Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "siren"
version = "0.2.3"
version = "0.2.4"
edition = "2021"
authors = ["Ben Sherriff <hello@bensherriff.com>"]
repository = "https://github.com/bensherriff/siren"
@@ -9,9 +9,12 @@ license = "GPL-3.0-or-later"
[dependencies]
dotenv = "0.15.0"
serde_json = "1.0"
log = "0.4.19"
serde_json = "1.0.107"
log = "0.4.20"
env_logger = "0.10.0"
diesel_migrations = { version = "2.1.0", features = ["postgres"] }
r2d2 = "0.8.10"
lazy_static = "1.4.0"
[dependencies.serenity]
version = "0.11.6"
@@ -23,23 +26,23 @@ version = "0.3.2"
features = ["builtin-queue", "yt-dlp"]
[dependencies.tokio]
version = "1.29.1"
version = "1.32.0"
features = ["macros", "rt-multi-thread"]
[dependencies.serde]
version = "1.0"
version = "1.0.188"
features = ["derive"]
[dependencies.reqwest]
version = "0.11.18"
version = "0.11.22"
default-features = false
features = ["json", "rustls-tls"]
[dependencies.diesel]
version = "2.1.0"
version = "2.1.2"
default-features = false
features = ["postgres", "32-column-tables", "serde_json", "r2d2", "with-deprecated"]
[dependencies.pyo3]
version = "0.19.1"
version = "0.19.2"
features = ["auto-initialize"]

5
Dockerfile
View File

@@ -1,10 +1,12 @@
FROM rust:1.70.0 as builder
# Builder
FROM rust:1.72.1-bookworm as builder
WORKDIR /siren
ADD src ./src/
ADD Cargo.toml ./
RUN apt-get update && apt-get install -y cmake && \
cargo build --release --bin siren
# Packages
FROM debian:bullseye-slim as packages
WORKDIR /packages
RUN apt-get update && apt-get install -y curl tar xz-utils && \
@@ -19,6 +21,7 @@ RUN apt-get update && apt-get install -y curl tar xz-utils && \
# curl -L https://download.pytorch.org/libtorch/cu117/libtorch-cxx11-abi-shared-with-deps-2.0.1%2Bcu117.zip > libtorch.zip && \
# unzip libtorch.zip && rm libtorch.zip
# Runner
FROM debian:bullseye-slim as runtime
WORKDIR /siren
RUN apt-get update && apt-get install -y libopus-dev libpq5 libpq-dev && apt-get auto-remove -y

30
docker-compose.yml
View File

@@ -3,38 +3,36 @@ version: '3'
services:
siren:
image: siren:${SIREN_VERSION:-latest}
container_name: siren
container_name: siren-service
build:
context: .
dockerfile: ./Dockerfile
args:
- VERSION=${SIREN_VERSION:-latest}
volumes:
- ./app:/siren
env_file:
- .env
environment:
DISCORD_TOKEN: ${DISCORD_TOKEN}
RUST_LOG: ${RUST_LOG}
OPENAI_API_KEY: ${OPENAI_API_KEY}
POSTGRES_USER: ${POSTGRES_USER}
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
POSTGRES_DB: ${POSTGRES_DB}
POSTGRES_HOST: db
DATABASE_HOST: db
DATABASE_PORT: 5432
depends_on:
- db
restart: unless-stopped
db:
image: postgres:latest
container_name: siren_db
container_name: siren-db
env_file:
- .env
environment:
POSTGRES_USER: ${POSTGRES_USER}
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
POSTGRES_DB: ${POSTGRES_DB}
POSTGRES_USER: ${DATABASE_USER}
POSTGRES_PASSWORD: ${DATABASE_PASSWORD}
POSTGRES_DB: ${DATABASE_NAME}
volumes:
- ./data:/var/lib/postgresql/data
- db:/var/lib/postgresql/data
- db_logs:/var/log
ports:
- "5432:5432"
- ${DATABASE_PORT:-5432}:5432
restart: unless-stopped
volumes:
db:
db_logs:

1
migrations/000000_create_messages/down.sql Normal file
View File

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

2
migrations/create_messages/up.sql → migrations/000000_create_messages/up.sql
View File

@@ -9,4 +9,4 @@ CREATE TABLE IF NOT EXISTS messages (
response TEXT NOT NULL,
request_tags TEXT[] NOT NULL,
response_tags TEXT[] NOT NULL
)
);

1
migrations/create_messages/down.sql
View File

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

22
src/commands/oai.rs
View File

@@ -2,8 +2,7 @@
use std::error::Error;
use std::fmt;
use diesel::{prelude::*, PgConnection, insert_into};
use diesel::r2d2::{Pool, ConnectionManager};
use diesel::{prelude::*, insert_into};
use log::{error, debug, trace, warn};
use serde::{Serialize, Deserialize};
@@ -13,7 +12,8 @@ use serenity::model::channel::Message;
use serenity::model::prelude::{ChannelType, PermissionOverwrite, PermissionOverwriteType};
use serenity::prelude::*;
use crate::database::models::{NewMessageDB, MessageDB};
use crate::db;
use crate::messages::{NewMessageDB, MessageDB};
pub struct OAI {
pub client: reqwest::Client,
@@ -187,26 +187,26 @@ impl OAI {
}
}
pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI, pool: &Pool<ConnectionManager<PgConnection>>) {
pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
debug!("Generating response for message: {}", msg.content);
let mut connection = db::connection().unwrap();
let guild_id = msg.guild_id.unwrap();
let channel_id = msg.channel_id;
let author_id = msg.author.id;
let mut connection = pool.get().unwrap();
// Parse out the bot mention from the message
let bot_mention: String = format!("<@{}>", ctx.cache.current_user_id().0);
let parsed_content = msg.content.replace(bot_mention.as_str(), "");
// Setup the request messages
let result: Result<Vec<MessageDB>, diesel::result::Error> = crate::database::schema::messages::table
let result: Result<Vec<MessageDB>, diesel::result::Error> = crate::schema::messages::table
.select(MessageDB::as_select())
.filter((crate::database::schema::messages::guild_id.eq(guild_id.0 as i64))
.and(crate::database::schema::messages::channel_id.eq(channel_id.0 as i64))
.and(crate::database::schema::messages::user_id.eq(author_id.0 as i64))
.filter((crate::schema::messages::guild_id.eq(guild_id.0 as i64))
.and(crate::schema::messages::channel_id.eq(channel_id.0 as i64))
.and(crate::schema::messages::user_id.eq(author_id.0 as i64))
)
.order(crate::database::schema::messages::created.asc())
.order(crate::schema::messages::created.asc())
.limit(oai.max_context_questions)
.load(&mut connection);
@@ -284,7 +284,7 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI, pool: &P
if !r.choices.is_empty() {
let res = r.choices[0].message.content.clone();
// Insert the message into the messages database table
if let Err(err) = insert_into(crate::database::schema::messages::table).values(NewMessageDB {
if let Err(err) = insert_into(crate::schema::messages::table).values(NewMessageDB {
id: &r.id,
guild_id: guild_id.0 as i64,
channel_id: response_channel.0 as i64,

61
src/database/mod.rs
View File

@@ -1,61 +0,0 @@
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();
if let Err(err) = diesel::sql_query("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
)").execute(&mut connection) {
error!("Could not create messages table: {}", err);
} else {
info!("Successfully created messages table");
}
// 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_host = env::var("POSTGRES_HOST").unwrap_or("localhost".to_string());
let database_url = format!("postgres://{}:{}@{}/{}", database_user, database_password, database_host, database_name);
let manager = ConnectionManager::<PgConnection>::new(database_url);
Pool::builder().build(manager).expect("Failed to create pool.")
}

39
src/db.rs Normal file
View File

@@ -0,0 +1,39 @@
use crate::error_handler::ServiceError;
use diesel::{r2d2::ConnectionManager, PgConnection};
use crate::diesel_migrations::MigrationHarness;
use lazy_static::lazy_static;
use log::{error, info};
use r2d2;
use std::env;
type Pool = r2d2::Pool<ConnectionManager<PgConnection>>;
pub type DbConnection = r2d2::PooledConnection<ConnectionManager<PgConnection>>;
pub const MIGRATIONS: diesel_migrations::EmbeddedMigrations = embed_migrations!();
lazy_static! {
static ref POOL: Pool = {
let username = env::var("DATABASE_USER").expect("DATABASE_USERNAME is not set");
let password = env::var("DATABASE_PASSWORD").expect("DATABASE_PASSWORD is not set");
let host = env::var("DATABASE_HOST").unwrap_or("localhost".to_string());
let name = env::var("DATABASE_NAME").expect("DATABASE_NAME is not set");
let port = env::var("DATABASE_PORT").unwrap_or("5432".to_string());
let url = format!("postgres://{}:{}@{}:{}/{}", username, password, host, port, name);
let manager = ConnectionManager::<PgConnection>::new(url);
Pool::builder().test_on_check_out(true).build(manager).expect("Failed to create db pool")
};
}
pub fn init() {
lazy_static::initialize(&POOL);
let mut pool: DbConnection = connection().expect("Failed to get db connection");
match pool.run_pending_migrations(MIGRATIONS) {
Ok(_) => info!("Database initialized"),
Err(err) => error!("Failed to initialize database; {}", err)
};
}
pub fn connection() -> Result<DbConnection, ServiceError> {
POOL.get()
.map_err(|e| ServiceError::new(500, format!("Failed getting db connection: {}", e)))
}

36
src/error_handler.rs Normal file
View File

@@ -0,0 +1,36 @@
use diesel::result::Error as DieselError;
use serde::{Deserialize, Serialize};
use std::fmt;
#[derive(Debug, Deserialize, Serialize)]
pub struct ServiceError {
pub error_status_code: u16,
pub error_message: String,
}
impl ServiceError {
pub fn new(error_status_code: u16, error_message: String) -> ServiceError {
ServiceError {
error_status_code,
error_message,
}
}
}
impl fmt::Display for ServiceError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str(self.error_message.as_str())
}
}
impl From<DieselError> for ServiceError {
fn from(error: DieselError) -> ServiceError {
match error {
DieselError::DatabaseError(_, err) => ServiceError::new(409, err.message().to_string()),
DieselError::NotFound => {
ServiceError::new(404, "The record was not found".to_string())
}
err => ServiceError::new(500, format!("Unknown Diesel error: {}", err)),
}
}
}

24
src/main.rs
View File

@@ -1,10 +1,12 @@
extern crate diesel;
#[macro_use]
extern crate diesel_migrations;
use std::collections::{HashSet, HashMap};
use std::env;
use std::sync::Arc;
use commands::audio::{create_response, AudioConfig, AudioConfigs};
use diesel::r2d2::{Pool, ConnectionManager};
use diesel::pg::PgConnection;
use dotenv::dotenv;
use log::{error, warn, info};
@@ -18,11 +20,13 @@ use serenity::prelude::*;
use songbird::SerenityInit;
mod commands;
mod database;
mod error_handler;
mod db;
mod messages;
mod schema;
struct Handler {
// Open AI Config
oai: Option<commands::oai::OAI>,
pool: Pool<ConnectionManager<PgConnection>>
oai: Option<commands::oai::OAI>
}
#[async_trait]
@@ -46,7 +50,7 @@ impl EventHandler for Handler {
Err(_) => false
};
if mentioned || bot_in_thread {
commands::oai::generate_response(&ctx, &msg, oai, &self.pool).await;
commands::oai::generate_response(&ctx, &msg, oai).await;
}
}
Err(why) => warn!("Could not check mentions: {:?}", why)
@@ -134,20 +138,18 @@ async fn main() {
Err(why) => panic!("Could not access application info: {:?}", why)
};
let pool = database::establish_connection();
database::run_migrations(&pool);
db::init();
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: 15 }),
pool
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: 15 })
}
}
Err(err) => {
warn!("Could not load OpenAI token: {}", err);
Handler { oai: None, pool }
Handler { oai: None }
}
};

3
src/messages/mod.rs Normal file
View File

@@ -0,0 +1,3 @@
mod model;
pub use model::*;

2
src/database/models.rs → src/messages/model.rs
View File

@@ -1,6 +1,6 @@
use diesel::prelude::*;
use super::schema::messages;
use crate::schema::messages;
#[derive(Queryable, Selectable)]
#[diesel(table_name = messages)]

0
src/database/schema.rs → src/schema.rs
View File