Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'evals' of https://github.com/buster-so/buster into evals

This commit is contained in:
Nate Kelley 2025-04-18 15:34:04 -06:00
parent 2eda72e60e 0d3f1589c1
commit b69bdda2c4
No known key found for this signature in database
GPG Key ID: FD90372AB8D98B4F
1 changed files with 39 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
api/libs/handlers/src/chats/restore_chat_handler.rs
View File

@ -4,9 +4,9 @@ use database::{
enums::AssetType,
models::{Message, MessageToFile},
pool::get_pg_pool,
schema::{chats, messages, messages_to_files},
schema::{chats, dashboard_files, messages, messages_to_files, metric_files},
};
use diesel::{insert_into, update, ExpressionMethods, QueryDsl};
use diesel::{update, ExpressionMethods, QueryDsl};
use diesel_async::RunQueryDsl;
use futures::future::try_join_all;
use middleware::AuthenticatedUser;
@ -132,12 +132,46 @@ pub async fn restore_chat_handler(
let (file_type, file_name, file_id, version_number) = restore_result??;
let last_message = last_message_result??;
// --- START: Fetch restored content from main table ---
let file_content: String = {
let mut conn = get_pg_pool().get().await?;
match request.asset_type {
AssetType::MetricFile => {
let content_json = metric_files::table
.filter(metric_files::id.eq(&file_id))
.select(metric_files::content)
.first::<serde_json::Value>(&mut conn)
.await?;
// Convert JSON Value to pretty YAML String
serde_yaml::to_string(&content_json)
.map_err(|e| anyhow!("Failed to convert metric content to YAML: {}", e))?
}
AssetType::DashboardFile => {
let content_json = dashboard_files::table
.filter(dashboard_files::id.eq(&file_id))
.select(dashboard_files::content)
.first::<serde_json::Value>(&mut conn)
.await?;
// Convert JSON Value to pretty YAML String for the message
serde_yaml::to_string(&content_json)
.map_err(|e| anyhow!("Failed to convert dashboard content to YAML: {}", e))?
}
// This case should be unreachable due to the check in the restore_task
_ => {
return Err(anyhow!(
"Unexpected asset type after successful restoration: {:?}",
request.asset_type
))
}
}
};
// --- END: Fetch restored content from main table ---
// Step 3: Construct message details
let tool_call_id = format!("call_{}", Uuid::new_v4().to_string().replace("-", ""));
let mut raw_llm_messages = if let Some(last_msg) = &last_message {
// Use clone here if last_message is Some(Message)
if let Ok(msgs) = serde_json::from_value::<Vec<Value>>(last_msg.raw_llm_messages.clone())
{
if let Ok(msgs) = serde_json::from_value::<Vec<Value>>(last_msg.raw_llm_messages.clone()) {
msgs
} else {
Vec::new()
@ -167,7 +201,7 @@ pub async fn restore_chat_handler(
"role": "tool",
"content": json!({
"message": format!("Successfully restored 1 {} files.", file_type),
"file_contents": file_name
"file_contents": file_content
}).to_string(),
"tool_call_id": tool_call_id
}));
@ -236,7 +270,6 @@ pub async fn restore_chat_handler(
let request_asset_type_clone = request.asset_type; // AssetType is likely Copy
let file_id_clone = file_id; // file_id is Uuid (Copy)
let insert_message_task = tokio::spawn(async move {
let mut conn = get_pg_pool().get().await?;
diesel::insert_into(messages::table)