From a2520c2efa02744fc780a51db2aa3d88c994e41d Mon Sep 17 00:00:00 2001 From: dal Date: Fri, 21 Mar 2025 11:12:14 -0600 Subject: [PATCH] Clean up agents library --- api/libs/agents/src/agent.rs | 9 ++++----- .../agents/src/tools/categories/file_tools/common.rs | 6 +++--- .../tools/categories/file_tools/create_dashboards.rs | 9 +++------ .../src/tools/categories/file_tools/create_metrics.rs | 7 ++----- .../tools/categories/file_tools/file_types/file.rs | 2 +- .../tools/categories/file_tools/modify_dashboards.rs | 9 +++------ .../src/tools/categories/file_tools/modify_metrics.rs | 11 ++++------- .../categories/file_tools/search_data_catalog.rs | 2 +- .../tools/categories/planning_tools/create_plan.rs | 7 ++----- 9 files changed, 23 insertions(+), 39 deletions(-) diff --git a/api/libs/agents/src/agent.rs b/api/libs/agents/src/agent.rs index e709b2cfc..843f2fbb7 100644 --- a/api/libs/agents/src/agent.rs +++ b/api/libs/agents/src/agent.rs @@ -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; } } diff --git a/api/libs/agents/src/tools/categories/file_tools/common.rs b/api/libs/agents/src/tools/categories/file_tools/common.rs index 1337dcaf5..33aede7f9 100644 --- a/api/libs/agents/src/tools/categories/file_tools/common.rs +++ b/api/libs/agents/src/tools/categories/file_tools/common.rs @@ -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, diff --git a/api/libs/agents/src/tools/categories/file_tools/create_dashboards.rs b/api/libs/agents/src/tools/categories/file_tools/create_dashboards.rs index 10b4b6569..5a7a74c1f 100644 --- a/api/libs/agents/src/tools/categories/file_tools/create_dashboards.rs +++ b/api/libs/agents/src/tools/categories/file_tools/create_dashboards.rs @@ -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 { diff --git a/api/libs/agents/src/tools/categories/file_tools/create_metrics.rs b/api/libs/agents/src/tools/categories/file_tools/create_metrics.rs index d7ff24c6f..55f7d17c4 100644 --- a/api/libs/agents/src/tools/categories/file_tools/create_metrics.rs +++ b/api/libs/agents/src/tools/categories/file_tools/create_metrics.rs @@ -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 { diff --git a/api/libs/agents/src/tools/categories/file_tools/file_types/file.rs b/api/libs/agents/src/tools/categories/file_tools/file_types/file.rs index 6b6a42c3a..8656ab062 100644 --- a/api/libs/agents/src/tools/categories/file_tools/file_types/file.rs +++ b/api/libs/agents/src/tools/categories/file_tools/file_types/file.rs @@ -30,7 +30,7 @@ pub struct FileWithId { #[derive(Debug, Serialize, Deserialize)] #[serde(untagged)] pub enum FileEnum { - Metric(MetricYml), + Metric(Box), Dashboard(DashboardYml), } diff --git a/api/libs/agents/src/tools/categories/file_tools/modify_dashboards.rs b/api/libs/agents/src/tools/categories/file_tools/modify_dashboards.rs index f0df48819..210239c12 100644 --- a/api/libs/agents/src/tools/categories/file_tools/modify_dashboards.rs +++ b/api/libs/agents/src/tools/categories/file_tools/modify_dashboards.rs @@ -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 { + async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result { let start_time = Instant::now(); debug!("Starting file modification execution"); diff --git a/api/libs/agents/src/tools/categories/file_tools/modify_metrics.rs b/api/libs/agents/src/tools/categories/file_tools/modify_metrics.rs index 707e2e7a9..72720a647 100644 --- a/api/libs/agents/src/tools/categories/file_tools/modify_metrics.rs +++ b/api/libs/agents/src/tools/categories/file_tools/modify_metrics.rs @@ -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 { + async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result { 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, diff --git a/api/libs/agents/src/tools/categories/file_tools/search_data_catalog.rs b/api/libs/agents/src/tools/categories/file_tools/search_data_catalog.rs index 35a80f702..50907bf51 100644 --- a/api/libs/agents/src/tools/categories/file_tools/search_data_catalog.rs +++ b/api/libs/agents/src/tools/categories/file_tools/search_data_catalog.rs @@ -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 { + async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result { let start_time = Instant::now(); // Fetch all non-deleted datasets diff --git a/api/libs/agents/src/tools/categories/planning_tools/create_plan.rs b/api/libs/agents/src/tools/categories/planning_tools/create_plan.rs index 3ca0656b6..a10a3abd2 100644 --- a/api/libs/agents/src/tools/categories/planning_tools/create_plan.rs +++ b/api/libs/agents/src/tools/categories/planning_tools/create_plan.rs @@ -38,7 +38,7 @@ impl ToolExecutor for CreatePlan { "create_plan".to_string() } - async fn execute(&self, params: Self::Params, tool_call_id: String) -> Result { + async fn execute(&self, params: Self::Params, _tool_call_id: String) -> Result { 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 {