mirror of https://github.com/buster-so/buster.git
Clean up agents library
This commit is contained in:
parent
f09dfdd543
commit
a2520c2efa
|
@ -554,15 +554,14 @@ impl Agent {
|
|||
let id = tool_call.id.clone().unwrap_or_else(|| {
|
||||
buffer.tool_calls
|
||||
.keys()
|
||||
.next()
|
||||
.map(|s| s.clone())
|
||||
.next().cloned()
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string())
|
||||
});
|
||||
|
||||
// Get or create the pending tool call
|
||||
let pending_call = buffer.tool_calls
|
||||
.entry(id.clone())
|
||||
.or_insert_with(PendingToolCall::new);
|
||||
.or_default();
|
||||
|
||||
// Update the pending call with the delta
|
||||
pending_call.update_from_delta(tool_call);
|
||||
|
@ -922,10 +921,10 @@ impl PendingToolCall {
|
|||
self.arguments.push_str(args);
|
||||
}
|
||||
}
|
||||
if let Some(_) = &tool_call.code_interpreter {
|
||||
if tool_call.code_interpreter.is_some() {
|
||||
self.code_interpreter = None;
|
||||
}
|
||||
if let Some(_) = &tool_call.retrieval {
|
||||
if tool_call.retrieval.is_some() {
|
||||
self.retrieval = None;
|
||||
}
|
||||
}
|
||||
|
|
|
@ -48,7 +48,7 @@ pub async fn validate_sql(
|
|||
};
|
||||
|
||||
// Try to execute the query using query_engine
|
||||
let results = match query_engine(&data_source_id, &sql.to_string(), None).await {
|
||||
let results = match query_engine(&data_source_id, sql, None).await {
|
||||
Ok(results) => results,
|
||||
Err(e) => return Err(anyhow!("SQL validation failed: {}", e)),
|
||||
};
|
||||
|
@ -590,7 +590,7 @@ pub async fn process_metric_file(
|
|||
Err(e) => return Err(format!("Invalid SQL query: {}", e)),
|
||||
};
|
||||
|
||||
let metric_yml_json = match serde_json::to_value(metric_yml.clone()) {
|
||||
let _metric_yml_json = match serde_json::to_value(metric_yml.clone()) {
|
||||
Ok(json) => json,
|
||||
Err(e) => return Err(format!("Failed to process metric: {}", e)),
|
||||
};
|
||||
|
@ -600,7 +600,7 @@ pub async fn process_metric_file(
|
|||
name: file_name.clone(),
|
||||
file_name: file_name.clone(),
|
||||
content: metric_yml.clone(),
|
||||
created_by: user_id.clone(),
|
||||
created_by: *user_id,
|
||||
verification: Verification::NotRequested,
|
||||
evaluation_obj: None,
|
||||
evaluation_summary: None,
|
||||
|
|
|
@ -107,7 +107,7 @@ async fn process_dashboard_file(
|
|||
content: dashboard_yml.clone(),
|
||||
filter: None,
|
||||
organization_id: Uuid::new_v4(),
|
||||
created_by: user_id.clone(),
|
||||
created_by: *user_id,
|
||||
created_at: Utc::now(),
|
||||
updated_at: Utc::now(),
|
||||
deleted_at: None,
|
||||
|
@ -130,13 +130,10 @@ impl ToolExecutor for CreateDashboardFilesTool {
|
|||
}
|
||||
|
||||
async fn is_enabled(&self) -> bool {
|
||||
match (
|
||||
matches!((
|
||||
self.agent.get_state_value("metrics_available").await,
|
||||
self.agent.get_state_value("plan_available").await,
|
||||
) {
|
||||
(Some(_), Some(_)) => true,
|
||||
_ => false,
|
||||
}
|
||||
), (Some(_), Some(_)))
|
||||
}
|
||||
|
||||
async fn execute(&self, params: Self::Params, tool_call_id: String) -> Result<Self::Output> {
|
||||
|
|
|
@ -74,13 +74,10 @@ impl ToolExecutor for CreateMetricFilesTool {
|
|||
}
|
||||
|
||||
async fn is_enabled(&self) -> bool {
|
||||
match (
|
||||
matches!((
|
||||
self.agent.get_state_value("data_context").await,
|
||||
self.agent.get_state_value("plan_available").await,
|
||||
) {
|
||||
(Some(_), Some(_)) => true,
|
||||
_ => false,
|
||||
}
|
||||
), (Some(_), Some(_)))
|
||||
}
|
||||
|
||||
async fn execute(&self, params: Self::Params, tool_call_id: String) -> Result<Self::Output> {
|
||||
|
|
|
@ -30,7 +30,7 @@ pub struct FileWithId {
|
|||
#[derive(Debug, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum FileEnum {
|
||||
Metric(MetricYml),
|
||||
Metric(Box<MetricYml>),
|
||||
Dashboard(DashboardYml),
|
||||
}
|
||||
|
||||
|
|
|
@ -59,16 +59,13 @@ impl ToolExecutor for ModifyDashboardFilesTool {
|
|||
}
|
||||
|
||||
async fn is_enabled(&self) -> bool {
|
||||
match (
|
||||
matches!((
|
||||
self.agent.get_state_value("dashboards_available").await,
|
||||
self.agent.get_state_value("plan_available").await,
|
||||
) {
|
||||
(Some(_), Some(_)) => true,
|
||||
_ => false,
|
||||
}
|
||||
), (Some(_), Some(_)))
|
||||
}
|
||||
|
||||
async fn execute(&self, params: Self::Params, tool_call_id: String) -> Result<Self::Output> {
|
||||
async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result<Self::Output> {
|
||||
let start_time = Instant::now();
|
||||
|
||||
debug!("Starting file modification execution");
|
||||
|
|
|
@ -60,16 +60,13 @@ impl ToolExecutor for ModifyMetricFilesTool {
|
|||
}
|
||||
|
||||
async fn is_enabled(&self) -> bool {
|
||||
match (
|
||||
matches!((
|
||||
self.agent.get_state_value("metrics_available").await,
|
||||
self.agent.get_state_value("plan_available").await,
|
||||
) {
|
||||
(Some(_), Some(_)) => true,
|
||||
_ => false,
|
||||
}
|
||||
), (Some(_), Some(_)))
|
||||
}
|
||||
|
||||
async fn execute(&self, params: Self::Params, tool_call_id: String) -> Result<Self::Output> {
|
||||
async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result<Self::Output> {
|
||||
let start_time = Instant::now();
|
||||
|
||||
debug!("Starting file modification execution");
|
||||
|
@ -157,7 +154,7 @@ impl ToolExecutor for ModifyMetricFilesTool {
|
|||
|
||||
// Process results and generate output message
|
||||
let duration = start_time.elapsed().as_millis() as i64;
|
||||
let output = ModifyFilesOutput {
|
||||
let _output = ModifyFilesOutput {
|
||||
message: String::new(),
|
||||
files: Vec::new(),
|
||||
duration,
|
||||
|
|
|
@ -270,7 +270,7 @@ impl ToolExecutor for SearchDataCatalogTool {
|
|||
type Output = SearchDataCatalogOutput;
|
||||
type Params = SearchDataCatalogParams;
|
||||
|
||||
async fn execute(&self, params: Self::Params, tool_call_id: String) -> Result<Self::Output> {
|
||||
async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result<Self::Output> {
|
||||
let start_time = Instant::now();
|
||||
|
||||
// Fetch all non-deleted datasets
|
||||
|
|
|
@ -38,7 +38,7 @@ impl ToolExecutor for CreatePlan {
|
|||
"create_plan".to_string()
|
||||
}
|
||||
|
||||
async fn execute(&self, params: Self::Params, tool_call_id: String) -> Result<Self::Output> {
|
||||
async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result<Self::Output> {
|
||||
self.agent
|
||||
.set_state_value(String::from("plan_available"), Value::Bool(true))
|
||||
.await;
|
||||
|
@ -50,10 +50,7 @@ impl ToolExecutor for CreatePlan {
|
|||
}
|
||||
|
||||
async fn is_enabled(&self) -> bool {
|
||||
match self.agent.get_state_value("data_context").await {
|
||||
Some(_) => true,
|
||||
None => false,
|
||||
}
|
||||
self.agent.get_state_value("data_context").await.is_some()
|
||||
}
|
||||
|
||||
async fn get_schema(&self) -> Value {
|
||||
|
|
Loading…
Reference in New Issue