ai_agent/
execution.rs

1//! Execution Engine - AI-powered task execution
2
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokio::sync::Mutex;
6use crate::types::{ExecutionContext, ExecutionDecision, ExecutionResult, ActionType, StepResult};
7use crate::models::LanguageModel;
8use crate::tools::{ToolRegistry, ToolCall, ToolArgs};
9use crate::errors::AgentError;
10
11/// Execution engine for running tasks
12pub struct ExecutionEngine {
13    model: Arc<dyn LanguageModel>,
14    tools: Arc<Mutex<ToolRegistry>>,
15    config: crate::config::ExecutionConfig,
16}
17
18impl ExecutionEngine {
19    pub fn new(model: Arc<dyn LanguageModel>, tools: Arc<Mutex<ToolRegistry>>, config: crate::config::ExecutionConfig) -> Self {
20        Self { model, tools, config }
21    }
22
23    /// Execute a task plan
24    pub async fn execute(&self, task_id: &str, plan: crate::types::TaskPlan) -> Result<ExecutionResult, AgentError> {
25        let start_time = std::time::Instant::now();
26        let mut context = ExecutionContext {
27            task_id: task_id.to_string(),
28            plan,
29            current_step: 0,
30            results: vec![],
31            variables: HashMap::new(),
32        };
33
34        loop {
35            // Check step limit
36            if context.current_step >= self.config.max_steps {
37                return Ok(ExecutionResult {
38                    success: false,
39                    summary: "Maximum steps exceeded".to_string(),
40                    details: "Task exceeded maximum allowed steps".to_string(),
41                    execution_time: start_time.elapsed().as_secs(),
42                });
43            }
44
45            // AI decides next action
46            let decision = self.make_execution_decision(&context).await?;
47
48            let action_taken = decision.action_type.clone();
49            let reasoning_used = decision.reasoning.clone();
50            let confidence_used = decision.confidence;
51
52            match action_taken {
53                ActionType::UseTool { ref tool_name, ref arguments } => {
54                    let tool_result = self.execute_tool(tool_name, arguments.clone()).await?;
55
56                    context.results.push(crate::types::ExecutionStep {
57                        step_number: context.current_step,
58                        action: crate::types::Action {
59                            action_type: action_taken.clone(),
60                            reasoning: reasoning_used.clone(),
61                            confidence: confidence_used,
62                        },
63                        result: Some(tool_result.clone()),
64                        timestamp: chrono::Utc::now(),
65                    });
66
67                    if !tool_result.success {
68                        return Ok(ExecutionResult {
69                            success: false,
70                            summary: "Tool execution failed".to_string(),
71                            details: tool_result.error.unwrap_or_default(),
72                            execution_time: start_time.elapsed().as_secs(),
73                        });
74                    }
75                }
76
77                ActionType::Complete { ref summary } => {
78                    return Ok(ExecutionResult {
79                        success: true,
80                        summary: summary.clone(),
81                        details: self.format_execution_details(&context),
82                        execution_time: start_time.elapsed().as_secs(),
83                    });
84                }
85
86                ActionType::Think { reasoning: _ } => {
87                    context.results.push(crate::types::ExecutionStep {
88                        step_number: context.current_step,
89                        action: crate::types::Action {
90                            action_type: action_taken.clone(),
91                            reasoning: reasoning_used.clone(),
92                            confidence: confidence_used,
93                        },
94                        result: None,
95                        timestamp: chrono::Utc::now(),
96                    });
97                }
98
99                ActionType::AskClarification { ref question } => {
100                    return Ok(ExecutionResult {
101                        success: false,
102                        summary: "Clarification needed".to_string(),
103                        details: format!("Question: {}", question),
104                        execution_time: start_time.elapsed().as_secs(),
105                    });
106                }
107            }
108
109            context.current_step += 1;
110        }
111    }
112
113    async fn make_execution_decision(&self, context: &ExecutionContext) -> Result<ExecutionDecision, AgentError> {
114        let prompt = self.build_execution_prompt(context);
115        let response = self.model.complete_with_tools(&prompt, &self.get_tool_definitions()).await?;
116
117        self.parse_execution_decision(&response.content)
118    }
119
120    fn build_execution_prompt(&self, context: &ExecutionContext) -> String {
121        format!(
122            "You are working on this task: {}
123
124Current progress: Step {} of estimated {}
125
126You have these tools available: {}
127
128You have complete freedom to decide your next action. You can:
129- Use a tool to accomplish part of the task
130- Think through the next steps
131- Complete the task if you believe it's done
132- Ask for clarification if needed
133
134Decide your next action and respond in this format:
135Action: [UseTool/Think/Complete/AskClarification]
136Tool: [tool name if using tool]
137Arguments: [tool arguments if using tool]
138Reasoning: [your reasoning]
139Confidence: [0.0-1.0]",
140            context.plan.understanding,
141            context.current_step + 1,
142            context.plan.estimated_steps.unwrap_or(0),
143            self.get_tool_descriptions()
144        )
145    }
146
147    fn get_tool_descriptions(&self) -> String {
148        let tools = self.tools.blocking_lock();
149        let mut descriptions = Vec::new();
150
151        for tool_name in tools.get_tool_names() {
152            if let Some(tool) = tools.get_tool(&tool_name) {
153                descriptions.push(format!("{}: {}", tool.name(), tool.description()));
154            }
155        }
156
157        descriptions.join("\n")
158    }
159
160    fn get_tool_definitions(&self) -> Vec<crate::models::ToolDefinition> {
161        let tools = self.tools.blocking_lock();
162        let mut definitions = Vec::new();
163
164        for tool_name in tools.get_tool_names() {
165            if let Some(tool) = tools.get_tool(&tool_name) {
166                let _parameters: Vec<serde_json::Value> = tool.parameters()
167                    .into_iter()
168                    .map(|p| {
169                        serde_json::json!({
170                            "type": "string",
171                            "description": p.description
172                        })
173                    })
174                    .collect();
175
176                definitions.push(crate::models::ToolDefinition {
177                    name: tool.name().to_string(),
178                    description: tool.description().to_string(),
179                    parameters: serde_json::json!({
180                        "type": "object",
181                        "properties": {},
182                        "required": []
183                    }),
184                });
185            }
186        }
187
188        definitions
189    }
190
191    async fn execute_tool(&self, tool_name: &str, arguments: HashMap<String, serde_json::Value>) -> Result<StepResult, AgentError> {
192        let tools = self.tools.lock().await;
193        let tool_call = ToolCall {
194            name: tool_name.to_string(),
195            args: ToolArgs::from_map(arguments),
196        };
197
198        match tools.execute(&tool_call).await {
199            Ok(result) => Ok(StepResult {
200                success: result.success,
201                output: serde_json::Value::String(result.content),
202                error: result.error,
203                execution_time: 0, // Simplified implementation
204            }),
205            Err(error) => Ok(StepResult {
206                success: false,
207                output: serde_json::Value::Null,
208                error: Some(error.to_string()),
209                execution_time: 0,
210            }),
211        }
212    }
213
214    fn format_execution_details(&self, context: &ExecutionContext) -> String {
215        format!(
216            "Task completed in {} steps. All objectives achieved.",
217            context.current_step
218        )
219    }
220
221    fn parse_execution_decision(&self, response: &str) -> Result<ExecutionDecision, AgentError> {
222        // Simple parsing logic for demonstration
223        let lines: Vec<&str> = response.lines().collect();
224        let mut action_type = ActionType::Think {
225            reasoning: "Continuing analysis".to_string(),
226        };
227        let mut reasoning = "AI reasoning".to_string();
228        let mut confidence = 0.7;
229
230        for line in lines {
231            if line.to_lowercase().starts_with("action:") {
232                let action = line[7..].trim().to_lowercase();
233                action_type = match action.as_str() {
234                    "usetool" => ActionType::UseTool {
235                        tool_name: "unknown".to_string(),
236                        arguments: HashMap::new(),
237                    },
238                    "complete" => ActionType::Complete {
239                        summary: "Task completed".to_string(),
240                    },
241                    "think" => ActionType::Think {
242                        reasoning: "Thinking through next steps".to_string(),
243                    },
244                    "askclarification" => ActionType::AskClarification {
245                        question: "Need clarification".to_string(),
246                    },
247                    _ => action_type,
248                };
249            } else if line.to_lowercase().starts_with("reasoning:") {
250                reasoning = line[10..].trim().to_string();
251            } else if line.to_lowercase().starts_with("confidence:") {
252                if let Ok(conf) = line[11..].trim().parse::<f32>() {
253                    confidence = conf;
254                }
255            }
256        }
257
258        Ok(ExecutionDecision {
259            action_type,
260            reasoning,
261            confidence,
262        })
263    }
264}