mirror of https://github.com/buster-so/buster.git
modify is now an overwrite
This commit is contained in:
dal
parent
3a19fa7c4b
commit
752dea2ba3
|
|
@ -11,13 +11,16 @@ use diesel::{upsert::excluded, ExpressionMethods, QueryDsl};
|
|||
use diesel_async::RunQueryDsl;
|
||||
use indexmap::IndexMap;
|
||||
use query_engine::data_types::DataType;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use tracing::{debug, error, info};
|
||||
use chrono::Utc;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::{
|
||||
common::{
|
||||
process_dashboard_file_modification, FailedFileModification, ModificationResult,
|
||||
ModifyFilesOutput, ModifyFilesParams,
|
||||
validate_metric_ids, FailedFileModification, ModificationResult,
|
||||
ModifyFilesOutput,
|
||||
},
|
||||
file_types::file::FileWithId,
|
||||
FileModificationTool,
|
||||
|
|
@ -28,15 +31,26 @@ use crate::{
|
|||
};
|
||||
|
||||
#[derive(Debug)]
|
||||
struct DashboardModificationBatch {
|
||||
struct DashboardUpdateBatch {
|
||||
pub files: Vec<DashboardFile>,
|
||||
pub ymls: Vec<DashboardYml>,
|
||||
pub failed_modifications: Vec<(String, String)>,
|
||||
pub modification_results: Vec<ModificationResult>,
|
||||
pub failed_updates: Vec<(String, String)>,
|
||||
pub update_results: Vec<ModificationResult>,
|
||||
pub validation_messages: Vec<String>,
|
||||
pub validation_results: Vec<Vec<IndexMap<String, DataType>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct FileUpdate {
|
||||
pub id: Uuid,
|
||||
pub yml_content: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct UpdateFilesParams {
|
||||
pub files: Vec<FileUpdate>,
|
||||
}
|
||||
|
||||
pub struct ModifyDashboardFilesTool {
|
||||
agent: Arc<Agent>,
|
||||
}
|
||||
|
|
@ -49,10 +63,142 @@ impl ModifyDashboardFilesTool {
|
|||
|
||||
impl FileModificationTool for ModifyDashboardFilesTool {}
|
||||
|
||||
/// Process a dashboard file update with complete new YAML content
|
||||
async fn process_dashboard_file_update(
|
||||
mut file: DashboardFile,
|
||||
yml_content: String,
|
||||
duration: i64,
|
||||
) -> Result<(
|
||||
DashboardFile,
|
||||
DashboardYml,
|
||||
Vec<ModificationResult>,
|
||||
String,
|
||||
Vec<IndexMap<String, DataType>>,
|
||||
)> {
|
||||
debug!(
|
||||
file_id = %file.id,
|
||||
file_name = %file.name,
|
||||
"Processing dashboard file update"
|
||||
);
|
||||
|
||||
let mut results = Vec::new();
|
||||
|
||||
// Create and validate new YML object
|
||||
match DashboardYml::new(yml_content) {
|
||||
Ok(new_yml) => {
|
||||
debug!(
|
||||
file_id = %file.id,
|
||||
file_name = %file.name,
|
||||
"Successfully parsed and validated new dashboard content"
|
||||
);
|
||||
|
||||
// Collect and validate metric IDs from rows
|
||||
let metric_ids: Vec<Uuid> = new_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() => {
|
||||
let error = format!("Invalid metric references: {:?}", missing_ids);
|
||||
error!(
|
||||
file_id = %file.id,
|
||||
file_name = %file.name,
|
||||
error = %error,
|
||||
"Metric validation error"
|
||||
);
|
||||
results.push(ModificationResult {
|
||||
file_id: file.id,
|
||||
file_name: file.name.clone(),
|
||||
success: false,
|
||||
error: Some(error.clone()),
|
||||
modification_type: "validation".to_string(),
|
||||
timestamp: Utc::now(),
|
||||
duration,
|
||||
});
|
||||
return Err(anyhow::anyhow!(error));
|
||||
}
|
||||
Err(e) => {
|
||||
let error = format!("Failed to validate metrics: {}", e);
|
||||
error!(
|
||||
file_id = %file.id,
|
||||
file_name = %file.name,
|
||||
error = %error,
|
||||
"Metric validation error"
|
||||
);
|
||||
results.push(ModificationResult {
|
||||
file_id: file.id,
|
||||
file_name: file.name.clone(),
|
||||
success: false,
|
||||
error: Some(error.clone()),
|
||||
modification_type: "validation".to_string(),
|
||||
timestamp: Utc::now(),
|
||||
duration,
|
||||
});
|
||||
return Err(anyhow::anyhow!(error));
|
||||
}
|
||||
Ok(_) => {
|
||||
// All metrics exist, continue with update
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update file record
|
||||
file.content = new_yml.clone();
|
||||
file.updated_at = Utc::now();
|
||||
// Also update the file name to match the YAML name
|
||||
file.name = new_yml.name.clone();
|
||||
|
||||
// Track successful update
|
||||
results.push(ModificationResult {
|
||||
file_id: file.id,
|
||||
file_name: file.name.clone(),
|
||||
success: true,
|
||||
error: None,
|
||||
modification_type: "content".to_string(),
|
||||
timestamp: Utc::now(),
|
||||
duration,
|
||||
});
|
||||
|
||||
// Return successful result with empty validation results
|
||||
// since dashboards don't have SQL to validate like metrics do
|
||||
Ok((
|
||||
file,
|
||||
new_yml.clone(),
|
||||
results,
|
||||
"Dashboard validation successful".to_string(),
|
||||
Vec::new(),
|
||||
))
|
||||
}
|
||||
Err(e) => {
|
||||
let error = format!("Failed to validate modified YAML: {}", e);
|
||||
error!(
|
||||
file_id = %file.id,
|
||||
file_name = %file.name,
|
||||
error = %error,
|
||||
"YAML validation error"
|
||||
);
|
||||
results.push(ModificationResult {
|
||||
file_id: file.id,
|
||||
file_name: file.name.clone(),
|
||||
success: false,
|
||||
error: Some(error.clone()),
|
||||
modification_type: "validation".to_string(),
|
||||
timestamp: Utc::now(),
|
||||
duration,
|
||||
});
|
||||
Err(anyhow::anyhow!(error))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolExecutor for ModifyDashboardFilesTool {
|
||||
type Output = ModifyFilesOutput;
|
||||
type Params = ModifyFilesParams;
|
||||
type Params = UpdateFilesParams;
|
||||
|
||||
fn get_name(&self) -> String {
|
||||
"update_dashboards".to_string()
|
||||
|
|
@ -61,16 +207,16 @@ impl ToolExecutor for ModifyDashboardFilesTool {
|
|||
async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result<Self::Output> {
|
||||
let start_time = Instant::now();
|
||||
|
||||
debug!("Starting file modification execution");
|
||||
debug!("Starting dashboard file update execution");
|
||||
|
||||
info!("Processing {} files for modification", params.files.len());
|
||||
info!("Processing {} files for update", params.files.len());
|
||||
|
||||
// Initialize batch processing structures
|
||||
let mut batch = DashboardModificationBatch {
|
||||
let mut batch = DashboardUpdateBatch {
|
||||
files: Vec::new(),
|
||||
ymls: Vec::new(),
|
||||
failed_modifications: Vec::new(),
|
||||
modification_results: Vec::new(),
|
||||
failed_updates: Vec::new(),
|
||||
update_results: Vec::new(),
|
||||
validation_messages: Vec::new(),
|
||||
validation_results: Vec::new(),
|
||||
};
|
||||
|
|
@ -78,11 +224,11 @@ impl ToolExecutor for ModifyDashboardFilesTool {
|
|||
// Get database connection
|
||||
let mut conn = get_pg_pool().get().await?;
|
||||
|
||||
// Process each file modification
|
||||
for modification in params.files {
|
||||
// Process each file update
|
||||
for file_update in params.files {
|
||||
// Get the dashboard file from database
|
||||
match dashboard_files::table
|
||||
.filter(dashboard_files::id.eq(modification.id))
|
||||
.filter(dashboard_files::id.eq(file_update.id))
|
||||
.filter(dashboard_files::deleted_at.is_null())
|
||||
.first::<DashboardFile>(&mut conn)
|
||||
.await
|
||||
|
|
@ -90,10 +236,10 @@ impl ToolExecutor for ModifyDashboardFilesTool {
|
|||
Ok(dashboard_file) => {
|
||||
let duration = start_time.elapsed().as_millis() as i64;
|
||||
|
||||
// Process the dashboard file modification
|
||||
match process_dashboard_file_modification(
|
||||
// Process the dashboard file update
|
||||
match process_dashboard_file_update(
|
||||
dashboard_file.clone(),
|
||||
&modification,
|
||||
file_update.yml_content,
|
||||
duration,
|
||||
)
|
||||
.await
|
||||
|
|
@ -123,20 +269,20 @@ impl ToolExecutor for ModifyDashboardFilesTool {
|
|||
|
||||
batch.files.push(dashboard_file);
|
||||
batch.ymls.push(dashboard_yml);
|
||||
batch.modification_results.extend(results);
|
||||
batch.update_results.extend(results);
|
||||
batch.validation_messages.push(validation_message);
|
||||
batch.validation_results.push(validation_results);
|
||||
}
|
||||
Err(e) => {
|
||||
batch
|
||||
.failed_modifications
|
||||
.push((modification.file_name.clone(), e.to_string()));
|
||||
.failed_updates
|
||||
.push((format!("Dashboard {}", file_update.id), e.to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
batch.failed_modifications.push((
|
||||
modification.file_name.clone(),
|
||||
batch.failed_updates.push((
|
||||
format!("Dashboard {}", file_update.id),
|
||||
format!("Failed to find dashboard file: {}", e),
|
||||
));
|
||||
}
|
||||
|
|
@ -180,7 +326,7 @@ impl ToolExecutor for ModifyDashboardFilesTool {
|
|||
message: format!(
|
||||
"Modified {} dashboard files and created new versions. {} failures.",
|
||||
batch.files.len(),
|
||||
batch.failed_modifications.len()
|
||||
batch.failed_updates.len()
|
||||
),
|
||||
duration,
|
||||
files: Vec::new(),
|
||||
|
|
@ -205,10 +351,10 @@ impl ToolExecutor for ModifyDashboardFilesTool {
|
|||
}
|
||||
}));
|
||||
|
||||
// Add failed modifications to output
|
||||
// Add failed updates to output
|
||||
output.failed_files.extend(
|
||||
batch
|
||||
.failed_modifications
|
||||
.failed_updates
|
||||
.into_iter()
|
||||
.map(|(file_name, error)| FailedFileModification { file_name, error }),
|
||||
);
|
||||
|
|
@ -230,31 +376,16 @@ impl ToolExecutor for ModifyDashboardFilesTool {
|
|||
"description": get_modify_dashboards_yml_description().await,
|
||||
"items": {
|
||||
"type": "object",
|
||||
"required": ["id", "modifications"],
|
||||
"required": ["id", "yml_content"],
|
||||
"strict": true,
|
||||
"properties": {
|
||||
"id": {
|
||||
"type": "string",
|
||||
"description": get_dashboard_modification_id_description().await
|
||||
},
|
||||
"modifications": {
|
||||
"type": "array",
|
||||
"description": get_modify_dashboards_modifications_description().await,
|
||||
"items": {
|
||||
"type": "object",
|
||||
"required": ["content_to_replace", "new_content"],
|
||||
"properties": {
|
||||
"content_to_replace": {
|
||||
"type": "string",
|
||||
"description": get_modify_dashboards_content_to_replace_description().await
|
||||
},
|
||||
"new_content": {
|
||||
"type": "string",
|
||||
"description": get_modify_dashboards_new_content_description().await
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
}
|
||||
"yml_content": {
|
||||
"type": "string",
|
||||
"description": get_dashboard_yml_description().await
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
|
|
@ -269,7 +400,7 @@ impl ToolExecutor for ModifyDashboardFilesTool {
|
|||
|
||||
async fn get_modify_dashboards_description() -> String {
|
||||
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
|
||||
return "Modifies existing dashboard configuration files by replacing specified content with new content. Before using this tool, carefully think through the dashboard format, specification, and structure to ensure your edits maintain consistency and meet requirements. If the dashboard fundamentally changes, the name should be updated to reflect the changes. However, if the core dashboard topic remains the same, the name should stay unchanged.".to_string();
|
||||
return "Updates existing dashboard configuration files with new YAML content. Provide the complete YAML content for the dashboard, replacing the entire existing file. The system will preserve version history and perform all necessary validations on the new content.".to_string();
|
||||
}
|
||||
|
||||
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
|
||||
|
|
@ -277,7 +408,7 @@ async fn get_modify_dashboards_description() -> String {
|
|||
Ok(message) => message,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get prompt system message: {}", e);
|
||||
"Modifies existing dashboard configuration files by replacing specified content with new content. Before using this tool, carefully think through the dashboard format, specification, and structure to ensure your edits maintain consistency and meet requirements. If the dashboard fundamentally changes, the name should be updated to reflect the changes. However, if the core dashboard topic remains the same, the name should stay unchanged.".to_string()
|
||||
"Updates existing dashboard configuration files with new YAML content. Provide the complete YAML content for the dashboard, replacing the entire existing file. The system will preserve version history and perform all necessary validations on the new content.".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -312,39 +443,9 @@ async fn get_dashboard_modification_id_description() -> String {
|
|||
}
|
||||
}
|
||||
|
||||
async fn get_modify_dashboards_file_name_description() -> String {
|
||||
async fn get_dashboard_yml_description() -> String {
|
||||
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
|
||||
return "Name of the dashboard file to modify".to_string();
|
||||
}
|
||||
|
||||
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
|
||||
match get_prompt_system_message(&client, "5e0761df-2668-40f3-874e-84eb54c66e4d").await {
|
||||
Ok(message) => message,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get prompt system message: {}", e);
|
||||
"Name of the dashboard file to modify".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_modify_dashboards_modifications_description() -> String {
|
||||
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
|
||||
return "List of content modifications to make to the dashboard file".to_string();
|
||||
}
|
||||
|
||||
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
|
||||
match get_prompt_system_message(&client, "ee5789a9-fd99-4afd-a5c4-88f2ebe58fe9").await {
|
||||
Ok(message) => message,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get prompt system message: {}", e);
|
||||
"List of content modifications to make to the dashboard file".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_modify_dashboards_new_content_description() -> String {
|
||||
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
|
||||
return "The new content to replace the existing content with. If fundamentally changing the dashboard's purpose or focus, update the name property accordingly. If the core topic remains the same, maintain the original name.".to_string();
|
||||
return "The complete new YAML content for the dashboard, following the dashboard schema specification. This will replace the entire existing content of the file.".to_string();
|
||||
}
|
||||
|
||||
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
|
||||
|
|
@ -352,153 +453,39 @@ async fn get_modify_dashboards_new_content_description() -> String {
|
|||
Ok(message) => message,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get prompt system message: {}", e);
|
||||
"The new content to replace the existing content with. If fundamentally changing the dashboard's purpose or focus, update the name property accordingly. If the core topic remains the same, maintain the original name.".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_modify_dashboards_content_to_replace_description() -> String {
|
||||
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
|
||||
return "The exact content in the file that should be replaced. Precise matching is crucial: the provided content must match exactly and be specific enough to target only the intended section, avoiding unintended replacements. This should typically be a small, specific snippet; replacing the entire file content is usually not intended unless the entire dashboard definition needs a complete overhaul. Use an empty string to append the new content to the end of the file."
|
||||
.to_string();
|
||||
}
|
||||
|
||||
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
|
||||
match get_prompt_system_message(&client, "7e89b1f9-30ed-4f0c-b4da-32ce03f31635").await {
|
||||
Ok(message) => message,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get prompt system message: {}", e);
|
||||
"The exact content in the file that should be replaced. Precise matching is crucial: the provided content must match exactly and be specific enough to target only the intended section, avoiding unintended replacements. This should typically be a small, specific snippet; replacing the entire file content is usually not intended unless the entire dashboard definition needs a complete overhaul. Use an empty string to append the new content to the end of the file.".to_string()
|
||||
"The complete new YAML content for the dashboard, following the dashboard schema specification. This will replace the entire existing content of the file.".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use super::*;
|
||||
use crate::tools::categories::file_tools::common::{
|
||||
apply_modifications_to_content, Modification, ModificationResult,
|
||||
};
|
||||
use chrono::Utc;
|
||||
use serde_json::json;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[test]
|
||||
fn test_apply_modifications_to_content() {
|
||||
let original_content =
|
||||
"name: test_dashboard\ntype: dashboard\ndescription: A test dashboard";
|
||||
|
||||
// Test single modification
|
||||
let mods1 = vec![Modification {
|
||||
content_to_replace: "type: dashboard".to_string(),
|
||||
new_content: "type: custom_dashboard".to_string(),
|
||||
}];
|
||||
let result1 = apply_modifications_to_content(original_content, &mods1, "test.yml").unwrap();
|
||||
assert_eq!(
|
||||
result1,
|
||||
"name: test_dashboard\ntype: custom_dashboard\ndescription: A test dashboard"
|
||||
);
|
||||
|
||||
// Test multiple non-overlapping modifications
|
||||
let mods2 = vec![
|
||||
Modification {
|
||||
content_to_replace: "test_dashboard".to_string(),
|
||||
new_content: "new_dashboard".to_string(),
|
||||
},
|
||||
Modification {
|
||||
content_to_replace: "A test dashboard".to_string(),
|
||||
new_content: "An updated dashboard".to_string(),
|
||||
},
|
||||
];
|
||||
let result2 = apply_modifications_to_content(original_content, &mods2, "test.yml").unwrap();
|
||||
assert_eq!(
|
||||
result2,
|
||||
"name: new_dashboard\ntype: dashboard\ndescription: An updated dashboard"
|
||||
);
|
||||
|
||||
// Test content not found
|
||||
let mods3 = vec![Modification {
|
||||
content_to_replace: "nonexistent content".to_string(),
|
||||
new_content: "new content".to_string(),
|
||||
}];
|
||||
let result3 = apply_modifications_to_content(original_content, &mods3, "test.yml");
|
||||
assert!(result3.is_err());
|
||||
assert!(result3
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("Content to replace not found"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_modification_result_tracking() {
|
||||
let result = ModificationResult {
|
||||
file_id: Uuid::new_v4(),
|
||||
file_name: "test.yml".to_string(),
|
||||
success: true,
|
||||
error: None,
|
||||
modification_type: "content".to_string(),
|
||||
timestamp: Utc::now(),
|
||||
duration: 0,
|
||||
};
|
||||
|
||||
assert!(result.success);
|
||||
assert!(result.error.is_none());
|
||||
|
||||
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() {
|
||||
// Test valid parameters
|
||||
let valid_params = json!({
|
||||
"files": [{
|
||||
"id": Uuid::new_v4().to_string(),
|
||||
"file_name": "test.yml",
|
||||
"modifications": [{
|
||||
"new_content": "new content",
|
||||
"line_numbers": [1, 2]
|
||||
}]
|
||||
"yml_content": "name: Test Dashboard\ndescription: A test dashboard\nrows:\n - id: 1\n items:\n - id: 00000000-0000-0000-0000-000000000001\n columnSizes: [12]"
|
||||
}]
|
||||
});
|
||||
let valid_args = serde_json::to_string(&valid_params).unwrap();
|
||||
let result = serde_json::from_str::<ModifyFilesParams>(&valid_args);
|
||||
let result = serde_json::from_str::<UpdateFilesParams>(&valid_args);
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Test missing required fields
|
||||
let missing_fields_params = json!({
|
||||
"files": [{
|
||||
"id": Uuid::new_v4().to_string(),
|
||||
"file_name": "test.yml"
|
||||
// missing modifications
|
||||
"id": Uuid::new_v4().to_string()
|
||||
// missing yml_content
|
||||
}]
|
||||
});
|
||||
let missing_args = serde_json::to_string(&missing_fields_params).unwrap();
|
||||
let result = serde_json::from_str::<ModifyFilesParams>(&missing_args);
|
||||
let result = serde_json::from_str::<UpdateFilesParams>(&missing_args);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_apply_modifications_append() {
|
||||
let original_content =
|
||||
"name: test_dashboard\ntype: dashboard\ndescription: A test dashboard";
|
||||
|
||||
// Test appending content with empty content_to_replace
|
||||
let mods = vec![Modification {
|
||||
content_to_replace: "".to_string(),
|
||||
new_content: "\nvisibility: public".to_string(),
|
||||
}];
|
||||
let result = apply_modifications_to_content(original_content, &mods, "test.yml").unwrap();
|
||||
assert_eq!(
|
||||
result,
|
||||
"name: test_dashboard\ntype: dashboard\ndescription: A test dashboard\nvisibility: public"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,14 +14,15 @@ use diesel_async::RunQueryDsl;
|
|||
use futures::future::join_all;
|
||||
use indexmap::IndexMap;
|
||||
use query_engine::{data_source_query_routes::query_engine::query_engine, data_types::DataType};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use tracing::{debug, error, info};
|
||||
use uuid::Uuid;
|
||||
use chrono::Utc;
|
||||
|
||||
use super::{
|
||||
common::{
|
||||
process_metric_file_modification, ModificationResult, ModifyFilesOutput, ModifyFilesParams,
|
||||
FailedFileModification,
|
||||
validate_sql, ModificationResult, ModifyFilesOutput, FailedFileModification,
|
||||
},
|
||||
file_types::file::FileWithId,
|
||||
FileModificationTool,
|
||||
|
|
@ -32,15 +33,26 @@ use crate::{
|
|||
};
|
||||
|
||||
#[derive(Debug)]
|
||||
struct MetricModificationBatch {
|
||||
struct MetricUpdateBatch {
|
||||
pub files: Vec<MetricFile>,
|
||||
pub ymls: Vec<MetricYml>,
|
||||
pub failed_modifications: Vec<(String, String)>,
|
||||
pub modification_results: Vec<ModificationResult>,
|
||||
pub failed_updates: Vec<(String, String)>,
|
||||
pub update_results: Vec<ModificationResult>,
|
||||
pub validation_messages: Vec<String>,
|
||||
pub validation_results: Vec<Vec<IndexMap<String, DataType>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct FileUpdate {
|
||||
pub id: Uuid,
|
||||
pub yml_content: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct UpdateFilesParams {
|
||||
pub files: Vec<FileUpdate>,
|
||||
}
|
||||
|
||||
pub struct ModifyMetricFilesTool {
|
||||
agent: Arc<Agent>,
|
||||
}
|
||||
|
|
@ -53,10 +65,120 @@ impl ModifyMetricFilesTool {
|
|||
|
||||
impl FileModificationTool for ModifyMetricFilesTool {}
|
||||
|
||||
/// Process a metric file update with complete new YAML content
|
||||
/// Returns updated file, YAML, validation results and messages if successful
|
||||
async fn process_metric_file_update(
|
||||
mut file: MetricFile,
|
||||
yml_content: String,
|
||||
duration: i64,
|
||||
) -> Result<(
|
||||
MetricFile,
|
||||
MetricYml,
|
||||
Vec<ModificationResult>,
|
||||
String,
|
||||
Vec<IndexMap<String, DataType>>,
|
||||
)> {
|
||||
debug!(
|
||||
file_id = %file.id,
|
||||
file_name = %file.name,
|
||||
"Processing metric file update"
|
||||
);
|
||||
|
||||
let mut results = Vec::new();
|
||||
|
||||
// Create and validate new YML object
|
||||
match MetricYml::new(yml_content) {
|
||||
Ok(new_yml) => {
|
||||
debug!(
|
||||
file_id = %file.id,
|
||||
file_name = %file.name,
|
||||
"Successfully parsed and validated new metric content"
|
||||
);
|
||||
|
||||
// Validate SQL and get dataset_id from the first dataset
|
||||
if new_yml.dataset_ids.is_empty() {
|
||||
let error = "Missing required field 'dataset_ids'".to_string();
|
||||
results.push(ModificationResult {
|
||||
file_id: file.id,
|
||||
file_name: file.name.clone(),
|
||||
success: false,
|
||||
error: Some(error.clone()),
|
||||
modification_type: "validation".to_string(),
|
||||
timestamp: Utc::now(),
|
||||
duration,
|
||||
});
|
||||
return Err(anyhow::anyhow!(error));
|
||||
}
|
||||
|
||||
let dataset_id = new_yml.dataset_ids[0];
|
||||
match validate_sql(&new_yml.sql, &dataset_id).await {
|
||||
Ok((message, validation_results, metadata)) => {
|
||||
// Update file record
|
||||
file.content = new_yml.clone();
|
||||
file.name = new_yml.name.clone();
|
||||
file.updated_at = Utc::now();
|
||||
file.data_metadata = metadata;
|
||||
|
||||
// Track successful update
|
||||
results.push(ModificationResult {
|
||||
file_id: file.id,
|
||||
file_name: file.name.clone(),
|
||||
success: true,
|
||||
error: None,
|
||||
modification_type: "content".to_string(),
|
||||
timestamp: Utc::now(),
|
||||
duration,
|
||||
});
|
||||
|
||||
Ok((file, new_yml.clone(), results, message, validation_results))
|
||||
}
|
||||
Err(e) => {
|
||||
let error = format!("SQL validation failed: {}", e);
|
||||
error!(
|
||||
file_id = %file.id,
|
||||
file_name = %file.name,
|
||||
error = %error,
|
||||
"SQL validation error"
|
||||
);
|
||||
results.push(ModificationResult {
|
||||
file_id: file.id,
|
||||
file_name: file.name.clone(),
|
||||
success: false,
|
||||
error: Some(error.clone()),
|
||||
modification_type: "sql_validation".to_string(),
|
||||
timestamp: Utc::now(),
|
||||
duration,
|
||||
});
|
||||
Err(anyhow::anyhow!(error))
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let error = format!("Failed to validate YAML: {}", e);
|
||||
error!(
|
||||
file_id = %file.id,
|
||||
file_name = %file.name,
|
||||
error = %error,
|
||||
"YAML validation error"
|
||||
);
|
||||
results.push(ModificationResult {
|
||||
file_id: file.id,
|
||||
file_name: file.name.clone(),
|
||||
success: false,
|
||||
error: Some(error.clone()),
|
||||
modification_type: "validation".to_string(),
|
||||
timestamp: Utc::now(),
|
||||
duration,
|
||||
});
|
||||
Err(anyhow::anyhow!(error))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolExecutor for ModifyMetricFilesTool {
|
||||
type Output = ModifyFilesOutput;
|
||||
type Params = ModifyFilesParams;
|
||||
type Params = UpdateFilesParams;
|
||||
|
||||
fn get_name(&self) -> String {
|
||||
"update_metrics".to_string()
|
||||
|
|
@ -65,16 +187,16 @@ impl ToolExecutor for ModifyMetricFilesTool {
|
|||
async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result<Self::Output> {
|
||||
let start_time = Instant::now();
|
||||
|
||||
debug!("Starting file modification execution");
|
||||
debug!("Starting file update execution");
|
||||
|
||||
info!("Processing {} files for modification", params.files.len());
|
||||
info!("Processing {} files for update", params.files.len());
|
||||
|
||||
// Initialize batch processing structures
|
||||
let mut batch = MetricModificationBatch {
|
||||
let mut batch = MetricUpdateBatch {
|
||||
files: Vec::new(),
|
||||
ymls: Vec::new(),
|
||||
failed_modifications: Vec::new(),
|
||||
modification_results: Vec::new(),
|
||||
failed_updates: Vec::new(),
|
||||
update_results: Vec::new(),
|
||||
validation_messages: Vec::new(),
|
||||
validation_results: Vec::new(),
|
||||
};
|
||||
|
|
@ -107,17 +229,17 @@ impl ToolExecutor for ModifyMetricFilesTool {
|
|||
.await
|
||||
{
|
||||
Ok(files) => {
|
||||
// Create futures for concurrent processing of file modifications
|
||||
let modification_futures = files
|
||||
// Create futures for concurrent processing of file updates
|
||||
let update_futures = files
|
||||
.into_iter()
|
||||
.filter_map(|file| {
|
||||
let modifications = file_map.get(&file.id)?;
|
||||
let file_update = file_map.get(&file.id)?;
|
||||
let start_time_elapsed = start_time.elapsed().as_millis() as i64;
|
||||
|
||||
Some(async move {
|
||||
let result = process_metric_file_modification(
|
||||
let result = process_metric_file_update(
|
||||
file.clone(),
|
||||
modifications,
|
||||
file_update.yml_content.clone(),
|
||||
start_time_elapsed,
|
||||
).await;
|
||||
|
||||
|
|
@ -125,14 +247,14 @@ impl ToolExecutor for ModifyMetricFilesTool {
|
|||
Ok((metric_file, metric_yml, results, validation_message, validation_results)) => {
|
||||
Ok((metric_file, metric_yml, results, validation_message, validation_results))
|
||||
}
|
||||
Err(e) => Err((modifications.file_name.clone(), e.to_string())),
|
||||
Err(e) => Err((file.name.clone(), e.to_string())),
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
// Wait for all futures to complete
|
||||
let results = join_all(modification_futures).await;
|
||||
let results = join_all(update_futures).await;
|
||||
|
||||
// Process results
|
||||
for result in results {
|
||||
|
|
@ -149,61 +271,14 @@ impl ToolExecutor for ModifyMetricFilesTool {
|
|||
.version_history
|
||||
.add_version(next_version, metric_yml.clone());
|
||||
|
||||
// Update metadata if SQL has changed
|
||||
// The SQL is already validated by process_metric_file_modification
|
||||
if results.iter().any(|r| r.modification_type == "content") {
|
||||
// Update the name field from the metric_yml
|
||||
// This is redundant but ensures the name is set correctly
|
||||
metric_file.name = metric_yml.name.clone();
|
||||
|
||||
// Check if we have a dataset to work with
|
||||
if !metric_yml.dataset_ids.is_empty() {
|
||||
let dataset_id = metric_yml.dataset_ids[0];
|
||||
|
||||
// Get data source for the dataset
|
||||
match datasets::table
|
||||
.filter(datasets::id.eq(dataset_id))
|
||||
.select(datasets::data_source_id)
|
||||
.first::<Uuid>(&mut conn)
|
||||
.await
|
||||
{
|
||||
Ok(data_source_id) => {
|
||||
// Execute query to get metadata
|
||||
match query_engine(
|
||||
&data_source_id,
|
||||
&metric_yml.sql,
|
||||
Some(100),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(query_result) => {
|
||||
// Update metadata
|
||||
metric_file.data_metadata =
|
||||
Some(query_result.metadata);
|
||||
debug!("Updated metadata for metric file {}", metric_file.id);
|
||||
}
|
||||
Err(e) => {
|
||||
debug!("Failed to execute SQL for metadata: {}", e);
|
||||
// Continue with the update even if metadata refresh fails
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
debug!("Failed to get data source ID: {}", e);
|
||||
// Continue with the update even if we can't get data source
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
batch.files.push(metric_file);
|
||||
batch.ymls.push(metric_yml);
|
||||
batch.modification_results.extend(results);
|
||||
batch.update_results.extend(results);
|
||||
batch.validation_messages.push(validation_message);
|
||||
batch.validation_results.push(validation_results);
|
||||
}
|
||||
Err((file_name, error)) => {
|
||||
batch.failed_modifications.push((file_name, error));
|
||||
batch.failed_updates.push((file_name, error));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -259,7 +334,7 @@ impl ToolExecutor for ModifyMetricFilesTool {
|
|||
message: format!(
|
||||
"Modified {} metric files and created new versions. {} failures.",
|
||||
batch.files.len(),
|
||||
batch.failed_modifications.len()
|
||||
batch.failed_updates.len()
|
||||
),
|
||||
duration,
|
||||
files: Vec::new(),
|
||||
|
|
@ -287,7 +362,7 @@ impl ToolExecutor for ModifyMetricFilesTool {
|
|||
// Add failed modifications to output
|
||||
output.failed_files.extend(
|
||||
batch
|
||||
.failed_modifications
|
||||
.failed_updates
|
||||
.into_iter()
|
||||
.map(|(file_name, error)| FailedFileModification { file_name, error }),
|
||||
);
|
||||
|
|
@ -309,38 +384,15 @@ impl ToolExecutor for ModifyMetricFilesTool {
|
|||
"description": get_modify_metrics_yml_description().await,
|
||||
"items": {
|
||||
"type": "object",
|
||||
"required": ["id", "file_name", "modifications"],
|
||||
"required": ["id", "yml_content"],
|
||||
"properties": {
|
||||
"id": {
|
||||
"type": "string",
|
||||
"description": get_metric_id_description().await
|
||||
},
|
||||
"file_name": {
|
||||
"yml_content": {
|
||||
"type": "string",
|
||||
"description": get_modify_metrics_file_name_description().await
|
||||
},
|
||||
"modifications": {
|
||||
"type": "array",
|
||||
"description": get_modify_metrics_modifications_description().await,
|
||||
"items": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"content_to_replace",
|
||||
"new_content"
|
||||
],
|
||||
"strict": true,
|
||||
"properties": {
|
||||
"content_to_replace": {
|
||||
"type": "string",
|
||||
"description": get_modify_metrics_content_to_replace_description().await
|
||||
},
|
||||
"new_content": {
|
||||
"type": "string",
|
||||
"description": get_modify_metrics_new_content_description().await
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
}
|
||||
"description": get_metric_yml_description().await
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
|
|
@ -355,7 +407,7 @@ impl ToolExecutor for ModifyMetricFilesTool {
|
|||
|
||||
async fn get_modify_metrics_description() -> String {
|
||||
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
|
||||
return "Modifies existing metric configuration files by replacing specified content with new content. Before using this tool, carefully review the metric YAML specification and thoughtfully plan your edits based on the visualization type and its specific configuration requirements. When modifying a metric's SQL, make sure to also update the name and/or time frame to reflect the changes. Each visualization has unique axis settings, formatting options, and data structure needs that must be considered when making modifications.".to_string();
|
||||
return "Updates existing metric configuration files with new YAML content. Provide the complete YAML content for the metric, replacing the entire existing file. The system will preserve version history and perform all necessary validations on the new content.".to_string();
|
||||
}
|
||||
|
||||
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
|
||||
|
|
@ -363,7 +415,7 @@ async fn get_modify_metrics_description() -> String {
|
|||
Ok(message) => message,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get prompt system message: {}", e);
|
||||
"Modifies existing metric configuration files by replacing specified content with new content. Before using this tool, carefully review the metric YAML specification and thoughtfully plan your edits based on the visualization type and its specific configuration requirements. When modifying a metric's SQL, make sure to also update the name and/or time frame to reflect the changes. Each visualization has unique axis settings, formatting options, and data structure needs that must be considered when making modifications.".to_string()
|
||||
"Updates existing metric configuration files with new YAML content. Provide the complete YAML content for the metric, replacing the entire existing file. The system will preserve version history and perform all necessary validations on the new content.".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -383,39 +435,9 @@ async fn get_modify_metrics_yml_description() -> String {
|
|||
}
|
||||
}
|
||||
|
||||
async fn get_modify_metrics_file_name_description() -> String {
|
||||
async fn get_metric_yml_description() -> String {
|
||||
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
|
||||
return "Name of the metric file to modify".to_string();
|
||||
}
|
||||
|
||||
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
|
||||
match get_prompt_system_message(&client, "5e9e0a31-760a-483f-8876-41f2027bf731").await {
|
||||
Ok(message) => message,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get prompt system message: {}", e);
|
||||
"Name of the metric file to modify".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_modify_metrics_modifications_description() -> String {
|
||||
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
|
||||
return "List of content modifications to make to the metric file".to_string();
|
||||
}
|
||||
|
||||
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
|
||||
match get_prompt_system_message(&client, "c56d3034-e527-45b6-aa2e-18fb5e3240de").await {
|
||||
Ok(message) => message,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get prompt system message: {}", e);
|
||||
"List of content modifications to make to the metric file".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_modify_metrics_new_content_description() -> String {
|
||||
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
|
||||
return "The new content to replace the existing content with. If modifying SQL, ensure name and time frame properties are also updated to reflect the changes.".to_string();
|
||||
return "The complete new YAML content for the metric, following the metric schema specification. This will replace the entire existing content of the file.".to_string();
|
||||
}
|
||||
|
||||
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
|
||||
|
|
@ -423,23 +445,7 @@ async fn get_modify_metrics_new_content_description() -> String {
|
|||
Ok(message) => message,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get prompt system message: {}", e);
|
||||
"The new content to replace the existing content with. If modifying SQL, ensure name and time frame properties are also updated to reflect the changes.".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_modify_metrics_content_to_replace_description() -> String {
|
||||
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
|
||||
return "The exact content in the file that should be replaced. Must match exactly and be specific enough to only match once. Use an empty string to append the new content to the end of the file."
|
||||
.to_string();
|
||||
}
|
||||
|
||||
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
|
||||
match get_prompt_system_message(&client, "ad7e79f0-dd3a-4239-9548-ee7f4ef3be5a").await {
|
||||
Ok(message) => message,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get prompt system message: {}", e);
|
||||
"The exact content in the file that should be replaced. Must match exactly and be specific enough to only match once. Use an empty string to append the new content to the end of the file.".to_string()
|
||||
"The complete new YAML content for the metric, following the metric schema specification. This will replace the entire existing content of the file.".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -461,58 +467,10 @@ async fn get_metric_id_description() -> String {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use crate::tools::file_tools::common::{apply_modifications_to_content, Modification};
|
||||
|
||||
use super::*;
|
||||
use chrono::Utc;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_apply_modifications_to_content() {
|
||||
let original_content = "name: test_metric\ntype: counter\ndescription: A test metric";
|
||||
|
||||
// Test single modification
|
||||
let mods1 = vec![Modification {
|
||||
content_to_replace: "type: counter".to_string(),
|
||||
new_content: "type: gauge".to_string(),
|
||||
}];
|
||||
let result1 = apply_modifications_to_content(original_content, &mods1, "test.yml").unwrap();
|
||||
assert_eq!(
|
||||
result1,
|
||||
"name: test_metric\ntype: gauge\ndescription: A test metric"
|
||||
);
|
||||
|
||||
// Test multiple non-overlapping modifications
|
||||
let mods2 = vec![
|
||||
Modification {
|
||||
content_to_replace: "test_metric".to_string(),
|
||||
new_content: "new_metric".to_string(),
|
||||
},
|
||||
Modification {
|
||||
content_to_replace: "A test metric".to_string(),
|
||||
new_content: "An updated metric".to_string(),
|
||||
},
|
||||
];
|
||||
let result2 = apply_modifications_to_content(original_content, &mods2, "test.yml").unwrap();
|
||||
assert_eq!(
|
||||
result2,
|
||||
"name: new_metric\ntype: counter\ndescription: An updated metric"
|
||||
);
|
||||
|
||||
// Test content not found
|
||||
let mods3 = vec![Modification {
|
||||
content_to_replace: "nonexistent content".to_string(),
|
||||
new_content: "new content".to_string(),
|
||||
}];
|
||||
let result3 = apply_modifications_to_content(original_content, &mods3, "test.yml");
|
||||
assert!(result3.is_err());
|
||||
assert!(result3
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("Content to replace not found"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_modification_result_tracking() {
|
||||
let result = ModificationResult {
|
||||
|
|
@ -544,27 +502,22 @@ mod tests {
|
|||
let valid_params = json!({
|
||||
"files": [{
|
||||
"id": Uuid::new_v4().to_string(),
|
||||
"file_name": "test.yml",
|
||||
"modifications": [{
|
||||
"new_content": "new content",
|
||||
"line_numbers": [1, 2]
|
||||
}]
|
||||
"yml_content": "name: Test Metric\ndescription: A test metric"
|
||||
}]
|
||||
});
|
||||
let valid_args = serde_json::to_string(&valid_params).unwrap();
|
||||
let result = serde_json::from_str::<ModifyFilesParams>(&valid_args);
|
||||
let result = serde_json::from_str::<UpdateFilesParams>(&valid_args);
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Test missing required fields
|
||||
let missing_fields_params = json!({
|
||||
"files": [{
|
||||
"id": Uuid::new_v4().to_string(),
|
||||
"file_name": "test.yml"
|
||||
// missing modifications
|
||||
"id": Uuid::new_v4().to_string()
|
||||
// missing yml_content
|
||||
}]
|
||||
});
|
||||
let missing_args = serde_json::to_string(&missing_fields_params).unwrap();
|
||||
let result = serde_json::from_str::<ModifyFilesParams>(&missing_args);
|
||||
let result = serde_json::from_str::<UpdateFilesParams>(&missing_args);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue