plan being returned better.

This commit is contained in:
dal 2025-03-10 12:02:11 -06:00
parent 090f010542
commit dbb05eab4f
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
2 changed files with 25 additions and 13 deletions

View File

@ -8,7 +8,7 @@ use crate::{agent::Agent, tools::ToolExecutor};
#[derive(Debug, Serialize, Deserialize)] #[derive(Debug, Serialize, Deserialize)]
pub struct CreatePlanOutput { pub struct CreatePlanOutput {
message: String, pub message: String,
} }
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]

View File

@ -3,9 +3,10 @@ use once_cell::sync::OnceCell;
use std::{collections::HashMap, sync::Mutex, time::Instant}; use std::{collections::HashMap, sync::Mutex, time::Instant};
use agents::{ use agents::{
tools::file_tools::{ tools::{file_tools::{
common::ModifyFilesOutput, create_dashboard_files::CreateDashboardFilesOutput, create_metric_files::CreateMetricFilesOutput, search_data_catalog::SearchDataCatalogOutput common::ModifyFilesOutput, create_dashboard_files::CreateDashboardFilesOutput,
}, create_metric_files::CreateMetricFilesOutput, search_data_catalog::SearchDataCatalogOutput,
}, planning_tools::CreatePlanOutput},
AgentExt, AgentMessage, AgentThread, BusterSuperAgent, AgentExt, AgentMessage, AgentThread, BusterSuperAgent,
}; };
@ -932,13 +933,29 @@ fn transform_tool_message(
"modify_metrics" => tool_modify_metrics(id.clone(), content)?, "modify_metrics" => tool_modify_metrics(id.clone(), content)?,
"create_dashboards" => tool_create_dashboards(id.clone(), content)?, "create_dashboards" => tool_create_dashboards(id.clone(), content)?,
"modify_dashboards" => tool_modify_dashboards(id.clone(), content)?, "modify_dashboards" => tool_modify_dashboards(id.clone(), content)?,
"create_plan" => return Ok(vec![]), // Return empty vec wrapped in Ok "create_plan" => tool_create_plan(id.clone(), content)?,
_ => return Err(anyhow::anyhow!("Unknown tool name: {}", name)), _ => return Err(anyhow::anyhow!("Unknown tool name: {}", name)),
}; };
Ok(messages) Ok(messages)
} }
fn tool_create_plan(id: String, content: String) -> Result<Vec<BusterReasoningMessage>> {
println!("MESSAGE_STREAM: Processing tool create plan message");
let buster_file = BusterReasoningMessage::Text(BusterReasoningText {
id,
reasoning_type: "text".to_string(),
title: "Plan".to_string(),
secondary_title: "".to_string(),
message: None,
message_chunk: None,
status: Some("completed".to_string()),
});
Ok(vec![buster_file])
}
// Update tool_create_metrics to require ID // Update tool_create_metrics to require ID
fn tool_create_metrics(id: String, content: String) -> Result<Vec<BusterReasoningMessage>> { fn tool_create_metrics(id: String, content: String) -> Result<Vec<BusterReasoningMessage>> {
println!("MESSAGE_STREAM: Processing tool create metrics message"); println!("MESSAGE_STREAM: Processing tool create metrics message");
@ -1317,12 +1334,7 @@ fn transform_assistant_tool_message(
.or(text.message_chunk.clone()); .or(text.message_chunk.clone());
text.message_chunk = None; text.message_chunk = None;
// Only set completed status for create_plan when message is present text.status = Some("loading".to_string());
if tool_call.function.name == "create_plan" && text.message.is_some() {
text.status = Some("completed".to_string());
} else {
text.status = Some("loading".to_string());
}
tracker.clear_chunk(text.id.clone()); tracker.clear_chunk(text.id.clone());
Some(BusterReasoningMessage::Text(text)) Some(BusterReasoningMessage::Text(text))