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' into api_search_rest_endpoint

This commit is contained in:
dal 2025-03-25 14:31:07 -06:00
parent 22322f0711 60e54d0257
commit df341a0ac5
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
20 changed files with 808 additions and 72 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
api/libs/database/src/models.rs
View File

@ -45,12 +45,12 @@ pub struct DashboardFile {
#[diesel(table_name = messages)]
pub struct Message {
pub id: Uuid,
pub request_message: String,
pub request_message: Option<String>,
pub response_messages: Value,
pub reasoning: Value,
pub title: String,
pub raw_llm_messages: Value,
pub final_reasoning_message: String,
pub final_reasoning_message: Option<String>,
pub chat_id: Uuid,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,

12
api/libs/handlers/src/chats/asset_messages.rs
View File

@ -38,7 +38,7 @@ pub async fn generate_asset_messages(
let file_message_id = Uuid::new_v4();
let file_message = Message {
id: file_message_id,
request_message: String::new(), // Empty request for auto-generated messages
request_message: None, // Empty request for auto-generated messages
chat_id: Uuid::nil(), // Will be set by caller
created_by: user.id,
created_at: Utc::now(),
@ -61,8 +61,8 @@ pub async fn generate_asset_messages(
]
}]),
reasoning: serde_json::Value::Array(vec![]),
final_reasoning_message: "Auto-generated file message".to_string(),
title: format!("View {}", asset_details.name),
final_reasoning_message: None,
title: format!("Chat with {}", asset_details.name),
raw_llm_messages: serde_json::json!([]),
feedback: None,
};
@ -71,7 +71,7 @@ pub async fn generate_asset_messages(
let text_message_id = Uuid::new_v4();
let text_message = Message {
id: text_message_id,
request_message: String::new(), // Empty request for auto-generated messages
request_message: None, // Empty request for auto-generated messages
chat_id: Uuid::nil(), // Will be set by caller
created_by: user.id,
created_at: Utc::now(),
@ -84,8 +84,8 @@ pub async fn generate_asset_messages(
"isFinalMessage": true
}]),
reasoning: serde_json::Value::Array(vec![]),
final_reasoning_message: "Auto-generated text message".to_string(),
title: format!("View {}", asset_details.name),
final_reasoning_message: None,
title: format!("Chat with {}", asset_details.name),
raw_llm_messages: serde_json::json!([]),
feedback: None,
};

20
api/libs/handlers/src/chats/get_chat_handler.rs
View File

@ -33,10 +33,10 @@ pub struct ChatWithUser {
#[derive(Queryable)]
pub struct MessageWithUser {
pub id: Uuid,
pub request_message: String,
pub request_message: Option<String>,
pub response_messages: Value,
pub reasoning: Value,
pub final_reasoning_message: String,
pub final_reasoning_message: Option<String>,
pub created_at: DateTime<Utc>,
pub user_id: Uuid,
pub user_name: Option<String>,
@ -257,11 +257,15 @@ pub async fn get_chat_handler(
.map(|arr| arr.to_vec())
.unwrap_or_default();
let request_message = ChatUserMessage {
request: msg.request_message,
sender_id: msg.user_id,
sender_name: msg.user_name.unwrap_or_else(|| "Unknown".to_string()),
sender_avatar,
let request_message = if let Some(request_message) = msg.request_message {
Some(ChatUserMessage {
request: request_message,
sender_id: msg.user_id,
sender_name: msg.user_name.unwrap_or_else(|| "Unknown".to_string()),
sender_avatar,
})
} else {
None
};
ChatMessage::new_with_messages(
@ -269,7 +273,7 @@ pub async fn get_chat_handler(
request_message,
response_messages,
reasoning,
Some(msg.final_reasoning_message),
msg.final_reasoning_message,
msg.created_at,
)
})

2
api/libs/handlers/src/chats/mod.rs
View File

@ -9,6 +9,7 @@ pub mod context_loaders;
pub mod list_chats_handler;
pub mod sharing;
pub mod asset_messages;
pub mod restore_chat_handler;
#[cfg(test)]
pub mod tests;
@ -19,6 +20,7 @@ pub use post_chat_handler::post_chat_handler;
pub use update_chats_handler::update_chats_handler;
pub use delete_chats_handler::delete_chats_handler;
pub use list_chats_handler::list_chats_handler;
pub use restore_chat_handler::restore_chat_handler;
pub use sharing::delete_chat_sharing_handler;
pub use sharing::create_chat_sharing_handler;
pub use sharing::update_chat_sharing_handler;

28
api/libs/handlers/src/chats/post_chat_handler.rs
View File

@ -237,12 +237,12 @@ pub async fn post_chat_handler(
for message in updated_messages {
let chat_message = ChatMessage::new_with_messages(
message.id,
ChatUserMessage {
Some(ChatUserMessage {
request: "".to_string(),
sender_id: user.id,
sender_name: user.name.clone().unwrap_or_default(),
sender_avatar: None,
},
}),
// Use the response_messages from the DB
serde_json::from_value(message.response_messages).unwrap_or_default(),
vec![],
@ -451,12 +451,12 @@ pub async fn post_chat_handler(
// Update chat_with_messages with final state
let message = ChatMessage::new_with_messages(
message_id,
ChatUserMessage {
Some(ChatUserMessage {
request: request.prompt.clone().unwrap_or_default(),
sender_id: user.id,
sender_name: user.name.clone().unwrap_or_default(),
sender_avatar: None,
},
}),
response_messages.clone(),
reasoning_messages.clone(),
Some(format!("Reasoned for {} seconds", reasoning_duration).to_string()),
@ -468,7 +468,7 @@ pub async fn post_chat_handler(
// Create and store message in the database with final state
let db_message = Message {
id: message_id,
request_message: request.prompt.unwrap_or_default(),
request_message: Some(request.prompt.unwrap_or_default()),
chat_id,
created_by: user.id,
created_at: Utc::now(),
@ -476,7 +476,7 @@ pub async fn post_chat_handler(
deleted_at: None,
response_messages: serde_json::to_value(&response_messages)?,
reasoning: serde_json::to_value(&reasoning_messages)?,
final_reasoning_message: format!("Reasoned for {} seconds", reasoning_duration),
final_reasoning_message: Some(format!("Reasoned for {} seconds", reasoning_duration)),
title: title.title.clone().unwrap_or_default(),
raw_llm_messages: serde_json::to_value(&raw_llm_messages)?,
feedback: None,
@ -753,7 +753,6 @@ pub enum BusterChatMessage {
file_type: String,
file_name: String,
version_number: i32,
version_id: String,
filter_version_id: Option<String>,
metadata: Option<Vec<BusterChatResponseFileMetadata>>,
},
@ -802,7 +801,6 @@ pub struct BusterFile {
pub file_type: String,
pub file_name: String,
pub version_number: i32,
pub version_id: String,
pub status: String,
pub file: BusterFileContent,
pub metadata: Option<Vec<BusterFileMetadata>>,
@ -943,7 +941,6 @@ pub async fn transform_message(
file_type: file_content.file_type.clone(),
file_name: file_content.file_name.clone(),
version_number: file_content.version_number,
version_id: file_content.version_id.clone(),
filter_version_id: None,
metadata: Some(vec![BusterChatResponseFileMetadata {
status: "completed".to_string(),
@ -1030,7 +1027,6 @@ pub async fn transform_message(
file_type: file_content.file_type.clone(),
file_name: file_content.file_name.clone(),
version_number: file_content.version_number,
version_id: file_content.version_id.clone(),
filter_version_id: None,
metadata: Some(vec![BusterChatResponseFileMetadata {
status: "completed".to_string(),
@ -1195,7 +1191,6 @@ fn tool_create_metrics(id: String, content: String) -> Result<Vec<BusterReasonin
file_type: "metric".to_string(),
file_name: file.name.clone(),
version_number: 1,
version_id: file.id.to_string(),
status: "completed".to_string(),
file: BusterFileContent {
text: Some(file.yml_content),
@ -1250,7 +1245,6 @@ fn tool_modify_metrics(id: String, content: String) -> Result<Vec<BusterReasonin
file_type: "metric".to_string(),
file_name: file.name.clone(),
version_number: 1,
version_id: file.id.to_string(),
status: "completed".to_string(),
file: BusterFileContent {
text: Some(file.yml_content),
@ -1306,7 +1300,6 @@ fn tool_create_dashboards(id: String, content: String) -> Result<Vec<BusterReaso
file_type: "dashboard".to_string(),
file_name: file.name.clone(),
version_number: 1,
version_id: file.id.to_string(),
status: "completed".to_string(),
file: BusterFileContent {
text: Some(file.yml_content),
@ -1361,7 +1354,6 @@ fn tool_modify_dashboards(id: String, content: String) -> Result<Vec<BusterReaso
file_type: "dashboard".to_string(),
file_name: file.name.clone(),
version_number: 1,
version_id: file.id.to_string(),
status: "completed".to_string(),
file: BusterFileContent {
text: Some(file.yml_content),
@ -2014,12 +2006,12 @@ async fn initialize_chat(
// Create new message
let message = ChatMessage::new_with_messages(
message_id,
ChatUserMessage {
Some(ChatUserMessage {
request: prompt_text,
sender_id: user.id,
sender_name: user.name.clone().unwrap_or_default(),
sender_avatar: None,
},
}),
Vec::new(),
Vec::new(),
None,
@ -2050,12 +2042,12 @@ async fn initialize_chat(
// Create initial message
let message = ChatMessage::new_with_messages(
message_id,
ChatUserMessage {
Some(ChatUserMessage {
request: prompt_text,
sender_id: user.id,
sender_name: user.name.clone().unwrap_or_default(),
sender_avatar: None,
},
}),
Vec::new(),
Vec::new(),
None,

305
api/libs/handlers/src/chats/restore_chat_handler.rs Normal file
View File

@ -0,0 +1,305 @@
use anyhow::{anyhow, Result};
use chrono::Utc;
use database::{
enums::AssetType,
models::{Message, MessageToFile},
pool::get_pg_pool,
schema::{messages, messages_to_files},
};
use diesel::{insert_into, ExpressionMethods, QueryDsl};
use diesel_async::RunQueryDsl;
use middleware::AuthenticatedUser;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use uuid::Uuid;
use crate::chats::get_chat_handler::get_chat_handler;
use crate::chats::types::ChatWithMessages;
// Import public handler types directly
use crate::dashboards::{update_dashboard_handler, DashboardUpdateRequest};
use crate::metrics::{update_metric_handler, UpdateMetricRequest};
/// Request structure for restoring an asset (metric or dashboard) version in a chat
#[derive(Debug, Serialize, Deserialize)]
pub struct ChatRestoreRequest {
/// ID of the asset to restore
pub asset_id: Uuid,
/// Type of asset to restore (metric_file or dashboard_file)
pub asset_type: AssetType,
/// Version number to restore
pub version_number: i32,
}
/// Handler for restoring a previous version of an asset (metric or dashboard) and documenting it in a chat
///
/// # Arguments
/// * `chat_id` - The UUID of the chat where the restoration will be documented
/// * `user` - The authenticated user making the request
/// * `request` - The restoration request with asset_id, asset_type, and version_number
///
/// # Returns
/// * `Result<ChatWithMessages>` - The updated chat with new messages documenting the restoration
///
/// # Process
/// 1. Restores the specified asset version using the appropriate handler
/// 2. Creates a text message in the chat documenting the restoration
/// 3. Creates a file message linking to the restored asset
/// 4. Returns the updated chat with all messages
pub async fn restore_chat_handler(
chat_id: &Uuid,
user: &AuthenticatedUser,
request: ChatRestoreRequest,
) -> Result<ChatWithMessages> {
let mut conn = get_pg_pool().get().await?;
// Step 1: Restore the asset using the appropriate handler
let (file_type, file_name, file_id, version_number) = match request.asset_type {
AssetType::MetricFile => {
// Create a metric update request with only the restore_to_version parameter
let metric_request = UpdateMetricRequest {
restore_to_version: Some(request.version_number),
..Default::default()
};
// Call the metric update handler through the public module function
let updated_metric =
update_metric_handler(&request.asset_id, user, metric_request).await?;
// Return the file information
(
"metric".to_string(),
updated_metric.name,
updated_metric.id,
updated_metric.versions.len() as i32, // Get version number from versions length
)
}
AssetType::DashboardFile => {
// Create a dashboard update request with only the restore_to_version parameter
let dashboard_request = DashboardUpdateRequest {
restore_to_version: Some(request.version_number),
..Default::default()
};
// Call the dashboard update handler through the public module function
let updated_dashboard =
update_dashboard_handler(request.asset_id, dashboard_request, user).await?;
// Return the file information
(
"dashboard".to_string(),
updated_dashboard.dashboard.name,
updated_dashboard.dashboard.id,
updated_dashboard.dashboard.version_number,
)
}
_ => {
return Err(anyhow!(
"Unsupported asset type for restoration: {:?}",
request.asset_type
))
}
};
// Step 2: Get the most recent message to copy raw_llm_messages
// Fetch the most recent message for the chat to extract raw_llm_messages
let last_message = messages::table
.filter(messages::chat_id.eq(chat_id))
.filter(messages::deleted_at.is_null())
.limit(1)
// We need to use order here to get the latest message
.then_order_by(messages::created_at.desc())
.first::<Message>(&mut conn)
.await
.ok();
// Create raw_llm_messages by copying from the previous message and adding restoration entries
let tool_call_id = format!("call_{}", Uuid::new_v4().to_string().replace("-", ""));
// Start with copied raw_llm_messages or an empty array
let mut raw_llm_messages = if let Some(last_msg) = &last_message {
if let Ok(msgs) = serde_json::from_value::<Vec<Value>>(last_msg.raw_llm_messages.clone()) {
msgs
} else {
Vec::new()
}
} else {
Vec::new()
};
// Add tool call message and tool response message
raw_llm_messages.push(json!({
"name": "buster_super_agent",
"role": "assistant",
"tool_calls": [
{
"id": tool_call_id,
"type": "function",
"function": {
"name": format!("restore_{}", file_type),
"arguments": json!({
"version_number": request.version_number
}).to_string()
}
}
]
}));
// Add the tool response
raw_llm_messages.push(json!({
"name": format!("restore_{}", file_type),
"role": "tool",
"content": json!({
"message": format!("Successfully restored 1 {} files.", file_type),
"file_contents": file_name
}).to_string(),
"tool_call_id": tool_call_id
}));
// Step 3: Create a message with text and file responses
let message_id = Uuid::new_v4();
let now = Utc::now();
let timestamp = now.timestamp();
// Create restoration message text response
let restoration_text = format!(
"Version {} was created by restoring version {}",
version_number, request.version_number
);
// Create file response message for the restored asset
// Create response messages array with both text and file response
let response_messages = json!([
// Text response message
{
"id": format!("chatcmpl-{}", Uuid::new_v4().to_string().replace("-", "")),
"type": "text",
"message": restoration_text,
"message_chunk": null,
"is_final_message": true
},
// File response message
{
"id": file_id.to_string(),
"type": "file",
"metadata": [
{
"status": "completed",
"message": format!("File {} completed", file_name),
"timestamp": timestamp
}
],
"file_name": file_name,
"file_type": file_type,
"version_number": version_number,
"filter_version_id": null
}
]);
// Create a Message object to insert
let message = Message {
id: message_id,
request_message: None, // Empty request message as per requirement
response_messages: response_messages,
reasoning: json!([]),
title: "Version Restoration".to_string(),
raw_llm_messages: Value::Array(raw_llm_messages.clone()),
final_reasoning_message: None,
chat_id: *chat_id,
created_at: now,
updated_at: now,
deleted_at: None,
created_by: user.id,
feedback: None,
};
// Insert the message
diesel::insert_into(messages::table)
.values(&message)
.execute(&mut conn)
.await?;
// Create the message-to-file association
let message_to_file = MessageToFile {
id: Uuid::new_v4(),
message_id: message_id,
file_id: file_id,
created_at: now,
updated_at: now,
deleted_at: None,
};
// Insert the message-to-file association into the database
insert_into(messages_to_files::table)
.values(&message_to_file)
.execute(&mut conn)
.await?;
// Return the updated chat with messages
get_chat_handler(chat_id, user, false).await
}
#[cfg(test)]
mod tests {
use super::*;
use crate::chats::get_chat_handler::get_chat_handler;
use database::enums::AssetType;
use middleware::test_utils::create_test_user;
use uuid::Uuid;
// This is a mock test to demonstrate the expected implementation
// Actual tests would need a test database with proper fixtures
#[tokio::test]
async fn test_restore_chat_handler() {
// Setup would include:
// 1. Create a test database connection
// 2. Create a test user, chat, and assets
// 3. Create version history for the assets
// Example test structure (not functional without test setup):
/*
// Create test user
let user = create_test_user();
// Create a test chat
let chat_id = Uuid::new_v4();
// Create a test dashboard with multiple versions
let dashboard_id = Uuid::new_v4();
// Create a restore request
let request = ChatRestoreRequest {
asset_id: dashboard_id,
asset_type: AssetType::DashboardFile,
version_number: 1, // Restore to version 1
};
// Call the handler
let result = restore_chat_handler(&chat_id, &user, request).await;
// Verify success
assert!(result.is_ok());
// Get the updated chat
let chat = result.unwrap();
// Verify messages were created
assert!(chat.messages.len() >= 2); // At least the restoration message and file message
// Verify one message contains restoration text
let has_restoration_message = chat.messages.values().any(|msg|
msg.message_type == "text" &&
msg.request_message.contains("by restoring version")
);
assert!(has_restoration_message);
// Verify one message is a file message
let has_file_message = chat.messages.values().any(|msg|
msg.message_type == "file" &&
msg.file_type == Some("dashboard".to_string())
);
assert!(has_file_message);
*/
}
}

1
api/libs/handlers/src/chats/streaming_parser.rs
View File

@ -277,7 +277,6 @@ impl StreamingParser {
file_type: file_type.clone(),
file_name: name.to_string(),
version_number: 1,
version_id: String::from("0203f597-5ec5-4fd8-86e2-8587fe1c23b6"),
status: "loading".to_string(),
file: BusterFileContent {
text: None,

2
api/libs/handlers/src/dashboards/update_dashboard_handler.rs
View File

@ -16,7 +16,7 @@ use uuid::Uuid;
use super::{get_dashboard_handler, BusterDashboardResponse, DashboardConfig};
#[derive(Debug, Serialize, Deserialize)]
#[derive(Debug, Serialize, Deserialize, Default)]
pub struct DashboardUpdateRequest {
/// New name for the dashboard (optional)
pub name: Option<String>,

8
api/libs/handlers/src/messages/types.rs
View File

@ -11,7 +11,7 @@ pub use message_feedback::MessageFeedback;
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ChatMessage {
pub id: Uuid,
pub request_message: ChatUserMessage,
pub request_message: Option<ChatUserMessage>,
pub response_message_ids: Vec<String>,
#[serde(default)]
pub response_messages: HashMap<String, Value>,
@ -39,12 +39,12 @@ impl ChatMessage {
) -> Self {
Self {
id: Uuid::new_v4(),
request_message: ChatUserMessage {
request_message: Some(ChatUserMessage {
request,
sender_id,
sender_name,
sender_avatar,
},
}),
response_message_ids: Vec::new(),
response_messages: HashMap::new(),
reasoning_message_ids: Vec::new(),
@ -56,7 +56,7 @@ impl ChatMessage {
pub fn new_with_messages(
id: Uuid,
request_message: ChatUserMessage,
request_message: Option<ChatUserMessage>,
response_messages: Vec<Value>,
reasoning_messages: Vec<Value>,
final_reasoning_message: Option<String>,

2
api/libs/handlers/src/metrics/update_metric_handler.rs
View File

@ -51,7 +51,7 @@ fn merge_json_objects(base: Value, update: Value) -> Result<Value> {
use crate::metrics::get_metric_handler::get_metric_handler;
use crate::metrics::types::BusterMetric;
#[derive(Debug, serde::Deserialize, serde::Serialize)]
#[derive(Debug, serde::Deserialize, serde::Serialize, Default)]
pub struct UpdateMetricRequest {
#[serde(alias = "title")]
pub name: Option<String>,

11
api/libs/streaming/src/parser.rs
View File

@ -419,16 +419,6 @@ impl StreamingParser {
.unwrap_or_default()
.to_string();
// Generate deterministic version ID
let version_id =
match self.generate_deterministic_uuid(&id, &file_name, &file_type) {
Ok(id) => id,
Err(e) => {
eprintln!("Failed to generate version ID: {}", e);
continue;
}
};
// Create file content
let file_content = FileContent {
text: Some(yml_content),
@ -442,7 +432,6 @@ impl StreamingParser {
file_type: file_type.clone(),
file_name: file_name.clone(),
version_number: 1,
version_id: version_id.to_string(),
status: "completed".to_string(),
file: file_content,
metadata: Some(vec![]),

1
api/libs/streaming/src/processors/create_dashboards_processor.rs
View File

@ -201,7 +201,6 @@ impl Processor for CreateDashboardsProcessor {
file_type: "dashboard".to_string(),
file_name: name.to_string(),
version_number: 1,
version_id: String::from("0203f597-5ec5-4fd8-86e2-8587fe1c23b6"),
status: "loading".to_string(),
file: FileContent {
text: None,

57
api/libs/streaming/src/processors/create_metrics_processor.rs
View File

@ -16,21 +16,26 @@ impl CreateMetricsProcessor {
}
/// Generate a deterministic UUID based on input parameters
fn generate_deterministic_uuid(&self, base_id: &str, name: &str, file_type: &str) -> Result<Uuid> {
fn generate_deterministic_uuid(
&self,
base_id: &str,
name: &str,
file_type: &str,
) -> Result<Uuid> {
use sha2::{Digest, Sha256};
// Create a deterministic string by combining the inputs
let combined = format!("{}:{}:{}", base_id, name, file_type);
// Hash the combined string
let mut hasher = Sha256::new();
hasher.update(combined.as_bytes());
let result = hasher.finalize();
// Use the first 16 bytes of the hash as the UUID
let mut bytes = [0u8; 16];
bytes.copy_from_slice(&result[0..16]);
Ok(Uuid::from_bytes(bytes))
}
}
@ -61,7 +66,12 @@ impl Processor for CreateMetricsProcessor {
self.process_with_context(id, json, None)
}
fn process_with_context(&self, id: String, json: &str, previous_output: Option<ProcessedOutput>) -> Result<Option<ProcessedOutput>> {
fn process_with_context(
&self,
id: String,
json: &str,
previous_output: Option<ProcessedOutput>,
) -> Result<Option<ProcessedOutput>> {
// Try to parse the JSON
if let Ok(value) = serde_json::from_str::<Value>(json) {
// Check if it's a metric file structure
@ -90,11 +100,12 @@ impl Processor for CreateMetricsProcessor {
let file_id_str = file_id.to_string();
// Get the previously processed content for this file
let previous_content = if let Some(prev_file) = previous_files.get(&file_id_str) {
prev_file.file.text_chunk.clone().unwrap_or_default()
} else {
String::new()
};
let previous_content =
if let Some(prev_file) = previous_files.get(&file_id_str) {
prev_file.file.text_chunk.clone().unwrap_or_default()
} else {
String::new()
};
// Calculate the new content (what wasn't in the previous content)
let new_content = if yml_content.len() > previous_content.len() {
@ -112,11 +123,14 @@ impl Processor for CreateMetricsProcessor {
file_type: "metric".to_string(),
file_name: name.to_string(),
version_number: 1,
version_id: String::from("0203f597-5ec5-4fd8-86e2-8587fe1c23b6"),
status: "loading".to_string(),
file: FileContent {
text: None,
text_chunk: if new_content.is_empty() { None } else { Some(new_content) },
text_chunk: if new_content.is_empty() {
None
} else {
Some(new_content)
},
modified: None,
},
metadata: None,
@ -196,13 +210,16 @@ mod tests {
assert_eq!(file_output.id, id);
assert_eq!(file_output.title, "Creating metric files...");
assert_eq!(file_output.files.len(), 1);
// Check the first file
let file_id = &file_output.file_ids[0];
let file = file_output.files.get(file_id).unwrap();
assert_eq!(file.file_name, "test_metric.yml");
assert_eq!(file.file.text, None);
assert_eq!(file.file.text_chunk, Some("name: Test Metric\ndescription: A test metric".to_string()));
assert_eq!(
file.file.text_chunk,
Some("name: Test Metric\ndescription: A test metric".to_string())
);
} else {
panic!("Expected ProcessedOutput::File");
}
@ -240,7 +257,10 @@ mod tests {
let file_id = &file_output.file_ids[0];
let file = file_output.files.get(file_id).unwrap();
assert_eq!(file.file.text, None);
assert_eq!(file.file.text_chunk, Some("name: Test Metric\n".to_string()));
assert_eq!(
file.file.text_chunk,
Some("name: Test Metric\n".to_string())
);
} else {
panic!("Expected ProcessedOutput::File");
}
@ -256,7 +276,10 @@ mod tests {
let file_id = &file_output.file_ids[0];
let file = file_output.files.get(file_id).unwrap();
assert_eq!(file.file.text, None);
assert_eq!(file.file.text_chunk, Some("description: A test metric".to_string()));
assert_eq!(
file.file.text_chunk,
Some("description: A test metric".to_string())
);
} else {
panic!("Expected ProcessedOutput::File");
}

1
api/libs/streaming/src/types.rs
View File

@ -77,7 +77,6 @@ pub struct File {
pub file_type: String,
pub file_name: String,
pub version_number: i32,
pub version_id: String,
pub status: String,
pub file: FileContent,
pub metadata: Option<Vec<FileMetadata>>,

4
api/migrations/2025-03-25-184029_make_request_message_and_final_reasoning_nullable/down.sql Normal file
View File

@ -0,0 +1,4 @@
-- This file should undo anything in `up.sql`
ALTER TABLE messages
ALTER COLUMN final_reasoning_message SET NOT NULL,
ALTER COLUMN request_message SET NOT NULL;

4
api/migrations/2025-03-25-184029_make_request_message_and_final_reasoning_nullable/up.sql Normal file
View File

@ -0,0 +1,4 @@
-- Your SQL goes here
ALTER TABLE messages
ALTER COLUMN final_reasoning_message DROP NOT NULL,
ALTER COLUMN request_message DROP NOT NULL;

3
api/src/routes/rest/routes/chats/mod.rs
View File

@ -8,6 +8,7 @@ mod get_chat;
mod get_chat_raw_llm_messages;
mod list_chats;
mod post_chat;
mod restore_chat;
mod sharing;
mod update_chat;
mod update_chats;
@ -17,6 +18,7 @@ pub use get_chat::get_chat_route;
pub use get_chat_raw_llm_messages::get_chat_raw_llm_messages;
pub use list_chats::list_chats_route;
pub use post_chat::post_chat_route;
pub use restore_chat::restore_chat_route;
pub use update_chat::update_chat_route;
pub use update_chats::update_chats_route;
@ -28,6 +30,7 @@ pub fn router() -> Router {
.route("/", delete(delete_chats_route))
.route("/:id", get(get_chat_route))
.route("/:id", put(update_chat_route))
.route("/:id/restore", put(restore_chat_route))
.route("/:id/raw_llm_messages", get(get_chat_raw_llm_messages))
.nest("/:id/sharing", sharing::router())
}

60
api/src/routes/rest/routes/chats/restore_chat.rs Normal file
View File

@ -0,0 +1,60 @@
use anyhow::Result;
use axum::extract::Path;
use axum::http::StatusCode;
use axum::Extension;
use axum::Json;
use handlers::chats::restore_chat_handler::{restore_chat_handler, ChatRestoreRequest};
use handlers::chats::types::ChatWithMessages;
use middleware::AuthenticatedUser;
use uuid::Uuid;
use crate::routes::rest::ApiResponse;
/// Restore a previous version of an asset in a chat
///
/// PUT /api/v1/chats/:id/restore
///
/// This endpoint allows restoring a previous version of a metric or dashboard file
/// and documenting the restoration in the chat history. It creates:
/// 1. A text message noting which version was restored
/// 2. A file message referencing the newly created version
pub async fn restore_chat_route(
Path(chat_id): Path<Uuid>,
Extension(user): Extension<AuthenticatedUser>,
Json(request): Json<ChatRestoreRequest>,
) -> Result<ApiResponse<ChatWithMessages>, (StatusCode, String)> {
match restore_chat_handler(&chat_id, &user, request).await {
Ok(chat) => Ok(ApiResponse::JsonData(chat)),
Err(e) => {
let error_message = e.to_string();
tracing::error!("Error restoring asset in chat {}: {}", chat_id, error_message);
// Map specific error messages to appropriate HTTP status codes
if error_message.contains("not found") || error_message.contains("Version") {
Err((StatusCode::NOT_FOUND, error_message))
} else if error_message.contains("permission") {
Err((StatusCode::FORBIDDEN, error_message))
} else if error_message.contains("Unsupported asset type") {
Err((StatusCode::BAD_REQUEST, error_message))
} else {
Err((StatusCode::INTERNAL_SERVER_ERROR, "Failed to restore version".to_string()))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
// Unit tests for the restore_chat_route function
// Integration tests are in tests/integration/chats/restore_chat_test.rs
#[tokio::test]
async fn test_restore_chat_route_error_handling() {
// Simple test to verify error handling in the route function
// This can be expanded with mock handlers in a real implementation
// For now, just ensure the function exists and compiles
assert!(true, "The restore_chat_route function is defined");
}
}

3
api/tests/integration/chats/mod.rs
View File

@ -1,4 +1,5 @@
pub mod sharing;
pub mod get_chat_test;
pub mod update_chat_test;
pub mod post_chat_test;
pub mod post_chat_test;
pub mod restore_chat_test;

352
api/tests/integration/chats/restore_chat_test.rs Normal file
View File

@ -0,0 +1,352 @@
use anyhow::Result;
use database::enums::AssetType;
use uuid::Uuid;
use chrono::Utc;
use serde_json::json;
use crate::common::{
assertions::response::ResponseAssertion,
fixtures::chats::ChatFixtureBuilder,
fixtures::dashboards::DashboardFixtureBuilder,
fixtures::metrics::MetricFixtureBuilder,
fixtures::users::UserFixtureBuilder,
http::client::TestClient,
};
#[tokio::test]
async fn test_restore_metric_in_chat() -> Result<()> {
// Create test client
let client = TestClient::new().await?;
// Create test user
let user = UserFixtureBuilder::new()
.with_name("Test User")
.with_email("test@example.com")
.create(&client)
.await?;
// Create a metric with initial content
let metric = MetricFixtureBuilder::new()
.with_name("Test Metric")
.with_created_by(user.id)
.with_organization_id(user.organization_id)
.with_sql("SELECT * FROM test_table")
.create(&client)
.await?;
// Create a chat
let chat = ChatFixtureBuilder::new()
.with_title("Test Chat")
.with_created_by(user.id)
.with_organization_id(user.organization_id)
.create(&client)
.await?;
// Update the metric to create version 2
let update_response = client
.put(&format!("/api/v1/metrics/{}", metric.id))
.json(&json!({
"name": "Updated Metric",
"sql": "SELECT * FROM updated_table"
}))
.with_auth(&user)
.send()
.await;
update_response.assert_status_ok()?;
// Now restore the metric to version 1 via the chat restoration endpoint
let restore_response = client
.put(&format!("/api/v1/chats/{}/restore", chat.id))
.json(&json!({
"asset_id": metric.id,
"asset_type": "metric_file",
"version_number": 1
}))
.with_auth(&user)
.send()
.await;
restore_response.assert_status_ok()?;
// Extract the updated chat from the response
let chat_with_messages = restore_response.json::<serde_json::Value>()?;
// Verify that messages were created in the chat
let messages = chat_with_messages["data"]["messages"].as_object().unwrap();
assert!(messages.len() >= 2, "Expected at least 2 messages in the chat");
// Verify that there's a message with restoration text in the response_messages
let has_restoration_message = messages.values().any(|msg| {
let response_messages = msg["response_messages"].as_array().unwrap_or(&Vec::new());
response_messages.iter().any(|rm| {
rm["type"].as_str() == Some("text") &&
rm["message"].as_str().unwrap_or("").contains("was created by restoring")
})
});
assert!(has_restoration_message, "Expected a restoration message in the chat");
// Verify that there's a file reference in the response_messages
let has_file_message = messages.values().any(|msg| {
let response_messages = msg["response_messages"].as_array().unwrap_or(&Vec::new());
response_messages.iter().any(|rm| {
rm["type"].as_str() == Some("file") &&
rm["file_type"].as_str() == Some("metric")
})
});
assert!(has_file_message, "Expected a file message referencing the metric");
// Get the metric to verify that a new version was created
let metric_response = client
.get(&format!("/api/v1/metrics/{}", metric.id))
.with_auth(&user)
.send()
.await;
metric_response.assert_status_ok()?;
let metric_data = metric_response.json::<serde_json::Value>()?;
let version = metric_data["data"]["version"].as_i64().unwrap();
// The version should be 3 (initial + update + restore)
assert_eq!(version, 3, "Expected version to be 3 after restoration");
Ok(())
}
#[tokio::test]
async fn test_restore_dashboard_in_chat() -> Result<()> {
// Create test client
let client = TestClient::new().await?;
// Create test user
let user = UserFixtureBuilder::new()
.with_name("Test User")
.with_email("test2@example.com")
.create(&client)
.await?;
// Create a dashboard with initial content
let dashboard = DashboardFixtureBuilder::new()
.with_name("Test Dashboard")
.with_created_by(user.id)
.with_organization_id(user.organization_id)
.create(&client)
.await?;
// Create a chat
let chat = ChatFixtureBuilder::new()
.with_title("Test Chat for Dashboard")
.with_created_by(user.id)
.with_organization_id(user.organization_id)
.create(&client)
.await?;
// Update the dashboard to create version 2
let update_response = client
.put(&format!("/api/v1/dashboards/{}", dashboard.id))
.json(&json!({
"name": "Updated Dashboard"
}))
.with_auth(&user)
.send()
.await;
update_response.assert_status_ok()?;
// Now restore the dashboard to version 1 via the chat restoration endpoint
let restore_response = client
.put(&format!("/api/v1/chats/{}/restore", chat.id))
.json(&json!({
"asset_id": dashboard.id,
"asset_type": "dashboard_file",
"version_number": 1
}))
.with_auth(&user)
.send()
.await;
restore_response.assert_status_ok()?;
// Extract the updated chat from the response
let chat_with_messages = restore_response.json::<serde_json::Value>()?;
// Verify that messages were created in the chat
let messages = chat_with_messages["data"]["messages"].as_object().unwrap();
assert!(messages.len() >= 2, "Expected at least 2 messages in the chat");
// Verify that there's a message with restoration text in the response_messages
let has_restoration_message = messages.values().any(|msg| {
let response_messages = msg["response_messages"].as_array().unwrap_or(&Vec::new());
response_messages.iter().any(|rm| {
rm["type"].as_str() == Some("text") &&
rm["message"].as_str().unwrap_or("").contains("was created by restoring")
})
});
assert!(has_restoration_message, "Expected a restoration message in the chat");
// Verify that there's a file reference in the response_messages
let has_file_message = messages.values().any(|msg| {
let response_messages = msg["response_messages"].as_array().unwrap_or(&Vec::new());
response_messages.iter().any(|rm| {
rm["type"].as_str() == Some("file") &&
rm["file_type"].as_str() == Some("dashboard")
})
});
assert!(has_file_message, "Expected a file message referencing the dashboard");
// Get the dashboard to verify that a new version was created
let dashboard_response = client
.get(&format!("/api/v1/dashboards/{}", dashboard.id))
.with_auth(&user)
.send()
.await;
dashboard_response.assert_status_ok()?;
let dashboard_data = dashboard_response.json::<serde_json::Value>()?;
let version = dashboard_data["data"]["dashboard"]["version"].as_i64().unwrap();
// The version should be 3 (initial + update + restore)
assert_eq!(version, 3, "Expected version to be 3 after restoration");
Ok(())
}
#[tokio::test]
async fn test_restore_wrong_version_in_chat() -> Result<()> {
// Create test client
let client = TestClient::new().await?;
// Create test user
let user = UserFixtureBuilder::new()
.with_name("Test User")
.with_email("test3@example.com")
.create(&client)
.await?;
// Create a metric
let metric = MetricFixtureBuilder::new()
.with_name("Test Metric")
.with_created_by(user.id)
.with_organization_id(user.organization_id)
.with_sql("SELECT * FROM test_table")
.create(&client)
.await?;
// Create a chat
let chat = ChatFixtureBuilder::new()
.with_title("Test Chat for Error Case")
.with_created_by(user.id)
.with_organization_id(user.organization_id)
.create(&client)
.await?;
// Try to restore a non-existent version (version 999)
let restore_response = client
.put(&format!("/api/v1/chats/{}/restore", chat.id))
.json(&json!({
"asset_id": metric.id,
"asset_type": "metric_file",
"version_number": 999
}))
.with_auth(&user)
.send()
.await;
// This should fail with a 404 Not Found
assert_eq!(restore_response.status().as_u16(), 404);
Ok(())
}
#[tokio::test]
async fn test_restore_invalid_asset_type_in_chat() -> Result<()> {
// Create test client
let client = TestClient::new().await?;
// Create test user
let user = UserFixtureBuilder::new()
.with_name("Test User")
.with_email("test4@example.com")
.create(&client)
.await?;
// Create a chat
let chat = ChatFixtureBuilder::new()
.with_title("Test Chat for Invalid Asset Type")
.with_created_by(user.id)
.with_organization_id(user.organization_id)
.create(&client)
.await?;
// Try to restore with an invalid asset type
let restore_response = client
.put(&format!("/api/v1/chats/{}/restore", chat.id))
.json(&json!({
"asset_id": Uuid::new_v4(),
"asset_type": "chat", // This is invalid for restoration
"version_number": 1
}))
.with_auth(&user)
.send()
.await;
// This should fail with a 400 Bad Request
assert_eq!(restore_response.status().as_u16(), 400);
Ok(())
}
#[tokio::test]
async fn test_restore_without_permission() -> Result<()> {
// Create test client
let client = TestClient::new().await?;
// Create test users
let owner = UserFixtureBuilder::new()
.with_name("Owner User")
.with_email("owner@example.com")
.create(&client)
.await?;
let other_user = UserFixtureBuilder::new()
.with_name("Other User")
.with_email("other@example.com")
.create(&client)
.await?;
// Create a metric owned by the owner
let metric = MetricFixtureBuilder::new()
.with_name("Owner's Metric")
.with_created_by(owner.id)
.with_organization_id(owner.organization_id)
.with_sql("SELECT * FROM owner_table")
.create(&client)
.await?;
// Create a chat owned by the owner
let chat = ChatFixtureBuilder::new()
.with_title("Owner's Chat")
.with_created_by(owner.id)
.with_organization_id(owner.organization_id)
.create(&client)
.await?;
// Try to restore as the other user who doesn't have permission
let restore_response = client
.put(&format!("/api/v1/chats/{}/restore", chat.id))
.json(&json!({
"asset_id": metric.id,
"asset_type": "metric_file",
"version_number": 1
}))
.with_auth(&other_user)
.send()
.await;
// This should fail with a 403 Forbidden or 404 Not Found
// (Depending on implementation, it may be Not Found if the user can't see the resources at all)
let status = restore_response.status().as_u16();
assert!(status == 403 || status == 404, "Expected status code 403 or 404, got {}", status);
Ok(())
}