diff --git a/api/libs/database/src/models.rs b/api/libs/database/src/models.rs index 6d2a9bbe2..790697b4b 100644 --- a/api/libs/database/src/models.rs +++ b/api/libs/database/src/models.rs @@ -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, 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, pub chat_id: Uuid, pub created_at: DateTime, pub updated_at: DateTime, diff --git a/api/libs/handlers/src/chats/asset_messages.rs b/api/libs/handlers/src/chats/asset_messages.rs index c51089e1c..ef5a06975 100644 --- a/api/libs/handlers/src/chats/asset_messages.rs +++ b/api/libs/handlers/src/chats/asset_messages.rs @@ -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, }; diff --git a/api/libs/handlers/src/chats/get_chat_handler.rs b/api/libs/handlers/src/chats/get_chat_handler.rs index 6ce43c7c2..f3b92d9e3 100644 --- a/api/libs/handlers/src/chats/get_chat_handler.rs +++ b/api/libs/handlers/src/chats/get_chat_handler.rs @@ -33,10 +33,10 @@ pub struct ChatWithUser { #[derive(Queryable)] pub struct MessageWithUser { pub id: Uuid, - pub request_message: String, + pub request_message: Option, pub response_messages: Value, pub reasoning: Value, - pub final_reasoning_message: String, + pub final_reasoning_message: Option, pub created_at: DateTime, pub user_id: Uuid, pub user_name: Option, @@ -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, ) }) diff --git a/api/libs/handlers/src/chats/mod.rs b/api/libs/handlers/src/chats/mod.rs index e6d0b5269..62f0e3b66 100644 --- a/api/libs/handlers/src/chats/mod.rs +++ b/api/libs/handlers/src/chats/mod.rs @@ -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; diff --git a/api/libs/handlers/src/chats/post_chat_handler.rs b/api/libs/handlers/src/chats/post_chat_handler.rs index cb85b8558..ff8633d38 100644 --- a/api/libs/handlers/src/chats/post_chat_handler.rs +++ b/api/libs/handlers/src/chats/post_chat_handler.rs @@ -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, metadata: Option>, }, @@ -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>, @@ -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 Result Result Result` - 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 { + 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::(&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::>(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); + */ + } +} diff --git a/api/libs/handlers/src/chats/streaming_parser.rs b/api/libs/handlers/src/chats/streaming_parser.rs index d3da52954..5b4ce3ea7 100644 --- a/api/libs/handlers/src/chats/streaming_parser.rs +++ b/api/libs/handlers/src/chats/streaming_parser.rs @@ -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, diff --git a/api/libs/handlers/src/dashboards/update_dashboard_handler.rs b/api/libs/handlers/src/dashboards/update_dashboard_handler.rs index 18c1a2b44..0977ae138 100644 --- a/api/libs/handlers/src/dashboards/update_dashboard_handler.rs +++ b/api/libs/handlers/src/dashboards/update_dashboard_handler.rs @@ -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, diff --git a/api/libs/handlers/src/messages/types.rs b/api/libs/handlers/src/messages/types.rs index fd4feddcd..77ccc9d2b 100644 --- a/api/libs/handlers/src/messages/types.rs +++ b/api/libs/handlers/src/messages/types.rs @@ -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, pub response_message_ids: Vec, #[serde(default)] pub response_messages: HashMap, @@ -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, response_messages: Vec, reasoning_messages: Vec, final_reasoning_message: Option, diff --git a/api/libs/handlers/src/metrics/update_metric_handler.rs b/api/libs/handlers/src/metrics/update_metric_handler.rs index 2d9ef51cc..7c0cad9d6 100644 --- a/api/libs/handlers/src/metrics/update_metric_handler.rs +++ b/api/libs/handlers/src/metrics/update_metric_handler.rs @@ -51,7 +51,7 @@ fn merge_json_objects(base: Value, update: Value) -> Result { 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, diff --git a/api/libs/streaming/src/parser.rs b/api/libs/streaming/src/parser.rs index 312f1c623..f30d1b1e7 100644 --- a/api/libs/streaming/src/parser.rs +++ b/api/libs/streaming/src/parser.rs @@ -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![]), diff --git a/api/libs/streaming/src/processors/create_dashboards_processor.rs b/api/libs/streaming/src/processors/create_dashboards_processor.rs index 839153fc3..60c82eab2 100644 --- a/api/libs/streaming/src/processors/create_dashboards_processor.rs +++ b/api/libs/streaming/src/processors/create_dashboards_processor.rs @@ -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, diff --git a/api/libs/streaming/src/processors/create_metrics_processor.rs b/api/libs/streaming/src/processors/create_metrics_processor.rs index 8192eb5b2..909833b3a 100644 --- a/api/libs/streaming/src/processors/create_metrics_processor.rs +++ b/api/libs/streaming/src/processors/create_metrics_processor.rs @@ -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 { + fn generate_deterministic_uuid( + &self, + base_id: &str, + name: &str, + file_type: &str, + ) -> Result { 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) -> Result> { + fn process_with_context( + &self, + id: String, + json: &str, + previous_output: Option, + ) -> Result> { // Try to parse the JSON if let Ok(value) = serde_json::from_str::(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"); } diff --git a/api/libs/streaming/src/types.rs b/api/libs/streaming/src/types.rs index 6b8c25fff..b88dfcece 100644 --- a/api/libs/streaming/src/types.rs +++ b/api/libs/streaming/src/types.rs @@ -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>, diff --git a/api/migrations/2025-03-25-184029_make_request_message_and_final_reasoning_nullable/down.sql b/api/migrations/2025-03-25-184029_make_request_message_and_final_reasoning_nullable/down.sql new file mode 100644 index 000000000..81e7f2ea4 --- /dev/null +++ b/api/migrations/2025-03-25-184029_make_request_message_and_final_reasoning_nullable/down.sql @@ -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; diff --git a/api/migrations/2025-03-25-184029_make_request_message_and_final_reasoning_nullable/up.sql b/api/migrations/2025-03-25-184029_make_request_message_and_final_reasoning_nullable/up.sql new file mode 100644 index 000000000..40e207584 --- /dev/null +++ b/api/migrations/2025-03-25-184029_make_request_message_and_final_reasoning_nullable/up.sql @@ -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; diff --git a/api/src/routes/rest/routes/chats/mod.rs b/api/src/routes/rest/routes/chats/mod.rs index ab83131dd..ab5b38b98 100644 --- a/api/src/routes/rest/routes/chats/mod.rs +++ b/api/src/routes/rest/routes/chats/mod.rs @@ -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()) } diff --git a/api/src/routes/rest/routes/chats/restore_chat.rs b/api/src/routes/rest/routes/chats/restore_chat.rs new file mode 100644 index 000000000..18fa38050 --- /dev/null +++ b/api/src/routes/rest/routes/chats/restore_chat.rs @@ -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, + Extension(user): Extension, + Json(request): Json, +) -> Result, (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"); + } +} \ No newline at end of file diff --git a/api/tests/integration/chats/mod.rs b/api/tests/integration/chats/mod.rs index 750c3b605..d15481158 100644 --- a/api/tests/integration/chats/mod.rs +++ b/api/tests/integration/chats/mod.rs @@ -1,4 +1,5 @@ pub mod sharing; pub mod get_chat_test; pub mod update_chat_test; -pub mod post_chat_test; \ No newline at end of file +pub mod post_chat_test; +pub mod restore_chat_test; \ No newline at end of file diff --git a/api/tests/integration/chats/restore_chat_test.rs b/api/tests/integration/chats/restore_chat_test.rs new file mode 100644 index 000000000..8fdcf389f --- /dev/null +++ b/api/tests/integration/chats/restore_chat_test.rs @@ -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::()?; + + // 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::()?; + 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::()?; + + // 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::()?; + 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(()) +} \ No newline at end of file