From 526c6c3539dbe2037b06a9f1fc57bae03bd4e6b6 Mon Sep 17 00:00:00 2001 From: dal Date: Wed, 16 Apr 2025 09:41:40 -0600 Subject: [PATCH] split out todos from plan tool --- .../src/agents/modes/create_plan_prompt.rs | 1 - .../create_plan_investigative.rs | 61 +++++----- .../create_plan_straightforward.rs | 70 ++++++------ .../categories/planning_tools/helpers/mod.rs | 7 ++ .../planning_tools/helpers/todo_generator.rs | 104 ++++++++++++++++++ .../tools/categories/planning_tools/mod.rs | 4 +- 6 files changed, 187 insertions(+), 60 deletions(-) create mode 100644 api/libs/agents/src/tools/categories/planning_tools/helpers/mod.rs create mode 100644 api/libs/agents/src/tools/categories/planning_tools/helpers/todo_generator.rs diff --git a/api/libs/agents/src/agents/modes/create_plan_prompt.rs b/api/libs/agents/src/agents/modes/create_plan_prompt.rs index 1399a3fd8..6c7e7d937 100644 --- a/api/libs/agents/src/agents/modes/create_plan_prompt.rs +++ b/api/libs/agents/src/agents/modes/create_plan_prompt.rs @@ -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** diff --git a/api/libs/agents/src/tools/categories/planning_tools/create_plan_investigative.rs b/api/libs/agents/src/tools/categories/planning_tools/create_plan_investigative.rs index bef2cc4aa..d4e612fff 100644 --- a/api/libs/agents/src/tools/categories/planning_tools/create_plan_investigative.rs +++ b/api/libs/agents/src/tools/categories/planning_tools/create_plan_investigative.rs @@ -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, + 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 = 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 = 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::>().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 } diff --git a/api/libs/agents/src/tools/categories/planning_tools/create_plan_straightforward.rs b/api/libs/agents/src/tools/categories/planning_tools/create_plan_straightforward.rs index 6f6d0f020..366c65f2f 100644 --- a/api/libs/agents/src/tools/categories/planning_tools/create_plan_straightforward.rs +++ b/api/libs/agents/src/tools/categories/planning_tools/create_plan_straightforward.rs @@ -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, + 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 = 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 = 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::>() - .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 } diff --git a/api/libs/agents/src/tools/categories/planning_tools/helpers/mod.rs b/api/libs/agents/src/tools/categories/planning_tools/helpers/mod.rs new file mode 100644 index 000000000..bd526e4b2 --- /dev/null +++ b/api/libs/agents/src/tools/categories/planning_tools/helpers/mod.rs @@ -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; diff --git a/api/libs/agents/src/tools/categories/planning_tools/helpers/todo_generator.rs b/api/libs/agents/src/tools/categories/planning_tools/helpers/todo_generator.rs new file mode 100644 index 000000000..88fc2cd58 --- /dev/null +++ b/api/libs/agents/src/tools/categories/planning_tools/helpers/todo_generator.rs @@ -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` 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> { + 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 = 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 = 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) +} \ No newline at end of file diff --git a/api/libs/agents/src/tools/categories/planning_tools/mod.rs b/api/libs/agents/src/tools/categories/planning_tools/mod.rs index 8fbc57a8e..f98b831af 100644 --- a/api/libs/agents/src/tools/categories/planning_tools/mod.rs +++ b/api/libs/agents/src/tools/categories/planning_tools/mod.rs @@ -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::*; \ No newline at end of file +pub use review_plan::*; +pub use helpers::*; \ No newline at end of file