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

Store messages into table, required docker container to create table

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

54
src/commands/oai.rs
View File

@@ -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<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
max_tokens: Option<u64>,
max_tokens: Option<i64>,
/// Value between -2.0 and 2.0
#[serde(skip_serializing_if = "Option::is_none")]
presence_penalty: Option<f64>,
@@ -62,7 +66,7 @@ enum Role {
struct ChatCompletionResponse {
id: String,
object: String,
created: u64,
created: i64,
model: String,
usage: Usage,
choices: Vec<Choice>
@@ -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<u64>,
code: Option<i64>,
message: Option<String>,
param: Option<String>,
#[serde(rename = "type")]
@@ -114,7 +118,7 @@ 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);
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::<OAIResponseEvent>(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<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);
@@ -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);
}

50
src/database/mod.rs
View File

@@ -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<ConnectionManager<PgConnection>>) {
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<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 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::<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>> {
let database_user = env::var("POSTGRES_USER").expect("Expected a user in the environment");

4
src/database/models.rs
View File

@@ -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,

2
src/database/schema.rs
View File

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

4
src/main.rs
View File

@@ -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) => {