Skip to main content

agent_base/engine/
reflexion.rs

1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4
5use crate::types::AgentResult;
6
7#[cfg(test)]
8use serde_json::json;
9
10#[async_trait]
11pub trait ReflexionHandler: Send + Sync {
12    async fn reflect_on_failure(
13        &self,
14        failed_action: &str,
15        error: &str,
16        context: &str,
17    ) -> AgentResult<ReflectionResult>;
18
19    async fn generate_alternatives(
20        &self,
21        reflection: &ReflectionResult,
22    ) -> AgentResult<Vec<AlternativeAction>>;
23
24    async fn should_retry(
25        &self,
26        reflection: &ReflectionResult,
27        retry_count: usize,
28    ) -> AgentResult<bool>;
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct ReflectionResult {
33    pub analysis: String,
34    pub root_cause: String,
35    pub confidence: f32,
36    pub suggested_fixes: Vec<String>,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct AlternativeAction {
41    pub description: String,
42    pub payload: Value,
43    pub priority: u32,
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct ReflexionContext {
48    pub objective: String,
49    pub failed_step_id: String,
50    pub failed_step_description: String,
51    pub step_payload: Value,
52    pub error: String,
53    pub previous_steps: Vec<StepHistoryEntry>,
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct StepHistoryEntry {
58    pub step_id: String,
59    pub description: String,
60    pub success: bool,
61    pub output: Option<String>,
62}
63
64impl ReflectionResult {
65    pub fn new(
66        analysis: impl Into<String>,
67        root_cause: impl Into<String>,
68        confidence: f32,
69        suggested_fixes: Vec<String>,
70    ) -> Self {
71        Self {
72            analysis: analysis.into(),
73            root_cause: root_cause.into(),
74            confidence,
75            suggested_fixes,
76        }
77    }
78
79    pub fn is_confident(&self) -> bool {
80        self.confidence >= 0.7
81    }
82}
83
84impl AlternativeAction {
85    pub fn new(
86        description: impl Into<String>,
87        payload: Value,
88        priority: u32,
89    ) -> Self {
90        Self {
91            description: description.into(),
92            payload,
93            priority,
94        }
95    }
96}
97
98impl ReflexionContext {
99    pub fn from_step(
100        objective: &str,
101        step: &crate::types::PlanStep,
102        error: &str,
103    ) -> Self {
104        Self {
105            objective: objective.to_string(),
106            failed_step_id: step.id.clone(),
107            failed_step_description: step.description.clone(),
108            step_payload: step.payload.clone(),
109            error: error.to_string(),
110            previous_steps: Vec::new(),
111        }
112    }
113
114    pub fn with_previous_steps(mut self, steps: Vec<StepHistoryEntry>) -> Self {
115        self.previous_steps = steps;
116        self
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use crate::types::PlanStep;
124
125    #[test]
126    fn test_reflection_result_new() {
127        let result = ReflectionResult::new("analysis", "root_cause", 0.8, vec!["fix1".to_string()]);
128
129        assert_eq!(result.analysis, "analysis");
130        assert_eq!(result.root_cause, "root_cause");
131        assert_eq!(result.confidence, 0.8);
132        assert_eq!(result.suggested_fixes.len(), 1);
133    }
134
135    #[test]
136    fn test_reflection_result_is_confident() {
137        let high_confidence = ReflectionResult::new("", "", 0.9, vec![]);
138        assert!(high_confidence.is_confident());
139
140        let low_confidence = ReflectionResult::new("", "", 0.5, vec![]);
141        assert!(!low_confidence.is_confident());
142
143        let threshold = ReflectionResult::new("", "", 0.7, vec![]);
144        assert!(threshold.is_confident());
145    }
146
147    #[test]
148    fn test_alternative_action_new() {
149        let action = AlternativeAction::new(
150            "description",
151            json!({"type":"ssh_command","command":"ls"}),
152            1,
153        );
154
155        assert_eq!(action.description, "description");
156        assert_eq!(action.priority, 1);
157    }
158
159    #[test]
160    fn test_reflexion_context_from_step() {
161        let step = PlanStep::new(
162            "step-1",
163            "check disk",
164            json!({"type":"ssh_command","command":"df -h","host_id":"host1"}),
165        );
166
167        let context = ReflexionContext::from_step("objective", &step, "error");
168
169        assert_eq!(context.objective, "objective");
170        assert_eq!(context.failed_step_id, "step-1");
171        assert_eq!(context.failed_step_description, "check disk");
172        assert_eq!(context.step_payload, json!({"type":"ssh_command","command":"df -h","host_id":"host1"}));
173        assert_eq!(context.error, "error");
174        assert!(context.previous_steps.is_empty());
175    }
176
177    #[test]
178    fn test_reflexion_context_with_previous_steps() {
179        let context = ReflexionContext {
180            objective: "objective".to_string(),
181            failed_step_id: "step-2".to_string(),
182            failed_step_description: "description".to_string(),
183            step_payload: json!({"type":"test"}),
184            error: "error".to_string(),
185            previous_steps: Vec::new(),
186        }
187        .with_previous_steps(vec![StepHistoryEntry {
188            step_id: "step-1".to_string(),
189            description: "previous step".to_string(),
190            success: true,
191            output: Some("done".to_string()),
192        }]);
193
194        assert_eq!(context.previous_steps.len(), 1);
195        assert_eq!(context.previous_steps[0].step_id, "step-1");
196    }
197
198    #[test]
199    fn test_reflection_result_serialization() {
200        let result = ReflectionResult::new("analysis", "root_cause", 0.8, vec!["fix1".to_string()]);
201
202        let json = serde_json::to_string(&result).unwrap();
203        let deserialized: ReflectionResult = serde_json::from_str(&json).unwrap();
204
205        assert_eq!(deserialized.analysis, "analysis");
206        assert_eq!(deserialized.root_cause, "root_cause");
207        assert_eq!(deserialized.confidence, 0.8);
208    }
209
210    #[test]
211    fn test_step_history_entry_serialization() {
212        let entry = StepHistoryEntry {
213            step_id: "step-1".to_string(),
214            description: "description".to_string(),
215            success: true,
216            output: Some("output".to_string()),
217        };
218
219        let json = serde_json::to_string(&entry).unwrap();
220        let deserialized: StepHistoryEntry = serde_json::from_str(&json).unwrap();
221
222        assert_eq!(deserialized.step_id, "step-1");
223        assert!(deserialized.success);
224        assert_eq!(deserialized.output, Some("output".to_string()));
225    }
226}