ai_agents_reasoning/
plan.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Serialize, Deserialize)]
4pub struct Plan {
5    pub id: String,
6    pub goal: String,
7    pub steps: Vec<PlanStep>,
8    pub status: PlanStatus,
9    pub created_at: chrono::DateTime<chrono::Utc>,
10}
11
12impl Plan {
13    pub fn new(goal: impl Into<String>) -> Self {
14        Self {
15            id: uuid::Uuid::new_v4().to_string(),
16            goal: goal.into(),
17            steps: Vec::new(),
18            status: PlanStatus::Pending,
19            created_at: chrono::Utc::now(),
20        }
21    }
22
23    pub fn with_steps(mut self, steps: Vec<PlanStep>) -> Self {
24        self.steps = steps;
25        self
26    }
27
28    pub fn add_step(&mut self, step: PlanStep) {
29        self.steps.push(step);
30    }
31
32    pub fn is_complete(&self) -> bool {
33        matches!(self.status, PlanStatus::Completed)
34    }
35
36    pub fn is_failed(&self) -> bool {
37        matches!(self.status, PlanStatus::Failed { .. })
38    }
39
40    pub fn pending_steps(&self) -> impl Iterator<Item = &PlanStep> {
41        self.steps.iter().filter(|s| s.status.is_pending())
42    }
43
44    pub fn completed_steps(&self) -> impl Iterator<Item = &PlanStep> {
45        self.steps.iter().filter(|s| s.status.is_completed())
46    }
47
48    pub fn next_executable_step(&self) -> Option<&PlanStep> {
49        self.steps.iter().find(|s| {
50            s.status.is_pending() && s.dependencies.iter().all(|dep| self.is_step_completed(dep))
51        })
52    }
53
54    fn is_step_completed(&self, step_id: &str) -> bool {
55        self.steps
56            .iter()
57            .find(|s| s.id == step_id)
58            .map(|s| s.status.is_completed())
59            .unwrap_or(false)
60    }
61}
62
63#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct PlanStep {
65    pub id: String,
66    pub description: String,
67    pub action: PlanAction,
68    #[serde(default)]
69    pub dependencies: Vec<String>,
70    #[serde(default)]
71    pub status: StepStatus,
72    #[serde(skip_serializing_if = "Option::is_none")]
73    pub result: Option<serde_json::Value>,
74}
75
76impl PlanStep {
77    pub fn new(description: impl Into<String>, action: PlanAction) -> Self {
78        Self {
79            id: uuid::Uuid::new_v4().to_string(),
80            description: description.into(),
81            action,
82            dependencies: Vec::new(),
83            status: StepStatus::Pending,
84            result: None,
85        }
86    }
87
88    pub fn with_id(mut self, id: impl Into<String>) -> Self {
89        self.id = id.into();
90        self
91    }
92
93    pub fn with_dependencies(mut self, deps: Vec<String>) -> Self {
94        self.dependencies = deps;
95        self
96    }
97
98    pub fn mark_running(&mut self) {
99        self.status = StepStatus::Running;
100    }
101
102    pub fn mark_completed(&mut self, result: Option<serde_json::Value>) {
103        self.status = StepStatus::Completed;
104        self.result = result;
105    }
106
107    pub fn mark_failed(&mut self, error: impl Into<String>) {
108        self.status = StepStatus::Failed {
109            error: error.into(),
110        };
111    }
112
113    pub fn mark_skipped(&mut self, reason: impl Into<String>) {
114        self.status = StepStatus::Skipped {
115            reason: reason.into(),
116        };
117    }
118}
119
120#[derive(Debug, Clone, Serialize, Deserialize)]
121#[serde(tag = "type", rename_all = "snake_case")]
122pub enum PlanAction {
123    Tool {
124        tool: String,
125        #[serde(default)]
126        args: serde_json::Value,
127    },
128    Skill {
129        skill: String,
130    },
131    Think {
132        prompt: String,
133    },
134    Respond {
135        template: String,
136    },
137}
138
139impl PlanAction {
140    pub fn tool(name: impl Into<String>, args: serde_json::Value) -> Self {
141        PlanAction::Tool {
142            tool: name.into(),
143            args,
144        }
145    }
146
147    pub fn skill(name: impl Into<String>) -> Self {
148        PlanAction::Skill { skill: name.into() }
149    }
150
151    pub fn think(prompt: impl Into<String>) -> Self {
152        PlanAction::Think {
153            prompt: prompt.into(),
154        }
155    }
156
157    pub fn respond(template: impl Into<String>) -> Self {
158        PlanAction::Respond {
159            template: template.into(),
160        }
161    }
162
163    pub fn is_tool(&self) -> bool {
164        matches!(self, PlanAction::Tool { .. })
165    }
166
167    pub fn is_skill(&self) -> bool {
168        matches!(self, PlanAction::Skill { .. })
169    }
170
171    pub fn is_think(&self) -> bool {
172        matches!(self, PlanAction::Think { .. })
173    }
174
175    pub fn is_respond(&self) -> bool {
176        matches!(self, PlanAction::Respond { .. })
177    }
178}
179
180#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
181#[serde(tag = "status", rename_all = "snake_case")]
182pub enum PlanStatus {
183    #[default]
184    Pending,
185    InProgress,
186    Completed,
187    Failed {
188        error: String,
189    },
190    Replanned {
191        reason: String,
192        new_plan_id: String,
193    },
194}
195
196impl PlanStatus {
197    pub fn is_pending(&self) -> bool {
198        matches!(self, PlanStatus::Pending)
199    }
200
201    pub fn is_in_progress(&self) -> bool {
202        matches!(self, PlanStatus::InProgress)
203    }
204
205    pub fn is_completed(&self) -> bool {
206        matches!(self, PlanStatus::Completed)
207    }
208
209    pub fn is_failed(&self) -> bool {
210        matches!(self, PlanStatus::Failed { .. })
211    }
212
213    pub fn is_replanned(&self) -> bool {
214        matches!(self, PlanStatus::Replanned { .. })
215    }
216
217    pub fn is_terminal(&self) -> bool {
218        matches!(
219            self,
220            PlanStatus::Completed | PlanStatus::Failed { .. } | PlanStatus::Replanned { .. }
221        )
222    }
223}
224
225#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
226#[serde(tag = "status", rename_all = "snake_case")]
227pub enum StepStatus {
228    #[default]
229    Pending,
230    Running,
231    Completed,
232    Failed {
233        error: String,
234    },
235    Skipped {
236        reason: String,
237    },
238}
239
240impl StepStatus {
241    pub fn is_pending(&self) -> bool {
242        matches!(self, StepStatus::Pending)
243    }
244
245    pub fn is_running(&self) -> bool {
246        matches!(self, StepStatus::Running)
247    }
248
249    pub fn is_completed(&self) -> bool {
250        matches!(self, StepStatus::Completed)
251    }
252
253    pub fn is_failed(&self) -> bool {
254        matches!(self, StepStatus::Failed { .. })
255    }
256
257    pub fn is_skipped(&self) -> bool {
258        matches!(self, StepStatus::Skipped { .. })
259    }
260
261    pub fn is_terminal(&self) -> bool {
262        matches!(
263            self,
264            StepStatus::Completed | StepStatus::Failed { .. } | StepStatus::Skipped { .. }
265        )
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    #[test]
274    fn test_plan_creation() {
275        let plan = Plan::new("Test goal");
276        assert_eq!(plan.goal, "Test goal");
277        assert!(plan.steps.is_empty());
278        assert!(plan.status.is_pending());
279    }
280
281    #[test]
282    fn test_plan_with_steps() {
283        let step1 = PlanStep::new("Step 1", PlanAction::think("Think about it"));
284        let step2 = PlanStep::new("Step 2", PlanAction::respond("Respond"));
285
286        let plan = Plan::new("Multi-step goal").with_steps(vec![step1, step2]);
287        assert_eq!(plan.steps.len(), 2);
288    }
289
290    #[test]
291    fn test_plan_step_dependencies() {
292        let step1 = PlanStep::new("Step 1", PlanAction::think("Think"))
293            .with_id("step1")
294            .with_dependencies(vec![]);
295
296        let step2 = PlanStep::new("Step 2", PlanAction::respond("Respond"))
297            .with_id("step2")
298            .with_dependencies(vec!["step1".to_string()]);
299
300        let mut plan = Plan::new("Goal").with_steps(vec![step1, step2]);
301
302        // First executable should be step1
303        let next = plan.next_executable_step().unwrap();
304        assert_eq!(next.id, "step1");
305
306        // Mark step1 as completed
307        plan.steps[0].mark_completed(None);
308
309        // Now step2 should be executable
310        let next = plan.next_executable_step().unwrap();
311        assert_eq!(next.id, "step2");
312    }
313
314    #[test]
315    fn test_plan_action_types() {
316        let tool = PlanAction::tool("search", serde_json::json!({"query": "test"}));
317        assert!(tool.is_tool());
318
319        let skill = PlanAction::skill("greeting");
320        assert!(skill.is_skill());
321
322        let think = PlanAction::think("Consider the options");
323        assert!(think.is_think());
324
325        let respond = PlanAction::respond("Final answer: {{ result }}");
326        assert!(respond.is_respond());
327    }
328
329    #[test]
330    fn test_plan_action_serde() {
331        let action = PlanAction::tool("http", serde_json::json!({"url": "https://example.com"}));
332        let json = serde_json::to_string(&action).unwrap();
333        let parsed: PlanAction = serde_json::from_str(&json).unwrap();
334        assert!(parsed.is_tool());
335    }
336
337    #[test]
338    fn test_plan_status() {
339        assert!(PlanStatus::Pending.is_pending());
340        assert!(!PlanStatus::Pending.is_terminal());
341
342        assert!(PlanStatus::InProgress.is_in_progress());
343        assert!(!PlanStatus::InProgress.is_terminal());
344
345        assert!(PlanStatus::Completed.is_completed());
346        assert!(PlanStatus::Completed.is_terminal());
347
348        let failed = PlanStatus::Failed {
349            error: "Error".to_string(),
350        };
351        assert!(failed.is_failed());
352        assert!(failed.is_terminal());
353
354        let replanned = PlanStatus::Replanned {
355            reason: "Better approach".to_string(),
356            new_plan_id: "plan2".to_string(),
357        };
358        assert!(replanned.is_replanned());
359        assert!(replanned.is_terminal());
360    }
361
362    #[test]
363    fn test_step_status() {
364        assert!(StepStatus::Pending.is_pending());
365        assert!(!StepStatus::Pending.is_terminal());
366
367        assert!(StepStatus::Running.is_running());
368        assert!(!StepStatus::Running.is_terminal());
369
370        assert!(StepStatus::Completed.is_completed());
371        assert!(StepStatus::Completed.is_terminal());
372
373        let failed = StepStatus::Failed {
374            error: "Error".to_string(),
375        };
376        assert!(failed.is_failed());
377        assert!(failed.is_terminal());
378
379        let skipped = StepStatus::Skipped {
380            reason: "Not needed".to_string(),
381        };
382        assert!(skipped.is_skipped());
383        assert!(skipped.is_terminal());
384    }
385
386    #[test]
387    fn test_plan_step_state_transitions() {
388        let mut step = PlanStep::new("Test step", PlanAction::think("Think"));
389
390        assert!(step.status.is_pending());
391
392        step.mark_running();
393        assert!(step.status.is_running());
394
395        step.mark_completed(Some(serde_json::json!({"answer": 42})));
396        assert!(step.status.is_completed());
397        assert!(step.result.is_some());
398    }
399
400    #[test]
401    fn test_plan_step_failure() {
402        let mut step = PlanStep::new("Test step", PlanAction::tool("http", serde_json::json!({})));
403
404        step.mark_running();
405        step.mark_failed("Connection timeout");
406
407        assert!(step.status.is_failed());
408        if let StepStatus::Failed { error } = &step.status {
409            assert_eq!(error, "Connection timeout");
410        }
411    }
412}
ai_agents_reasoning/plan.rs

ai_agents_reasoning/
plan.rs
