mirror of https://github.com/buster-so/buster.git
split out todos from plan tool
This commit is contained in:
parent
55476793d5
commit
526c6c3539
|
@ -133,7 +133,6 @@ To determine whether to use a Straightforward Plan or an Investigative Plan, con
|
|||
- When creating a plan that involves generating assets (visualizations and dashboards), do not include a separate step for delivering these assets, as they are automatically displayed to the user upon creation.
|
||||
- Assume that all datasets required for the plan are available, as their availability has already been confirmed in the previous step.
|
||||
- If the user's request includes aspects that are not supported (e.g., specific visualizations, forecasts, etc.), do not include these in the step-by-step plan. Instead, mention them in the note section of the plan, and specify that they should be addressed in the final response to the user.
|
||||
- The tools used for creating plans include a `todos` argument. This argument is a list of short summary points. **Crucially, each step in the generated plan must correspond to exactly one item in the `todos` list.** These `todos` serve as a concise overview of the plan's execution steps. Do not include any review steps in the `todos` list, as reviews are handled separately.
|
||||
|
||||
**Examples**
|
||||
|
||||
|
|
|
@ -3,7 +3,9 @@ use async_trait::async_trait;
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
use tracing::warn;
|
||||
|
||||
use super::helpers::generate_todos_from_plan;
|
||||
use crate::{agent::Agent, tools::ToolExecutor};
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
|
@ -14,9 +16,7 @@ pub struct CreatePlanInvestigativeOutput {
|
|||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CreatePlanInvestigativeInput {
|
||||
#[serde(rename = "plan")]
|
||||
_plan: String,
|
||||
todos: Vec<String>,
|
||||
plan: String,
|
||||
}
|
||||
|
||||
pub struct CreatePlanInvestigative {
|
||||
|
@ -43,22 +43,39 @@ impl ToolExecutor for CreatePlanInvestigative {
|
|||
.set_state_value(String::from("plan_available"), Value::Bool(true))
|
||||
.await;
|
||||
|
||||
let todos_state_objects: Vec<Value> = params
|
||||
.todos
|
||||
.iter()
|
||||
.map(|item| {
|
||||
let mut map = serde_json::Map::new();
|
||||
map.insert("completed".to_string(), Value::Bool(false));
|
||||
map.insert("todo".to_string(), Value::String(item.clone()));
|
||||
Value::Object(map)
|
||||
})
|
||||
.collect();
|
||||
let mut todos_string = String::new();
|
||||
|
||||
self.agent
|
||||
.set_state_value(String::from("todos"), Value::Array(todos_state_objects))
|
||||
.await;
|
||||
match generate_todos_from_plan(
|
||||
¶ms.plan,
|
||||
self.agent.get_user_id(),
|
||||
self.agent.get_session_id(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(todos_state_objects) => {
|
||||
let formatted_todos: Vec<String> = todos_state_objects
|
||||
.iter()
|
||||
.filter_map(|val| val.as_object())
|
||||
.filter_map(|obj| obj.get("todo"))
|
||||
.filter_map(|todo_val| todo_val.as_str())
|
||||
.map(|todo_str| format!("[ ] {}", todo_str))
|
||||
.collect();
|
||||
todos_string = formatted_todos.join("\n");
|
||||
|
||||
let todos_string = params.todos.iter().map(|item| format!("[ ] {}", item)).collect::<Vec<_>>().join("\n");
|
||||
self.agent
|
||||
.set_state_value(String::from("todos"), Value::Array(todos_state_objects))
|
||||
.await;
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"Failed to generate todos from plan using LLM: {}. Proceeding without todos.",
|
||||
e
|
||||
);
|
||||
self.agent
|
||||
.set_state_value(String::from("todos"), Value::Array(vec![]))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(CreatePlanInvestigativeOutput { success: true, todos: todos_string })
|
||||
}
|
||||
|
@ -71,19 +88,13 @@ impl ToolExecutor for CreatePlanInvestigative {
|
|||
"parameters": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"plan",
|
||||
"todos"
|
||||
"plan"
|
||||
],
|
||||
"properties": {
|
||||
"plan": {
|
||||
"type": "string",
|
||||
"description": get_plan_investigative_description().await
|
||||
},
|
||||
"todos": {
|
||||
"type": "array",
|
||||
"description": "Ordered todo points summarizing the steps of the plan. There should be max one todo for each step in the plan, in order. For example, if the plan has two steps, plan_todos should have two items, each summarizing a step. Do not include review or response steps—these will be handled by a separate agent.",
|
||||
"items": { "type": "string" },
|
||||
},
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
}
|
||||
|
|
|
@ -3,7 +3,9 @@ use async_trait::async_trait;
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
use tracing::warn;
|
||||
|
||||
use super::helpers::generate_todos_from_plan;
|
||||
use crate::{agent::Agent, tools::ToolExecutor};
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
|
@ -14,9 +16,7 @@ pub struct CreatePlanStraightforwardOutput {
|
|||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CreatePlanStraightforwardInput {
|
||||
#[serde(rename = "plan")]
|
||||
_plan: String,
|
||||
todos: Vec<String>,
|
||||
plan: String,
|
||||
}
|
||||
|
||||
pub struct CreatePlanStraightforward {
|
||||
|
@ -43,32 +43,41 @@ impl ToolExecutor for CreatePlanStraightforward {
|
|||
.set_state_value(String::from("plan_available"), Value::Bool(true))
|
||||
.await;
|
||||
|
||||
let todos_state_objects: Vec<Value> = params
|
||||
.todos
|
||||
.iter()
|
||||
.map(|item| {
|
||||
let mut map = serde_json::Map::new();
|
||||
map.insert("completed".to_string(), Value::Bool(false));
|
||||
map.insert("todo".to_string(), Value::String(item.clone()));
|
||||
Value::Object(map)
|
||||
})
|
||||
.collect();
|
||||
let mut todos_string = String::new();
|
||||
|
||||
self.agent
|
||||
.set_state_value(String::from("todos"), Value::Array(todos_state_objects))
|
||||
.await;
|
||||
match generate_todos_from_plan(
|
||||
¶ms.plan,
|
||||
self.agent.get_user_id(),
|
||||
self.agent.get_session_id(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(todos_state_objects) => {
|
||||
let formatted_todos: Vec<String> = todos_state_objects
|
||||
.iter()
|
||||
.filter_map(|val| val.as_object())
|
||||
.filter_map(|obj| obj.get("todo"))
|
||||
.filter_map(|todo_val| todo_val.as_str())
|
||||
.map(|todo_str| format!("[ ] {}", todo_str))
|
||||
.collect();
|
||||
todos_string = formatted_todos.join("\n");
|
||||
|
||||
let todos_string = params
|
||||
.todos
|
||||
.iter()
|
||||
.map(|item| format!("[ ] {}", item))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
self.agent
|
||||
.set_state_value(String::from("todos"), Value::Array(todos_state_objects))
|
||||
.await;
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"Failed to generate todos from plan using LLM: {}. Proceeding without todos.",
|
||||
e
|
||||
);
|
||||
self.agent
|
||||
.set_state_value(String::from("todos"), Value::Array(vec![]))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(CreatePlanStraightforwardOutput {
|
||||
success: true,
|
||||
todos: todos_string,
|
||||
})
|
||||
Ok(CreatePlanStraightforwardOutput { success: true, todos: todos_string })
|
||||
}
|
||||
|
||||
async fn get_schema(&self) -> Value {
|
||||
|
@ -79,18 +88,13 @@ impl ToolExecutor for CreatePlanStraightforward {
|
|||
"parameters": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"plan",
|
||||
"todos"
|
||||
"plan"
|
||||
],
|
||||
"properties": {
|
||||
"plan": {
|
||||
"type": "string",
|
||||
"description": get_plan_straightforward_description().await
|
||||
},
|
||||
"todos": {
|
||||
"type": "array",
|
||||
"description": "Ordered todo points summarizing the steps of the plan. There should be max one todo for each step in the plan, in order. For example, if the plan has two steps, plan_todos should have two items, each summarizing a step. Do not include review or response steps—these will be handled by a separate agent.", "items": { "type": "string" },
|
||||
},
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
}
|
||||
|
|
|
@ -0,0 +1,7 @@
|
|||
// Intentionally left empty or for re-exporting modules
|
||||
|
||||
// Declare the new module
|
||||
pub mod todo_generator;
|
||||
|
||||
// Re-export the function for easier access
|
||||
pub use todo_generator::generate_todos_from_plan;
|
|
@ -0,0 +1,104 @@
|
|||
use anyhow::Result;
|
||||
use litellm::{AgentMessage, ChatCompletionRequest, LiteLLMClient, Metadata, ResponseFormat};
|
||||
use serde_json::Value;
|
||||
use tracing::{error, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Generates a list of todo items (as JSON Values for agent state) from a plan string using an LLM.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `plan` - The plan string generated by the primary LLM.
|
||||
/// * `user_id` - The ID of the user.
|
||||
/// * `session_id` - The ID of the current session.
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// A `Result` containing a `Vec<Value>` where each `Value` is a JSON object representing a todo item
|
||||
/// (`{"completed": false, "todo": "..."}`), or an error if generation or parsing fails.
|
||||
pub async fn generate_todos_from_plan(
|
||||
plan: &str,
|
||||
user_id: Uuid,
|
||||
session_id: Uuid,
|
||||
) -> Result<Vec<Value>> {
|
||||
let llm_client = LiteLLMClient::new(None, None);
|
||||
|
||||
let prompt = format!(
|
||||
r#"
|
||||
Given the following plan, extract the main actionable steps and return them as a JSON list of concise todo strings. Focus on the core actions described in each step. Do not include any introductory text, summary, or review steps. Only include the main tasks to be performed.
|
||||
|
||||
Plan:
|
||||
"""
|
||||
{}
|
||||
"""
|
||||
|
||||
Return ONLY a valid JSON array of strings, where each string is a short todo item corresponding to a main step in the plan.
|
||||
Example format: `["Create 11 visualizations", "Create dashboard"]`
|
||||
"#,
|
||||
plan
|
||||
);
|
||||
|
||||
let request = ChatCompletionRequest {
|
||||
model: "gemini-2.0-flash-001".to_string(),
|
||||
messages: vec![AgentMessage::User { id: None, content: prompt, name: None }],
|
||||
stream: Some(false),
|
||||
response_format: Some(ResponseFormat { type_: "json_object".to_string(), json_schema: None }),
|
||||
metadata: Some(Metadata {
|
||||
generation_name: "generate_todos_from_plan".to_string(),
|
||||
user_id: user_id.to_string(),
|
||||
session_id: session_id.to_string(),
|
||||
trace_id: Uuid::new_v4().to_string(),
|
||||
}),
|
||||
max_completion_tokens: Some(1024),
|
||||
temperature: Some(0.0),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let response = llm_client.chat_completion(request).await?;
|
||||
|
||||
let content = match response.choices.get(0).and_then(|c| c.message.get_content()) {
|
||||
Some(content) => content,
|
||||
None => return Err(anyhow::anyhow!("LLM response for todo generation was empty or malformed")),
|
||||
};
|
||||
|
||||
// Assuming the LLM returns a JSON object like `{"todos": ["...", "..."]}` or just the array `["", ""]`
|
||||
let parsed_value: Value = serde_json::from_str(&content).map_err(|e| {
|
||||
error!("Failed to parse LLM JSON response for todos: {}. Content: {}", e, content);
|
||||
anyhow::anyhow!("Failed to parse LLM JSON response for todos: {}", e)
|
||||
})?;
|
||||
|
||||
let todo_strings: Vec<String> = match parsed_value {
|
||||
Value::Array(arr) => arr
|
||||
.into_iter()
|
||||
.filter_map(|v| v.as_str().map(String::from))
|
||||
.collect(),
|
||||
Value::Object(mut map) => map
|
||||
.remove("todos") // Attempt to extract from a common pattern like {"todos": [...]}
|
||||
.and_then(|v| v.as_array().cloned())
|
||||
.map(|arr| {
|
||||
arr.into_iter()
|
||||
.filter_map(|v| v.as_str().map(String::from))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_else(|| {
|
||||
warn!("LLM todo response was object but did not contain a 'todos' array or it was not an array of strings. Content: {}", content);
|
||||
vec![]
|
||||
}),
|
||||
_ => {
|
||||
warn!("Unexpected JSON structure for todos from LLM. Content: {}", content);
|
||||
return Err(anyhow::anyhow!("Unexpected JSON structure for todos from LLM"));
|
||||
}
|
||||
};
|
||||
|
||||
let todos_state_objects: Vec<Value> = todo_strings
|
||||
.into_iter()
|
||||
.map(|item| {
|
||||
let mut map = serde_json::Map::new();
|
||||
map.insert("completed".to_string(), Value::Bool(false));
|
||||
map.insert("todo".to_string(), Value::String(item));
|
||||
Value::Object(map)
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(todos_state_objects)
|
||||
}
|
|
@ -1,7 +1,9 @@
|
|||
pub mod create_plan_investigative;
|
||||
pub mod create_plan_straightforward;
|
||||
pub mod review_plan;
|
||||
pub mod helpers;
|
||||
|
||||
pub use create_plan_investigative::*;
|
||||
pub use create_plan_straightforward::*;
|
||||
pub use review_plan::*;
|
||||
pub use review_plan::*;
|
||||
pub use helpers::*;
|
Loading…
Reference in New Issue