From 504b27f607b7bd71efbaa55f9bf772d61f904f53 Mon Sep 17 00:00:00 2001 From: Andi Pieper Date: Wed, 21 Feb 2024 12:16:01 +0100 Subject: [PATCH 1/6] do not record messages without a parent message. --- src/db/ai_help.rs | 15 ++++++++++++ tests/api/ai_help_history.rs | 45 ++++++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/src/db/ai_help.rs b/src/db/ai_help.rs index 448442fc..26e7ed15 100644 --- a/src/db/ai_help.rs +++ b/src/db/ai_help.rs @@ -140,6 +140,21 @@ pub fn add_help_history_message( conn: &mut PgConnection, mut message: AIHelpHistoryMessageInsert, ) -> Result { + // Check if the parent message exists in the database first. + // We could end up with a non-existing parent message if the + // user enables history after the conversation has already + // started. Do not try to store the message in that case. + if let Some(parent_id) = message.parent_id { + let parent_message_count = ai_help_history_messages::table + .filter(ai_help_history_messages::message_id.eq(&parent_id)) + .count() + .get_result::(conn)?; + if parent_message_count == 0 { + return Ok(Utc::now().naive_utc()); + } + } + + // Ok, we are good on the parent message. let updated_at = update(ai_help_history::table) .filter( ai_help_history::user_id diff --git a/tests/api/ai_help_history.rs b/tests/api/ai_help_history.rs index c00f3f25..93570b61 100644 --- a/tests/api/ai_help_history.rs +++ b/tests/api/ai_help_history.rs @@ -5,9 +5,12 @@ use actix_web::test; use anyhow::Error; use async_openai::types::ChatCompletionRequestMessage; use async_openai::types::Role::{Assistant, User}; +use diesel::prelude::*; +use diesel::ExpressionMethods; use rumba::ai::help::RefDoc; use rumba::db::ai_help::{add_help_history, add_help_history_message}; use rumba::db::model::{AIHelpHistoryMessageInsert, SettingsInsert}; +use rumba::db::schema::ai_help_history_messages; use rumba::db::settings::create_or_update_settings; use serde_json::Value::Null; use uuid::Uuid; @@ -95,6 +98,48 @@ async fn test_history() -> Result<(), Error> { Ok(()) } +#[actix_rt::test] +#[stubr::mock(port = 4321)] +async fn test_history_message_without_parent() -> Result<(), Error> { + let pool = reset()?; + let app = test_app_with_login(&pool).await.unwrap(); + let service = test::init_service(app).await; + let mut _logged_in_client = TestHttpClient::new(service).await; + + let mut conn = pool.get()?; + // Make sure that history is enabled + create_or_update_settings( + &mut conn, + SettingsInsert { + user_id: 1, + ai_help_history: Some(true), + ..Default::default() + }, + )?; + + // create a message with a non-existing parent + let message = AIHelpHistoryMessageInsert { + user_id: 1, + chat_id: CHAT_ID, + message_id: MESSAGE_ID, + parent_id: Some(Uuid::from_u128(2)), + created_at: None, + sources: None, + request: None, + response: None, + }; + let res = add_help_history_message(&mut conn, message); + // We return OK, but the message should not be recorded. + assert!(res.is_ok()); + let message_count = ai_help_history_messages::table + .filter(ai_help_history_messages::user_id.eq(1)) + .count() + .get_result::(&mut conn)?; + assert_eq!(message_count, 0); + drop_stubr(stubr).await; + Ok(()) +} + fn normalize_digits(s: &str) -> String { let mut result = String::new(); From 0446a253f36ae314ef8224211da6f7def81c42fa Mon Sep 17 00:00:00 2001 From: Andi Pieper Date: Wed, 21 Feb 2024 15:06:06 +0100 Subject: [PATCH 2/6] removed debug logs, clippy fixes --- src/api/ai_help.rs | 35 +++++++++++++++++++++------- src/db/ai_help.rs | 15 ------------ tests/api/ai_help_history.rs | 45 ------------------------------------ 3 files changed, 27 insertions(+), 68 deletions(-) diff --git a/src/api/ai_help.rs b/src/api/ai_help.rs index 71f46625..c858be2f 100644 --- a/src/api/ai_help.rs +++ b/src/api/ai_help.rs @@ -26,7 +26,7 @@ use crate::{ delete_full_help_history, delete_help_history, get_count, help_history, help_history_get_message, list_help_history, update_help_history_label, AI_HELP_LIMIT, }, - model::{AIHelpHistoryMessage, AIHelpHistoryMessageInsert, Settings}, + model::{AIHelpHistoryMessage, AIHelpHistoryMessageInsert, Settings, UserQuery}, settings::get_settings, SupaPool, }, @@ -195,7 +195,7 @@ fn record_question( pool: &Data, message: &ChatCompletionRequestMessage, history_enabled: bool, - user_id: i64, + user: &UserQuery, help_ids: HelpIds, ) -> Result, ApiError> { if !history_enabled { @@ -207,11 +207,22 @@ fn record_question( message_id, parent_id, } = help_ids; - if let Err(err) = add_help_history(&mut conn, user_id, chat_id) { + + // Check if our message has a parent message. + // If not, we are in a conversation that started while history was disabled. + // We do not store either the history/chat nor the message itself in this case. + if let Some(parent_id) = parent_id { + let parent_message = help_history_get_message(&mut conn, user, &parent_id)?; + if parent_message.is_none() { + return Ok(None); + } + } + + if let Err(err) = add_help_history(&mut conn, user.id, chat_id) { error!("AI Help log: {err}"); } let insert = AIHelpHistoryMessageInsert { - user_id, + user_id: user.id, chat_id, message_id, parent_id, @@ -233,20 +244,28 @@ fn record_sources( pool: &Data, sources: &Vec, history_enabled: bool, - user_id: i64, + user: &UserQuery, help_ids: HelpIds, ) -> Result, ApiError> { if !history_enabled { return Ok(None); } let mut conn = pool.get()?; + let HelpIds { chat_id, message_id, parent_id, } = help_ids; + if let Some(parent_id) = parent_id { + let parent_message = help_history_get_message(&mut conn, user, &parent_id)?; + if parent_message.is_none() { + return Ok(None); + } + } + let insert = AIHelpHistoryMessageInsert { - user_id, + user_id: user.id, chat_id, message_id, parent_id, @@ -393,7 +412,7 @@ pub async fn ai_help( &diesel_pool, question, history_enabled(&settings), - user.id, + &user, help_ids, )?; } @@ -405,7 +424,7 @@ pub async fn ai_help( &diesel_pool, &sources, history_enabled(&settings), - user.id, + &user, help_ids, )? { Some(x) => Utc.from_utc_datetime(&x), diff --git a/src/db/ai_help.rs b/src/db/ai_help.rs index 26e7ed15..448442fc 100644 --- a/src/db/ai_help.rs +++ b/src/db/ai_help.rs @@ -140,21 +140,6 @@ pub fn add_help_history_message( conn: &mut PgConnection, mut message: AIHelpHistoryMessageInsert, ) -> Result { - // Check if the parent message exists in the database first. - // We could end up with a non-existing parent message if the - // user enables history after the conversation has already - // started. Do not try to store the message in that case. - if let Some(parent_id) = message.parent_id { - let parent_message_count = ai_help_history_messages::table - .filter(ai_help_history_messages::message_id.eq(&parent_id)) - .count() - .get_result::(conn)?; - if parent_message_count == 0 { - return Ok(Utc::now().naive_utc()); - } - } - - // Ok, we are good on the parent message. let updated_at = update(ai_help_history::table) .filter( ai_help_history::user_id diff --git a/tests/api/ai_help_history.rs b/tests/api/ai_help_history.rs index 93570b61..c00f3f25 100644 --- a/tests/api/ai_help_history.rs +++ b/tests/api/ai_help_history.rs @@ -5,12 +5,9 @@ use actix_web::test; use anyhow::Error; use async_openai::types::ChatCompletionRequestMessage; use async_openai::types::Role::{Assistant, User}; -use diesel::prelude::*; -use diesel::ExpressionMethods; use rumba::ai::help::RefDoc; use rumba::db::ai_help::{add_help_history, add_help_history_message}; use rumba::db::model::{AIHelpHistoryMessageInsert, SettingsInsert}; -use rumba::db::schema::ai_help_history_messages; use rumba::db::settings::create_or_update_settings; use serde_json::Value::Null; use uuid::Uuid; @@ -98,48 +95,6 @@ async fn test_history() -> Result<(), Error> { Ok(()) } -#[actix_rt::test] -#[stubr::mock(port = 4321)] -async fn test_history_message_without_parent() -> Result<(), Error> { - let pool = reset()?; - let app = test_app_with_login(&pool).await.unwrap(); - let service = test::init_service(app).await; - let mut _logged_in_client = TestHttpClient::new(service).await; - - let mut conn = pool.get()?; - // Make sure that history is enabled - create_or_update_settings( - &mut conn, - SettingsInsert { - user_id: 1, - ai_help_history: Some(true), - ..Default::default() - }, - )?; - - // create a message with a non-existing parent - let message = AIHelpHistoryMessageInsert { - user_id: 1, - chat_id: CHAT_ID, - message_id: MESSAGE_ID, - parent_id: Some(Uuid::from_u128(2)), - created_at: None, - sources: None, - request: None, - response: None, - }; - let res = add_help_history_message(&mut conn, message); - // We return OK, but the message should not be recorded. - assert!(res.is_ok()); - let message_count = ai_help_history_messages::table - .filter(ai_help_history_messages::user_id.eq(1)) - .count() - .get_result::(&mut conn)?; - assert_eq!(message_count, 0); - drop_stubr(stubr).await; - Ok(()) -} - fn normalize_digits(s: &str) -> String { let mut result = String::new(); From c44a11b017683892f3175de1702ce278388d165e Mon Sep 17 00:00:00 2001 From: Andi Pieper Date: Wed, 21 Feb 2024 23:40:08 +0100 Subject: [PATCH 3/6] refactored for clean happy path --- src/api/ai_help.rs | 28 ++++++++++++++--------- src/db/ai_help.rs | 56 ++++++++++++++++++++++++++++++++-------------- 2 files changed, 56 insertions(+), 28 deletions(-) diff --git a/src/api/ai_help.rs b/src/api/ai_help.rs index c858be2f..3549532c 100644 --- a/src/api/ai_help.rs +++ b/src/api/ai_help.rs @@ -24,7 +24,8 @@ use crate::{ ai_help::{ add_help_history, add_help_history_message, create_or_increment_total, delete_full_help_history, delete_help_history, get_count, help_history, - help_history_get_message, list_help_history, update_help_history_label, AI_HELP_LIMIT, + help_history_get_message, list_help_history, update_help_history, + update_help_history_label, AI_HELP_LIMIT, }, model::{AIHelpHistoryMessage, AIHelpHistoryMessageInsert, Settings, UserQuery}, settings::get_settings, @@ -208,19 +209,24 @@ fn record_question( parent_id, } = help_ids; - // Check if our message has a parent message. - // If not, we are in a conversation that started while history was disabled. - // We do not store either the history/chat nor the message itself in this case. - if let Some(parent_id) = parent_id { - let parent_message = help_history_get_message(&mut conn, user, &parent_id)?; - if parent_message.is_none() { - return Ok(None); - } - } + // If a `parent_id` is present, we execute an update on the help history + // record because one of these are true: + // * We created a history record at the beginning of the conversation. + // * History was switched off, we did not create a record and the update + // will simply not match/change any record. + // + // With no `parent_id`, we create a new record because at this point, history + // _is_ enabled and we are at the start of a new conversation. + let res = if parent_id.is_some() { + update_help_history(&mut conn, user.id, chat_id) + } else { + add_help_history(&mut conn, user.id, chat_id) + }; - if let Err(err) = add_help_history(&mut conn, user.id, chat_id) { + if let Err(err) = res { error!("AI Help log: {err}"); } + let insert = AIHelpHistoryMessageInsert { user_id: user.id, chat_id, diff --git a/src/db/ai_help.rs b/src/db/ai_help.rs index 448442fc..e7332646 100644 --- a/src/db/ai_help.rs +++ b/src/db/ai_help.rs @@ -128,35 +128,54 @@ pub fn add_help_history( }; insert_into(ai_help_history::table) .values(history) - .on_conflict(ai_help_history::chat_id) - .do_update() - .set(ai_help_history::updated_at.eq(diesel::dsl::now)) .execute(conn)?; - Ok(()) } -pub fn add_help_history_message( +pub fn update_help_history( conn: &mut PgConnection, - mut message: AIHelpHistoryMessageInsert, -) -> Result { - let updated_at = update(ai_help_history::table) + user_id: i64, + chat_id: Uuid, +) -> Result<(), DbError> { + update(ai_help_history::table) .filter( ai_help_history::user_id - .eq(message.user_id) - .and(ai_help_history::chat_id.eq(message.chat_id)), + .eq(user_id) + .and(ai_help_history::chat_id.eq(chat_id)), ) .set(ai_help_history::updated_at.eq(diesel::dsl::now)) - .returning(ai_help_history::updated_at) - .get_result::(conn)?; - message.created_at = Some(updated_at); - insert_into(ai_help_history_messages::table) + .execute(conn)?; + Ok(()) +} + +pub fn add_help_history_message( + conn: &mut PgConnection, + message: AIHelpHistoryMessageInsert, +) -> Result { + //message.created_at = Some(Utc::now().naive_utc()); + let res = insert_into(ai_help_history_messages::table) .values(&message) .on_conflict(ai_help_history_messages::message_id) .do_update() .set(&message) - .execute(conn)?; - Ok(updated_at) + .returning(ai_help_history_messages::created_at) + .get_result::(conn); + match res { + Ok(created_at) => Ok(created_at), + Err(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::ForeignKeyViolation, + _info, + )) => { + // We have reached an edge case where the ai help history record + // is missing. Most likely the user had history turned off + // at the start of the conversation, causing a foreign key violation + // now. + // Handling that specific case here spares us to do a pre-flight check + // on every call, keeping the happy path fast. + Ok(Utc::now().naive_utc()) + } + Err(e) => Err(e.into()), + } } pub fn help_history_get_message( @@ -255,7 +274,10 @@ pub fn update_help_history_label( .eq(user.id) .and(ai_help_history::chat_id.eq(chat_id)), ) - .set(ai_help_history::label.eq(label)) + .set(( + ai_help_history::label.eq(label), + ai_help_history::updated_at.eq(diesel::dsl::now), + )) .execute(conn)?; Ok(()) } From 2ac79b0a709d255cbb14afc0cf32bbcde28937c0 Mon Sep 17 00:00:00 2001 From: Andi Pieper Date: Wed, 21 Feb 2024 23:55:49 +0100 Subject: [PATCH 4/6] removed dead code --- src/api/ai_help.rs | 8 -------- 1 file changed, 8 deletions(-) diff --git a/src/api/ai_help.rs b/src/api/ai_help.rs index 3549532c..50ff6668 100644 --- a/src/api/ai_help.rs +++ b/src/api/ai_help.rs @@ -222,7 +222,6 @@ fn record_question( } else { add_help_history(&mut conn, user.id, chat_id) }; - if let Err(err) = res { error!("AI Help log: {err}"); } @@ -257,18 +256,11 @@ fn record_sources( return Ok(None); } let mut conn = pool.get()?; - let HelpIds { chat_id, message_id, parent_id, } = help_ids; - if let Some(parent_id) = parent_id { - let parent_message = help_history_get_message(&mut conn, user, &parent_id)?; - if parent_message.is_none() { - return Ok(None); - } - } let insert = AIHelpHistoryMessageInsert { user_id: user.id, From 27829c2f4b04ded11335685c786f4545b614ea3d Mon Sep 17 00:00:00 2001 From: Andi Pieper Date: Thu, 22 Feb 2024 00:52:57 +0100 Subject: [PATCH 5/6] move history record logic into message record handling --- src/api/ai_help.rs | 24 +++--------------------- src/db/ai_help.rs | 33 +++++++++++++++++++++++---------- 2 files changed, 26 insertions(+), 31 deletions(-) diff --git a/src/api/ai_help.rs b/src/api/ai_help.rs index 50ff6668..62446136 100644 --- a/src/api/ai_help.rs +++ b/src/api/ai_help.rs @@ -22,10 +22,9 @@ use crate::{ db::{ self, ai_help::{ - add_help_history, add_help_history_message, create_or_increment_total, - delete_full_help_history, delete_help_history, get_count, help_history, - help_history_get_message, list_help_history, update_help_history, - update_help_history_label, AI_HELP_LIMIT, + add_help_history_message, create_or_increment_total, delete_full_help_history, + delete_help_history, get_count, help_history, help_history_get_message, + list_help_history, update_help_history_label, AI_HELP_LIMIT, }, model::{AIHelpHistoryMessage, AIHelpHistoryMessageInsert, Settings, UserQuery}, settings::get_settings, @@ -209,23 +208,6 @@ fn record_question( parent_id, } = help_ids; - // If a `parent_id` is present, we execute an update on the help history - // record because one of these are true: - // * We created a history record at the beginning of the conversation. - // * History was switched off, we did not create a record and the update - // will simply not match/change any record. - // - // With no `parent_id`, we create a new record because at this point, history - // _is_ enabled and we are at the start of a new conversation. - let res = if parent_id.is_some() { - update_help_history(&mut conn, user.id, chat_id) - } else { - add_help_history(&mut conn, user.id, chat_id) - }; - if let Err(err) = res { - error!("AI Help log: {err}"); - } - let insert = AIHelpHistoryMessageInsert { user_id: user.id, chat_id, diff --git a/src/db/ai_help.rs b/src/db/ai_help.rs index e7332646..8fff2ba1 100644 --- a/src/db/ai_help.rs +++ b/src/db/ai_help.rs @@ -128,6 +128,9 @@ pub fn add_help_history( }; insert_into(ai_help_history::table) .values(history) + .on_conflict(ai_help_history::chat_id) + .do_update() + .set(ai_help_history::updated_at.eq(diesel::dsl::now)) .execute(conn)?; Ok(()) } @@ -152,7 +155,23 @@ pub fn add_help_history_message( conn: &mut PgConnection, message: AIHelpHistoryMessageInsert, ) -> Result { - //message.created_at = Some(Utc::now().naive_utc()); + // If a `parent_id` is present, we execute an update on the help history + // record because one of these are true: + // * We created a history record at the beginning of the conversation. + // * History was switched off, we did not create a record and the update + // will simply not match/change any record. + // + // With no `parent_id`, we create a new record because at this point, history + // _is_ enabled and we are at the start of a new conversation. + let res = if message.parent_id.is_some() { + update_help_history(conn, message.user_id, message.chat_id) + } else { + add_help_history(conn, message.user_id, message.chat_id) + }; + if let Err(err) = res { + error!("AI Help log: {err}"); + } + let res = insert_into(ai_help_history_messages::table) .values(&message) .on_conflict(ai_help_history_messages::message_id) @@ -162,18 +181,12 @@ pub fn add_help_history_message( .get_result::(conn); match res { Ok(created_at) => Ok(created_at), + // Ignore foreign key violations deliberately + // because of the edge cases described above Err(diesel::result::Error::DatabaseError( diesel::result::DatabaseErrorKind::ForeignKeyViolation, _info, - )) => { - // We have reached an edge case where the ai help history record - // is missing. Most likely the user had history turned off - // at the start of the conversation, causing a foreign key violation - // now. - // Handling that specific case here spares us to do a pre-flight check - // on every call, keeping the happy path fast. - Ok(Utc::now().naive_utc()) - } + )) => Ok(Utc::now().naive_utc()), Err(e) => Err(e.into()), } } From 590f99a8a6b18080036f29d14bceae069d119e3c Mon Sep 17 00:00:00 2001 From: Florian Dieminger Date: Fri, 23 Feb 2024 10:05:36 +0100 Subject: [PATCH 6/6] revert user -> user_id and don't updade label ts --- src/api/ai_help.rs | 14 +++++++------- src/db/ai_help.rs | 5 +---- 2 files changed, 8 insertions(+), 11 deletions(-) diff --git a/src/api/ai_help.rs b/src/api/ai_help.rs index 62446136..0c97975b 100644 --- a/src/api/ai_help.rs +++ b/src/api/ai_help.rs @@ -26,7 +26,7 @@ use crate::{ delete_help_history, get_count, help_history, help_history_get_message, list_help_history, update_help_history_label, AI_HELP_LIMIT, }, - model::{AIHelpHistoryMessage, AIHelpHistoryMessageInsert, Settings, UserQuery}, + model::{AIHelpHistoryMessage, AIHelpHistoryMessageInsert, Settings}, settings::get_settings, SupaPool, }, @@ -195,7 +195,7 @@ fn record_question( pool: &Data, message: &ChatCompletionRequestMessage, history_enabled: bool, - user: &UserQuery, + user_id: i64, help_ids: HelpIds, ) -> Result, ApiError> { if !history_enabled { @@ -209,7 +209,7 @@ fn record_question( } = help_ids; let insert = AIHelpHistoryMessageInsert { - user_id: user.id, + user_id, chat_id, message_id, parent_id, @@ -231,7 +231,7 @@ fn record_sources( pool: &Data, sources: &Vec, history_enabled: bool, - user: &UserQuery, + user_id: i64, help_ids: HelpIds, ) -> Result, ApiError> { if !history_enabled { @@ -245,7 +245,7 @@ fn record_sources( } = help_ids; let insert = AIHelpHistoryMessageInsert { - user_id: user.id, + user_id, chat_id, message_id, parent_id, @@ -392,7 +392,7 @@ pub async fn ai_help( &diesel_pool, question, history_enabled(&settings), - &user, + user.id, help_ids, )?; } @@ -404,7 +404,7 @@ pub async fn ai_help( &diesel_pool, &sources, history_enabled(&settings), - &user, + user.id, help_ids, )? { Some(x) => Utc.from_utc_datetime(&x), diff --git a/src/db/ai_help.rs b/src/db/ai_help.rs index 8fff2ba1..e0ca699d 100644 --- a/src/db/ai_help.rs +++ b/src/db/ai_help.rs @@ -287,10 +287,7 @@ pub fn update_help_history_label( .eq(user.id) .and(ai_help_history::chat_id.eq(chat_id)), ) - .set(( - ai_help_history::label.eq(label), - ai_help_history::updated_at.eq(diesel::dsl::now), - )) + .set(ai_help_history::label.eq(label)) .execute(conn)?; Ok(()) }