diff --git a/api/src/utils/tools/file_tools/create_files.rs b/api/src/utils/tools/file_tools/create_files.rs index 9f9798139..a67376973 100644 --- a/api/src/utils/tools/file_tools/create_files.rs +++ b/api/src/utils/tools/file_tools/create_files.rs @@ -1,14 +1,14 @@ use std::time::Instant; -use anyhow::Result; +use anyhow::{anyhow, Result}; use async_trait::async_trait; use chrono::Utc; use diesel::insert_into; use diesel_async::RunQueryDsl; use serde::{Deserialize, Serialize}; use serde_json::Value; +use tracing::debug; use uuid::Uuid; -use tracing::{debug, error}; use crate::{ database_dep::{ @@ -21,9 +21,9 @@ use crate::{ }; use super::{ - file_types::{dashboard_yml::DashboardYml, file::FileEnum, metric_yml::MetricYml}, + common::{validate_metric_ids, validate_sql}, + file_types::{dashboard_yml::DashboardYml, metric_yml::MetricYml}, FileModificationTool, - common::{validate_sql, validate_metric_ids}, }; use litellm::ToolCall; @@ -64,6 +64,92 @@ impl CreateFilesTool { impl FileModificationTool for CreateFilesTool {} +/// Process a metric file creation request +/// Returns Ok((MetricFile, MetricYml)) if successful, or an error message if failed +async fn process_metric_file(file: FileParams) -> Result<(MetricFile, MetricYml), String> { + debug!("Processing metric file creation: {}", file.name); + + let metric_yml = MetricYml::new(file.yml_content.clone()) + .map_err(|e| format!("Failed to parse metric YAML: {}", e))?; + + let metric_id = metric_yml.id.ok_or_else(|| { + "Metric YML file does not have an id. This should never happen.".to_string() + })?; + + // Validate SQL + if let Err(e) = validate_sql(&metric_yml.sql, &metric_id).await { + return Err(format!("SQL validation failed: {}", e)); + } + + let metric_file = MetricFile { + id: metric_id, + name: metric_yml.title.clone(), + file_name: format!("{}.yml", file.name), + content: serde_json::to_value(metric_yml.clone()) + .map_err(|e| format!("Failed to serialize metric content: {}", e))?, + 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, + }; + + Ok((metric_file, metric_yml)) +} + +/// Process a dashboard file creation request +/// Returns Ok((DashboardFile, DashboardYml)) if successful, or an error message if failed +async fn process_dashboard_file(file: FileParams) -> Result<(DashboardFile, DashboardYml), String> { + debug!("Processing dashboard file creation: {}", file.name); + + let dashboard_yml = DashboardYml::new(file.yml_content.clone()) + .map_err(|e| format!("Failed to parse dashboard YAML: {}", e))?; + + let dashboard_id = dashboard_yml.id.ok_or_else(|| { + "Dashboard YML file does not have an id. This should never happen.".to_string() + })?; + + // Collect and validate metric IDs from rows + let metric_ids: Vec = dashboard_yml + .rows + .iter() + .flat_map(|row| row.items.iter()) + .map(|item| item.id) + .collect(); + + if !metric_ids.is_empty() { + match validate_metric_ids(&metric_ids).await { + Ok(missing_ids) if !missing_ids.is_empty() => { + return Err(format!("Referenced metrics not found: {:?}", missing_ids)); + } + Err(e) => { + return Err(format!("Failed to validate metric IDs: {}", e)); + } + Ok(_) => (), // All metrics exist + } + } + + let dashboard_file = DashboardFile { + id: dashboard_id, + name: dashboard_yml.name.clone(), + file_name: format!("{}.yml", file.name), + content: serde_json::to_value(dashboard_yml.clone()) + .map_err(|e| format!("Failed to serialize dashboard content: {}", e))?, + filter: None, + organization_id: Uuid::new_v4(), + created_by: Uuid::new_v4(), + created_at: Utc::now(), + updated_at: Utc::now(), + deleted_at: None, + }; + + Ok((dashboard_file, dashboard_yml)) +} + #[async_trait] impl ToolExecutor for CreateFilesTool { type Output = CreateFilesOutput; @@ -84,10 +170,7 @@ impl ToolExecutor for CreateFilesTool { match serde_json::from_str(&tool_call.function.arguments.clone()) { Ok(params) => params, Err(e) => { - return Err(anyhow::anyhow!( - "Failed to parse create files parameters: {}", - e - )); + return Err(anyhow!("Failed to parse create files parameters: {}", e)); } }; @@ -104,92 +187,32 @@ impl ToolExecutor for CreateFilesTool { // First pass - validate and prepare all records for file in files { match file.file_type.as_str() { - "metric" => { - match MetricYml::new(file.yml_content.clone()) { - Ok(metric_yml) => { - if let Some(metric_id) = &metric_yml.id { - // Validate SQL before creating the record - if let Err(e) = validate_sql(&metric_yml.sql, metric_id).await { - failed_files.push((file.name, format!("SQL validation failed: {}", e))); - continue; - } - - 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())); - } + "metric" => match process_metric_file(file.clone()).await { + Ok((metric_file, metric_yml)) => { + metric_records.push(metric_file); + metric_ymls.push(metric_yml); } - } - "dashboard" => { - match DashboardYml::new(file.yml_content.clone()) { - Ok(dashboard_yml) => { - if let Some(dashboard_id) = &dashboard_yml.id { - // Collect and validate metric IDs from rows - let metric_ids: Vec = dashboard_yml.rows - .iter() - .flat_map(|row| row.items.iter()) - .map(|item| item.id) - .collect(); - - if !metric_ids.is_empty() { - match validate_metric_ids(&metric_ids).await { - Ok(missing_ids) if !missing_ids.is_empty() => { - failed_files.push((file.name, format!("Referenced metrics not found: {:?}", missing_ids))); - continue; - } - Err(e) => { - failed_files.push((file.name, format!("Failed to validate metric IDs: {}", e))); - continue; - } - Ok(_) => (), // All metrics exist - } - } - - let dashboard_file = DashboardFile { - id: dashboard_id.clone(), - name: dashboard_yml.name.clone(), - 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())); - } + Err(e) => { + failed_files.push((file.name, e)); } - } + }, + "dashboard" => match process_dashboard_file(file.clone()).await { + Ok((dashboard_file, dashboard_yml)) => { + dashboard_records.push(dashboard_file); + dashboard_ymls.push(dashboard_yml); + } + Err(e) => { + failed_files.push((file.name, e)); + } + }, _ => { - failed_files.push((file.name, format!("Invalid file type: {}. Currently only `metric` and `dashboard` types are supported.", file.file_type))); + failed_files.push(( + file.name, + format!( + "Invalid file type: {}. Currently only `metric` and `dashboard` types are supported.", + file.file_type + ), + )); } } } @@ -197,7 +220,7 @@ impl ToolExecutor for CreateFilesTool { // Second pass - bulk insert records let mut conn = match get_pg_pool().get().await { Ok(conn) => conn, - Err(e) => return Err(anyhow::anyhow!(e)), + Err(e) => return Err(anyhow!(e)), }; // Insert metric files @@ -210,7 +233,10 @@ impl ToolExecutor for CreateFilesTool { Ok(_) => { for (i, yml) in metric_ymls.into_iter().enumerate() { created_files.push(CreateFilesFile { - name: metric_records[i].file_name.trim_end_matches(".yml").to_string(), + name: metric_records[i] + .file_name + .trim_end_matches(".yml") + .to_string(), file_type: "metric".to_string(), yml_content: serde_yaml::to_string(&yml).unwrap_or_default(), }); @@ -237,7 +263,10 @@ impl ToolExecutor for CreateFilesTool { Ok(_) => { for (i, yml) in dashboard_ymls.into_iter().enumerate() { created_files.push(CreateFilesFile { - name: dashboard_records[i].file_name.trim_end_matches(".yml").to_string(), + name: dashboard_records[i] + .file_name + .trim_end_matches(".yml") + .to_string(), file_type: "dashboard".to_string(), yml_content: serde_yaml::to_string(&yml).unwrap_or_default(), });