use crate::types::{Layer2Result, TaskId};
use crate::workflow_engine::{Dag, Node};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum DecompositionStrategy {
Sequential,
Parallel,
Hybrid,
Hierarchical,
}
impl Default for DecompositionStrategy {
fn default() -> Self {
Self::Hybrid
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SubTask {
pub id: String,
pub name: String,
pub description: String,
pub priority: u32,
pub dependencies: Vec<String>,
pub estimated_complexity: u32,
pub tool: Option<String>,
pub tool_args: Option<serde_json::Value>,
pub validation_criteria: Vec<String>,
pub fallback: Option<Box<SubTask>>,
}
impl SubTask {
pub fn new(
id: impl Into<String>,
name: impl Into<String>,
description: impl Into<String>,
) -> Self {
Self {
id: id.into(),
name: name.into(),
description: description.into(),
priority: 0,
dependencies: Vec::new(),
estimated_complexity: 5,
tool: None,
tool_args: None,
validation_criteria: Vec::new(),
fallback: None,
}
}
pub fn with_priority(mut self, priority: u32) -> Self {
self.priority = priority;
self
}
pub fn with_dependency(mut self, dep_id: impl Into<String>) -> Self {
self.dependencies.push(dep_id.into());
self
}
pub fn with_tool(mut self, tool: impl Into<String>, args: serde_json::Value) -> Self {
self.tool = Some(tool.into());
self.tool_args = Some(args);
self
}
pub fn with_validation(mut self, criteria: impl Into<String>) -> Self {
self.validation_criteria.push(criteria.into());
self
}
pub fn with_fallback(mut self, fallback: SubTask) -> Self {
self.fallback = Some(Box::new(fallback));
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionPlan {
pub id: String,
pub original_task: String,
pub strategy: DecompositionStrategy,
pub subtasks: Vec<SubTask>,
pub execution_order: Vec<String>,
pub estimated_steps: u32,
pub risk_level: RiskLevel,
pub created_at: chrono::DateTime<chrono::Utc>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub enum RiskLevel {
Low,
Medium,
High,
Critical,
}
impl Default for RiskLevel {
fn default() -> Self {
Self::Medium
}
}
impl ExecutionPlan {
pub fn new(original_task: impl Into<String>) -> Self {
Self {
id: TaskId::new().to_string(),
original_task: original_task.into(),
strategy: DecompositionStrategy::default(),
subtasks: Vec::new(),
execution_order: Vec::new(),
estimated_steps: 1,
risk_level: RiskLevel::default(),
created_at: chrono::Utc::now(),
}
}
pub fn add_subtask(&mut self, subtask: SubTask) -> &mut Self {
self.subtasks.push(subtask);
self
}
pub fn compute_execution_order(&mut self) -> Layer2Result<()> {
let mut dag = Dag::new();
for subtask in &self.subtasks {
let node = Node::new(&subtask.id, &subtask.name);
dag.add_node(node)?;
}
for subtask in &self.subtasks {
for dep in &subtask.dependencies {
dag.add_edge(dep, &subtask.id)?;
}
}
if dag.has_cycle() {
return Err(anyhow::anyhow!(
"Circular dependency detected in execution plan"
));
}
self.execution_order = dag.topological_sort()?;
self.estimated_steps = self.subtasks.len() as u32;
Ok(())
}
pub fn to_dag(&self) -> Layer2Result<Dag> {
let mut dag = Dag::new();
for subtask in &self.subtasks {
let mut node = Node::new(&subtask.id, &subtask.name);
node.config = serde_json::json!({
"description": subtask.description,
"tool": subtask.tool,
"tool_args": subtask.tool_args,
});
dag.add_node(node)?;
}
for subtask in &self.subtasks {
for dep in &subtask.dependencies {
dag.add_edge(dep, &subtask.id)?;
}
}
Ok(dag)
}
}
pub struct TaskDecomposer {
strategy: DecompositionStrategy,
max_depth: u32,
min_granularity: u32,
}
impl Default for TaskDecomposer {
fn default() -> Self {
Self {
strategy: DecompositionStrategy::Hybrid,
max_depth: 3,
min_granularity: 1,
}
}
}
impl TaskDecomposer {
pub fn new() -> Self {
Self::default()
}
pub fn with_strategy(mut self, strategy: DecompositionStrategy) -> Self {
self.strategy = strategy;
self
}
pub fn with_max_depth(mut self, depth: u32) -> Self {
self.max_depth = depth;
self
}
pub fn decompose(&self, task: &str) -> Layer2Result<ExecutionPlan> {
let mut plan = ExecutionPlan::new(task);
plan.strategy = self.strategy;
let complexity = self.analyze_complexity(task);
plan.risk_level = self.estimate_risk(task, complexity);
let subtasks = match self.strategy {
DecompositionStrategy::Sequential => self.decompose_sequential(task),
DecompositionStrategy::Parallel => self.decompose_parallel(task),
DecompositionStrategy::Hierarchical => self.decompose_hierarchical(task, 0),
DecompositionStrategy::Hybrid => self.decompose_hybrid(task),
};
for subtask in subtasks {
plan.add_subtask(subtask);
}
plan.compute_execution_order()?;
Ok(plan)
}
fn analyze_complexity(&self, task: &str) -> u32 {
let mut complexity = 1u32;
let task_lower = task.to_lowercase();
if task_lower.contains("implement") || task_lower.contains("create") {
complexity += 2;
}
if task_lower.contains("refactor") || task_lower.contains("rewrite") {
complexity += 2;
}
if task_lower.contains("integrate") || task_lower.contains("connect") {
complexity += 1;
}
if task_lower.contains("test") || task_lower.contains("verify") {
complexity += 1;
}
if task_lower.contains("and") || task_lower.contains("then") {
complexity += 1;
}
if task_lower.contains("multiple") || task_lower.contains("several") {
complexity += 1;
}
let word_count = task.split_whitespace().count();
if word_count > 20 {
complexity += 1;
}
if word_count > 50 {
complexity += 1;
}
complexity.min(10)
}
fn estimate_risk(&self, task: &str, complexity: u32) -> RiskLevel {
let task_lower = task.to_lowercase();
if task_lower.contains("delete")
|| task_lower.contains("remove")
|| task_lower.contains("drop")
{
return RiskLevel::Critical;
}
if task_lower.contains("production")
|| task_lower.contains("live")
|| task_lower.contains("deploy")
{
return RiskLevel::High;
}
if task_lower.contains("database") || task_lower.contains("migration") {
return RiskLevel::High;
}
match complexity {
1..=3 => RiskLevel::Low,
4..=6 => RiskLevel::Medium,
7..=8 => RiskLevel::High,
_ => RiskLevel::Critical,
}
}
fn decompose_sequential(&self, task: &str) -> Vec<SubTask> {
let steps = self.extract_steps(task);
let mut subtasks = Vec::new();
let mut prev_id: Option<String> = None;
for (i, step) in steps.into_iter().enumerate() {
let id = format!("step_{}", i + 1);
let mut subtask = SubTask::new(&id, format!("Step {}", i + 1), step);
subtask.priority = i as u32;
if let Some(prev) = prev_id {
subtask = subtask.with_dependency(prev);
}
prev_id = Some(id);
subtasks.push(subtask);
}
if subtasks.is_empty() {
subtasks.push(SubTask::new("step_1", "Execute task", task));
}
subtasks
}
fn decompose_parallel(&self, task: &str) -> Vec<SubTask> {
let parts = self.extract_parallel_parts(task);
let mut subtasks = Vec::new();
for (i, part) in parts.into_iter().enumerate() {
let id = format!("parallel_{}", i + 1);
let subtask = SubTask::new(&id, format!("Task {}", i + 1), part);
subtasks.push(subtask);
}
if subtasks.is_empty() {
subtasks.push(SubTask::new("parallel_1", "Execute task", task));
}
subtasks
}
fn decompose_hierarchical(&self, task: &str, depth: u32) -> Vec<SubTask> {
if depth >= self.max_depth {
return vec![SubTask::new(&format!("leaf_{}", depth), "Execute", task)];
}
let main_steps = self.extract_steps(task);
let mut subtasks = Vec::new();
for (i, step) in main_steps.into_iter().enumerate() {
let id = format!("h{}_{}", depth, i + 1);
let mut subtask = SubTask::new(&id, format!("Phase {}", i + 1), step.clone());
subtask.estimated_complexity = self.analyze_complexity(&step);
if subtask.estimated_complexity > 5 && depth < self.max_depth - 1 {
let sub_subtasks = self.decompose_hierarchical(&step, depth + 1);
for (j, sub_sub) in sub_subtasks.into_iter().enumerate() {
let mut sub_sub_id = sub_sub;
sub_sub_id.id = format!("{}_{}", id, j + 1);
sub_sub_id.dependencies.push(id.clone());
subtasks.push(sub_sub_id);
}
}
subtasks.push(subtask);
}
subtasks
}
fn decompose_hybrid(&self, task: &str) -> Vec<SubTask> {
let complexity = self.analyze_complexity(task);
if complexity <= 3 {
vec![SubTask::new("execute", "Execute task", task)]
} else if complexity <= 6 {
self.decompose_sequential(task)
} else {
self.decompose_hierarchical(task, 0)
}
}
fn extract_steps(&self, task: &str) -> Vec<String> {
let mut steps = Vec::new();
let sentences: Vec<&str> = task
.split(&['.', ';', '\n'][..])
.map(|s| s.trim())
.filter(|s| !s.is_empty())
.collect();
if sentences.len() > 1 {
steps = sentences.into_iter().map(|s| s.to_string()).collect();
} else {
let mut then_parts: Vec<&str> = task.split("and then").collect();
if then_parts.len() == 1 {
then_parts = task.split("then").collect();
}
if then_parts.len() == 1 {
then_parts = task.split("after that").collect();
}
if then_parts.len() > 1 {
steps = then_parts
.into_iter()
.map(|s| s.trim().to_string())
.collect();
} else {
steps.push(task.to_string());
}
}
steps
}
fn extract_parallel_parts(&self, task: &str) -> Vec<String> {
let mut parts: Vec<&str> = task.split(", and ").collect();
if parts.len() == 1 {
parts = task.split(" and ").collect();
}
if parts.len() == 1 {
parts = task.split(", ").collect();
}
let parts: Vec<&str> = parts
.into_iter()
.map(|s| s.trim())
.filter(|s| !s.is_empty() && s.len() > 3)
.collect();
if parts.len() > 1 {
parts.into_iter().map(|s| s.to_string()).collect()
} else {
vec![task.to_string()]
}
}
}
#[derive(Debug, Clone)]
pub struct PlanResult {
pub plan: ExecutionPlan,
pub quality_score: u32,
pub suggestions: Vec<String>,
}
impl PlanResult {
pub fn new(plan: ExecutionPlan) -> Self {
let quality_score = Self::calculate_quality(&plan);
let suggestions = Self::generate_suggestions(&plan);
Self {
plan,
quality_score,
suggestions,
}
}
fn calculate_quality(plan: &ExecutionPlan) -> u32 {
let mut score = 100u32;
if plan.subtasks.is_empty() {
score = 0;
} else if plan.subtasks.len() == 1 {
score -= 20; }
if plan.execution_order.len() != plan.subtasks.len() {
score -= 30;
}
let has_validation = plan
.subtasks
.iter()
.any(|s| !s.validation_criteria.is_empty());
if !has_validation {
score -= 10;
}
let has_fallback = plan.subtasks.iter().any(|s| s.fallback.is_some());
if !has_fallback && plan.risk_level >= RiskLevel::High {
score -= 15;
}
score
}
fn generate_suggestions(plan: &ExecutionPlan) -> Vec<String> {
let mut suggestions = Vec::new();
if plan.subtasks.len() == 1 {
suggestions.push("Consider breaking down the task into smaller subtasks".to_string());
}
if plan.risk_level >= RiskLevel::High {
suggestions.push("High-risk task: consider adding validation steps".to_string());
}
let has_fallback = plan.subtasks.iter().any(|s| s.fallback.is_some());
if !has_fallback && !plan.subtasks.is_empty() {
suggestions
.push("Consider adding fallback strategies for critical subtasks".to_string());
}
suggestions
}
}
#[cfg(test)]
mod tests {
use super::{
DecompositionStrategy, ExecutionPlan, PlanResult, RiskLevel, SubTask, TaskDecomposer,
};
#[test]
fn test_subtask_creation() {
let subtask = SubTask::new("test_1", "Test", "Test subtask");
assert_eq!(subtask.id, "test_1");
assert_eq!(subtask.name, "Test");
}
#[test]
fn test_subtask_with_dependencies() {
let subtask = SubTask::new("test_2", "Test", "Test").with_dependency("test_1");
assert_eq!(subtask.dependencies.len(), 1);
}
#[test]
fn test_execution_plan_creation() {
let plan = ExecutionPlan::new("Test task");
assert!(!plan.original_task.is_empty());
assert!(plan.subtasks.is_empty());
}
#[test]
fn test_task_decomposer() {
let decomposer = TaskDecomposer::new();
let plan = decomposer
.decompose("Create a file and write some content")
.unwrap();
assert!(!plan.subtasks.is_empty());
assert!(!plan.execution_order.is_empty());
}
#[test]
fn test_complexity_analysis() {
let decomposer = TaskDecomposer::new();
let simple = decomposer.analyze_complexity("Read a file");
assert!(simple <= 3);
let complex = decomposer.analyze_complexity(
"Implement a complete authentication system with OAuth2 integration",
);
assert!(complex > 3);
}
#[test]
fn test_risk_estimation() {
let decomposer = TaskDecomposer::new();
let low = decomposer.estimate_risk("Read a file", 2);
assert_eq!(low, RiskLevel::Low);
let critical = decomposer.estimate_risk("Delete the production database", 5);
assert_eq!(critical, RiskLevel::Critical);
}
#[test]
fn test_sequential_decomposition() {
let decomposer = TaskDecomposer::new().with_strategy(DecompositionStrategy::Sequential);
let plan = decomposer
.decompose("First step. Second step. Third step.")
.unwrap();
assert!(plan.subtasks.len() >= 3);
for i in 1..plan.subtasks.len() {
assert!(plan.subtasks[i]
.dependencies
.contains(&plan.subtasks[i - 1].id));
}
}
#[test]
fn test_parallel_decomposition() {
let decomposer = TaskDecomposer::new().with_strategy(DecompositionStrategy::Parallel);
let plan = decomposer
.decompose("Task A and Task B and Task C")
.unwrap();
assert!(plan.subtasks.len() >= 2);
let has_deps: bool = plan.subtasks.iter().any(|s| !s.dependencies.is_empty());
assert!(!has_deps);
}
#[test]
fn test_plan_result_quality() {
let mut plan = ExecutionPlan::new("Test task");
plan.add_subtask(SubTask::new("s1", "Step 1", "First step"));
plan.add_subtask(SubTask::new("s2", "Step 2", "Second step").with_dependency("s1"));
plan.compute_execution_order().unwrap();
let result = PlanResult::new(plan);
assert!(result.quality_score > 0);
}
#[test]
fn test_dag_conversion() {
let mut plan = ExecutionPlan::new("Test task");
plan.add_subtask(SubTask::new("s1", "Step 1", "First step"));
plan.add_subtask(SubTask::new("s2", "Step 2", "Second step").with_dependency("s1"));
plan.compute_execution_order().unwrap();
let dag_result = plan.to_dag();
assert!(dag_result.is_ok());
}
}