-
Notifications
You must be signed in to change notification settings - Fork 16
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat(ai-help): delete ai history after six months #437
Changes from 7 commits
8a8cbee
f500c28
42718cc
ace4596
59f6ad5
da70a77
f59fa38
534df25
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,35 @@ | ||
use crate::db::schema::ai_help_history; | ||
use crate::db::Pool; | ||
use crate::diesel::QueryDsl; | ||
use crate::{api::error::ApiError, settings::SETTINGS}; | ||
use actix_web::web::Data; | ||
use chrono::Utc; | ||
use diesel::{ExpressionMethods, RunQueryDsl}; | ||
use std::ops::Sub; | ||
use std::time::Duration; | ||
|
||
/// This removes old AI history records from the database. It is meant to be called from a | ||
/// cron job calling the respective endpoint in the admin API. | ||
pub async fn do_delete_old_ai_history(pool: Data<Pool>) -> Result<(), ApiError> { | ||
let mut conn = pool.get()?; | ||
let history_deletion_period_in_sec = SETTINGS | ||
.ai | ||
.as_ref() | ||
.map(|ai| ai.history_deletion_period_in_sec) | ||
.ok_or(ApiError::Generic( | ||
"ai.history_deletion_period_in_sec missing from configuration".to_string(), | ||
))?; | ||
|
||
let oldest_timestamp = Utc::now() | ||
.sub(Duration::from_secs(history_deletion_period_in_sec)) | ||
.naive_utc(); | ||
|
||
let affected_rows = diesel::delete( | ||
ai_help_history::table.filter(ai_help_history::updated_at.lt(oldest_timestamp)), | ||
) | ||
.execute(&mut conn)?; | ||
info!( | ||
"Deleted old AI history before {oldest_timestamp}: {affected_rows} old record(s) deleted." | ||
); | ||
Ok(()) | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,5 +1,6 @@ | ||
pub mod ai_explain; | ||
pub mod ai_help; | ||
pub mod ai_history; | ||
pub mod documents; | ||
pub mod error; | ||
pub mod fxa_webhook; | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -5,11 +5,19 @@ use actix_web::test; | |
use anyhow::Error; | ||
use async_openai::types::ChatCompletionRequestMessage; | ||
use async_openai::types::Role::{Assistant, User}; | ||
use chrono::{NaiveDateTime, Utc}; | ||
use diesel::dsl::count; | ||
use diesel::prelude::*; | ||
use diesel::{insert_into, ExpressionMethods, RunQueryDsl}; | ||
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::model::{AIHelpHistoryInsert, AIHelpHistoryMessageInsert, SettingsInsert}; | ||
use rumba::db::schema::ai_help_history; | ||
use rumba::db::settings::create_or_update_settings; | ||
use rumba::settings::SETTINGS; | ||
use serde_json::Value::Null; | ||
use std::ops::Sub; | ||
use std::time::Duration; | ||
use uuid::Uuid; | ||
|
||
const CHAT_ID: Uuid = Uuid::nil(); | ||
|
@@ -95,6 +103,105 @@ async fn test_history() -> Result<(), Error> { | |
Ok(()) | ||
} | ||
|
||
#[actix_rt::test] | ||
#[stubr::mock(port = 4321)] | ||
async fn test_history_deletion() -> 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()?; | ||
|
||
let history_deletion_period_in_sec = &SETTINGS | ||
.ai | ||
.as_ref() | ||
.map(|ai| ai.history_deletion_period_in_sec) | ||
.expect("ai.history_deletion_period_in_sec missing"); | ||
|
||
// Add an old chat history entry, double our configured period ago. | ||
let ts = Utc::now() | ||
.sub(Duration::from_secs(history_deletion_period_in_sec * 2)) | ||
.naive_utc(); | ||
let history = AIHelpHistoryInsert { | ||
user_id: 1, | ||
chat_id: Uuid::from_u128(1), | ||
created_at: Some(ts), | ||
updated_at: Some(ts), | ||
label: "old entry".to_string(), | ||
}; | ||
insert_into(ai_help_history::table) | ||
.values(history) | ||
.execute(&mut conn)?; | ||
|
||
// Add a newer chat history entry, half of our configured period ago. | ||
let ts = Utc::now() | ||
// .checked_sub_months(Months::new(2)) | ||
.sub(Duration::from_secs(history_deletion_period_in_sec / 2)) | ||
.naive_utc(); | ||
let history = AIHelpHistoryInsert { | ||
user_id: 1, | ||
chat_id: Uuid::from_u128(2), | ||
created_at: Some(ts), | ||
updated_at: Some(ts), | ||
label: "new entry".to_string(), | ||
}; | ||
insert_into(ai_help_history::table) | ||
.values(history) | ||
.execute(&mut conn)?; | ||
|
||
// check the history count before we run the delete job | ||
let rec_count = ai_help_history::table | ||
.filter(ai_help_history::user_id.eq(1)) | ||
.select(count(ai_help_history::user_id)) | ||
.first::<i64>(&mut conn)?; | ||
|
||
assert_eq!(2, rec_count); | ||
|
||
// Now, run the delete job. | ||
let req = test::TestRequest::post() | ||
.uri("/admin-api/ai-history/") | ||
.insert_header(( | ||
"Authorization", | ||
format!("Bearer {}", SETTINGS.auth.admin_update_bearer_token), | ||
)) | ||
.to_request(); | ||
|
||
let res = test::call_service(&service, req).await; | ||
assert!(res.response().status().is_success()); | ||
|
||
// Check database that the old entry is gone. | ||
// Loop until we see the old entry is gone because the | ||
// delete job runs asynchonously. | ||
let mut retry = 0; | ||
const MAX_RETRIES: u32 = 40; | ||
Comment on lines
+175
to
+176
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Curious why you use There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes, u8 would suffice here. It is test code though, so I we can be a bit more relaxed about this. |
||
let mut records: Vec<NaiveDateTime>; | ||
loop { | ||
records = ai_help_history::table | ||
.filter(ai_help_history::user_id.eq(1)) | ||
.select(ai_help_history::updated_at) | ||
.get_results(&mut conn)?; | ||
if records.len() == 1 { | ||
break; | ||
} | ||
|
||
actix_rt::time::sleep(Duration::from_millis(5)).await; | ||
retry += 1; | ||
if retry > MAX_RETRIES { | ||
break; | ||
} | ||
} | ||
Comment on lines
+178
to
+192
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Maybe this could be made more readable by using the There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes, probably. Lets look into retry crate next time we have flakyness issues. |
||
|
||
assert_eq!(1, records.len()); | ||
// we only compare up to millisecond precision | ||
assert_eq!( | ||
ts.timestamp_micros(), | ||
records.get(0).unwrap().timestamp_micros() | ||
); | ||
|
||
drop_stubr(stubr).await; | ||
Ok(()) | ||
} | ||
|
||
fn normalize_digits(s: &str) -> String { | ||
let mut result = String::new(); | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just wondering out loud if this would be "more" idiomatic:
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, it looks more rustful.