From 94c1635a348a72b4ab01076b8b5c9136c281a4f4 Mon Sep 17 00:00:00 2001 From: dal Date: Fri, 7 Feb 2025 01:29:49 -0700 Subject: [PATCH] refactor(file_tools): Enhance file modification and creation processes - Implement robust line-based content modification for metric and dashboard files - Add comprehensive error handling and validation for file modifications - Improve modification tracking with detailed modification results - Optimize file processing with batch insertion and validation - Add extensive test coverage for modification and validation logic --- .../utils/tools/file_tools/create_files.rs | 228 ++-- .../utils/tools/file_tools/modify_files.rs | 980 +++++++++++++++++- 2 files changed, 1069 insertions(+), 139 deletions(-) diff --git a/api/src/utils/tools/file_tools/create_files.rs b/api/src/utils/tools/file_tools/create_files.rs index fe3c56979..6b3df6e9a 100644 --- a/api/src/utils/tools/file_tools/create_files.rs +++ b/api/src/utils/tools/file_tools/create_files.rs @@ -62,38 +62,125 @@ impl ToolExecutor for CreateFilesTool { }; let files = params.files; - let mut created_files = vec![]; - let mut failed_files = vec![]; + // Separate files by type and validate/prepare them + let mut metric_records = vec![]; + let mut dashboard_records = vec![]; + let mut metric_ymls = vec![]; + let mut dashboard_ymls = vec![]; + + // First pass - validate and prepare all records for file in files { match file.file_type.as_str() { - "metric" => match create_metric_file(file.clone()).await { - Ok(f) => { - created_files.push(f); - continue; + "metric" => { + match MetricYml::new(file.yml_content.clone()) { + Ok(metric_yml) => { + if let Some(metric_id) = &metric_yml.id { + let metric_file = MetricFile { + id: metric_id.clone(), + name: metric_yml.title.clone(), + file_name: format!("{}.yml", file.name), + content: serde_json::to_value(metric_yml.clone()).unwrap(), + created_by: Uuid::new_v4(), + verification: Verification::NotRequested, + evaluation_obj: None, + evaluation_summary: None, + evaluation_score: None, + organization_id: Uuid::new_v4(), + created_at: Utc::now(), + updated_at: Utc::now(), + deleted_at: None, + }; + metric_records.push(metric_file); + metric_ymls.push(metric_yml); + } else { + failed_files.push((file.name, "Metric YML file does not have an id. This should never happen.".to_string())); + } + } + Err(e) => { + failed_files.push((file.name, e.to_string())); + } } - Err(e) => { - failed_files.push((file.name, e.to_string())); - continue; + } + "dashboard" => { + match DashboardYml::new(file.yml_content.clone()) { + Ok(dashboard_yml) => { + if let Some(dashboard_id) = &dashboard_yml.id { + let dashboard_file = DashboardFile { + id: dashboard_id.clone(), + name: dashboard_yml.name.clone().unwrap_or_else(|| "New Dashboard".to_string()), + file_name: format!("{}.yml", file.name), + content: serde_json::to_value(dashboard_yml.clone()).unwrap(), + filter: None, + organization_id: Uuid::new_v4(), + created_by: Uuid::new_v4(), + created_at: Utc::now(), + updated_at: Utc::now(), + deleted_at: None, + }; + dashboard_records.push(dashboard_file); + dashboard_ymls.push(dashboard_yml); + } else { + failed_files.push((file.name, "Dashboard YML file does not have an id. This should never happen.".to_string())); + } + } + Err(e) => { + failed_files.push((file.name, e.to_string())); + } } - }, - "dashboard" => match create_dashboard_file(file.clone()).await { - Ok(f) => { - created_files.push(f); - continue; - } - Err(e) => { - failed_files.push((file.name, e.to_string())); - continue; - } - }, + } _ => { failed_files.push((file.name, format!("Invalid file type: {}. Currently only `metric` and `dashboard` types are supported.", file.file_type))); - continue; } - }; + } + } + + // Second pass - bulk insert records + let mut conn = match get_pg_pool().get().await { + Ok(conn) => conn, + Err(e) => return Err(anyhow::anyhow!(e)), + }; + + // Insert metric files + if !metric_records.is_empty() { + match insert_into(metric_files::table) + .values(&metric_records) + .execute(&mut conn) + .await + { + Ok(_) => { + created_files.extend(metric_ymls.into_iter().map(FileEnum::Metric)); + } + Err(e) => { + failed_files.extend( + metric_records + .iter() + .map(|r| (r.file_name.clone(), format!("Failed to create metric file: {}", e))), + ); + } + } + } + + // Insert dashboard files + if !dashboard_records.is_empty() { + match insert_into(dashboard_files::table) + .values(&dashboard_records) + .execute(&mut conn) + .await + { + Ok(_) => { + created_files.extend(dashboard_ymls.into_iter().map(FileEnum::Dashboard)); + } + Err(e) => { + failed_files.extend( + dashboard_records + .iter() + .map(|r| (r.file_name.clone(), format!("Failed to create dashboard file: {}", e))), + ); + } + } } let message = if failed_files.is_empty() { @@ -164,103 +251,6 @@ impl ToolExecutor for CreateFilesTool { } } -async fn create_metric_file(file: FileParams) -> Result { - let metric_yml = match MetricYml::new(file.yml_content) { - Ok(metric_file) => metric_file, - Err(e) => return Err(e), - }; - - let mut conn = match get_pg_pool().get().await { - Ok(conn) => conn, - Err(e) => return Err(anyhow::anyhow!(e)), - }; - - let metric_id = match &metric_yml.id { - Some(id) => id, - None => { - return Err(anyhow::anyhow!( - "Metric YML file does not have an id. This should never happen." - )) - } - }; - - let metric_file_record = MetricFile { - id: metric_id.clone(), - name: metric_yml.title.clone(), - file_name: format!("{}.yml", file.name), - content: serde_json::to_value(metric_yml.clone()).unwrap(), - created_by: Uuid::new_v4(), - verification: Verification::NotRequested, - evaluation_obj: None, - evaluation_summary: None, - evaluation_score: None, - organization_id: Uuid::new_v4(), - created_at: Utc::now(), - updated_at: Utc::now(), - deleted_at: None, - }; - - match insert_into(metric_files::table) - .values(&metric_file_record) - .execute(&mut conn) - .await - { - Ok(metric_file_record) => metric_file_record, - Err(e) => return Err(anyhow::anyhow!("Failed to create metric file: {}", e)), - }; - - Ok(FileEnum::Metric(metric_yml)) -} - -async fn create_dashboard_file(file: FileParams) -> Result { - let dashboard_yml = match DashboardYml::new(file.yml_content) { - Ok(dashboard_file) => dashboard_file, - Err(e) => return Err(e), - }; - - let mut conn = match get_pg_pool().get().await { - Ok(conn) => conn, - Err(e) => return Err(anyhow::anyhow!(e)), - }; - - let dashboard_id = match &dashboard_yml.id { - Some(id) => id, - None => { - return Err(anyhow::anyhow!( - "Dashboard YML file does not have an id. This should never happen." - )) - } - }; - - let dashboard_file_record = DashboardFile { - id: dashboard_id.clone(), - name: dashboard_yml - .name - .clone() - .unwrap_or_else(|| "New Dashboard".to_string()), - file_name: format!("{}.yml", file.name), - content: serde_json::to_value(dashboard_yml.clone()).unwrap(), - filter: None, - organization_id: Uuid::new_v4(), - created_by: Uuid::new_v4(), - created_at: Utc::now(), - updated_at: Utc::now(), - deleted_at: None, - }; - - match insert_into(dashboard_files::table) - .values(&dashboard_file_record) - .returning(dashboard_files::all_columns) - .execute(&mut conn) - .await - { - Ok(_) => (), - Err(e) => return Err(anyhow::anyhow!(e)), - }; - - Ok(FileEnum::Dashboard(dashboard_yml)) -} - #[cfg(test)] mod tests { use super::*; diff --git a/api/src/utils/tools/file_tools/modify_files.rs b/api/src/utils/tools/file_tools/modify_files.rs index a8a5b3129..f7d66c0c6 100644 --- a/api/src/utils/tools/file_tools/modify_files.rs +++ b/api/src/utils/tools/file_tools/modify_files.rs @@ -1,12 +1,32 @@ use anyhow::Result; use async_trait::async_trait; +use chrono::Utc; +use diesel::{upsert::excluded, ExpressionMethods, QueryDsl}; +use diesel_async::RunQueryDsl; use serde_json::Value; use serde::{Deserialize, Serialize}; +use uuid::Uuid; +use tracing::{debug, error, info, warn}; -use crate::utils::{clients::ai::litellm::ToolCall, tools::ToolExecutor}; -use super::FileModificationTool; +use crate::{ + database::{ + enums::Verification, + lib::get_pg_pool, + models::{DashboardFile, MetricFile}, + schema::{dashboard_files, metric_files}, + }, + utils::{clients::ai::litellm::ToolCall, tools::ToolExecutor}, +}; +use super::{ + file_types::{ + dashboard_yml::DashboardYml, + file::FileEnum, + metric_yml::MetricYml, + }, + FileModificationTool, +}; -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Clone)] struct Modification { new_content: String, line_numbers: Vec, @@ -14,18 +34,190 @@ struct Modification { #[derive(Debug, Serialize, Deserialize)] struct FileModification { - file: String, + id: Uuid, + file_type: String, + file_name: String, modifications: Vec, } #[derive(Debug, Serialize, Deserialize)] struct ModifyFilesParams { - files_with_modifications: Vec, + files: Vec, +} + +#[derive(Debug, Serialize)] +struct ModificationResult { + file_id: Uuid, + file_type: String, + file_name: String, + success: bool, + original_lines: Vec, + adjusted_lines: Vec, + error: Option, + modification_type: String, + timestamp: chrono::DateTime, +} + +#[derive(Debug)] +struct FileModificationBatch { + metric_files: Vec, + dashboard_files: Vec, + metric_ymls: Vec, + dashboard_ymls: Vec, + failed_modifications: Vec<(String, String)>, + modification_results: Vec, +} + +#[derive(Debug)] +struct LineAdjustment { + original_start: i64, + original_end: i64, + new_length: i64, + offset: i64, +} + +impl LineAdjustment { + fn new(original_start: i64, original_end: i64, new_length: i64) -> Self { + let original_length = original_end - original_start + 1; + let offset = new_length - original_length; + + Self { + original_start, + original_end, + new_length, + offset, + } + } + + fn adjust_line_number(&self, line: i64) -> i64 { + if line < self.original_start { + line + } else { + line + self.offset + } + } +} + +// Helper functions for line number validation and modification +fn validate_line_numbers(line_numbers: &[i64]) -> Result<()> { + // Check if empty + if line_numbers.is_empty() { + return Err(anyhow::anyhow!("Line numbers array cannot be empty")); + } + + // Check if we have exactly 2 numbers for the range + if line_numbers.len() != 2 { + return Err(anyhow::anyhow!( + "Line numbers must specify a range with exactly 2 numbers [start,end], got {} numbers", + line_numbers.len() + )); + } + + let start = line_numbers[0]; + let end = line_numbers[1]; + + // Check if starts with at least 1 + if start < 1 { + return Err(anyhow::anyhow!( + "Line numbers must be 1-indexed, got starting line {}", + start + )); + } + + // Check if end is greater than or equal to start + if end < start { + return Err(anyhow::anyhow!( + "End line {} must be greater than or equal to start line {}", + end, + start + )); + } + + Ok(()) +} + +// Helper function to expand range into sequential numbers +fn expand_line_range(line_numbers: &[i64]) -> Vec { + if line_numbers.len() != 2 { + return line_numbers.to_vec(); + } + let start = line_numbers[0]; + let end = line_numbers[1]; + (start..=end).collect() +} + +fn apply_modifications_to_content(content: &str, modifications: &[Modification], file_name: &str) -> Result { + let mut lines: Vec<&str> = content.lines().collect(); + let mut modified_lines = lines.clone(); + let mut total_offset = 0; + + // Validate and sort modifications by starting line number + let mut sorted_modifications = modifications.to_vec(); + sorted_modifications.sort_by_key(|m| m.line_numbers[0]); + + // Check for overlapping modifications + for window in sorted_modifications.windows(2) { + let first_end = window[0].line_numbers[1]; + let second_start = window[1].line_numbers[0]; + if second_start <= first_end { + return Err(anyhow::anyhow!( + "Overlapping modifications in file '{}': line {} overlaps with line {}", + file_name, + first_end, + second_start + )); + } + } + + // Apply modifications and track adjustments + for modification in &sorted_modifications { + // Validate line numbers + validate_line_numbers(&modification.line_numbers)?; + + // Expand range into sequential numbers for processing + let line_range = expand_line_range(&modification.line_numbers); + + // Adjust line numbers based on previous modifications + let original_start = line_range[0] as usize - 1; + let original_end = line_range[line_range.len() - 1] as usize - 1; + let adjusted_start = (original_start as i64 + total_offset) as usize; + + // Validate line numbers are within bounds + if original_end >= lines.len() { + return Err(anyhow::anyhow!( + "Line numbers out of bounds in file '{}': file has {} lines, but modification attempts to modify line {}", + file_name, + lines.len(), + original_end + 1 + )); + } + + // Split new content into lines + let new_lines: Vec<&str> = modification.new_content.lines().collect(); + + // Calculate the change in number of lines + let old_length = original_end - original_start + 1; + let new_length = new_lines.len(); + total_offset += new_length as i64 - old_length as i64; + + // Apply the modification + let prefix = modified_lines[..adjusted_start].to_vec(); + let suffix = modified_lines[adjusted_start + old_length..].to_vec(); + + modified_lines = [ + prefix, + new_lines, + suffix, + ].concat(); + } + + Ok(modified_lines.join("\n")) } #[derive(Debug, Serialize)] pub struct ModifyFilesOutput { - success: bool, + message: String, + files: Vec, } pub struct ModifyFilesTool; @@ -41,11 +233,218 @@ impl ToolExecutor for ModifyFilesTool { } async fn execute(&self, tool_call: &ToolCall) -> Result { - let params: ModifyFilesParams = serde_json::from_str(&tool_call.function.arguments.clone())?; - // TODO: Implement actual file modification logic - let output = ModifyFilesOutput { - success: true, + debug!("Starting file modification execution"); + let params: ModifyFilesParams = match serde_json::from_str(&tool_call.function.arguments.clone()) { + Ok(params) => params, + Err(e) => { + return Ok(ModifyFilesOutput { + message: format!("Failed to parse modification parameters: {}", e), + files: Vec::new(), + }); + } }; + + info!("Processing {} files for modification", params.files.len()); + + // Initialize batch processing structures + let mut batch = FileModificationBatch { + metric_files: Vec::new(), + dashboard_files: Vec::new(), + metric_ymls: Vec::new(), + dashboard_ymls: Vec::new(), + failed_modifications: Vec::new(), + modification_results: Vec::new(), + }; + + // Group files by type and fetch existing records + let mut metric_ids = Vec::new(); + let mut dashboard_ids = Vec::new(); + let mut file_map = std::collections::HashMap::new(); + + for file in ¶ms.files { + match file.file_type.as_str() { + "metric" => metric_ids.push(file.id), + "dashboard" => dashboard_ids.push(file.id), + _ => { + batch.failed_modifications.push(( + file.file_name.clone(), + format!("Invalid file type: {}", file.file_type), + )); + continue; + } + } + file_map.insert(file.id, file); + } + + // Get database connection + let mut conn = match get_pg_pool().get().await { + Ok(conn) => conn, + Err(e) => { + return Ok(ModifyFilesOutput { + message: format!("Failed to connect to database: {}", e), + files: Vec::new(), + }); + } + }; + + // Fetch metric files + if !metric_ids.is_empty() { + use crate::database::schema::metric_files::dsl::*; + match metric_files + .filter(id.eq_any(metric_ids)) + .filter(deleted_at.is_null()) + .load::(&mut conn) + .await + { + Ok(files) => { + for file in files { + if let Some(modifications) = file_map.get(&file.id) { + match process_metric_file(file, modifications).await { + Ok((metric_file, metric_yml, results)) => { + batch.metric_files.push(metric_file); + batch.metric_ymls.push(metric_yml); + batch.modification_results.extend(results); + } + Err(e) => { + batch.failed_modifications.push(( + modifications.file_name.clone(), + e.to_string(), + )); + } + } + } + } + } + Err(e) => { + return Ok(ModifyFilesOutput { + message: format!("Failed to fetch metric files: {}", e), + files: Vec::new(), + }); + } + } + } + + // Fetch dashboard files + if !dashboard_ids.is_empty() { + use crate::database::schema::dashboard_files::dsl::*; + match dashboard_files + .filter(id.eq_any(dashboard_ids)) + .filter(deleted_at.is_null()) + .load::(&mut conn) + .await + { + Ok(files) => { + for file in files { + if let Some(modifications) = file_map.get(&file.id) { + match process_dashboard_file(file, modifications).await { + Ok((dashboard_file, dashboard_yml, results)) => { + batch.dashboard_files.push(dashboard_file); + batch.dashboard_ymls.push(dashboard_yml); + batch.modification_results.extend(results); + } + Err(e) => { + batch.failed_modifications.push(( + modifications.file_name.clone(), + e.to_string(), + )); + } + } + } + } + } + Err(e) => { + return Ok(ModifyFilesOutput { + message: format!("Failed to fetch dashboard files: {}", e), + files: Vec::new(), + }); + } + } + } + + // Process results and generate output message + let mut output = ModifyFilesOutput { + message: String::new(), + files: Vec::new(), + }; + + // Update metric files in database + if !batch.metric_files.is_empty() { + use diesel::insert_into; + match insert_into(metric_files::table) + .values(&batch.metric_files) + .on_conflict(metric_files::id) + .do_update() + .set(( + metric_files::content.eq(excluded(metric_files::content)), + metric_files::updated_at.eq(Utc::now()), + metric_files::verification.eq(Verification::NotRequested), + )) + .execute(&mut conn) + .await + { + Ok(_) => { + output.files.extend(batch.metric_ymls.into_iter().map(FileEnum::Metric)); + } + Err(e) => { + batch.failed_modifications.push(( + "metric_files".to_string(), + format!("Failed to update metric files: {}", e), + )); + } + } + } + + // Update dashboard files in database + if !batch.dashboard_files.is_empty() { + use diesel::insert_into; + match insert_into(dashboard_files::table) + .values(&batch.dashboard_files) + .on_conflict(dashboard_files::id) + .do_update() + .set(( + dashboard_files::content.eq(excluded(dashboard_files::content)), + dashboard_files::updated_at.eq(Utc::now()), + )) + .execute(&mut conn) + .await + { + Ok(_) => { + output + .files + .extend(batch.dashboard_ymls.into_iter().map(FileEnum::Dashboard)); + } + Err(e) => { + batch.failed_modifications.push(( + "dashboard_files".to_string(), + format!("Failed to update dashboard files: {}", e), + )); + } + } + } + + // Generate final message + if batch.failed_modifications.is_empty() { + output.message = format!("Successfully modified {} files.", output.files.len()); + } else { + let success_msg = if !output.files.is_empty() { + format!("Successfully modified {} files. ", output.files.len()) + } else { + String::new() + }; + + let failures: Vec = batch + .failed_modifications + .iter() + .map(|(name, error)| format!("Failed to modify '{}': {}", name, error)) + .collect(); + + output.message = format!( + "{}Failed to modify {} files: {}", + success_msg, + batch.failed_modifications.len(), + failures.join("; ") + ); + } Ok(output) } @@ -56,17 +455,26 @@ impl ToolExecutor for ModifyFilesTool { "strict": true, "parameters": { "type": "object", - "required": ["files_with_modifications"], + "required": ["files"], "properties": { - "files_with_modifications": { + "files": { "type": "array", "items": { "type": "object", - "required": ["file", "modifications"], + "required": ["id", "file_type", "file_name", "modifications"], "properties": { - "file": { + "id": { "type": "string", - "description": "The path to a yml file that needs to be modified." + "description": "The UUID of the file to modify" + }, + "file_type": { + "type": "string", + "enum": ["metric", "dashboard"], + "description": "The type of file to modify" + }, + "file_name": { + "type": "string", + "description": "The name of the file being modified" }, "modifications": { "type": "array", @@ -81,10 +489,11 @@ impl ToolExecutor for ModifyFilesTool { "line_numbers": { "type": "array", "items": { - "type": "number", - "description": "Line numbers in the yml file." + "type": "number" }, - "description": "Array of line numbers that need to be replaced, includes start and end line. If continuous lines are being edited please keep them together. i.e. [20, 34]" + "minItems": 2, + "maxItems": 2, + "description": "Array containing exactly 2 numbers [start,end] specifying the range of lines to replace. For example, [1,5] replaces lines 1 through 5. For a single line, use [n,n] (e.g., [3,3] for line 3)." } }, "additionalProperties": false @@ -94,12 +503,543 @@ impl ToolExecutor for ModifyFilesTool { }, "additionalProperties": false }, - "description": "Array of objects containing file paths and their respective modifications." + "description": "Array of files to modify with their modifications." } }, "additionalProperties": false }, - "description": "Makes multiple line-level modifications to one or more existing YAML files in a single call. If you need to update SQL, chart config, or other sections within a file, use this." + "description": "Makes multiple line-level modifications to one or more existing YAML files in a single call. Line numbers are specified as [start,end] ranges. If you need to update SQL, chart config, or other sections within a file, use this." }) } +} + +async fn process_metric_file( + mut file: MetricFile, + modification: &FileModification, +) -> Result<(MetricFile, MetricYml, Vec)> { + debug!( + file_id = %file.id, + file_name = %modification.file_name, + "Processing metric file modifications" + ); + + let mut results = Vec::new(); + + // Parse existing content + let current_yml: MetricYml = match serde_json::from_value(file.content.clone()) { + Ok(yml) => yml, + Err(e) => { + let error = format!("Failed to parse existing metric YAML: {}", e); + error!( + file_id = %file.id, + file_name = %modification.file_name, + error = %error, + "YAML parsing error" + ); + results.push(ModificationResult { + file_id: file.id, + file_type: "metric".to_string(), + file_name: modification.file_name.clone(), + success: false, + original_lines: vec![], + adjusted_lines: vec![], + error: Some(error.clone()), + modification_type: "parsing".to_string(), + timestamp: Utc::now(), + }); + return Err(anyhow::anyhow!(error)); + } + }; + + // Convert to YAML string for line modifications + let current_content = match serde_yaml::to_string(¤t_yml) { + Ok(content) => content, + Err(e) => { + let error = format!("Failed to serialize metric YAML: {}", e); + error!( + file_id = %file.id, + file_name = %modification.file_name, + error = %error, + "YAML serialization error" + ); + results.push(ModificationResult { + file_id: file.id, + file_type: "metric".to_string(), + file_name: modification.file_name.clone(), + success: false, + original_lines: vec![], + adjusted_lines: vec![], + error: Some(error.clone()), + modification_type: "serialization".to_string(), + timestamp: Utc::now(), + }); + return Err(anyhow::anyhow!(error)); + } + }; + + // Track original line numbers before modifications + let mut original_lines = Vec::new(); + let mut adjusted_lines = Vec::new(); + + // Apply modifications + for m in &modification.modifications { + original_lines.extend(m.line_numbers.clone()); + } + + // Apply modifications and track results + match apply_modifications_to_content(¤t_content, &modification.modifications, &modification.file_name) { + Ok(modified_content) => { + // Create and validate new YML object + match MetricYml::new(modified_content) { + Ok(new_yml) => { + debug!( + file_id = %file.id, + file_name = %modification.file_name, + "Successfully modified and validated metric file" + ); + + // Update file record + file.content = serde_json::to_value(&new_yml)?; + file.updated_at = Utc::now(); + file.verification = Verification::NotRequested; + + // Track successful modification + results.push(ModificationResult { + file_id: file.id, + file_type: "metric".to_string(), + file_name: modification.file_name.clone(), + success: true, + original_lines: original_lines.clone(), + adjusted_lines: adjusted_lines.clone(), + error: None, + modification_type: "content".to_string(), + timestamp: Utc::now(), + }); + + Ok((file, new_yml, results)) + } + Err(e) => { + let error = format!("Failed to validate modified YAML: {}", e); + error!( + file_id = %file.id, + file_name = %modification.file_name, + error = %error, + "YAML validation error" + ); + results.push(ModificationResult { + file_id: file.id, + file_type: "metric".to_string(), + file_name: modification.file_name.clone(), + success: false, + original_lines, + adjusted_lines: vec![], + error: Some(error.clone()), + modification_type: "validation".to_string(), + timestamp: Utc::now(), + }); + Err(anyhow::anyhow!(error)) + } + } + } + Err(e) => { + let error = format!("Failed to apply modifications: {}", e); + error!( + file_id = %file.id, + file_name = %modification.file_name, + error = %error, + "Modification application error" + ); + results.push(ModificationResult { + file_id: file.id, + file_type: "metric".to_string(), + file_name: modification.file_name.clone(), + success: false, + original_lines, + adjusted_lines: vec![], + error: Some(error.clone()), + modification_type: "modification".to_string(), + timestamp: Utc::now(), + }); + Err(anyhow::anyhow!(error)) + } + } +} + +async fn process_dashboard_file( + mut file: DashboardFile, + modification: &FileModification, +) -> Result<(DashboardFile, DashboardYml, Vec)> { + let mut results = Vec::new(); + + // Parse existing content + let current_yml: DashboardYml = match serde_json::from_value(file.content.clone()) { + Ok(yml) => yml, + Err(e) => { + let error = format!("Failed to parse existing dashboard YAML: {}", e); + error!( + file_id = %file.id, + file_name = %modification.file_name, + error = %error, + "YAML parsing error" + ); + results.push(ModificationResult { + file_id: file.id, + file_type: "dashboard".to_string(), + file_name: modification.file_name.clone(), + success: false, + original_lines: vec![], + adjusted_lines: vec![], + error: Some(error.clone()), + modification_type: "parsing".to_string(), + timestamp: Utc::now(), + }); + return Err(anyhow::anyhow!(error)); + } + }; + + // Convert to YAML string for line modifications + let current_content = match serde_yaml::to_string(¤t_yml) { + Ok(content) => content, + Err(e) => { + let error = format!("Failed to serialize dashboard YAML: {}", e); + error!( + file_id = %file.id, + file_name = %modification.file_name, + error = %error, + "YAML serialization error" + ); + results.push(ModificationResult { + file_id: file.id, + file_type: "dashboard".to_string(), + file_name: modification.file_name.clone(), + success: false, + original_lines: vec![], + adjusted_lines: vec![], + error: Some(error.clone()), + modification_type: "serialization".to_string(), + timestamp: Utc::now(), + }); + return Err(anyhow::anyhow!(error)); + } + }; + + // Track original line numbers before modifications + let mut original_lines = Vec::new(); + let mut adjusted_lines = Vec::new(); + + // Apply modifications + for m in &modification.modifications { + original_lines.extend(m.line_numbers.clone()); + } + + // Apply modifications and track results + match apply_modifications_to_content(¤t_content, &modification.modifications, &modification.file_name) { + Ok(modified_content) => { + // Create and validate new YML object + match DashboardYml::new(modified_content) { + Ok(new_yml) => { + // Update file record + file.content = match serde_json::to_value(&new_yml) { + Ok(content) => content, + Err(e) => { + let error = format!("Failed to serialize modified dashboard YAML: {}", e); + error!( + file_id = %file.id, + file_name = %modification.file_name, + error = %error, + "YAML serialization error" + ); + results.push(ModificationResult { + file_id: file.id, + file_type: "dashboard".to_string(), + file_name: modification.file_name.clone(), + success: false, + original_lines, + adjusted_lines: vec![], + error: Some(error.clone()), + modification_type: "serialization".to_string(), + timestamp: Utc::now(), + }); + return Err(anyhow::anyhow!(error)); + } + }; + file.updated_at = Utc::now(); + + // Track successful modification + results.push(ModificationResult { + file_id: file.id, + file_type: "dashboard".to_string(), + file_name: modification.file_name.clone(), + success: true, + original_lines: original_lines.clone(), + adjusted_lines: adjusted_lines.clone(), + error: None, + modification_type: "content".to_string(), + timestamp: Utc::now(), + }); + + Ok((file, new_yml, results)) + } + Err(e) => { + let error = format!("Failed to validate modified dashboard YAML: {}", e); + error!( + file_id = %file.id, + file_name = %modification.file_name, + error = %error, + "YAML validation error" + ); + results.push(ModificationResult { + file_id: file.id, + file_type: "dashboard".to_string(), + file_name: modification.file_name.clone(), + success: false, + original_lines, + adjusted_lines: vec![], + error: Some(error.clone()), + modification_type: "validation".to_string(), + timestamp: Utc::now(), + }); + return Err(anyhow::anyhow!(error)); + } + } + } + Err(e) => { + let error = format!("Failed to apply modifications: {}", e); + error!( + file_id = %file.id, + file_name = %modification.file_name, + error = %error, + "Modification application error" + ); + results.push(ModificationResult { + file_id: file.id, + file_type: "dashboard".to_string(), + file_name: modification.file_name.clone(), + success: false, + original_lines, + adjusted_lines: vec![], + error: Some(error.clone()), + modification_type: "modification".to_string(), + timestamp: Utc::now(), + }); + return Err(anyhow::anyhow!(error)); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Utc; + use serde_json::json; + + #[test] + fn test_validate_line_numbers() { + // Test valid range + assert!(validate_line_numbers(&[1, 5]).is_ok()); + assert!(validate_line_numbers(&[1, 1]).is_ok()); // Single line + + // Test empty array + let empty_err = validate_line_numbers(&[]).unwrap_err(); + assert!(empty_err.to_string().contains("cannot be empty")); + + // Test wrong number of elements + let wrong_len_err = validate_line_numbers(&[1, 2, 3]).unwrap_err(); + assert!(wrong_len_err.to_string().contains("exactly 2 numbers")); + + // Test invalid range (end < start) + let invalid_range_err = validate_line_numbers(&[5, 3]).unwrap_err(); + assert!(invalid_range_err.to_string().contains("must be greater than or equal to")); + + // Test starting below 1 + let invalid_start_err = validate_line_numbers(&[0, 2]).unwrap_err(); + assert!(invalid_start_err.to_string().contains("must be 1-indexed")); + } + + #[test] + fn test_apply_modifications_to_content() { + let original_content = "line1\nline2\nline3\nline4\nline5"; + + // Test single modification replacing two lines + let mods1 = vec![Modification { + new_content: "new line2\nnew line3".to_string(), + line_numbers: vec![2, 3], // Replace lines 2-3 + }]; + let result1 = apply_modifications_to_content(original_content, &mods1, "test.yml").unwrap(); + assert_eq!(result1.trim_end(), "line1\nnew line2\nnew line3\nline4\nline5"); + + // Test multiple non-overlapping modifications + let mods2 = vec![ + Modification { + new_content: "new line2".to_string(), + line_numbers: vec![2, 2], // Single line replacement + }, + Modification { + new_content: "new line4".to_string(), + line_numbers: vec![4, 4], // Single line replacement + }, + ]; + let result2 = apply_modifications_to_content(original_content, &mods2, "test.yml").unwrap(); + assert_eq!(result2.trim_end(), "line1\nnew line2\nline3\nnew line4\nline5"); + + // Test multiple modifications with line shifts + let content_with_more_lines = "line1\nline2\nline3\nline4\nline5\nline6\nline7\nline8\nline9\nline10"; + let mods_with_shifts = vec![ + Modification { + new_content: "new line2\nnew line2.1\nnew line2.2".to_string(), + line_numbers: vec![2, 3], // Replace lines 2-3 with 3 lines (net +1 line) + }, + Modification { + new_content: "new line6".to_string(), + line_numbers: vec![6, 7], // Replace lines 6-7 with 1 line (net -1 line) + }, + Modification { + new_content: "new line9\nnew line9.1".to_string(), + line_numbers: vec![9, 9], // Replace line 9 with 2 lines (net +1 line) + }, + ]; + let result_with_shifts = apply_modifications_to_content(content_with_more_lines, &mods_with_shifts, "test.yml").unwrap(); + assert_eq!( + result_with_shifts.trim_end(), + "line1\nnew line2\nnew line2.1\nnew line2.2\nline4\nline5\nnew line6\nline8\nnew line9\nnew line9.1\nline10" + ); + + // Test overlapping modifications (should fail) + let mods3 = vec![ + Modification { + new_content: "new lines".to_string(), + line_numbers: vec![2, 3], // Replace lines 2-3 + }, + Modification { + new_content: "overlap".to_string(), + line_numbers: vec![3, 4], // Overlaps with previous modification + }, + ]; + let result3 = apply_modifications_to_content(original_content, &mods3, "test.yml"); + assert!(result3.is_err()); + assert!(result3.unwrap_err().to_string().contains("overlaps")); + + // Test out of bounds modification + let mods4 = vec![Modification { + new_content: "new line".to_string(), + line_numbers: vec![6, 6], // Try to modify line 6 in a 5-line file + }]; + let result4 = apply_modifications_to_content(original_content, &mods4, "test.yml"); + assert!(result4.is_err()); + assert!(result4.unwrap_err().to_string().contains("out of bounds")); + + // Test broader line ranges with sequential modifications + let content_with_many_lines = "line1\nline2\nline3\nline4\nline5\nline6\nline7\nline8\nline9\nline10\nline11\nline12"; + let broad_range_mods = vec![ + Modification { + new_content: "new block 1-6\nnew content".to_string(), + line_numbers: vec![1, 6], // Replace lines 1-6 with 2 lines (net -4) + }, + Modification { + new_content: "new block 9-11\nextra line\nmore content".to_string(), + line_numbers: vec![9, 11], // Replace lines 9-11 with 3 lines (no net change) + }, + ]; + let result_broad_range = apply_modifications_to_content(content_with_many_lines, &broad_range_mods, "test.yml").unwrap(); + assert_eq!( + result_broad_range.trim_end(), + "new block 1-6\nnew content\nline7\nline8\nnew block 9-11\nextra line\nmore content\nline12" + ); + + // Test overlapping broad ranges (should fail) + let overlapping_broad_mods = vec![ + Modification { + new_content: "new content 1-6".to_string(), + line_numbers: vec![1, 6], + }, + Modification { + new_content: "overlap 4-8".to_string(), + line_numbers: vec![4, 8], + }, + ]; + let result_overlapping = apply_modifications_to_content(content_with_many_lines, &overlapping_broad_mods, "test.yml"); + assert!(result_overlapping.is_err()); + assert!(result_overlapping.unwrap_err().to_string().contains("overlaps")); + } + + #[test] + fn test_modification_result_tracking() { + let result = ModificationResult { + file_id: Uuid::new_v4(), + file_type: "metric".to_string(), + file_name: "test.yml".to_string(), + success: true, + original_lines: vec![1, 2, 3], + adjusted_lines: vec![1, 2], + error: None, + modification_type: "content".to_string(), + timestamp: Utc::now(), + }; + + // Test successful modification result + assert!(result.success); + assert_eq!(result.file_type, "metric"); + assert_eq!(result.original_lines, vec![1, 2, 3]); + assert_eq!(result.adjusted_lines, vec![1, 2]); + assert!(result.error.is_none()); + + // Test error modification result + let error_result = ModificationResult { + success: false, + error: Some("Failed to parse YAML".to_string()), + ..result + }; + assert!(!error_result.success); + assert!(error_result.error.is_some()); + assert_eq!(error_result.error.unwrap(), "Failed to parse YAML"); + } + + #[test] + fn test_tool_parameter_validation() { + let tool = ModifyFilesTool; + + // Test valid parameters + let valid_params = json!({ + "files": [{ + "id": Uuid::new_v4().to_string(), + "file_type": "metric", + "file_name": "test.yml", + "modifications": [{ + "new_content": "test content", + "line_numbers": [1, 2, 3] + }] + }] + }); + let valid_args = serde_json::to_string(&valid_params).unwrap(); + let result = serde_json::from_str::(&valid_args); + assert!(result.is_ok()); + + // Test invalid file type + let invalid_type_params = json!({ + "files": [{ + "id": Uuid::new_v4().to_string(), + "file_type": "invalid", + "file_name": "test.yml", + "modifications": [{ + "new_content": "test content", + "line_numbers": [1, 2, 3] + }] + }] + }); + let invalid_args = serde_json::to_string(&invalid_type_params).unwrap(); + let result = serde_json::from_str::(&invalid_args); + assert!(result.is_ok()); // Type validation happens during execution + + // Test missing required fields + let missing_fields_params = json!({ + "files": [{ + "id": Uuid::new_v4().to_string(), + "file_type": "metric" + // missing file_name and modifications + }] + }); + let missing_args = serde_json::to_string(&missing_fields_params).unwrap(); + let result = serde_json::from_str::(&missing_args); + assert!(result.is_err()); + } } \ No newline at end of file