Clean up agents library

This commit is contained in:
dal 2025-03-21 11:12:14 -06:00
parent f09dfdd543
commit a2520c2efa
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
9 changed files with 23 additions and 39 deletions

View File

@ -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;
}
}

View File

@ -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,

View File

@ -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> {

View File

@ -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> {

View File

@ -30,7 +30,7 @@ pub struct FileWithId {
#[derive(Debug, Serialize, Deserialize)]
#[serde(untagged)]
pub enum FileEnum {
Metric(MetricYml),
Metric(Box<MetricYml>),
Dashboard(DashboardYml),
}

View File

@ -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");

View File

@ -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,

View File

@ -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

View File

@ -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 {