Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

modify files going good

This commit is contained in:
dal 2025-04-04 09:44:10 -06:00
parent fddc909f0f
commit 38f1c9c54c
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
8 changed files with 56 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
api/libs/agents/src/tools/categories/file_tools/common.rs
View File

@ -442,6 +442,7 @@ definitions:
metric_value_aggregate:
type: string
enum: ["sum", "average", "median", "max", "min", "count", "first"]
description: "Optional - only used when the user specifically requests it, otherwise leave blank"
required:
- selected_chart_type
- metric_column_id

1
api/libs/agents/src/tools/categories/file_tools/create_dashboards.rs
View File

@ -261,6 +261,7 @@ impl ToolExecutor for CreateDashboardFilesTool {
results: None,
created_at: dashboard_records[i].created_at,
updated_at: dashboard_records[i].updated_at,
version_number: dashboard_records[i].version_history.get_version_number(),
});
}
}

1
api/libs/agents/src/tools/categories/file_tools/create_metrics.rs
View File

@ -176,6 +176,7 @@ impl ToolExecutor for CreateMetricFilesTool {
results: Some(results_vec[i].1.clone()),
created_at: metric_records[i].created_at,
updated_at: metric_records[i].updated_at,
version_number: metric_records[i].version_history.get_version_number(),
});
}
}

1
api/libs/agents/src/tools/categories/file_tools/file_types/file.rs
View File

@ -24,6 +24,7 @@ pub struct FileWithId {
#[serde(skip_serializing_if = "Option::is_none")]
pub results: Option<Vec<IndexMap<String, DataType>>>,
pub created_at: DateTime<Utc>,
pub version_number: i32,
pub updated_at: DateTime<Utc>,
}

57
api/libs/agents/src/tools/categories/file_tools/modify_dashboards.rs
View File

@ -1,5 +1,5 @@
use std::{env, sync::Arc};
use std::time::Instant;
use std::{env, sync::Arc};
use anyhow::Result;
use async_trait::async_trait;
@ -16,8 +16,8 @@ use tracing::{debug, error, info};
use super::{
common::{
ModificationResult,
ModifyFilesOutput, ModifyFilesParams, process_dashboard_file_modification,
process_dashboard_file_modification, ModificationResult, ModifyFilesOutput,
ModifyFilesParams,
},
file_types::file::FileWithId,
FileModificationTool,
@ -59,10 +59,13 @@ impl ToolExecutor for ModifyDashboardFilesTool {
}
async fn is_enabled(&self) -> bool {
matches!((
self.agent.get_state_value("dashboards_available").await,
self.agent.get_state_value("plan_available").await,
), (Some(_), Some(_)))
matches!(
(
self.agent.get_state_value("dashboards_available").await,
self.agent.get_state_value("plan_available").await,
),
(Some(_), Some(_))
)
}
async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result<Self::Output> {
@ -96,22 +99,38 @@ 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(
dashboard_file,
dashboard_file.clone(),
&modification,
duration,
)
.await
{
Ok((
dashboard_file,
mut dashboard_file,
dashboard_yml,
results,
validation_message,
validation_results,
)) => {
// Calculate next version number from existing version history
let next_version =
match dashboard_file.version_history.get_latest_version() {
Some(version) => version.version_number + 1,
None => 1,
};
// Add new version to history
dashboard_file
.version_history
.add_version(next_version, dashboard_yml.clone());
// Ensure the name field is updated
// This is redundant but ensures the name is set correctly
dashboard_file.name = dashboard_yml.name.clone();
batch.files.push(dashboard_file);
batch.ymls.push(dashboard_yml);
batch.modification_results.extend(results);
@ -134,7 +153,7 @@ impl ToolExecutor for ModifyDashboardFilesTool {
}
}
// Update dashboard files in database
// Update dashboard files in database with version history
if !batch.files.is_empty() {
use diesel::insert_into;
match insert_into(dashboard_files::table)
@ -144,12 +163,16 @@ impl ToolExecutor for ModifyDashboardFilesTool {
.set((
dashboard_files::content.eq(excluded(dashboard_files::content)),
dashboard_files::updated_at.eq(excluded(dashboard_files::updated_at)),
// Add version history and name to ensure they're updated
dashboard_files::version_history.eq(excluded(dashboard_files::version_history)),
// Explicitly set name even though it's in content to ensure it's updated in case content parsing fails
dashboard_files::name.eq(excluded(dashboard_files::name)),
))
.execute(&mut conn)
.await
{
Ok(_) => {
debug!("Successfully updated dashboard files in database");
debug!("Successfully updated dashboard files with versioning");
}
Err(e) => {
error!("Failed to update dashboard files in database: {}", e);
@ -164,7 +187,10 @@ impl ToolExecutor for ModifyDashboardFilesTool {
// Generate output
let duration = start_time.elapsed().as_millis() as i64;
let mut output = ModifyFilesOutput {
message: format!("Modified {} dashboard files", batch.files.len()),
message: format!(
"Modified {} dashboard files and created new versions",
batch.files.len()
),
duration,
files: Vec::new(),
};
@ -183,6 +209,7 @@ impl ToolExecutor for ModifyDashboardFilesTool {
results: Some(batch.validation_results[i].clone()),
created_at: file.created_at,
updated_at: file.updated_at,
version_number: file.version_history.get_version_number(),
}
}));
@ -335,7 +362,8 @@ async fn get_modify_dashboards_new_content_description() -> 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. Must match exactly.".to_string();
return "The exact content in the file that should be replaced. Must match exactly."
.to_string();
}
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
@ -432,7 +460,6 @@ mod tests {
#[test]
fn test_tool_parameter_validation() {
// Test valid parameters
let valid_params = json!({
"files": [{

3
api/libs/agents/src/tools/categories/file_tools/modify_metrics.rs
View File

@ -148,7 +148,7 @@ impl ToolExecutor for ModifyMetricFilesTool {
// 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];
@ -279,6 +279,7 @@ impl ToolExecutor for ModifyMetricFilesTool {
results: Some(batch.validation_results[i].clone()),
created_at: file.created_at,
updated_at: file.updated_at,
version_number: file.version_history.get_version_number(),
}
}));

6
api/libs/database/src/types/version_history.rs
View File

@ -74,6 +74,12 @@ impl VersionHistory {
self.0.values().max_by_key(|v| v.version_number)
}
pub fn get_version_number(&self) -> i32 {
self.get_latest_version()
.map(|version| version.version_number)
.unwrap_or(1)
}
/// Updates the content of the latest version without creating a new version
///
/// This is used when we want to overwrite the current version instead of creating a new one.

4
api/libs/handlers/src/chats/post_chat_handler.rs
View File

@ -1298,8 +1298,8 @@ fn tool_modify_metrics(id: String, content: String) -> Result<Vec<BusterReasonin
let buster_file = BusterFile {
id: file_id.clone(),
file_type: "metric".to_string(),
file_name: file.name.clone(),
version_number: 1,
file_name: file.name.clone(), // Use the updated name from the file
version_number: file.version_number,
status: "completed".to_string(),
file: BusterFileContent {
text: Some(file.yml_content),