From a3ee00ff849f1672bb628a6fa951a5dec43e0b91 Mon Sep 17 00:00:00 2001
From: dal <dallin@buster.so>
Date: Fri, 31 Jan 2025 10:16:35 -0700
Subject: [PATCH] refactor: Enhance file creation tool with error handling and
 type validation

- Added robust error handling for JSON parameter parsing in `CreateFilesTool`
- Updated file type description to clarify naming conventions
- Prepared for file creation logic implementation
- Moved file-related types to a new `file_types` module
- Removed redundant `types.rs` file
---
 .../utils/tools/file_tools/create_files.rs    |  17 +-
 .../file_tools/file_types/dashboard_file.rs   |  99 ++++++
 .../{types.rs => file_types/file.rs}          |   5 +-
 .../file_tools/file_types/metric_file.rs      | 294 ++++++++++++++++++
 .../utils/tools/file_tools/file_types/mod.rs  |   3 +
 api/src/utils/tools/file_tools/mod.rs         |   4 +-
 6 files changed, 415 insertions(+), 7 deletions(-)
 create mode 100644 api/src/utils/tools/file_tools/file_types/dashboard_file.rs
 rename api/src/utils/tools/file_tools/{types.rs => file_types/file.rs} (82%)
 create mode 100644 api/src/utils/tools/file_tools/file_types/metric_file.rs
 create mode 100644 api/src/utils/tools/file_tools/file_types/mod.rs

diff --git a/api/src/utils/tools/file_tools/create_files.rs b/api/src/utils/tools/file_tools/create_files.rs
index 430fdf4f2..402176563 100644
--- a/api/src/utils/tools/file_tools/create_files.rs
+++ b/api/src/utils/tools/file_tools/create_files.rs
@@ -27,8 +27,19 @@ impl ToolExecutor for CreateFilesTool {
 
     async fn execute(&self, tool_call: &ToolCall) -> Result<Value> {
         let params: CreateFilesParams =
-            serde_json::from_str(&tool_call.function.arguments.clone())?;
+            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
+                    ));
+                }
+            };
 
+        let files = params.files;
+
+        for file in files {}
         Ok(Value::Array(vec![]))
     }
 
@@ -48,12 +59,12 @@ impl ToolExecutor for CreateFilesTool {
                             "properties": {
                                 "name": {
                                     "type": "string",
-                                    "description": "The name of the file to be created"
+                                    "description": "The name of the file to be created. This should exlude the file extension. (i.e. '.yml')"
                                 },
                                 "file_type": {
                                     "type": "string",
                                     "enum": ["metric", "dashboard"],
-                                    "description": "The type of file to create"
+                                    "description": ""
                                 },
                                 "yml_content": {
                                     "type": "string",
diff --git a/api/src/utils/tools/file_tools/file_types/dashboard_file.rs b/api/src/utils/tools/file_tools/file_types/dashboard_file.rs
new file mode 100644
index 000000000..7acdf3493
--- /dev/null
+++ b/api/src/utils/tools/file_tools/file_types/dashboard_file.rs
@@ -0,0 +1,99 @@
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct DashboardFile {
+    id: Option<Uuid>,
+    name: Option<String>,
+    rows: Vec<Row>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Row {
+    items: Vec<RowItem>,    // max number of items in a row is 4, min is 1
+    row_height: u32,        // max is 550, min is 320
+    column_sizes: Vec<u32>, // max sum of elements is 12 min is 3
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct RowItem {
+    // This id is the id of the metric or item reference that goes here in the dashboard.
+    id: Uuid,
+}
+
+impl DashboardFile {
+    pub fn new(yml_content: String) -> Result<Self> {
+        // Oftentimes if the llm is creating a new dashboard, it will not include the id.
+        let mut file: DashboardFile = match serde_yaml::from_str(&yml_content) {
+            Ok(file) => file,
+            Err(e) => return Err(anyhow::anyhow!("Error parsing YAML: {}", e)),
+        };
+
+        // If the id is not provided, we will generate a new one.
+        if file.id.is_none() {
+            file.id = Some(Uuid::new_v4());
+        }
+
+        // If the name is not provided, we will use a default name.
+        if file.name.is_none() {
+            file.name = Some(String::from("New Dashboard"));
+        }
+
+        // Validate the file
+        match file.validate() {
+            Ok(_) => Ok(file),
+            Err(e) => Err(anyhow::anyhow!("Error compiling file: {}", e)),
+        }
+    }
+
+    pub fn validate(&self) -> Result<()> {
+        // Validate the id and name
+        if self.id.is_none() {
+            return Err(anyhow::anyhow!("Dashboard file id is required"));
+        }
+
+        if self.name.is_none() {
+            return Err(anyhow::anyhow!("Dashboard file name is required"));
+        }
+
+        // Validate each row
+        for row in &self.rows {
+            // Check row height constraints
+            if row.row_height < 320 || row.row_height > 550 {
+                return Err(anyhow::anyhow!(
+                    "Row height must be between 320 and 550, got {}",
+                    row.row_height
+                ));
+            }
+
+            // Check number of items constraint
+            if row.items.len() < 1 || row.items.len() > 4 {
+                return Err(anyhow::anyhow!(
+                    "Number of items in row must be between 1 and 4, got {}",
+                    row.items.len()
+                ));
+            }
+
+            // Check column sizes sum to valid amount
+            let sum: u32 = row.column_sizes.iter().sum();
+            if sum < 3 || sum > 12 {
+                return Err(anyhow::anyhow!(
+                    "Sum of column sizes must be between 3 and 12, got {}",
+                    sum
+                ));
+            }
+
+            // Check column sizes match number of items
+            if row.column_sizes.len() != row.items.len() {
+                return Err(anyhow::anyhow!(
+                    "Number of column sizes ({}) must match number of items ({})",
+                    row.column_sizes.len(),
+                    row.items.len()
+                ));
+            }
+        }
+
+        Ok(())
+    }
+}
diff --git a/api/src/utils/tools/file_tools/types.rs b/api/src/utils/tools/file_tools/file_types/file.rs
similarity index 82%
rename from api/src/utils/tools/file_tools/types.rs
rename to api/src/utils/tools/file_tools/file_types/file.rs
index e873ec653..31027173e 100644
--- a/api/src/utils/tools/file_tools/types.rs
+++ b/api/src/utils/tools/file_tools/file_types/file.rs
@@ -1,9 +1,10 @@
+use anyhow::Result;
 use serde::{Deserialize, Serialize};
+use uuid::Uuid;
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct File {
     pub name: String,
-    #[serde(rename = "type")]
     pub file_type: String,
     pub yml_content: String,
-}
+}
\ No newline at end of file
diff --git a/api/src/utils/tools/file_tools/file_types/metric_file.rs b/api/src/utils/tools/file_tools/file_types/metric_file.rs
new file mode 100644
index 000000000..bc0d9280a
--- /dev/null
+++ b/api/src/utils/tools/file_tools/file_types/metric_file.rs
@@ -0,0 +1,294 @@
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct MetricFile {
+    pub id: Option<Uuid>,
+    pub title: String,
+    pub description: Option<String>,
+    pub sql: String,
+    pub chart_config: ChartConfig,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(tag = "selectedChartType")]
+pub enum ChartConfig {
+    #[serde(rename = "bar")]
+    Bar(BarLineChartConfig),
+    #[serde(rename = "line")]
+    Line(BarLineChartConfig),
+    #[serde(rename = "scatter")]
+    Scatter(ScatterChartConfig),
+    #[serde(rename = "pie")]
+    Pie(PieChartConfig),
+    #[serde(rename = "combo")]
+    Combo(ComboChartConfig),
+    #[serde(rename = "metric")]
+    Metric(MetricChartConfig),
+    #[serde(rename = "table")]
+    Table(TableChartConfig),
+}
+
+// Base chart config shared by all chart types
+#[derive(Debug, Serialize, Deserialize)]
+pub struct BaseChartConfig {
+    pub column_label_formats: std::collections::HashMap<String, ColumnLabelFormat>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub column_settings: Option<std::collections::HashMap<String, ColumnSettings>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub colors: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub show_legend: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub grid_lines: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub show_legend_headline: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub goal_lines: Option<Vec<GoalLine>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub trendlines: Option<Vec<Trendline>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub disable_tooltip: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub y_axis_config: Option<serde_json::Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub x_axis_config: Option<serde_json::Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub category_axis_style_config: Option<serde_json::Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub y2_axis_config: Option<serde_json::Value>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ColumnLabelFormat {
+    pub column_type: String,
+    pub style: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub display_name: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub number_separator_style: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub minimum_fraction_digits: Option<i32>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maximum_fraction_digits: Option<i32>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub multiplier: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub prefix: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub suffix: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub replace_missing_data_with: Option<serde_json::Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub compact_numbers: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub currency: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub date_format: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub use_relative_time: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub is_utc: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub convert_number_to: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ColumnSettings {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub show_data_labels: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub show_data_labels_as_percentage: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub column_visualization: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub line_width: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub line_style: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub line_type: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub line_symbol_size: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub bar_roundness: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub line_symbol_size_dot: Option<f64>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct GoalLine {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub show: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub value: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub show_goal_line_label: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub goal_line_label: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub goal_line_color: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Trendline {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub show: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub show_trendline_label: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub trendline_label: Option<String>,
+    pub r#type: String,
+    pub column_id: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub trend_line_color: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct BarLineChartConfig {
+    #[serde(flatten)]
+    pub base: BaseChartConfig,
+    pub bar_and_line_axis: BarAndLineAxis,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub bar_layout: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub bar_sort_by: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub bar_group_type: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub bar_show_total_at_top: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub line_group_type: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct BarAndLineAxis {
+    pub x: Vec<String>,
+    pub y: Vec<String>,
+    pub category: Vec<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tooltip: Option<Vec<String>>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ScatterChartConfig {
+    #[serde(flatten)]
+    pub base: BaseChartConfig,
+    pub scatter_axis: ScatterAxis,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub scatter_dot_size: Option<Vec<f64>>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ScatterAxis {
+    pub x: Vec<String>,
+    pub y: Vec<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub category: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub size: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tooltip: Option<Vec<String>>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct PieChartConfig {
+    #[serde(flatten)]
+    pub base: BaseChartConfig,
+    pub pie_chart_axis: PieChartAxis,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub pie_display_label_as: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub pie_show_inner_label: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub pie_inner_label_aggregate: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub pie_inner_label_title: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub pie_label_position: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub pie_donut_width: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub pie_minimum_slice_percentage: Option<f64>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct PieChartAxis {
+    pub x: Vec<String>,
+    pub y: Vec<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tooltip: Option<Vec<String>>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ComboChartConfig {
+    #[serde(flatten)]
+    pub base: BaseChartConfig,
+    pub combo_chart_axis: ComboChartAxis,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ComboChartAxis {
+    pub x: Vec<String>,
+    pub y: Vec<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub y2: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub category: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tooltip: Option<Vec<String>>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct MetricChartConfig {
+    #[serde(flatten)]
+    pub base: BaseChartConfig,
+    pub metric_column_id: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub metric_value_aggregate: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub metric_header: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub metric_sub_header: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub metric_value_label: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct TableChartConfig {
+    #[serde(flatten)]
+    pub base: BaseChartConfig,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub table_column_order: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub table_column_widths: Option<std::collections::HashMap<String, f64>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub table_header_background_color: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub table_header_font_color: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub table_column_font_color: Option<String>,
+}
+
+impl MetricFile {
+    pub fn new(yml_content: String) -> Result<Self> {
+        let mut file: MetricFile = match serde_yaml::from_str(&yml_content) {
+            Ok(file) => file,
+            Err(e) => return Err(anyhow::anyhow!("Error parsing YAML: {}", e)),
+        };
+
+        if file.id.is_none() {
+            file.id = Some(Uuid::new_v4());
+        }
+
+        match file.validate() {
+            Ok(_) => Ok(file),
+            Err(e) => Err(anyhow::anyhow!("Error compiling file: {}", e)),
+        }
+    }
+
+    // TODO: Add validation logic here
+    pub fn validate(&self) -> Result<()> {
+        Ok(())
+    }
+}
diff --git a/api/src/utils/tools/file_tools/file_types/mod.rs b/api/src/utils/tools/file_tools/file_types/mod.rs
new file mode 100644
index 000000000..87a7aa98e
--- /dev/null
+++ b/api/src/utils/tools/file_tools/file_types/mod.rs
@@ -0,0 +1,3 @@
+pub mod dashboard_file;
+pub mod file;
+pub mod metric_file;
diff --git a/api/src/utils/tools/file_tools/mod.rs b/api/src/utils/tools/file_tools/mod.rs
index ccb7336a6..5d60e2a4d 100644
--- a/api/src/utils/tools/file_tools/mod.rs
+++ b/api/src/utils/tools/file_tools/mod.rs
@@ -1,10 +1,10 @@
 mod bulk_modify_files;
 mod create_files;
+pub mod file_types;
 mod open_files;
 mod search_data_catalog;
 mod search_files;
 mod send_to_user;
-mod types;
 
 pub use bulk_modify_files::BulkModifyFilesTool;
 pub use create_files::CreateFilesTool;
@@ -12,4 +12,4 @@ pub use open_files::OpenFilesTool;
 pub use search_data_catalog::SearchDataCatalogTool;
 pub use search_files::SearchFilesTool;
 pub use send_to_user::SendToUserTool;
-pub use types::*;
+